Update app.py

app.py

@@ -8,22 +8,24 @@ import torch
 # =========================
 # Model setup (CPU-safe, Multi-language)
 # =========================
-# Use mBART for multilingual support (English + Arabic)
-SUMMARIZER_MODEL = "facebook/mbart-large-50-many-to-many-mmt"
-QA_MODEL = "google/flan-t5-base"  # Better for question generation
+# Use different models for English and Arabic
+EN_SUMMARIZER_MODEL = "sshleifer/distilbart-cnn-12-6"  # English summarization
+AR_SUMMARIZER_MODEL = "csebuetnlp/mT5_multilingual_XLSum"  # Multilingual (includes Arabic)
+QA_MODEL = "google/flan-t5-small"  # Question generation
 
 print("Loading models... This may take a minute on first run.")
 
-# Summarizer with mBART (supports Arabic)
-summarizer_tokenizer = AutoTokenizer.from_pretrained(SUMMARIZER_MODEL)
-summarizer_model = AutoModelForSeq2SeqLM.from_pretrained(SUMMARIZER_MODEL)
-summarizer = pipeline(
+# English summarizer
+en_summarizer = pipeline(
     "summarization",
-    model=summarizer_model,
-    tokenizer=summarizer_tokenizer,
+    model=EN_SUMMARIZER_MODEL,
     device=-1  # CPU only
 )
 
+# Multilingual summarizer (for Arabic and other languages)
+ar_tokenizer = AutoTokenizer.from_pretrained(AR_SUMMARIZER_MODEL)
+ar_model = AutoModelForSeq2SeqLM.from_pretrained(AR_SUMMARIZER_MODEL)
+
 # Question generator
 question_generator = pipeline(
     "text2text-generation",
@@ -31,7 +33,7 @@ question_generator = pipeline(
     device=-1  # CPU only
 )
 
-CHUNK_SIZE = 512  # Conservative for mBART
+CHUNK_SIZE = 512  # Conservative chunk size
 
 # =========================
 # Language Detection
@@ -63,8 +65,14 @@ def clean_text(text: str) -> str:
             result.append(s.strip())
     return " ".join(result)
 
-def chunk_text(text: str, tokenizer):
+def chunk_text(text: str, language: str):
     """Token-aware chunking to avoid model overflow."""
+    # Use appropriate tokenizer based on language
+    if language == "ar_AR":
+        tokenizer = ar_tokenizer
+    else:
+        tokenizer = AutoTokenizer.from_pretrained(EN_SUMMARIZER_MODEL)
+    
     tokens = tokenizer.encode(text, add_special_tokens=False)
     chunks = []
     for i in range(0, len(tokens), CHUNK_SIZE):
@@ -188,28 +196,34 @@ def summarize_long_text(text: str, summary_length: str, language: str, progress=
     headings_section = extract_possible_headings(text)
     
     progress(0.1, desc="Chunking text...")
-    chunks = chunk_text(text, summarizer_tokenizer)
+    chunks = chunk_text(text, language)
     
     summaries = []
     progress(0.2, desc="Summarizing chunks...")
     
-    # Set language tokens for mBART
-    src_lang = language
-    tgt_lang = language
-    
     for i in progress.tqdm(range(len(chunks))):
         chunk = chunks[i]
         try:
-            # For mBART, we need to set source and target language
-            summarizer_tokenizer.src_lang = src_lang
-            
-            summary = summarizer(
-                chunk,
-                max_length=length_params["max"],
-                min_length=length_params["min"],
-                do_sample=False,
-                forced_bos_token_id=summarizer_tokenizer.lang_code_to_id[tgt_lang]
-            )[0]["summary_text"]
+            if language == "ar_AR":
+                # Use mT5 for Arabic
+                inputs = ar_tokenizer(chunk, return_tensors="pt", max_length=512, truncation=True)
+                summary_ids = ar_model.generate(
+                    inputs["input_ids"],
+                    max_length=length_params["max"],
+                    min_length=length_params["min"],
+                    length_penalty=2.0,
+                    num_beams=4,
+                    early_stopping=True
+                )
+                summary = ar_tokenizer.decode(summary_ids[0], skip_special_tokens=True)
+            else:
+                # Use distilbart for English
+                summary = en_summarizer(
+                    chunk,
+                    max_length=length_params["max"],
+                    min_length=length_params["min"],
+                    do_sample=False
+                )[0]["summary_text"]
             
             cleaned = clean_text(summary)
             chunk_label = f"**Chunk {i+1}:**" if language == "en_XX" else f"**الجزء {i+1}:**"