07/01/23-15:16

app.py

@@ -20,13 +20,16 @@ def generate_code(input_text,modelName):
         result = codeGenTokenizer.decode(generated_ids[0], skip_special_tokens=True)
         return extract_code(result)
     elif(modelName == "mistral-7b"):
-        input_ids = mistralTokenizer(input_text, return_tensors="pt").input_ids
+        input_ids = mistralTokenizer(generate_prompt_mistral(input_text), return_tensors="pt").input_ids
         generated_ids = mistralModel.generate(input_ids, max_length=128)
         result = mistralTokenizer.decode(generated_ids[0], skip_special_tokens=True)
         return result
     else:
         return None
     
+def generate_prompt_mistral(text):
+    system_msg = "Below is an instruction that describes a programming task. Write a response code that appropriately completes the request.\n"
+    return f"<s>[INST] {system_msg}\n{text} [/INST]"
 
 def respond(message, chat_history, additional_inputs):
     return  f"Here's an example code:\n\n```python\n{generate_code(message,additional_inputs)}\n```"