Update app.py

app.py

@@ -4,8 +4,11 @@ from transformers import AutoTokenizer, BigBirdForSequenceClassification
 from scipy.special import softmax
 import torch
 
+# Import AJCC API logic
+from ajcc_api import stage_cancer
+
 # Initialize FastAPI
-app = FastAPI(title="TNM Endpoint", version="1.0")
+app = FastAPI(title="TNM + AJCC Endpoint", version="2.0")
 
 # Models (TNM) from Hugging Face
 MODEL_T = "jkefeli/CancerStage_Classifier_T"
@@ -15,41 +18,52 @@ MODEL_M = "jkefeli/CancerStage_Classifier_M"
 # Load tokenizer once
 tokenizer = AutoTokenizer.from_pretrained("yikuan8/Clinical-BigBird")
 
-# Load models once (CPU mode)
+# Load models once
 model_T = BigBirdForSequenceClassification.from_pretrained(MODEL_T)
 model_N = BigBirdForSequenceClassification.from_pretrained(MODEL_N)
 model_M = BigBirdForSequenceClassification.from_pretrained(MODEL_M)
 
 class Report(BaseModel):
     text: str
+    cancer_type: str = "colon"   # default, you can change it (breast, lung, etc.)
 
-def predict_stage(text, model):
+def predict_class(text, model):
     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=2048)
     with torch.no_grad():
         outputs = model(**inputs)
         probs = softmax(outputs.logits.numpy(), axis=1)
         pred_class = probs.argmax(axis=1)[0]
-    return {"class": int(pred_class), "probs": probs.tolist()}
+    return int(pred_class)
 
 @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(report: Report = Body(...)):
+@app.post("/predict_full")
+def predict_full(report: Report = Body(...)):
     text = report.text
+    cancer = report.cancer_type.lower()
+
     try:
-        t_result = predict_stage(text, model_T)
-        n_result = predict_stage(text, model_N)
-        m_result = predict_stage(text, model_M)
+        # 1) Predict numeric TNM classes
+        t_class = predict_class(text, model_T)
+        n_class = predict_class(text, model_N)
+        m_class = predict_class(text, model_M)
+
+        # 2) Convert numeric classes → TNM labels
+        T = f"T{t_class}"
+        N = f"N{n_class}"
+        M = f"M{m_class}"
+
+        # 3) Compute AJCC Stage
+        stage = stage_cancer(cancer_type=cancer, T=T, N=N, M=M)
 
         return {
             "input": text,
-            "TNM_prediction": {
-                "T": t_result,
-                "N": n_result,
-                "M": m_result
-            }
+            "TNM_prediction": {"T": T, "N": N, "M": M},
+            "AJCC_stage": stage,
+            "cancer_type_used": cancer
         }
+
     except Exception as e:
         return {"error": str(e)}