Update AlphaS2S.py

AlphaS2S.py

@@ -1,15 +1,11 @@
 import tensorflow as tf
 from tensorflow.keras import layers, Model
-!pip install sentencepiece
-
+import numpy as np
+import tensorflow.keras.backend as K
+from tensorflow.keras import mixed_precision
 import sentencepiece as spm
-import os, json, numpy as np, tensorflow as tf
-from tensorflow.keras import layers, Model
+import os, json
 import requests
-from tensorflow import keras
-from tensorflow.keras import layers
-import tensorflow.keras.backend as K
-
 
 print('1')
 
@@ -17,9 +13,10 @@ tf.get_logger().setLevel("ERROR")
 SEED = 42
 tf.random.set_seed(SEED)
 np.random.seed(SEED)
-max_len = 100
-# TPU 초기화
+max_len = 200 # 기존 코드에서 200으로 설정됨
+batch_size = 128
 
+# TPU 초기화 (기존 코드와 동일)
 try:
     resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu="local")
     tf.tpu.experimental.initialize_tpu_system(resolver)
@@ -32,14 +29,13 @@ 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):
@@ -75,9 +71,6 @@ unk_id = sp.piece_to_id("<unk>")
 vocab_size = sp.get_piece_size()
 print(f"✅ Vocabulary size: {vocab_size}")
 
-max_len = 200
-batch_size = 128
-
 def text_to_ids(text):
     return sp.encode(text, out_type=int)
 
@@ -85,6 +78,10 @@ def ids_to_text(ids):
     return sp.decode(ids)
 
 
+# =======================
+# 2) 데이터셋 생성 함수 (기존 코드와 동일)
+# =======================
+
 def jsonl_stream(file_path):
     with open(file_path, "r", encoding="utf-8") as f:
         for line in f:
@@ -103,12 +100,24 @@ def jsonl_stream(file_path):
                     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>")
+                
+                # 인코더 입력은 <start> 프롬프트 <sep> 부분, 디코더 입력은 <sep> 응답 <end> 부분
+                # (Unified Input: 인코더/디코더 입력 모두 full_input을 사용)
+                input_text = full
+                
+                # 타겟 시퀀스는 응답 시작 부분부터 <end>까지이며, 입력보다 한 칸 시프트됨
+                # 여기서 target_text는 응답 부분만 추출하여 타겟 마스킹에 사용됩니다.
+                target_text_raw = full[sep_index + len("<sep>"):]
+
+                input_ids  = text_to_ids(input_text) # 전체 시퀀스
+                target_ids_raw = text_to_ids(target_text_raw) # 응답 부분만
+                
+                # 길이 처리 및 마스킹 로직은 기존 코드를 그대로 유지
+                full_input = input_ids[:max_len]
+                target_ids = target_ids_raw[:max_len - len(input_ids)]
+                
                 available_len = max_len - len(input_ids)
-
+                
                 if available_len <= 0:
                     input_ids = input_ids[-max_len:]
                     target_ids = []
@@ -121,30 +130,49 @@ def jsonl_stream(file_path):
                 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 = full_input[1:] + [end_id] 
                 target_seq = target_seq[:max_len]
+                
+                # 마스킹된 타겟 생성 (프롬프트/패딩 부분은 pad_id로 대체)
                 masked_target = [
                     t if m == 1 else pad_id
                     for t, m in zip(target_seq, target_mask)
                 ]
+
+                # AlphaS2S는 인코더/디코더 입력으로 같은 시퀀스를 사용
+                # 입력 시퀀스 = full_input
+                # 타겟 시퀀스 = masked_target
                 yield (
                     tf.convert_to_tensor(full_input, dtype=tf.int32),
-                    tf.convert_to_tensor(masked_target, dtype=tf.int32)
+                    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),
+        tf.TensorSpec(shape=(max_len,), dtype=tf.int32), # enc_inputs
+        tf.TensorSpec(shape=(max_len,), dtype=tf.int32), # dec_inputs
+        tf.TensorSpec(shape=(max_len,), dtype=tf.int32), # target
     ),
 )
 
+# 학습을 위해 딕셔너리 형태로 맵핑
+def map_fn(enc_input, dec_input, dec_target):
+    return {"enc_inputs": enc_input, "dec_inputs": dec_input}, dec_target
+
+dataset = dataset.map(map_fn, num_parallel_calls=tf.data.AUTOTUNE)
 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)
-    
+
+# =======================
+# 3) 모델 레이어 (기존 코드 유지)
+# =======================
+
 class SwiGLU(layers.Layer):
     def __init__(self, d_model, d_ff):
         super().__init__()
@@ -210,11 +238,13 @@ class LoU(layers.Layer):
         remaining_seq = seq[1:]
         remaining_alpha = alpha_seq[1:]
         elems = (remaining_seq, remaining_alpha)
+        # tf.scan을 사용하여 시계열 EMA 계산
         ema_seq = tf.scan(fn=step, elems=elems, initializer=init)
         ema_seq = tf.concat([tf.expand_dims(init, 0), ema_seq], axis=0)
         ema = tf.transpose(ema_seq, perm=[1, 0, 2])
         return ema
 
+    # LoU는 원래 Uni-directional Attention/Recurrent Block 역할
     def call(self, x, z):
         x_f32 = tf.cast(x, tf.float32)
         residual = x_f32
@@ -223,9 +253,9 @@ class LoU(layers.Layer):
         q = self.Q(x_f32)
         k = self.K(x_f32)
         V = self.V(x_f32)
-        # 기존 코드:
-        # g_q = tf.nn.sigmoid(q)
-        # g_k = tf.nn.sigmoid(k)
+        
+        # Unidirectional Masking: 미래 정보를 막는 Look-ahead Mask를 수동으로 적용해야 하지만,
+        # 기존 LoU 구현은 Self-Attention이 아니므로 Skip.
 
         g_q = (tf.nn.tanh(q) + 1.0) / 2.0
         g_k = (tf.nn.tanh(k) + 1.0) / 2.0
@@ -238,47 +268,76 @@ class LoU(layers.Layer):
         score_norm = score_ema / denom
         score_clipped = tf.clip_by_value(score_norm, -self.clip_value, self.clip_value)
         x_comb = score_clipped * V
+        
+        # LoU 블록에서는 x_comb + residual 후 CrossBlock을 통과
         out = self.norm(x_comb + residual)
-        out = self.cross(out, z)
+        out = self.cross(out, z) # z는 인코더 출력 (enc_out)
         out = self.glu(out)
         return tf.cast(out, x.dtype)
         
+# =======================
+# 4) AlphaS2S 모델 (기존 코드 유지)
+# =======================
+
 class AlphaS2S(tf.keras.Model):
-    def __init__(self, num_layers, d_model, num_heads, input_vocab_size, target_vocab_size, max_len=100, dropout=0.1):
+    def __init__(self, num_layers, d_model, num_heads, input_vocab_size, target_vocab_size, max_len=200, dropout=0.1):
         super().__init__()
         self.max_len = max_len
         self.d_model = d_model
+        
+        # 인코더와 디코더 임베딩 및 위치 임베딩은 모두 max_len을 사용
         self.enc_embedding = layers.Embedding(input_vocab_size, d_model)
         self.enc_pos_embedding = layers.Embedding(max_len, d_model)
         self.dec_embedding = layers.Embedding(target_vocab_size, d_model)
         self.dec_pos_embedding = layers.Embedding(max_len, d_model)
-        self.enc_layers = [EncoderBlock(d_model, num_heads, dropout) for _ in range(num_layers)]
+        
+        # EncoderBlock과 LoU는 기존 코드와 동일한 구조
+        self.enc_layers = [EncoderBlock(d_model, num_heads, d_model * 4, dropout) for _ in range(num_layers)]
         self.dec_layers = [LoU(d_model) for _ in range(num_layers)]
+        
         self.final_layer = layers.Dense(target_vocab_size, use_bias=False)
+        
     def call(self, inputs, training=False):
-        enc_inputs = inputs["enc_inputs"]
+        # enc_inputs와 dec_inputs는 동일한 시퀀스 (Unified Input)
+        enc_inputs = inputs["enc_inputs"] 
         dec_inputs = inputs["dec_inputs"]
+        
         enc_pos = tf.range(tf.shape(enc_inputs)[1])[tf.newaxis, :]
         dec_pos = tf.range(tf.shape(dec_inputs)[1])[tf.newaxis, :]
+        
+        # 인코더 실행
         x = self.enc_embedding(enc_inputs) + self.enc_pos_embedding(enc_pos)
+        # Note: 마스크 없음 -> Bi-directional (BERT-like Encoder)
         for layer in self.enc_layers: x = layer(x, training=training)
-        enc_out = x
+        enc_out = x # 인코더의 최종 출력 (디코더의 'z' 입력)
+        
+        # 디코더 실행
         y = self.dec_embedding(dec_inputs) + self.dec_pos_embedding(dec_pos)
-        for layer in self.dec_layers: y = layer(y, enc_out, training=training)
+        # Note: LoU는 내부적으로 EMA를 사용하며, 일반적인 Cross-Attention 블록의 역할을 수행
+        for layer in self.dec_layers: y = layer(y, enc_out, training=training) 
+        
         return self.final_layer(y)
 
+# =======================
+# 5) 학습 설정 및 실행
+# =======================
+
 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)
+    # mixed_bfloat16 사용 시 나눗셈 시 NaN 방지
+    sum_mask = tf.reduce_sum(mask)
+    safe_sum_mask = tf.where(sum_mask == 0.0, 1.0, sum_mask)
+    masked_loss = tf.reduce_sum(loss * mask) / safe_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))  # 수치 안정성 확보
-
+    sum_mask = tf.reduce_sum(mask)
+    safe_sum_mask = tf.where(sum_mask == 0.0, 1.0, sum_mask)
+    avg_loss = tf.reduce_sum(loss * mask) / safe_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(
@@ -288,67 +347,115 @@ def create_lr_schedule(initial_lr=5e-5, decay_steps=10000, decay_rate=0.9):
         staircase=False
     )
 
-chat_model = AlphaS2S(num_layers=4, d_model=160, num_heads=8,
-                         input_vocab_size=chat_vocab_size, target_vocab_size=chat_vocab_size)
-dummy_input = {
-    "enc_inputs": tf.zeros((1, max_len), dtype=tf.int32),
-    "dec_inputs": tf.zeros((1, max_len), dtype=tf.int32)
-}
-_ = chat_model(dummy_input)
-loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
-
-
-# 옵티마이저 설정
-optimizer = tf.keras.optimizers.Adam(
-    learning_rate=create_lr_schedule(),
-    beta_1=0.9,
-    beta_2=0.95,
-    epsilon=1e-8,
-    clipnorm=1.0
-)
+with strategy.scope():
+    # ⚠️ 수정: chat_vocab_size 대신 정의된 vocab_size 사용
+    chat_model = AlphaS2S(num_layers=4, d_model=160, num_heads=8,
+                             input_vocab_size=vocab_size, target_vocab_size=vocab_size, max_len=max_len)
+    
+    dummy_input = {
+        "enc_inputs": tf.zeros((1, max_len), dtype=tf.int32),
+        "dec_inputs": tf.zeros((1, max_len), dtype=tf.int32)
+    }
+    _ = chat_model(dummy_input)
+    
+    loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction='none')
 
-# 모델 컴파일
-chat_model.compile(
-    optimizer=optimizer,
-    loss=masked_loss,
-    metrics=[
-        masked_perplexity
-    ]
-)
 
-history = chat_model.fit(dataset, epochs=1, verbose=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
+    )
+
+    # 모델 컴파일
+    chat_model.compile(
+        optimizer=optimizer,
+        loss=masked_loss,
+        metrics=[
+            masked_perplexity
+        ]
+    )
+
+    print("✅ 모델 컴파일 완료, 학습 시작...")
+    # ⚠️ 학습 실행
+    history = chat_model.fit(dataset, epochs=1, verbose=1)
+
 # 가중치 저장
 chat_model.save_weights("chat_model.weights.h5")
-print("모델 가중치 저장 완료!")
+print("\n✅ 모델 가중치 저장 완료!")
+
+# =======================
+# 6) 추론 함수 (기존 코드 유지)
+# =======================
 
-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}")
+def generate_text_topp(model, prompt, max_len=200, max_gen=100, p=0.9, temperature=0.8, min_len=20):
+    # 인코더 입력은 <start> Prompt <sep> 만 사용
+    model_input = text_to_ids(f"<start> {prompt} <sep>")
     model_input = model_input[:max_len]
     generated = list(model_input)
+    
     for step in range(max_gen):
-        if len(generated) > max_len:
+        current_len = len(generated)
+        
+        # 현재까지 생성된 시퀀스를 입력으로 사용
+        if current_len > max_len:
             input_seq = generated[-max_len:]
         else:
             input_seq = 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])
-        logits = model(input_tensor, training=False)
+        
+        # 모델 추론 (enc_inputs, dec_inputs 모두 동일한 시퀀스를 사용)
+        dummy_input = {
+            "enc_inputs": input_tensor,
+            "dec_inputs": input_tensor
+        }
+        logits = model(dummy_input, training=False)
+        
+        # 다음 토큰의 로짓은 시퀀스의 마지막 토큰 위치에서 가져옴 (0-based index: current_len - 1)
+        # 하지만 패딩 후 input_tensor의 실제 시퀀스 길이는 len(input_seq)
         next_token_logits = logits[0, len(input_seq) - 1].numpy()
+        
+        # 특수 토큰 생성 억제
         next_token_logits[end_id] -= 5.0
         next_token_logits[pad_id] -= 10.0
+        
         probs = tf.nn.softmax(next_token_logits / temperature).numpy()
         sorted_indices = np.argsort(probs)[::-1]
         sorted_probs = probs[sorted_indices]
+        
+        # Top-p (Nucleus) Sampling
         cumulative_probs = np.cumsum(sorted_probs)
         cutoff = np.searchsorted(cumulative_probs, p)
         top_indices = sorted_indices[:cutoff + 1]
         top_probs = sorted_probs[:cutoff + 1]
         top_probs /= np.sum(top_probs)
         next_token_id = np.random.choice(top_indices, p=top_probs)
+
         if next_token_id == end_id and len(generated) >= min_len:
             break
+            
         generated.append(int(next_token_id))
-    return ids_to_text(generated)
+
+    # <start> 토큰 제거 및 <sep> 이전 부분 제거
+    try:
+        sep_index = generated.index(sep_id)
+        # <sep> 이후부터 <end> 이전까지의 응답만 반환
+        result_ids = generated[sep_index + 1:]
+        try:
+            end_index = result_ids.index(end_id)
+            result_ids = result_ids[:end_index]
+        except ValueError:
+            pass
+        return ids_to_text(result_ids)
+    except ValueError:
+        return ids_to_text(generated) # <sep>이 없으면 전체 반환
 
 print("\n\n===== 생성 결과 =====")  
-print(generate_text_topp(chat_model, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))
+# 모델이 1 epoch만 학습되었으므로 의미 있는 결과가 아닐 수 있습니다.
+print(generate_text_topp(chat_model, "지난 2년 동안 출연연이 국가가 필요한 연구를", p=0.9))