Update app.py

app.py

@@ -4,6 +4,14 @@ from PyPDF2 import PdfReader
 import docx
 import os
 import re
+from datetime import datetime
+
+# Page config
+st.set_page_config(
+    page_title="Document Translator (NLLB-200)",
+    page_icon="📄",
+    layout="wide"
+)
 
 # Load NLLB model and tokenizer
 @st.cache_resource
@@ -19,7 +27,92 @@ def initialize_models():
     tokenizer, model = load_translation_model()
     return {"nllb": (tokenizer, model)}
 
-# Function to extract text from different file types
+def split_long_sentence(sentence, max_length=200):
+    """Split long sentences into smaller chunks at appropriate break points."""
+    if len(sentence) <= max_length:
+        return [sentence]
+    
+    chunks = []
+    current_chunk = ""
+    words = sentence.split()
+    
+    for word in words:
+        if len(current_chunk) + len(word) + 1 <= max_length:
+            current_chunk += (" " + word if current_chunk else word)
+        else:
+            chunks.append(current_chunk)
+            current_chunk = word
+    
+    if current_chunk:
+        chunks.append(current_chunk)
+    
+    return chunks
+
+def preprocess_idioms(text, src_lang, tgt_lang):
+    if src_lang == "en" and tgt_lang == "hi":
+        idiom_map = {
+            # Common English-Hindi idiom mappings
+            "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": "थोड़े से धैर्य से",
+            "break the ice": "बातचीत की शुरुआत करना",
+            "on cloud nine": "सातवें आसमान पर होना",
+            "once in a blue moon": "कभी-कभार",
+            "beating around the bush": "इधर-उधर की बात करना",
+            "burning the midnight oil": "रात-रात भर जागकर काम करना",
+            "calm before the storm": "तूफान से पहले की शांति",
+            "cost an arm and a leg": "बहुत महंगा होना",
+            "blessing in disguise": "छुपा हुआ वरदान",
+            "kill two birds with one stone": "एक पंथ दो काज",
+            "a piece of cake": "बहुत आसान काम",
+            "under the weather": "तबीयत ठीक न होना",
+            "pull yourself together": "खुद को संभालो",
+            "rise and shine": "जल्दी उठो और तैयार हो जाओ",
+            "time flies": "समय पंख लगाकर उड़ता है",
+            "actions speak louder than words": "कथनी से करनी बड़ी",
+            "all ears": "पूरा ध्यान से सुन रहा हूं",
+            "back to square one": "वापस शुरुआत में",
+            "better late than never": "देर आये दुरुस्त आये",
+            "cry over spilled milk": "बीती बात पर पछताना",
+            "down to earth": "सरल स्वभाव का",
+            "every cloud has a silver lining": "हर मुसीबत में कोई न कोई अच्छाई छिपी होती है",
+            "food for thought": "सोचने वाली बात",
+            "give someone the benefit of the doubt": "शक का फायदा देना",
+            "hit the nail on the head": "सटीक बात कहना",
+            "in hot water": "मुसीबत में होना"
+        }
+        
+        # Sort idioms by length (longest first) to handle overlapping phrases
+        sorted_idioms = sorted(idiom_map.keys(), key=len, reverse=True)
+        
+        # Replace idioms with their translations
+        for idiom in sorted_idioms:
+            pattern = r'\b' + re.escape(idiom) + r'\b'
+            text = re.sub(pattern, idiom_map[idiom], text, flags=re.IGNORECASE)
+            
+    elif src_lang == "en" and tgt_lang == "mr":
+        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": "थोड्या धीराने",
+            "break the ice": "संभाषणाची सुरुवात करणे",
+            "on cloud nine": "आनंदात असणे",
+            "once in a blue moon": "क्वचितच",
+            "burning the midnight oil": "रात्रंदिवस मेहनत करणे",
+            "better late than never": "उशीर का होईना पण योग्य वेळी"
+        }
+        for idiom, translation in idiom_map.items():
+            pattern = r'\b' + re.escape(idiom) + r'\b'
+            text = re.sub(pattern, translation, text, flags=re.IGNORECASE)
+    
+    return text
+
 def extract_text(file):
     ext = os.path.splitext(file.name)[1].lower()
     
@@ -43,7 +136,6 @@ def extract_text(file):
     else:
         raise ValueError("Unsupported file format. Please upload PDF, DOCX, or TXT files.")
 
-# Translation function
 def translate_text(text, src_lang, tgt_lang, models):
     if src_lang == tgt_lang:
         return text
@@ -61,86 +153,16 @@ def translate_text(text, src_lang, tgt_lang, models):
     # Preprocess for idioms
     preprocessed_text = preprocess_idioms(text, src_lang, tgt_lang)
     
-    # Split text into manageable chunks
-    sentences = preprocessed_text.split("\n")
-    translated_text = ""
+    # Split text into smaller chunks (sentences)
+    sentences = re.split(r'(?<=[.!?])\s+', preprocessed_text)
+    translated_text = []
     
     for sentence in sentences:
         if sentence.strip():
-            inputs = tokenizer(sentence, return_tensors="pt", padding=True, truncation=True, max_length=512)
-            # Use lang_code_to_id instead of get_lang_id
-            translated = model.generate(
-                **inputs,
-                forced_bos_token_id=tokenizer.lang_code_to_id[tgt_lang_code],
-                max_length=512
-            )
-            translated_sentence = tokenizer.decode(translated[0], skip_special_tokens=True)
-            translated_text += translated_sentence + "\n"
-    
-    return translated_text
-
-# Function to save text as a file
-def save_text_to_file(text, original_filename, prefix="translated"):
-    output_filename = f"{prefix}_{os.path.basename(original_filename)}.txt"
-    with open(output_filename, "w", encoding="utf-8") as f:
-        f.write(text)
-    return output_filename
-
-# Main processing function
-def process_document(file, source_lang, target_lang, models):
-    try:
-        # Extract text from uploaded file
-        text = extract_text(file)
-        
-        # Translate the text
-        translated_text = translate_text(text, source_lang, target_lang, models)
-        
-        # Save the result (success or error) to a file
-        if translated_text.startswith("Error:"):
-            output_file = save_text_to_file(translated_text, file.name, prefix="error")
-        else:
-            output_file = save_text_to_file(translated_text, file.name)
-        
-        return output_file, translated_text
-    except Exception as e:
-        # Save error message to a file
-        error_message = f"Error: {str(e)}"
-        output_file = save_text_to_file(error_message, file.name, prefix="error")
-        return output_file, error_message
-
-# Streamlit interface
-def main():
-    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
-    models = initialize_models()
-    
-    # File uploader
-    uploaded_file = st.file_uploader("Upload Document", type=["pdf", "docx", "txt"])
-    
-    # Language selection
-    col1, col2 = st.columns(2)
-    with col1:
-        source_lang = st.selectbox("Source Language", ["en", "hi", "mr"], index=0)
-    with col2:
-        target_lang = st.selectbox("Target Language", ["en", "hi", "mr"], index=1)
-    
-    if uploaded_file is not None and st.button("Translate"):
-        with st.spinner("Translating..."):
-            output_file, result_text = process_document(uploaded_file, source_lang, target_lang, models)
-            
-            # Display result
-            st.text_area("Translated Text", result_text, height=300)
+            chunks = split_long_sentence(sentence, max_length=200)
             
-            # Provide download button
-            with open(output_file, "rb") as file:
-                st.download_button(
-                    label="Download Translated Document",
-                    data=file,
-                    file_name=os.path.basename(output_file),
-                    mime="text/plain"
-                )
-
-if __name__ == "__main__":
-    main()
+            for chunk in chunks:
+                try:
+                    inputs = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True, max_length=512)
+                    translated = model.generate(
+                        **inputs