google

app.py

@@ -1,19 +1,55 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
+# Supported summarization models
+model_choices = {
+    "Pegasus (google/pegasus-xsum)": "google/pegasus-xsum",
+    "BART (facebook/bart-large-cnn)": "facebook/bart-large-cnn",
+    "T5 (t5-small)": "t5-small"
+}
 
-summarizer = pipeline("summarization", model="google/pegasus-xsum")
+# Cache for loaded models/tokenizers
+model_cache = {}
 
-def process(text):
-    summary = summarizer(text, max_length=10, min_length=5, do_sample=False, clean_up_tokenization_spaces=True, truncation=True)
-    summary = summary[:65]
-    return summary
+def load_model(model_name):
+    if model_name not in model_cache:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        model_cache[model_name] = (tokenizer, model)
+    return model_cache[model_name]
 
-demo = gr.Interface(
-    fn=process, 
-    inputs=gr.Textbox(label="Input Text"), 
-    outputs="json"
-)
+# Summarization function
+def summarize_text(input_text, model_label):
+    if not input_text.strip():
+        return "Please enter some text."
+
+    model_name = model_choices[model_label]
+    tokenizer, model = load_model(model_name)
+
+    if "t5" in model_name.lower():
+        input_text = "summarize: " + input_text
 
-demo.launch(share=True)
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
+    summary_ids = model.generate(
+        inputs["input_ids"],
+        max_length=20,  # Approximate for 65 characters
+        min_length=5,
+        do_sample=False
+    )
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+
+    return summary[:65]  # Ensure character limit
+
+# Gradio UI
+iface = gr.Interface(
+    fn=summarize_text,
+    inputs=[
+        gr.Textbox(lines=6, label="Enter text to summarize"),
+        gr.Dropdown(choices=list(model_choices.keys()), label="Choose summarization model", value="Pegasus (google/pegasus-xsum)")
+    ],
+    outputs=gr.Textbox(lines=2, label="Summary (max 65 characters)"),
+    title="Short Text Summarizer",
+    description="Summarizes input text to under 65 characters using a selected model."
+)
 
+iface.launch()