cpd_semantic/embed/
catalog.rs

1// catalog.rs — the embedding models jscpd has calibrated, with the settings
2// each one gets by default.
3//
4// Models score similarity on different scales: a pair of functions that one
5// model scores 0.9 another scores 0.5, so thresholds tuned for one model do
6// not carry over to another. Each model here got its thresholds on Rosetta
7// Code, where the pairs to find are known: the similarities at which its
8// precision equals that of jina-embeddings-v2-base-code at 0.6 across
9// languages and 0.75 within one. The calibration put the prefix a model's
10// card gives for code before every function, so jscpd does the same.
11//
12// The near-best margin and the group floor of the rules were tuned with
13// jina-embeddings-v2-base-code only. For the other models they are scaled
14// from it by the model's gap between its two thresholds, which stands for
15// how spread out its scores are: the margin is 0.05 per 0.15 of gap, and the
16// floor sits as far above the same-language threshold (0.05 per 0.15).
17
18use super::models::{self, LocalModel};
19use crate::search::Thresholds;
20
21/// An embedding model jscpd has calibrated.
22#[derive(Debug)]
23pub struct KnownModel {
24    /// The name `--semantic-model` takes, and `--semantic-models` shows.
25    pub name: &'static str,
26    /// Hugging Face repository id.
27    pub id: &'static str,
28    /// Names Ollama serves the model under; the first is the one to pull.
29    pub ollama: &'static [&'static str],
30    /// Lowest cosine similarity of a pair across languages.
31    pub threshold: f32,
32    /// Lowest cosine similarity of a pair within one language.
33    pub same_threshold: f32,
34    /// Put before every function text.
35    pub prefix: &'static str,
36    pub license: &'static str,
37    /// The files, when jscpd runs the model itself; the others need an
38    /// embeddings API.
39    pub local: Option<&'static LocalModel>,
40}
41
42/// The default model first, then the other one jscpd runs itself, then the
43/// ones that need an embeddings API.
44pub static KNOWN_MODELS: &[KnownModel] = &[
45    KnownModel {
46        name: "CodeRankEmbed",
47        id: models::CODERANKEMBED.id,
48        ollama: &[],
49        threshold: 0.4125,
50        same_threshold: 0.6375,
51        prefix: "",
52        license: "MIT",
53        local: Some(&models::CODERANKEMBED),
54    },
55    KnownModel {
56        name: "jina-embeddings-v2-base-code",
57        id: models::JINA_V2_BASE_CODE.id,
58        ollama: &["unclemusclez/jina-embeddings-v2-base-code"],
59        threshold: 0.6,
60        same_threshold: 0.75,
61        prefix: "",
62        license: "Apache-2.0",
63        local: Some(&models::JINA_V2_BASE_CODE),
64    },
65    KnownModel {
66        name: "jina-code-embeddings-0.5b",
67        id: "jinaai/jina-code-embeddings-0.5b",
68        ollama: &[],
69        threshold: 0.5625,
70        same_threshold: 0.7125,
71        prefix: "Candidate code snippet:\n",
72        license: "CC-BY-NC-4.0",
73        local: None,
74    },
75    KnownModel {
76        name: "Qwen3-Embedding-0.6B",
77        id: "Qwen/Qwen3-Embedding-0.6B",
78        ollama: &[
79            "qwen3-embedding:0.6b",
80            "qwen3-embedding:0.6b-fp16",
81            "qwen3-embedding:0.6b-q8_0",
82        ],
83        threshold: 0.5875,
84        same_threshold: 0.7625,
85        prefix: "Instruct: Given a code snippet, retrieve code that implements the same functionality\nQuery:",
86        license: "Apache-2.0",
87        local: None,
88    },
89    KnownModel {
90        name: "SFR-Embedding-Code-400M_R",
91        id: "Salesforce/SFR-Embedding-Code-400M_R",
92        ollama: &[],
93        threshold: 0.7375,
94        same_threshold: 0.8375,
95        prefix: "",
96        license: "CC-BY-NC-4.0",
97        local: None,
98    },
99    KnownModel {
100        name: "gte-modernbert-base",
101        id: "Alibaba-NLP/gte-modernbert-base",
102        ollama: &[],
103        threshold: 0.6875,
104        same_threshold: 0.85,
105        prefix: "",
106        license: "Apache-2.0",
107        local: None,
108    },
109    KnownModel {
110        name: "codesage-small-v2",
111        id: "codesage/codesage-small-v2",
112        ollama: &[],
113        threshold: 0.3125,
114        same_threshold: 0.5625,
115        prefix: "",
116        license: "Apache-2.0",
117        local: None,
118    },
119    KnownModel {
120        name: "granite-embedding-english-r2",
121        id: "ibm-granite/granite-embedding-english-r2",
122        ollama: &[],
123        threshold: 0.8625,
124        same_threshold: 0.925,
125        prefix: "",
126        license: "Apache-2.0",
127        local: None,
128    },
129    KnownModel {
130        name: "bge-m3",
131        id: "BAAI/bge-m3",
132        ollama: &["bge-m3", "bge-m3:567m", "bge-m3:567m-fp16"],
133        threshold: 0.7,
134        same_threshold: 0.8375,
135        prefix: "",
136        license: "MIT",
137        local: None,
138    },
139];
140
141impl KnownModel {
142    /// The thresholds of the rules on this model's scale: its calibrated
143    /// two, and the near-best margin and group floor scaled from
144    /// [`Thresholds::REFERENCE`].
145    pub fn thresholds(&self) -> Thresholds {
146        on_scale(self.threshold, self.same_threshold)
147    }
148}
149
150/// The thresholds of a model whose calibrated thresholds are `across` and
151/// `within`.
152fn on_scale(across: f32, within: f32) -> Thresholds {
153    let reference = Thresholds::REFERENCE;
154    let spread = (within - across) / (reference.within - reference.across);
155    Thresholds {
156        across,
157        within,
158        near_best: reference.near_best * spread,
159        group_floor: (within + (reference.group_floor - reference.within) * spread).min(1.0),
160    }
161}
162
163/// The model `name` stands for: its name, its Hugging Face id or one of its
164/// Ollama names, in any letter case, with or without Ollama's `:latest`
165/// tag. A model of another owner under the same name is not this model.
166pub fn find(name: &str) -> Option<&'static KnownModel> {
167    let wanted = key(name);
168    KNOWN_MODELS.iter().find(|m| {
169        std::iter::once(m.name)
170            .chain([m.id])
171            .chain(m.ollama.iter().copied())
172            .any(|known| key(known) == wanted)
173    })
174}
175
176/// A name without `:latest`, in lower case.
177fn key(name: &str) -> String {
178    let name = name.trim();
179    name.strip_suffix(":latest").unwrap_or(name).to_lowercase()
180}
181
182/// The thresholds of the rules for `model`: its calibrated ones, or
183/// [`Thresholds::REFERENCE`] for a model jscpd does not know, with
184/// `across` and `within` put in where they are set. `across` set alone
185/// moves `within` along, keeping the model's gap between the two, at most
186/// 1.
187pub fn thresholds(model: &str, across: Option<f32>, within: Option<f32>) -> Thresholds {
188    let own = find(model).map_or(Thresholds::REFERENCE, KnownModel::thresholds);
189    let across_given = across.unwrap_or(own.across);
190    Thresholds {
191        across: across_given,
192        within: within.unwrap_or(((across_given - own.across) + own.within).min(1.0)),
193        ..own
194    }
195}
196
197/// The names of the models jscpd runs itself, for messages.
198pub fn local_names() -> String {
199    KNOWN_MODELS
200        .iter()
201        .filter(|m| m.local.is_some())
202        .map(|m| m.name)
203        .collect::<Vec<_>>()
204        .join(", ")
205}
206
207#[cfg(test)]
208mod tests {
209    use super::*;
210
211    #[test]
212    fn a_model_is_found_by_name_id_or_ollama_name() {
213        for name in [
214            "CodeRankEmbed",
215            " coderankembed ",
216            "nomic-ai/CodeRankEmbed",
217            "NOMIC-AI/coderankembed",
218        ] {
219            assert_eq!(find(name).unwrap().id, "nomic-ai/CodeRankEmbed", "{name}");
220        }
221        let jina = find("unclemusclez/jina-embeddings-v2-base-code:latest").unwrap();
222        assert_eq!(jina.name, "jina-embeddings-v2-base-code");
223        assert_eq!(
224            find("qwen3-embedding:0.6b").unwrap().name,
225            "Qwen3-Embedding-0.6B"
226        );
227        assert_eq!(find("bge-m3:latest").unwrap().id, "BAAI/bge-m3");
228        assert_eq!(
229            find("jina-code-embeddings-0.5b").unwrap().id,
230            "jinaai/jina-code-embeddings-0.5b"
231        );
232        // Another owner, size or quantization is another model.
233        assert!(find("myorg/CodeRankEmbed").is_none());
234        assert!(find("someone/bge-m3").is_none());
235        assert!(find("qwen3-embedding:8b").is_none());
236        assert!(find("nomic-embed-text").is_none());
237        assert!(find("").is_none());
238    }
239
240    #[test]
241    fn thresholds_follow_the_model_and_keep_its_gap() {
242        let model = "CodeRankEmbed";
243        let own = thresholds(model, None, None);
244        assert_eq!((own.across, own.within), (0.4125, 0.6375));
245        assert_eq!(thresholds(model, Some(0.4125), None), own);
246        let moved = thresholds(model, Some(0.5), None);
247        assert!((moved.within - 0.725).abs() < 1e-6, "{moved:?}");
248        assert_eq!(
249            thresholds(model, Some(0.9), None).within,
250            1.0,
251            "never above identical"
252        );
253        let both = thresholds(model, Some(0.5), Some(0.55));
254        assert_eq!((both.across, both.within), (0.5, 0.55));
255        assert_eq!(
256            thresholds("some/unknown-model", None, None),
257            Thresholds::REFERENCE
258        );
259    }
260
261    #[test]
262    fn the_group_rules_scale_with_the_model() {
263        let close = |a: f32, b: f32| (a - b).abs() < 1e-6;
264        let jina = find("jina-embeddings-v2-base-code").unwrap().thresholds();
265        assert!(
266            close(jina.near_best, 0.05) && close(jina.group_floor, 0.8),
267            "{jina:?}"
268        );
269        // A gap of 0.225, one and a half times the reference's 0.15.
270        let coderank = find("CodeRankEmbed").unwrap().thresholds();
271        assert!(close(coderank.near_best, 0.075), "{coderank:?}");
272        assert!(close(coderank.group_floor, 0.7125), "{coderank:?}");
273        for m in KNOWN_MODELS {
274            let t = m.thresholds();
275            assert!(
276                t.within < t.group_floor && t.group_floor <= 1.0,
277                "{}: {t:?}",
278                m.name
279            );
280            // At the model's own threshold, the other one is its own too.
281            assert_eq!(
282                thresholds(m.name, Some(m.threshold), None).within,
283                m.same_threshold
284            );
285        }
286    }
287
288    #[test]
289    fn every_model_is_sound() {
290        let mut taken = std::collections::HashSet::new();
291        for m in KNOWN_MODELS {
292            assert!(0.0 < m.threshold && m.threshold < m.same_threshold && m.same_threshold <= 1.0);
293            let names: std::collections::HashSet<String> = [m.name, m.id]
294                .into_iter()
295                .chain(m.ollama.iter().copied())
296                .map(key)
297                .collect();
298            for name in names {
299                assert!(taken.insert(name.clone()), "{name} names two models");
300            }
301            assert_eq!(
302                m.id.rsplit('/').next(),
303                Some(m.name),
304                "the name is the repository's"
305            );
306            if let Some(local) = m.local {
307                assert_eq!(local.id, m.id);
308            }
309        }
310        assert_eq!(KNOWN_MODELS[0].name, super::super::DEFAULT_LOCAL_MODEL);
311        assert_eq!(local_names(), "CodeRankEmbed, jina-embeddings-v2-base-code");
312    }
313}
cpd_semantic/embed/catalog.rs

cpd_semantic/embed/
catalog.rs
