Update app.py

app.py

@@ -86,6 +86,8 @@ def build_chat():
     prompt = conv.get_prompt() 
     return prompt
 
+from fastchat.model import get_conversation_template
+
 @spaces.GPU
 def generate(
     message: str,
@@ -97,15 +99,18 @@ def generate(
     top_k: int = 50,
     repetition_penalty: float = 1.2,
 ) -> Iterator[str]:
-    conversation = []
-    if system_prompt:
-        conversation.append({"role": "system", "content": system_prompt})
-    for user, assistant in chat_history:
-        conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
-    conversation.append({"role": "user", "content": message})
-
-    chat = tokenizer.apply_chat_template(conversation, tokenize=False)
-    inputs = tokenizer(chat, return_tensors="pt", add_special_tokens=False).to("cuda")
+    conv = get_conversation_template("vicuna")
+    conv.append_message(conv.roles[0], message)
+    conv.append_message(conv.roles[1], None)
+    prompt = conv.get_prompt()
+    # if system_prompt:
+    #     conversation.append({"role": "system", "content": system_prompt})
+    # for user, assistant in chat_history:
+    #     conversation.extend([{"role": "user", "content": user}, {"role": "assistant", "content": assistant}])
+    # conversation.append({"role": "user", "content": message})
+
+    # chat = tokenizer.apply_chat_template(conversation, tokenize=False)
+    inputs = tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to("cuda")
     if len(inputs) > MAX_INPUT_TOKEN_LENGTH:
         inputs = inputs[-MAX_INPUT_TOKEN_LENGTH:]
         gr.Warning("Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")