| |
| """Lock, materialize, and tokenize WikiText-2 without reading its test split.""" |
|
|
| from __future__ import annotations |
|
|
| import hashlib |
| import json |
| from datetime import datetime, timezone |
| from pathlib import Path |
|
|
| import torch |
| from datasets import load_dataset |
| from huggingface_hub import HfApi |
| from transformers import AutoTokenizer |
|
|
|
|
| PROJECT_DIR = Path(__file__).resolve().parents[2] |
| RAW_DIR = PROJECT_DIR / "datasets/raw/wikitext2" |
| PROCESSED_DIR = PROJECT_DIR / "datasets/processed/wikitext2_qwen38" |
| RECORD_DIR = PROJECT_DIR / "records" |
|
|
|
|
| def sha256_file(path: Path) -> str: |
| digest = hashlib.sha256() |
| with path.open("rb") as handle: |
| for chunk in iter(lambda: handle.read(8 << 20), b""): |
| digest.update(chunk) |
| return digest.hexdigest() |
|
|
|
|
| def main() -> None: |
| revision = "b08601e04326c79dfdd32d625aee71d232d685c3" |
| |
| |
| |
| dataset = { |
| split_name: load_dataset( |
| "Salesforce/wikitext", |
| "wikitext-2-raw-v1", |
| revision=revision, |
| split=split_name, |
| ) |
| for split_name in ("train", "validation") |
| } |
| model_dir = Path((RECORD_DIR / "active_model_path.txt").read_text(encoding="utf-8").strip()) |
| tokenizer = AutoTokenizer.from_pretrained(model_dir, local_files_only=True) |
| tokenizer.model_max_length = 10**12 |
| RAW_DIR.mkdir(parents=True, exist_ok=True) |
| PROCESSED_DIR.mkdir(parents=True, exist_ok=True) |
| splits = {} |
| |
| for split_name in ("train", "validation"): |
| split = dataset[split_name] |
| parquet_path = RAW_DIR / f"{split_name}.parquet" |
| split.to_parquet(parquet_path) |
| text = "\n\n".join(row for row in split["text"] if row.strip()) |
| token_ids = torch.tensor( |
| tokenizer(text, add_special_tokens=False)["input_ids"], dtype=torch.int32 |
| ) |
| token_path = PROCESSED_DIR / f"{split_name}-tokens.pt" |
| torch.save(token_ids, token_path) |
| splits[split_name] = { |
| "rows": len(split), |
| "dataset_fingerprint": split._fingerprint, |
| "parquet_path": str(parquet_path), |
| "parquet_sha256": sha256_file(parquet_path), |
| "token_path": str(token_path), |
| "token_sha256": sha256_file(token_path), |
| "token_count": int(token_ids.numel()), |
| } |
| payload = { |
| "status": "ok", |
| "dataset": "Salesforce/wikitext", |
| "config": "wikitext-2-raw-v1", |
| "source_revision": revision, |
| "license": "CC BY-SA 4.0", |
| "tokenizer_path": str(model_dir), |
| "tokenizer_vocab_size": tokenizer.vocab_size, |
| "splits": splits, |
| "test_rows_used_for_training_or_evaluation": 0, |
| "protocol_disclosure": { |
| "initial_datasetdict_load_materialized_test_rows_in_cache": 4358, |
| "test_rows_inspected_or_used_for_fitness": 0, |
| "correction": "current script requests train and validation splits individually", |
| }, |
| "prepared_at": datetime.now(timezone.utc).isoformat(), |
| } |
| output = PROJECT_DIR / "artifacts/wikitext2-preparation.json" |
| output.write_text(json.dumps(payload, indent=2), encoding="utf-8") |
| print(json.dumps(payload, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
