More FIM templates

Oleksiy Syvokon created 1 week ago

Change summary

crates/ollama/src/ollama.rs                          | 16 -------
crates/ollama/src/ollama_edit_prediction_delegate.rs | 29 +++++++++++--
2 files changed, 23 insertions(+), 22 deletions(-)

Detailed changes

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<Item = &'a str>,
-) -> 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};

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!("<fim_prefix>{prefix}<fim_suffix>{suffix}<fim_middle>")
         }
@@ -402,3 +410,12 @@ struct OllamaGenerateOptions {
 struct OllamaGenerateResponse {
     response: String,
 }
+pub fn pick_recommended_edit_prediction_model<'a>(
+    available_models: impl IntoIterator<Item = &'a str>,
+) -> 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))
+}