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

```bash
# 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/
```

```python
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

```bash
# 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
```

```python
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:

```json
{"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

```python
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:

```json
{"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

```bash
# 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
```

```python
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)

```python
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:

```json
{
    "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

```python
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

```bash
# 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
```

```python
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

```bash
# 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

```python
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.