Account for FIM tokens in prompt

huggingface · Jan 20, 2024 · e531c61 · e531c61
1 parent 585ea3a
commit e531c61
Showing 1 changed file with 18 additions and 12 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("");
@@ -592,6 +592,9 @@ impl Backend {
     async fn get_completions(&self, params: CompletionParams) -> Result<CompletionResult> {
         let request_id = Uuid::new_v4();
         let span = info_span!("completion_request", %request_id);
+        self.client
+        .log_message(MessageType::INFO, format!("get_completions"))
+        .await;
         async move {
             let document_map = self.document_map.read().await;
 
@@ -734,10 +737,14 @@ impl LanguageServer for Backend {
 
     async fn did_change(&self, params: DidChangeTextDocumentParams) {
         let uri = params.text_document.uri.to_string();
-        self.client
-            .log_message(MessageType::INFO, format!("{uri} changed"))
-            .await;
+        // ignore the output scheme
+        if uri.starts_with("output:") {
+            return;
+        }
         let mut document_map = self.document_map.write().await;
+        self.client
+        .log_message(MessageType::INFO, format!("{uri} changed"))
+        .await;
         let doc = document_map.get_mut(&uri);
         if let Some(doc) = doc {
             for change in &params.content_changes {
@@ -851,4 +858,3 @@ async fn main() {
 
     Server::new(stdin, stdout, socket).serve(service).await;
 }
-