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

feat: add tokens to clear (#19)

Browse source
This commit is contained in:
Luc Georges 2023-09-20 19:24:38 +02:00 committed by GitHub
parent 5378a67ce8
commit 7f9c7855d5
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 10 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -125,7 +125,7 @@ struct CompletionParams {
fim: FimParams, fim: FimParams,
api_token: Option<String>, api_token: Option<String>,
model: String, model: String,
model_eos: String, tokens_to_clear: Vec<String>,
tokenizer_path: Option<String>, tokenizer_path: Option<String>,
context_window: usize, context_window: usize,
tls_skip_verify_insecure: bool, tls_skip_verify_insecure: bool,
@ -257,11 +257,15 @@ async fn request_completion(
} }
} }
fn parse_generations(generations: Vec<Generation>, eos: &str) -> Vec<Completion> { fn parse_generations(generations: Vec<Generation>, tokens_to_clear: &[String]) -> Vec<Completion> {
generations generations
.into_iter() .into_iter()
.map(|g| Completion { .map(|g| {
generated_text: g.generated_text.replace(eos, ""), let mut generated_text = g.generated_text;
for token in tokens_to_clear {
generated_text = generated_text.replace(token, "")
}
Completion { generated_text }
}) })
.collect() .collect()
} }
@ -380,7 +384,7 @@ impl Backend {
) )
.await?; .await?;
Ok(parse_generations(result, &params.model_eos)) Ok(parse_generations(result, &params.tokens_to_clear))
} }
} }
@ -391,7 +395,7 @@ impl LanguageServer for Backend {
Ok(InitializeResult { Ok(InitializeResult {
server_info: Some(ServerInfo { server_info: Some(ServerInfo {
name: "llm-ls".to_owned(), name: "llm-ls".to_owned(),
version: Some("0.0.3".to_owned()), version: Some("0.1.0".to_owned()),
}), }),
capabilities: ServerCapabilities { capabilities: ServerCapabilities {
text_document_sync: Some(TextDocumentSyncCapability::Kind( text_document_sync: Some(TextDocumentSyncCapability::Kind(