Update main.py

main.py

@@ -3,35 +3,44 @@ from fastapi import FastAPI
 from pydantic import BaseModel
 from transformers import T5ForConditionalGeneration, T5Tokenizer
 import torch
+import threading
+import uvicorn
 
-# Load your fine-tuned model
-model_path = "./t5-summarizer"  # Path inside Docker container
-model = T5ForConditionalGeneration.from_pretrained(model_path)
+# 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 = model.to(device)
+model.to(device)
 
+# FastAPI setup
 app = FastAPI()
-
 class TextInput(BaseModel):
     text: str
 
 @app.post("/summarize/")
 def summarize_text(input: TextInput):
-    input_text = "summarize: " + input.text.strip().replace("\n", " ")
-
-    inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True).to(device)
-    summary_ids = model.generate(inputs, 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}
-
-# Gradio UI setup
-gr.Interface(
-    fn=lambda text: summarize_text(TextInput(text=text))["summary"],  # Ensure it returns summary
-    inputs=gr.Textbox(label="Input Text"),
-    outputs=gr.Textbox(label="Summarized Text"),
-    flagging=False  # Disable flagging to prevent permission issues
-).launch()
-
+    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)}
+
+def run_fastapi():
+    uvicorn.run(app, host="0.0.0.0", port=8000)
+
+# Gradio UI
+iface = gr.Interface(
+    fn=lambda text: summarize_text(TextInput(text=text))["summary"],
+    inputs=gr.Textbox(lines=10, placeholder="Paste 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
+)
+
+# Start FastAPI in background, then launch Gradio
+threading.Thread(target=run_fastapi, daemon=True).start()
+iface.launch(server_name="0.0.0.0", server_port=7860)