oak-team/oak-gpui
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Add ep distill command, for generating edit prediction training examples (#44670)

Browse Source 
Release Notes:

- N/A

---------

Co-authored-by: Oleksiy Syvokon <oleksiy@zed.dev>
Co-authored-by: Agus Zubiaga <agus@zed.dev>
This commit is contained in:
Max Brunsfeld
2025-12-11 14:57:58 -08:00
committed by GitHub
co-authored by Oleksiy Syvokon Agus Zubiaga
parent fc6ca38989
commit 1cd7563f04
11 changed files with 149 additions and 113 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/edit_prediction_cli/Cargo.toml
+1 -1
View File
@@ -52,7 +52,7 @@ sqlez_macros.workspace = true
terminal_view.workspace = true
util.workspace = true
watch.workspace = true
edit_prediction = { workspace = true, features = ["eval-support"] }
edit_prediction = { workspace = true, features = ["cli-support"] }
wasmtime.workspace = true
zeta_prompt.workspace = true
zlog.workspace = true
crates/edit_prediction_cli/src/distill.rs
+14
View File
@@ -0,0 +1,14 @@
use std::mem;
use crate::example::Example;
pub async fn run_distill(example: &mut Example) {
let [prediction]: [_; 1] = mem::take(&mut example.predictions)
.try_into()
.expect("Run predict first with a single repetition");
example.expected_patch = prediction.actual_patch;
example.prompt = None;
example.predictions = Vec::new();
example.score = Vec::new();
}
crates/edit_prediction_cli/src/example.rs
+9 -14
View File
@@ -25,6 +25,7 @@ pub struct Example {
pub name: String,
pub repository_url: String,
pub revision: String,
#[serde(default)]
pub uncommitted_diff: String,
pub cursor_path: Arc<Path>,
pub cursor_position: String,
@@ -195,9 +196,9 @@ pub fn read_examples(inputs: &[PathBuf]) -> Vec<Example> {
.enumerate()
.map(|(line_ix, line)| {
let mut example =
serde_json::from_str::<Example>(line).unwrap_or_else(|_| {
serde_json::from_str::<Example>(line).unwrap_or_else(|error| {
panic!(
"Failed to parse example on {}:{}",
"Failed to parse example on {}:{}\n{error}",
path.display(),
line_ix + 1
)
@@ -264,12 +265,12 @@ fn parse_markdown_example(id: String, input: &str) -> Result<Example> {
state: None,
};
let mut name = String::new();
let mut text = String::new();
let mut block_info: CowStr = "".into();
#[derive(PartialEq)]
enum Section {
Start,
UncommittedDiff,
EditHistory,
CursorPosition,
@@ -278,14 +279,16 @@ fn parse_markdown_example(id: String, input: &str) -> Result<Example> {
Other,
}
let mut current_section = Section::Other;
let mut current_section = Section::Start;
for event in parser {
match event {
Event::Text(line) => {
text.push_str(&line);
if let Some((field, value)) = line.split_once('=') {
if let Section::Start = current_section
&& let Some((field, value)) = line.split_once('=')
{
match field.trim() {
REPOSITORY_URL_FIELD => {
example.repository_url = value.trim().to_string();
@@ -297,14 +300,6 @@ fn parse_markdown_example(id: String, input: &str) -> Result<Example> {
}
}
}
Event::End(TagEnd::Heading(HeadingLevel::H1)) => {
if !name.is_empty() {
anyhow::bail!(
"Found multiple H1 headings. There should only be one with the name of the example."
);
}
name = mem::take(&mut text);
}
Event::End(TagEnd::Heading(HeadingLevel::H2)) => {
let title = mem::take(&mut text);
current_section = if title.eq_ignore_ascii_case(UNCOMMITTED_DIFF_HEADING) {
@@ -363,7 +358,7 @@ fn parse_markdown_example(id: String, input: &str) -> Result<Example> {
Section::ExpectedPatch => {
example.expected_patch = mem::take(&mut text);
}
Section::Other => {}
Section::Start | Section::Other => {}
}
}
_ => {}
crates/edit_prediction_cli/src/format_prompt.rs
+73 -74
View File
@@ -2,9 +2,13 @@ use crate::{
PromptFormat,
example::{Example, ExamplePrompt},
headless::EpAppState,
load_project::run_load_project,
retrieve_context::run_context_retrieval,
};
use edit_prediction::{EditPredictionStore, zeta2::zeta2_prompt_input};
use edit_prediction::{
EditPredictionStore,
zeta2::{zeta2_output_for_patch, zeta2_prompt_input},
};
use gpui::AsyncApp;
use std::sync::Arc;
use zeta_prompt::format_zeta_prompt;
@@ -15,11 +19,20 @@ pub async fn run_format_prompt(
app_state: Arc<EpAppState>,
mut cx: AsyncApp,
) {
run_context_retrieval(example, app_state, cx.clone()).await;
run_context_retrieval(example, app_state.clone(), cx.clone()).await;
let prompt = match prompt_format {
PromptFormat::Teacher => TeacherPrompt::format(example),
match prompt_format {
PromptFormat::Teacher => {
let prompt = TeacherPrompt::format_prompt(example);
example.prompt = Some(ExamplePrompt {
input: prompt,
expected_output: example.expected_patch.clone(), // TODO
format: prompt_format,
});
}
PromptFormat::Zeta2 => {
run_load_project(example, app_state, cx.clone()).await;
let ep_store = cx
.update(|cx| EditPredictionStore::try_global(cx).unwrap())
.unwrap();
@@ -41,30 +54,28 @@ pub async fn run_format_prompt(
)
})
.unwrap();
format_zeta_prompt(&input)
let prompt = format_zeta_prompt(&input);
let expected_output = zeta2_output_for_patch(&input, &example.expected_patch.clone());
example.prompt = Some(ExamplePrompt {
input: prompt,
expected_output,
format: prompt_format,
});
}
};
example.prompt = Some(ExamplePrompt {
input: prompt,
expected_output: example.expected_patch.clone(), // TODO
format: prompt_format,
});
}
pub trait PromptFormatter {
fn format(example: &Example) -> String;
}
pub trait PromptParser {
/// Return unified diff patch of prediction given raw LLM response
fn parse(example: &Example, response: &str) -> String;
}
pub struct TeacherPrompt;
impl PromptFormatter for TeacherPrompt {
fn format(example: &Example) -> String {
impl TeacherPrompt {
const PROMPT: &str = include_str!("teacher.prompt.md");
pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";
pub(crate) const EDITABLE_REGION_END: &str = "<|editable_region_end|>";
/// Truncate edit history to this number of last lines
const MAX_HISTORY_LINES: usize = 128;
pub fn format_prompt(example: &Example) -> String {
let edit_history = Self::format_edit_history(&example.edit_history);
let context = Self::format_context(example);
let editable_region = Self::format_editable_region(example);
@@ -76,15 +87,46 @@ impl PromptFormatter for TeacherPrompt {
prompt
}
}
impl TeacherPrompt {
const PROMPT: &str = include_str!("teacher.prompt.md");
pub(crate) const EDITABLE_REGION_START: &str = "<|editable_region_start|>\n";
pub(crate) const EDITABLE_REGION_END: &str = "<|editable_region_end|>";
pub fn parse(example: &Example, response: &str) -> String {
// Ideally, we should always be able to find cursor position in the retrieved context.
// In reality, sometimes we don't find it for these reasons:
// 1. `example.cursor_position` contains _more_ context than included in the retrieved context
// (can be fixed by getting cursor coordinates at the load_example stage)
// 2. Context retriever just didn't include cursor line.
//
// In that case, fallback to using `cursor_position` as excerpt.
let cursor_file = &example
.buffer
.as_ref()
.expect("`buffer` should be filled in in the context collection step")
.content;
/// Truncate edit history to this number of last lines
const MAX_HISTORY_LINES: usize = 128;
// Extract updated (new) editable region from the model response
let new_editable_region = extract_last_codeblock(response);
// Reconstruct old editable region we sent to the model
let old_editable_region = Self::format_editable_region(example);
let old_editable_region = Self::extract_editable_region(&old_editable_region);
if !cursor_file.contains(&old_editable_region) {
panic!("Something's wrong: editable_region is not found in the cursor file")
}
// Apply editable region to a larger context and compute diff.
// This is needed to get a better context lines around the editable region
let edited_file = cursor_file.replace(&old_editable_region, &new_editable_region);
let diff = language::unified_diff(&cursor_file, &edited_file);
let diff = indoc::formatdoc! {"
--- a/{path}
+++ b/{path}
{diff}",
path = example.cursor_path.to_string_lossy(),
diff = diff,
};
diff
}
fn format_edit_history(edit_history: &str) -> String {
// Strip comments ("garbage lines") from edit history
@@ -157,49 +199,6 @@ impl TeacherPrompt {
}
}
impl PromptParser for TeacherPrompt {
fn parse(example: &Example, response: &str) -> String {
// Ideally, we should always be able to find cursor position in the retrieved context.
// In reality, sometimes we don't find it for these reasons:
// 1. `example.cursor_position` contains _more_ context than included in the retrieved context
// (can be fixed by getting cursor coordinates at the load_example stage)
// 2. Context retriever just didn't include cursor line.
//
// In that case, fallback to using `cursor_position` as excerpt.
let cursor_file = &example
.buffer
.as_ref()
.expect("`buffer` should be filled in in the context collection step")
.content;
// Extract updated (new) editable region from the model response
let new_editable_region = extract_last_codeblock(response);
// Reconstruct old editable region we sent to the model
let old_editable_region = Self::format_editable_region(example);
let old_editable_region = Self::extract_editable_region(&old_editable_region);
if !cursor_file.contains(&old_editable_region) {
panic!("Something's wrong: editable_region is not found in the cursor file")
}
// Apply editable region to a larger context and compute diff.
// This is needed to get a better context lines around the editable region
let edited_file = cursor_file.replace(&old_editable_region, &new_editable_region);
let diff = language::unified_diff(&cursor_file, &edited_file);
let diff = indoc::formatdoc! {"
--- a/{path}
+++ b/{path}
{diff}
",
path = example.cursor_path.to_string_lossy(),
diff = diff,
};
diff
}
}
fn extract_last_codeblock(text: &str) -> String {
let mut last_block = None;
let mut search_start = 0;
@@ -221,7 +220,7 @@ fn extract_last_codeblock(text: &str) -> String {
}
if let Some(end_pos) = text[backtick_end..].find(&closing_backticks) {
let code_block = &text[backtick_end + 1..backtick_end + end_pos - 1];
let code_block = &text[backtick_end + 1..backtick_end + end_pos];
last_block = Some(code_block.to_string());
search_start = backtick_end + end_pos + backtick_count;
} else {
@@ -250,7 +249,7 @@ mod tests {
`````
"};
let last_block = extract_last_codeblock(text);
assert_eq!(last_block, "last block");
assert_eq!(last_block, "last block\n");
}
#[test]
crates/edit_prediction_cli/src/main.rs
+9
View File
@@ -1,4 +1,5 @@
mod anthropic_client;
mod distill;
mod example;
mod format_prompt;
mod headless;
@@ -16,6 +17,7 @@ use reqwest_client::ReqwestClient;
use serde::{Deserialize, Serialize};
use std::{path::PathBuf, sync::Arc};
use crate::distill::run_distill;
use crate::example::{read_examples, write_examples};
use crate::format_prompt::run_format_prompt;
use crate::load_project::run_load_project;
@@ -54,6 +56,9 @@ enum Command {
Predict(PredictArgs),
/// Computes a score based on actual and expected patches
Score(PredictArgs),
/// Prepares a distillation dataset by copying expected outputs to
/// predicted outputs and removing actual outputs and prompts.
Distill,
/// Print aggregated scores
Eval(PredictArgs),
/// Remove git repositories and worktrees
@@ -87,6 +92,7 @@ enum PredictionProvider {
Zeta1,
Zeta2,
Teacher,
TeacherNonBatching,
}
impl EpArgs {
@@ -175,6 +181,9 @@ fn main() {
)
.await;
}
Command::Distill => {
run_distill(example).await;
}
Command::Score(args) | Command::Eval(args) => {
run_scoring(example, &args, app_state, cx).await;
}
crates/edit_prediction_cli/src/predict.rs
+11 -5
View File
@@ -2,7 +2,7 @@ use crate::{
PredictionProvider, PromptFormat,
anthropic_client::AnthropicClient,
example::{Example, ExamplePrediction},
format_prompt::{PromptParser, TeacherPrompt, run_format_prompt},
format_prompt::{TeacherPrompt, run_format_prompt},
headless::EpAppState,
load_project::run_load_project,
paths::{LATEST_EXAMPLE_RUN_DIR, RUN_DIR},
@@ -30,20 +30,24 @@ pub async fn run_prediction(
return;
}
run_load_project(example, app_state.clone(), cx.clone()).await;
run_context_retrieval(example, app_state.clone(), cx.clone()).await;
let provider = provider.unwrap();
if matches!(provider, PredictionProvider::Teacher) {
if matches!(
provider,
PredictionProvider::Teacher | PredictionProvider::TeacherNonBatching
) {
if example.prompt.is_none() {
run_format_prompt(example, PromptFormat::Teacher, app_state.clone(), cx).await;
}
let batched = true;
let batched = matches!(provider, PredictionProvider::Teacher);
return predict_anthropic(example, repetition_count, batched).await;
}
run_load_project(example, app_state.clone(), cx.clone()).await;
if matches!(
provider,
PredictionProvider::Zeta1 | PredictionProvider::Zeta2
@@ -75,7 +79,9 @@ pub async fn run_prediction(
PredictionProvider::Zeta2 => edit_prediction::EditPredictionModel::Zeta2,
PredictionProvider::Sweep => edit_prediction::EditPredictionModel::Sweep,
PredictionProvider::Mercury => edit_prediction::EditPredictionModel::Mercury,
PredictionProvider::Teacher => unreachable!(),
PredictionProvider::Teacher | PredictionProvider::TeacherNonBatching => {
unreachable!()
}
};
store.set_edit_prediction_model(model);
})
crates/edit_prediction_cli/src/teacher.prompt.md
+1
View File
@@ -18,6 +18,7 @@ Focus on:
Rules:
- Do not just mechanically apply patterns - reason about what changes make sense given the context and the programmer's apparent goals.
- Do not just fix syntax errors - look for the broader refactoring pattern and apply it systematically throughout the code.
- Keep existing formatting unless it's absolutely necessary
Input format:
- You receive small code fragments called context (structs, field definitions, function signatures, etc.). They may or may not be relevant.