social-data-attribution-w2/docs/TRACKSTAR_DATA_ARTIFACTS.md

SOC-156 Data Artifacts Guide

Run ID: 20260326T163642Z_1102443 Corpus: 5,678,621 docs (10K docs/bin stratified sample from HCAI-Lab/dolma3-6t-sample-10000-docs) Model: allenai/Olmo-3-1025-7B (OLMo3 7B Base) Preconditioner: Mixed (SOC-152), coefficient 0.0877, from 100K random sample Parameters: fp32, projection_dim=16, token_batch_size=4096, truncation=true, unit_normalize=true

Cloud-only (2026-05-23 update): every artifact below is on HuggingFace. The HF paths in the Cloud column are the active reference. The PACE paths below are kept as historical record of where each artifact was originally built — they are no longer the access surface and may not exist on disk after PACE cleanup. Always start from the HF path. The migration preserved file content bit-exact — see PRESERVATION_GAP_REPORT.md for verification details.

Artifact Summary

Artifact	Cloud (HF, primary)	Path on PACE (historical)	Size	Description
Full score matrices	HCAI-Lab/trackstar-scores-base-olmes-4bench (bucket, public)	/storage/.../trackstar/scores_full/base/20260326T163642Z_1102443/	~396 GB	All docs × all queries, per-shard numpy
Top-K results	HCAI-Lab/dolma3-trackstar-influence-scores (dataset, private; top2k_*.parquet + influence_scores_full.parquet)	~/dev/data-attribution-soc156/runs/trackstar/results/base/20260326T163642Z_1102443/ (top-100 JSONL)	155 MB	Ranked influence per benchmark
Gradient index	HCAI-Lab/trackstar-gradient-index-base (bucket, private)	/storage/.../trackstar/builds/base/20260326T163642Z_1102443/sample_10000_docs/	1.3 TB	Reusable training gradients (316 shards)
Query indices	HCAI-Lab/trackstar-query-gradients-base (bucket, private; base/20260326T163642Z_1102443/queries_*/)	/storage/.../trackstar/query_builds/base/20260326T163642Z_1102443/	17 GB	Per-benchmark query gradients
Training shards	HCAI-Lab/dolma3-6t-sample-10000-docs-trackstar-shards (dataset, private)	/storage/.../trackstar/shards_10k/sample_10000_docs/	41 GB	JSONL for resolving positional doc IDs (shard_NNNN:INDEX) → source text
Query JSONL	HCAI-Lab/base-query-data (dataset, public)	~/dev/data-attribution-soc156/queries/base/	~50 MB	Query prompts + metadata
Preconditioner	HCAI-Lab/trackstar-preconditioners (bucket, public; olmo-3-1025-7b/)	~/scratch/soc152/output_production/mixed_preconditioner/	282 MB	TrackStar mixed preconditioner

1. Full Score Matrices (per-shard numpy)

Cloud (primary): HCAI-Lab/trackstar-scores-base-olmes-4bench HF Bucket (public). Layout: <benchmark_dir>/shard_NNNN.npy + <benchmark_dir>/shard_NNNN_doc_ids.json + <benchmark_dir>/query_ids.json. Use hf buckets cp / hf buckets sync to fetch.

Path on PACE (historical): /storage/ice-shared/cs7634/staff/TDA/trackstar/scores_full/base/20260326T163642Z_1102443/

Benchmark	Directory	Queries	Output size
GSM8K	queries_gsm8k/	1,319	29 GB
MMLU Social Sciences	queries_mmlu_social_science/	3,077	66 GB
MMLU STEM	queries_mmlu_stem/	3,018	64 GB
SocialIQA	queries_socialiqa/	10,000	~211 GB

Each benchmark directory contains 633 files:

• query_ids.json -- ordered list of query IDs
• shard_0000.npy through shard_0315.npy -- score matrices, shape (shard_docs, n_queries), dtype float32
• shard_0000_doc_ids.json through shard_0315_doc_ids.json -- ordered list of doc IDs per shard

Interpreting scores:

• scores[i, j] = influence of training document doc_ids[i] on query query_ids[j]
• Positive score = training doc pushes the model toward the query's behavior
• Negative score = training doc pushes the model away from the query's behavior
• Magnitude = strength of influence
• Typical range: -0.03 to +0.04

Loading one shard

# 1) Fetch the 3 files for one query benchmark from the HF bucket
hf buckets cp hf://buckets/HCAI-Lab/trackstar-scores-base-olmes-4bench/queries_gsm8k/query_ids.json ./queries_gsm8k/
hf buckets cp hf://buckets/HCAI-Lab/trackstar-scores-base-olmes-4bench/queries_gsm8k/shard_0000.npy ./queries_gsm8k/
hf buckets cp hf://buckets/HCAI-Lab/trackstar-scores-base-olmes-4bench/queries_gsm8k/shard_0000_doc_ids.json ./queries_gsm8k/

import numpy as np
import json

# Cloud-fetched: ./queries_gsm8k/
# PACE-internal historical: /storage/ice-shared/cs7634/staff/TDA/trackstar/scores_full/base/20260326T163642Z_1102443/queries_gsm8k
benchmark_dir = "./queries_gsm8k"

with open(f"{benchmark_dir}/query_ids.json") as f:
    query_ids = json.load(f)  # ["0", "1", ..., "1318"]

scores = np.load(f"{benchmark_dir}/shard_0000.npy")  # (17971, 1319), float32
with open(f"{benchmark_dir}/shard_0000_doc_ids.json") as f:
    doc_ids = json.load(f)  # ["shard_0000:0", ..., "shard_0000:17970"]

# Example: top-5 most influential docs for query 0
query_idx = 0
top5 = np.argsort(scores[:, query_idx])[-5:][::-1]
for idx in top5:
    print(f"{doc_ids[idx]}: {scores[idx, query_idx]:.6f}")

Loading all shards for one benchmark

# Sync the full benchmark dir (~80 GB for gsm8k, ~211 GB for socialiqa)
hf buckets sync ./queries_gsm8k hf://buckets/HCAI-Lab/trackstar-scores-base-olmes-4bench/queries_gsm8k --download

import numpy as np
import json
from pathlib import Path

# Cloud-fetched local cache (see hf buckets sync above)
benchmark_dir = Path("./queries_gsm8k")
# PACE-internal historical: Path("/storage/ice-shared/cs7634/staff/TDA/trackstar/scores_full/base/20260326T163642Z_1102443/queries_gsm8k")

with open(benchmark_dir / "query_ids.json") as f:
    query_ids = json.load(f)

all_scores = []
all_doc_ids = []
for shard_path in sorted(benchmark_dir.glob("shard_*.npy")):
    shard_name = shard_path.stem
    all_scores.append(np.load(shard_path))
    with open(benchmark_dir / f"{shard_name}_doc_ids.json") as f:
        all_doc_ids.extend(json.load(f))

scores = np.concatenate(all_scores, axis=0)  # (5678621, 1319)
# scores[i, j] = influence of doc all_doc_ids[i] on query query_ids[j]

Note: the full concatenated matrix for SocialIQA (5.68M x 10K x 4 bytes) is ~211 GB. Load shard-by-shard for memory-constrained analysis.

2. Top-100 Ranked Results (JSONL)

Path: ~/dev/data-attribution-soc156/runs/trackstar/results/base/20260326T163642Z_1102443/

File	Queries	Rows	Size
queries_gsm8k_top100.jsonl	1,319	131,900	12 MB
queries_mmlu_social_science_top100.jsonl	3,077	307,700	28 MB
queries_mmlu_stem_top100.jsonl	3,018	301,800	27 MB
queries_socialiqa_top100.jsonl	10,000	1,000,000	89 MB

Each row:

{"query_id": "0", "doc_id": "shard_0072:6152", "score": 0.0224134624004364, "rank": 1}

100 rows per query, sorted by score descending (rank 1 = most influential).

Loading

import json
import pandas as pd

rows = []
with open("queries_gsm8k_top100.jsonl") as f:
    for line in f:
        rows.append(json.loads(line))
df = pd.DataFrame(rows)
# df columns: query_id, doc_id, score, rank

3. Doc ID Mapping

Doc IDs in score outputs use positional format: shard_NNNN:INDEX.

• shard_0072:6152 means the 6153rd document (0-indexed) in shard_0072.jsonl

Resolving to original doc ID and text

Cloud (primary): HCAI-Lab/dolma3-6t-sample-10000-docs-trackstar-shards HF Dataset (private). Use hf download to fetch the 316 shard_NNNN.jsonl files. Required for resolving any positional doc ID in the score matrices above.

Path on PACE (historical): /storage/ice-shared/cs7634/staff/TDA/trackstar/shards_10k/sample_10000_docs/

316 files (shard_0000.jsonl through shard_0315.jsonl), ~18K docs each. Plain JSONL with two fields:

{"id": "c929f3a2-4a12-4c85-87ac-10481f1d6624", "text": "Full document text here..."}

• id is the original Dolma corpus UUID
• text is the full document content

Lookup function

# Fetch the training shards once (~46 GB) into a local cache
hf download HCAI-Lab/dolma3-6t-sample-10000-docs-trackstar-shards --repo-type dataset --local-dir ./trackstar_shards

import json
import linecache

# Cloud-fetched local cache (see hf download above)
SHARDS_DIR = "./trackstar_shards"
# PACE-internal historical: "/storage/ice-shared/cs7634/staff/TDA/trackstar/shards_10k/sample_10000_docs"

def resolve_doc(doc_id_str):
    """Resolve 'shard_0072:6152' to {'id': 'uuid...', 'text': '...'}"""
    shard_name, index = doc_id_str.rsplit(":", 1)
    index = int(index)
    shard_path = f"{SHARDS_DIR}/{shard_name}.jsonl"
    with open(shard_path) as f:
        for i, line in enumerate(f):
            if i == index:
                return json.loads(line)
    return None

doc = resolve_doc("shard_0072:6152")
print(doc["id"])    # original Dolma UUID
print(doc["text"])  # full document text

Batch lookup (preload shard index)

import json

def build_shard_index(shard_path):
    """Build byte-offset index for fast random access."""
    offsets = []
    with open(shard_path, "rb") as f:
        while True:
            pos = f.tell()
            line = f.readline()
            if not line:
                break
            offsets.append(pos)
    return offsets

def lookup_by_offset(shard_path, offsets, index):
    with open(shard_path) as f:
        f.seek(offsets[index])
        return json.loads(f.readline())

4. Query ID Mapping

Query IDs in score outputs are positional indices ("0", "1", etc.) corresponding to row order in the source query JSONL files.

Query JSONL files: ~/dev/data-attribution-soc156/queries/base/

File	Queries	Source
olmes_gsm8k.jsonl	1,319	HCAI-Lab/base-query-data
olmes_mmlu_social_science.jsonl	3,077	HCAI-Lab/base-query-data
olmes_mmlu_stem.jsonl	3,018	HCAI-Lab/base-query-data
olmes_socialiqa.jsonl	10,000	HCAI-Lab/base-query-data

Each row:

{
    "text": "Question: Janet's ducks lay 16 eggs per day...\nAnswer: 18",
    "query_id": "gsm8k:0",
    "is_correct": true,
    "predicted_answer": "18",
    "correct_answer": "18"
}

Fields:

• text -- full OLMES evaluation prompt concatenated with the answer
• query_id -- benchmark-scoped ID (e.g., gsm8k:0, socialiqa:4231)
• is_correct -- whether OLMo3-7B answered this query correctly
• predicted_answer / correct_answer -- model output vs ground truth

Mapping positional ID to query metadata

import json

def load_queries(path):
    with open(path) as f:
        return [json.loads(line) for line in f]

queries = load_queries("queries/base/olmes_gsm8k.jsonl")
# queries[0] corresponds to query_id "0" in the score outputs
# queries[0]["query_id"] = "gsm8k:0"
# queries[0]["text"] = full prompt text
# queries[0]["is_correct"] = True/False

5. Gradient Index (Bergson build artifacts)

Cloud (primary): HCAI-Lab/trackstar-gradient-index-base HF Bucket (private, 1.3 TB). 316 subdirs shard_0000/ .. shard_0315/, each with the layout below.

Path on PACE (historical): /storage/ice-shared/cs7634/staff/TDA/trackstar/builds/base/20260326T163642Z_1102443/sample_10000_docs/

316 subdirectories (shard_0000/ through shard_0315/), each containing:

File	Description
gradients.bin	Memory-mapped gradient array (~3.8 GB per shard)
info.json	Shape metadata: num_grads, dtype definition
data.hf/	HF Dataset with length and loss columns (no doc IDs)
index_config.json	Full build parameters (model, precision, projection_dim, etc.)
normalizers.pth	Gradient normalizer state
preconditioners.pth	Preconditioner matrices
preconditioners_eigen.pth	Eigendecompositions
processor_config.json	Gradient processor config
preprocess_config.json	Preprocessing config

Total: ~1.2 TB across 316 shards. This is the reusable artifact. Scoring new query sets against it is CPU-only (no GPU needed).

Loading with Bergson Attributor

# Fetch just one shard (~4 GB) to evaluate locally
hf buckets sync ./gradient_index_shard_0000 hf://buckets/HCAI-Lab/trackstar-gradient-index-base/shard_0000 --download

from bergson.query.attributor import Attributor

# Cloud-fetched local cache (see hf buckets sync above)
train_index = Attributor(
    "./gradient_index_shard_0000",
    device="cpu",
    unit_norm=True,
)
# PACE-internal historical: "/storage/ice-shared/cs7634/staff/TDA/trackstar/builds/base/20260326T163642Z_1102443/sample_10000_docs/shard_0000"
# train_index.grads is a dict of {module_name: tensor} with shape (17971, 256) per module

6. Query Gradient Indices

Cloud (primary): HCAI-Lab/trackstar-query-gradients-base HF Bucket (private). SOC-156 base builds at base/20260326T163642Z_1102443/queries_{gsm8k,mmlu_social_science,mmlu_stem,socialiqa}/. Same Bergson layout as §5.

Path on PACE (historical): /storage/ice-shared/cs7634/staff/TDA/trackstar/query_builds/base/20260326T163642Z_1102443/

Directory	Benchmark	Queries
queries_gsm8k/	GSM8K	1,319
queries_mmlu_social_science/	MMLU Social Sciences	3,077
queries_mmlu_stem/	MMLU STEM	3,018
queries_socialiqa/	SocialIQA	10,000

Same Bergson build format as the training index. ~1 GB each. These preserve per-query gradient vectors.

7. Preconditioner

Cloud (primary): HCAI-Lab/trackstar-preconditioners HF Bucket (public). Three model variants — use olmo-3-1025-7b/ to match the SOC-156 base run. Important: the preconditioner must match the model used to build the gradients (SOC-162 finding).

Path on PACE (historical): ~/scratch/soc152/output_production/mixed_preconditioner/

Built in SOC-152 from 100K random document sample. TrackStar method with mixing coefficient 0.0877, target_downweight_components=1000.

Source data: HCAI-Lab/dolma3-6t-preconditioner-100k (100K docs, 251.5M tokens).

End-to-End Example: Find Most Influential Docs for a Query

The example below uses local paths after fetching from cloud. The "Fetch from cloud" step shows how to populate the local cache; substitute PACE paths in benchmark_dir / shards_dir if you have direct PACE access instead.

Fetch from cloud (HF) first

# 1. Query JSONL (one-time, ~50 MB)
hf download HCAI-Lab/base-query-data --repo-type dataset --local-dir ./queries/base

# 2. Score matrix for one benchmark (gsm8k example, ~29 GB)
mkdir -p ./scores/queries_gsm8k
hf buckets sync hf://buckets/HCAI-Lab/trackstar-scores-base-olmes-4bench/queries_gsm8k ./scores/queries_gsm8k

# 3. Training shards (one-time, ~41 GB) for doc-id resolution
hf download HCAI-Lab/dolma3-6t-sample-10000-docs-trackstar-shards --repo-type dataset --local-dir ./training_shards

Then load and analyse

import json
import numpy as np

# 1. Pick a query
queries = []
with open("queries/base/olmes_gsm8k.jsonl") as f:
    queries = [json.loads(line) for line in f]

query_idx = 42
print(f"Query: {queries[query_idx]['query_id']}")
print(f"Correct: {queries[query_idx]['is_correct']}")
print(f"Text: {queries[query_idx]['text'][:200]}...")

# 2. Load scores for this query across all shards
benchmark_dir = "./scores/queries_gsm8k"  # or PACE path if you have access
all_scores = []
all_doc_ids = []
for shard_idx in range(316):
    shard_name = f"shard_{shard_idx:04d}"
    scores = np.load(f"{benchmark_dir}/{shard_name}.npy")
    all_scores.append(scores[:, query_idx])
    with open(f"{benchmark_dir}/{shard_name}_doc_ids.json") as f:
        all_doc_ids.extend(json.load(f))

scores = np.concatenate(all_scores)  # (5678621,)

# 3. Find top-10 most influential docs
top10 = np.argsort(scores)[-10:][::-1]
shards_dir = "./training_shards"  # or PACE path

for rank, idx in enumerate(top10, 1):
    doc_id = all_doc_ids[idx]
    shard_name, local_idx = doc_id.rsplit(":", 1)
    with open(f"{shards_dir}/{shard_name}.jsonl") as f:
        for i, line in enumerate(f):
            if i == int(local_idx):
                doc = json.loads(line)
                break
    print(f"\nRank {rank}: score={scores[idx]:.6f}")
    print(f"  Dolma ID: {doc['id']}")
    print(f"  Text: {doc['text'][:150]}...")

If you only need the top-100 / top-2K rather than the full per-doc scores, use HCAI-Lab/dolma3-trackstar-influence-scores instead — that dataset has the pre-ranked results per benchmark and skips the 29 GB download.