Update app.py

app.py

@@ -2,225 +2,117 @@ import os
 import gradio as gr
 import torch
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
-import fitz  # PyMuPDF for PDF processing
-import docx2txt  # For DOCX processing
-from fpdf import FPDF  # For creating PDF outputs
 
-# Load model and tokenizer
-model_name = "facebook/mbart-large-50-many-to-many-mmt"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-
-# Set device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model = model.to(device)
-
-# Language mappings
+# First, let's create a simpler interface without complex schema handling
+# Define languages
 LANGUAGES = {
-    # Major Global Languages
     "English": "en_XX",
-    "Spanish": "es_XX",
-    "French": "fr_XX",
-    "German": "de_DE",
-    "Russian": "ru_RU",
-    "Chinese": "zh_CN",
-    "Japanese": "ja_XX",
-    "Arabic": "ar_AR",
-    
-    # Major Indian Languages
     "Hindi": "hi_IN",
     "Bengali": "bn_IN",
-    "Gujarati": "gu_IN", 
-    "Marathi": "mr_IN",
     "Tamil": "ta_IN",
     "Telugu": "te_IN",
     "Malayalam": "ml_IN",
     "Urdu": "ur_PK"
 }
 
-# Define translation function first
-def translate(text: str, source_lang: str, target_lang: str, max_length: int = 1024) -> str:
-    """
-    Translate text from source language to target language.
-    
-    Args:
-        text: Text to translate
-        source_lang: Source language name
-        target_lang: Target language name
-        max_length: Maximum length of input text
-        
-    Returns:
-        str: Translated text
-    """
+# Initialize model and tokenizer
+model_name = "facebook/mbart-large-50-many-to-many-mmt"
+tokenizer = None
+model = None
+
+def load_model():
+    global tokenizer, model
+    if tokenizer is None:
+        tokenizer = AutoTokenizer.from_pretrained(model_name)
+    if model is None:
+        model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+        if torch.cuda.is_available():
+            model = model.to("cuda")
+
+def translate_text(text, source_lang, target_lang):
+    """Simple translation function"""
     if not text:
-        return "No text provided for translation."
+        return "Please enter some text to translate."
     
     try:
+        load_model()
+        
         # Get language codes
         src_lang = LANGUAGES.get(source_lang)
         tgt_lang = LANGUAGES.get(target_lang)
         
-        if not src_lang or not tgt_lang:
-            return "Source or target language not supported."
-        
-        # Set tokenizer source language
+        # Set source language
         tokenizer.src_lang = src_lang
         
-        # Prepare input
-        inputs = tokenizer(text, return_tensors="pt", max_length=max_length, truncation=True)
-        inputs = {k: v.to(device) for k, v in inputs.items()}
+        # Tokenize
+        inputs = tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512)
+        if torch.cuda.is_available():
+            inputs = {k: v.to("cuda") for k, v in inputs.items()}
         
         # Generate translation
         with torch.no_grad():
             generated_tokens = model.generate(
                 **inputs,
                 forced_bos_token_id=tokenizer.lang_to_id[tgt_lang],
-                max_length=max_length,
-                num_beams=5,
+                max_length=512,
+                num_beams=4,
                 early_stopping=True
             )
         
-        # Decode translation
+        # Decode
         translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
         return translation
     
     except Exception as e:
-        return f"Translation error: {str(e)}"
-
-# File handling functions
-def extract_text_from_pdf(file_path: str) -> str:
-    """Extract text from a PDF file"""
-    text = ""
-    try:
-        doc = fitz.open(file_path)
-        for page in doc:
-            text += page.get_text()
-        return text
-    except Exception as e:
-        return f"Error extracting PDF text: {str(e)}"
-
-def extract_text_from_docx(file_path: str) -> str:
-    """Extract text from a DOCX file"""
-    try:
-        return docx2txt.process(file_path)
-    except Exception as e:
-        return f"Error extracting DOCX text: {str(e)}"
-
-def extract_text_from_txt(file_path: str) -> str:
-    """Extract text from a TXT file"""
-    try:
-        with open(file_path, 'r', encoding='utf-8') as file:
-            return file.read()
-    except UnicodeDecodeError:
-        try:
-            with open(file_path, 'r', encoding='latin-1') as file:
-                return file.read()
-        except Exception as e:
-            return f"Error extracting TXT text: {str(e)}"
-    except Exception as e:
-        return f"Error extracting TXT text: {str(e)}"
-
-def save_as_pdf(text: str, output_path: str) -> str:
-    """Save text as PDF"""
-    pdf = FPDF()
-    pdf.add_page()
-    pdf.set_font("Arial", size=12)
-    
-    try:
-        # Try UTF-8 first
-        pdf.multi_cell(0, 10, text)
-    except Exception:
-        try:
-            # Fall back to latin-1 with replacement
-            encoded_text = text.encode('latin-1', 'replace').decode('latin-1')
-            pdf.multi_cell(0, 10, encoded_text)
-        except Exception as e:
-            return f"Error creating PDF: {str(e)}"
-    
-    try:
-        pdf.output(output_path)
-        return output_path
-    except Exception as e:
-        return f"Error saving PDF: {str(e)}"
-
-def process_file(file, source_lang: str, target_lang: str) -> tuple[str | None, str]:
-    """Process uploaded file and translate its content"""
-    if file is None:
-        return None, "No file uploaded."
-        
-    try:
-        # Save uploaded file temporarily
-        temp_file_path = file.name
-        
-        # Extract text based on file type
-        if temp_file_path.lower().endswith('.pdf'):
-            text = extract_text_from_pdf(temp_file_path)
-        elif temp_file_path.lower().endswith('.docx'):
-            text = extract_text_from_docx(temp_file_path)
-        elif temp_file_path.lower().endswith('.txt'):
-            text = extract_text_from_txt(temp_file_path)
-        else:
-            return None, "Unsupported file format. Please upload PDF, DOCX, or TXT files."
-        
-        # Translate the extracted text
-        translated_text = translate(text, source_lang, target_lang)
-        
-        # Save translation as PDF
-        output_pdf_path = os.path.join(os.path.dirname(temp_file_path), 
-                                     f"translated_{os.path.basename(temp_file_path)}.pdf")
-        result = save_as_pdf(translated_text, output_pdf_path)
-        
-        if isinstance(result, str) and result.startswith("Error"):
-            return None, result
-            
-        return output_pdf_path, translated_text
-    
-    except Exception as e:
-        return None, f"Error processing file: {str(e)}"
+        return f"Translation Error: {str(e)}"
 
 # Create Gradio interface
-with gr.Blocks(title="Indian Language Translator") as demo:
-    gr.Markdown("# Indian & Global Language Translator")
-    gr.Markdown("Translate text with understanding of idioms and cultural expressions")
+with gr.Blocks(title="Simple Language Translator") as demo:
+    gr.Markdown("# Simple Language Translator")
     
-    with gr.Tab("Text Translation"):
-        with gr.Row():
-            source_lang_text = gr.Dropdown(list(LANGUAGES.keys()), value="English", label="Source Language")
-            target_lang_text = gr.Dropdown(list(LANGUAGES.keys()), value="Hindi", label="Target Language")
-        
-        with gr.Row():
-            input_text = gr.Textbox(label="Enter text to translate", lines=5, placeholder="Type or paste text here...")
-            output_text = gr.Textbox(label="Translation", lines=5)
+    with gr.Row():
+        with gr.Column():
+            input_text = gr.Textbox(
+                label="Input Text",
+                placeholder="Enter text to translate...",
+                lines=5
+            )
+            source_lang = gr.Dropdown(
+                choices=list(LANGUAGES.keys()),
+                value="English",
+                label="Source Language"
+            )
+            target_lang = gr.Dropdown(
+                choices=list(LANGUAGES.keys()),
+                value="Hindi",
+                label="Target Language"
+            )
+            translate_btn = gr.Button("Translate")
         
-        translate_btn = gr.Button("Translate Text", variant="primary")
-        translate_btn.click(
-            fn=translate,
-            inputs=[input_text, source_lang_text, target_lang_text],
-            outputs=output_text
-        )
+        with gr.Column():
+            output_text = gr.Textbox(
+                label="Translation",
+                lines=5
+            )
     
-    with gr.Tab("Document Translation"):
-        with gr.Row():
-            source_lang_doc = gr.Dropdown(list(LANGUAGES.keys()), value="English", label="Source Language")
-            target_lang_doc = gr.Dropdown(list(LANGUAGES.keys()), value="Hindi", label="Target Language")
-        
-        file_input = gr.File(label="Upload Document (PDF, DOCX, TXT)", file_types=[".pdf", ".docx", ".txt"])
-        with gr.Row():
-            output_file = gr.File(label="Translated PDF")
-            output_preview = gr.Textbox(label="Translation Preview", lines=8)
-        
-        translate_doc_btn = gr.Button("Translate Document", variant="primary")
-        translate_doc_btn.click(
-            fn=process_file,
-            inputs=[file_input, source_lang_doc, target_lang_doc],
-            outputs=[output_file, output_preview]
-        )
+    # Set up translation event
+    translate_btn.click(
+        fn=translate_text,
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text
+    )
     
-    gr.Markdown("### Supported File Types: PDF, DOCX, TXT")
-    gr.Markdown("### Features:")
-    gr.Markdown("- Supports major Indian languages including Hindi, Bengali, Tamil, Telugu, Malayalam")
-    gr.Markdown("- Context-aware translation that understands idioms and cultural expressions")
-    gr.Markdown("- Document translation with PDF output")
+    # Add examples
+    gr.Examples(
+        examples=[
+            ["Hello, how are you?", "English", "Hindi"],
+            ["नमस्ते, कैसे हैं आप?", "Hindi", "English"],
+        ],
+        inputs=[input_text, source_lang, target_lang],
+        outputs=output_text,
+        fn=translate_text,
+        cache_examples=True,
+    )
 
 if __name__ == "__main__":
     demo.launch(share=True)