Update app.py

app.py

@@ -4,7 +4,8 @@ from tensorflow.keras import layers
 import gradio as gr
 
 # --- 1. 환경 설정 및 모델 구조 정의 ---
-TOKENIZER_PATH = "tokenizer.model" # 파일 이름만 사용
+# 파일 이름만 사용 (현재 작업 디렉토리에 파일이 있어야 함)
+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>")
@@ -114,11 +115,10 @@ 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")
 
@@ -189,38 +189,46 @@ class InferenceEngine:
 
 engine = InferenceEngine(blocklm, head, sp)
 
-# --- 4. Gradio UI (Manual Layout) ---
-with gr.Blocks(title="RWKV Chatbot") as demo:
-    gr.Markdown("## 🤖 Dynamic RWKV LLM Chat")
-    
-    chatbot = gr.Chatbot(label="Chat History")
-    msg = gr.Textbox(placeholder="질문을 입력하세요...", label="Input")
+# --- 4. Gradio UI (단순 텍스트 입출력 방식) ---
+with gr.Blocks(title="RWKV Text Generator") as demo:
+    gr.Markdown("## 🖋️ Dynamic RWKV Text Generation")
+    gr.Markdown("질문을 입력하고 Generate를 누르면 답변이 아래 텍스트 박스에 실시간으로 생성됩니다.")
     
     with gr.Row():
-        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:"
+        with gr.Column():
+            input_text = gr.Textbox(lines=5, label="Input Prompt", placeholder="여기에 질문이나 문장을 입력하세요...")
+            with gr.Row():
+                temp_slider = gr.Slider(0, 2, value=0.7, label="Temperature")
+                top_p_slider = gr.Slider(0, 1, value=0.92, label="Top-P")
+            with gr.Row():
+                penalty_slider = gr.Slider(1, 2, value=1.2, label="Penalty")
+                max_tokens = gr.Slider(1, 2048, value=512, step=1, label="Max Tokens")
+            
+            submit_btn = gr.Button("Generate", variant="primary")
+            clear_btn = gr.Button("Clear")
+
+        with gr.Column():
+            output_text = gr.Textbox(lines=15, label="Generated Output", interactive=False)
+
+    def run_generation(prompt, tokens, temp, top_p, penalty):
+        if not prompt.strip():
+            return "프롬프트를 입력해주세요."
         
-        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
+        full_prompt = f"Question: {prompt}\nAnswer:"
+        current_output = ""
+        
+        for chunk in engine.generate(full_prompt, int(tokens), temp, 40, top_p, penalty):
+            current_output += chunk
+            yield current_output
+
+    # 버튼 클릭 및 엔터 키 입력 이벤트
+    submit_btn.click(
+        fn=run_generation, 
+        inputs=[input_text, max_tokens, temp_slider, top_p_slider, penalty_slider], 
+        outputs=output_text
     )
-    clear.click(lambda: None, None, chatbot, queue=False)
+    
+    clear_btn.click(lambda: ("", ""), outputs=[input_text, output_text])
 
 if __name__ == "__main__":
     demo.queue().launch()