Update main.py

main.py

@@ -33,6 +33,7 @@ file_handler.setFormatter(formatter)
 logger = logging.getLogger('gunicorn.error')
 logger.addHandler(file_handler)
 
+llm_prompt_format = os.getenv('LLM_PROMPT_FORMAT', None)
 model_path = os.environ.get('LLAMACPP_PATH', None)
 
 
@@ -75,19 +76,35 @@ async def read_device():
 def create_generated_text(messages: list[dict[str, str]] = Body(...), temperature: float = Body(default=1.0)):
     input_text = ''
     
-    for message in messages:
-        if message['role'] == 'system' or message['role'] == 'user':
-            input_text += f"<start_of_turn>user\n{message['content']}<end_of_turn>\n"
-        elif message['role'] == 'assistant':
-            input_text += f"<start_of_turn>model\n{message['content']}<end_of_turn>\n"
+    if llm_prompt_format == 'Llama':
+        for message in messages:
+            if message['role'] == 'system':
+                input_text += f"<|start_header_id|>system<|end_header_id|>\n\n{message['content']}<|eot_id|>"
+            elif message['role'] == 'user':
+                input_text += f"<|start_header_id|>user<|end_header_id|>\n\n{message['content']}<|eot_id|>"
+            elif message['role'] == 'assistant':
+                input_text += f"<|start_header_id|>assistant<|end_header_id|>\n\n{message['content']}<|eot_id|>"
+
+        input_text += '<|start_header_id|>assistant<|end_header_id|>\n\n'
+        pattern = r'<|start_header_id|>assistant<|end_header_id|>\n\n(.+?)(?:(?:<|eot_id|>)|$)'
+
+    else:
+        for message in messages:
+            if message['role'] == 'system' or message['role'] == 'user':
+                input_text += f"<start_of_turn>user\n{message['content']}<end_of_turn>\n"
+            elif message['role'] == 'assistant':
+                input_text += f"<start_of_turn>model\n{message['content']}<end_of_turn>\n"
+
+        input_text += '<start_of_turn>model\n'
+        pattern = r'<start_of_turn>model\n(.+?)(?:(?:<end_of_turn>)|$)'
 
     if len(input_text) > 0:
         llm = Llama(model_path=model_path, n_ctx=8192, n_gpu_layers=-1, n_batch=32, verbose=False)
         choices = []
         
         try:
-            for choice in llm(input_text + '<start_of_turn>model\n', max_tokens=2048, temperature=temperature, top_p=0.95, echo=True)['choices']:
-                matches = re.findall(r'<start_of_turn>model\n(.+?)(?:(?:<end_of_turn>)|$)', choice['text'], re.DOTALL)
+            for choice in llm(input_text, max_tokens=2048, temperature=temperature, top_p=0.95, echo=True)['choices']:
+                matches = re.findall(pattern, choice['text'], re.DOTALL)
 
                 if len(matches) > 0:
                     choices.append({'role': 'assistant', 'content': matches[len(matches) - 1]})