from transformers import pipeline
import gradio as gr

# Load pipelines
sentiment_analyzer = pipeline("sentiment-analysis", model="distilbert-base-uncased-finetuned-sst-2-english")
ner_tagger = pipeline("ner", model="dslim/bert-base-NER", grouped_entities=True)

# Translation models for different languages
translation_models = {
    "French": "Helsinki-NLP/opus-mt-en-fr",
    "German": "Helsinki-NLP/opus-mt-en-de",
    "Spanish": "Helsinki-NLP/opus-mt-en-es",
    "Hindi": "Helsinki-NLP/opus-mt-en-hi",
    "Tamil": "ItchyChin/OrpoLlama-3-8B-memorize-translate-tamil-20241009",
    "Japanese": "Helsinki-NLP/opus-mt-en-ja",
    "Chinese": "Helsinki-NLP/opus-mt-en-zh",
    "Russian": "Helsinki-NLP/opus-mt-en-ru",
    "Arabic": "Helsinki-NLP/opus-mt-en-ar"
}

# Main function
def analyze_text(sentence, target_language):
    # Perform Sentiment Analysis
    sentiment = sentiment_analyzer(sentence)

    # Perform Named Entity Recognition
    ner = ner_tagger(sentence)

    # Perform Translation
    model_id = translation_models[target_language]
    translator = pipeline("translation", model=model_id)
    translated = translator(sentence, max_length=40)[0]['translation_text']

    return sentiment, ner, translated

# Gradio UI
interface = gr.Interface(
    fn=analyze_text,
    inputs=[
        gr.Textbox(label="Enter an English Sentence"),
        gr.Dropdown(choices=list(translation_models.keys()), label="Translate to Language")
    ],
    outputs=[
        gr.JSON(label="Sentiment Analysis"),
        gr.JSON(label="Named Entities"),
        gr.Textbox(label="Translation Result")
    ],
    title="🌍 NLP Translator + NER + Sentiment",
    description="This tool analyzes a sentence for sentiment, named entities, and translates it into a chosen language using Hugging Face Transformers."
)

interface.launch()