Update app.py

app.py

@@ -1,22 +1,39 @@
-from transformers import M2M100Tokenizer, M2M100ForConditionalGeneration
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import streamlit as st
 from PyPDF2 import PdfReader
 import docx
 import os
+import re
 
-# Load M2M100 model and tokenizer
+# Load NLLB model and tokenizer
 @st.cache_resource
 def load_translation_model():
-    model_name = "facebook/m2m100_418M"
-    tokenizer = M2M100Tokenizer.from_pretrained(model_name)
-    model = M2M100ForConditionalGeneration.from_pretrained(model_name)
+    model_name = "facebook/nllb-200-distilled-600M"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
     return tokenizer, model
 
 # Initialize model
 @st.cache_resource
 def initialize_models():
     tokenizer, model = load_translation_model()
-    return {"m2m100": (tokenizer, model)}
+    return {"nllb": (tokenizer, model)}
+
+# Preprocess text to handle idioms (basic mapping for demonstration)
+def preprocess_idioms(text, src_lang, tgt_lang):
+    if src_lang == "en" and tgt_lang == "hi":
+        idiom_map = {
+            "no piece of cake": "कोई आसान काम नहीं",
+            "bite the bullet": "दांतों तले उंगली दबाना",
+            "tackle it head-on": "इसे पूरे मन से हाथ में लेना",
+            "fell into place": "ठीक हो गया",
+            "see the light at the end of the tunnel": "मुश्किलों के अंत में उम्मीद की किरण दिखाई देना",
+            "with a little perseverance": "थोड़े से धैर्य से"
+        }
+        for idiom, translation in idiom_map.items():
+            text = re.sub(r'\b' + idiom + r'\b', translation, text, flags=re.IGNORECASE)
+    # Add more mappings for other language pairs (e.g., en-mr) as needed
+    return text
 
 # Function to extract text from different file types
 def extract_text(file):
@@ -47,8 +64,8 @@ def translate_text(text, src_lang, tgt_lang, models):
     if src_lang == tgt_lang:
         return text
 
-    # Language codes for M2M100 (simplified to match user input)
-    lang_map = {"en": "en", "hi": "hi", "mr": "mr"}
+    # Language codes for NLLB
+    lang_map = {"en": "eng_Latn", "hi": "hin_Deva", "mr": "mar_Deva"}
 
     if src_lang not in lang_map or tgt_lang not in lang_map:
         return "Error: Unsupported language combination"
@@ -56,19 +73,19 @@ def translate_text(text, src_lang, tgt_lang, models):
     src_lang_code = lang_map[src_lang]
     tgt_lang_code = lang_map[tgt_lang]
 
-    tokenizer, model = models["m2m100"]
+    tokenizer, model = models["nllb"]
     
-    # Set source language
-    tokenizer.src_lang = src_lang_code
+    # Preprocess for idioms
+    preprocessed_text = preprocess_idioms(text, src_lang, tgt_lang)
     
     # Split text into manageable chunks
-    sentences = text.split("\n")
+    sentences = preprocessed_text.split("\n")
     translated_text = ""
     
     for sentence in sentences:
         if sentence.strip():
-            inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=512)
-            translated = model.generate(**inputs, forced_bos_token_id=tokenizer.get_lang_id(tgt_lang_code))
+            inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=512, src_lang=src_lang_code)
+            translated = model.generate(**inputs, forced_bos_token_id=tokenizer.get_lang_id(tgt_lang_code), max_length=512)
             translated_sentence = tokenizer.decode(translated[0], skip_special_tokens=True)
             translated_text += translated_sentence + "\n"
     
@@ -105,7 +122,7 @@ def process_document(file, source_lang, target_lang, models):
 
 # Streamlit interface
 def main():
-    st.title("Document Translator (M2M100)")
+    st.title("Document Translator (NLLB-200)")
     st.write("Upload a document (PDF, DOCX, or TXT) and select source and target languages (English, Hindi, Marathi).")
     
     # Initialize models