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

# Load the model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("kriton/greek-text-summarization")
model = AutoModelForSeq2SeqLM.from_pretrained("kriton/greek-text-summarization")

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

# Define the summarization function
def generate_summary(article):
    inputs = tokenizer(
        'summarize: ' + article, 
        return_tensors="pt", 
        max_length=1024, 
        truncation=True,
        padding="max_length",
    )

    outputs = model.generate(
        inputs["input_ids"], 
        max_length=512, 
        min_length=130, 
        length_penalty=3.0, 
        num_beams=8, 
        early_stopping=True,
        repetition_penalty=3.0,
    )

    return tokenizer.decode(outputs[0], skip_special_tokens=True)

# Create Gradio interface
iface = gr.Interface(
    fn=generate_summary,  # Function to run
    inputs=gr.Textbox(label="Enter Greek Article", placeholder="Type or paste your article here..."),  # Input component
    outputs=gr.Textbox(label="Summary", interactive=True),  # Output component
    title="Greek Text Summarization",  # Title for the UI
    description="This app uses a pre-trained Greek summarization model to generate a brief summary of your input text.",  # Description
    allow_flagging="never"  # Optional: Disable flagging feature
)

# Launch the interface
iface.launch()