Create app.py

app.py

@@ -0,0 +1,242 @@
+import sentencepiece as spm
+import os, numpy as np, tensorflow as tf
+from tensorflow.keras import layers
+import gradio as gr
+
+# --- 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
+end_id = sp.piece_to_id("</s>")
+vocab_size = sp.get_piece_size()
+
+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):
+        last_x, aa, bb, pp = state
+        t_type = x.dtype
+        tm_w, tm_k, tm_v, tm_r = tf.cast(self.tm_w, t_type), tf.cast(self.tm_k, t_type), tf.cast(self.tm_v, t_type), tf.cast(self.tm_r, t_type)
+        dx = x * tm_w + last_x * (1 - tm_w)
+        w = tf.cast(self.time_decay, t_type) + tf.cast(self.w_proj(dx), t_type)
+        w = -tf.exp(tf.cast(w, tf.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)
+        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)
+        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
+
+# --- 2. 모델 로드 및 초기화 ---
+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)]
+
+# 가중치 구조 생성을 위한 Dummy Call
+_o, _s = blocklm(tf.constant([[0]]), get_init_state())
+_ = head(_o)
+
+# 가중치 파일 로드
+blocklm.load_weights("blocklm.weights.h5")
+head.load_weights("head.weights.h5")
+
+# --- 3. 추론 엔진 정의 (기존 유지) ---
+class InferenceEngine:
+    def __init__(self, model, head, sp):
+        self.model = model
+        self.head = head
+        self.sp = sp
+        self.pad_id = sp.piece_to_id("<pad>") if sp.piece_to_id("<pad>") != -1 else 0
+        self.eos_id = sp.piece_to_id("</s>") if sp.piece_to_id("</s>") != -1 else sp.piece_to_id("[EOS]")
+
+    def apply_repetition_penalty(self, logits, generated_ids, penalty, window):
+        if not generated_ids: return logits
+        recent_ids = set(generated_ids[-window:])
+        for token_id in recent_ids:
+            if logits[token_id] > 0: logits[token_id] /= penalty
+            else: logits[token_id] *= penalty
+        return logits
+
+    def sample(self, logits, temp, top_k, top_p):
+        if temp <= 0: return np.argmax(logits)
+        logits = logits / temp
+        if top_k > 0:
+            indices_to_remove = logits < np.sort(logits)[-min(top_k, logits.shape[-1])]
+            logits[indices_to_remove] = -float('inf')
+        
+        probs = tf.nn.softmax(logits).numpy()
+        sorted_indices = np.argsort(probs)[::-1]
+        sorted_probs = probs[sorted_indices]
+        cumulative_probs = np.cumsum(sorted_probs)
+        idx_to_remove = cumulative_probs > top_p
+        if np.any(idx_to_remove):
+            cutoff_idx = max(1, np.where(idx_to_remove)[0][0] + 1)
+            probs[sorted_indices[cutoff_idx:]] = 0
+            if np.sum(probs) > 0: probs /= np.sum(probs)
+            else: probs[sorted_indices[0]] = 1.0
+        return np.random.choice(len(probs), p=probs)
+
+    @tf.function(reduce_retracing=True)
+    def model_step(self, token_id, states):
+        out, next_states = self.model(token_id, states)
+        logits = self.head(out)
+        return logits, next_states
+
+    def generate_stream(self, prompt, max_new_tokens, temperature, top_k, top_p, penalty, window):
+        input_ids = self.sp.encode(prompt)
+        states = get_init_state()
+        generated = []
+
+        if len(input_ids) > 1:
+            for i in range(len(input_ids) - 1):
+                _, states = self.model_step(tf.constant([[input_ids[i]]]), states)
+        
+        curr_token_id = input_ids[-1]
+        prev_text = ""
+        
+        for _ in range(max_new_tokens):
+            logits_out, states = self.model_step(tf.constant([[curr_token_id]]), states)
+            logits = logits_out[0, 0].numpy()
+            logits = self.apply_repetition_penalty(logits, input_ids + generated, penalty, window)
+            logits[self.pad_id] = -float('inf')
+
+            next_id = int(self.sample(logits, temperature, top_k, top_p))
+            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
+
+engine = InferenceEngine(blocklm, head, sp)
+
+# --- 4. Gradio 인터페이스 구성 ---
+def chat_response(message, history, max_tokens, temp, top_p, top_k, penalty):
+    # 대화 맥락을 포함한 프롬프트 생성
+    # 간단한 구조: Question: {msg}\nAnswer:
+    full_prompt = f"Question: {message}\nAnswer:"
+    
+    partial_message = ""
+    for delta in engine.generate_stream(
+        full_prompt, 
+        max_new_tokens=max_tokens,
+        temperature=temp,
+        top_k=top_k,
+        top_p=top_p,
+        penalty=penalty,
+        window=64
+    ):
+        partial_message += delta
+        yield partial_message
+
+# Gradio 테마 및 레이아웃 설정
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🚀 Dynamic Engine Chatbot")
+    gr.Markdown("동적 파라미터를 사용하는 모델을 위한 실시간 스트리밍 채팅 UI입니다.")
+    
+    with gr.Row():
+        with gr.Column(scale=4):
+            chatbot = gr.ChatInterface(
+                fn=chat_response,
+                additional_inputs=[
+                    gr.Slider(1, 2048, value=512, step=1, label="Max New Tokens"),
+                    gr.Slider(0.0, 2.0, value=0.7, step=0.1, label="Temperature"),
+                    gr.Slider(0.0, 1.0, value=0.92, step=0.01, label="Top-P"),
+                    gr.Slider(0, 100, value=40, step=1, label="Top-K"),
+                    gr.Slider(1.0, 2.0, value=1.2, step=0.05, label="Repetition Penalty"),
+                ],
+                examples=[["What is AI?"], ["Hello."]],
+            )
+        
+    gr.Markdown("---")
+    gr.Markdown("### 🛠 Model Info")
+    gr.Markdown(f"- **D_Model**: {d_model} | **Layers**: {n_layers} | **Vocab**: {vocab_size}")
+
+if __name__ == "__main__":
+    # share=True를 설정하면 외부 공유 링크가 생성됩니다.
+    demo.queue().launch(share=True)