Spaces:
Sleeping
Sleeping
| 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 | |
| # ========================================== | |
| 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 | |
| # ========================================== | |
| 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 | |
| } | |
| 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) | |
| ] | |
| } | |
| } |