import gradio as gr
from transformers import AutoModelForQuestionAnswering, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForTokenClassification
import torch
import pandas as pd

# Load the models and tokenizers
QA_model = AutoModelForQuestionAnswering.from_pretrained("guldasta/xtreme-MLQA.hi.hi")
QA_token = AutoTokenizer.from_pretrained("guldasta/xtreme-MLQA.hi.hi")

Translation_model = AutoModelForSeq2SeqLM.from_pretrained("guldasta/Helsinki-NLP-opus-mt-mul-en-finetuned-hi-to-en")
Translation_token = AutoTokenizer.from_pretrained("guldasta/Helsinki-NLP-opus-mt-mul-en-finetuned-hi-to-en")

TokenClass_Model = AutoModelForTokenClassification.from_pretrained("guldasta/roberta-ner-multilingua-HiNER-Org")
TokenClass_Token = AutoTokenizer.from_pretrained("guldasta/roberta-ner-multilingua-HiNER-Org")

# Define the functions
def get_ans(question, context):
    model_inputs = QA_token(question, context, return_tensors="pt")
    outputs = QA_model(**model_inputs)
    start_logits = outputs.start_logits
    end_logits = outputs.end_logits
    start_indx = torch.argmax(start_logits)
    end_indx = torch.argmax(end_logits)
    answer = QA_token.decode(model_inputs['input_ids'][0][start_indx:end_indx + 1])
    return answer

def translate_text(text):
    model_inputs = Translation_token(text, return_tensors="pt")
    outputs = Translation_model.generate(**model_inputs)
    translated_text = Translation_token.decode(outputs[0], skip_special_tokens=True)
    return translated_text

def get_pos(text):
    inputs = TokenClass_Token(text, return_tensors="pt")
    outputs = TokenClass_Model(**inputs)
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=2)
    tokens = TokenClass_Token.convert_ids_to_tokens(inputs["input_ids"][0])
    predicted_labels = [TokenClass_Model.config.id2label[label.item()] for label in predictions[0]]

    filtered_tokens = []
    filtered_labels = []
    for token, label in zip(tokens, predicted_labels):
        if token not in TokenClass_Token.all_special_tokens:
            filtered_tokens.append(token)
            filtered_labels.append(label)

    df = pd.DataFrame({
        "Token": filtered_tokens,
        "Entity": filtered_labels
    })
    return df

# Create Gradio interfaces
qa_interface = gr.Interface(
    fn=get_ans,
    inputs=[
        gr.Textbox(lines=5, placeholder="Type a paragraph or context here...", label="Paragraph"),
        gr.Textbox(lines=1, placeholder="Type your question here...", label="Question"),
    ],
    outputs=gr.Textbox(label="Answer"),
    title="Question Answering",
    description="Enter a paragraph and ask a question to get the answer."
)

translation_interface = gr.Interface(
    fn=translate_text,
    inputs=gr.Textbox(lines=2, placeholder="Type your text here...", label="Text to Translate"),
    outputs=gr.Textbox(label="Translated Text"),
    title="Text Translation",
    description="Enter text to translate from Hindi to English."
)

pos_tag_interface = gr.Interface(
    fn=get_pos,
    inputs=gr.Textbox(lines=2, placeholder="Type your text here...", label="Text for POS Tagging"),
    outputs=gr.Dataframe(headers=["Token", "Entity"]),
    title="POS Tagging",
    description="Enter text to get POS tags and named entities."
)

# Combine all interfaces into a single Gradio app
app = gr.TabbedInterface(
    [qa_interface, translation_interface, pos_tag_interface],
    ["Question Answering", "Translation", "POS Tagging"]
)

# Launch the app
app.launch()