Update Mo.py

Mo.py

@@ -1,23 +1,19 @@
-import tensorflow as tf
-from tensorflow.keras import layers, Model
-import numpy as np
-import tensorflow.keras.backend as K
-from tensorflow.keras import mixed_precision
+!pip install sentencepiece
 import sentencepiece as spm
-import os, json
+import os, json, numpy as np, tensorflow as tf
+from tensorflow.keras import layers, Model
 import requests
-import gradio as gr
+from tensorflow import keras
+from tensorflow.keras import layers
+import tensorflow.keras.backend as K
 
 print('1')
-
 tf.get_logger().setLevel("ERROR")
 SEED = 42
 tf.random.set_seed(SEED)
 np.random.seed(SEED)
-max_len = 512 # 기존 코드에서 200으로 설정됨
-batch_size = 128
 
-# TPU 초기화 (기존 코드와 동일)
+# TPU 초기화
 try:
     resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="local")
     tf.tpu.experimental.initialize_tpu_system(resolver)
@@ -30,15 +26,15 @@ except Exception as e:
     strategy = tf.distribute.get_strategy()
     on_tpu = False
 
-# Mixed precision (기존 코드와 동일)
+# Mixed precision
+from tensorflow.keras import mixed_precision
 policy = mixed_precision.Policy("mixed_bfloat16" if on_tpu else "float32")
 mixed_precision.set_global_policy(policy)
 print("✅ Mixed precision:", policy)
 
 # =======================
-# 1) 파일 다운로드 및 토크나이저 초기화 (기존 코드와 동일)
+# 1) 파일 다운로드
 # =======================
-
 def download_file(url, save_path):
     r = requests.get(url, stream=True)
     r.raise_for_status()
@@ -47,13 +43,13 @@ def download_file(url, save_path):
             f.write(chunk)
     print(f"✅ {save_path} 저장됨")
 
-MODEL_PATH = "model.weights.h5"
+DATA_PATH = "corpus.txt"
 TOKENIZER_PATH = "ko_unigram.model"
 
-if not os.path.exists(MODEL_PATH):
+if not os.path.exists(DATA_PATH):
     download_file(
-        "https://huggingface.co/Yuchan5386/Model_Prototype/resolve/main/model.weights.h5?download=true",
-        MODEL_PATH
+        "https://huggingface.co/datasets/Yuchan5386/Prototype/resolve/main/corpus_ko.txt?download=true",
+        DATA_PATH
     )
 
 if not os.path.exists(TOKENIZER_PATH):
@@ -72,12 +68,52 @@ unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
+max_len = 512
+batch_size = 128
+
 def text_to_ids(text):
     return sp.encode(text, out_type=int)
 
 def ids_to_text(ids):
     return sp.decode(ids)
 
+def txt_stream(file_path):
+    with open(file_path, "r", encoding="utf-8") as f:
+        for line in f:
+            text = line.strip()
+            if not text:
+                continue
+
+            ids = text_to_ids(text)
+            ids = ids[:max_len - 1]  # 마지막에 <end> 넣기 위해 -1
+
+            full_input = ids + [end_id]
+            pad_len = max_len - len(full_input)
+            full_input += [pad_id] * pad_len
+
+            # target = next-token shifted sequence
+            target = full_input[1:] + [pad_id]
+            yield (
+                tf.convert_to_tensor(full_input, dtype=tf.int32),
+                tf.convert_to_tensor(target, dtype=tf.int32)
+            )
+
+
+LIMIT = 500000  # 원하는 만큼
+
+dataset = tf.data.Dataset.from_generator(
+    lambda: txt_stream(DATA_PATH),
+    output_signature=(
+        tf.TensorSpec(shape=(max_len,), dtype=tf.int32),
+        tf.TensorSpec(shape=(max_len,), dtype=tf.int32),
+    )
+)
+
+dataset = dataset.take(LIMIT).shuffle(2000, seed=SEED).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)
+
+with strategy.scope():
+    dist_dataset = strategy.experimental_distribute_dataset(dataset)
+
 class SwiGLU(layers.Layer):
     def __init__(self, d_model, d_ff):
         super().__init__()
@@ -180,24 +216,67 @@ class ReLM(tf.keras.Model):
         logits = tf.matmul(x, embedding_matrix, transpose_b=True)
         return tf.cast(logits, tf.float32)
 
+loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
+
+def masked_loss(y_true, y_pred):
+    loss = loss_fn(y_true, y_pred)
+    mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
+    masked_loss = tf.reduce_sum(loss * mask) / tf.reduce_sum(mask)
+    return masked_loss
+
+def masked_perplexity(y_true, y_pred):
+    loss = loss_fn(y_true, y_pred)
+    mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
+    avg_loss = tf.reduce_sum(loss * mask) / tf.reduce_sum(mask)
+    return tf.exp(tf.minimum(avg_loss, 10.0))  # 수치 안정성 확보
+
+def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
+    return tf.keras.optimizers.schedules.ExponentialDecay(
+        initial_learning_rate=initial_lr,
+        decay_steps=decay_steps,
+        decay_rate=decay_rate,
+        staircase=False
+    )
 
+# 모델 생성
 model = ReLM(
     vocab_size=vocab_size,
     max_seq_len=max_len,
     d_model=256,
     n_layers=1
 )
+
+# 옵티마이저 설정
+optimizer = tf.keras.optimizers.Adam(
+    learning_rate=create_lr_schedule(),
+    beta_1=0.9,
+    beta_2=0.95,
+    epsilon=1e-8,
+    clipnorm=1.0
+)
+
+# 모델 컴파일
+model.compile(
+    optimizer=optimizer,
+    loss=masked_loss,
+    metrics=[
+        masked_perplexity
+    ]
+)
+
+# 더미 인풋으로 모델 초기화
 dummy_input = np.zeros((1, max_len), dtype=np.int32)
-_ = model(dummy_input)
+model(dummy_input)
 model.summary()
-model.load_weights(MODEL_PATH)
-print("모델 가중치 로드 완료!")
-# =======================
-# 6) 추론 함수 (기존 코드 유지)
-# 더미 인풋으로 모델 초기화
 
+history = model.fit(dataset, epochs=1, verbose=1)
 
-def generate_text_topp(model, prompt, max_len=512, max_gen=512, p=0.9, temperature=0.8, min_len=20):
+
+# 가중치 저장
+model.save_weights("model.weights.h5")
+print("모델 가중치 저장 완료!")
+
+def generate_text_topp(model, prompt, max_len=150, max_gen=150, 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)
@@ -226,23 +305,5 @@ def generate_text_topp(model, prompt, max_len=512, max_gen=512, p=0.9, temperatu
         generated.append(int(next_token_id))
     return ids_to_text(generated)
 
-def gr_generate(prompt, max_len=512, max_gen=512, p=0.8, temperature=0.8):
-    return generate_text_topp(model, prompt, max_len=max_len, p=p, temperature=temperature)
-
-# Gradio 인터페이스 정의
-iface = gr.Interface(
-    fn=gr_generate,
-    inputs=[
-        gr.Textbox(label="Prompt 입력", placeholder="여기에 문장 입력...", lines=2),
-        gr.Slider(20, 512, value=150, step=1, label="Max length"),
-        gr.Slider(0.1, 1.0, value=0.8, step=0.05, label="Top-p"),
-        gr.Slider(0.1, 2.0, value=0.8, step=0.05, label="Temperature")
-    ],
-    outputs=[
-        gr.Textbox(label="생성 결과", lines=10)
-    ],
-    title="Cuma LM 텍스트 생성",
-    description="간단한 Gradio UI로 Cuma 모델 텍스트 생성 테스트"
-)
-
-iface.launch()
+print("\n\n===== 생성 결과 =====")  
+print(generate_text_topp(model, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))