data-sampler

"use strict";
/**
 * @file dataset.ts
 * @description HuggingFace Datasets Server REST API client.
 *
 * Uses the HF Datasets Server (datasets-server.huggingface.co) to discover
 * dataset configs/splits and fetch sample rows. Uses only Node.js built-in
 * fetch() — no additional npm dependencies.
 *
 * Endpoints used:
 *   - GET /splits  → discover available configs and splits
 *   - GET /rows    → fetch consecutive rows by offset/length
 *   - GET /search  → BM25 full-text search across rows
 *
 * @module dataset
 */
Object.defineProperty(exports, "__esModule", { value: true });
exports.fetchSplits = fetchSplits;
exports.fetchRows = fetchRows;
exports.searchRows = searchRows;
exports.loadDataset = loadDataset;
// ---------------------------------------------------------------------------
// Constants
// ---------------------------------------------------------------------------
/** Base URL for the HuggingFace Datasets Server API. */
const HF_DATASETS_SERVER_BASE = "https://datasets-server.huggingface.co";
/** Soft ceiling: how many rows we request in a single call. */
const MAX_ROWS_PER_REQUEST = 100;
/** Minimum allowed max_samples. */
const MIN_SAMPLES = 1;
// ---------------------------------------------------------------------------
// HTTP helpers
// ---------------------------------------------------------------------------
/**
 * Perform a fetch with error classification.
 *
 * @param url - Fully qualified URL to fetch.
 * @param signal - Optional AbortSignal for cancellation.
 * @returns The parsed JSON response body.
 */
async function apiFetch(url, signal) {
    let response;
    try {
        response = await fetch(url, { signal });
    }
    catch (err) {
        if (err instanceof DOMException && err.name === "AbortError") {
            throw err; // let callers handle cancellation
        }
        throw new Error(`Network error while contacting HuggingFace Datasets Server: ${err instanceof Error ? err.message : String(err)}`);
    }
    if (response.status === 404) {
        throw new Error(`Resource not found on HuggingFace Datasets Server (HTTP 404). ` +
            `Check that the dataset name and config are correct.`);
    }
    if (response.status === 429) {
        throw new Error("Rate limited by HuggingFace Datasets Server (HTTP 429). " +
            "Please wait a moment before trying again.");
    }
    if (!response.ok) {
        const body = await response.text().catch(() => "");
        throw new Error(`HuggingFace Datasets Server returned HTTP ${response.status}` +
            (body ? `: ${body.slice(0, 500)}` : ""));
    }
    return response.json();
}
// ---------------------------------------------------------------------------
// Public API functions
// ---------------------------------------------------------------------------
/**
 * Fetch available configs and splits for a given dataset.
 *
 * @param dataset - HF dataset name (e.g. "squad", "ibm/duorc").
 * @param signal - Optional AbortSignal.
 * @returns The parsed splits response.
 */
async function fetchSplits(dataset, signal) {
    const url = `${HF_DATASETS_SERVER_BASE}/splits?dataset=${encodeURIComponent(dataset)}`;
    return apiFetch(url, signal);
}
/**
 * Fetch a consecutive block of rows from a dataset split.
 *
 * @param dataset - HF dataset name.
 * @param config - Dataset config/subset name.
 * @param split - Split name ("train", "test", "validation").
 * @param offset - Row offset to start from (0-based).
 * @param length - Number of rows to fetch (max 100).
 * @param signal - Optional AbortSignal.
 * @returns An object with the rows array and total row count.
 */
async function fetchRows(dataset, config, split, offset, length, signal) {
    const clamped = Math.min(Math.max(1, length), MAX_ROWS_PER_REQUEST);
    const url = `${HF_DATASETS_SERVER_BASE}/rows` +
        `?dataset=${encodeURIComponent(dataset)}` +
        `&config=${encodeURIComponent(config)}` +
        `&split=${encodeURIComponent(split)}` +
        `&offset=${Math.max(0, offset)}` +
        `&length=${clamped}`;
    const data = await apiFetch(url, signal);
    return {
        rows: data.rows,
        totalRows: data.num_rows_total ?? data.rows.length,
    };
}
/**
 * BM25 full-text search across rows in a dataset split.
 *
 * @param dataset - HF dataset name.
 * @param config - Dataset config/subset name.
 * @param split - Split name.
 * @param query - Search query text.
 * @param length - Number of results to return (max 100).
 * @param signal - Optional AbortSignal.
 * @returns An object with matching rows and an estimate of total matches.
 */
async function searchRows(dataset, config, split, query, length, signal) {
    const clamped = Math.min(Math.max(1, length), MAX_ROWS_PER_REQUEST);
    const url = `${HF_DATASETS_SERVER_BASE}/search` +
        `?dataset=${encodeURIComponent(dataset)}` +
        `&config=${encodeURIComponent(config)}` +
        `&split=${encodeURIComponent(split)}` +
        `&query=${encodeURIComponent(query)}` +
        `&offset=0` +
        `&length=${clamped}`;
    const data = await apiFetch(url, signal);
    return {
        rows: data.rows,
        totalRows: data.total ?? data.rows.length,
    };
}
// ---------------------------------------------------------------------------
// Formatting helpers
// ---------------------------------------------------------------------------
/**
 * Format a single sample's data into a text representation.
 *
 * @param sample - The sample with index and data.
 * @param fmt - The requested format style.
 * @param template - Optional template with {text} placeholder.
 * @returns A formatted text string for this one sample.
 */
function formatSingleSample(sample, fmt, template, inputColumns, outputColumn) {
    // Build the key-value lines from the sample data.
    const dataLines = Object.entries(sample.data)
        .filter(([, value]) => value != null)
        .map(([key, value]) => {
        const display = typeof value === "object" && value !== null
            ? JSON.stringify(value)
            : String(value);
        return `${key}: ${display}`;
    })
        .join("\n");
    switch (fmt) {
        case "json": {
            return JSON.stringify(sample.data, null, 2);
        }
        case "prompt": {
            return `=== Sample ${sample.index} ===\n${dataLines}\n---`;
        }
        case "fewshot": {
            const inputCols = inputColumns ??
                Object.keys(sample.data).filter((k) => k !== outputColumn);
            const parts = [];
            for (const col of inputCols) {
                const value = sample.data[col];
                if (value == null)
                    continue;
                const display = typeof value === "object" && value !== null
                    ? JSON.stringify(value)
                    : String(value);
                const label = col.charAt(0).toUpperCase() + col.slice(1);
                parts.push(`${label}: ${display}`);
            }
            if (outputColumn != null && sample.data[outputColumn] != null) {
                const value = sample.data[outputColumn];
                const display = typeof value === "object" && value !== null
                    ? JSON.stringify(value)
                    : String(value);
                const label = outputColumn.charAt(0).toUpperCase() + outputColumn.slice(1);
                parts.push(`${label}: ${display}`);
            }
            return parts.join("\n");
        }
        case "text":
        default: {
            if (template != null && template.length > 0) {
                return template.replace("{text}", dataLines);
            }
            return dataLines;
        }
    }
}
// ---------------------------------------------------------------------------
// Orchestration
// ---------------------------------------------------------------------------
/**
 * Main entry point: load samples from a HuggingFace dataset.
 *
 * Orchestration flow:
 * 1. If no config was provided, call `/splits` to discover available configs.
 *    - Exactly one config → use it automatically.
 *    - Multiple configs → throw with a message listing choices.
 *    - No configs → throw with available splits info.
 * 2. If a query is provided, call `/search`; otherwise call `/rows`.
 * 3. Format the returned rows into the standard result structure.
 *
 * @param params - LoadDatasetParams from the tool invocation.
 * @param signal - Optional AbortSignal from LM Studio for cancellation.
 * @returns A LoadDatasetResult with samples, metadata, and formatted text.
 */
async function loadDataset(params, signal) {
    const dataset = params.dataset;
    const split = params.split ?? "train";
    const maxSamples = Math.max(MIN_SAMPLES, Math.min(MAX_ROWS_PER_REQUEST, params.max_samples ?? 5));
    const format = params.format ?? "text";
    const template = params.template;
    // ---- Step 1: resolve config ----
    // The HuggingFace Datasets Server uses "config" as the internal subset name.
    // Datasets without subsets still expose a config (often "default").
    let config = params.config;
    if (config == null || config.length === 0) {
        const splitsData = await fetchSplits(dataset, signal);
        // Filter splits that match the requested split name.
        const matching = splitsData.splits.filter((s) => s.split === split);
        if (matching.length === 0) {
            const available = [
                ...new Set(splitsData.splits.map((s) => `${s.config}/${s.split}`)),
            ];
            throw new Error(`No config found for dataset "${dataset}" with split "${split}". ` +
                `Available splits: ${available.join(", ") || "none found"}.`);
        }
        const uniqueConfigs = [...new Set(matching.map((s) => s.config))];
        if (uniqueConfigs.length > 1) {
            throw new Error(`Dataset "${dataset}" has multiple configs for split "${split}": ` +
                `${uniqueConfigs.join(", ")}. ` +
                `Please specify the "config" parameter to select one.`);
        }
        config = uniqueConfigs[0];
    }
    // ---- Step 2: fetch rows ----
    let rows;
    let totalRows;
    if (params.indices != null && params.indices.trim().length > 0) {
        // Parse comma/space-separated indices string
        const parsed = params.indices
            .split(/[, ]+/)
            .map((s) => s.trim())
            .filter((s) => s.length > 0)
            .map((s) => parseInt(s, 10))
            .filter((n) => !isNaN(n) && n >= 0);
        if (parsed.length === 0) {
            throw new Error(`Invalid indices string: "${params.indices}". Expected comma/space-separated numbers.`);
        }
        // Sort and deduplicate
        const uniqueSorted = [...new Set(parsed)].sort((a, b) => a - b);
        const minIdx = uniqueSorted[0];
        const maxIdx = uniqueSorted[uniqueSorted.length - 1];
        const rangeLength = maxIdx - minIdx + 1;
        if (rangeLength > MAX_ROWS_PER_REQUEST) {
            throw new Error(`Index range [${minIdx}..${maxIdx}] spans ${rangeLength} rows, ` +
                `which exceeds the maximum of ${MAX_ROWS_PER_REQUEST} rows per request.`);
        }
        const rowsResult = await fetchRows(dataset, config, split, minIdx, rangeLength, signal);
        totalRows = rowsResult.totalRows;
        // Build a map for O(1) lookup, then select only requested indices
        const rowsMap = new Map(rowsResult.rows.map((r) => [r.row_idx, r]));
        rows = uniqueSorted
            .map((idx) => rowsMap.get(idx))
            .filter((r) => r != null);
    }
    else if (params.query != null && params.query.length > 0) {
        const searchResult = await searchRows(dataset, config, split, params.query, maxSamples, signal);
        rows = searchResult.rows;
        totalRows = searchResult.totalRows;
    }
    else {
        const rowsResult = await fetchRows(dataset, config, split, 0, maxSamples, signal);
        rows = rowsResult.rows;
        totalRows = rowsResult.totalRows;
    }
    // ---- Step 3: build samples ----
    const samples = rows.map((r) => ({
        index: r.row_idx,
        data: Object.freeze({ ...r.row }),
    }));
    // ---- Step 4: build formatted text ----
    let outCol = params.output_column;
    if (format === "fewshot" && outCol == null && samples.length > 0) {
        const keys = Object.keys(samples[0].data);
        outCol = keys[keys.length - 1];
    }
    const formatted = samples
        .map((s) => formatSingleSample(s, format, template, params.input_columns, outCol))
        .join("\n\n");
    return {
        dataset,
        split,
        config,
        samples,
        totalRows,
        count: samples.length,
        format,
        formatted,
    };
}
//# sourceMappingURL=dataset.js.map