Update app.py

app.py

@@ -8,7 +8,7 @@ from lit_llama import LLaMA, Tokenizer
 from lit_llama.utils import EmptyInitOnDevice
 
 
-class ChatDoctor:
+class ChatBot:
     def __init__(self, model, tokenizer, fabric):
         self.model = model
         self.tokenizer = tokenizer
@@ -26,17 +26,18 @@ class ChatDoctor:
             "환자의 문의 내용에 대해 답변하세요. 환자의 질병을 진단하고, 가능하면 처방을 하세요. \n\n"
             f"### 문의:\n{example['instruction']}\n\n### 응답:"
     )
-    
-    # This method generates the chatbot's responses.
+
+    # default generation
     @torch.no_grad()
     def generate(
-        self, 
-        idx, 
-        max_new_tokens, 
+        self,
+        idx,
+        max_new_tokens,
         max_seq_length=None,
-        temperature=0.8, 
-        top_k=None, 
-        eos_id=None
+        temperature=0.8,
+        top_k=None,
+        eos_id=None,
+        repetition_penalty=1.1,
     ):
         T = idx.size(0)
         T_new = T + max_new_tokens
@@ -86,14 +87,9 @@ class ChatDoctor:
 
         return idx
 
-
-    # This method handles user's messages and updates the conversation history.
-    def user(self, user_message, history):
-        # The user's message is added to the history with None as the bot's response.
-        return "", history + [[user_message, None]]
-
-    # This method generates and handles bot's responses.
-    def bot(self, history, max_new_tokens, top_k, temperature):
+    # LLM generation 함수
+    def ans(self, user_message, history, max_new_tokens, top_k, temperature):
+        history = history + [[user_message, None]]
         instruction = history[-1][0].strip()
         sample = { "instruction" : instruction, "input" : None }
         prompt = self.generate_prompt(sample)
@@ -106,17 +102,15 @@ class ChatDoctor:
             top_k=top_k,
             eos_id=self.tokenizer.eos_id
         )
-        
+
         self.model.reset_cache()
 
         response = self.tokenizer.decode(y)
         response = response.split('응답:')[1].strip()
-        
-        # The history is updated with the bot's response.
-        history[-1][1] = response
-        
-        return history
 
+        # history 업데이트
+        history[-1][1] = response
+        return response
 
 def load_model():
     # Settings for inference
@@ -143,82 +137,66 @@ def load_model():
 
     return model, tokenizer, fabric
 
-
-def setup_gradio_ui(chat_doctor):
-    with gr.Blocks(theme=gr.themes.Monochrome()) as demo:
-        gr.Markdown(
-        """
-        # ChatDoctor-KR Demo
-        
-        last modified : 23.05.18
-        """)
-
-        chatbot = gr.Chatbot(label="ChatDoctor-KR")
-        msg = gr.Textbox(lines=1, placeholder="질문 입력 후 엔터를 누르세요.", label="질문")
-        clear = gr.Button("클리어")
-
-        gr.Markdown(
-        """
-        ## Parameters
-        """)
-        
-        max_new_tokens = gr.Slider(
-            minimum=1, 
-            maximum=512, 
-            step=1, 
-            value=512, 
-            label="max_new_tokens", 
-            info="The number of new tokens to generate",
-            interactive=True
-        )
-
-        top_k = gr.Slider(
+# theme 'Taithrah/Minimal' 'abidlabs/dracula_test' 'JohnSmith9982/small_and_pretty'
+def setup_gradio_ui(chat_bot, css):
+    gr.ChatInterface(
+        fn=chat_bot.ans,
+        css=css,
+        textbox=gr.Textbox(placeholder="질문을 입력해주세요.", container=False, scale=7),
+        chatbot=gr.Chatbot(height=600, value=[[None, "안녕하세요. 무엇이 궁금하신가요?"]], avatar_images=["asset/human.png", "asset/bot.jpg"]),
+        title="의료용 챗봇 데모",
+        theme='soft',
+        examples=[["두통이 너무 심해요."], ["배가 아프고 토할것 같아요."], ["허리가 끊어질 듯이 아파요."]],
+        submit_btn=gr.Button(value="전송", icon="send.png", elem_id="green"),
+        retry_btn=gr.Button(value="다시보내기 (재질문)↩", elem_id="blue"),
+        undo_btn=gr.Button(value="이전챗 삭제 ❌", elem_id="blue"),
+        clear_btn=gr.Button(value="전챗 삭제 💫", elem_id="blue"),
+        additional_inputs=[
+            gr.Slider(
             minimum=1,
-            maximum=300, 
-            step=1, 
-            value=200,
-            label="top_k",
-            info="If specified, only sample among the tokens with the k highest probabilities",
-            interactive=True
-        )
-
-        temperature = gr.Slider(
-            minimum=0.1, 
-            maximum=1.0, 
-            step=0.1, 
-            value=0.8,
-            label="temperature",
-            info="Scales the predicted logits by 1 / temperature",
+            maximum=512,
+            step=1,
+            value=512,
+            label="max_new_tokens",
+            info="최대 생성 가능 토큰 수",
             interactive=True
-        )
-
-        with gr.Accordion(label="Open for More!", open=False):
-            gr.Markdown("Blah Blah ...")
-
-        submit_result = msg.submit(
-            chat_doctor.user, [msg, chatbot], [msg, chatbot], queue=False 
-        )
-        submit_result.then(
-            chat_doctor.bot, [chatbot, max_new_tokens, top_k, temperature], chatbot
-        )
-
-        # This part clears the chatbot history when the clear button is clicked.
-        clear.click(lambda: None, None, chatbot, queue=False)
-    
-    demo.queue()
-    
-    demo.launch(share=True, server_name="0.0.0.0")
-
+            ),
+
+            gr.Slider(
+                minimum=1,
+                maximum=300,
+                step=1,
+                value=150,
+                label="top_k",
+                info="확률이 가장 높은 토큰 k개 샘플링",
+                interactive=True
+            ),
+
+            gr.Slider(
+                minimum=0.1,
+                maximum=1.0,
+                step=0.1,
+                value=0.5,
+                label="temperature",
+                info="1에 가까울수록 다양한 답변 생성",
+                interactive=True
+            )
+        ]
+    ).queue().launch()
 
 def main():
-    # Load model and tokenizer
+    # 모델, 토크나이저 로드
     model, tokenizer, fabric = load_model()
 
-    # ChatDoctor instance
-    chat_doctor = ChatDoctor(model, tokenizer, fabric)
+    # 챗봇 객체 생성
+    chat_bot = ChatBot(model, tokenizer, fabric)
+
+    # ui
+    css = """
+    #green {background-color: #00EF91}
+    #blue {background-color: #B9E2FA}
+    """
+    setup_gradio_ui(chat_bot, css)
 
-    # Gradio UI setup and launch
-    setup_gradio_ui(chat_doctor)
-    
 if __name__ == "__main__":
-    main()
+    main()