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/Cargo.toml
+1 -1
View File
@@ -12,7 +12,7 @@ workspace = true
path = "src/edit_prediction.rs"
[features]
eval-support = []
cli-support = []
[dependencies]
ai_onboarding.workspace = true
crates/edit_prediction/src/edit_prediction.rs
+13 -13
View File
@@ -55,7 +55,7 @@ pub mod open_ai_response;
mod prediction;
pub mod sweep_ai;
#[cfg(any(test, feature = "test-support", feature = "eval-support"))]
#[cfg(any(test, feature = "test-support", feature = "cli-support"))]
pub mod udiff;
mod zed_edit_prediction_delegate;
@@ -158,7 +158,7 @@ pub struct EditPredictionStore {
use_context: bool,
options: ZetaOptions,
update_required: bool,
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
eval_cache: Option<Arc<dyn EvalCache>>,
edit_prediction_model: EditPredictionModel,
pub sweep_ai: SweepAi,
@@ -505,7 +505,7 @@ impl EditPredictionStore {
},
),
update_required: false,
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
eval_cache: None,
edit_prediction_model: EditPredictionModel::Zeta2,
sweep_ai: SweepAi::new(cx),
@@ -554,7 +554,7 @@ impl EditPredictionStore {
.is_some()
}
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
pub fn with_eval_cache(&mut self, cache: Arc<dyn EvalCache>) {
self.eval_cache = Some(cache);
}
@@ -1590,8 +1590,8 @@ impl EditPredictionStore {
client: Arc<Client>,
llm_token: LlmApiToken,
app_version: Version,
#[cfg(feature = "eval-support")] eval_cache: Option<Arc<dyn EvalCache>>,
#[cfg(feature = "eval-support")] eval_cache_kind: EvalCacheEntryKind,
#[cfg(feature = "cli-support")] eval_cache: Option<Arc<dyn EvalCache>>,
#[cfg(feature = "cli-support")] eval_cache_kind: EvalCacheEntryKind,
) -> Result<(open_ai::Response, Option<EditPredictionUsage>)> {
let url = if let Some(predict_edits_url) = PREDICT_EDITS_URL.as_ref() {
http_client::Url::parse(&predict_edits_url)?
@@ -1601,7 +1601,7 @@ impl EditPredictionStore {
.build_zed_llm_url("/predict_edits/raw", &[])?
};
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
let cache_key = if let Some(cache) = eval_cache {
use collections::FxHasher;
use std::hash::{Hash, Hasher};
@@ -1635,7 +1635,7 @@ impl EditPredictionStore {
)
.await?;
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
if let Some((cache, request, key)) = cache_key {
cache.write(key, &request, &serde_json::to_string_pretty(&response)?);
}
@@ -1767,7 +1767,7 @@ impl EditPredictionStore {
}
}
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
pub fn set_context_for_buffer(
&mut self,
project: &Entity<Project>,
@@ -1892,10 +1892,10 @@ pub struct ZedUpdateRequiredError {
minimum_version: Version,
}
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
pub type EvalCacheKey = (EvalCacheEntryKind, u64);
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum EvalCacheEntryKind {
Context,
@@ -1903,7 +1903,7 @@ pub enum EvalCacheEntryKind {
Prediction,
}
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
impl std::fmt::Display for EvalCacheEntryKind {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
@@ -1914,7 +1914,7 @@ impl std::fmt::Display for EvalCacheEntryKind {
}
}
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
pub trait EvalCache: Send + Sync {
fn read(&self, key: EvalCacheKey) -> Option<String>;
fn write(&self, key: EvalCacheKey, input: &str, value: &str);
crates/edit_prediction/src/udiff.rs
+1 -1
View File
@@ -138,7 +138,7 @@ pub fn apply_diff_to_string(diff_str: &str, text: &str) -> Result<String> {
DiffEvent::Hunk { hunk, .. } => {
let hunk_offset = text
.find(&hunk.context)
.ok_or_else(|| anyhow!("couldn't result hunk {:?}", hunk.context))?;
.ok_or_else(|| anyhow!("couldn't resolve hunk {:?}", hunk.context))?;
for edit in hunk.edits.iter().rev() {
let range = (hunk_offset + edit.range.start)..(hunk_offset + edit.range.end);
text.replace_range(range, &edit.text);
crates/edit_prediction/src/zeta2.rs
+16 -4
View File
@@ -1,4 +1,4 @@
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
use crate::EvalCacheEntryKind;
use crate::open_ai_response::text_from_response;
use crate::prediction::EditPredictionResult;
@@ -44,7 +44,7 @@ pub fn request_prediction_with_zeta2(
let llm_token = store.llm_token.clone();
let app_version = AppVersion::global(cx);
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
let eval_cache = store.eval_cache.clone();
let request_task = cx.background_spawn({
@@ -95,9 +95,9 @@ pub fn request_prediction_with_zeta2(
client,
llm_token,
app_version,
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
eval_cache,
#[cfg(feature = "eval-support")]
#[cfg(feature = "cli-support")]
EvalCacheEntryKind::Prediction,
)
.await;
@@ -226,3 +226,15 @@ pub fn zeta2_prompt_input(
};
(editable_offset_range, prompt_input)
}
#[cfg(feature = "cli-support")]
pub fn zeta2_output_for_patch(input: &zeta_prompt::ZetaPromptInput, patch: &str) -> String {
eprintln!("{}", patch);
eprintln!("---------------------");
eprintln!("{}", input.cursor_excerpt);
crate::udiff::apply_diff_to_string(
patch,
&input.cursor_excerpt[input.editable_range_in_excerpt.clone()],
)
.unwrap()
}