Upload 4 files

openlm.weights.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5db6b49d3d0ba41030b618650f72874a2ebf8150ebcac0371cc611e40f3d81ef
+size 32381304

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46f7f887c9d36c6bde12637bc280a8a121c57a1ef6a3034e0e8f417ac4bfa6e1
+size 517610

추론.py

@@ -0,0 +1,198 @@
+import numpy as np
+import tensorflow as tf
+import sentencepiece as spm
+from tensorflow.keras import layers, Model
+
+# 1. 토크나이저 로드
+TOKENIZER_PATH = "tokenizer.model"
+sp = spm.SentencePieceProcessor(TOKENIZER_PATH)
+pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
+bos_id = sp.piece_to_id("<s>") if sp.piece_to_id("<s>") != -1 else 1
+eos_id = sp.piece_to_id("</s>")
+vocab_size = sp.get_piece_size()
+
+# 2. 모델 인스턴스 생성 (학습 때와 동일한 하이퍼파라미터 필수)
+# 하드웨어 사양에 맞춰 Strategy 범위 밖에서 단일 장치(CPU/GPU)용으로 생성 가능합니다.
+d_model = 256
+n_layers = 6
+max_len = 1024
+
+
+class RotaryPositionalEmbedding(layers.Layer):
+    def __init__(self, dim, max_seq_len=2048):
+        super().__init__()
+        self.dim = dim
+        self.max_seq_len = max_seq_len
+        # 미리 주파수 정보를 계산해 둡니다.
+        inv_freq = 1.0 / (10000 ** (tf.range(0, dim, 2, dtype=tf.float32) / dim))
+        t = tf.range(max_seq_len, dtype=tf.float32)
+        freqs = tf.einsum('i,j->ij', t, inv_freq)
+        self.emb_sin = tf.exp(1.0) # 더미 초기화
+        self.emb_cos = tf.exp(1.0)
+        
+        # 상수로 저장하여 매번 계산하지 않게 함
+        self.sin_cached = tf.sin(freqs)[tf.newaxis, tf.newaxis, :, :]
+        self.cos_cached = tf.cos(freqs)[tf.newaxis, tf.newaxis, :, :]
+
+    def call(self, x):
+        seq_len = tf.shape(x)[2]
+        t_type = x.dtype
+        
+        # 캐싱된 테이블에서 현재 시퀀스 길이만큼만 잘라서 사용
+        emb_sin = tf.cast(self.sin_cached[:, :, :seq_len, :], t_type)
+        emb_cos = tf.cast(self.cos_cached[:, :, :seq_len, :], t_type)
+        
+        x1 = x[..., ::2]
+        x2 = x[..., 1::2]
+        out = tf.stack([x1 * emb_cos - x2 * emb_sin, 
+                        x1 * emb_sin + x2 * emb_cos], axis=-1)
+        return tf.reshape(out, tf.shape(x))
+
+class SwiGLU(layers.Layer):
+    def __init__(self, d_model, d_ff):
+        super().__init__()
+
+        self.w12 = layers.Dense(d_ff, use_bias=False)
+        self.w3 = layers.Dense(d_model, use_bias=False)
+
+    def call(self, x):
+        x_proj = self.w12(x)
+        x_gate, x_val = tf.split(x_proj, 2, axis=-1)
+        return self.w3(tf.nn.silu(x_gate) * x_val)
+
+class TransformerBlock(layers.Layer):
+    def __init__(self, d_model, d_ff, num_heads=4, dropout_rate=0.1):
+        super().__init__()
+        self.ln1 = layers.LayerNormalization(epsilon=1e-5)
+        self.mha = layers.MultiHeadAttention(num_heads=num_heads, key_dim=d_model // num_heads)
+        self.rope = RotaryPositionalEmbedding(d_model // num_heads)
+        
+        self.ln2 = layers.LayerNormalization(epsilon=1e-5)
+        self.ffn = SwiGLU(d_model, d_ff)
+        self.dropout = layers.Dropout(dropout_rate)
+
+    def call(self, x, training=False):
+        # Attention Path
+        skip = x
+        x = self.ln1(x)
+
+        b, s, d = tf.unstack(tf.shape(x))
+        h = self.mha._num_heads
+        dh = d // h
+        
+        qkv = tf.reshape(x, [b, s, h, dh])
+        qkv = tf.transpose(qkv, [0, 2, 1, 3]) # [b, h, s, dh]
+        qkv_rope = self.rope(qkv)
+        qkv_rope = tf.transpose(qkv_rope, [0, 2, 1, 3])
+        qkv_rope = tf.reshape(qkv_rope, [b, s, d])
+        
+        attn_out = self.mha(query=qkv_rope, value=x, key=qkv_rope, use_causal_mask=True, training=training)
+        attn_out = attn_out
+        x = skip + self.dropout(attn_out, training=training)
+        
+        # FFN Path
+        x = x + self.dropout(self.ffn(self.ln2(x)), training=training)
+        return x
+    
+class OpenLM(tf.keras.Model):
+    def __init__(self, vocab_size, d_model=256, n_layers=6, d_ff=1024, dropout_rate=0.1):
+        super().__init__()
+        self.d_model = d_model
+        # 1. 임베딩 레이어 정의
+        self.token_embedding = layers.Embedding(vocab_size, d_model)
+        
+        self.blocks = [TransformerBlock(d_model, d_ff, dropout_rate=dropout_rate) for _ in range(n_layers)]
+        self.ln_f = layers.LayerNormalization(epsilon=1e-5)
+        # lm_head 레이어는 삭제합니다.
+
+    def call(self, x, training=False):
+        # 입력 임베딩
+        x = self.token_embedding(x)
+        
+        for block in self.blocks:
+            x = block(x, training=training)
+        
+        # 최종 노멀라이제이션
+        x = self.ln_f(x)
+        
+        # 2. 가중치 공유 (Weight Tying) 구현
+        # embedding 가중치를 가져옵니다. shape: [vocab_size, d_model]
+        weights = self.token_embedding.embeddings
+        
+        # 연산을 위해 x를 float32로 올리거나, 가중치를 x의 타입에 맞춥니다.
+        # 여기서는 안정성을 위해 x와 가중치를 모두 float32로 캐스팅하여 연산합니다.
+        x = tf.cast(x, tf.float32)
+        weights = tf.cast(weights, tf.float32)
+        
+        # 3. 행렬곱 수행: [batch, seq, d_model] @ [d_model, vocab_size]
+        # transpose_b=True 옵션을 주면 weights를 [vocab_size, d_model]에서 [d_model, vocab_size]로 간주해 곱합니다.
+        logits = tf.matmul(x, weights, transpose_b=True)
+        
+        return logits # 이미 float32이므로 그대로 반환
+
+lm = OpenLM(vocab_size=vocab_size, d_model=d_model, n_layers=n_layers)
+
+# 가중치 로드를 위해 더미 입력으로 빌드
+dummy_input = tf.zeros((1, 1), dtype=tf.int32)
+_ = lm(dummy_input)
+
+# 3. 가중치 로드
+lm.load_weights("openlm.weights.h5")
+print("✅ 모델 가중치 로드 완료")
+
+def generate_text(model, tokenizer, prompt, max_new_tokens=100, temperature=0.8, top_p=0.9):
+    input_ids = tokenizer.encode_as_ids(prompt)
+    input_ids = tf.expand_dims(tf.convert_to_tensor(input_ids, dtype=tf.int32), 0)
+
+    for _ in range(max_new_tokens):
+        # 1. 모델 예측
+        curr_input = input_ids[:, -max_len:]
+        logits = model(curr_input, training=False)
+        next_token_logits = logits[:, -1, :] # [Batch, Vocab]
+
+        # 2. Temperature 적용
+        if temperature > 0:
+            next_token_logits = next_token_logits / temperature
+        else:
+            # temperature가 0이면 가장 높은 확률만 선택(Greedy)
+            next_token = tf.argmax(next_token_logits, axis=-1, output_type=tf.int32)
+            next_token = tf.reshape(next_token, [1, 1])
+            input_ids = tf.concat([input_ids, next_token], axis=-1)
+            if next_token[0, 0].numpy() == eos_id: break
+            continue
+
+        # 3. Top-p (Nucleus) Filtering
+        # 내림차순 정렬
+        sorted_logits = tf.sort(next_token_logits, direction='DESCENDING', axis=-1)
+        sorted_probs = tf.nn.softmax(sorted_logits, axis=-1)
+        cumulative_probs = tf.cumsum(sorted_probs, axis=-1)
+
+        # 누적 확률이 top_p를 넘는 로짓들을 찾아 매우 낮은 값으로 마스킹
+        # (첫 번째 토큰은 제외하기 위해 1칸 shift)
+        sorted_indices_to_remove = cumulative_probs > top_p
+        # 첫 번째(가장 높은 확률) 토큰은 무조건 유지
+        mask_shifted = tf.concat([tf.zeros_like(sorted_indices_to_remove[:, :1]), 
+                                 sorted_indices_to_remove[:, :-1]], axis=-1)
+        
+        # 정렬된 상태에서 마스킹 기준이 되는 '최소 로짓 값' 결정
+        # 마스킹될 위치의 로짓 중 가장 큰 값을 찾아 그보다 작은 로짓은 다 제거
+        min_logit_to_keep = tf.reduce_min(tf.where(mask_shifted, 1e10, sorted_logits), axis=-1, keepdims=True)
+        
+        # 원본 로짓에 마스크 적용
+        next_token_logits = tf.where(next_token_logits < min_logit_to_keep, -1e10, next_token_logits)
+
+        # 4. 샘플링
+        next_token = tf.random.categorical(next_token_logits, num_samples=1)
+        next_token = tf.cast(next_token, tf.int32)
+        
+        input_ids = tf.concat([input_ids, next_token], axis=-1)
+
+        if next_token[0, 0].numpy() == eos_id:
+            break
+
+    return tokenizer.decode_ids(input_ids[0].numpy().tolist())
+
+# --- 실행 예시 ---
+prompt_text = "Question: What is AI?\nAnswer:"
+result = generate_text(lm, sp, prompt_text, max_new_tokens=50, temperature=0.7)
+print(f"\n🚀 생성 결과:\n{result}")