from transformers import pipeline
import gradio as gr

# Input language translators (to English)
input_translators = {
    "Hindi": pipeline("translation_hi_to_en", model="Helsinki-NLP/opus-mt-hi-en"),
    "French": pipeline("translation_fr_to_en", model="Helsinki-NLP/opus-mt-fr-en"),
    "German": pipeline("translation_de_to_en", model="Helsinki-NLP/opus-mt-de-en"),
    "English": None  # No translation needed
}

# Summarization model
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Output translators (English → target)
output_translators = {
    "None": None,
    "Hindi": pipeline("translation_en_to_hi", model="Helsinki-NLP/opus-mt-en-hi"),
    "French": pipeline("translation_en_to_fr", model="Helsinki-NLP/opus-mt-en-fr"),
    "German": pipeline("translation_en_to_de", model="Helsinki-NLP/opus-mt-en-de"),
}

def summarize_multilang(text, input_lang, output_lang):
    # Step 1: Translate to English (if needed)
    if input_lang != "English":
        translator = input_translators[input_lang]
        text = translator(text)[0]['translation_text']
    
    # Step 2: Summarize
    summary = summarizer(text, max_length=130, min_length=30, do_sample=False)[0]['summary_text']
    
    # Step 3: Translate summary (if needed)
    if output_lang != "None":
        summary = output_translators[output_lang](summary)[0]['translation_text']
    
    return summary

# Gradio interface
gr.Interface(
    fn=summarize_multilang,
    inputs=[
        gr.Textbox(lines=10, label="Input Text"),
        gr.Dropdown(choices=["English", "Hindi", "French", "German"], label="Input Language"),
        gr.Dropdown(choices=["None", "Hindi", "French", "German"], label="Translate Summary To")
    ],
    outputs=gr.Textbox(label="Final Summary"),
    title="SummarAI",
    description="Supports input in Hindi, French, German, or English. Summarizes and optionally translates the summary."
).launch()