Update app.py

app.py

@@ -1,100 +1,223 @@
 """
-model.py — DistilBERT multi-head model definition and loader
+Workout Coach — FastAPI Inference App
+Runs DistilBERT classification + Claude debrief generation
+Designed for Hugging Face Spaces with Docker
 """
 
+from fastapi import FastAPI, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel, Field
+from contextlib import asynccontextmanager
+from typing import Optional, Dict
 import torch
-import torch.nn as nn
-from transformers import DistilBertModel, DistilBertTokenizer
+import anthropic
+import os
 import logging
 
+from model import MultiHeadDistilBERT, load_model
+from inference import predict, decode_predictions, build_prompt
+
+# ─────────────────────────────────────────────
+# LOGGING
+# ─────────────────────────────────────────────
+
+logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
+# ─────────────────────────────────────────────
+# LABEL MAPS
+# ─────────────────────────────────────────────
+
+MOOD_MAP = {
+    0: "accomplished", 1: "anxious",    2: "distracted",
+    3: "energized",    4: "fatigued",   5: "frustrated",
+    6: "neutral",      7: "positive",
+}
+
+EXERTION_MAP = {0: "low", 1: "moderate", 2: "high"}
+
+COMPLETION_MAP = {0: "partial", 1: "full"}
+
+SORENESS_MAP = {
+    0:  "none",
+    1:  "biceps_mild",       2:  "biceps_moderate",
+    3:  "back_mild",         4:  "back_moderate",     5:  "back_severe",
+    6:  "chest_mild",        7:  "chest_moderate",    8:  "chest_severe",
+    9:  "legs_mild",         10: "legs_moderate",     11: "legs_severe",
+    12: "shoulder_mild",     13: "shoulder_moderate", 14: "shoulder_severe",
+    15: "triceps_mild",      16: "triceps_moderate",
+}
+
+# ─────────────────────────────────────────────
+# APP STATE — model loaded once at startup
+# ─────────────────────────────────────────────
+
+app_state = {}
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Load model and tokenizer once at startup, clean up at shutdown."""
+    logger.info("Loading DistilBERT model...")
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logger.info(f"Using device: {device}")
+
+    model, tokenizer = load_model(
+        model_path=os.getenv("MODEL_PATH", "best_overall_model.pt"),
+        device=device,
+    )
 
-class MultiHeadDistilBERT(nn.Module):
-    """
-    Multi-head DistilBERT classifier.
-    4 independent classification heads sharing one BERT backbone:
-      - mood       (8 classes)
-      - exertion   (3 classes)
-      - soreness   (17 classes — combined region + severity)
-      - completion (2 classes)
-    """
+    app_state["model"]     = model
+    app_state["tokenizer"] = tokenizer
+    app_state["device"]    = device
 
-    def __init__(
-        self,
-        num_moods:              int = 8,
-        num_exertion_levels:    int = 3,
-        num_soreness_classes:   int = 17,
-        num_completion_statuses:int = 2,
-    ):
-        super().__init__()
-
-        self.bert       = DistilBertModel.from_pretrained("distilbert-base-uncased")
-        hidden_size     = self.bert.config.hidden_size  # 768
-
-        self.dropout      = nn.Dropout(0.3)
-        self.head_dropout = nn.Dropout(0.1)
-
-        # Simple heads for easy tasks
-        self.mood_head       = nn.Linear(hidden_size, num_moods)
-        self.completion_head = nn.Linear(hidden_size, num_completion_statuses)
-
-        # Deeper head for exertion
-        self.exertion_head = nn.Sequential(
-            nn.Linear(hidden_size, 128),
-            nn.ReLU(),
-            nn.Dropout(0.2),
-            nn.Linear(128, num_exertion_levels),
-        )
+    # Anthropic client — reads ANTHROPIC_API_KEY from environment
+    app_state["anthropic_client"] = anthropic.Anthropic(
+        api_key=os.getenv("ANTHROPIC_API_KEY")
+    )
 
-        # Deeper head for soreness (hardest task — 17 classes)
-        self.soreness_head = nn.Sequential(
-            nn.Linear(hidden_size, 256),
-            nn.ReLU(),
-            nn.Dropout(0.3),
-            nn.Linear(256, num_soreness_classes),
-        )
+    logger.info("Model and clients loaded successfully.")
+    yield
 
-    def forward(self, input_ids, attention_mask):
-        outputs    = self.bert(input_ids=input_ids, attention_mask=attention_mask)
-        cls_output = self.dropout(outputs.last_hidden_state[:, 0, :])
-        x          = self.head_dropout(cls_output)
+    # Cleanup
+    app_state.clear()
+    logger.info("App shutdown complete.")
 
-        return (
-            self.mood_head(x),
-            self.exertion_head(x),
-            self.soreness_head(x),
-            self.completion_head(x),
-        )
+
+app = FastAPI(
+    title="Workout Coach Inference API",
+    description="DistilBERT classification + Claude debrief generation",
+    version="1.0.0",
+    lifespan=lifespan,
+)
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# ─────────────────────────────────────────────
+# REQUEST / RESPONSE SCHEMAS
+# ─────────────────────────────────────────────
+
+class SessionRequest(BaseModel):
+    # Free-text input from the user — fed into DistilBERT
+    user_text: str = Field(..., min_length=5, max_length=500,
+                           example="That was really tough, chest is killing me but I feel accomplished.")
+
+    # UI form fields — collected separately in the app
+    duration_minutes: int   = Field(..., ge=1, le=300, example=45)
+    workout_type:     str   = Field(..., example="upper_body_push")
+    user_goal:        str   = Field(..., example="muscle_gain")
+
+    # Optional — whether to generate the Claude debrief
+    generate_debrief: bool  = Field(default=True)
 
 
-def load_model(
-    model_path: str,
-    device: torch.device,
-    num_moods:               int = 8,
-    num_exertion_levels:     int = 3,
-    num_soreness_classes:    int = 17,
-    num_completion_statuses: int = 2,
-):
+class BertLabels(BaseModel):
+    mood:       str
+    exertion:   str
+    soreness:   str
+    completion: str
+
+
+class SessionResponse(BaseModel):
+    bert_labels: BertLabels
+    debrief:     Optional[str] = None
+
+
+class HealthResponse(BaseModel):
+    # model_config suppresses Pydantic's warning about field names
+    # starting with "model_" conflicting with its protected namespace
+    model_config = {"protected_namespaces": ()}
+
+    status:       str
+    device:       str
+    model_loaded: bool
+
+# ─────────────────────────────────────────────
+# ROUTES
+# ─────────────────────────────────────────────
+
+@app.get("/health", response_model=HealthResponse)
+def health():
+    """Health check — confirms model is loaded and ready."""
+    return {
+        "status":       "ok",
+        "device":       str(app_state.get("device", "unknown")),
+        "model_loaded": "model" in app_state,
+    }
+
+
+@app.post("/classify", response_model=SessionResponse)
+def classify_session(req: SessionRequest):
     """
-    Instantiate the model, load saved weights, set to eval mode.
-    Returns (model, tokenizer).
+    Runs DistilBERT inference on user_text and optionally
+    generates a Claude debrief using the classified labels
+    combined with the session form data.
     """
-    logger.info(f"Loading model weights from: {model_path}")
-
-    model = MultiHeadDistilBERT(
-        num_moods=num_moods,
-        num_exertion_levels=num_exertion_levels,
-        num_soreness_classes=num_soreness_classes,
-        num_completion_statuses=num_completion_statuses,
+    model     = app_state["model"]
+    tokenizer = app_state["tokenizer"]
+    device    = app_state["device"]
+    client    = app_state["anthropic_client"]
+
+    # ── Step 1: DistilBERT inference ─────────────────────────
+    try:
+        raw_preds = predict(req.user_text, model, tokenizer, device)
+    except Exception as e:
+        logger.error(f"Inference error: {e}")
+        raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}")
+
+    # ── Step 2: Decode integer labels → strings ───────────────
+    bert_labels = decode_predictions(
+        raw_preds, MOOD_MAP, EXERTION_MAP, SORENESS_MAP, COMPLETION_MAP
     )
 
-    state_dict = torch.load(model_path, map_location=device, weights_only=True)
-    model.load_state_dict(state_dict)
-    model.to(device)
-    model.eval()
+    # ── Step 3: Optionally generate Claude debrief ────────────
+    debrief = None
+    if req.generate_debrief:
+        prompt = build_prompt(
+            bert_labels=bert_labels,
+            user_text=req.user_text,
+            duration_minutes=req.duration_minutes,
+            workout_type=req.workout_type,
+            user_goal=req.user_goal,
+        )
+        try:
+            message = client.messages.create(
+                model="claude-sonnet-4-6",
+                max_tokens=400,
+                messages=[{"role": "user", "content": prompt}],
+            )
+            debrief = message.content[0].text
+        except Exception as e:
+            logger.error(f"Claude API error: {e}")
+            # Debrief failure is non-fatal — return labels without debrief
+            debrief = None
+
+    return SessionResponse(
+        bert_labels=BertLabels(**bert_labels),
+        debrief=debrief,
+    )
 
-    tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased")
 
-    logger.info("Model loaded and set to eval mode.")
-    return model, tokenizer
+@app.post("/classify/labels-only", response_model=BertLabels)
+def classify_labels_only(req: SessionRequest):
+    """
+    Runs only DistilBERT inference. Skips Claude.
+    Useful for storing labels to DB without generating a debrief yet.
+    """
+    model     = app_state["model"]
+    tokenizer = app_state["tokenizer"]
+    device    = app_state["device"]
+
+    try:
+        raw_preds   = predict(req.user_text, model, tokenizer, device)
+        bert_labels = decode_predictions(
+            raw_preds, MOOD_MAP, EXERTION_MAP, SORENESS_MAP, COMPLETION_MAP
+        )
+        return BertLabels(**bert_labels)
+    except Exception as e:
+        logger.error(f"Inference error: {e}")
+        raise HTTPException(status_code=500, detail=f"Inference failed: {str(e)}")