|
import gradio as gr |
|
from transformers import pipeline |
|
|
|
|
|
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() |
|
|