Update app.py

app.py

@@ -122,6 +122,35 @@ def is_partial_stop(output, stop_str):
 tokenizer.chat_template = "{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
 
 
+# def prepare_history_for_model(history):
+#     """
+#     Converts the history to a tokenized prompt in the format expected by the model.
+#     Params:
+#       history: dialogue history
+#     Returns:
+#       Tokenized prompt
+#     """
+#     messages = []
+#     for idx, (user_msg, model_msg) in enumerate(history):
+#         # skip the last assistant message if its empty, the tokenizer will do the formating
+#         if idx == len(history) - 1 and not model_msg:
+#             messages.append({"role": "User", "content": user_msg})
+#             break
+#         if user_msg:
+#             messages.append({"role": "User", "content": user_msg})
+#         if model_msg:
+#             messages.append({"role": "Assistant", "content": model_msg})
+#     input_token = tokenizer.apply_chat_template(
+#         messages,
+#         add_generation_prompt=True,
+#         tokenize=True,
+#         return_tensors="pt",
+#         return_dict=True
+#     )
+#     return input_token
+
+
+
 def prepare_history_for_model(history):
     """
     Converts the history to a tokenized prompt in the format expected by the model.
@@ -131,15 +160,23 @@ def prepare_history_for_model(history):
       Tokenized prompt
     """
     messages = []
+    
+    # Add instruction
+    instruction = "Generate quotes for AWS RDS services"
+    messages.append({"role": "Instruction", "content": instruction})
+
     for idx, (user_msg, model_msg) in enumerate(history):
-        # skip the last assistant message if its empty, the tokenizer will do the formating
+        # Assuming the user message contains the product information
+        if user_msg:
+            messages.append({"role": "Input", "content": user_msg})
+        
+        # Skip the last assistant message if it's empty
         if idx == len(history) - 1 and not model_msg:
-            messages.append({"role": "User", "content": user_msg})
             break
-        if user_msg:
-            messages.append({"role": "User", "content": user_msg})
+        
         if model_msg:
-            messages.append({"role": "Assistant", "content": model_msg})
+            messages.append({"role": "Output", "content": model_msg})
+
     input_token = tokenizer.apply_chat_template(
         messages,
         add_generation_prompt=True,
@@ -150,6 +187,7 @@ def prepare_history_for_model(history):
     return input_token
 
 
+
 def generate(history, temperature, max_new_tokens, top_p, repetition_penalty, assisted):
     """
     Generates the assistant's reponse given the chatbot history and generation parameters