app.py

from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import streamlit as st
from PyPDF2 import PdfReader
import docx
import os
import re
import asyncio
from concurrent.futures import ThreadPoolExecutor
import torch
# Replace pytesseract with easyocr
import easyocr
from PIL import Image
import numpy as np

# Set up async environment for torch
if torch.cuda.is_available():
    torch.multiprocessing.set_start_method('spawn', force=True)

# Initialize asyncio event loop
try:
    loop = asyncio.get_event_loop()
except RuntimeError:
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

# Initialize EasyOCR reader
@st.cache_resource
def load_ocr_reader():
    try:
        return easyocr.Reader(['en'])  # Initialize for English
    except Exception as e:
        st.error(f"Error loading OCR reader: {str(e)}")
        return None

# Modified extract_text_from_image function with better error handling
def extract_text_from_image(image_file):
    try:
        # Get the OCR reader
        reader = load_ocr_reader()
        if reader is None:
            raise Exception("Failed to initialize OCR reader")
        
        # Read the image using PIL
        image = Image.open(image_file)
        
        # Convert to numpy array
        image_np = np.array(image)
        
        # Perform OCR
        results = reader.readtext(image_np)
        
        if not results:
            return "No text was detected in the image."
        
        # Extract text from results
        text = "\n".join([result[1] for result in results])
        return text.strip()
    except Exception as e:
        raise Exception(f"Error extracting text from image: {str(e)}")

# Modified extract_text function to support all file types
def extract_text(file):
    try:
        ext = os.path.splitext(file.name)[1].lower()
        
        if ext == ".pdf":
            try:
                reader = PdfReader(file)
                text = ""
                for page in reader.pages:
                    text += page.extract_text() + "\n"
                return text.strip()
            except Exception as e:
                raise Exception(f"Error reading PDF file: {str(e)}")
        
        elif ext == ".docx":
            try:
                doc = docx.Document(file)
                text = ""
                for para in doc.paragraphs:
                    text += para.text + "\n"
                return text.strip()
            except Exception as e:
                raise Exception(f"Error reading DOCX file: {str(e)}")
        
        elif ext == ".txt":
            try:
                return file.read().decode("utf-8").strip()
            except Exception as e:
                raise Exception(f"Error reading TXT file: {str(e)}")
        
        elif ext in [".jpg", ".jpeg", ".png"]:
            try:
                return extract_text_from_image(file)
            except Exception as e:
                raise Exception(f"Error processing image file: {str(e)}")
        
        else:
            raise ValueError("Unsupported file format. Please upload PDF, DOCX, TXT, or image files (JPG, JPEG, PNG).")
    except Exception as e:
        raise Exception(f"Error extracting text from file: {str(e)}")        

# Load NLLB model and tokenizer with error handling
@st.cache_resource
def load_translation_model():
    try:
        model_name = "facebook/nllb-200-distilled-600M"
        tokenizer = AutoTokenizer.from_pretrained(model_name)
        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
        return tokenizer, model
    except Exception as e:
        st.error(f"Error loading model: {str(e)}")
        return None, None

# Initialize model
@st.cache_resource
def initialize_models():
    tokenizer, model = load_translation_model()
    if tokenizer is None or model is None:
        st.error("Failed to initialize models")
        return None
    return {"nllb": (tokenizer, model)}



# Enhanced idiom mapping with more comprehensive translations
def preprocess_idioms(text, src_lang, tgt_lang):
   
    idiom_map = {}
    
    if src_lang == "en" and tgt_lang == "hi":
        idiom_map = {
           "no piece of cake": "कोई आसान काम नहीं",
            "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": "थोड़े से धैर्य से",
            
            # Additional common idioms
            "break a leg": "बहुत बहुत शुभकामनाएं",
            "hit the nail on the head": "बिल्कुल सही बात कहना",
            "once in a blue moon": "बहुत कम, कभी-कभार",
            "under the weather": "तबीयत ठीक नहीं",
            "cost an arm and a leg": "बहुत महंगा",
            "beating around the bush": "इधर-उधर की बात करना",
            "call it a day": "काम समाप्त करना",
            "burn the midnight oil": "रात-रात भर जागकर काम करना",
            "get the ball rolling": "शुरुआत करना",
            "pull yourself together": "खुद को संभालो",
            "shoot yourself in the foot": "अपना ही नुकसान करना",
            "take it with a grain of salt": "संदेह से लेना",
            "the last straw": "सहनशीलता की आखिरी सीमा",
            "time flies": "समय पंख लगाकर उड़ता है",
            "wrap your head around": "समझने की कोशिश करना",
            "cut corners": "काम में छोटा रास्ता अपनाना",
            "back to square one": "फिर से शुरू से",
            "blessing in disguise": "छिपा हुआ वरदान",
            "cry over spilled milk": "बीती बात पर पछताना",
            "keep your chin up": "हिम्मत रखना",
            
            # Work-related idioms
            "think outside the box": "नए तरीके से सोचना",
            "raise the bar": "मानक ऊंचा करना",
            "learning curve": "सीखने की प्रक्रिया",
            "up and running": "चालू और कार्यरत",
            "back to the drawing board": "फिर से योजना बनाना",
            
            # Project-related phrases
            "running into issues": "समस्याओं का सामना करना",
            "iron out the bugs": "खामियां दूर करना",
            "in the pipeline": "विचाराधीन",
            "moving forward": "आगे बढ़ते हुए",
            "touch base": "संपर्क में रहना",
            
            # Technical phrases
            "user-friendly": "उपयोगकर्ता के अनुकूल",
            "cutting-edge": "अत्याधुनिक",
            "state of the art": "अत्याधुनिक तकनीक",
            "proof of concept": "व्यवहार्यता का प्रमाण",
            "game changer": "खेल बदलने वाला"
        }
    elif src_lang == "en" and tgt_lang == "mr":
        idiom_map = {
            "no piece of cake": "सोपं काम नाही",
            "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 a leg": "खूप शुभेच्छा",
            "hit the nail on the head": "अगदी बरोबर बोललात",
            "once in a blue moon": "क्वचितच, कधीतरी",
            "under the weather": "तब्येत ठीक नसणे",
            "cost an arm and a leg": "खूप महाग",
            "beating around the bush": "गोल गोल फिरवणे",
            "call it a day": "दिवसाचं काम संपवणे",
            "burn the midnight oil": "रात्रंदिवस मेहनत करणे",
            "get the ball rolling": "सुरुवात करणे",
            "pull yourself together": "स्वतःला सावरा",
            "shoot yourself in the foot": "स्वतःचेच पाय स्वतः कापणे",
            "take it with a grain of salt": "साशंक दृष्टीने पाहणे",
            "the last straw": "सहनशक्तीची शेवटची मर्यादा",
            "time flies": "वेळ पंख लावून उडतो",
            "wrap your head around": "समजून घेण्याचा प्रयत्न करणे",
            "cut corners": "कमी वेळात काम उरकणे",
            "back to square one": "पुन्हा सुरुवातीला",
            "blessing in disguise": "आशीर्वाद लपलेला",
            "cry over spilled milk": "झालेल्या गोष्टीसाठी रडत बसणे",
            "keep your chin up": "धीर धरा",

            # Work-related idioms
            "think outside the box": "वेगळ्या पद्धतीने विचार करणे",
            "raise the bar": "पातळी उंचावणे",
            "learning curve": "शिकण्याची प्रक्रिया",
            "up and running": "सुरू आणि कार्यरत",
            "back to the drawing board": "पुन्हा नव्याने योजना आखणे",
            
            # Project-related phrases
            "running into issues": "अडचणींना सामोरे जाणे",
            "iron out the bugs": "त्रुटी दूर करणे",
            "in the pipeline": "विचाराधीन",
            "moving forward": "पुढे जाताना",
            "touch base": "संपर्कात राहणे",
            
            # Technical phrases
            "user-friendly": "वापरकर्त्यास सोयीस्कर",
            "cutting-edge": "अत्याधुनिक",
            "state of the art": "सर्वोत्कृष्ट तंत्रज्ञान",
            "proof of concept": "संकल्पनेची सिद्धता",
            "game changer": "खेळ बदलणारी गोष्ट"
        }

    if idiom_map:
        sorted_idioms = sorted(idiom_map.keys(), key=len, reverse=True)
        pattern = '|'.join(map(re.escape, sorted_idioms))
        
        def replace_idiom(match):
            return idiom_map[match.group(0).lower()]
        
        text = re.sub(pattern, replace_idiom, text, flags=re.IGNORECASE)
    
    return text

# Async translation function with fixed idiom processing
async def translate_text_async(text, src_lang, tgt_lang, models):
    if src_lang == tgt_lang:
        return text

    # Updated language mapping handling
    src_lang_simple = src_lang.lower()
    tgt_lang_simple = tgt_lang.lower()
    
    lang_map = {"english": "eng_Latn", "hindi": "hin_Deva", "marathi": "mar_Deva"}

    if src_lang_simple not in lang_map or tgt_lang_simple not in lang_map:
        return "Error: Unsupported language combination"

    try:
        # Process idioms first
        preprocessed_text = preprocess_idioms(text, src_lang_simple[:2], tgt_lang_simple[:2])
        
        tgt_lang_code = lang_map[tgt_lang_simple]
        tokenizer, model = models["nllb"]
        
        chunks = []
        current_chunk = ""
        
        # Split text into chunks while preserving sentences
        for sentence in re.split('([.!?।]+)', preprocessed_text):
            if sentence.strip():
                if len(current_chunk) + len(sentence) < 450:
                    current_chunk += sentence
                else:
                    if current_chunk:
                        chunks.append(current_chunk)
                    current_chunk = sentence
        
        if current_chunk:
            chunks.append(current_chunk)
        
        translated_text = ""
        
        # Translate each chunk
        for chunk in chunks:
            if chunk.strip():
                inputs = tokenizer(chunk, return_tensors="pt", padding=True, truncation=True, max_length=512)
                tgt_lang_id = tokenizer.convert_tokens_to_ids(tgt_lang_code)
                
                translated = model.generate(
                    **inputs,
                    forced_bos_token_id=tgt_lang_id,
                    max_length=512,
                    num_beams=5,
                    length_penalty=1.0,
                    no_repeat_ngram_size=3
                )
                
                translated_chunk = tokenizer.decode(translated[0], skip_special_tokens=True)
                translated_text += translated_chunk + " "
        
        return translated_text.strip()
    except Exception as e:
        return f"Error during translation: {str(e)}"

# Synchronous wrapper for translation
def translate_text(text, src_lang, tgt_lang, models):
    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)
    try:
        return loop.run_until_complete(translate_text_async(text, src_lang, tgt_lang, models))
    finally:
        loop.close()

def save_text_to_file(text, original_filename, prefix="translated"):
    try:
        # Get the original file extension and base name
        base_name = os.path.splitext(os.path.basename(original_filename))[0]
        output_filename = f"{prefix}_{base_name}.txt"
        
        # Save all translations as text files for simplicity and build speed
        with open(output_filename, "w", encoding="utf-8") as f:
            f.write(text)
        
        return output_filename
    except Exception as e:
        st.error(f"Error saving file: {str(e)}")
        return None

# Modified process_document function to handle multiple formats
def process_document(file, source_lang, target_lang, models):
    try:
        text = extract_text(file)
        translated_text = translate_text(text, source_lang, target_lang, models)
        
        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)
        
        if output_file is None:
            raise Exception("Failed to save output file")
        
        return output_file, translated_text
    except Exception as e:
        error_message = f"Error: {str(e)}"
        output_file = save_text_to_file(error_message, file.name, prefix="error")
        return output_file, error_message


# Modified main function to ensure proper language handling
def main():
    st.title("Document Translation Toolkit")

    # Initialize models with error handling
    models = initialize_models()
    if models is None:
        st.error("Failed to initialize translation models. Please try again.")
        return
    
    # Create tabs for different translation modes
    tab1, tab2 = st.tabs(["Document Translation", "Text Translation"])
    
    # Document Translation Tab
    with tab1:
        st.subheader("Document Translation")
        st.write("Upload a document (PDF, DOCX, TXT, or Image) and select languages.")
        
        uploaded_file = st.file_uploader(
            "Upload Document", 
            type=["pdf", "docx", "txt", "jpg", "jpeg", "png"],
            key="doc_uploader"
        )
        
        col1, col2 = st.columns(2)
        with col1:
            source_lang = st.selectbox(
                "Source Language",
                ["English", "Hindi", "Marathi"],
                index=0,
                key="doc_src"
            )
        with col2:
            target_lang = st.selectbox(
                "Target Language",
                ["English", "Hindi", "Marathi"],
                index=1,
                key="doc_tgt"
            )
        
        if uploaded_file is not None and st.button("Translate Document"):
            try:
                with st.spinner("Translating..."):
                    # Extract and show input text
                    input_text = extract_text(uploaded_file)
                    st.subheader("Input Text")
                    st.text_area("Original Text", input_text, height=200)
                    
                    # Translate and show output text
                    output_file, result_text = process_document(
                        uploaded_file,
                        source_lang.lower(),
                        target_lang.lower(),
                        models
                    )
                    
                    st.subheader("Translated Text")
                    st.text_area("Translation", result_text, height=200)
                    
                    # Provide download button with correct MIME type
                    if output_file and os.path.exists(output_file):
                        with open(output_file, "rb") as file:
                            # Set appropriate MIME type based on file extension
                            ext = os.path.splitext(output_file)[1].lower()
                            mime_types = {
                                '.pdf': 'application/pdf',
                                '.docx': 'application/vnd.openxmlformats-officedocument.wordprocessingml.document',
                                '.txt': 'text/plain',
                                '.jpg': 'image/jpeg',
                                '.jpeg': 'image/jpeg',
                                '.png': 'image/png'
                            }
                            mime_type = mime_types.get(ext, 'text/plain')
                            
                            st.download_button(
                                label="Download Translated Document",
                                data=file,
                                file_name=os.path.basename(output_file),
                                mime=mime_type
                            )
                    else:
                        st.error("Failed to generate output file")
            except Exception as e:
                st.error(f"An error occurred during translation: {str(e)}")
    
    # Text Translation Tab
    with tab2:
        st.subheader("Text Translation")
        st.write("Enter text directly for translation.")
        
        col1, col2 = st.columns(2)
        with col1:
            text_source_lang = st.selectbox(
                "Source Language",
                ["English", "Hindi", "Marathi"],
                index=0,
                key="text_src"
            )
        with col2:
            text_target_lang = st.selectbox(
                "Target Language",
                ["English", "Hindi", "Marathi"],
                index=1,
                key="text_tgt"
            )
        
        input_text = st.text_area("Enter text to translate", height=150)
        
        if input_text and st.button("Translate Text"):
            try:
                with st.spinner("Translating..."):
                    # Translate the input text
                    translated_text = translate_text(
                        input_text,
                        text_source_lang.lower(),
                        text_target_lang.lower(),
                        models
                    )
                    
                    # Show translation result
                    st.text_area("Translation", translated_text, height=150)
                    
                    # Add download button for translated text
                    st.download_button(
                        label="Download Translation",
                        data=translated_text,
                        file_name="translation.txt",
                        mime="text/plain"
                    )
            except Exception as e:
                st.error(f"An error occurred during translation: {str(e)}")

if __name__ == "__main__":
    try:
        main()
    except Exception as e:
        st.error(f"Application error: {str(e)}")