Sync latest aetheris source code

aetheris/model.py

@@ -30,17 +30,22 @@ class HybridMambaMoE(nn.Module):
         nn.init.normal_(self.embedding.weight, mean=0.0, std=0.02)
 
     def resize_token_embeddings(self, new_vocab_size: int):
-        """Resize embedding and lm_head for new tokens. New embeddings initialized from mean of existing."""
+        """Resize embedding and lm_head. Handles both growing and shrinking."""
         old_vocab_size = self.embedding.num_embeddings
         if new_vocab_size == old_vocab_size:
             return
-        old_weight = self.embedding.weight.data
-        mean_embed = old_weight.mean(dim=0)
+        old_weight = self.embedding.weight.data.clone()
         self.embedding = nn.Embedding(new_vocab_size, self.config.d_model)
-        self.embedding.weight.data[:old_vocab_size] = old_weight
-        self.embedding.weight.data[old_vocab_size:] = mean_embed.unsqueeze(0).expand(
-            new_vocab_size - old_vocab_size, -1
-        )
+        if new_vocab_size > old_vocab_size:
+            # Growing: copy old weights, init new with mean
+            mean_embed = old_weight.mean(dim=0)
+            self.embedding.weight.data[:old_vocab_size] = old_weight
+            self.embedding.weight.data[old_vocab_size:] = mean_embed.unsqueeze(0).expand(
+                new_vocab_size - old_vocab_size, -1
+            )
+        else:
+            # Shrinking: truncate (caller should copy correct rows afterward)
+            self.embedding.weight.data[:] = old_weight[:new_vocab_size]
         self.lm_head = nn.Linear(self.config.d_model, new_vocab_size, bias=False)
         self.lm_head.weight = self.embedding.weight  # Re-tie weights
         self.config.vocab_size = new_vocab_size