mirror of
https://github.com/zed-industries/zed.git
synced 2026-10-15 00:32:45 +00:00
edit_prediction: Send only format-specific FIM stop sequences (#65360)
## Summary The FIM edit prediction path (`open_ai_compatible_api` and `ollama`) always sent a fixed list of 14 stop sequences, covering the control tokens of every supported prompt format, regardless of the configured `prompt_format`. OpenAI-compatible servers that enforce the OpenAI limit of 4 `stop` entries (e.g. vLLM) reject every such request with a 400, so no predictions appear and the error is only visible in the log. This change replaces `get_fim_stop_tokens()` with `fim_stop_tokens(prompt_format)`, which returns only the markers that `format_fim_prompt` uses for the selected format, plus the format's end-of-text token where applicable, and never more than 4 entries. The prompt format is already resolved (including `Infer`) before `fim::request_prediction` is called, so no settings or schema changes are needed. `clean_fim_completion` is unchanged and still truncates the completion at all known markers client-side as a safety net. Closes #65359 ## Testing - Added `fim_stop_tokens_respect_openai_stop_limit`, which asserts that every FIM prompt format sends between 1 and 4 stop sequences. - Added `fim_stop_tokens_match_prompt_markers`, which asserts that every non end-of-text stop token actually appears in the prompt built by `format_fim_prompt` for that format, so the two can't drift apart. - `cargo test -p edit_prediction --lib fim` To reproduce manually: run vLLM with a Qwen coder model, configure `edit_predictions.provider` as `open_ai_compatible_api` with `prompt_format: "qwen"`, and type in a source file. Before this change every request fails with a 400 about too many stop sequences; after it, predictions appear and the request body contains exactly 4 `stop` entries. ## Self-review - [x] I've reviewed my diff for quality, security, reliability, and performance. - ~~[ ] UI changes follow the [checklist](https://zed.dev/docs/development/ui-checklist).~~ - [x] Tests cover the new or changed behavior. Release Notes: - Fixed OpenAI-compatible edit predictions failing on servers that limit stop sequences to 4. --------- Co-authored-by: zed-zippy[bot] <234243425+zed-zippy[bot]@users.noreply.github.com>
This commit is contained in:
parent
129c68ff4b
commit
efed1f245a
1 changed files with 78 additions and 18 deletions
|
|
@ -98,7 +98,10 @@ pub fn request_prediction(
|
||||||
let prefix = editable_text[..cursor_in_editable].to_string();
|
let prefix = editable_text[..cursor_in_editable].to_string();
|
||||||
let suffix = editable_text[cursor_in_editable..].to_string();
|
let suffix = editable_text[cursor_in_editable..].to_string();
|
||||||
let prompt = format_fim_prompt(prompt_format, &prefix, &suffix);
|
let prompt = format_fim_prompt(prompt_format, &prefix, &suffix);
|
||||||
let stop_tokens = get_fim_stop_tokens();
|
let stop_tokens = fim_stop_tokens(prompt_format)
|
||||||
|
.iter()
|
||||||
|
.map(|token| token.to_string())
|
||||||
|
.collect();
|
||||||
|
|
||||||
let max_tokens = settings.max_output_tokens;
|
let max_tokens = settings.max_output_tokens;
|
||||||
|
|
||||||
|
|
@ -219,23 +222,40 @@ fn format_fim_prompt(
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn get_fim_stop_tokens() -> Vec<String> {
|
// OpenAI-compatible servers such as vLLM reject requests with more than 4 stop
|
||||||
vec![
|
// sequences, so each format may only list up to 4 tokens.
|
||||||
"<|endoftext|>".to_string(),
|
fn fim_stop_tokens(prompt_format: EditPredictionPromptFormat) -> &'static [&'static str] {
|
||||||
"<|file_separator|>".to_string(),
|
match prompt_format {
|
||||||
"<|fim_pad|>".to_string(),
|
EditPredictionPromptFormat::CodeLlama => &["<PRE>", "<SUF>", "<MID>"],
|
||||||
"<|fim_prefix|>".to_string(),
|
EditPredictionPromptFormat::DeepseekCoder => {
|
||||||
"<|fim_middle|>".to_string(),
|
&["<|fim▁begin|>", "<|fim▁hole|>", "<|fim▁end|>"]
|
||||||
"<|fim_suffix|>".to_string(),
|
}
|
||||||
"<fim_prefix>".to_string(),
|
EditPredictionPromptFormat::Qwen => &[
|
||||||
"<fim_middle>".to_string(),
|
"<|endoftext|>",
|
||||||
"<fim_suffix>".to_string(),
|
"<|fim_prefix|>",
|
||||||
"<PRE>".to_string(),
|
"<|fim_middle|>",
|
||||||
"<SUF>".to_string(),
|
"<|fim_suffix|>",
|
||||||
"<MID>".to_string(),
|
],
|
||||||
"[PREFIX]".to_string(),
|
EditPredictionPromptFormat::CodeGemma => &[
|
||||||
"[SUFFIX]".to_string(),
|
"<|file_separator|>",
|
||||||
]
|
"<|fim_prefix|>",
|
||||||
|
"<|fim_middle|>",
|
||||||
|
"<|fim_suffix|>",
|
||||||
|
],
|
||||||
|
EditPredictionPromptFormat::Codestral => &["[PREFIX]", "[SUFFIX]"],
|
||||||
|
EditPredictionPromptFormat::Glm => &[
|
||||||
|
"<|endoftext|>",
|
||||||
|
"<|code_prefix|>",
|
||||||
|
"<|code_suffix|>",
|
||||||
|
"<|code_middle|>",
|
||||||
|
],
|
||||||
|
_ => &[
|
||||||
|
"<|endoftext|>",
|
||||||
|
"<fim_prefix>",
|
||||||
|
"<fim_middle>",
|
||||||
|
"<fim_suffix>",
|
||||||
|
],
|
||||||
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
fn clean_fim_completion(response: &str) -> String {
|
fn clean_fim_completion(response: &str) -> String {
|
||||||
|
|
@ -315,6 +335,46 @@ mod tests {
|
||||||
);
|
);
|
||||||
}
|
}
|
||||||
|
|
||||||
|
const FIM_PROMPT_FORMATS: [EditPredictionPromptFormat; 7] = [
|
||||||
|
EditPredictionPromptFormat::CodeLlama,
|
||||||
|
EditPredictionPromptFormat::StarCoder,
|
||||||
|
EditPredictionPromptFormat::DeepseekCoder,
|
||||||
|
EditPredictionPromptFormat::Qwen,
|
||||||
|
EditPredictionPromptFormat::CodeGemma,
|
||||||
|
EditPredictionPromptFormat::Codestral,
|
||||||
|
EditPredictionPromptFormat::Glm,
|
||||||
|
];
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn fim_stop_tokens_respect_openai_stop_limit() {
|
||||||
|
for format in FIM_PROMPT_FORMATS {
|
||||||
|
let stop_tokens = fim_stop_tokens(format);
|
||||||
|
assert!(!stop_tokens.is_empty(), "{format:?} has no stop tokens");
|
||||||
|
assert!(
|
||||||
|
stop_tokens.len() <= 4,
|
||||||
|
"{format:?} has {} stop tokens, but OpenAI allows at most 4",
|
||||||
|
stop_tokens.len()
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn fim_stop_tokens_match_prompt_markers() {
|
||||||
|
let end_of_text_tokens = ["<|endoftext|>", "<|file_separator|>"];
|
||||||
|
for format in FIM_PROMPT_FORMATS {
|
||||||
|
let prompt = format_fim_prompt(format, "a", "b");
|
||||||
|
for token in fim_stop_tokens(format) {
|
||||||
|
if end_of_text_tokens.contains(token) {
|
||||||
|
continue;
|
||||||
|
}
|
||||||
|
assert!(
|
||||||
|
prompt.contains(token),
|
||||||
|
"stop token {token:?} for {format:?} does not appear in prompt {prompt:?}"
|
||||||
|
);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
#[test]
|
#[test]
|
||||||
fn infer_prompt_format_returns_none_for_unsupported_models() {
|
fn infer_prompt_format_returns_none_for_unsupported_models() {
|
||||||
assert_eq!(infer_prompt_format("llama3:70b"), None);
|
assert_eq!(infer_prompt_format("llama3:70b"), None);
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue