| from __future__ import annotations |
| import hashlib, os, time, uuid |
| from pathlib import Path |
| from fastapi import FastAPI, HTTPException |
| from pydantic import BaseModel, Field |
| import torch |
| from model import OncoCT1B, OncoCTConfig |
|
|
| app = FastAPI(title="OncoCT-1B Research Inference", version="0.1.0") |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| MODEL = OncoCT1B(OncoCTConfig()).to(DEVICE).eval() |
| CHECKPOINT = os.getenv("ONCOCT_CHECKPOINT") |
| if CHECKPOINT and Path(CHECKPOINT).exists(): |
| state = torch.load(CHECKPOINT, map_location=DEVICE, weights_only=False) |
| MODEL.load_state_dict(state.get("model", state), strict=False) |
|
|
| class InferenceRequest(BaseModel): |
| volume_tokens_path: str = Field(..., description="Path to a preprocessed research tensor; do not expose raw filesystem paths in production") |
| study_time: float = 0.0 |
| request_id: str | None = None |
|
|
| @app.get("/health") |
| def health(): return {"status": "ok", "device": str(DEVICE), "model_loaded": CHECKPOINT is not None} |
|
|
| @app.post("/infer") |
| def infer(req: InferenceRequest): |
| rid = req.request_id or str(uuid.uuid4()) |
| started = time.time() |
| p = Path(req.volume_tokens_path) |
| if not p.exists(): raise HTTPException(404, "tensor input not found") |
| item = torch.load(p, map_location="cpu", weights_only=True) |
| tokens = item["volume_tokens"].float().unsqueeze(0).to(DEVICE) |
| t = torch.tensor([[req.study_time]], device=DEVICE) |
| with torch.inference_mode(): out = MODEL(tokens, t) |
| abstain = float(torch.sigmoid(out["abstain_logit"])[0].cpu()) |
| response_prob = torch.softmax(out["response_logits"], -1)[0].cpu().tolist() |
| audit_hash = hashlib.sha256(f"{rid}|{p}|{time.time_ns()}".encode()).hexdigest() |
| return {"request_id": rid, "status": "abstain" if abstain >= 0.5 else "review_required", "abstention_score": abstain, "response_probabilities": response_prob, "evidence": {"volume_log1p": float(out["volume_log1p"][0].cpu()), "long_axis_mm": float(out["long_axis_mm"][0].cpu())}, "latency_ms": (time.time()-started)*1000, "audit_hash": audit_hash, "clinical_notice": "Research output only. Requires qualified radiologist/oncologist review; does not diagnose or recommend treatment."} |
|
|
