AlphaS2S.py

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
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 = 150 # 기존 코드에서 200으로 설정됨
batch_size = 128

# 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 (기존 코드와 동일)
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*2):
            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/TinyInst/resolve/main/output.jsonl?download=true",
        DATA_PATH
    )

if not os.path.exists(TOKENIZER_PATH):
    download_file(
        "https://huggingface.co/datasets/Yuchan5386/TinyInst/resolve/main/ko_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}")

def text_to_ids(text):
    return sp.encode(text, out_type=int)

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:
            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>")
                
                # 인코더 입력은 <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 = []
                    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]
                
                # 마스킹된 타겟 생성 (프롬프트/패딩 부분은 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(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), # 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__()
        self.proj = layers.Dense(d_ff)
        self.out = 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))
        
class gMLPBlock(layers.Layer):
    def __init__(self, d_model, seq_len, dropout=0.1):
        super().__init__()
        self.d_model = d_model
        self.seq_len = seq_len
        self.norm = layers.LayerNormalization(epsilon=1e-6)
        
        # FFN: Channel Expansion
        # d_model * 4로 확장
        self.channel_proj = layers.Dense(d_model * 4, use_bias=True)
        self.dropout = layers.Dropout(dropout)
        
        # Spatial Gating Unit (SGU)
        self.sgu_norm = layers.LayerNormalization(epsilon=1e-6)
        self.sgu_proj = layers.Dense(seq_len, use_bias=False)
        
        # 출력 차원을 d_model * 2 (U의 차원)로 설정
        self.sgu_final = layers.Dense(d_model * 2, use_bias=True) 
        
        self.out_proj = layers.Dense(d_model, use_bias=True)

    def call(self, x, training=False):
        # 1. Norm and Channel Expansion
        residual = x
        x_norm = self.norm(x)
        x_proj = self.channel_proj(x_norm) # Shape: (B, L, 4*D)
        
        # 2. Split (U and V streams)
        u, v = tf.split(x_proj, 2, axis=-1) # u, v Shape: (B, L, 2*D)
        
        # 3. Spatial Gating Unit (SGU)
        v_norm = self.sgu_norm(v)
        v_norm_T = tf.transpose(v_norm, perm=[0, 2, 1]) # (B, 2D, L)
        
        # 💡 토큰 믹싱 발생 (시퀀스 축으로 Dense 적용)
        v_proj = self.sgu_proj(v_norm_T) # (B, 2D, L)
        v_proj_T = tf.transpose(v_proj, perm=[0, 2, 1]) # (B, L, 2D)
        
        # 4. Activation and Gate Generation
        # 표준 gMLP는 U에 GELU를 적용하고 V는 선형 게이트로 사용
        # 여기서는 U에 GELU를 적용
        u_act = tf.nn.gelu(u)
        v_gate = self.sgu_final(v_proj_T) # Shape: (B, L, 2*D) 
        
        # 5. Gating and Contraction
        z = u_act * v_gate # 게이팅
        z = self.dropout(z, training=training)
        out = self.out_proj(z) # Shape: (B, L, D)
        
        # 6. Residual Connection
        return residual + out

class CrossBlock(layers.Layer):
    def __init__(self, clip_value=5.0, eps=1e-6): # 💡 d_model 인자 추가
        super().__init__()
        self.clip_value = clip_value
        self.eps = eps
        self.attn = layers.MultiHeadAttention(8, 20)
        # 💡 수정: 출력 차원을 1에서 d_model로 변경
    def call(self, x, z):
        y = self.attn(x, z, z)
        return y

class LoU(layers.Layer):
    def __init__(self, d_model, clip_value=5.0, eps=1e-6):
        super().__init__()
        self.d_model = d_model
        self.clip_value = float(clip_value)
        self.mha = layers.MultiHeadAttention(8, 20)
        self.norm1 = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
        self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32')
        
        self.glu = SwiGLU(d_model, 350)
        self.cross = CrossBlock()

    def call(self, x, z):
        x_f32 = tf.cast(x, tf.float32)
        residual = x_f32
        x = self.norm1(x)

        x_comb = self.mha(x, x, x, use_causal_mask=True)
        
        out = self.norm(x_comb + residual)
        out = self.cross(out, z)
        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=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)
        
        # EncoderBlock과 LoU는 기존 코드와 동일한 구조
        self.enc_layers = [gMLPBlock(d_model, seq_len=max_len) 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와 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 # 인코더의 최종 출력 (디코더의 '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) 
        
        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)
    # 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)
    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(
        initial_learning_rate=initial_lr,
        decay_steps=decay_steps,
        decay_rate=decay_rate,
        staircase=False
    )

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')


    # 옵티마이저 설정
    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
        ]
    )
    chat_model.summary()
    print("✅ 모델 컴파일 완료, 학습 시작...")
    # ⚠️ 학습 실행
    history = chat_model.fit(dataset, epochs=1, verbose=1)

# 가중치 저장
chat_model.save_weights("chat_model.weights.h5")
print("\n✅ 모델 가중치 저장 완료!")

# =======================
# 6) 추론 함수 (기존 코드 유지)
# =======================

def generate_text_topp(model, prompt, max_len=150, 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):
        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])
        
        # 모델 추론 (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))

    # <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===== 생성 결과 =====")  
# 모델이 1 epoch만 학습되었으므로 의미 있는 결과가 아닐 수 있습니다.
print(generate_text_topp(chat_model, "제가 이따가 버스를 타야 해서 준비 좀 해야겠어요. 재미있는 대화였습니다!", p=0.9))