Update main.py

main.py

@@ -1,43 +1,42 @@
-from fastapi import FastAPI, UploadFile, File, Form
-from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import FileResponse
-from fastapi.staticfiles import StaticFiles
-from summarizer import generate_summary
-from pdf_reader import extract_text_from_pdf
+from fastapi import FastAPI
+from pydantic import BaseModel
+from transformers import BartForConditionalGeneration, BartTokenizer
+import torch
 
-app = FastAPI()
+app = FastAPI(title="Multilingual Text Summarizer")
 
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+# Model for English and Russian summarization
+MODEL_NAME = "facebook/bart-large-cnn"  # English
+tokenizer_en = BartTokenizer.from_pretrained(MODEL_NAME)
+model_en = BartForConditionalGeneration.from_pretrained(MODEL_NAME)
 
-app.mount("/static", StaticFiles(directory="."), name="static")
-
-@app.get("/")
-async def root():
-    return FileResponse("index.html")
-
-
-@app.post("/summarize/text")
-async def summarize_text(data: dict):
-    text = data.get("text", "")
-    lang = data.get("lang", "en")
-    if not text.strip():
-        return {"summary": "No text provided."}
-
-    summary = generate_summary(text, lang=lang)
-    return {"summary": summary}
+MODEL_NAME_RU = "IlyaGusev/mbart_ru_sum_gazeta"  # Russian
+tokenizer_ru = BartTokenizer.from_pretrained(MODEL_NAME_RU)
+model_ru = BartForConditionalGeneration.from_pretrained(MODEL_NAME_RU)
 
+class TextRequest(BaseModel):
+    text: str
+    lang: str  # "en" or "ru"
 
-@app.post("/summarize/pdf")
-async def summarize_pdf(file: UploadFile = File(...), lang: str = Form("en")):
-    pdf_bytes = await file.read()
-    text = extract_text_from_pdf(pdf_bytes)
-    if not text.strip():
-        return {"summary": "PDF is empty or could not extract text."}
-
-    summary = generate_summary(text, lang=lang)
+@app.get("/")
+def root():
+    return {"message": "Multilingual Text Summarizer is running!"}
+
+@app.post("/summarize/")
+def summarize(request: TextRequest):
+    if request.lang.lower() == "ru":
+        tokenizer = tokenizer_ru
+        model = model_ru
+    else:
+        tokenizer = tokenizer_en
+        model = model_en
+
+    inputs = tokenizer([request.text], max_length=1024, return_tensors="pt", truncation=True)
+    summary_ids = model.generate(
+        inputs["input_ids"],
+        num_beams=4,
+        max_length=150,
+        early_stopping=True
+    )
+    summary = tokenizer.decode(summary_ids[0], skip_special_tokens=True)
     return {"summary": summary}