Merge pull request #608 from janhq/fix/llama3_chat_template

fix: add llama3 chat template
janhq · May 23, 2024 · 7cf63fa · 7cf63fa
2 parents c549d2e + d8b4523
commit 7cf63fa
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 0 deletions.
diff --git a/cortex-js/src/infrastructure/commanders/prompt-constants.ts b/cortex-js/src/infrastructure/commanders/prompt-constants.ts
@@ -17,6 +17,12 @@ export const ZEPHYR_JINJA = `{% for message in messages %}
 {% endif %}
 {% endfor %}`;
 
+export const LLAMA_3_JINJA = `{% set loop_messages = messages %}{% for message in loop_messages %}{% set content = '<|start_header_id|>' + message['role'] + '<|end_header_id|>
+
+'+ message['content'] | trim + '<|eot_id|>' %}{% if loop.index0 == 0 %}{% set content = bos_token + content %}{% endif %}{{ content }}{% endfor %}{% if add_generation_prompt %}{{ '<|start_header_id|>assistant<|end_header_id|>
+
+' }}{% endif %}`;
+
 //// Corresponding prompt template
 export const OPEN_CHAT_3_5 = `GPT4 Correct User: {prompt}<|end_of_turn|>GPT4 Correct Assistant:`;
 
@@ -35,3 +41,5 @@ export const LLAMA_2 = `[INST] <<SYS>>
 {system_message}
 <</SYS>>
 {prompt}[/INST]`;
+export const LLAMA_3 =
+  '<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\n{system_message}<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{prompt}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n';
diff --git a/cortex-js/src/infrastructure/commanders/usecases/models.cli.usecases.ts b/cortex-js/src/infrastructure/commanders/usecases/models.cli.usecases.ts
@@ -14,6 +14,8 @@ import { Inject, Injectable } from '@nestjs/common';
 import { Presets, SingleBar } from 'cli-progress';
 import {
   LLAMA_2,
+  LLAMA_3,
+  LLAMA_3_JINJA,
   OPEN_CHAT_3_5,
   OPEN_CHAT_3_5_JINJA,
   ZEPHYR,
@@ -262,6 +264,9 @@ export class ModelsCliUsecases {
       case OPEN_CHAT_3_5_JINJA:
         return OPEN_CHAT_3_5;
 
+      case LLAMA_3_JINJA:
+        return LLAMA_3;
+
       default:
         // console.log(
         //   'Unknown jinja code:',