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")


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


with gr.Blocks() as 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="openalex topic predicted"),
        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()