Update Mo.py

Mo.py

@@ -67,7 +67,7 @@ unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
-max_len = 200
+max_len = 128
 batch_size = 96
 
 def text_to_ids(text):
@@ -126,7 +126,7 @@ class SwiGLU(layers.Layer):
         
 
 class MHLA(layers.Layer):
-    def __init__(self, embed_dim, num_heads=8, dropout=0.0):
+    def __init__(self, embed_dim, num_heads=2, dropout=0.0):
         super().__init__()
         assert embed_dim % num_heads == 0, "embed_dim must be divisible by num_heads"
         self.embed_dim = embed_dim
@@ -187,7 +187,7 @@ class Lo(layers.Layer):
 class Block(layers.Layer):
     def __init__(self, d_model):
         super().__init__()
-        self.lou = MHLA(d_model, 8)
+        self.lou = MHLA(d_model, 2)
         self.glu = SwiGLU(d_model, 1048)
         self.lo = Lo(d_model)