dataset.cpp

#include "dataset.hpp"
#include <algorithm>
#include <fstream>
#include <set>

std::string Dataset::extract_field(const std::string& line, const std::string& key) {
    // Search for  "key":  or  "key" :
    std::string search = "\"" + key + "\"";
    auto pos = line.find(search);
    if (pos == std::string::npos) return {};

    pos += search.size();
    // Skip whitespace and colon
    while (pos < line.size() && (line[pos] == ' ' || line[pos] == ':')) ++pos;
    if (pos >= line.size() || line[pos] != '"') return {};
    ++pos; // skip opening quote

    std::string result;
    while (pos < line.size() && line[pos] != '"') {
        if (line[pos] == '\\' && pos + 1 < line.size()) {
            ++pos;
            switch (line[pos]) {
                case 'n':  result += '\n'; break;
                case 't':  result += '\t'; break;
                case '"':  result += '"';  break;
                case '\\': result += '\\'; break;
                default:   result += line[pos]; break;
            }
        } else {
            result += line[pos];
        }
        ++pos;
    }
    return result;
}

int Dataset::parse_file(const std::string& path) {
    std::ifstream f(path);
    if (!f) return 0;
    int count = 0;
    std::string line;
    while (std::getline(f, line)) {
        if (line.empty() || line[0] != '{') continue;
        auto prompt = extract_field(line, "prompt");
        auto answer = extract_field(line, "answer");
        if (!prompt.empty() && !answer.empty()) {
            examples_.push_back({prompt, answer});
            ++count;
        }
    }
    return count;
}

void Dataset::rebuild_labels() {
    std::set<std::string> seen;
    labels_.clear();
    for (const auto& ex : examples_) {
        if (seen.insert(ex.answer).second)
            labels_.push_back(ex.answer);
    }
    std::sort(labels_.begin(), labels_.end());
}

int Dataset::load(const std::string& path) {
    examples_.clear();
    labels_.clear();
    int n = parse_file(path);
    rebuild_labels();
    return n;
}

int Dataset::append(const std::string& path) {
    int n = parse_file(path);
    rebuild_labels();
    return n;
}

std::vector<std::string> Dataset::all_prompts() const {
    std::vector<std::string> out;
    out.reserve(examples_.size());
    for (const auto& ex : examples_) out.push_back(ex.prompt);
    return out;
}

int Dataset::label_index(const std::string& answer) const {
    auto it = std::find(labels_.begin(), labels_.end(), answer);
    return it == labels_.end() ? -1 : static_cast<int>(it - labels_.begin());
}