Update main.py

main.py

@@ -1,29 +1,44 @@
+# main.py
+
 from fastapi import FastAPI
 from pydantic import BaseModel
-from transformers import BartForConditionalGeneration, BartTokenizer
+from transformers import BartForConditionalGeneration, BartTokenizer, MBartForConditionalGeneration, MBartTokenizer
 import torch
 
+# FastAPI қосымшасы
 app = FastAPI(title="Multilingual Text Summarizer")
 
-# 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)
+# ==========================
+# Модельдерді жүктеу
+# ==========================
+
+# English BART
+MODEL_NAME_EN = "facebook/bart-large-cnn"
+tokenizer_en = BartTokenizer.from_pretrained(MODEL_NAME_EN)
+model_en = BartForConditionalGeneration.from_pretrained(MODEL_NAME_EN)
 
-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)
+# Russian MBart
+MODEL_NAME_RU = "IlyaGusev/mbart_ru_sum_gazeta"
+tokenizer_ru = MBartTokenizer.from_pretrained(MODEL_NAME_RU)
+model_ru = MBartForConditionalGeneration.from_pretrained(MODEL_NAME_RU)
 
+# ==========================
+# API Request схемасы
+# ==========================
 class TextRequest(BaseModel):
     text: str
-    lang: str  # "en" or "ru"
+    lang: str  # "en" немесе "ru"
 
+# ==========================
+# Routes
+# ==========================
 @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
@@ -31,12 +46,16 @@ def summarize(request: TextRequest):
         tokenizer = tokenizer_en
         model = model_en
 
+    # Tokenize
     inputs = tokenizer([request.text], max_length=1024, return_tensors="pt", truncation=True)
+
+    # Summarize
     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}