Update vit_zana_v3.py

vit_zana_v3.py

@@ -21,7 +21,7 @@ class PentachoraEmbedding(nn.Module):
     
     def __init__(self, vertices: torch.Tensor):
         super().__init__()
-        assert vertices.shape == (5, 128), f"Expected shape (5, 128), got {vertices.shape}"
+        #assert vertices.shape == (5, 128), f"Expected shape (5, 128), got {vertices.shape}"
         
         self.embed_dim = vertices.shape[-1]
         
@@ -135,7 +135,8 @@ class BaselineViT(nn.Module):
     
     def __init__(
         self,
-        pentachora_list: list,  # List of torch.Tensor, each [5, 128]
+        pentachora_list: list,  # List of torch.Tensor, each [5, vocab_dim]
+        vocab_dim: int = 256,
         img_size: int = 32,
         patch_size: int = 4,
         embed_dim: int = 512,
@@ -155,13 +156,12 @@ class BaselineViT(nn.Module):
         # Validate each pentachora
         for i, penta in enumerate(pentachora_list):
             assert isinstance(penta, torch.Tensor), f"Item {i} is not a tensor"
-            assert penta.shape == (5, 128), f"Item {i} has shape {penta.shape}, expected (5, 128)"
         
         self.num_classes = len(pentachora_list)
         self.embed_dim = embed_dim
         self.num_patches = (img_size // patch_size) ** 2
         self.similarity_mode = similarity_mode
-        self.pentachora_dim = 128  # Always 128 from vocab
+        self.pentachora_dim = vocab_dim
         
         # Create individual pentachora embeddings from list
         self.class_pentachora = nn.ModuleList([
@@ -239,15 +239,15 @@ class BaselineViT(nn.Module):
         """
         if self.similarity_mode == 'rose':
             # Stack all vertices into single tensor for batch Rose scoring
-            all_vertices = torch.stack([penta.vertices for penta in self.class_pentachora])  # [100, 5, 128]
+            all_vertices = torch.stack([penta.vertices for penta in self.class_pentachora])  # [100, 5, vocab_dim]
             # Expand features for batch computation
-            features_exp = features.unsqueeze(1).expand(-1, self.num_classes, -1)  # [B, 100, 128]
+            features_exp = features.unsqueeze(1).expand(-1, self.num_classes, -1)  # [B, 100, vocab_dim]
             # Compute Rose scores in parallel
-            return PentachoronStabilizer.rose_score_magnitude(features_exp.reshape(-1, 128), all_vertices.repeat(features.shape[0], 1, 1)).reshape(features.shape[0], -1)
+            return PentachoronStabilizer.rose_score_magnitude(features_exp.reshape(-1, self.embed_dim), all_vertices.repeat(features.shape[0], 1, 1)).reshape(features.shape[0], -1)
         else:
             # Stack all centroids
-            centroids = torch.stack([penta.centroid_norm for penta in self.class_pentachora])  # [100, 128]
-            features_norm = F.normalize(features, dim=-1)  # [B, 128]
+            centroids = torch.stack([penta.centroid_norm for penta in self.class_pentachora])  # [100, vocab_dim]
+            features_norm = F.normalize(features, dim=-1)  # [B, vocab_dim]
             return torch.matmul(features_norm, centroids.T)  # [B, 100]
 
     
@@ -319,5 +319,5 @@ class BaselineViT(nn.Module):
 if __name__ == "__main__":
     print("BaselineViT requires:")
     print("  1. PentachoronStabilizer loaded externally")
-    print("  2. pentachora_batch tensor [num_classes, 5, 128]")
+    print("  2. pentachora_batch tensor [num_classes, 5, vocab_dim]")
     print("\nNo random initialization. No fallbacks.")