From 06eb07fa90032bac6b2432b48a1e1ab0209247eb Mon Sep 17 00:00:00 2001
From: Kevsnz <15689588+Kevsnz@users.noreply.github.com>
Date: Sat, 3 Feb 2024 16:58:21 +0500
Subject: [PATCH] Added space for better infills and proper stop token

---
 src/prompts/processors/models.ts | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/prompts/processors/models.ts b/src/prompts/processors/models.ts
index 6133028..058905f 100644
--- a/src/prompts/processors/models.ts
+++ b/src/prompts/processors/models.ts
@@ -19,7 +19,7 @@ export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix:
     }
 
     // Stable code FIM
-    if (args.format === 'stable-code')  {
+    if (args.format === 'stable-code') {
         return {
             prompt: `<fim_prefix>${args.prefix}<fim_suffix>${args.suffix}<fim_middle>`,
             stop: [`<|endoftext|>`]
@@ -28,7 +28,7 @@ export function adaptPrompt(args: { format: ModelFormat, prefix: string, suffix:
 
     // Codellama FIM
     return {
-        prompt: `<PRE> ${args.prefix} <SUF>${args.suffix} <MID>`,
-        stop: [`<PRE>`, `<SUF>`, `<MID>`, `<END>`]
+        prompt: `<PRE> ${args.prefix} <SUF> ${args.suffix} <MID>`,
+        stop: [`<END>`, `<EOD>`, `<EOT>`]
     };
 }
\ No newline at end of file