Account for FIM tokens in prompt

huggingface · Jan 23, 2024 · 50e52e1 · 50e52e1
1 parent 585ea3a
commit 50e52e1
Showing 1 changed file with 8 additions and 9 deletions.
diff --git a/crates/llm-ls/src/main.rs b/crates/llm-ls/src/main.rs
@@ -306,7 +306,7 @@ fn build_prompt(
 ) -> Result<String> {
     let t = Instant::now();
     if fim.enabled {
-        let mut token_count = context_window;
+        let mut remaining_token_count = context_window - 3; // account for FIM tokens
         let mut before_iter = text.lines_at(pos.line as usize + 1).reversed();
         let mut after_iter = text.lines_at(pos.line as usize);
         let mut before_line = before_iter.next();
@@ -332,10 +332,10 @@ fn build_prompt(
                 } else {
                     before_line.len()
                 };
-                if tokens > token_count {
+                if tokens > remaining_token_count {
                     break;
                 }
-                token_count -= tokens;
+                remaining_token_count -= tokens;
                 before.push(before_line);
             }
             if let Some(after_line) = after_line {
@@ -348,10 +348,10 @@ fn build_prompt(
                 } else {
                     after_line.len()
                 };
-                if tokens > token_count {
+                if tokens > remaining_token_count {
                     break;
                 }
-                token_count -= tokens;
+                remaining_token_count -= tokens;
                 after.push_str(&after_line);
             }
             before_line = before_iter.next();
@@ -369,7 +369,7 @@ fn build_prompt(
         info!(prompt, build_prompt_ms = time, "built prompt in {time} ms");
         Ok(prompt)
     } else {
-        let mut token_count = context_window;
+        let mut remaining_token_count = context_window;
         let mut before = vec![];
         let mut first = true;
         for mut line in text.lines_at(pos.line as usize + 1).reversed() {
@@ -387,10 +387,10 @@ fn build_prompt(
             } else {
                 line.len()
             };
-            if tokens > token_count {
+            if tokens > remaining_token_count {
                 break;
             }
-            token_count -= tokens;
+            remaining_token_count -= tokens;
             before.push(line);
         }
         let prompt = before.into_iter().rev().collect::<Vec<_>>().join("");
@@ -851,4 +851,3 @@ async fn main() {
 
     Server::new(stdin, stdout, socket).serve(service).await;
 }
-