import gradio as gr
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from transformers import pipeline, AutoTokenizer, AutoModelForTokenClassification

# 正確載入 safetensors 模型
model = pipeline(
    "ner",
    model=AutoModelForTokenClassification.from_pretrained("ckiplab/bert-base-chinese-ner", use_safetensors=True),
    tokenizer=AutoTokenizer.from_pretrained("ckiplab/bert-base-chinese-ner"),
    aggregation_strategy="simple"
)

def analyze(sentence: str):
    result = model(sentence)
    return " ".join([f"{r['word']}/{r['entity_group']}" for r in result])

demo = gr.Interface(fn=analyze, inputs="text", outputs="text", title="實體辨識")

app = FastAPI()
app = gr.mount_gradio_app(app, demo, path="/")

@app.post("/analyze")
async def api_analyze(request: Request):
    payload = await request.json()
    return JSONResponse(content={"result": analyze(payload.get("sentence", ""))})