xtask/
suite.rs

1use anyhow::{Context, Result, ensure};
2use serde::{Deserialize, Serialize};
3use serde_json::Value;
4use std::collections::HashSet;
5
6#[derive(Clone, Debug, Serialize, Deserialize)]
7pub struct Configuration {
8    pub id: String,
9    pub label: String,
10    pub args: Vec<String>,
11    pub store_bits: Option<u8>,
12    pub reproducible_cut: bool,
13}
14
15#[derive(Clone, Debug, Serialize, Deserialize)]
16pub struct Case {
17    pub id: String,
18    pub kind: String,
19    #[serde(default, skip_serializing_if = "Option::is_none")]
20    pub prompt: Option<String>,
21    #[serde(default, skip_serializing_if = "Option::is_none")]
22    pub repeat_text: Option<String>,
23    #[serde(default, skip_serializing_if = "Option::is_none")]
24    pub repeat: Option<usize>,
25    #[serde(default, skip_serializing_if = "Option::is_none")]
26    pub suffix: Option<String>,
27    pub max_tokens: usize,
28    pub stop: String,
29    #[serde(default, skip_serializing_if = "Option::is_none")]
30    pub max_ctx: Option<usize>,
31    #[serde(default, skip_serializing_if = "Option::is_none")]
32    pub rounds: Option<usize>,
33}
34
35impl Case {
36    pub fn prompt(&self) -> String {
37        if let Some(prompt) = &self.prompt {
38            return prompt.clone();
39        }
40
41        self.repeat_text
42            .as_deref()
43            .unwrap_or("")
44            .repeat(self.repeat.unwrap_or(0))
45            + self.suffix.as_deref().unwrap_or("")
46    }
47}
48
49#[derive(Deserialize)]
50pub struct Suite {
51    pub rounds: usize,
52    pub max_ctx: usize,
53    pub configurations: Vec<Configuration>,
54    pub cases: Vec<Case>,
55}
56
57pub fn select<T: Clone>(
58    items: &[T],
59    selection: Option<&str>,
60    id: impl Fn(&T) -> &str,
61) -> Result<Vec<T>> {
62    let Some(selection) = selection else {
63        return Ok(items.to_vec());
64    };
65    let mut seen = HashSet::new();
66
67    selection
68        .split(',')
69        .map(|name| {
70            ensure!(seen.insert(name), "repeated selection: {name}");
71
72            items
73                .iter()
74                .find(|v| id(v) == name)
75                .cloned()
76                .with_context(|| format!("unknown selection: {name}"))
77        })
78        .collect()
79}
80
81pub fn schedule(
82    configs: &[Configuration],
83    cases: &[Case],
84    rounds: usize,
85) -> Vec<(usize, usize, usize)> {
86    let mut jobs = Vec::new();
87
88    for round in 0..rounds {
89        for (case_index, case) in cases.iter().enumerate() {
90            if case.kind == "svg" || round >= case.rounds.unwrap_or(rounds) {
91                continue;
92            }
93
94            for index in 0..configs.len() {
95                jobs.push((round, (index + round) % configs.len(), case_index));
96            }
97        }
98    }
99
100    for (index, case) in cases.iter().enumerate() {
101        if case.kind != "svg" {
102            continue;
103        }
104
105        jobs.extend((0..configs.len()).map(|c| (0, c, index)));
106    }
107
108    jobs
109}
110
111pub fn only_higher_caps(old: &Value, new: &Value) -> bool {
112    let (mut old_other, mut new_other) = (old.clone(), new.clone());
113    let (Some(a), Some(b)) = (old_other.as_object_mut(), new_other.as_object_mut()) else {
114        return false;
115    };
116    let (Some(a), Some(b)) = (a.remove("cases"), b.remove("cases")) else {
117        return false;
118    };
119
120    if old_other != new_other {
121        return false;
122    }
123
124    let (Some(a), Some(b)) = (a.as_array(), b.as_array()) else {
125        return false;
126    };
127
128    if a.len() != b.len() {
129        return false;
130    }
131
132    let mut increased = false;
133
134    for (a, b) in a.iter().zip(b) {
135        let (mut a, mut b) = (a.clone(), b.clone());
136        let (Some(a), Some(b)) = (a.as_object_mut(), b.as_object_mut()) else {
137            return false;
138        };
139        let (Some(x), Some(y)) = (
140            a.remove("max_tokens").and_then(|v| v.as_u64()),
141            b.remove("max_tokens").and_then(|v| v.as_u64()),
142        ) else {
143            return false;
144        };
145
146        if a != b || y < x {
147            return false;
148        }
149
150        increased |= y > x;
151    }
152
153    increased
154}
xtask/suite.rs

xtask/
suite.rs
