sources_to_hplt_ids.py (3763B)
1 from typing import Any 2 import argparse 3 import collections 4 import json 5 import pathlib 6 import zipfile 7 8 import datasets 9 10 from hplt_align.common import map_minhash 11 12 13 def parse_command_argument() -> argparse.Namespace: 14 """ Parse command line arguments. """ 15 parser = argparse.ArgumentParser(description="Update a huggingface dataset using a minhash input and output.") 16 parser.add_argument("output", type=pathlib.Path, help="Path to the output dataset splits.") 17 parser.add_argument("dataset", type=pathlib.Path, help="Path to the input dataset splits.") 18 parser.add_argument("sources", type=pathlib.Path, help="Path to the source jsonl input file.") 19 parser.add_argument("minhash", type=pathlib.Path, help="Path to the output zip of the minhash script.") 20 21 args: argparse.Namespace = parser.parse_args() 22 return args 23 24 25 def read_sources(sources_path: pathlib.Path) -> dict[int, int]: 26 """Returns a dictionary mapping UCDP ids to line ids.""" 27 sources: dict[int, int] = {} 28 with sources_path.open("r") as sources_file: 29 for line_number, source_line in enumerate(sources_file): 30 source: dict[str, Any] = json.loads(source_line) 31 sources[source["id"]] = line_number + 1 # Line numbers start at 1 32 return sources 33 34 35 def match_filter(match: dict[str, Any]) -> bool: 36 # Filter out non-English matches and short texts and long texts 37 return match["qid"].startswith("deduplicated/eng_Latn/") and 100 <= len(match["text"]) <= 10000 38 39 40 def read_minhash(minhash_path: pathlib.Path) -> dict[int, list[dict[str, Any]]]: 41 """Returns a dictionary mapping each line id to a list of HPLT match in order of approximate Jacard similarity.""" 42 minhash: dict[int, list[dict[str, Any]]] = collections.defaultdict(list) 43 map_minhash(minhash_path, lambda match: minhash[abs(match["tid"])].append(match)) 44 45 for tid, matches in minhash.items(): 46 minhash[tid] = sorted(list(filter(match_filter, matches)), key=lambda match: match["sim"], reverse=True) 47 return minhash 48 49 50 def process_split(output_path: pathlib.Path, split_path: pathlib.Path, sources: dict[int, int], minhash: dict[int, list[dict[str, Any]]], source_count: dict[str, int]) -> None: 51 datasets.load_from_disk(str(split_path))\ 52 .filter(input_columns=["id"], load_from_cache_file=False, 53 function = lambda eid: eid in sources and minhash[sources[eid]] and source_count[minhash[sources[eid]][0]["text"]]==1)\ 54 .map(input_columns=["id"], load_from_cache_file=False, 55 function = lambda eid: {"source_article": minhash[sources[eid]][0]["qid"]})\ 56 .save_to_disk(str(output_path)) 57 58 59 def main(output_path: pathlib.Path, dataset_path: pathlib.Path, sources_path: pathlib.Path, minhash_path: pathlib.Path) -> None: 60 output_path.mkdir(parents=True) 61 sources: dict[int, int] = read_sources(sources_path) 62 minhash: dict[int, list[dict[str, Any]]] = read_minhash(minhash_path) 63 64 # If two documents are mapped to the same HPLT document, it indicates a document of poor quality, count them in order to remove repeated ones. 65 source_count: dict[str, int] = collections.Counter( 66 minhash[sources[sample["id"]]][0]["text"] 67 for split_path in dataset_path.iterdir() 68 for sample in datasets.load_from_disk(str(split_path)) 69 if sample["id"] in sources and minhash[sources[sample["id"]]]) 70 for split_path in dataset_path.iterdir(): 71 process_split(output_path / split_path.name, split_path, sources, minhash, source_count) 72 73 74 if __name__ == "__main__": 75 args: argparse.Namespace = parse_command_argument() 76 main(output_path=args.output, dataset_path=args.dataset, sources_path=args.sources, minhash_path=args.minhash)