Update Mo.py

Mo.py

@@ -1,19 +1,22 @@
-!pip install sentencepiece
-import sentencepiece as spm
-import os, json, numpy as np, tensorflow as tf
+import tensorflow as tf
 from tensorflow.keras import layers, Model
-import requests
-from tensorflow import keras
-from tensorflow.keras import layers
+import numpy as np
 import tensorflow.keras.backend as K
+from tensorflow.keras import mixed_precision
+import sentencepiece as spm
+import os, json
+import requests
 
 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)
@@ -26,15 +29,15 @@ except Exception as e:
     strategy = tf.distribute.get_strategy()
     on_tpu = False
 
-# Mixed precision
-from tensorflow.keras import mixed_precision
+# 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()
@@ -43,13 +46,13 @@ def download_file(url, save_path):
             f.write(chunk)
     print(f"✅ {save_path} 저장됨")
 
-DATA_PATH = "corpus.txt"
+MODEL_PATH = "model.weights.h5"
 TOKENIZER_PATH = "ko_unigram.model"
 
-if not os.path.exists(DATA_PATH):
+if not os.path.exists(MODEL_PATH):
     download_file(
-        "https://huggingface.co/datasets/Yuchan5386/Prototype/resolve/main/corpus_ko.txt?download=true",
-        DATA_PATH
+        "https://huggingface.co/Yuchan5386/Model_Prototype/resolve/main/model.weights.h5?download=true",
+        MODEL_PATH
     )
 
 if not os.path.exists(TOKENIZER_PATH):
@@ -68,52 +71,12 @@ 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__()
@@ -216,67 +179,24 @@ 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)
-
-
-# 가중치 저장
-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):
+def generate_text_topp(model, prompt, max_len=512, max_gen=512, 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)
@@ -306,4 +226,4 @@ def generate_text_topp(model, prompt, max_len=150, max_gen=150, p=0.9, temperatu
     return ids_to_text(generated)
 
 print("\n\n===== 생성 결과 =====")  
-print(generate_text_topp(model, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))
+print(generate_text_topp(model, "지난 2년 동안", p=0.8))