import gradio as gr
from transformers import pipeline

# 加载模型
print("正在加载病理检测NER模型...")
ner = pipeline(
    "token-classification", 
    model="OpenMed/OpenMed-NER-PathologyDetect-BigMed-560M",
    aggregation_strategy="max"
)
print("模型加载完成!")

# 处理函数
def process_text(text):
    if not text:
        return "请输入医学文本"
        
    results = ner(text)
    output = ""
    
    for result in results:
        entity = result["entity_group"]
        word = result["word"]
        score = round(result["score"], 2)
        output += f"检测到病理实体: {word} (类型: {entity}, 置信度: {score})\n"
        
    if not output:
        output = "未检测到任何病理相关实体"
        
    return output

# 创建界面
demo = gr.Interface(
    fn=process_text,
    inputs=gr.Textbox(placeholder="请输入医学文本...", lines=5),
    outputs="text",
    title="OpenMed 病理检测 NER 模型演示",
    description="使用OpenMed-NER-PathologyDetect-BigMed-560M模型识别文本中的病理实体"
)

# 启动服务
demo.launch()