yams/tests/benchmarks/beir_loader.h

#pragma once

#include <nlohmann/json.hpp>
#include <spdlog/spdlog.h>
#include <filesystem>
#include <fstream>
#include <map>
#include <sstream>
#include <string>
#include <vector>
#include <yams/core/types.h>

namespace yams::bench {

struct BEIRDocument {
    std::string id;
    std::string title;
    std::string text;
    /// Per-document metadata from the corpus line's `metadata` object (scalars stringified).
    /// Applied to the ingested document after ingest so extractors see keys such as
    /// `event_time`, `session_id` and `round_index`.
    std::map<std::string, std::string> metadata;
};

/// Copy the `metadata` object of a BEIR corpus line into `doc.metadata`. Strings are kept as
/// is, numbers and booleans are stringified, other value kinds are dropped. A missing or
/// non-object `metadata` leaves the map empty.
inline void readBEIRDocumentMetadata(const nlohmann::json& line, BEIRDocument& doc) {
    const auto it = line.find("metadata");
    if (it == line.end() || !it->is_object()) {
        return;
    }
    for (const auto& [key, value] : it->items()) {
        if (value.is_string()) {
            doc.metadata[key] = value.get<std::string>();
        } else if (value.is_number() || value.is_boolean()) {
            doc.metadata[key] = value.dump();
        }
    }
}

struct BEIRQuery {
    std::string id;
    std::string text;
};

struct BEIRQrel {
    std::string docId;
    int score = 0;
};

struct BenchmarkKgEntity {
    std::string docId;
    std::string entityKey;
    std::string label;
    std::string type = "benchmark_entity";
    std::string relation = "mentioned_in";
    std::string extractor = "benchmark_manifest";
    float weight = 1.0F;
};

struct BEIRDataset {
    std::string name;
    std::map<std::string, BEIRDocument> documents;
    std::map<std::string, BEIRQuery> queries;
    std::multimap<std::string, BEIRQrel> qrels;
    std::vector<BenchmarkKgEntity> kgEntities;
    std::filesystem::path basePath;
};

inline yams::Result<BEIRDataset> loadBEIRDataset(const std::string& datasetName,
                                                 const std::filesystem::path& cacheDir) {
    namespace fs = std::filesystem;
    using nlohmann::json;

    static const std::map<std::string, std::string> DATASET_ZIP_NAMES = {
        {"touche-2020", "webis-touche2020"},
        {"dbpedia-entity", "dbpedia-entity"},
    };

    std::string zipName = datasetName;
    auto zipIt = DATASET_ZIP_NAMES.find(datasetName);
    if (zipIt != DATASET_ZIP_NAMES.end()) {
        zipName = zipIt->second;
    }

    fs::path datasetDir = cacheDir;
    if (datasetDir.empty()) {
        const char* home = std::getenv("HOME");
        if (!home) {
            return yams::Error{yams::ErrorCode::InvalidArgument, "HOME not set and no cacheDir"};
        }
        datasetDir = fs::path(home) / ".cache" / "yams" / "benchmarks" / datasetName;
    }

    BEIRDataset dataset;
    dataset.name = datasetName;
    dataset.basePath = datasetDir;

    fs::path corpusFile = datasetDir / "corpus.jsonl";
    fs::path queriesFile = datasetDir / "queries.jsonl";
    fs::path qrelsFile = datasetDir / "qrels" / "test.tsv";

    if (!fs::exists(corpusFile) || !fs::exists(queriesFile) || !fs::exists(qrelsFile)) {
        if (datasetName == "longmemeval_s") {
            return yams::Error{yams::ErrorCode::NotFound,
                               "LongMemEval_S dataset not found. Convert from HuggingFace:\n"
                               "  uv pip install datasets\n"
                               "  python scripts/prepare_longmemeval_s.py\n"
                               "Expected location: ~/.cache/yams/benchmarks/longmemeval_s/"};
        }
        std::ostringstream oss;
        oss << datasetName << " dataset not found. Please download manually:\n"
            << "  mkdir -p ~/.cache/yams/benchmarks && cd ~/.cache/yams/benchmarks\n"
            << "  curl -L -o " << zipName << ".zip "
            << "https://public.ukp.informatik.tu-darmstadt.de/thakur/BEIR/datasets/" << zipName
            << ".zip\n"
            << "  unzip " << zipName << ".zip";
        if (datasetName != zipName) {
            oss << " && mv " << zipName << " " << datasetName;
        }
        return yams::Error{yams::ErrorCode::NotFound, oss.str()};
    }

    std::ifstream corpusIn(corpusFile);
    std::string line;
    while (std::getline(corpusIn, line)) {
        if (line.empty())
            continue;
        try {
            auto j = json::parse(line);
            BEIRDocument doc;
            doc.id = j.value("_id", "");
            if (doc.id.empty())
                doc.id = j.value("id", "");
            doc.title = j.value("title", "");
            doc.text = j.value("text", "");
            readBEIRDocumentMetadata(j, doc);
            if (!doc.id.empty() && !doc.text.empty()) {
                dataset.documents[doc.id] = doc;
            }
        } catch (const json::exception& e) {
            spdlog::warn("BEIR loader: failed to parse corpus line: {}", e.what());
        }
    }
    corpusIn.close();

    std::ifstream queriesIn(queriesFile);
    while (std::getline(queriesIn, line)) {
        if (line.empty())
            continue;
        try {
            auto j = json::parse(line);
            BEIRQuery q;
            q.id = j.value("_id", "");
            if (q.id.empty())
                q.id = j.value("id", "");
            q.text = j.value("text", "");
            if (!q.id.empty() && !q.text.empty()) {
                dataset.queries[q.id] = q;
            }
        } catch (const json::exception& e) {
            spdlog::warn("BEIR loader: failed to parse query line: {}", e.what());
        }
    }
    queriesIn.close();

    std::ifstream qrelsIn(qrelsFile);
    bool firstLine = true;
    while (std::getline(qrelsIn, line)) {
        if (line.empty() || line[0] == '#')
            continue;
        if (firstLine) {
            firstLine = false;
            if (line.find("query-id") != std::string::npos)
                continue;
        }
        std::istringstream iss(line);
        std::string queryId, docId, scoreStr;
        if (std::getline(iss, queryId, '\t') && std::getline(iss, docId, '\t') &&
            std::getline(iss, scoreStr, '\t')) {
            try {
                int score = std::stoi(scoreStr);
                dataset.qrels.emplace(queryId, BEIRQrel{.docId = docId, .score = score});
            } catch (const std::exception& e) {
                spdlog::warn("BEIR loader: bad qrel score '{}': {}", scoreStr, e.what());
            }
        }
    }
    qrelsIn.close();

    spdlog::info("BEIR loader: {} -> {} docs, {} queries, {} qrels", datasetName,
                 dataset.documents.size(), dataset.queries.size(), dataset.qrels.size());

    return dataset;
}

} // namespace yams::bench