from transformers import T5ForConditionalGeneration, T5Tokenizer
import gradio as gr

model = T5ForConditionalGeneration.from_pretrained("abhi-codes/finetuned_flank_t5_for_summarization")
tokenizer = T5Tokenizer.from_pretrained("abhi-codes/finetuned_flank_t5_for_summarization")

examples = [
    ["""Tom: Did you submit the report?
Anika: Not yet, I'm fixing the charts.
Tom: The deadline is 5 pm.
Anika: I know. I'll send it by 4:30.
Tom: Great, please copy me on the email."""],
    ["""Nora: Are you picking up the groceries today?
Eli: Yes, after work.
Nora: Please get milk, eggs, and bread.
Eli: Got it. Anything else?
Nora: Bananas if they look fresh.
Eli: Okay, I'll be home around 6:30"""],
    ["""Priya: Did you call the dentist?
Karan: Yes, they had an opening tomorrow at 11.
Priya: Great. Did you book it?
Karan: Yes, I confirmed it.
Priya: Thanks. I'll leave work early to go."""]
]


def summarize(input):
    input = "summarize: "+ input
    model_inputs = tokenizer(input, return_tensors="pt", max_length=512, truncation=True,padding = 'max_length')
    summary_ids = model.generate(
        input_ids=model_inputs["input_ids"],
        attention_mask=model_inputs["attention_mask"],
        max_new_tokens=128,
        num_beams=4,
        no_repeat_ngram_size=3
    )
    return tokenizer.decode(summary_ids[0], skip_special_tokens=True)
  
demo = gr.Interface(
    fn=summarize,
    inputs=[
        gr.Textbox(
            lines=8,
            label="Dialogue",
            placeholder="Paste a conversation here"
        )],
        outputs=[
        gr.Textbox(
            lines=2,
            label="Summary"
        ),
    ],
    title="Dialogue Summarizer",
    description=(
        "Enter a chat-style conversation and the model will generate a short summary. "
        "For best results, write each message on a new line with the speaker name."
    ),
    examples=examples,
    flagging_mode="never"
    
)
demo.launch()