Upload AuriStream base model code

modeling_auristream.py

@@ -251,13 +251,11 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
         super().__init__(config)
         self.config = config
         
-        # Transformer components
-        self.transformer = nn.ModuleDict(dict(
-            wte=nn.Embedding(config.vocab_size, config.n_embd),
-            drop=nn.Dropout(config.dropout),
-            h=nn.ModuleList([Block(config) for _ in range(config.n_layer)]),
-            ln_f=RMSNorm(config.n_embd, bias=config.bias),
-        ))
+        # Transformer components (no wrapper to match weight keys)
+        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
+        self.drop = nn.Dropout(config.dropout)
+        self.h = nn.ModuleList([Block(config) for _ in range(config.n_layer)])
+        self.ln_f = RMSNorm(config.n_embd, bias=config.bias)
         
         # Multi-token prediction heads
         if hasattr(config, 'n_pred_steps') and config.n_pred_steps > 1:
@@ -269,7 +267,7 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
             self.future_heads = None
         
         # Output head
-        self.coch_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
+        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
         
         # Initialize weights
         self.apply(self._init_weights)
@@ -279,10 +277,10 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
                 torch.nn.init.normal_(p, mean=0.0, std=0.02/math.sqrt(2 * config.n_layer))
 
     def get_input_embeddings(self):
-        return self.transformer.wte
+        return self.wte
 
     def set_input_embeddings(self, value):
-        self.transformer.wte = value
+        self.wte = value
 
     def get_num_params(self, non_embedding=True):
         """Return the number of parameters in the model."""
@@ -319,14 +317,14 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
             labels = tgt
             
         # Get embeddings
-        tok_emb = self.transformer.wte(input_ids)
-        x = self.transformer.drop(tok_emb)
+        tok_emb = self.wte(input_ids)
+        x = self.drop(tok_emb)
         
         # Collect hidden states if requested
         all_hidden_states = []
         
         # Forward through transformer blocks
-        for block in self.transformer.h:
+        for block in self.h:
             if output_hidden_states:
                 all_hidden_states.append(x)
             x = block(x)
@@ -335,8 +333,8 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
             all_hidden_states.append(x)
         
         # Final layer norm and output head
-        x = self.transformer.ln_f(x)
-        logits = self.coch_head(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
         
         # Compute loss if labels provided
         loss = None
@@ -438,42 +436,42 @@ class AuriStreamModel(AuriStreamPreTrainedModel):
         b, t = seq.size()
         
         # Encode conditioning sequence into KV cache
-        tok_emb = self.transformer.wte(seq)
-        x = self.transformer.drop(tok_emb)
+        tok_emb = self.wte(seq)
+        x = self.drop(tok_emb)
         
         k_list = []
         v_list = []
-        for block in self.transformer.h:
+        for block in self.h:
             x, k, v = block(x, return_kv=True)
             k_list.append(k)
             v_list.append(v)
         
         k_cache = torch.stack(k_list, dim=0)
         v_cache = torch.stack(v_list, dim=0)
-        x = self.transformer.ln_f(x)
+        x = self.ln_f(x)
         
         # First prediction
-        logits = self.coch_head(x[:, [-1]])
+        logits = self.lm_head(x[:, [-1]])
         predictions = [self.sample_logits(logits, temperature=temp, top_k=top_k, top_p=top_p)]
         all_logits.append(logits)
         
         # Generate remaining tokens
         for i in range(n_tokens - 1):
-            tok_emb = self.transformer.wte(predictions[-1])
-            x = self.transformer.drop(tok_emb)
+            tok_emb = self.wte(predictions[-1])
+            x = self.drop(tok_emb)
             
             k_list = []
             v_list = []
-            for block_idx, block in enumerate(self.transformer.h):
+            for block_idx, block in enumerate(self.h):
                 x, k, v = block(x, k_cache=k_cache[block_idx], v_cache=v_cache[block_idx])
                 k_list.append(k)
                 v_list.append(v)
             
-            x = self.transformer.ln_f(x)
+            x = self.ln_f(x)
             k_cache = torch.stack(k_list, dim=0)
             v_cache = torch.stack(v_list, dim=0)
             
-            logits = self.coch_head(x)
+            logits = self.lm_head(x)
             predictions.append(self.sample_logits(logits, temperature=temp, top_k=top_k, top_p=top_p))
             all_logits.append(logits)