Update app.py

app.py

@@ -52,6 +52,56 @@ def get_nllb_lang_token(lang_code: str) -> str:
     """Get the correct token format for NLLB model."""
     return f"___{lang_code}___"
 
+def extract_text_from_file(uploaded_file) -> str:
+    """Extract text content from uploaded file based on its type."""
+    file_extension = Path(uploaded_file.name).suffix.lower()
+    
+    if file_extension == '.pdf':
+        return extract_from_pdf(uploaded_file)
+    elif file_extension == '.docx':
+        return extract_from_docx(uploaded_file)
+    elif file_extension == '.txt':
+        return uploaded_file.getvalue().decode('utf-8')
+    else:
+        raise ValueError(f"Unsupported file format: {file_extension}")
+
+def extract_from_pdf(file) -> str:
+    """Extract text from PDF file."""
+    pdf_reader = PyPDF2.PdfReader(file)
+    text = ""
+    for page in pdf_reader.pages:
+        text += page.extract_text() + "\n"
+    return text.strip()
+
+def extract_from_docx(file) -> str:
+    """Extract text from DOCX file."""
+    doc = docx.Document(file)
+    text = ""
+    for paragraph in doc.paragraphs:
+        text += paragraph.text + "\n"
+    return text.strip()
+
+def batch_process_text(text: str, max_length: int = 512) -> list:
+    """Split text into batches for processing."""
+    words = text.split()
+    batches = []
+    current_batch = []
+    current_length = 0
+    
+    for word in words:
+        if current_length + len(word) + 1 > max_length:
+            batches.append(" ".join(current_batch))
+            current_batch = [word]
+            current_length = len(word)
+        else:
+            current_batch.append(word)
+            current_length += len(word) + 1
+    
+    if current_batch:
+        batches.append(" ".join(current_batch))
+    
+    return batches
+
 @st.cache_resource
 def load_models():
     """Load and cache the translation and context interpretation models."""
@@ -76,7 +126,7 @@ def load_models():
         nllb_tokenizer = AutoTokenizer.from_pretrained(
             "facebook/nllb-200-distilled-600M",
             token=HF_TOKEN,
-            use_fast=False,  # Use slow tokenizer for better compatibility
+            use_fast=False,
             trust_remote_code=True
         )
         nllb_model = AutoModelForSeq2SeqLM.from_pretrained(
@@ -114,7 +164,34 @@ def load_models():
         st.error(f"PyTorch version: {torch.__version__}")
         raise e
 
-# [Previous file handling functions remain the same]
+@torch.no_grad()
+def interpret_context(text: str, gemma_tuple: Tuple) -> str:
+    """Use Gemma model to interpret context and understand regional nuances."""
+    tokenizer, model = gemma_tuple
+    
+    batches = batch_process_text(text)
+    interpreted_batches = []
+    
+    for batch in batches:
+        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()}
+        
+        outputs = model.generate(
+            **inputs,
+            max_length=512,
+            do_sample=True,
+            temperature=0.3,
+            pad_token_id=tokenizer.eos_token_id,
+            num_return_sequences=1
+        )
+        
+        interpreted_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        interpreted_text = interpreted_text.replace(prompt, "").strip()
+        interpreted_batches.append(interpreted_text)
+    
+    return " ".join(interpreted_batches)
 
 @torch.no_grad()
 def translate_text(text: str, source_lang: str, target_lang: str, nllb_tuple: Tuple) -> str:
@@ -124,15 +201,12 @@ def translate_text(text: str, source_lang: str, target_lang: str, nllb_tuple: Tu
     batches = batch_process_text(text)
     translated_batches = []
     
-    # Get target language token
     target_lang_token = get_nllb_lang_token(target_lang)
     
     for batch in batches:
-        # Prepare input text
         inputs = tokenizer(batch, return_tensors="pt", max_length=512, truncation=True)
         inputs = {k: v.to(model.device) for k, v in inputs.items()}
         
-        # Get target language token ID
         target_lang_id = tokenizer.convert_tokens_to_ids(target_lang_token)
         
         outputs = model.generate(
@@ -190,6 +264,17 @@ def correct_grammar(text: str, target_lang: str, mt5_tuple: Tuple) -> str:
     
     return " ".join(corrected_batches)
 
+def save_as_docx(text: str) -> io.BytesIO:
+    """Save translated text as a DOCX file."""
+    doc = docx.Document()
+    doc.add_paragraph(text)
+    
+    docx_buffer = io.BytesIO()
+    doc.save(docx_buffer)
+    docx_buffer.seek(0)
+    
+    return docx_buffer
+
 def main():
     st.title("🌐 Document Translation App")