Update app.py

app.py

@@ -5,40 +5,37 @@ import torch
 
 app = FastAPI(
     title="Medical Text Classifier",
-    description="Simple medical text classifier using a lightweight BioBERT model.",
+    description="Classifier using Bio_ClinicalBERT",
     version="1.0"
 )
 
-MODEL_NAME = "d4data/biobert-v1.1-finetuned-MedICAL"
+MODEL_NAME = "emilyalsentzer/Bio_ClinicalBERT"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=4   # تقدر تغيّر عدد اللابلز
+)
 
-# Example labels (you can change them to cancer types)
-LABELS = [
-    "disease_related",
-    "treatment_related",
-    "test_related",
-    "symptom_related"
-]
+LABELS = ["disease", "treatment", "test", "symptom"]
 
 class Input(BaseModel):
     text: str
 
 @app.get("/")
 def home():
-    return {"status": "Medical classifier running successfully"}
+    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 = probs.argmax().item()
+    label_id = torch.argmax(probs).item()
 
     return {
         "input": data.text,
-        "predicted_label": LABELS[label_id] if label_id < len(LABELS) else label_id,
+        "predicted_label": LABELS[label_id],
         "confidence": float(probs.max())
     }