import gradio as gr
from transformers import pipeline

# Define the models
model1 = pipeline("text-classification", model="albertmartinez/bert-sdg-classification")
model2 = pipeline("text-classification", model="albertmartinez/bert-multilingual-sdg-classification")
model3 = pipeline("text-classification", model="albertmartinez/distilbert-multilingual-sdg-classification")
model4 = pipeline("text-classification", model="albertmartinez/xlm-roberta-large-sdg-classification")

def classify_text(text, model):
    result = model(text, top_k=16, truncation=True, max_length=512)
    return {p["label"]: p["score"] for p in result}

def classify_all(text):
    return [
        {p["label"]: p["score"] for p in model1(text, top_k=16, truncation=True, max_length=512)},
        {p["label"]: p["score"] for p in model2(text, top_k=16, truncation=True, max_length=512)},
        {p["label"]: p["score"] for p in model3(text, top_k=16, truncation=True, max_length=512)},
        {p["label"]: p["score"] for p in model4(text, top_k=16, truncation=True, max_length=512)}
    ]


ifaceall = gr.Interface(
    fn=classify_all,
    inputs=gr.Textbox(lines=2, label="Text", placeholder="Enter text here..."),
    outputs=[gr.Label(label="bert"), gr.Label(label="bert-multilingual"), gr.Label(label="distilbert-multilingual"), gr.Label(label="xlm-roberta-large")],
    title="",
    description="Enter a text and see the text classification result!"
)

# Interface for the first model
iface1 = gr.Interface(
    fn=lambda text: classify_text(text, model1),
    inputs=gr.Textbox(lines=2, label="Text", placeholder="Enter text here..."),
    outputs=gr.Label(label="Top SDG Predicted"),
    title="albertmartinez/bert-sdg-classification",
    description="Enter a text and see the text classification result!"
)

# Interface for the second model
iface2 = gr.Interface(
    fn=lambda text: classify_text(text, model2),
    inputs=gr.Textbox(lines=2, label="Text", placeholder="Enter text here..."),
    outputs=gr.Label(label="Top SDG Predicted"),
    title="albertmartinez/bert-multilingual-sdg-classification",
    description="Enter a text and see the text classification result!"
)

# Interface for the three model
iface3 = gr.Interface(
    fn=lambda text: classify_text(text, model3),
    inputs=gr.Textbox(lines=2, label="Text", placeholder="Enter text here..."),
    outputs=gr.Label(label="Top SDG Predicted"),
    title="albertmartinez/distilbert-multilingual-sdg-classification",
    description="Enter a text and see the text classification result!"
)

# Interface for the four model
iface4 = gr.Interface(
    fn=lambda text: classify_text(text, model4),
    inputs=gr.Textbox(lines=2, label="Text", placeholder="Enter text here..."),
    outputs=gr.Label(label="Top SDG Predicted"),
    title="albertmartinez/xlm-roberta-large-sdg-classification",
    description="Enter a text and see the text classification result!"
)

# Combine both interfaces into a tabbed interface
io = gr.TabbedInterface(
    interface_list=[ifaceall, iface1, iface2, iface3, iface4],
    tab_names=["ALL", "bert-sdg-classification", "bert-multilingual-sdg-classification", "distilbert-multilingual-sdg-classification", "xlm-roberta-large-sdg-classification"],
    title="Sustainable Development Goals (SDG) Text Classifier App",
    theme='base'
)

io.launch()