spring-ai/semantic-cache/src/main/java/com/ankurm/semanticcache/Dataset.java

package com.ankurm.semanticcache;

import java.io.BufferedReader;
import java.io.InputStreamReader;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

/** The labelled questions in {@code dataset.tsv}: 30 support intents, four phrasings each, plus 20 off-topic questions. */
public final class Dataset {

    public record Row(String intent, String role, String text) { }

    public final List<Row> rows = new ArrayList<>();

    public Dataset() {
        try (var in = Dataset.class.getResourceAsStream("/dataset.tsv");
             var r = new BufferedReader(new InputStreamReader(in, StandardCharsets.UTF_8))) {
            String line;
            while ((line = r.readLine()) != null) {
                if (line.startsWith("#") || line.isBlank()) continue;
                String[] p = line.split("\t");
                // every off-topic question is its own intent: an answer cached for one is wrong for the next
                String intent = p[0].equals("ood") ? "ood-" + rows.stream().filter(x -> x.role().equals("ood")).count() : p[0];
                rows.add(new Row(intent, p[1], p[2]));
            }
        } catch (Exception e) {
            throw new IllegalStateException(e);
        }
    }

    public List<Row> seeds() { return rows.stream().filter(r -> r.role().equals("seed")).toList(); }

    public List<Row> paraphrases() { return rows.stream().filter(r -> r.role().equals("para")).toList(); }

    public List<Row> offTopic() { return rows.stream().filter(r -> r.role().equals("ood")).toList(); }

    /** text -> intent, so the scripted model knows which answer a question deserves. */
    public Map<String, String> intentByText() {
        Map<String, String> m = new LinkedHashMap<>();
        rows.forEach(r -> m.put(r.text(), r.intent()));
        return m;
    }
}