| !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 |
| |
| from tensorflow.keras import mixed_precision |
| policy = mixed_precision.Policy('mixed_float16') |
| mixed_precision.set_global_policy(policy) |
|
|
|
|
| print('1') |
| tf.get_logger().setLevel("ERROR") |
| SEED = 42 |
| tf.random.set_seed(SEED) |
| np.random.seed(SEED) |
|
|
| |
| 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 |
|
|
| |
| from tensorflow.keras import mixed_precision |
| policy = mixed_precision.Policy("mixed_bfloat16" if on_tpu else "float32") |
| mixed_precision.set_global_policy(policy) |
| print("โ
Mixed precision:", policy) |
|
|
| |
| |
| |
| 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 = "corpus.txt" |
| TOKENIZER_PATH = "ko_unigram.model" |
|
|
| if not os.path.exists(DATA_PATH): |
| download_file( |
| "https://huggingface.co/datasets/OpenLab-NLP/corpus-ko-22m/resolve/main/shuffled_corpus.txt?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", |
| 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 = 128 |
| batch_size = 96 |
|
|
| 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] |
|
|
| full_input = ids + [end_id] |
| pad_len = max_len - len(full_input) |
| full_input += [pad_id] * pad_len |
|
|
| |
| target = full_input[1:] + [pad_id] |
| yield ( |
| tf.convert_to_tensor(full_input, dtype=tf.int32), |
| tf.convert_to_tensor(target, dtype=tf.int32) |
| ) |
|
|
|
|
| steps_per_epoch = 23119910 // batch_size |
| LIMIT = 23119910 |
|
|
| 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__() |
| 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 Lo(layers.Layer): |
| def __init__(self, d_model): |
| super().__init__() |
| self.d = layers.Dense(64, activation='silu', dtype='float16') |
| self.w = layers.Dense(d_model, dtype='float16') |
| self.norm = layers.LayerNormalization(epsilon=1e-5, dtype='float32') |
|
|
| def call(self, x): |
| p = self.d(x) |
| p = self.w(p) |
| p = self.norm(p) |
| return tf.cast(p, x.dtype) + x |
|
|
|
|
| class Block(layers.Layer): |
| def __init__(self, d_model): |
| super().__init__() |
| self.mha = layers.MultiHeadAttention(8, 384//8) |
| self.glu = SwiGLU(d_model, 1048) |
| self.lo = Lo(d_model) |
|
|
| def call(self, x): |
| x = self.mha(x) |
| x = self.glu(x) |
| x = self.lo(x) |
| return x |
|
|
| class LaSLM(tf.keras.Model): |
| def __init__(self, vocab_size, max_seq_len, d_model, n_layers, dropout_rate=0.1): |
| super().__init__() |
| self.token_embedding = layers.Embedding(vocab_size, d_model, dtype=policy.compute_dtype) |
| self.pos_embedding = layers.Embedding(max_seq_len, d_model, dtype=policy.compute_dtype) |
| self.blocks = [Block(d_model) for _ in range(n_layers)] |
| self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype='float32') |
|
|
| 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, :] |
| x = self.token_embedding(x) + self.pos_embedding(positions) |
| for block in self.blocks: |
| x = block(x) |
| x = self.ln_f(x) |
| embedding_matrix = tf.cast(self.token_embedding.embeddings, x.dtype) |
| logits = tf.matmul(x, embedding_matrix, transpose_b=True) |
| 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) |
|
|
|
|
| with strategy.scope(): |
| model = LaSLM(vocab_size=vocab_size, max_seq_len=max_len, d_model=384, n_layers=3) |
| 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) |
|
|
| |
| history = model.fit(dist_dataset, epochs=1, steps_per_epoch=steps_per_epoch, verbose=1) |
| model.save_weights("tf_model.weights.h5") |
| print("โ
๋ชจ๋ธ ๊ฐ์ค์น ์ ์ฅ ์๋ฃ!") |
|
|
| def generate_text_topp(model, prompt, max_len=500, max_gen=500, 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) |
| for step in range(max_gen): |
| if len(generated) > 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) |
| 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] |
| 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) |
|
|
| print("\n\n===== ์์ฑ ๊ฒฐ๊ณผ =====") |
| print(generate_text_topp(model, "์ง๋ 2๋
๋์ ์ถ์ฐ์ฐ์ด ๊ตญ๊ฐ๊ฐ ํ์ํ ์ฐ๊ตฌ๋ฅผ", p=0.9)) |
