import gradio as gr
from transformers import pipeline

# Define the models
model = pipeline("text-classification",
                 model="OpenAlex/bert-base-multilingual-cased-finetuned-openalex-topic-classification-title-abstract")

model2 = pipeline("text-classification",
                 model="albertmartinez/openalex-topic-classification-title-abstract")


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

demo = gr.Interface(
    fn=classify_text,
    inputs=[gr.Textbox(lines=5, label="Text", placeholder="<TITLE> {title}\n<ABSTRACT> {abstract}",
                       value="<TITLE> {title}\n<ABSTRACT> {abstract}"),
            gr.Number(label="top_k", value=10, precision=0)],
    outputs=[gr.Label(label="Model 1: OpenAlex"),
             gr.Label(label="Model 2: AlbertMartinez")],
    title="OpenAlex Topic Classification",
    description="Enter a text and see the topic classification result!",
    flagging_mode="never",
    api_name="classify"
)

if __name__ == "__main__":
    print(gr.__version__)
    demo.launch()