diff --git a/crates/ollama/src/ollama.rs b/crates/ollama/src/ollama.rs index 24b87530ec88626b7795a83b4a82f92d4cb9a144..26d4be781c49d6e8bf6d7550b8785c6604f6aabf 100644 --- a/crates/ollama/src/ollama.rs +++ b/crates/ollama/src/ollama.rs @@ -4,22 +4,6 @@ pub use ollama_edit_prediction_delegate::OllamaEditPredictionDelegate; use anyhow::{Context, Result}; -pub const RECOMMENDED_EDIT_PREDICTION_MODELS: [&str; 4] = [ - "qwen2.5-coder:3b-base", - "qwen2.5-coder:7b-base", - "qwen2.5-coder:3b", - "qwen2.5-coder:7b", -]; - -pub fn pick_recommended_edit_prediction_model<'a>( - available_models: impl IntoIterator, -) -> Option<&'static str> { - let available: std::collections::HashSet<&str> = available_models.into_iter().collect(); - - RECOMMENDED_EDIT_PREDICTION_MODELS - .into_iter() - .find(|recommended| available.contains(recommended)) -} use futures::{AsyncBufReadExt, AsyncReadExt, StreamExt, io::BufReader, stream::BoxStream}; use http_client::{AsyncBody, HttpClient, HttpRequestExt, Method, Request as HttpRequest}; use serde::{Deserialize, Serialize}; diff --git a/crates/ollama/src/ollama_edit_prediction_delegate.rs b/crates/ollama/src/ollama_edit_prediction_delegate.rs index 78e8edf7908babc0e762fb530e0f586284110eac..e01ce909af7b067ec5dc93d5dc83a46e5006a1e3 100644 --- a/crates/ollama/src/ollama_edit_prediction_delegate.rs +++ b/crates/ollama/src/ollama_edit_prediction_delegate.rs @@ -16,7 +16,7 @@ use std::{ }; use text::ToOffset; -use crate::{OLLAMA_API_URL, get_models, pick_recommended_edit_prediction_model}; +use crate::{OLLAMA_API_URL, get_models}; pub const DEBOUNCE_TIMEOUT: Duration = Duration::from_millis(150); @@ -26,6 +26,13 @@ const EXCERPT_OPTIONS: EditPredictionExcerptOptions = EditPredictionExcerptOptio target_before_cursor_over_total_bytes: 0.66, }; +pub const RECOMMENDED_EDIT_PREDICTION_MODELS: [&str; 4] = [ + "qwen2.5-coder:3b-base", + "qwen2.5-coder:3b", + "qwen2.5-coder:7b-base", + "qwen2.5-coder:7b", +]; + #[derive(Clone)] struct CurrentCompletion { snapshot: BufferSnapshot, @@ -80,7 +87,7 @@ impl OllamaEditPredictionDelegate { raw: true, stream: false, options: Some(OllamaGenerateOptions { - num_predict: Some(256), + num_predict: Some(64), temperature: Some(0.2), stop: Some(get_stop_tokens()), }), @@ -311,11 +318,9 @@ fn format_fim_prompt(model: &str, prefix: &str, suffix: &str) -> String { } "deepseek-coder" | "deepseek-coder-v2" => { // DeepSeek uses special Unicode characters for FIM tokens - format!( - "<\u{ff5c}fim\u{2581}begin\u{ff5c}>{prefix}<\u{ff5c}fim\u{2581}hole\u{ff5c}>{suffix}<\u{ff5c}fim\u{2581}end\u{ff5c}>" - ) + format!("<|fim▁begin|>{prefix}<|fim▁hole|>{suffix}<|fim▁end|>") } - "qwen2.5-coder" | "qwen-coder" => { + "qwen2.5-coder" | "qwen-coder" | "qwen" => { format!("<|fim_prefix|>{prefix}<|fim_suffix|>{suffix}<|fim_middle|>") } "codegemma" => { @@ -324,6 +329,9 @@ fn format_fim_prompt(model: &str, prefix: &str, suffix: &str) -> String { "codestral" | "mistral" => { format!("[SUFFIX]{suffix}[PREFIX]{prefix}") } + "glm" | "glm-4" | "glm-4.5" => { + format!("<|code_prefix|>{prefix}<|code_suffix|>{suffix}<|code_middle|>") + } _ => { format!("{prefix}{suffix}") } @@ -402,3 +410,12 @@ struct OllamaGenerateOptions { struct OllamaGenerateResponse { response: String, } +pub fn pick_recommended_edit_prediction_model<'a>( + available_models: impl IntoIterator, +) -> Option<&'static str> { + let available: std::collections::HashSet<&str> = available_models.into_iter().collect(); + + RECOMMENDED_EDIT_PREDICTION_MODELS + .into_iter() + .find(|recommended| available.contains(recommended)) +}