Update app.py

app.py

@@ -2,12 +2,17 @@ import streamlit as st
 import PyPDF2
 import docx
 import io
-from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, T5ForConditionalGeneration
+from transformers import AutoModelForCausalLM, AutoTokenizer, AutoModelForSeq2SeqLM, MT5ForConditionalGeneration
 import torch
 from pathlib import Path
 import tempfile
 from typing import Union, Tuple
 import os
+from datetime import datetime, timezone
+
+# Display current information
+st.sidebar.text(f"Current Time (UTC): {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')}")
+st.sidebar.text(f"User: {os.environ.get('USER', 'gauravchand')}")
 
 # Get Hugging Face token from environment variables
 HF_TOKEN = os.environ.get('HF_TOKEN')
@@ -32,43 +37,67 @@ MT5_LANG_CODES = {
 @st.cache_resource
 def load_models():
     """Load and cache the translation, context interpretation, and grammar correction models."""
-    # Load Gemma model for context interpretation
-    gemma_tokenizer = AutoTokenizer.from_pretrained(
-        "google/gemma-2b",
-        token=HF_TOKEN
-    )
-    gemma_model = AutoModelForCausalLM.from_pretrained(
-        "google/gemma-2b",
-        device_map="auto",
-        torch_dtype=torch.float16,
-        token=HF_TOKEN
-    )
-    
-    # Load NLLB model for translation
-    nllb_tokenizer = AutoTokenizer.from_pretrained(
-        "facebook/nllb-200-distilled-600M",
-        token=HF_TOKEN
-    )
-    nllb_model = AutoModelForSeq2SeqLM.from_pretrained(
-        "facebook/nllb-200-distilled-600M",
-        device_map="auto",
-        torch_dtype=torch.float16,
-        token=HF_TOKEN
-    )
-    
-    # Load MT5 model for grammar correction
-    mt5_tokenizer = AutoTokenizer.from_pretrained(
-        "google/mt5-small",
-        token=HF_TOKEN
-    )
-    mt5_model = T5ForConditionalGeneration.from_pretrained(
-        "google/mt5-small",
-        device_map="auto",
-        torch_dtype=torch.float16,
-        token=HF_TOKEN
-    )
+    try:
+        # Set device
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        
+        # Load Gemma model for context interpretation
+        gemma_tokenizer = AutoTokenizer.from_pretrained(
+            "google/gemma-2b",
+            token=HF_TOKEN,
+            trust_remote_code=True
+        )
+        gemma_model = AutoModelForCausalLM.from_pretrained(
+            "google/gemma-2b",
+            token=HF_TOKEN,
+            torch_dtype=torch.float16,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True
+        )
+        
+        # Load NLLB model for translation
+        nllb_tokenizer = AutoTokenizer.from_pretrained(
+            "facebook/nllb-200-distilled-600M",
+            token=HF_TOKEN,
+            trust_remote_code=True
+        )
+        nllb_model = AutoModelForSeq2SeqLM.from_pretrained(
+            "facebook/nllb-200-distilled-600M",
+            token=HF_TOKEN,
+            torch_dtype=torch.float16,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True
+        )
+        
+        # Load MT5 model for grammar correction
+        mt5_tokenizer = AutoTokenizer.from_pretrained(
+            "google/mt5-small",
+            token=HF_TOKEN,
+            trust_remote_code=True
+        )
+        mt5_model = MT5ForConditionalGeneration.from_pretrained(
+            "google/mt5-small",
+            token=HF_TOKEN,
+            torch_dtype=torch.float16,
+            device_map="auto" if torch.cuda.is_available() else None,
+            trust_remote_code=True
+        )
+        
+        # Move models to device if not using device_map="auto"
+        if not torch.cuda.is_available():
+            gemma_model = gemma_model.to(device)
+            nllb_model = nllb_model.to(device)
+            mt5_model = mt5_model.to(device)
+        
+        return (gemma_tokenizer, gemma_model), (nllb_tokenizer, nllb_model), (mt5_tokenizer, mt5_model)
     
-    return (gemma_tokenizer, gemma_model), (nllb_tokenizer, nllb_model), (mt5_tokenizer, mt5_model)
+    except Exception as e:
+        st.error(f"Error loading models: {str(e)}")
+        st.error("Detailed error information:")
+        st.error(f"Python version: {sys.version}")
+        st.error(f"PyTorch version: {torch.__version__}")
+        st.error(f"Transformers version: {transformers.__version__}")
+        raise e
 
 def extract_text_from_file(uploaded_file) -> str:
     """Extract text content from uploaded file based on its type."""
@@ -99,43 +128,87 @@ def extract_from_docx(file) -> str:
         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
+
+@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
     
-    prompt = f"""Analyze the following text for context and cultural nuances, 
-    maintaining the core meaning while identifying any idiomatic expressions or 
-    cultural references: {text}"""
+    # Split text into batches
+    batches = batch_process_text(text)
+    interpreted_batches = []
     
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    outputs = model.generate(
-        **inputs,
-        max_length=1024,
-        temperature=0.3,
-        pad_token_id=tokenizer.eos_token_id
-    )
+    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}"""
+        
+        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,
+            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_batches.append(interpreted_text)
     
-    interpreted_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return interpreted_text
+    return " ".join(interpreted_batches)
 
+@torch.no_grad()
 def translate_text(text: str, source_lang: str, target_lang: str, nllb_tuple: Tuple) -> str:
     """Translate text using NLLB model."""
     tokenizer, model = nllb_tuple
     
-    inputs = tokenizer(text, return_tensors="pt").to(model.device)
-    forced_bos_token_id = tokenizer.lang_code_to_id[target_lang]
+    # Split text into batches
+    batches = batch_process_text(text)
+    translated_batches = []
     
-    outputs = model.generate(
-        **inputs,
-        forced_bos_token_id=forced_bos_token_id,
-        max_length=1024,
-        temperature=0.7,
-        num_beams=5
-    )
+    for batch in batches:
+        inputs = tokenizer(batch, return_tensors="pt", max_length=512, truncation=True)
+        inputs = {k: v.to(model.device) for k, v in inputs.items()}
+        
+        forced_bos_token_id = tokenizer.lang_code_to_id[target_lang]
+        
+        outputs = model.generate(
+            **inputs,
+            forced_bos_token_id=forced_bos_token_id,
+            max_length=512,
+            temperature=0.7,
+            num_beams=5,
+            num_return_sequences=1
+        )
+        
+        translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        translated_batches.append(translated_text)
     
-    translated_text = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
-    return translated_text
+    return " ".join(translated_batches)
 
+@torch.no_grad()
 def correct_grammar(text: str, target_lang: str, mt5_tuple: Tuple) -> str:
     """
     Correct grammar using MT5 model for all supported languages.
@@ -146,29 +219,40 @@ def correct_grammar(text: str, target_lang: str, mt5_tuple: Tuple) -> str:
     
     # Language-specific prompts for grammar correction
     prompts = {
-        'en': f"grammar: {text}",
-        'hi': f"व्याकरण सुधार: {text}",
-        'mr': f"व्याकरण सुधारणा: {text}"
+        'en': "grammar: ",
+        'hi': "व्याकरण सुधार: ",
+        'mr': "व्याकरण सुधारणा: "
     }
     
-    prompt = prompts[lang_code]
-    inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(model.device)
-    
-    outputs = model.generate(
-        **inputs,
-        max_length=512,
-        num_beams=5,
-        temperature=0.7,
-        top_p=0.9,
-        do_sample=True
-    )
-    
-    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Split text into batches
+    batches = batch_process_text(text)
+    corrected_batches = []
     
-    # Clean up any artifacts from the model output
-    corrected_text = corrected_text.replace("grammar:", "").replace("व्याकरण सुधार:", "").replace("व्याकरण सुधारणा:", "").strip()
+    for batch in batches:
+        prompt = prompts[lang_code] + 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,
+            num_beams=5,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            num_return_sequences=1
+        )
+        
+        corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Clean up any artifacts from the model output
+        for prefix in prompts.values():
+            corrected_text = corrected_text.replace(prefix, "")
+        corrected_text = corrected_text.strip()
+        
+        corrected_batches.append(corrected_text)
     
-    return corrected_text
+    return " ".join(corrected_batches)
 
 def save_as_docx(text: str) -> io.BytesIO:
     """Save translated text as a DOCX file."""
@@ -191,7 +275,7 @@ def main():
         except Exception as e:
             st.error(f"Error loading models: {str(e)}")
             st.error("Please check if the HF_TOKEN is valid and has the necessary permissions.")
-            st.stop()
+            return
     
     # File upload
     uploaded_file = st.file_uploader(