Update app.py

app.py

@@ -3,8 +3,8 @@ import os, numpy as np, tensorflow as tf
 from tensorflow.keras import layers
 import gradio as gr
 
-# --- 1. 환경 설정 및 모델 구조 정의 (기존 유지) ---
-TOKENIZER_PATH = "tokenizer.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
 end_id = sp.piece_to_id("</s>")
@@ -106,7 +106,7 @@ class LM(tf.keras.Model):
             new_states.extend(b_state)
         return self.ln_f(x), new_states
 
-# --- 2. 모델 로드 및 초기화 ---
+# --- 2. 초기화 및 가중치 로드 ---
 d_model, n_layers = 512, 10
 blocklm = LM(d_model, n_layers)
 head = Head(vocab_size)
@@ -114,15 +114,15 @@ 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
+# 구조 생성을 위한 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. 추론 엔진 정의 (기존 유지) ---
+# --- 3. 추론 엔진 ---
 class InferenceEngine:
     def __init__(self, model, head, sp):
         self.model = model
@@ -131,7 +131,7 @@ class InferenceEngine:
         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):
+    def apply_repetition_penalty(self, logits, generated_ids, penalty, window=64):
         if not generated_ids: return logits
         recent_ids = set(generated_ids[-window:])
         for token_id in recent_ids:
@@ -145,7 +145,6 @@ class InferenceEngine:
         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]
@@ -164,27 +163,22 @@ class InferenceEngine:
         logits = self.head(out)
         return logits, next_states
 
-    def generate_stream(self, prompt, max_new_tokens, temperature, top_k, top_p, penalty, window):
+    def generate(self, prompt, max_new_tokens, temp, top_k, top_p, penalty):
         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.apply_repetition_penalty(logits, input_ids + generated, penalty)
             logits[self.pad_id] = -float('inf')
-
-            next_id = int(self.sample(logits, temperature, top_k, top_p))
+            next_id = int(self.sample(logits, temp, 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):]
@@ -195,48 +189,38 @@ class InferenceEngine:
 
 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:"
+# --- 4. Gradio UI (Manual Layout) ---
+with gr.Blocks(title="RWKV Chatbot") as demo:
+    gr.Markdown("## 🤖 Dynamic RWKV LLM Chat")
     
-    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입니다.")
+    chatbot = gr.Chatbot(label="Chat History")
+    msg = gr.Textbox(placeholder="질문을 입력하세요...", label="Input")
     
     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."]],
-            )
+        temp_slider = gr.Slider(0, 2, value=0.7, label="Temperature")
+        top_p_slider = gr.Slider(0, 1, value=0.92, label="Top-P")
+        penalty_slider = gr.Slider(1, 2, value=1.2, label="Penalty")
+        max_tokens = gr.Slider(1, 1024, value=512, step=1, label="Max Tokens")
+    
+    clear = gr.Button("Clear")
+
+    def user(user_message, history):
+        return "", history + [[user_message, None]]
+
+    def bot(history, temp, top_p, penalty, tokens):
+        user_message = history[-1][0]
+        full_prompt = f"Question: {user_message}\nAnswer:"
         
-    gr.Markdown("---")
-    gr.Markdown("### 🛠 Model Info")
-    gr.Markdown(f"- **D_Model**: {d_model} | **Layers**: {n_layers} | **Vocab**: {vocab_size}")
+        history[-1][1] = ""
+        for chunk in engine.generate(full_prompt, tokens, temp, 40, top_p, penalty):
+            history[-1][1] += chunk
+            yield history
+
+    # 이벤트 연결: 엔터를 치거나 전송 시 작동
+    msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
+        bot, [chatbot, temp_slider, top_p_slider, penalty_slider, max_tokens], chatbot
+    )
+    clear.click(lambda: None, None, chatbot, queue=False)
 
 if __name__ == "__main__":
-    # share=True를 설정하면 외부 공유 링크가 생성됩니다.
-    demo.queue().launch(share=True)
+    demo.queue().launch()