429 lines
20 KiB
Python
429 lines
20 KiB
Python
#!/usr/bin/env python3
|
|
"""Prepare Slovak Aya train/eval/test data for LLaMA-Factory."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from collections import Counter
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
from itertools import islice
|
|
import json
|
|
from pathlib import Path
|
|
import sqlite3
|
|
import sys
|
|
import unicodedata
|
|
from typing import Any, Iterator
|
|
|
|
|
|
HF_DATASET = "CohereLabs/aya_collection_language_split"
|
|
HF_REVISION = "a3af2fde4b4cb5b2775830b11244a1a20b5f004f"
|
|
HF_CONFIG = "slovak"
|
|
HF_SPLIT = "train"
|
|
SPLITS = ("train", "eval", "test")
|
|
|
|
|
|
def normalized(text: str, *, fold_case: bool = False) -> str:
|
|
text = " ".join(unicodedata.normalize("NFKC", text).split())
|
|
return text.casefold() if fold_case else text
|
|
|
|
|
|
def digest_json(value: Any) -> str:
|
|
payload = json.dumps(value, ensure_ascii=False, separators=(",", ":"))
|
|
return hashlib.sha256(payload.encode("utf-8")).hexdigest()
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
result = hashlib.sha256()
|
|
with path.open("rb") as source:
|
|
for chunk in iter(lambda: source.read(1024 * 1024), b""):
|
|
result.update(chunk)
|
|
return result.hexdigest()
|
|
|
|
|
|
def write_json(path: Path, content: Any) -> None:
|
|
with path.open("x", encoding="utf-8", newline="\n") as target:
|
|
json.dump(content, target, ensure_ascii=False, indent=2, sort_keys=True)
|
|
target.write("\n")
|
|
|
|
|
|
def local_rows(path: Path) -> Iterator[Any]:
|
|
with path.open("r", encoding="utf-8-sig") as source:
|
|
for line_number, line in enumerate(source, 1):
|
|
if not line.strip():
|
|
continue
|
|
try:
|
|
yield json.loads(line)
|
|
except json.JSONDecodeError as error:
|
|
raise ValueError(f"Invalid JSON at {path}:{line_number}: {error}") from error
|
|
|
|
|
|
def open_source(args: argparse.Namespace) -> tuple[Iterator[Any], dict[str, Any]]:
|
|
if args.local_jsonl:
|
|
path = args.local_jsonl.expanduser().resolve(strict=True)
|
|
if not path.is_file():
|
|
raise ValueError(f"Local source is not a file: {path}")
|
|
return local_rows(path), {
|
|
"kind": "local_jsonl",
|
|
"path": str(path),
|
|
"sha256": file_sha256(path),
|
|
"bytes": path.stat().st_size,
|
|
"revision": None,
|
|
}
|
|
|
|
try:
|
|
from datasets import load_dataset
|
|
from huggingface_hub import HfApi
|
|
except ImportError as error:
|
|
raise RuntimeError("Install datasets and huggingface_hub for Hugging Face input.") from error
|
|
|
|
info = HfApi().dataset_info(HF_DATASET, revision=args.revision or HF_REVISION)
|
|
revision = info.sha
|
|
if not revision:
|
|
raise RuntimeError("Hugging Face did not return an immutable dataset revision.")
|
|
data = load_dataset(
|
|
HF_DATASET, HF_CONFIG, split=HF_SPLIT, revision=revision, streaming=True
|
|
)
|
|
return iter(data), {
|
|
"kind": "huggingface",
|
|
"dataset": HF_DATASET,
|
|
"config": HF_CONFIG,
|
|
"split": HF_SPLIT,
|
|
"requested_revision": args.revision,
|
|
"revision": revision,
|
|
"streaming": True,
|
|
}
|
|
|
|
|
|
def convert_row(row: Any) -> tuple[dict[str, str], dict[str, Any]] | None:
|
|
if not isinstance(row, dict):
|
|
return None
|
|
if "inputs" in row or "targets" in row:
|
|
instruction, query, response = row.get("inputs"), "", row.get("targets")
|
|
excluded = {"inputs", "targets"}
|
|
else:
|
|
instruction, query, response = row.get("instruction"), row.get("input", ""), row.get("output")
|
|
query = "" if query is None else query
|
|
excluded = {"instruction", "input", "output"}
|
|
if not all(isinstance(value, str) for value in (instruction, query, response)):
|
|
return None
|
|
if not instruction.strip() or not response.strip():
|
|
return None
|
|
metadata = {key: value for key, value in row.items() if key not in excluded}
|
|
return {
|
|
"instruction": instruction.strip(),
|
|
"input": query.strip(),
|
|
"output": response.strip(),
|
|
}, metadata
|
|
|
|
|
|
def create_database(path: Path) -> sqlite3.Connection:
|
|
connection = sqlite3.connect(path)
|
|
connection.executescript(
|
|
"""
|
|
PRAGMA journal_mode = DELETE;
|
|
PRAGMA synchronous = NORMAL;
|
|
PRAGMA temp_store = FILE;
|
|
PRAGMA cache_size = -32768;
|
|
CREATE TABLE prompt_groups (
|
|
group_id TEXT PRIMARY KEY,
|
|
order_key TEXT NOT NULL,
|
|
size INTEGER NOT NULL,
|
|
split TEXT NOT NULL DEFAULT 'train'
|
|
);
|
|
CREATE TABLE records (
|
|
position INTEGER PRIMARY KEY,
|
|
stable_id TEXT NOT NULL UNIQUE,
|
|
group_id TEXT NOT NULL,
|
|
instruction TEXT NOT NULL,
|
|
input TEXT NOT NULL,
|
|
output TEXT NOT NULL,
|
|
metadata TEXT NOT NULL,
|
|
source_row_index INTEGER NOT NULL,
|
|
source_dataset TEXT NOT NULL,
|
|
source_task TEXT NOT NULL
|
|
);
|
|
"""
|
|
)
|
|
return connection
|
|
|
|
|
|
SYSTEM_TEXT = "Below is an instruction that describes a task. Write a response that appropriately completes the request.\n\n"
|
|
USER_TEXT = "### Instruction:\n{}\n\n### Response:\n"
|
|
|
|
|
|
class SFTLengthFilter:
|
|
"""Keep a whole pair only when its exact Alpaca length fits both model configs."""
|
|
def __init__(self, config_dir: Path):
|
|
import yaml
|
|
from transformers import AutoTokenizer
|
|
self.profiles = {}
|
|
self.excluded_by_model = Counter()
|
|
self.max_accepted_tokens = Counter()
|
|
for name in ("mistral", "olmo"):
|
|
settings = yaml.safe_load((config_dir / f"{name}.yaml").read_text(encoding="utf-8-sig"))
|
|
if settings.get("template") != "alpaca":
|
|
raise ValueError("Length filtering currently requires the Alpaca template for both models.")
|
|
limit = int(settings["cutoff_len"])
|
|
if limit < 1:
|
|
raise ValueError("cutoff_len must be positive.")
|
|
tokenizer = AutoTokenizer.from_pretrained(settings["model_name_or_path"],
|
|
revision=settings["model_revision"], use_fast=True, fix_mistral_regex=True,
|
|
model_max_length=10**9)
|
|
backend = tokenizer.backend_tokenizer
|
|
backend.no_truncation()
|
|
backend.no_padding()
|
|
if tokenizer.eos_token_id is None:
|
|
raise ValueError(f"Missing EOS token: {name}")
|
|
self.profiles[name] = {
|
|
"model_name_or_path": settings["model_name_or_path"],
|
|
"model_revision": settings["model_revision"], "template": "alpaca",
|
|
"cutoff_len": limit, "tokenizer_class": type(tokenizer).__name__,
|
|
"system_tokens": len(tokenizer.encode(SYSTEM_TEXT, add_special_tokens=False)),
|
|
"ending_tokens": 1 + len(tokenizer.encode("\n\n", add_special_tokens=False)),
|
|
"backend": backend,
|
|
}
|
|
|
|
def lengths(self, examples: list[dict[str, str]]) -> list[dict[str, int]]:
|
|
prompts = [USER_TEXT.format("\n".join(part for part in (row["instruction"], row["input"]) if part))
|
|
for row in examples]
|
|
responses = [row["output"] for row in examples]
|
|
result = [{} for _ in examples]
|
|
for name, profile in self.profiles.items():
|
|
backend = profile["backend"]
|
|
encode = getattr(backend, "encode_batch_fast", backend.encode_batch)
|
|
inputs = encode(prompts, add_special_tokens=False)
|
|
targets = encode(responses, add_special_tokens=False)
|
|
for item, source, target in zip(result, inputs, targets):
|
|
item[name] = profile["system_tokens"] + len(source) + len(target) + profile["ending_tokens"]
|
|
return result
|
|
|
|
def over_limit(self, lengths: dict[str, int]) -> list[str]:
|
|
return [name for name, length in lengths.items() if length > self.profiles[name]["cutoff_len"]]
|
|
|
|
def report(self) -> dict:
|
|
return {"enabled": True,
|
|
"models": {name: {key: value for key, value in profile.items() if key != "backend"}
|
|
for name, profile in self.profiles.items()},
|
|
"excluded_by_model": {name: self.excluded_by_model[name] for name in self.profiles},
|
|
"max_accepted_tokens_before_pair_dedup": {name: self.max_accepted_tokens[name] for name in self.profiles},
|
|
"excluded_log": "excluded_long.jsonl"}
|
|
|
|
|
|
def ingest(connection: sqlite3.Connection, rows: Iterator[Any], seed: int,
|
|
length_filter: SFTLengthFilter, excluded) -> Counter:
|
|
counts: Counter = Counter()
|
|
indexed = enumerate(rows)
|
|
next_progress = 10000
|
|
while batch := list(islice(indexed, 512)):
|
|
candidates = []
|
|
for source_index, row in batch:
|
|
counts["source_rows"] += 1
|
|
converted = convert_row(row)
|
|
if converted is None:
|
|
counts["invalid_or_empty_rows"] += 1
|
|
else:
|
|
candidates.append((source_index, *converted))
|
|
measured = length_filter.lengths([data for _, data, _ in candidates])
|
|
for (source_index, data, metadata), lengths in zip(candidates, measured):
|
|
over = length_filter.over_limit(lengths)
|
|
if over:
|
|
counts["filtered_long_rows"] += 1
|
|
length_filter.excluded_by_model.update(over)
|
|
excluded.write(json.dumps({"source_row_index": source_index, "source_id": metadata.get("id"),
|
|
"dataset_name": metadata.get("dataset_name"), "task_type": metadata.get("task_type"),
|
|
"token_lengths": lengths, "over_limit_models": over}, ensure_ascii=False) + "\n")
|
|
continue
|
|
for name, length in lengths.items():
|
|
length_filter.max_accepted_tokens[name] = max(length_filter.max_accepted_tokens[name], length)
|
|
prompt_key = [normalized(data["instruction"], fold_case=True), normalized(data["input"], fold_case=True)]
|
|
group_id = digest_json(prompt_key)
|
|
stable_id = digest_json([prompt_key, normalized(data["output"])])
|
|
cursor = connection.execute(
|
|
"""INSERT OR IGNORE INTO records
|
|
(stable_id, group_id, instruction, input, output, metadata,
|
|
source_row_index, source_dataset, source_task)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
|
|
(stable_id, group_id, data["instruction"], data["input"], data["output"],
|
|
json.dumps(metadata, ensure_ascii=False, separators=(",", ":")), source_index,
|
|
str(metadata.get("dataset_name") or "unknown"), str(metadata.get("task_type") or "unknown")))
|
|
if cursor.rowcount == 0:
|
|
counts["duplicate_pairs_removed"] += 1
|
|
else:
|
|
counts["retained_rows"] += 1
|
|
connection.execute(
|
|
"""INSERT INTO prompt_groups (group_id, order_key, size) VALUES (?, ?, 1)
|
|
ON CONFLICT(group_id) DO UPDATE SET size = size + 1""",
|
|
(group_id, digest_json([seed, group_id])))
|
|
if counts["source_rows"] >= next_progress:
|
|
connection.commit()
|
|
print(f"Read {counts['source_rows']:,}; filtered long {counts['filtered_long_rows']:,}; retained {counts['retained_rows']:,}", flush=True)
|
|
next_progress += 10000
|
|
connection.commit()
|
|
for key in ("source_rows", "invalid_or_empty_rows", "filtered_long_rows", "duplicate_pairs_removed", "retained_rows"):
|
|
counts.setdefault(key, 0)
|
|
if not counts["retained_rows"]:
|
|
raise ValueError("The source contains no usable examples after validation and length filtering.")
|
|
connection.execute("CREATE INDEX prompt_group_order ON prompt_groups(order_key, group_id)")
|
|
connection.execute("CREATE INDEX records_group ON records(group_id)")
|
|
connection.commit()
|
|
counts["prompt_groups"] = connection.execute("SELECT COUNT(*) FROM prompt_groups").fetchone()[0]
|
|
counts["groups_with_multiple_responses"] = connection.execute("SELECT COUNT(*) FROM prompt_groups WHERE size > 1").fetchone()[0]
|
|
return counts
|
|
|
|
|
|
def assign_splits(connection: sqlite3.Connection, eval_size: int, test_size: int) -> dict[str, dict[str, int]]:
|
|
remaining = {"eval": eval_size, "test": test_size}
|
|
assignments: list[tuple[str, str]] = []
|
|
# Selection depends only on seed, normalized prompt, and the number of its
|
|
# distinct responses. An entire prompt group always stays in one split.
|
|
for group_id, size in connection.execute(
|
|
"SELECT group_id, size FROM prompt_groups ORDER BY order_key, group_id"
|
|
):
|
|
split = next((name for name in ("eval", "test") if size <= remaining[name]), None)
|
|
if split is not None:
|
|
remaining[split] -= size
|
|
assignments.append((split, group_id))
|
|
if not any(remaining.values()):
|
|
break
|
|
connection.executemany("UPDATE prompt_groups SET split = ? WHERE group_id = ?", assignments)
|
|
connection.commit()
|
|
result = {name: {"rows": 0, "prompt_groups": 0} for name in SPLITS}
|
|
for split, group_count, row_count in connection.execute(
|
|
"SELECT split, COUNT(*), SUM(size) FROM prompt_groups GROUP BY split"
|
|
):
|
|
result[split] = {"rows": row_count, "prompt_groups": group_count}
|
|
if not result["train"]["rows"]:
|
|
raise ValueError("No rows remain for train. Reduce --eval-size and/or --test-size.")
|
|
# A record belongs to one prompt group; each group has exactly one split.
|
|
conflicts = connection.execute(
|
|
"""SELECT COUNT(*) FROM (
|
|
SELECT r.group_id FROM records r JOIN prompt_groups g USING(group_id)
|
|
GROUP BY r.group_id HAVING COUNT(DISTINCT g.split) != 1
|
|
)"""
|
|
).fetchone()[0]
|
|
if conflicts:
|
|
raise RuntimeError("Prompt-disjoint split verification failed.")
|
|
if result["eval"]["rows"] > eval_size or result["test"]["rows"] > test_size:
|
|
raise RuntimeError("A holdout exceeds its requested row limit.")
|
|
return result
|
|
|
|
|
|
def emit_splits(connection: sqlite3.Connection, output: Path, expected: dict[str, dict[str, int]]) -> dict[str, Any]:
|
|
files: dict[str, Any] = {}
|
|
for split in SPLITS:
|
|
data_path = output / f"{split}.jsonl"
|
|
metadata_path = output / f"{split}.metadata.jsonl"
|
|
row_count = 0
|
|
with data_path.open("x", encoding="utf-8", newline="\n") as data_file, metadata_path.open(
|
|
"x", encoding="utf-8", newline="\n"
|
|
) as metadata_file:
|
|
cursor = connection.execute(
|
|
"""SELECT r.stable_id, r.group_id, r.instruction, r.input, r.output,
|
|
r.metadata, r.source_row_index
|
|
FROM records r JOIN prompt_groups g USING(group_id)
|
|
WHERE g.split = ? ORDER BY g.order_key, r.stable_id""", (split,)
|
|
)
|
|
for stable_id, group_id, instruction, query, response, metadata, source_index in cursor:
|
|
item = {"instruction": instruction, "input": query, "output": response}
|
|
provenance = {
|
|
"row_index": row_count,
|
|
"stable_id": stable_id,
|
|
"prompt_group_id": group_id,
|
|
"source_row_index": source_index,
|
|
"source_metadata": json.loads(metadata),
|
|
}
|
|
provenance["source_id"] = provenance["source_metadata"].get("id")
|
|
data_file.write(json.dumps(item, ensure_ascii=False, separators=(",", ":")) + "\n")
|
|
metadata_file.write(json.dumps(provenance, ensure_ascii=False, separators=(",", ":")) + "\n")
|
|
row_count += 1
|
|
if row_count != expected[split]["rows"]:
|
|
raise RuntimeError(f"Unexpected {split} output count: {row_count}")
|
|
for path in (data_path, metadata_path):
|
|
files[path.name] = {"sha256": file_sha256(path), "bytes": path.stat().st_size, "rows": row_count}
|
|
return files
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--revision", help="Hugging Face revision; resolved to an immutable commit before loading.")
|
|
parser.add_argument("--local-jsonl", type=Path, help="Offline JSONL input in Aya or Alpaca format.")
|
|
parser.add_argument("--eval-size", type=int, default=1000, help="Maximum number of eval records (default: 1000).")
|
|
parser.add_argument("--test-size", type=int, default=1000, help="Maximum number of test records (default: 1000).")
|
|
parser.add_argument("--seed", type=int, default=42)
|
|
parser.add_argument("--keep-db", action="store_true", help="Keep the temporary SQLite database.")
|
|
args = parser.parse_args()
|
|
if min(args.eval_size, args.test_size) < 0:
|
|
parser.error("Holdout sizes must be nonnegative.")
|
|
if args.local_jsonl and args.revision:
|
|
parser.error("--revision applies only to Hugging Face input.")
|
|
output = args.output_dir.expanduser().resolve()
|
|
if output.exists() and (not output.is_dir() or any(output.iterdir())):
|
|
parser.error(f"Output must be a new or empty directory: {output}")
|
|
# Resolve source before creating output, so authentication or dependency
|
|
# errors do not produce a misleading empty data directory.
|
|
rows, source = open_source(args)
|
|
length_filter = SFTLengthFilter(Path(__file__).resolve().parents[1] / "configs")
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
database_path = output / "preparation.sqlite3"
|
|
connection = create_database(database_path)
|
|
try:
|
|
excluded_path = output / "excluded_long.jsonl"
|
|
with excluded_path.open("x", encoding="utf-8", newline="\n") as excluded:
|
|
counts = ingest(connection, rows, args.seed, length_filter, excluded)
|
|
splits = assign_splits(connection, args.eval_size, args.test_size)
|
|
if sum(item["rows"] for item in splits.values()) != counts["retained_rows"]:
|
|
raise RuntimeError("Split counts do not sum to the retained record count.")
|
|
files = emit_splits(connection, output, splits)
|
|
files[excluded_path.name] = {"sha256": file_sha256(excluded_path),
|
|
"bytes": excluded_path.stat().st_size,
|
|
"rows": counts["filtered_long_rows"]}
|
|
registration = {
|
|
f"aya_{split}": {
|
|
"file_name": f"{split}.jsonl", "formatting": "alpaca",
|
|
"columns": {"prompt": "instruction", "query": "input", "response": "output"},
|
|
} for split in SPLITS
|
|
}
|
|
info_path = output / "dataset_info.json"
|
|
write_json(info_path, registration)
|
|
files[info_path.name] = {"sha256": file_sha256(info_path), "bytes": info_path.stat().st_size}
|
|
breakdown = [
|
|
{"split": split, "dataset_name": dataset, "task_type": task, "rows": rows_count, "prompt_groups": groups_count}
|
|
for split, dataset, task, rows_count, groups_count in connection.execute(
|
|
"""SELECT g.split, r.source_dataset, r.source_task, COUNT(*), COUNT(DISTINCT r.group_id)
|
|
FROM records r JOIN prompt_groups g USING(group_id)
|
|
GROUP BY g.split, r.source_dataset, r.source_task
|
|
ORDER BY g.split, r.source_dataset, r.source_task"""
|
|
)
|
|
]
|
|
manifest = {
|
|
"format_version": 1,
|
|
"prepared_at_utc": datetime.now(timezone.utc).isoformat(),
|
|
"preparation_script_sha256": file_sha256(Path(__file__).resolve()),
|
|
"source": source,
|
|
"seed": args.seed,
|
|
"counts": dict(counts),
|
|
"length_filter": length_filter.report(),
|
|
"requested_holdout_max_rows": {"eval": args.eval_size, "test": args.test_size},
|
|
"splits": splits,
|
|
"source_breakdown": breakdown,
|
|
"files": files,
|
|
}
|
|
write_json(output / "manifest.json", manifest)
|
|
finally:
|
|
connection.close()
|
|
if not args.keep_db:
|
|
# Only our explicitly named staging file inside this new output is removed.
|
|
database_path.unlink()
|
|
print(json.dumps({"output_dir": str(output), "source": source, "counts": dict(counts), "splits": splits}, ensure_ascii=False, indent=2))
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except (ValueError, RuntimeError) as error:
|
|
print(f"ERROR: {error}", file=sys.stderr)
|
|
raise SystemExit(1)
|