from transformers.pipelines import pipeline
import gradio as gr
from huggingface_hub import login
import os
login(os.environ['HF_Token'])
predictor = pipeline(
    "sentences_sim",
    model="minskiter/simbert-chinese-bert-wwm-ext",
    device="cpu",
    trust_remote_code=True,
    use_auth_token=True
)

def sim_predictor_gradio(text_a,text_b):
    sim = predictor((text_a, text_b))
    return {text_b[:10]+"...": sim}

demo = gr.Interface(
    fn=sim_predictor_gradio, 
    inputs=[
        gr.Textbox(lines=5, label="文本1"),
        gr.Textbox(lines=5, label="文本2")
    ], 
    outputs=gr.Label(num_top_classes=1, label="相似度")
)
demo.launch()