Update Mo.py

Mo.py

@@ -48,7 +48,7 @@ TOKENIZER_PATH = "ko_unigram.model"
 
 if not os.path.exists(DATA_PATH):
     download_file(
-        "https://huggingface.co/datasets/Yuchan5386/Prototype/resolve/main/corpus_ko.txt?download=true",
+        "https://huggingface.co/datasets/Yuchan5386/1/resolve/main/shuffled_corpus.txt?download=true",
         DATA_PATH
     )
 
@@ -68,7 +68,7 @@ unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
-max_len = 512
+max_len = 256
 batch_size = 256
 
 def text_to_ids(text):
@@ -99,7 +99,7 @@ def txt_stream(file_path):
             )
 
 
-LIMIT = 2000000  # 원하는 만큼
+LIMIT = 36757266
 
 dataset = tf.data.Dataset.from_generator(
     lambda: txt_stream(DATA_PATH),
@@ -117,7 +117,7 @@ with strategy.scope():
 class SwiGLU(layers.Layer):
     def __init__(self, d_model, d_ff):
         super().__init__()
-        self.proj = layers.Dense(2048)
+        self.proj = layers.Dense(960)
         self.out = layers.Dense(d_model)
     def call(self, x):
         x_proj = self.proj(x)
@@ -257,7 +257,7 @@ def masked_perplexity(y_true, y_pred, eps=0.1):
 # 모델 생성 & 컴파일
 # =======================
 with strategy.scope():
-    model = ReLM(vocab_size=vocab_size, max_seq_len=max_len, d_model=512, n_layers=9)
+    model = ReLM(vocab_size=vocab_size, max_seq_len=max_len, d_model=384, n_layers=3)
     dummy_input = tf.zeros((batch_size, max_len), dtype=tf.int32)
     _ = model(dummy_input, training=False)
     model.summary()
@@ -271,7 +271,7 @@ with strategy.scope():
 model.save_weights("tf_model.weights.h5")
 print("✅ 모델 가중치 저장 완료!")
 
-def generate_text_topp(model, prompt, max_len=512, max_gen=512, p=0.9, temperature=0.8, min_len=20):
+def generate_text_topp(model, prompt, max_len=500, max_gen=500, p=0.9, temperature=0.8, min_len=20):
     model_input = text_to_ids(f"<start> {prompt}")
     model_input = model_input[:max_len]
     generated = list(model_input)