Update app.py

app.py

@@ -3,48 +3,42 @@ from pydantic import BaseModel
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
-app = FastAPI()
+app = FastAPI(
+    title="Medical Text Classifier",
+    description="Simple medical text classifier using a lightweight BioBERT model.",
+    version="1.0"
+)
 
-MODEL_NAME = "monologg/distilbiobert"
+MODEL_NAME = "d4data/biobert-v1.1-finetuned-MedICAL"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(
-    MODEL_NAME,
-    num_labels=12  # عدد أنواع السرطان (أنت تتحكم به)
-)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
 
+# Example labels (you can change them to cancer types)
 LABELS = [
- "breast_cancer",
- "lung_cancer",
- "prostate_cancer",
- "colon_cancer",
- "lymphoma",
- "melanoma",
- "thyroid_cancer",
- "kidney_cancer",
- "pancreatic_cancer",
- "ovarian_cancer",
- "cervical_cancer",
- "brain_tumor"
+    "disease_related",
+    "treatment_related",
+    "test_related",
+    "symptom_related"
 ]
 
 class Input(BaseModel):
     text: str
 
+@app.get("/")
+def home():
+    return {"status": "Medical classifier running successfully"}
+
 @app.post("/predict")
 def predict(data: Input):
     inputs = tokenizer(data.text, return_tensors="pt", truncation=True)
     outputs = model(**inputs)
-
-    probs = torch.nn.functional.softmax(outputs.logits, dim=1)
-    label_id = torch.argmax(probs).item()
-    confidence = float(torch.max(probs))
+    
+    probs = torch.softmax(outputs.logits, dim=1)
+    label_id = probs.argmax().item()
 
     return {
-        "prediction": LABELS[label_id],
-        "confidence": confidence
+        "input": data.text,
+        "predicted_label": LABELS[label_id] if label_id < len(LABELS) else label_id,
+        "confidence": float(probs.max())
     }
-
-@app.get("/")
-def home():
-    return {"status": "Model is running"}