Update main.py

main.py

@@ -6,41 +6,58 @@ import torch
 import threading
 import uvicorn
 
-# Load model & tokenizer
+# 1. Load model & tokenizer
 model_path = "./t5-summarizer"
 tokenizer = T5Tokenizer.from_pretrained(model_path, legacy=False)
 model = T5ForConditionalGeneration.from_pretrained(model_path)
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
-# FastAPI setup
+# 2. FastAPI setup
 app = FastAPI()
 class TextInput(BaseModel):
     text: str
 
 @app.post("/summarize/")
 def summarize_text(input: TextInput):
-    inputs = tokenizer("summarize: " + input.text.replace("\n"," "),
-                       return_tensors="pt", max_length=512, truncation=True)
-    outputs = model.generate(inputs.input_ids.to(device),
-                             max_length=150, min_length=30,
-                             length_penalty=2.0, num_beams=4, early_stopping=True)
-    return {"summary": tokenizer.decode(outputs[0], skip_special_tokens=True)}
+    inputs = tokenizer(
+        "summarize: " + input.text.replace("\n", " "),
+        return_tensors="pt",
+        max_length=512,
+        truncation=True
+    ).to(device)
+    summary_ids = model.generate(
+        inputs.input_ids,
+        max_length=150,
+        min_length=30,
+        length_penalty=2.0,
+        num_beams=4,
+        early_stopping=True
+    )
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+    return {"summary": summary}
 
 def run_fastapi():
     uvicorn.run(app, host="0.0.0.0", port=8000)
 
-# Gradio UI
+# 3. Gradio UI
+def summarize_ui(text):
+    return summarize_text(TextInput(text=text))["summary"]
+
 iface = gr.Interface(
-    fn=lambda text: summarize_text(TextInput(text=text))["summary"],
-    inputs=gr.Textbox(lines=10, placeholder="Paste text here..."),
+    fn=summarize_ui,
+    inputs=gr.Textbox(lines=10, placeholder="Paste your text here..."),
     outputs=gr.Textbox(label="Summary"),
     title="Text Summarizer",
-    description="Fine-tuned T5 summarizer",
-    flagging_mode="never",               # Disable flagging
-    examples=[["Your example text here..."]]  # Pre-load examples
+    description="Fine-tuned T5 summarizer on CNN/DailyMail v3.0.0",
+    examples=[
+        ["Scientists have recently discovered a new species of frog in the Amazon rainforest..."],
+        ["The global economy is expected to grow at a slower pace this year..."],
+        ["In a thrilling final match, the underdog team scored a last-minute goal..."]
+    ],
+    allow_flagging="never"      # Disable flagging properly :contentReference[oaicite:3]{index=3}
 )
 
-# Start FastAPI in background, then launch Gradio
+# 4. Run both servers
 threading.Thread(target=run_fastapi, daemon=True).start()
 iface.launch(server_name="0.0.0.0", server_port=7860)