Update app.py

app.py

@@ -2,6 +2,32 @@ import gradio as gr
 import torch
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
+from typing import List, Literal, Sequence, TypedDict
+
+Role = Literal["system", "user", "assistant"]
+
+class Message(TypedDict):
+    role: Role
+    content: str
+
+Dialog = Sequence[Message]
+
+class ChatFormat:
+    def encode_header(self, message: Message) -> str:
+        return f"{message['role']}\n\n"
+
+    def encode_message(self, message: Message) -> str:
+        header = self.encode_header(message)
+        return f"{header}{message['content'].strip()}"
+
+    def encode_dialog_prompt(self, dialog: Dialog) -> str:
+        dialog_str = ""
+        for message in dialog:
+            dialog_str += self.encode_message(message)
+        dialog_str += self.encode_header({"role": "assistant", "content": ""})
+        return dialog_str
+
+
 class MedS_Llama3:
     def __init__(self, model_path: str):
         # 加载模型到CPU
@@ -19,10 +45,21 @@ class MedS_Llama3:
         )
         self.tokenizer.pad_token = self.tokenizer.eos_token
         self.model.eval()
+        self.prompt_engine = ChatFormat()
         print('Model and tokenizer loaded on CPU!')
-
+        
+    def __build_inputs_for_llama3(self, query: str, instruction: str) -> str:
+        input_ss = [
+            {"role": 'system', "content": instruction},
+            {"role": 'user', "content": query}
+        ]
+        return self.prompt_engine.encode_dialog_prompt(input_ss)
+    
     def chat(self, query: str, instruction: str, max_output_tokens: int) -> str:
-        input_sentence = f"{instruction}\n\n{query}"
+
+        formatted_query = f"Input:\n{query}\nOutput:\n"
+        input_sentence = self.__build_inputs_for_llama3(formatted_query, instruction)
+        
         input_tokens = self.tokenizer(
             input_sentence,
             return_tensors="pt",