import torch
import gradio as gr
from transformers import pipeline


pipe=pipeline("summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16, device="cpu")
#pipe=pipeline("summarization", model="facebook/bart-large-cnn", torch_dtype=torch.bfloat16, device="cpu")
from pydantic import BaseModel, PydanticUserError, ConfigDict
from pydantic import BaseModel, ConfigDict

class MyModel(BaseModel):
    request: 'starlette.requests.Request'
    model_config = ConfigDict(arbitrary_types_allowed=True)
from pydantic_core import core_schema
from starlette.requests import Request

def get_pydantic_core_schema(request_type, handler):
    return core_schema.any_schema()

Request.__get_pydantic_core_schema__ = get_pydantic_core_schema
#pipe = pipeline("summarization", model="sshleifer/distilbart-cnn-6-6")
text_summary=pipeline(task="summarization", model="sshleifer/distilbart-cnn-12-6", torch_dtype=torch.bfloat16)


def summary(input):
    #max_input_len=1024
    #if len(full_text.split()) > max_input_len:
    #    full_text = " ".join(full_text.split()[:max_input_len])
    output=text_summary(input)
    return output[0]['summary_text']

gr.close_all()

demo=gr.Interface(fn=summary,
                  inputs=[gr.Textbox(label='Input Text to Summarize', lines=1)],
                  outputs=[gr.Textbox(label='Summarized Text', lines=4)],
                  title='KS Text Summarizer',
                  description='This application will be used to summarize a corpus')

demo.launch()