Update app.py

app.py

@@ -23,7 +23,7 @@ st.set_page_config(
     layout="wide"
 )
 
-# Display current information in sidebar with proper formatting
+# Display current information in sidebar
 current_time = datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')
 st.sidebar.markdown("""
 ### System Information
@@ -76,8 +76,8 @@ def load_models():
         nllb_tokenizer = AutoTokenizer.from_pretrained(
             "facebook/nllb-200-distilled-600M",
             token=HF_TOKEN,
-            trust_remote_code=True,
-            use_fast=False  # Use slow tokenizer to avoid warnings
+            src_lang="eng_Latn",
+            trust_remote_code=True
         )
         nllb_model = AutoModelForSeq2SeqLM.from_pretrained(
             "facebook/nllb-200-distilled-600M",
@@ -89,14 +89,12 @@ def load_models():
         
         # Load MT5 model for grammar correction
         mt5_tokenizer = AutoTokenizer.from_pretrained(
-            "google/mt5-small",
+            "google/mt5-base",  # Changed to base model for better performance
             token=HF_TOKEN,
-            trust_remote_code=True,
-            legacy=False,  # Use new behavior
-            use_fast=False  # Use slow tokenizer to avoid warnings
+            trust_remote_code=True
         )
         mt5_model = MT5ForConditionalGeneration.from_pretrained(
-            "google/mt5-small",
+            "google/mt5-base",  # Changed to base model for better performance
             token=HF_TOKEN,
             torch_dtype=torch.float16,
             device_map="auto" if torch.cuda.is_available() else None,
@@ -177,9 +175,7 @@ def interpret_context(text: str, gemma_tuple: Tuple) -> str:
     interpreted_batches = []
     
     for batch in batches:
-        prompt = f"""Analyze the following text for context and cultural nuances, 
-        maintaining the core meaning while identifying any idiomatic expressions or 
-        cultural references: {batch}"""
+        prompt = f"""Analyze and maintain the core meaning of this text: {batch}"""
         
         inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
@@ -194,6 +190,8 @@ def interpret_context(text: str, gemma_tuple: Tuple) -> str:
         )
         
         interpreted_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Remove the prompt from the output
+        interpreted_text = interpreted_text.replace(prompt, "").strip()
         interpreted_batches.append(interpreted_text)
     
     return " ".join(interpreted_batches)
@@ -207,17 +205,12 @@ def translate_text(text: str, source_lang: str, target_lang: str, nllb_tuple: Tu
     translated_batches = []
     
     for batch in batches:
-        # Add source language token to input
-        batch_with_lang = f"{source_lang} {batch}"
-        inputs = tokenizer(batch_with_lang, return_tensors="pt", max_length=512, truncation=True)
+        inputs = tokenizer(batch, return_tensors="pt", max_length=512, truncation=True)
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
         
-        # Add target language token
-        target_lang_token = tokenizer(target_lang, add_special_tokens=False)["input_ids"][0]
-        
         outputs = model.generate(
             **inputs,
-            forced_bos_token_id=target_lang_token,
+            forced_bos_token_id=tokenizer.lang_code_to_id[target_lang],
             max_length=512,
             do_sample=True,
             temperature=0.7,
@@ -236,35 +229,36 @@ def correct_grammar(text: str, target_lang: str, mt5_tuple: Tuple) -> str:
     tokenizer, model = mt5_tuple
     lang_code = MT5_LANG_CODES[target_lang]
     
+    # Language-specific prompts for grammar correction
     prompts = {
-        'en': "grammar: ",
-        'hi': "व्याकरण सुधार: ",
-        'mr': "व्याकरण सुधारणा: "
+        'en': "Fix grammar: ",
+        'hi': "व्याकरण: ",
+        'mr': "व्याकरण: "
     }
     
     batches = batch_process_text(text)
     corrected_batches = []
     
     for batch in batches:
-        prompt = prompts[lang_code] + batch
-        inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True)
+        # Prepare input with target language prefix
+        input_text = f"{prompts[lang_code]}{batch}"
+        inputs = tokenizer(input_text, return_tensors="pt", max_length=512, truncation=True)
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
         
         outputs = model.generate(
             **inputs,
             max_length=512,
             num_beams=5,
-            do_sample=True,
-            temperature=0.7,
-            top_p=0.9,
-            num_return_sequences=1
+            length_penalty=1.0,
+            early_stopping=True,
+            do_sample=False  # Disable sampling for more stable output
         )
         
         corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        # Clean up the output
         for prefix in prompts.values():
             corrected_text = corrected_text.replace(prefix, "")
-        corrected_text = corrected_text.strip()
-        
+        corrected_text = corrected_text.replace("<extra_id_0>", "").replace("<extra_id_1>", "").strip()
         corrected_batches.append(corrected_text)
     
     return " ".join(corrected_batches)