Update app.py

app.py

@@ -1,44 +1,33 @@
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
 
 # Load the model and tokenizer
-tokenizer = AutoTokenizer.from_pretrained("kriton/greek-text-summarization")
-model = AutoModelForSeq2SeqLM.from_pretrained("kriton/greek-text-summarization")
+model_name = 'IMISLab/GreekWiki-umt5-base'
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 # Set up the summarizer pipeline
-summarizer = pipeline("summarization", model=model, tokenizer=tokenizer)
+summarizer = pipeline(
+    'summarization',
+    model=model,
+    tokenizer=tokenizer,
+    device=-1,  # -1 for CPU; set to 0 for GPU if available
+    max_new_tokens=220,
+    truncation=True
+)
 
 # 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)
+def generate_summary(text):
+    output = summarizer('summarize: ' + text)
+    return output[0]['summary_text']
 
 # 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
+    fn=generate_summary,  # The function that Gradio will use
+    inputs=gr.Textbox(label="Input Text", lines=5, placeholder="Enter the text to summarize..."),
+    outputs=gr.Textbox(label="Summary"),
+    live=True
 )
 
-# Launch the interface
+# Launch the Gradio interface
 iface.launch()