Upload 3 files

blocklm.weights.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:53ee854cce43a1b6a1f226719e047907aa253e05e2e5fc0ca1eec1e87c9bb861
+size 216686536

head.weights.h5

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

xxx.py

@@ -0,0 +1,333 @@
+import sentencepiece as spm
+import os, numpy as np, tensorflow as tf
+from tensorflow.keras import layers
+
+# --- 환경 설정 ---
+TOKENIZER_PATH = r"C:\Users\yuchan\write1\openlm\tokenizer.model"
+sp = spm.SentencePieceProcessor(TOKENIZER_PATH)
+pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
+end_id = sp.piece_to_id("</s>")
+
+max_len = 1024
+vocab_size = sp.get_piece_size()
+
+def text_to_ids(text):
+    return sp.encode(text, out_type=int)
+
+def ids_to_text(ids):
+    return sp.decode(ids)
+
+class TimeMix(layers.Layer):
+    def __init__(self, d_model, layer_id, n_layers):
+        super().__init__()
+        self.d_model = d_model
+        ratio = (layer_id / (n_layers - 1)) if n_layers > 1 else 0.5
+        
+        # 기본 베이스 파라미터
+        decay_speed = np.arange(d_model)
+        self.time_decay = tf.Variable(-5 + 8 * (decay_speed / (d_model - 1)) ** (0.7 + 1.3 * ratio), dtype=tf.float32)
+        self.time_first = tf.Variable(np.ones(d_model) * np.log(0.3), dtype=tf.float32)
+
+        # --- 동적 프로젝션 레이어 (추가됨) ---
+        self.w_proj = layers.Dense(d_model, kernel_initializer='zeros', use_bias=False)
+        self.r_proj = layers.Dense(d_model, kernel_initializer='zeros', use_bias=False)
+        self.k_proj = layers.Dense(d_model, kernel_initializer='zeros', use_bias=False)
+        self.v_proj = layers.Dense(d_model, kernel_initializer='zeros', use_bias=False)
+        
+        self.key = layers.Dense(d_model, use_bias=False)
+        self.value = layers.Dense(d_model, use_bias=False)
+        self.receptance = layers.Dense(d_model, use_bias=False)
+        self.output_projection = layers.Dense(d_model, use_bias=False)
+
+        # 믹싱 계수
+        self.tm_w = tf.Variable(1 - (ratio ** 0.5), dtype=tf.float32)
+        self.tm_k = tf.Variable(1 - (ratio ** 0.5), dtype=tf.float32)
+        self.tm_v = tf.Variable(1 - (ratio ** 0.5), dtype=tf.float32)
+        self.tm_r = tf.Variable(1 - (ratio ** 0.2), dtype=tf.float32)
+
+    def call(self, x, state):
+        # state: [last_x, aa, bb, pp]
+        last_x, aa, bb, pp = state
+        t_type = x.dtype
+        
+        # 믹싱 계수 캐스팅
+        tm_w = tf.cast(self.tm_w, t_type)
+        tm_k = tf.cast(self.tm_k, t_type)
+        tm_v = tf.cast(self.tm_v, t_type)
+        tm_r = tf.cast(self.tm_r, t_type)
+
+        # 1. 동적 파라미터 생성을 위한 dx 계산
+        dx = x * tm_w + last_x * (1 - tm_w)
+
+        # 2. 동적 w, r, k, v 계산 (핵심 변경 사항)
+        # 추론 시에도 매번 x에 따라 파라미터가 변합니다.
+        w = tf.cast(self.time_decay, t_type) + tf.cast(self.w_proj(dx), t_type)
+        w = -tf.exp(tf.cast(w, tf.float32)) # 감쇠율 (float32 정밀도)
+        
+        r = self.receptance(x * tm_r + last_x * (1 - tm_r)) + self.r_proj(dx)
+        k = self.key(x * tm_k + last_x * (1 - tm_k)) + self.k_proj(dx)
+        v = self.value(x * tm_v + last_x * (1 - tm_v)) + self.v_proj(dx)
+
+        # 3. RNN 모드 WKV 연산
+        u = tf.cast(self.time_first, tf.float32)
+        kv, vv = tf.cast(k, tf.float32), tf.cast(v, tf.float32)
+
+        ww = u + kv
+        p = tf.maximum(pp, ww)
+        e1, e2 = tf.exp(pp - p), tf.exp(ww - p)
+        wkv = (e1 * aa + e2 * vv) / (e1 * bb + e2 + 1e-12)
+
+        # 다음 스테이트를 위한 업데이트 (동적 w 적용)
+        ww_next = w + pp
+        p_next = tf.maximum(ww_next, kv)
+        e1_next, e2_next = tf.exp(ww_next - p_next), tf.exp(kv - p_next)
+        
+        new_state = [x, e1_next * aa + e2_next * vv, e1_next * bb + e2_next, p_next]
+        return self.output_projection(tf.nn.sigmoid(r) * tf.cast(wkv, t_type)), new_state
+
+class ChannelMix(layers.Layer):
+    def __init__(self, d_model, layer_id, n_layers):
+        super().__init__()
+        ratio = (layer_id / (n_layers - 1)) if n_layers > 1 else 0.5
+        self.time_mix_k = tf.Variable(1 - (ratio ** 0.5), dtype=tf.float32)
+        self.time_mix_r = tf.Variable(1 - (ratio ** 0.5), dtype=tf.float32)
+        self.key = layers.Dense(int(d_model * 4.25), use_bias=False)
+        self.receptance = layers.Dense(d_model, use_bias=False)
+        self.value = layers.Dense(d_model, use_bias=False)
+
+    def call(self, x, last_x):
+        t_type = x.dtype
+        tm_k, tm_r = tf.cast(self.time_mix_k, t_type), tf.cast(self.time_mix_r, t_type)
+        k = self.key(x * tm_k + last_x * (1 - tm_k))
+        r = self.receptance(x * tm_r + last_x * (1 - tm_r))      
+        kv = self.value(tf.square(tf.nn.relu(k)))
+        return tf.nn.sigmoid(r) * kv, x
+
+class Block(layers.Layer):
+    def __init__(self, d_model, layer_id, n_layers):
+        super().__init__()
+        self.ln = layers.LayerNormalization(epsilon=1e-5)
+        self.time_mix = TimeMix(d_model, layer_id, n_layers)
+        self.channel_mix = ChannelMix(d_model, layer_id, n_layers)
+    def call(self, x, state):
+        ln_x = self.ln(x)
+        tm_out, tm_state = self.time_mix(ln_x, state[:4])
+        x = x + tm_out
+        cm_out, cm_last_x = self.channel_mix(ln_x, state[4])
+        x = x + cm_out
+        return x, tm_state + [cm_last_x]
+
+class Head(tf.keras.Model):
+    def __init__(self, vocab_size):
+        super().__init__()
+        self.lm_head = layers.Dense(vocab_size, use_bias=False, name="output_head")
+    def call(self, x):
+        return tf.cast(self.lm_head(x), tf.float32)
+    
+class LM(tf.keras.Model):
+    def __init__(self, d_model, n_layers):
+        super().__init__()
+        self.token_embedding = layers.Embedding(vocab_size, d_model)
+        self.blocks = [Block(d_model, i, n_layers) for i in range(n_layers)]
+        self.ln_f = layers.LayerNormalization(epsilon=1e-5, dtype=tf.float32)
+    def call(self, x, states):
+        x = self.token_embedding(x)
+        new_states = []
+        for i, block in enumerate(self.blocks):
+            x, b_state = block(x, states[i*5 : (i+1)*5])
+            new_states.extend(b_state)
+        return self.ln_f(x), new_states
+
+# --- 초기화 및 로드 ---
+d_model, n_layers = 512, 10
+blocklm = LM(d_model, n_layers)
+head = Head(vocab_size)
+
+def get_init_state():
+    return [tf.zeros((1, 1, d_model)) if i%5!=3 else tf.ones((1, 1, d_model))*-1e30 for i in range(n_layers*5)]
+
+_o, _s = blocklm(tf.constant([[0]]), get_init_state())
+_ = head(_o)
+blocklm.load_weights(r"C:\Users\yuchan\write1\blocklm.weights.h5")
+head.load_weights(r"C:\Users\yuchan\write1\head.weights.h5")
+
+
+import numpy as np
+import tensorflow as tf
+
+class InferenceEngine:
+    def __init__(self, model, head, sp, config):
+        self.model = model
+        self.head = head
+        self.sp = sp
+        self.config = config
+        # SentencePiece에서 특수 토큰 ID 추출
+        self.pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
+        # EOS 토큰이 없으면 보통 </s>나 [EOS]를 찾습니다.
+        self.eos_id = sp.piece_to_id("</s>")
+        if self.eos_id == -1:
+            self.eos_id = sp.piece_to_id("[EOS]")
+
+    def apply_repetition_penalty(self, logits, generated_ids, penalty):
+        """
+        Logits Scaling 방식의 패널티 부여 (HuggingFace 스타일)
+        최근 생성된 토큰들이 다시 나올 확률을 억제합니다.
+        """
+        if not generated_ids:
+            return logits
+            
+        # 설정된 윈도우 크기 내에서 등장한 토큰들만 패널티 대상
+        recent_ids = set(generated_ids[-self.config.get('penalty_window', 64):])
+        
+        for token_id in recent_ids:
+            score = logits[token_id]
+            if score > 0:
+                # 양수 스코어는 낮추고
+                logits[token_id] /= penalty
+            else:
+                # 음수 스코어는 더 낮게 (절댓값을 키움)
+                logits[token_id] *= penalty
+        return logits
+
+    def sample(self, logits, temp, top_k, top_p):
+        """Temperature, Top-K, Top-P가 결합된 수치 안정적 샘플링"""
+        # 1. Temperature 적용
+        if temp > 0:
+            logits = logits / temp
+        else:
+            # Greedy Search (가장 높은 확률 선택)
+            return np.argmax(logits)
+
+        # 2. Top-K 필터링
+        if top_k > 0:
+            top_k = min(top_k, logits.shape[-1])
+            indices_to_remove = logits < np.sort(logits)[-top_k]
+            logits[indices_to_remove] = -float('inf')
+
+        # 3. Top-P (Nucleus) 필터링
+        # 확률 도메인에서 계산하기 위해 softmax 적용
+        probs = tf.nn.softmax(logits).numpy()
+        sorted_indices = np.argsort(probs)[::-1]
+        sorted_probs = probs[sorted_indices]
+        cumulative_probs = np.cumsum(sorted_probs)
+
+        # 문맥 유지를 위해 최소 1개 토큰은 남기고 나머지는 제거
+        idx_to_remove = cumulative_probs > top_p
+        if np.any(idx_to_remove):
+            cutoff_idx = np.where(idx_to_remove)[0][0] + 1
+            cutoff_idx = max(1, cutoff_idx)
+            probs[sorted_indices[cutoff_idx:]] = 0
+            # 확률 재정규화
+            if np.sum(probs) > 0:
+                probs = probs / np.sum(probs)
+            else:
+                # 폴백: 가장 높은 확률 토큰에 1 몰아주기
+                probs[sorted_indices[0]] = 1.0
+
+        # 4. 최종 샘플링
+        return np.random.choice(len(probs), p=probs)
+
+    @tf.function(reduce_retracing=True)
+    def model_step(self, token_id, states):
+        """TPU/GPU 가속을 위한 Graph 모드 추론 실행부"""
+        out, next_states = self.model(token_id, states)
+        logits = self.head(out)
+        return logits, next_states
+
+    def generate(self, prompt, 
+                 max_new_tokens=512, 
+                 temperature=0.7, 
+                 top_k=40, 
+                 top_p=0.9, 
+                 repetition_penalty=1.2):
+        """
+        스트리밍 방식으로 텍스트를 생성하는 제너레이터.
+        """
+        # 입력 텍스트 인코딩
+        input_ids = self.sp.encode(prompt)
+        states = get_init_state() # 외부 정의된 초기 상태 함수 호출
+        generated = []
+
+        # [1] Prefill 단계: 프롬프트의 마지막 토큰 전까지 상태를 미리 계산
+        if len(input_ids) > 1:
+            for i in range(len(input_ids) - 1):
+                _, states = self.model_step(tf.constant([[input_ids[i]]]), states)
+        
+        # [2] Decoding 단계: 마지막 토큰을 시작으로 다음 토큰들 생성
+        curr_token_id = input_ids[-1]
+        prev_text = ""
+        
+        for _ in range(max_new_tokens):
+            curr_token_tensor = tf.constant([[curr_token_id]])
+            logits_out, states = self.model_step(curr_token_tensor, states)
+            
+            # (batch, seq, vocab) 구조에서 현재 토큰의 로짓만 추출
+            logits = logits_out[0, 0].numpy()
+
+            # 후처리: 패널티 및 패딩 방지
+            logits = self.apply_repetition_penalty(logits, input_ids + generated, repetition_penalty)
+            logits[self.pad_id] = -float('inf')
+
+            # 샘플링 실행
+            next_id = int(self.sample(logits, temperature, top_k, top_p))
+            
+            # --- 중단 조건: EOS 토큰 감지 ---
+            if next_id == self.eos_id:
+                break
+
+            generated.append(next_id)
+            
+            # --- 띄어쓰기 유지 디코딩 로직 ---
+            # 지금까지 생성된 토큰들을 통째로 디코딩한 뒤, 
+            # 이전에 출력한 텍스트를 제외한 나머지(증분)만 추출
+            full_text = self.sp.decode(generated)
+            new_part = full_text[len(prev_text):]
+            
+            if new_part:
+                yield new_part
+                prev_text = full_text
+
+            # 다음 루프를 위해 현재 토큰 업데이트
+            curr_token_id = next_id
+
+# --- 설정 및 엔진 초기화 ---
+config = {
+    'penalty_window': 64,  # 최근 64개 토큰 이내 반복 억제
+}
+
+# InferenceEngine 인스턴스 생성 (blocklm, head, sp는 외부 정의 상태)
+engine = InferenceEngine(blocklm, head, sp, config)
+
+# --- 실행 루프 (Main Loop) ---
+print("===== Advanced Dynamic RWKV Engine (No Omissions) =====")
+while True:
+    try:
+        user_input = input("\n[User]: ").strip()
+        if user_input.lower() in ['exit', 'quit']:
+            print("프로그램을 종료합니다.")
+            break
+        
+        if not user_input:
+            continue
+            
+        # 모델에게 답변 형식을 명확히 인지시키는 프롬프트 템플릿
+        # 문맥 유지를 위해 질문 뒤에 Answer:를 붙이고 공백을 유도합니다.
+        full_prompt = f"Question: {user_input}\nAnswer:" 
+        
+        print("[AI]:", end=" ", flush=True)
+        
+        # 생성 시작
+        for delta in engine.generate(
+            full_prompt, 
+            max_new_tokens=1024,
+            temperature=0.7, 
+            top_p=0.92, 
+            repetition_penalty=1.2
+        ):
+            print(delta, end="", flush=True)
+            
+        print() # 생성 종료 후 줄바꿈
+        
+    except KeyboardInterrupt:
+        print("\n중단되었습니다.")
+        break