opentslm 0.1.0

Rust implementation of OpenTSLM using Burn, WGPU, and llama.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129

//! TSQA dataset loader — mirrors `TSQADataset` / `QADataset` in Python.
//!
//! **Source**: HuggingFace [`ChengsenWang/TSQA`](https://huggingface.co/datasets/ChengsenWang/TSQA),
//! served as JSONL files downloaded by the [`downloader`](crate::data::downloader).
//!
//! In the Rust pipeline the WISDM-W accelerometer dataset (which the
//! downloader uses) is converted into TSQA-format MCQ rows: the model must
//! identify which of four activity labels matches the x-axis accelerometer
//! signal shown in a 4-second window.
//!
//! # JSONL row schema
//!
//! ```json
//! { "Question": "...", "Answer": "...", "Task": "activity classification",
//!   "Series": "[0.1, -0.2, ...]" }
//! ```
//!
//! # Splits
//!
//! 80 / 10 / 10 (train / val / test) with seed 42, matching the Python splits.
//!
//! # Curriculum stage
//!
//! Stage 1 — MCQ on time series.

use std::path::Path;

use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};

use crate::data::batch::{normalize, Sample};

// ── Row schema ────────────────────────────────────────────────────────────────

/// Deserialised row from the TSQA JSONL file.
#[derive(Debug, Deserialize)]
struct TsqaRow {
    #[serde(rename = "Question")]
    question: String,
    #[serde(rename = "Answer")]
    answer: String,
    #[serde(rename = "Task")]
    task: String,
    #[serde(rename = "Series")]
    series: String, // JSON-encoded array
}

// ── Public API ────────────────────────────────────────────────────────────────

/// Train / val / test splits for the TSQA dataset.
pub struct TsqaSplits {
    pub train: Vec<Sample>,
    pub val:   Vec<Sample>,
    pub test:  Vec<Sample>,
}

/// Load the TSQA dataset from a JSONL file downloaded from HuggingFace.
///
/// Expected file: `<data_dir>/tsqa/train.jsonl`
/// (Download with `huggingface-cli download ChengsenWang/TSQA --local-dir <data_dir>/tsqa`)
pub fn load_tsqa_splits(data_dir: &Path) -> Result<TsqaSplits> {
    let file = data_dir.join("tsqa").join("train.jsonl");
    let text = std::fs::read_to_string(&file)
        .with_context(|| format!("Cannot read TSQA dataset at {file:?}"))?;

    let rows: Vec<TsqaRow> = text
        .lines()
        .filter(|l| !l.trim().is_empty())
        .enumerate()
        .map(|(i, line)| {
            serde_json::from_str(line)
                .with_context(|| format!("TSQA line {i}: parse error"))
        })
        .collect::<Result<_>>()?;

    let samples: Vec<Sample> = rows.iter().map(row_to_sample).collect::<Result<Vec<_>>>()?.into_iter().collect();

    // 80 / 10 / 10 split (seeded deterministically)
    split_80_10_10(samples)
}

fn row_to_sample(row: &TsqaRow) -> Result<Sample> {
    // Parse series JSON
    let series: Vec<f32> = serde_json::from_str(&row.series)
        .with_context(|| "Failed to parse TSQA series field")?;

    // Normalise
    let (series_norm, mean, std) = normalize(&series);

    let pre_prompt = row.question.trim().to_string();
    let post_prompt = format!("Predict the {} Answer:", row.task);
    let ts_text = format!(
        "This is the time series, it has mean {mean:.4} and std {std:.4}."
    );

    Ok(Sample {
        pre_prompt,
        time_series_text: vec![ts_text],
        time_series: vec![series_norm],
        post_prompt,
        answer: row.answer.trim().to_string(),
        label: Some(row.answer.trim().to_string()),
    })
}

fn split_80_10_10(samples: Vec<Sample>) -> Result<TsqaSplits> {
    let n = samples.len();
    // Deterministic shuffle with seed 42 (matches Python random_state=42)
    let shuffled = deterministic_shuffle(samples, 42);

    let test_n = (n as f64 * 0.10).round() as usize;
    let val_n = (n as f64 * 0.10).round() as usize;
    let train_n = n - test_n - val_n;

    let mut it = shuffled.into_iter();
    let train: Vec<_> = it.by_ref().take(train_n).collect();
    let val: Vec<_> = it.by_ref().take(val_n).collect();
    let test: Vec<_> = it.collect();

    Ok(TsqaSplits { train, val, test })
}

fn deterministic_shuffle(mut v: Vec<Sample>, seed: u64) -> Vec<Sample> {
    use rand::{seq::SliceRandom, SeedableRng};
    use rand::rngs::StdRng;
    let mut rng = StdRng::seed_from_u64(seed);
    v.shuffle(&mut rng);
    v
}