| """Reconstruct reviewed requests from pinned source data and text-free label sidecars.""" |
|
|
| |
| import argparse |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
|
|
| def sha256_text(text): |
| return hashlib.sha256(text.encode()).hexdigest() |
|
|
|
|
| def hydrate_records(sidecar, sources): |
| repository = sidecar["source"]["repository"] |
| result = [] |
| if repository == "CohereLabs/aya_dataset": |
| lookup = {} |
| for row in sources: |
| if not row.get("user_id"): |
| continue |
| key = ( |
| sha256_text(row["inputs"]), |
| row["language_code"], |
| "aya-author-" + sha256_text(row["user_id"]), |
| ) |
| lookup[key] = row["inputs"] |
| for row in sidecar["records"]: |
| key = (row["source_text_sha256"], row["language"], row["group_id"]) |
| if key not in lookup: |
| raise ValueError( |
| "Reviewed Aya request or author differs from the pinned source" |
| ) |
| result.append({**row, "text": lookup[key]}) |
| elif repository == "databricks/databricks-dolly-15k": |
| for row in sidecar["records"]: |
| source = sources[row["source_row"]] |
| instruction, context = ( |
| source["instruction"].strip(), |
| source["context"].strip(), |
| ) |
| text = instruction + ("\n\n" + context if context else "") |
| if sha256_text(text) != row["source_text_sha256"]: |
| raise ValueError( |
| "Reviewed Dolly request differs from the pinned source" |
| ) |
| result.append( |
| {**row, "text": text, "instruction": instruction, "context": context} |
| ) |
| else: |
| raise ValueError("Unknown reviewed source") |
| return result |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser() |
| parser.add_argument("--sidecar", type=Path, required=True) |
| parser.add_argument("--source", type=Path, required=True) |
| parser.add_argument("--output", type=Path, required=True) |
| args = parser.parse_args() |
| if args.output.exists(): |
| raise ValueError("Refusing to overwrite reviewed annotations") |
| sidecar = json.loads(args.sidecar.read_text()) |
| if ( |
| hashlib.sha256(args.source.read_bytes()).hexdigest() |
| != sidecar["source"]["file_sha256"] |
| ): |
| raise ValueError("Source file hash differs from the reviewed revision") |
| if sidecar["source"]["repository"] == "CohereLabs/aya_dataset": |
| from pyarrow import parquet |
|
|
| sources = parquet.read_table( |
| args.source, columns=["inputs", "language_code", "user_id"] |
| ).to_pylist() |
| else: |
| sources = [ |
| json.loads(line) |
| for line in args.source.read_text().split("\n") |
| if line.strip() |
| ] |
| rows = hydrate_records(sidecar, sources) |
| content = "".join( |
| json.dumps(row, ensure_ascii=False, sort_keys=True) + "\n" for row in rows |
| ) |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| args.output.write_text(content) |
| print( |
| json.dumps( |
| {"rows": len(rows), "sha256": hashlib.sha256(content.encode()).hexdigest()} |
| ) |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
