Update app.py

app.py

@@ -5,10 +5,13 @@ import torch
 
 app = FastAPI()
 
-MODEL_NAME = "kamalkraj/distilbiobert_cancer_classification"
+MODEL_NAME = "monologg/distilbiobert"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
+model = AutoModelForSequenceClassification.from_pretrained(
+    MODEL_NAME,
+    num_labels=12  # عدد أنواع السرطان (أنت تتحكم به)
+)
 
 LABELS = [
  "breast_cancer",
@@ -28,10 +31,6 @@ LABELS = [
 class Input(BaseModel):
     text: str
 
-@app.get("/")
-def home():
-    return {"status": "Cancer Type Classifier is running!"}
-
 @app.post("/predict")
 def predict(data: Input):
     inputs = tokenizer(data.text, return_tensors="pt", truncation=True)
@@ -45,3 +44,7 @@ def predict(data: Input):
         "prediction": LABELS[label_id],
         "confidence": confidence
     }
+
+@app.get("/")
+def home():
+    return {"status": "Model is running"}