app.py

import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
import fitz  # PyMuPDF for PDF handling
from io import BytesIO

# Load IndicTrans model for Indian languages (example: English to Hindi)
model_name = "ai4bharat/indictrans2-en-indic-1b"  # Supports multiple Indian languages
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
translator = pipeline("translation", model=model, tokenizer=tokenizer)

# Language mapping for Indian languages and others
language_map = {
    "French": "fr",  # Using OPUS model for non-Indian languages
    "Spanish": "es",
    "German": "de",
    "Hindi": "hi",
    "Tamil": "ta",
    "Telugu": "te",
    "Bengali": "bn",
    "Gujarati": "gu",
    "Marathi": "mr",
    "Kannada": "kn",
    "Malayalam": "ml",
    "Punjabi": "pa",
}

def translate_text(input_text, target_language):
    """Translate text with context awareness"""
    target_lang_code = language_map[target_language]
    
    # For Indian languages, use IndicTrans
    if target_lang_code in ["hi", "ta", "te", "bn", "gu", "mr", "kn", "ml", "pa"]:
        translated = translator(input_text, src_lang="en", tgt_lang=target_lang_code)[0]['translation_text']
    else:
        # For non-Indian languages, switch to OPUS model (example: English to French)
        opus_translator = pipeline("translation", model=f"Helsinki-NLP/opus-mt-en-{target_lang_code}")
        translated = opus_translator(input_text)[0]['translation_text']
    
    return translated

def process_input(input_data, target_language):
    """Handle both text and document inputs"""
    if isinstance(input_data, str):  # Direct text input
        text = input_data
    else:  # File input (assuming text-based document)
        doc = fitz.open(stream=input_data.read(), filetype="pdf")
        text = ""
        for page in doc:
            text += page.get_text()
    
    # Translate the extracted text
    translated_text = translate_text(text, target_language)
    
    # Create PDF output
    pdf_output = fitz.open()
    page = pdf_output.new_page()
    page.insert_text((50, 50), translated_text)
    
    # Save PDF to bytes
    pdf_bytes = BytesIO()
    pdf_output.save(pdf_bytes)
    pdf_bytes.seek(0)
    
    return pdf_bytes

# Gradio Interface
with gr.Blocks(title="Context-Aware Translator with Indian Languages") as demo:
    gr.Markdown("# Context-Aware Language Translator")
    gr.Markdown("Translate text or upload a document into Indian languages or others, and get a PDF output.")
    
    with gr.Row():
        with gr.Column():
            input_type = gr.Radio(["Text", "Document"], label="Input Type", value="Text")
            text_input = gr.Textbox(lines=5, label="Enter Text", visible=True)
            file_input = gr.File(label="Upload Document", visible=False)
            target_lang = gr.Dropdown(
                choices=list(language_map.keys()),
                label="Target Language",
                value="Hindi"
            )
            submit_btn = gr.Button("Translate")
        
        with gr.Column():
            output_pdf = gr.File(label="Download Translated PDF")

    # Dynamic visibility based on input type
    def update_visibility(choice):
        return (
            gr.update(visible=(choice == "Text")),
            gr.update(visible=(choice == "Document"))
        )
    
    input_type.change(
        fn=update_visibility,
        inputs=input_type,
        outputs=[text_input, file_input]
    )

    # Process the input and generate output
    submit_btn.click(
        fn=process_input,
        inputs=[gr.State(value=None, _js="() => document.querySelector('input[name=\"input_type\"]:checked').value === 'Text' ? document.querySelector('#text_input textarea').value : document.querySelector('#file_input input').files[0]"), target_lang],
        outputs=output_pdf
    )

demo.launch()