import re
import gradio as gr
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

MODEL_ID = "lihtmcad/roberta-clf-bearing"

tok = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)
model.eval()

def _split_sentences(text: str) -> list[str]:
    parts = re.split(r'(?<=[.!?])\s+', text.strip())
    return [s.strip() for s in parts if len(s.split()) >= 5]

@torch.no_grad()
def _score(text: str) -> float:
    enc = tok(text, return_tensors="pt", truncation=True, max_length=512)
    return round(torch.softmax(model(**enc).logits, dim=-1)[0][1].item(), 4)

def classify(text: str) -> dict:
    prob = _score(text)
    return {"ai_prob": prob, "label": "AI" if prob >= 0.5 else "Human"}

def attribute(text: str) -> dict:
    sents = _split_sentences(text)
    if not sents:
        return {"mean_ai_prob": None, "label": None, "sentences": []}
    scored = []
    for s in sents:
        p = _score(s)
        scored.append({"text": s, "ai_prob": p,
                       "risk": "high" if p >= 0.7 else
                               "medium" if p >= 0.4 else "low"})
    mean_p = round(sum(r["ai_prob"] for r in scored) / len(scored), 4)
    return {
        "mean_ai_prob": mean_p,
        "label": "AI" if mean_p >= 0.5 else "Human",
        "sentences": scored,
    }

with gr.Blocks(title="AI Text Classifier") as demo:
    gr.Markdown("## AI Text Classifier — Bearing/Tribology Domain\n"
                "`roberta-base` fine-tuned on domain academic sentence pairs.  \n"
                "API: `/gradio_api/call/predict` 段落判定 | `/gradio_api/call/attribute` 句子归因")

    with gr.Tab("段落判定"):
        t1 = gr.Textbox(lines=4, label="Input text")
        o1 = gr.JSON(label="Result")
        gr.Button("Classify").click(classify, t1, o1, api_name="predict")

    with gr.Tab("句子归因"):
        t2 = gr.Textbox(lines=6, label="Input paragraph")
        o2 = gr.JSON(label="Sentence-level attribution")
        gr.Button("Attribute").click(attribute, t2, o2, api_name="attribute")

demo.launch()