Spaces:
Sleeping
Sleeping
| import contextlib | |
| import gzip | |
| import logging | |
| import re | |
| import subprocess | |
| import tempfile | |
| from collections import defaultdict | |
| from pathlib import Path | |
| from typing import Callable, Dict, Iterable, List, NamedTuple, Type | |
| from cc_net.jsonql import open_remote_file, open_write | |
| from cc_net.process_wet_file import CCSegmentsReader | |
| from typing import Sequence | |
| import functools | |
| import multiprocessing | |
| BUFFER_SIZE = "32G" | |
| SORT_PARALLEL = 8 | |
| KNOWN_VERSIONS = ["v1.0.0", "v1.0.beta", "v1.0.alpha"] | |
| class NormalizedBitextPtr(NamedTuple): | |
| lang_pair: str | |
| line_no: int | |
| segment: str | |
| digest: str | |
| ptr_start: int | |
| ptr_end: int | |
| score: float | |
| class Bitext(NamedTuple): | |
| lang_pair: str | |
| line_no: int | |
| score: float | |
| text: str | |
| class SimpleBitext(NamedTuple): | |
| line_no: int | |
| score: float | |
| text: str | |
| WEB_PAT = re.compile(r"https?:[^ \n]* ") | |
| WEB_REPL = "WEB " | |
| WEB2_PAT = re.compile(r"https?:[^ \n]*\n") | |
| WEB2_REPL = "WEB\n" | |
| def clean_content(raw_content: str) -> str: | |
| # We need to clean all the content, because otherwise there is no way for | |
| # the user to know if we need to clean it or not. | |
| par = raw_content | |
| par = par.replace("</s>", ". ") | |
| par = par.replace("\t", " ") | |
| par = re.sub(WEB_PAT, WEB_REPL, par, count=0) | |
| par = re.sub(WEB2_PAT, WEB2_REPL, par, count=0) | |
| return par | |
| def get_typed_parser(cls: Type) -> Callable: | |
| types = cls.__annotations__.values() | |
| def parser(line: str) -> NamedTuple: | |
| parts = line.rstrip("\n").split("\t") | |
| assert len(parts) == len( | |
| types | |
| ), f"Print size mismatch expected the following columns {cls.__annotations__} got: {parts}" | |
| return cls(*(t(p) for t, p in zip(types, parts))) | |
| return parser | |
| def open_read(file: Path) -> Iterable[str]: | |
| if file.suffix == ".gz": | |
| reader = gzip.open(file, "rt") | |
| else: | |
| reader = open(file, "rt") | |
| with reader as f: | |
| for line in f: | |
| yield line | |
| def dl(outdir: Path = Path("data"), version: str = KNOWN_VERSIONS[0], parallelism: int = 8): | |
| """ | |
| Download bitext pointers from FAIR dataset and extract corresponding CC snippets. | |
| - version: Specific version to download | |
| - outdir: Directory where the data should go. Files will be in {outdir}/{version}/raw/ | |
| """ | |
| assert version in KNOWN_VERSIONS, f"Unknown version {version}, chose from {KNOWN_VERSIONS}" | |
| metadata_dir = f"https://dl.fbaipublicfiles.com/laser/CCMatrix/{version}" | |
| file_list = [l.strip() for l in open_remote_file(metadata_dir + "/list.txt")] | |
| outdir.mkdir(exist_ok=True) | |
| outdir = outdir / version / "raw" | |
| outdir.mkdir(exist_ok=True, parents=True) | |
| dlf = functools.partial(dl_file, metadata_dir, outdir) | |
| # list(map(dlf, file_list)) | |
| with multiprocessing.Pool(parallelism) as pool: | |
| pool.map(dlf, file_list) | |
| def get_documents(segment: str) -> Dict[str, str]: | |
| return {d["digest"]: d["raw_content"] for d in CCSegmentsReader([segment])} | |
| def dl_file(metadata_dir: str, outdir: Path, file: str): | |
| metadata = "/".join((metadata_dir, file)) | |
| parser = get_typed_parser(NormalizedBitextPtr) | |
| found_bitext, missed_bitext, skipped_line = 0, 0, 0 | |
| segment = "" | |
| segment_downloads: Dict[str, int] = defaultdict(int) | |
| raw_documents: Dict[str, str] = {} | |
| cleaned_documents: Dict[str, str] = {} | |
| outfile = outdir / file | |
| if outfile.exists(): | |
| return | |
| o = FileWriterWithTmp(outfile) | |
| for i, line in enumerate(open_remote_file(metadata)): | |
| try: | |
| bitext: NormalizedBitextPtr = parser(line) | |
| # Add some more assert in case the line is invalid but still parse | |
| assert bitext.segment.startswith("crawl-data/") | |
| assert bitext.digest.startswith("sha1:") | |
| except AssertionError: | |
| logging.error(f"Skipping line {i}: {line}") | |
| skipped_line += 1 | |
| continue | |
| if not segment or bitext.segment != segment: | |
| segment = bitext.segment | |
| segment_downloads[segment] += 1 | |
| # Load segment in RAM, purge document cache | |
| raw_documents = get_documents(segment) | |
| cleaned_documents = {} | |
| raw_doc = raw_documents.get(bitext.digest) | |
| if raw_doc is None: | |
| logging.error(f"Document not found: {bitext.digest} in {segment}") | |
| missed_bitext += 1 | |
| continue | |
| clean_doc = cleaned_documents.get(bitext.digest) | |
| if clean_doc is None: | |
| clean_doc = clean_content(raw_doc) | |
| cleaned_documents[bitext.digest] = clean_doc | |
| text = clean_doc[bitext.ptr_start : bitext.ptr_end] | |
| score = getattr(bitext, "score", 0.0) | |
| bt = Bitext(bitext.lang_pair, bitext.line_no, score, text) | |
| print(*bt, sep="\t", file=o) | |
| o.close(True) | |
| logging.info(f"Found {found_bitext} sentences, missed {missed_bitext} sentences.") | |
| if skipped_line > 0: | |
| logging.error(f"Skipped {skipped_line} unparsable lines") | |
| expected_dl = len(segment_downloads) | |
| actual_dl = sum(segment_downloads.values()) | |
| if actual_dl != expected_dl: | |
| logging.error( | |
| f"Some segments where downloaded twice. Total dl: {actual_dl}, distinct dl: {expected_dl}" | |
| ) | |
| def _tmp(file: Path) -> Path: | |
| tmp_dir = file.parent | |
| prefix = file.name.split(".", 1)[0] + "." | |
| suffix = ".tmp." + file.name[len(prefix) :] | |
| _, tmp_path = tempfile.mkstemp(dir=tmp_dir, prefix=prefix, suffix=suffix) | |
| return Path(tmp_path) | |
| class FileWriterWithTmp: | |
| def __init__(self, file: Path): | |
| self.file = file | |
| self.tmp_file = _tmp(file) | |
| # We don't want to make FileWriterWithTmp a ContextManager | |
| self.handle = open_write(self.tmp_file).__enter__() | |
| def write(self, data) -> int: | |
| return self.handle.write(data) | |
| def close(self, success: bool = False): | |
| self.handle.close() | |
| if success: | |
| self.tmp_file.rename(self.file) | |
| def transpose_file(outdir: Path, file: Path) -> None: | |
| sentinel_file = file.with_suffix(".transposed") | |
| if sentinel_file.exists(): | |
| return | |
| outputs: Dict[str, FileWriterWithTmp] = {} | |
| parser = get_typed_parser(Bitext) | |
| success = False | |
| try: | |
| for line in open_read(file): | |
| bt: Bitext = parser(line) | |
| lang_pair = bt.lang_pair | |
| if bt.lang_pair not in outputs: | |
| assert ( | |
| "/" in lang_pair | |
| ), f"Invalid lang pair '{lang_pair}' should be 'src-trg/src' or 'src-trg/trg'" | |
| (outdir / f"{lang_pair}").mkdir(exist_ok=True, parents=True) | |
| o = FileWriterWithTmp(outdir / f"{lang_pair}_{file.name}") | |
| outputs[lang_pair] = o | |
| simple_bt = SimpleBitext(bt.line_no, bt.score, bt.text) | |
| print(*simple_bt, sep="\t", file=outputs[lang_pair]) | |
| success = True | |
| finally: | |
| for o in outputs.values(): | |
| o.close(success) | |
| if success: | |
| sentinel_file.write_text("\n".join(str(o.file) for o in outputs.values())) | |
| # file.unlink() | |
| def sort_files(outdir: Path, lang_pair_dir: Path, lang: str) -> Path: | |
| out = outdir / lang_pair_dir.name / f"{lang}.txt" | |
| if out.exists(): | |
| return out | |
| files: List[Path] = [] | |
| for f in lang_pair_dir.iterdir(): | |
| if not f.suffix == ".gz": | |
| continue | |
| if f.name.split("_")[0] != lang: | |
| continue | |
| files.append(f) | |
| print(f"Found {len(files)} files for lang '{lang}' in {lang_pair_dir}: {files}") | |
| assert len(files) > 0 | |
| (outdir / lang_pair_dir.name).mkdir(exist_ok=True, parents=True) | |
| tmp_out = _tmp(out) | |
| unzipped_files = [] | |
| for f in files: | |
| subprocess.check_call(["gunzip", "-k", str(f)]) | |
| unzipped_files.append(str(f)[:-3]) | |
| sort_cmd = [ | |
| "sort", | |
| "-nk1", | |
| f"--parallel={SORT_PARALLEL}", | |
| f"--buffer-size={BUFFER_SIZE}", | |
| "--output", | |
| str(tmp_out), | |
| ] + unzipped_files | |
| subprocess.check_call(sort_cmd) | |
| tmp_out.rename(out) | |
| return out | |
| def finalize( | |
| outdir: Path = Path("data"), version: str = KNOWN_VERSIONS[0], pairs: Sequence[str] = [] | |
| ) -> None: | |
| """From the downloaded raw text files, extract the bitexts, sorted by language pair. | |
| Assumes 'dl' has been run with the same outdir and version before. | |
| - version: Specific version to download | |
| - outdir: Directory where the data should go. Files will be in {outdir}/{version}/bitext/ | |
| - pairs: List of language pairs you are interested in. Defaults to all. | |
| """ | |
| raw_dir = outdir / version / "raw" | |
| if not raw_dir.is_dir(): | |
| cmd = f"python {__file__} dl --outdir {outdir} --version {version}" | |
| assert raw_dir.is_dir(), f"Dir not found {raw_dir}. Did you run following command?\n{cmd}" | |
| raw_files = list(raw_dir.glob("*.gz")) | |
| split_dir = outdir / version / "split_by_lang" | |
| split_dir.mkdir(exist_ok=True, parents=True) | |
| tr = functools.partial(transpose_file, split_dir) | |
| with multiprocessing.Pool() as pool: | |
| pool.map(tr, raw_files) | |
| bitext_dir = outdir / version / "bitext" | |
| bitext_dir.mkdir(exist_ok=True, parents=True) | |
| if pairs: | |
| pair_dirs = [] | |
| for pair in pairs: | |
| assert ( | |
| len(pair.split("-")) == 2 | |
| ), f"Invalid pair '{pair}', should be 'src-trg'" | |
| pair_dir = split_dir / pair | |
| assert ( | |
| pair_dir.is_dir() | |
| ), f"Dir {pair_dir} not found for lang pair '{pair}'. Is the pair valid ?" | |
| pair_dirs.append(pair_dir) | |
| else: | |
| pair_dirs = [d for d in split_dir.iterdir() if d.is_dir()] | |
| for pair_dir in pair_dirs: | |
| src, trg = pair_dir.name.split("-") | |
| src_file = sort_files(bitext_dir, pair_dir, src) | |
| trg_file = sort_files(bitext_dir, pair_dir, trg) | |
| validate(src_file, trg_file) | |
| def validate(src_file: Path, trg_file: Path) -> None: | |
| """Checks that the segments in the given batch are valid.""" | |
| lines_src, lines_trg, found_pairs = 0, 0, 0 | |
| parser = get_typed_parser(SimpleBitext) | |
| with open(src_file) as src_f, open(trg_file) as trg_f: | |
| src_l = src_f.readline() | |
| trg_l = trg_f.readline() | |
| while src_l and trg_l: | |
| src: SimpleBitext = parser(src_l) | |
| trg: SimpleBitext = parser(trg_l) | |
| if src.line_no <= trg.line_no: | |
| lines_src += 1 | |
| src_l = src_f.readline() | |
| if trg.line_no <= src.line_no: | |
| lines_trg += 1 | |
| trg_l = trg_f.readline() | |
| if trg.line_no == src.line_no: | |
| found_pairs += 1 | |
| if found_pairs == lines_src and found_pairs == lines_trg: | |
| logging.info( | |
| f"Validated {src_file} and {trg_file}. Found {found_pairs} bitexts." | |
| ) | |
| else: | |
| logging.error( | |
| f"Validated {src_file} and {trg_file}. " | |
| f"Found {found_pairs} bitexts, from {lines_src} in {src_file} and {lines_trg} in {trg_file}" | |
| ) | |
| if __name__ == "__main__": | |
| import func_argparse | |
| func_argparse.main(dl, finalize) | |