Update app.py

app.py

@@ -9,9 +9,16 @@ from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from dataclasses import dataclass
 
 
-chatml_template = """{% for message in messages %}
-    {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
-{% endfor %}"""
+# chatml_template = """{% for message in messages %}
+#     {{'<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>' + '\n'}}
+# {% endfor %}"""
+
+prompt_format = '''<|im_start|>system
+{system_message}<|im_end|>
+<|im_start|>user
+{prompt}<|im_end|>
+<|im_start|>assistant
+'''
 
 
 system_prompt = '''You are given a partial input text for another AI chat interface. 
@@ -25,16 +32,16 @@ Answers should be only the completions themselves. If you have nothing as a comp
 
 Examples: 
 (1)
-User: "Help me write a sentiment analysis pipeline"
-Assistant: "using huggingface;using NLTK;using python"
+User: Help me write a sentiment analysis pipeline
+Assistant: using huggingface;using NLTK;using python
 
 (2)
-User: "My name is"
-Assistant: "<NOTHING>" (nothing much to contribute at this point. return nothing)
+User: My name is
+Assistant: <NOTHING> (nothing much to contribute at this point. return nothing)
 
 (3)
-User: "Help me find a present for my"
-Assistant: "girlfriend;mother;father;friend"
+User: Help me find a present for my
+Assistant: girlfriend;mother;father;friend
 '''
 
 # setup
@@ -90,20 +97,20 @@ def set_past_key_values():
     return detach_past_kv(model(tokenized.to(model.device)).past_key_values)
     
 
-@spaces.GPU
+# @spaces.GPU
 def generate(text, past_key_values):
-    messages = [
-        *start_messages,
-        {'role': 'user', 'content': text}
-    ]
-
+    # messages = [
+    #     *start_messages,
+    #     {'role': 'user', 'content': text}
+    # ]
+    
     cur_generate_kwargs = deepcopy(generate_kwargs)
     
     if past_key_values:
         past_key_values = past_kv_to_device(past_key_values, pipe.model.device, pipe.model.dtype)
         cur_generate_kwargs.update({'past_key_values': past_key_values})
     
-    response = pipe(messages, **cur_generate_kwargs)[0]['generated_text']
+    response = pipe(prompt_format.format(system_message=system_prompt, prompt=text), **cur_generate_kwargs)[0]['generated_text']
     print(response)
     return response[-1]['content']