diff --git a/cc_pseudo_crawl/language_annotation/python_scripts/detect_html_lang_attrib.py b/cc_pseudo_crawl/language_annotation/python_scripts/detect_html_lang_attrib.py index e69de29b..8d3cf9ac 100644 --- a/cc_pseudo_crawl/language_annotation/python_scripts/detect_html_lang_attrib.py +++ b/cc_pseudo_crawl/language_annotation/python_scripts/detect_html_lang_attrib.py @@ -0,0 +1,105 @@ +import os +import logging +import subprocess +import sys +from argparse import ArgumentParser +from pathlib import Path + +import datasets +import lxml.html +from datasets import config, load_from_disk +from datasets.utils.logging import set_verbosity_info + +set_verbosity_info() +logger = logging.getLogger(__name__) + +# For `soup.decode_content` that can hit the limit +sys.setrecursionlimit(10000) + + +def get_args(): + parser = ArgumentParser() + parser.add_argument( + "--dataset-path", + type=str, + required=True, + help="path to the parquet dataset folder", + ) + parser.add_argument("--save-path", type=str, help="Where to save the datasets.") + parser.add_argument("--use-datasets-caching", action="store_true") + parser.add_argument( + "--num-proc", type=int, default=1, help="Number of procs use for preprocessing." + ) + parser.add_argument( + "--num-examples", + type=int, + default=None, + help="Optional argument to select a subset (used for debugging purposes). Example `10`.", + ) + args = parser.parse_args() + + return args + + +def main(): + # Setup logging + logging.basicConfig( + format="%(asctime)s - %(levelname)s - %(name)s - %(message)s", + datefmt="%m/%d/%Y %H:%M:%S", + level=logging.INFO, + ) + args = get_args() + logger.info( + f"** The job is runned with the following arguments: **\n{args}\n **** " + ) + + if os.path.isdir(args.save_path): + logger.info(f"Seed id {args.save_path.split('/')[-1]} already processed") + return + + if not args.use_datasets_caching: + datasets.set_caching_enabled(False) + else: + logger.info( + f"the datasets results will be cached at {config.HF_DATASETS_CACHE}." + ) + + ds = load_from_disk(args.dataset_path) + logger.info(f"the dataset is {ds}") + + if args.num_examples: + ds = ds.select([i for i in range(args.num_examples)]) + + def detect_lang(example): + if example["text"] is None or len(example["text"]) == 0: + example["html_lang_attr"] = None + else: + root = lxml.html.fromstring(example["html_str"]) + root_lang = root.attrib.get("lang") + example["html_lang_attr"] = root_lang + return example + + ds = ds.map( + detect_lang, + batched=False, + num_proc=args.num_proc, + ) + + if args.save_path: + save_path = Path(args.save_path) + else: + save_path = Path(args.dataset_path) + + logger.info( + f"Lang attribute detected for {len([e for e in ds['train']['html_lang_attr'] if e is not None])} rows out of {len(ds['train'])} rows." + ) + + save_path_tmp = f"{str(save_path.absolute())}.tmp" + logger.info(f"Saving the dataset at {save_path_tmp}") + ds.save_to_disk(save_path_tmp) + logger.info(f"Moving the saved dataset to {str(save_path.absolute())}") + subprocess.run(["mv", save_path_tmp, str(save_path.absolute())]) + + +if __name__ == "__main__": + main() diff --git a/cc_pseudo_crawl/language_annotation/slurm_scripts/02_detect_html_lang_attrib.slurm b/cc_pseudo_crawl/language_annotation/slurm_scripts/02_detect_html_lang_attrib.slurm index e69de29b..02f250da 100644 --- a/cc_pseudo_crawl/language_annotation/slurm_scripts/02_detect_html_lang_attrib.slurm +++ b/cc_pseudo_crawl/language_annotation/slurm_scripts/02_detect_html_lang_attrib.slurm @@ -0,0 +1,35 @@ +#!/bin/bash +#SBATCH --job-name=pseudo_crawl_extract_lang_tag +#SBATCH --nodes=1 +#SBATCH --ntasks-per-node=1 # crucial - only 1 task per dist per node! +#SBATCH --cpus-per-task=40 # number of cores per tasks +#SBATCH --hint=nomultithread # we get physical cores not logical +#SBATCH --partition=cpu_p1 +#SBATCH --time 20:00:00 # maximum execution time (HH:MM:SS) +#SBATCH --output=/gpfsscratch/rech/six/commun/pseudo_crawl/seeds_batch_2/logs/extract_lang_tag_V5/%x-%j.out # output file name #TODO change path if necessary +#SBATCH --array=35,341,297 +#SBATCH --account=six@cpu + +set -x -e + +source $six_ALL_CCFRWORK/start-prod +conda activate thomas_data_tooling + +DATA_TOOLING_REPO=$WORK/repos/sync_data_tooling/data_tooling + +DATASET_PATH=/gpfsscratch/rech/six/urd43gx/crawl/shards/shard_"$SLURM_ARRAY_TASK_ID" +SAVE_DATASET_DIR=$six_ALL_CCFRSCRATCH/pseudo_crawl/fasttext_annotation/seeds_batch_1/datasets-lang-annotation/bigscience-catalogue-data +SAVE_DATASET_PATH=$SAVE_DATASET_DIR/seed_id="$SLURM_ARRAY_TASK_ID" +echo $DATASET_PATH +pushd $DATA_TOOLING_REPO + +mkdir -p $SAVE_DATASET_DIR + +export HF_DATASETS_OFFLINE=1 +export HF_DATASETS_CACHE=$SCRATCH/to_delete + +python -m cc_pseudo_crawl.language_annotation.python_scripts.detect_html_lang_attrib \ + --dataset-path $DATASET_PATH \ + --num-proc 40 \ + --save-path $SAVE_DATASET_PATH \ + --use-datasets-caching