Create app.py

app.py

@@ -1,75 +1,56 @@
-
-import os
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
+from fastapi import FastAPI, Body
 from pydantic import BaseModel
 from huggingface_hub import InferenceClient
+import os
+
+# FastAPI app
+app = FastAPI(title="TNM Cancer Stage Endpoint", version="1.0")
+
+# HF token from environment (set in Space secrets)
+HF_TOKEN = os.getenv("HF_TOKEN")
 
-# === Config ===
-HF_TOKEN = os.getenv("HF_TOKEN")  # add this as a Secret in your Space
+# Models (يمكنك تعديلها من Variables في Hugging Face Space)
 MODEL_T = os.getenv("MODEL_T", "jkefeli/CancerStage_Classifier_T")
 MODEL_N = os.getenv("MODEL_N", "jkefeli/CancerStage_Classifier_N")
 MODEL_M = os.getenv("MODEL_M", "jkefeli/CancerStage_Classifier_M")
 
-# Initialize FastAPI
-app = FastAPI(title="TNM Endpoint", version="1.0.0")
+# Hugging Face client
+client = InferenceClient(token=HF_TOKEN)
 
-# CORS (optional): allow all origins by default; tighten for production
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-# Init Hugging Face Inference clients (remote inference, no heavy compute on your Space)
-clients = {
-    "T": InferenceClient(model=MODEL_T, token=HF_TOKEN),
-    "N": InferenceClient(model=MODEL_N, token=HF_TOKEN),
-    "M": InferenceClient(model=MODEL_M, token=HF_TOKEN),
-}
-
-class ReportInput(BaseModel):
+class Report(BaseModel):
     text: str
 
-@app.get("/healthz")
-def healthz():
-    return {"status": "ok", "models": {"T": MODEL_T, "N": MODEL_N, "M": MODEL_M}}
-
-def _classify(client: InferenceClient, text: str):
-    # Uses HF Inference API task: text-classification
-    # Returns best label with score
-    try:
-        outputs = client.text_classification(text, wait_for_model=True)
-        if not outputs:
-            raise ValueError("Empty response from model")
-        best = max(outputs, key=lambda x: x.get("score", 0))
-        return {"label": best.get("label", "UNKNOWN"), "score": float(best.get("score", 0.0)), "raw": outputs}
-    except Exception as e:
-        raise HTTPException(status_code=502, detail=f"Inference error: {e}")
+@app.get("/")
+def health_check():
+    return {
+        "status": "running",
+        "models": {"T": MODEL_T, "N": MODEL_N, "M": MODEL_M}
+    }
 
 @app.post("/predict_tnm")
-def predict_tnm(input: ReportInput):
-    text = (input.text or "").strip()
-    if not text:
-        raise HTTPException(status_code=400, detail="Empty 'text'")
-    # Optionally hard-truncate very long inputs to avoid API limits
-    if len(text) > 20000:
-        text = text[:20000]
-
-    preds = {}
-    for key, client in clients.items():
-        preds[key] = _classify(client, text)
+def predict_tnm(report: Report = Body(...)):
+    text = report.text
 
-    t = preds["T"]["label"]
-    n = preds["N"]["label"]
-    m = preds["M"]["label"]
-    tnm_string = f"{t} {n} {m}"
+    try:
+        # Call each model
+        pred_T = client.text_classification(text, model=MODEL_T)
+        pred_N = client.text_classification(text, model=MODEL_N)
+        pred_M = client.text_classification(text, model=MODEL_M)
+
+        # Extract top labels
+        t_label = pred_T[0]["label"] if pred_T else None
+        n_label = pred_N[0]["label"] if pred_N else None
+        m_label = pred_M[0]["label"] if pred_M else None
+
+        return {
+            "input_chars": len(text),
+            "tnm": {
+                "T": pred_T[0] if pred_T else None,
+                "N": pred_N[0] if pred_N else None,
+                "M": pred_M[0] if pred_M else None,
+            },
+            "tnm_string": f"{t_label} {n_label} {m_label}"
+        }
 
-    return {
-        "input_chars": len(text),
-        "tnm": preds,
-        "tnm_string": tnm_string,
-        "meta": {"models": {"T": MODEL_T, "N": MODEL_N, "M": MODEL_M}}
-    }
+    except Exception as e:
+        return {"error": str(e)}