!pip install sentencepiece
import sentencepiece as spm

# 불러오기
import os, json, numpy as np, tensorflow as tf
import requests
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-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 = 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 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):
        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.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))

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)

    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 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.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):
        # Embedding
        x = self.token_embedding(x)  # (batch, seq_len, d_model)
        for block in self.blocks:
            x = block(x, training=training)

        x = self.ln_f(x)  # (batch, seq_len, d_model)
        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)
    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 = 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])

    # 학습
    history = model.fit(dist_dataset, epochs=1, verbose=1)

# =======================
# 가중치 저장
# =======================
model.save_weights("tf_model.weights.h5")
print("✅ 모델 가중치 저장 완료!")

# =======================
# 샘플 생성 함수
# =======================
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)
    
    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)

# =======================
# 테스트 생성
# =======================
prompt = "안녕하세요! 한국 밴드에 대해 궁금한 것이 있어요!"
sample_text = generate_text_topp(model, prompt, p=0.9)
print("\n===== 생성 결과 =====\n")
print(sample_text)