Update Test.py

Test.py

@@ -1,758 +1,362 @@
 !pip install sentencepiece
-
 import sentencepiece as spm
 
+# 불러오기
 import os, json, numpy as np, tensorflow as tf
-
-from tensorflow.keras import layers, Model
-
 import requests
-
-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)
 
-
-
 # TPU 초기화
-
 try:
-
     resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="local")
-
     tf.tpu.experimental.initialize_tpu_system(resolver)
-
     strategy = tf.distribute.TPUStrategy(resolver)
-
     print("✅ TPU 초기화 완료:", resolver.cluster_spec().as_dict())
-
     on_tpu = True
-
 except Exception as e:
-
     print("⚠️ TPU 미사용, GPU/CPU로 진행:", e)
-
     strategy = tf.distribute.get_strategy()
-
     on_tpu = False
 
-
-
 # Mixed precision
-
 from tensorflow.keras import mixed_precision
-
+import tensorflow as tf
+from tensorflow.keras import layers, activations, initializers
 policy = mixed_precision.Policy("mixed_bfloat16" if on_tpu else "float32")
-
 mixed_precision.set_global_policy(policy)
-
 print("✅ Mixed precision:", policy)
 
-
-
 # =======================
-
 # 1) 파일 다운로드
-
 # =======================
-
 def download_file(url, save_path):
-
     r = requests.get(url, stream=True)
-
     r.raise_for_status()
-
     with open(save_path, "wb") as f:
-
         for chunk in r.iter_content(8192):
-
             f.write(chunk)
-
     print(f"✅ {save_path} 저장됨")
 
-
-
 DATA_PATH = "converted.jsonl"
-
 TOKENIZER_PATH = "ko_unigram.model"
 
-
-
 if not os.path.exists(DATA_PATH):
-
     download_file(
-
         "https://huggingface.co/datasets/Yuchan5386/SFT/resolve/main/data_shuffled_1.jsonl?download=true",
-
         DATA_PATH
-
     )
 
-
-
 if not os.path.exists(TOKENIZER_PATH):
-
     download_file(
-
-        "https://huggingface.co/Yuchan5386/inlam-100m/resolve/main/ko_unigram.model?download=true",
-
+        "https://huggingface.co/Yuchan5386/inlam-70m-instruct/resolve/main/unigram.model?download=true",
         TOKENIZER_PATH
-
     )
 
-
-
 sp = spm.SentencePieceProcessor(TOKENIZER_PATH)
 
-
-
 pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
-
 start_id = sp.piece_to_id("<start>")
-
 sep_id = sp.piece_to_id("<sep>")
-
 end_id = sp.piece_to_id("<end>")
-
 unk_id = sp.piece_to_id("<unk>")
-
 vocab_size = sp.get_piece_size()
-
 print(f"✅ Vocabulary size: {vocab_size}")
 
-
-
-max_len = 200
-
+max_len = 1024
 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 jsonl_stream(file_path):
-
     with open(file_path, "r", encoding="utf-8") as f:
-
         for line in f:
-
             data = json.loads(line)
-
             conversations = data.get("conversations", [])
-
             for i in range(0, len(conversations) - 1, 2):
-
                 human_msg = conversations[i]
-
                 gpt_msg   = conversations[i + 1]
-
                 if human_msg.get("from") != "human" or gpt_msg.get("from") != "gpt":
-
                     continue
-
                 prompt   = human_msg.get("value", "").strip()
-
                 response = gpt_msg.get("value", "").strip()
-
                 full = f"<start> {prompt} <sep> {response} <end>"
-
                 if "<sep>" not in full:
-
                     continue
-
                 sep_index  = full.index("<sep>")
-
                 input_text = full[:sep_index + len("<sep>")].strip()
-
                 target_text = full[sep_index + len("<sep>"):].strip()
 
-
-
                 input_ids  = text_to_ids(input_text)
-
                 target_ids = text_to_ids(target_text + " <end>")
 
-
-
                 available_len = max_len - len(input_ids)
-
                 if available_len <= 0:
-
                     input_ids = input_ids[-max_len:]
-
                     target_ids = []
-
                     target_mask = [0] * len(input_ids)
-
                 else:
-
                     target_ids = target_ids[:available_len]
-
                     target_mask = [0] * len(input_ids) + [1] * len(target_ids)
 
-
-
                 full_input = input_ids + target_ids
-
                 pad_len = max_len - len(full_input)
-
                 full_input += [pad_id] * pad_len
-
                 target_mask += [0] * pad_len
 
-
-
                 target_seq = full_input[1:] + [end_id]
-
                 target_seq = target_seq[:max_len]
 
-
-
                 masked_target = [
-
                     t if m == 1 else pad_id
-
                     for t, m in zip(target_seq, target_mask)
-
                 ]
 
-
-
                 yield (
-
                     tf.convert_to_tensor(full_input, dtype=tf.int32),
-
                     tf.convert_to_tensor(masked_target, dtype=tf.int32)
-
                 )
 
-
-
 dataset = tf.data.Dataset.from_generator(
-
     lambda: jsonl_stream(DATA_PATH),
-
     output_signature=(
-
         tf.TensorSpec(shape=(max_len,), dtype=tf.int32),
-
         tf.TensorSpec(shape=(max_len,), dtype=tf.int32),
-
     ),
-
 )
-
 dataset = dataset.shuffle(1000, seed=SEED).batch(batch_size, drop_remainder=True).prefetch(tf.data.AUTOTUNE)
 
-
-
 with strategy.scope():
-
     dist_dataset = strategy.experimental_distribute_dataset(dataset)
-
-
-
-class Lo(layers.Layer):
-
-    def __init__(self, d_model):
-
-        super().__init__()
-
-        self.proj = layers.Dense(d_model, use_bias=True, dtype='bfloat16')
-
-        self.p = layers.Dense(128, use_bias=True, dtype='bfloat16')
-
-        
+    
+class RotaryPositionalEmbedding(tf.keras.layers.Layer):
+    def __init__(self, dim):  
+        super().__init__()  
+        inv_freq = 1.0 / (10000 ** (np.arange(0, dim, 2) / dim))  
+        self.inv_freq = tf.constant(inv_freq, dtype=tf.float32)
 
     def call(self, x):
-
-        x = self.proj(x)
-
-        x = tf.nn.gelu(x)
-
-        x = self.p(x)
-
-        return x
-
-        
-
-class LoSoU(layers.Layer):
-
-    def __init__(self, d_model):
-
+        b, h, s, d = tf.unstack(tf.shape(x))
+        t = tf.range(s, dtype=tf.float32)
+        freqs = tf.einsum('i,j->ij', t, self.inv_freq)
+        dtype = x.dtype
+        emb_sin = tf.cast(tf.sin(freqs), dtype)
+        emb_cos = tf.cast(tf.cos(freqs), dtype)
+        emb_cos = tf.reshape(emb_cos, [1,1,s,-1])
+        emb_sin = tf.reshape(emb_sin, [1,1,s,-1])
+        x1, x2 = x[..., ::2], x[..., 1::2]
+        x_rot = tf.stack([x1*emb_cos - x2*emb_sin, x1*emb_sin + x2*emb_cos], axis=-1)
+        x_rot = tf.reshape(x_rot, tf.shape(x))
+        return x_rot
+
+class SwiGLU(tf.keras.layers.Layer):
+    def __init__(self, d_model, d_ff):
         super().__init__()
-
-        self.Q = layers.Dense(128)
-
-        self.K = layers.Dense(128)
-
-        self.V = Lo(d_model)
-
-        self.O = layers.Dense(d_model)
-
-        self.proj = layers.Dense(d_model, use_bias=True)
-
-
-
+        self.proj = tf.keras.layers.Dense(d_ff)
+        self.out  = tf.keras.layers.Dense(d_model)
     def call(self, x):
+        x_proj = self.proj(x)
+        x_val, x_gate = tf.split(x_proj, 2, axis=-1)
+        return self.out(x_val * tf.nn.silu(x_gate))
 
-        residual = x   # 🔹 원본 저장
-
-        q = self.Q(x)
-
-        k = self.K(x)
-
-        V = self.V(x)
-
-
-
-        g_q = tf.nn.sigmoid(q)
-
-        g_k = tf.nn.sigmoid(k)
-
-
-
-        score = g_q * g_k
-
-        score = tf.cumsum(score, axis=1)  # (B, L, D)
-
-        x = score * V
-
-
-
-        out = self.proj(x)  # 🔹 residual과 같은 차원으로 통일
-
-
-
-        a, b = tf.split(out, 2, axis=-1)
-
-        out = self.O(tf.nn.silu(a) * b)
-
-
-
-        return out + residual   # ✅ 잔차 연결 안정화
-
-
-
-      
-
-class Block(layers.Layer):
-
-    def __init__(self, d_model, r, hyper_n, num_heads, num_groups):
-
+class FlashAttentionMHA(layers.Layer):
+    def __init__(self, d_model, num_heads=8, dropout_rate=0.1):
         super().__init__()
+        self.d_model = d_model
+        self.num_heads = num_heads
+        self.dh = d_model // num_heads
+
+        self.q_proj = layers.Dense(d_model, use_bias=False)
+        self.k_proj = layers.Dense(d_model, use_bias=False)
+        self.v_proj = layers.Dense(d_model, use_bias=False)
+        self.out_proj = layers.Dense(d_model, use_bias=False)
+        self.dropout = layers.Dropout(dropout_rate)
+        self.rope = RotaryPositionalEmbedding(self.dh)
+
+    @tf.function(jit_compile=True)
+    def call(self, x, training=False, causal=False):
+        B, N, D = tf.shape(x)[0], tf.shape(x)[1], x.shape[2]
+
+        # Q,K,V: (B, N, num_heads, dh)
+        Q = tf.reshape(self.q_proj(x), [B, N, self.num_heads, self.dh])
+        K = tf.reshape(self.k_proj(x), [B, N, self.num_heads, self.dh])
+        V = tf.reshape(self.v_proj(x), [B, N, self.num_heads, self.dh])
+
+        # transpose for attention: (B, num_heads, N, dh)
+        Q = tf.transpose(Q, [0,2,1,3])
+        K = tf.transpose(K, [0,2,1,3])
+        V = tf.transpose(V, [0,2,1,3])
+
+        # ROPE 적용
+        Q = self.rope(Q)
+        K = self.rope(K)
+
+        # Scaled dot-product
+        scale = tf.cast(self.dh ** -0.5, x.dtype)
+        Q = Q * scale
+        attn_scores = tf.matmul(Q, K, transpose_b=True)
+
+        if causal:
+            mask = tf.linalg.band_part(tf.ones((N,N), dtype=x.dtype), -1, 0)
+            attn_scores = attn_scores * mask - 1e9 * (1 - mask)
+
+        attn_weights = tf.nn.softmax(attn_scores, axis=-1)
+        attn_weights = self.dropout(attn_weights, training=training)
+        out = tf.matmul(attn_weights, V)  # (B, h, N, dh)
+        out = tf.transpose(out, [0,2,1,3])
+        out = tf.reshape(out, [B, N, D])
+        out = self.out_proj(out)
+        return out
+
+
+class GPTBlock(tf.keras.layers.Layer):
+    def __init__(self, d_model, d_ff, num_heads=12, dropout_rate=0.1, adapter_dim=64):
+        super().__init__()
+        self.ln1 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
+        self.mha = FlashAttentionMHA(d_model, num_heads, dropout_rate=dropout_rate)
+        self.dropout1 = tf.keras.layers.Dropout(dropout_rate)
+        self.adapter_down = tf.keras.layers.Dense(adapter_dim, activation='gelu')
+        self.adapter_up   = tf.keras.layers.Dense(d_model)
+        self.ln2 = tf.keras.layers.LayerNormalization(epsilon=1e-5)
+        self.ffn = SwiGLU(d_model, d_ff)
+        self.dropout2 = tf.keras.layers.Dropout(dropout_rate)
 
-        self.losou = [LoSoU(d_model) for _ in range(hyper_n)]
-
-
-
-    def call(self, x):
-
-        for losou in self.losou:
-
-            x = losou(x)     
-
+    def call(self, x, training=False):
+        x_norm = self.ln1(x)
+        attn_out = self.mha(x_norm, training=training, causal=True)
+        attn_out = self.dropout1(attn_out, training=training)
+        adapter_out = self.adapter_up(self.adapter_down(attn_out))
+        attn_out = attn_out + adapter_out
+        x = x + attn_out
+        ffn_out = self.ffn(self.ln2(x))
+        x = x + self.dropout2(ffn_out, training=training)
         return x
 
- 
-
-class Sequen(tf.keras.Model):
-
-    def __init__(self, vocab_size, max_seq_len, d_model, n_layers, dropout_rate=0.1):
-
+class InLaM(tf.keras.Model):
+    def __init__(self, vocab_size, seq_len, d_model, d_ff, n_layers, num_heads=12, dropout_rate=0.1):
         super().__init__()
-
-        self.token_embedding = layers.Embedding(vocab_size, d_model)
-
-        self.pos_embedding = layers.Embedding(max_seq_len, d_model)
-
-        self.blocks = [Block(d_model, r=204, hyper_n=3, num_heads=8, num_groups=2) for _ in range(n_layers)]
-
-
-
-        # ✅ 마지막도 RMSNorm으로
-
-        self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype="bfloat16")
-
-
-
+        self.vocab_size = vocab_size
+        self.d_model = d_model
+        
+        # Embedding 레이어 (bfloat16)
+        self.token_embedding = tf.keras.layers.Embedding(vocab_size, d_model, dtype="bfloat16")
+        
+        # Transformer Blocks
+        self.blocks = [GPTBlock(d_model, d_ff, num_heads, dropout_rate) for _ in range(n_layers)]
+        
+        # Final LayerNorm
+        self.ln_f = tf.keras.layers.LayerNormalization(epsilon=1e-5, dtype="bfloat16")
     def call(self, x, training=False):
-
-        batch_size, seq_len = tf.shape(x)[0], tf.shape(x)[1]
-
-        positions = tf.range(seq_len)[tf.newaxis, :]  # (1, seq_len)
-
-
-
-        x = self.token_embedding(x) + self.pos_embedding(positions)  # (batch, seq_len, d_model)
-
+        # Embedding
+        x = self.token_embedding(x)  # (batch, seq_len, d_model)
         for block in self.blocks:
-
-            x = block(x)
-
-            
+            x = block(x, training=training)
 
         x = self.ln_f(x)  # (batch, seq_len, d_model)
-
-
-
-        # ✅ embedding weight tying
-
-        embedding_matrix = self.token_embedding.embeddings
-
-        logits = tf.matmul(x, embedding_matrix, transpose_b=True)  # (batch, seq_len, vocab_size)
-
+        embed_weights = self.token_embedding.weights[0]  # (vocab_size, d_model)
+        logits = tf.matmul(x, embed_weights, transpose_b=True)  # (batch, seq_len, vocab_size)
+        
+        # float32로 캐스팅 (손실 계산 등에서 안정성 확보)
         return tf.cast(logits, tf.float32)
 
-
-
+# =======================
+# 손실/메트릭 정의
+# =======================
 def smoothed_loss_keras(y_true, y_pred, eps=0.1):
-
     y_true = tf.cast(y_true, tf.int32)
-
     mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-
     vocab = tf.shape(y_pred)[-1]
-
     y_true_oh = tf.one_hot(y_true, depth=vocab, dtype=tf.float32)
-
     y_true_ls = (1.0 - eps) * y_true_oh + eps / tf.cast(vocab, tf.float32)
-
     log_probs = tf.nn.log_softmax(y_pred, axis=-1)
-
-    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1) * mask
-
+    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1)
+    per_tok = per_tok * mask
     return tf.reduce_sum(per_tok) / (tf.reduce_sum(mask) + 1e-8)
 
-
-
 def masked_accuracy(y_true, y_pred):
-
+    y_true = tf.cast(y_true, tf.int32)
     mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
-
     pred_id = tf.argmax(y_pred, axis=-1, output_type=tf.int32)
-
     acc = tf.cast(tf.equal(y_true, pred_id), tf.float32) * mask
-
     return tf.reduce_sum(acc) / (tf.reduce_sum(mask) + 1e-8)
 
-    
-
-# =======================
+def masked_perplexity(y_true, y_pred, eps=0.1):
+    y_true = tf.cast(y_true, tf.int32)
+    mask = tf.cast(tf.not_equal(y_true, pad_id), tf.float32)
+    vocab = tf.shape(y_pred)[-1]
+    y_true_oh = tf.one_hot(y_true, depth=vocab, dtype=tf.float32)
+    y_true_ls = (1.0 - eps) * y_true_oh + eps / tf.cast(vocab, tf.float32)
+    log_probs = tf.nn.log_softmax(y_pred, axis=-1)
+    per_tok = -tf.reduce_sum(y_true_ls * log_probs, axis=-1)
+    per_tok = per_tok * mask
+    mean_loss = tf.reduce_sum(per_tok) / (tf.reduce_sum(mask) + 1e-8)
+    return tf.exp(mean_loss)
 
-# 모델 생성 & 학습
 
 # =======================
-
+# 모델 생성 & 컴파일
+# =======================
 with strategy.scope():
-
-    model = Sequen(vocab_size, max_seq_len=max_len, d_model=384, n_layers=12, dropout_rate=0.1)
-
+    model = InLaM(vocab_size=vocab_size, seq_len=max_len, d_model=768, d_ff=768*4, n_layers=12)
     dummy_input = tf.zeros((batch_size, max_len), dtype=tf.int32)
-
     _ = model(dummy_input, training=False)
-
     model.summary()
 
+    optimizer = tf.keras.optimizers.Adam(1e-4, beta_1=0.9, beta_2=0.95, epsilon=1e-8, clipnorm=1.0)
+    model.compile(optimizer=optimizer, loss=smoothed_loss_keras, metrics=[masked_accuracy, masked_perplexity])
 
-
-    optimizer = tf.keras.optimizers.Adam(3e-4, beta_1=0.9, beta_2=0.95, epsilon=1e-8, clipnorm=1.0)
-
-    model.compile(optimizer=optimizer, loss=smoothed_loss_keras, metrics=[masked_accuracy])
-
+    # 학습
     history = model.fit(dist_dataset, epochs=1, verbose=1)
 
-
-
 # =======================
-
 # 가중치 저장
-
 # =======================
-
-model.save_weights("Sequen.weights.h5")
-
+model.save_weights("tf_model.weights.h5")
 print("✅ 모델 가중치 저장 완료!")
 
-
-
 # =======================
-
-@tf.function(input_signature=[
-
-    tf.TensorSpec(shape=(1, None), dtype=tf.int32),  # input_ids
-
-    tf.TensorSpec(shape=(vocab_size,), dtype=tf.int32),  # token_counts
-
-    tf.TensorSpec(shape=(), dtype=tf.int32),  # current_length
-
-    tf.TensorSpec(shape=(), dtype=tf.float32),  # temperature
-
-    tf.TensorSpec(shape=(), dtype=tf.float32),  # repetition_penalty
-
-    tf.TensorSpec(shape=(), dtype=tf.float32),  # top_p
-
-    tf.TensorSpec(shape=(), dtype=tf.int32),  # top_k
-
-    tf.TensorSpec(shape=(), dtype=tf.int32),  # min_len
-
-    tf.TensorSpec(shape=(), dtype=tf.int32),  # step
-
-])
-
-def generate_step(input_ids, token_counts, current_length, temperature, repetition_penalty, top_p, top_k, min_len, step):
-
-    pad_len = max_len - tf.shape(input_ids)[1]
-
-    input_padded = tf.pad(input_ids, [[0,0],[0,pad_len]], constant_values=pad_id)
-
-    logits = model(input_padded, training=False)
-
-    next_logits = logits[0, current_length - 1]
-
-
-
-    penalty = tf.pow(repetition_penalty, tf.cast(token_counts, tf.float32))
-
-    next_logits = next_logits / penalty
-
-
-
-    # 최소 길이와 pad 마스킹
-
-    if current_length < min_len:
-
-        next_logits = tf.tensor_scatter_nd_update(next_logits, [[end_id]], [-1e9])
-
-    next_logits = tf.tensor_scatter_nd_update(next_logits, [[pad_id]], [-1e9])
-
-
-
-    # top-k 필터링
-
-    if top_k > 0:
-
-        kth_val = tf.math.top_k(next_logits, k=top_k).values[-1]
-
-        mask = next_logits < kth_val
-
-        next_logits = tf.where(mask, -1e9, next_logits)
-
-
-
-    # top-p (nucleus) 필터링 + temperature
-
-    next_logits = next_logits / temperature
-
-    probs = tf.nn.softmax(next_logits)
-
-    sorted_probs, sorted_idx = tf.math.top_k(probs, k=vocab_size)
-
-    cum_probs = tf.cumsum(sorted_probs)
-
-    cutoff_mask = cum_probs <= top_p
-
-    cutoff_idx = tf.reduce_sum(tf.cast(cutoff_mask, tf.int32)) + 1
-
-    cutoff_idx = tf.minimum(cutoff_idx, vocab_size)
-
-    filtered_idx = sorted_idx[:cutoff_idx]
-
-    filtered_probs = sorted_probs[:cutoff_idx]
-
-    filtered_probs = filtered_probs / tf.reduce_sum(filtered_probs)
-
-
-
-    # 🔹 50%는 argmax, 50%는 샘플링
-
-    rand_val = tf.random.uniform([], 0.1, 1)
-
-    def sample():
-
-        sampled_id = tf.random.categorical(tf.math.log([filtered_probs]), 1)[0,0]
-
-        return filtered_idx[sampled_id]
-
-    def argmax():
-
-        return filtered_idx[tf.argmax(filtered_probs)]
-
-    sampled_id = tf.cond(rand_val < 0, argmax, sample)
-
-    sampled_id = tf.cast(sampled_id, tf.int32)
-
-
-
-    # token_counts 업데이트
-
-    token_counts = tf.tensor_scatter_nd_add(token_counts, [[sampled_id]], [1])
-
-    return sampled_id, token_counts
-
-
-
-
-
-# =====================
-
-# 스트리밍 생성기 (CPU 최적화 버전)
-
-# =====================
-
-def generate_text_streaming(model, prompt, max_len=115, max_gen=100,
-
-                            temperature=0.75, min_len=20,
-
-                            repetition_penalty=1.2, top_p=0.9, top_k=50):
-
+# 샘플 생성 함수
+# =======================
+def generate_text_topp(model, prompt, max_len=115, max_gen=98, p=0.9, temperature=0.68, min_len=20):
     model_input = text_to_ids(f"<start> {prompt} <sep>")
-
     model_input = model_input[:max_len]
-
     generated = list(model_input)
-
-    start_output_idx = len(model_input)
-
-
-
-    # TF 변수로 토큰 카운트 관리
-
-    token_counts_np = np.zeros(vocab_size, dtype=np.int32)
-
-    for t in generated:
-
-        token_counts_np[t] += 1
-
-    token_counts = tf.Variable(token_counts_np, dtype=tf.int32)
-
-
-
-    prev_decoded = ""
-
-
-
+    
     for step in range(max_gen):
+        input_seq = generated[-max_len:] if len(generated) > max_len else generated
+        input_padded = np.pad(input_seq, (0, max_len - len(input_seq)), constant_values=pad_id)
+        input_tensor = tf.convert_to_tensor([input_padded], dtype=tf.int32)
+        
+        logits = model(input_tensor, training=False).numpy()[0, len(input_seq)-1]
+        logits[end_id] -= 5.0
+        logits[pad_id] -= 10.0
+        
+        probs = tf.nn.softmax(logits / temperature).numpy()
+        sorted_idx = np.argsort(probs)[::-1]
+        sorted_probs = probs[sorted_idx]
+        cumulative = np.cumsum(sorted_probs)
+        cutoff = np.searchsorted(cumulative, p)
+        top_idx = sorted_idx[:cutoff + 1]
+        top_probs = sorted_probs[:cutoff + 1] / sorted_probs[:cutoff + 1].sum()
+        
+        next_token = int(np.random.choice(top_idx, p=top_probs))
+        if next_token == end_id and len(generated) >= min_len:
+            break
+        generated.append(next_token)
+    
+    return ids_to_text(generated)
 
-        input_tensor = tf.expand_dims(generated, axis=0)  # [1, seq_len]
-
-
-
-        sampled_id, token_counts = generate_step(
-
-            input_tensor,
-
-            token_counts,
-
-            tf.constant(len(generated), dtype=tf.int32),
-
-            tf.constant(temperature, dtype=tf.float32),
-
-            tf.constant(repetition_penalty, dtype=tf.float32),
-
-            tf.constant(top_p, dtype=tf.float32),
-
-            tf.constant(top_k, dtype=tf.int32),
-
-            tf.constant(min_len, dtype=tf.int32),
-
-            tf.constant(step, dtype=tf.int32)
-
-        )
-
-
-
-        sampled_id = int(sampled_id.numpy())
-
-        generated.append(sampled_id)
-
-
-
-        # 디코딩은 출력 시점에만
-
-        if len(generated) > start_output_idx:
-
-            decoded_full = sp.decode(generated[start_output_idx:])
-
-            decoded_full = decoded_full.replace("▁", " ").strip()
-
-            for t in ["<start>", "<sep>", "<end>"]:
-
-                decoded_full = decoded_full.replace(t, "")
-
-            decoded_full = decoded_full.lstrip(",!?.는은 ")
-
-
-
-            new_output = decoded_full[len(prev_decoded):]
-
-            if new_output:
-
-                yield new_output
-
-                prev_decoded = decoded_full
-
-
-
-            # 종료 조건
-
-            if len(generated) >= min_len and (sampled_id == end_id or decoded_full.endswith(('.', '!', '?'))):
-
-                break
-
-
-
-
-
-
-
-for token in generate_text_streaming(
-
-    model, '안녕하세요',
-
-    max_len=max_len,
-
-    max_gen=115,
-
-    temperature=0.8,
-
-    min_len=10,
-
-    repetition_penalty=1.1,
-
-    top_p=0.9,
-
-    top_k=32
-
-):
-
-    print(token, end="", flush=True)  
-
-
-
-
-이 학습 코드가 왜 NaN을 뱉는지 설명해
+# =======================
+# 테스트 생성
+# =======================
+prompt = "안녕하세요! 한국 밴드에 대해 궁금한 것이 있어요!"
+sample_text = generate_text_topp(model, prompt, p=0.9)
+print("\n===== 생성 결과 =====\n")
+print(sample_text)