gered/llm-ls
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

Account for FIM tokens in prompt (#62)

Browse source
This commit is contained in:
Markus Hennerbichler 2024-01-28 22:16:01 +00:00 committed by GitHub
parent f40e8cc6ea
commit a5f2e87315
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 8 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
crates/llm-ls/src/main.rs
View file

@ -306,7 +306,7 @@ fn build_prompt(
) -> Result<String> { ) -> Result<String> {
let t = Instant::now(); let t = Instant::now();
if fim.enabled { 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 before_iter = text.lines_at(pos.line as usize + 1).reversed();
let mut after_iter = text.lines_at(pos.line as usize); let mut after_iter = text.lines_at(pos.line as usize);
let mut before_line = before_iter.next(); let mut before_line = before_iter.next();
@ -332,10 +332,10 @@ fn build_prompt(
} else { } else {
before_line.len() before_line.len()
}; };
if tokens > token_count { if tokens > remaining_token_count {
break; break;
} }
token_count -= tokens; remaining_token_count -= tokens;
before.push(before_line); before.push(before_line);
} }
if let Some(after_line) = after_line { if let Some(after_line) = after_line {
@ -348,10 +348,10 @@ fn build_prompt(
} else { } else {
after_line.len() after_line.len()
}; };
if tokens > token_count { if tokens > remaining_token_count {
break; break;
} }
token_count -= tokens; remaining_token_count -= tokens;
after.push_str(&after_line); after.push_str(&after_line);
} }
before_line = before_iter.next(); before_line = before_iter.next();
@ -369,7 +369,7 @@ fn build_prompt(
info!(prompt, build_prompt_ms = time, "built prompt in {time} ms"); info!(prompt, build_prompt_ms = time, "built prompt in {time} ms");
Ok(prompt) Ok(prompt)
} else { } else {
let mut token_count = context_window; let mut remaining_token_count = context_window;
let mut before = vec![]; let mut before = vec![];
let mut first = true; let mut first = true;
for mut line in text.lines_at(pos.line as usize + 1).reversed() { for mut line in text.lines_at(pos.line as usize + 1).reversed() {
@ -387,10 +387,10 @@ fn build_prompt(
} else { } else {
line.len() line.len()
}; };
if tokens > token_count { if tokens > remaining_token_count {
break; break;
} }
token_count -= tokens; remaining_token_count -= tokens;
before.push(line); before.push(line);
} }
let prompt = before.into_iter().rev().collect::<Vec<_>>().join(""); let prompt = before.into_iter().rev().collect::<Vec<_>>().join("");
@ -851,4 +851,3 @@ async fn main() {
Server::new(stdin, stdout, socket).serve(service).await; Server::new(stdin, stdout, socket).serve(service).await;
} }