Spaces:
Sleeping
Sleeping
File size: 5,738 Bytes
7e6d950 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | import torch
import torch.nn as nn
import os
# ==========================================
# 1. CONFIGURATION (V5.4 Standard)
# ==========================================
AEG_CONFIG = {
"num_phonemes": 64,
"num_errors": 9,
"num_major_types": 3,
"num_positions": 3,
"num_structures": 8,
"max_seq_len": 10,
"hidden_size": 128
}
class AEGPolicyNetwork_V5_4(nn.Module):
def __init__(self):
super(AEGPolicyNetwork_V5_4, self).__init__()
# Embeddings for categorical inputs
self.emb_error = nn.Embedding(AEG_CONFIG['num_errors'], 8)
self.emb_major = nn.Embedding(AEG_CONFIG['num_major_types'], 4)
self.emb_pos = nn.Embedding(AEG_CONFIG['num_positions'], 4)
# Input Dimension Calculation:
# Cluster(64) + Trap(64) + Feat(8) + Emb_Err(8) + Emb_Maj(4) + Emb_Pos(4) + Flags(6) + Stats(4) = 162
input_dim = 64 + 64 + 8 + 8 + 4 + 4 + 6 + 4
# Shared Encoder (Patient Embedding)
self.encoder = nn.Sequential(
nn.Linear(input_dim, 256),
nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.2),
nn.Linear(256, AEG_CONFIG['hidden_size']), nn.ReLU()
)
# Head A: Volume (Scalar)
self.head_volume = nn.Linear(AEG_CONFIG['hidden_size'], 1)
# Head B: Difficulty Progression (LSTM Sequence)
self.lstm_diff = nn.LSTM(input_size=AEG_CONFIG['hidden_size'], hidden_size=64, batch_first=True)
self.fc_diff = nn.Linear(64, 5) # 5 Difficulty levels
self.emb_diff_cond = nn.Embedding(5, 8) # Embedding for the predicted difficulty to feed into other heads
# Head D: Format Selector
self.head_format = nn.Sequential(
nn.Linear(AEG_CONFIG['hidden_size'] + 8 + 8 + 4, 64),
nn.ReLU(),
nn.Linear(64, 8) # 8 Activity Formats
)
# Head E: Targets (Residual Anchor)
self.head_targets = nn.Sequential(
nn.Linear(64 + 8 + 64, 128),
nn.ReLU(),
nn.Linear(128, 64)
)
# Head F: Structure Scaffolding
self.head_struct = nn.Sequential(
nn.Linear(64 + 8 + 8, 64),
nn.ReLU(),
nn.Linear(64, AEG_CONFIG['num_structures'])
)
# Head C: Safety/Forbidden (Residual Bypass)
self.head_forbidden = nn.Linear(AEG_CONFIG['hidden_size'], 64)
# Head C2: Dynamic Parameters
self.head_dynamic = nn.Linear(AEG_CONFIG['hidden_size'], 5)
def forward(self, error_idx, major_idx, pos_idx, cluster, trap, feat, flags, stats):
# 1. Embed Categoricals
e_err = self.emb_error(error_idx)
e_maj = self.emb_major(major_idx)
e_pos = self.emb_pos(pos_idx)
# 2. Shared Encoding
x_cat = torch.cat([cluster, trap, feat, e_err, e_maj, e_pos, flags, stats], dim=1)
h_shared = self.encoder(x_cat)
# 3. Head A: Volume
pred_vol = self.head_volume(h_shared)
# 4. Head B: Difficulty (LSTM)
# Repeat shared vector for sequence length
lstm_input = h_shared.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
lstm_out, _ = self.lstm_diff(lstm_input)
pred_diff_logits = self.fc_diff(lstm_out)
# Get predicted difficulty indices to condition other heads
cond_diff = torch.argmax(pred_diff_logits, dim=2)
cond_vec = self.emb_diff_cond(cond_diff)
# Prepare repeated embeddings for sequence generation
h_rep = h_shared.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
e_err_rep = e_err.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
e_maj_rep = e_maj.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
cluster_rep = cluster.unsqueeze(1).repeat(1, AEG_CONFIG['max_seq_len'], 1)
# 5. Head D: Format
fmt_input = torch.cat([h_rep, cond_vec, e_err_rep, e_maj_rep], dim=2)
pred_fmt = self.head_format(fmt_input)
# 6. Head E: Targets (With Residual Anchor)
# Output = Net(x) + (Input_Cluster * 5.0)
target_input = torch.cat([lstm_out, cond_vec, cluster_rep], dim=2)
pred_targ = self.head_targets(target_input) + (cluster_rep * 5.0)
# 7. Head F: Structure
struct_input = torch.cat([lstm_out, cond_vec, e_err_rep], dim=2)
pred_struct = self.head_struct(struct_input)
# 8. Head C: Forbidden (With Residual Bypass)
# Output = Net(x) + (Input_Trap * 10.0)
pred_forbid = self.head_forbidden(h_shared) + (trap * 10.0)
# 9. Head C2: Dynamic
pred_dyn = self.head_dynamic(h_shared)
return pred_vol, pred_diff_logits, pred_fmt, pred_targ, pred_struct, pred_forbid, pred_dyn
# ==========================================
# 2. LOADER UTILITY
# ==========================================
def load_aeg_model(model_path, device):
"""
Initializes the model and loads weights safely.
"""
print(f"--- Loading AEG V5.4 from {model_path} ---")
model = AEGPolicyNetwork_V5_4().to(device)
if os.path.exists(model_path):
try:
state_dict = torch.load(model_path, map_location=device)
model.load_state_dict(state_dict)
model.eval() # Set to inference mode
print("✅ AEG Weights Loaded Successfully")
except Exception as e:
print(f"❌ CRITICAL: Failed to load AEG weights: {e}")
raise e
else:
print(f"⚠️ WARNING: Model file not found at {model_path}. Running with random weights (Debugging Mode).")
return model |