import gradio as gr
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline, Seq2SeqTrainer, Seq2SeqTrainingArguments

model_path = 'T5_samsum'

# Load the model
model = AutoModelForSeq2SeqLM.from_pretrained(model_path)

# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained(model_path)

# Create the summarization pipeline
summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)

# Define the summarization function
def summarize_dialogue(dialogue):
    summary = summarizer(dialogue, max_length=150, min_length=50, do_sample=False)
    return summary[0]['summary_text']

# Create the Gradio interface
iface = gr.Interface(
    fn=summarize_dialogue,
    inputs=gr.Textbox(lines=10, placeholder="Enter the dialogue here..."),
    outputs="text",
    title="Dialogue Summarizer",
    description="Enter a dialogue and this app will generate a summary using a pre-trained model."
)

# Launch the app
iface.launch()