from fastapi import FastAPI
from pydantic import BaseModel
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

app = FastAPI(
    title="Medical Text Classifier",
    description="Classifier using Bio_ClinicalBERT",
    version="1.0"
)

MODEL_NAME = "emilyalsentzer/Bio_ClinicalBERT"

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=4   # تقدر تغيّر عدد اللابلز
)

LABELS = ["disease", "treatment", "test", "symptom"]

class Input(BaseModel):
    text: str

@app.get("/")
def home():
    return {"status": "Bio_ClinicalBERT is running"}

@app.post("/predict")
def predict(data: Input):
    inputs = tokenizer(data.text, return_tensors="pt", truncation=True)
    outputs = model(**inputs)

    probs = torch.softmax(outputs.logits, dim=1)
    label_id = torch.argmax(probs).item()

    return {
        "input": data.text,
        "predicted_label": LABELS[label_id],
        "confidence": float(probs.max())
    }