Update app/main.py

app/main.py

@@ -29,16 +29,19 @@ def extract_restext(response, is_chat=False):
   return response['choices'][0]['text' if is_chat else 'message'].strip()
 
 def ask_llama(llm: Llama, question: str, max_new_tokens=200, temperature=0.5, repeat_penalty=2.0):
-  result = extract_restext(llm.create_chat_completion({"role": "user", "content": question}, max_tokens=max_new_tokens, temperature=temperature, repeat_penalty=repeat_penalty, stop=["<|eot_id|>", "<|end_of_text|>"]), is_chat=True)
+  prompt = f"""<|begin_of_text|>
+<|start_header_id|> user  <|end_header_id|> {question} <|eot_id|>
+<|start_header_id|> assistant <|end_header_id|>"""
+  result = extract_restext(llm.complete(prompt, max_tokens=max_new_tokens, temperature=temperature, repeat_penalty=repeat_penalty, stop=["<|eot_id|>", "<|end_of_text|>"])).replace("<|eot_id|>", "").replace("<|end_of_text|>", "")
   return result
 
-def chat_llama(llm: Llama, chat_history: dict, max_new_tokens=200, temperature=0.5, repeat_penalty=2.0):
-  result = extract_restext(llm.create_chat_completion(chat_history, max_tokens=max_new_tokens, temperature=temperature, repeat_penalty=repeat_penalty, stop=["<|eot_id|>", "<|end_of_text|>"]), is_chat=True)
-  return result
+# def chat_llama(llm: Llama, chat_history: dict, max_new_tokens=200, temperature=0.5, repeat_penalty=2.0):
+#   result = extract_restext(llm.create_chat_completion(chat_history, max_tokens=max_new_tokens, temperature=temperature, repeat_penalty=repeat_penalty, stop=["<|eot_id|>", "<|end_of_text|>"]), is_chat=True)
+#   return result
 
 # TESTING THE MODEL
 print("Testing model...")
-assert ask_llama(PHllm, ["Hello!, How are you today?"], max_new_tokens=5) #Just checking that it can run
+assert ask_llama(PHllm, "Hello!, How are you today?", max_new_tokens=5) #Just checking that it can run
 print("Checking Translators.")
 assert t.translate("Hello!") == "สวัสดี!"
 assert e.translate("สวัสดี!") == "Hello!"