1 files changed, 104 insertions, 0 deletions
diff --git a/src/data.rs b/src/data.rs
new file mode 100644
index 0000000..2224583
--- /dev/null
+++ b/src/data.rs
@@ -0,0 +1,104 @@
+use anyhow::Result;
+use serde::Deserialize;
+use std::path::Path;
+use tokenizers::Tokenizer;
+
+#[derive(Debug, Deserialize)]
+struct DollyRow {
+    instruction: String,
+    #[serde(default)]
+    context: String,
+    response: String,
+}
+
+/// A tokenized example. `mask[i] = 1.0` marks the positions that count
+/// towards the loss: only the answer, never the prompt.
+#[derive(Debug, Clone)]
+pub struct Example {
+    pub ids: Vec<i64>,
+    pub mask: Vec<f32>,
+}
+
+fn encode(tok: &Tokenizer, text: &str) -> Result<Vec<i64>> {
+    Ok(tok
+        .encode(text, false)
+        .map_err(|e| anyhow::anyhow!("{e}"))?
+        .get_ids()
+        .iter()
+        .map(|&i| i as i64)
+        .collect())
+}
+
+pub fn load_dolly(
+    path: impl AsRef<Path>,
+    tok: &Tokenizer,
+    max_len: usize,
+    limit: usize,
+) -> Result<Vec<Example>> {
+    let raw = std::fs::read_to_string(path)?;
+    let mut out = Vec::new();
+
+    for line in raw.lines() {
+        if out.len() >= limit {
+            break;
+        }
+        let row: DollyRow = match serde_json::from_str(line) {
+            Ok(r) => r,
+            Err(_) => continue,
+        };
+
+        // ChatML, the SmolLM2 chat template
+        let user = if row.context.trim().is_empty() {
+            row.instruction.clone()
+        } else {
+            format!("{}\n\n{}", row.instruction, row.context)
+        };
+        let prompt = format!("<|im_start|>user\n{user}<|im_end|>\n<|im_start|>assistant\n");
+        let answer = format!("{}<|im_end|>", row.response);
+
+        let p = encode(tok, &prompt)?;
+        let a = encode(tok, &answer)?;
+        if p.len() + a.len() > max_len || a.is_empty() {
+            continue;
+        }
+
+        let mut ids = p.clone();
+        ids.extend_from_slice(&a);
+        let mut mask = vec![0.0; p.len()];
+        mask.extend(std::iter::repeat(1.0).take(a.len()));
+
+        out.push(Example { ids, mask });
+    }
+    Ok(out)
+}
+
+/// Pads a batch to its longest example. Padding has mask 0, so it adds
+/// nothing to the loss.
+pub struct Batch {
+    pub ids: Vec<i64>,
+    pub mask: Vec<f32>,
+    pub batch: usize,
+    pub seq: usize,
+}
+
+pub fn collate(examples: &[Example]) -> Batch {
+    let seq = examples.iter().map(|e| e.ids.len()).max().unwrap_or(1);
+    let batch = examples.len();
+    let mut ids = Vec::with_capacity(batch * seq);
+    let mut mask = Vec::with_capacity(batch * seq);
+
+    for e in examples {
+        ids.extend_from_slice(&e.ids);
+        mask.extend_from_slice(&e.mask);
+        for _ in e.ids.len()..seq {
+            ids.push(0);
+            mask.push(0.0);
+        }
+    }
+    Batch {
+        ids,
+        mask,
+        batch,
+        seq,
+    }
+}