import gradio as gr
from transformers import MarianMTModel, MarianTokenizer, GPT2LMHeadModel, GPT2Tokenizer, AutoTokenizer, AutoModelForSequenceClassification
import torch

# Translation
def translate(text, target_language):
    language_codes = {
        "Spanish": "es",
        "French (European)": "fr",
        "French (Canadian)": "fr",
        "Italian": "it",
        "Ukrainian": "uk",
        "Portuguese (Brazilian)": "pt_BR",
        "Portuguese (European)": "pt",
        "Russian": "ru",
        "Chinese": "zh",
        "Dutch": "nl",
        "German": "de",
        "Arabic": "ar",
        "Hebrew": "he",
        "Greek": "el"
    }
    

# Text Generation
def generate_text(prompt):
    text_gen = pipeline("text-generation", model="gpt2")
    generated_text = text_gen(prompt, max_length=max_length, do_sample=True)[0]["generated_text"]
    return generated_text


    

# Text Classification
def classify_text(text):
    classifier = pipeline("zero-shot-classification")
    result = classifier(text, labels.split(','))
    scores = result["scores"]
    predictions = result["labels"]
    sorted_predictions = [pred for _, pred in sorted(zip(scores, predictions), reverse=True)]
    return sorted_predictions

    

# Sentiment Analysis
def sentiment_analysis(text):
    model_name = "distilbert-base-uncased-finetuned-sst-2-english"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    sentiment_scores = torch.softmax(outputs.logits, dim=1)
    sentiment = "positive" if sentiment_scores[0, 1] > sentiment_scores[0, 0] else "negative"
    return sentiment

language_options = [
    "Spanish", "French (European)", "French (Canadian)", "Italian", "Ukrainian",
    "Portuguese (Brazilian)", "Portuguese (European)", "Russian", "Chinese",
    "Dutch", "German", "Arabic", "Hebrew", "Greek"
]

iface = gr.Interface(
    [translate, generate_text, classify_text, sentiment_analysis],
    inputs=[
        gr.inputs.Textbox(lines=5, label="Enter text to translate:"),
        gr.inputs.Dropdown(choices=language_options, label="Target Language"),
        gr.inputs.Textbox(lines=5, label="Enter text for text generation:"),
        gr.inputs.Textbox(lines=5, label="Enter text for text classification:"),
        gr.inputs.Textbox(lines=5, label="Enter text for sentiment analysis:"),
    ],
    outputs=[
        gr.outputs.Textbox(label="Translated Text"),
        gr.outputs.Textbox(label="Generated Text"),
        gr.outputs.Textbox(label="Classification Result"),
        gr.outputs.Textbox(label="Sentiment Result"),
    ],
)

iface.launch()