Create main.py

main.py

@@ -0,0 +1,202 @@
+import torch
+import numpy as np
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel
+from typing import List, Optional, Dict, Any
+
+from core.tensor_bridge import TensorBridge
+from core.amlm import load_amlm_model
+from core.aeg import load_aeg_model
+from core.triage import TriageEngineV2
+
+# ==========================================
+# 1. SETUP & CONFIG
+# ==========================================
+app = FastAPI(title="Clario Training Module", version="1.0.0")
+
+# Global Model Containers
+MODELS = {}
+DEVICE = "cpu" # Force CPU for Hugging Face Spaces (unless GPU upgraded)
+
+# Paths to weights (Uploaded in Phase 5)
+AMLM_PATH = "models/amlm_v5.pth"
+AEG_PATH = "models/aeg_v5.pth"
+
+# Activity Format Map (Must match AEG Training Data)
+FORMAT_MAP = [
+    "Auditory_Bombardment", "Phoneme_Isolation", "Minimal_Pairs",
+    "Syllable_Chaining", "Carrier_Phrases", "Pacing",
+    "Shadowing", "Speed_Drills"
+]
+STRUCTURE_MAP = ["CV", "VC", "CVC", "VCV", "CVCV", "CCV", "VCC", "Complex"]
+
+# ==========================================
+# 2. DATA MODELS (Pydantic)
+# ==========================================
+class TensorUpdatePayload(BaseModel):
+    current_tensor: Optional[List[List[List[float]]]] = None # 3x64x9 List
+    new_results: List[Dict[str, Any]] # From Backend or Practice
+    
+class UserStats(BaseModel):
+    severity: float # 0.0 - 1.0
+    fatigue: float # 0.0 - 1.0
+    age: float # Normalized 0.0 - 1.0
+    delta: float # Velocity (-0.5 to +0.5)
+
+class StrategyGeneratePayload(BaseModel):
+    current_tensor: List[List[List[float]]] # The raw observed tensor
+    user_stats: UserStats
+    force_easier: bool = False # The "Frustration Valve"
+
+# ==========================================
+# 3. LIFECYCLE EVENTS
+# ==========================================
+@app.on_event("startup")
+async def startup_event():
+    print("🚀 Initializing Clario Training Module...")
+    try:
+        # Load Cognitive Core
+        MODELS['amlm'] = load_amlm_model(AMLM_PATH, DEVICE)
+        MODELS['aeg'] = load_aeg_model(AEG_PATH, DEVICE)
+        MODELS['bridge'] = TensorBridge()
+        MODELS['triage'] = TriageEngineV2(DEVICE)
+        print("✅ All Systems Operational.")
+    except Exception as e:
+        print(f"❌ Startup Failed: {e}")
+
+# ==========================================
+# 4. HELPER FUNCTIONS
+# ==========================================
+def _prepare_amlm_input(raw_tensor_list):
+    """
+    Converts raw 3x64x9 list into Batchx6x64x9 Tensor for AMLM.
+    Generates visibility mask based on non-zero entries.
+    """
+    # 1. Convert to Numpy (3, 64, 9)
+    val_map = np.array(raw_tensor_list, dtype=np.float32)
+    
+    # 2. Generate Mask (Assumption: Val > 0 means "Observed Error")
+    # Zeros are treated as "Ambiguous/Untested" for the Inference Engine
+    vis_map = (val_map > 0.0).astype(np.float32)
+    
+    # 3. Concatenate (6, 64, 9)
+    input_numpy = np.concatenate([val_map, vis_map], axis=0)
+    
+    # 4. Add Batch Dim & Convert to Torch
+    return torch.FloatTensor(input_numpy).unsqueeze(0).to(DEVICE)
+
+# ==========================================
+# 5. ENDPOINTS
+# ==========================================
+
+@app.post("/tensor/update")
+async def update_tensor(payload: TensorUpdatePayload):
+    """
+    THE FAST LOOP: Updates the patient's digital twin with new error data.
+    """
+    bridge = MODELS['bridge']
+    
+    # Run Bridge Logic
+    updated_tensor = bridge.update_tensor(
+        payload.current_tensor,
+        payload.new_results
+    )
+    
+    return {
+        "status": "success",
+        "message": "Tensor updated successfully.",
+        "updated_tensor": updated_tensor
+    }
+
+@app.post("/strategy/generate")
+async def generate_strategy(payload: StrategyGeneratePayload):
+    """
+    THE SLOW LOOP: Generates a full therapy session blueprint.
+    """
+    amlm = MODELS['amlm']
+    aeg = MODELS['aeg']
+    triage = MODELS['triage']
+    
+    # 1. AMLM Inference (Fill Gaps)
+    # We use the raw tensor to infer the full clinical image
+    amlm_input = _prepare_amlm_input(payload.current_tensor)
+    with torch.no_grad():
+        diagnosis_map, _ = amlm(amlm_input)
+        # Squeeze batch dim -> (3, 64, 9) Numpy
+        diag_np = diagnosis_map[0].cpu().numpy()
+        
+    # 2. Triage (Find Focus)
+    all_diagnoses = triage.scan_tensor(diag_np)
+    
+    if not all_diagnoses:
+        return {"status": "healthy", "message": "No significant errors detected."}
+        
+    focus_diag = all_diagnoses[0] # Highest priority issue
+    
+    # 3. Prepare Stats (Frustration Valve Logic)
+    stats = payload.user_stats
+    effective_fatigue = 1.0 if payload.force_easier else stats.fatigue
+    
+    # 4. AEG Input Preparation
+    aeg_inputs = triage.prepare_aeg_input(
+        focus_diag, 
+        all_diagnoses, 
+        [stats.severity, effective_fatigue, stats.age, stats.delta]
+    )
+    
+    # 5. AEG Inference (Generate Blueprint)
+    with torch.no_grad():
+        p_vol, p_diff, p_fmt, p_targ, p_struct, p_forbid, p_dyn = aeg(
+            aeg_inputs['error_idx'], aeg_inputs['major_idx'], aeg_inputs['pos_idx'],
+            aeg_inputs['cluster'], aeg_inputs['trap'], aeg_inputs['feat'],
+            aeg_inputs['flags'], aeg_inputs['stats']
+        )
+        
+    # 6. Decode Outputs
+    vol = int(torch.clamp(p_vol, 3, 10).item())
+    
+    # Extract Sequences (Slice by volume)
+    diff_seq = torch.argmax(p_diff, dim=2)[0].tolist()[:vol]
+    # Difficulty is 0-4 internally, map to 1-5 for UI
+    diff_seq_ui = [d + 1 for d in diff_seq]
+    
+    fmt_indices = torch.argmax(p_fmt, dim=2)[0].tolist()[:vol]
+    fmt_names = [FORMAT_MAP[i] for i in fmt_indices]
+    
+    struct_indices = torch.argmax(p_struct, dim=2)[0].tolist()[:vol]
+    struct_names = [STRUCTURE_MAP[i] for i in struct_indices]
+    
+    # Target Phonemes (Multi-hot)
+    target_logits = p_targ[0][:vol] # (Vol, 64)
+    target_phonemes = []
+    
+    from core.triage import PHONEMES # Helper to map back
+    
+    for step_logits in target_logits:
+        # Sigmoid > 0.5 implies active target
+        active_indices = (torch.sigmoid(step_logits) > 0.5).nonzero(as_tuple=True)[0].tolist()
+        targets = [PHONEMES[i] for i in active_indices]
+        if not targets: targets = focus_diag.active_phonemes # Fallback to main targets
+        target_phonemes.append(targets)
+
+    return {
+        "status": "generated",
+        "focus_diagnosis": {
+            "name": focus_diag.error_name,
+            "type": focus_diag.major_type,
+            "severity": focus_diag.severity
+        },
+        "blueprint": {
+            "volume": vol,
+            "plan": [
+                {
+                    "step": i+1,
+                    "difficulty": diff_seq_ui[i],
+                    "format": fmt_names[i],
+                    "structure": struct_names[i],
+                    "targets": target_phonemes[i]
+                }
+                for i in range(vol)
+            ]
+        }
+    }