Revamp cc2dataset warc text extraction

.gitignore

@@ -11,4 +11,7 @@ wandb
 **/*.parquet
 **/*.tar.gz
 .mypy_cache
-.pytest_cache
+.pytest_cache
+assets/*.bin
+build/*
+assets/*.out

assets/spark_slurm.sh

@@ -0,0 +1,18 @@
+#!/bin/bash
+#SBATCH --partition=spark
+#SBATCH --job-name=bild-cc-pile
+#SBATCH --nodes 5
+#SBATCH --ntasks-per-node 1
+#SBATCH --cpus-per-task=48
+#SBATCH --mem=0 # 0 means use all available memory (in MB)
+#SBATCH --output=%x_%j.out
+#SBATCH --comment laion
+#SBATCH --exclusive
+
+# wget https://dlcdn.apache.org/spark/spark-3.3.1/spark-3.3.1-bin-hadoop3.tgz && tar xf spark-3.3.1-bin-hadoop3.tgz
+
+# step 1: get environment variables
+# step 2: setup rank 0 to be the master and start the indexing python file (the last two operations happen in parallel)
+# step 3: start workers on all nodes
+
+srun --comment laion bash worker_spark_on_slurm.sh

assets/worker_spark_on_slurm.sh

@@ -0,0 +1,50 @@
+#!/bin/bash
+#
+# get environment variables
+GLOBAL_RANK=$SLURM_PROCID
+CPUS=$SLURM_CPUS_PER_TASK
+MEM=$SLURM_MEM_PER_NODE # seems to be in MB
+MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1)
+LOCAL_IP=$(hostname -I | awk '{print $1}')
+
+# set some environment variables for the indexing script
+export MASTER_ADDR=$MASTER_ADDR
+export MEMORY=$MEM
+export CPUS=$CPUS
+export SPARK_LOCAL_IP=$LOCAL_IP
+
+# setup the master node
+if [ $GLOBAL_RANK == 0 ]
+then
+    # print out some info
+    echo -e "MASTER ADDR: $MASTER_ADDR\tGLOBAL RANK: $GLOBAL_RANK\tCPUS PER TASK: $CPUS\tMEM PER NODE: $MEM"
+
+    # then start the spark master node in the background
+    ./spark-3.3.1-bin-hadoop3/sbin/start-master.sh -p 7079 -h $LOCAL_IP
+
+fi
+
+sleep 10
+
+# then start the spark worker node in the background
+MEM_IN_GB=$(($MEM / 1000))
+# concat a "G" to the end of the memory string
+MEM_IN_GB="$MEM_IN_GB"G
+echo "MEM IN GB: $MEM_IN_GB"
+
+export SPARK_LOCAL_DIRS=/scratch/local
+export SPARK_WORKER_DIR=/scratch/work
+
+./spark-3.3.1-bin-hadoop3/sbin/start-worker.sh -c $CPUS -m $MEM_IN_GB "spark://$MASTER_ADDR:7079"
+echo "Hello from worker $GLOBAL_RANK"
+
+
+sleep 10
+
+if [ $GLOBAL_RANK == 0 ]
+then
+    # then start some script
+    echo "hi"
+fi
+
+sleep 10000

cc2dataset/__init__.py

@@ -1,3 +1,6 @@
 """cc2dataset"""
 
-from cc2dataset.main import process_wat, cc2dataset
+from cc2dataset.main import cc2dataset
+from cc2dataset.warc_utils import process_warc, extract_documents_from_warc
+from cc2dataset.wat_utils import process_wat
+

cc2dataset/index_utils.py

@@ -0,0 +1,61 @@
+import fsspec
+from multiprocessing.pool import ThreadPool
+import random
+
+def get_cc_links(source_cc_protocol,ccfile):
+    """Get cc wat links"""
+    if source_cc_protocol == "s3":
+        fs, p = fsspec.core.url_to_fs("s3://commoncrawl/crawl-data/")
+        if ccfile == "warc":
+            links = ["s3://" + e for e in fs.glob(p + "/*/warc.paths.gz")]
+        elif ccfile == "wat":
+            links = ["s3://" + e for e in fs.glob(p + "/*/wat.paths.gz")]
+        elif ccfile == "wet":
+            links = ["s3://" + e for e in fs.glob(p + "/*/wet.paths.gz")]
+        else:
+            raise ValueError(f"Unknown ccfile: {ccfile}")
+        return links
+    elif source_cc_protocol == "http":
+        fs, p = fsspec.core.url_to_fs("https://commoncrawl.org/the-data/get-started/")
+        a = fs.open(p).read()
+        l = a.splitlines()
+        l = [e.decode("utf8").replace("[WARC] ", "") for e in l]
+        l = [e for e in l if "<li>s3://commoncrawl/crawl-data/" in e]
+        l = [
+            e.split(" ")[0].replace("<li>s3://commoncrawl/", "https://data.commoncrawl.org/").replace("<wbr>", "")
+            for e in l
+        ]
+        if ccfile == "warc":
+            l = [(e + "/warc.paths.gz").replace("//warc", "/warc") for e in l]
+        elif ccfile == "wat":
+            l = [(e + "/wat.paths.gz").replace("//wat", "/wat") for e in l]
+        elif ccfile == "wet":
+            l = [(e + "/wet.paths.gz").replace("//wet", "/wet") for e in l]
+        return l
+    else:
+        raise ValueError(f"Unknown protocol {source_cc_protocol}")
+
+
+def read_index_file(file_index):
+    with fsspec.open(file_index, "rb", compression="gzip") as f:
+        crawlfiles = [a.decode("utf8").strip() for a in f.readlines()]
+    return crawlfiles
+
+
+def read_index_files(crawl_count, file_count, source_cc_protocol, ccfile, crawl_index_list, shuffle):
+    """Read all wat index files"""
+    cc_links = get_cc_links(source_cc_protocol,ccfile)
+    if crawl_index_list is not None:
+        cc_links = [cc_links[i] for i in crawl_index_list]
+    if crawl_count is not None:
+        cc_links = cc_links[-crawl_count:]  # pylint: disable=invalid-unary-operand-type
+    all_files = []
+    with ThreadPool(16) as pool:
+        for wats in pool.imap_unordered(read_index_file, cc_links):
+            all_files.extend(wats)
+    if file_count is not None:
+        all_files = random.choices(all_files, k=file_count)
+    if shuffle:
+        # shuffle to increase duplication over each part hence reduce size of each part after duplication
+        random.shuffle(all_files)
+    return all_files

cc2dataset/lang_utils.py

@@ -0,0 +1,86 @@
+import fasttext
+from pathlib import Path
+
+def load_fasttext_model(path_fasttext_model):
+    print("Loading fasttext model from ", path_fasttext_model)
+    return fasttext.load_model(path_fasttext_model)
+
+def get_fasttext_info(line, model_lang_id):
+    """The line should be in lower case and without \n in it."""
+    pred = model_lang_id.predict(line)
+    lang_pred_fasttext_id = pred[0][0].replace("__label__", "")
+    score_pred = pred[1][0]
+    return lang_pred_fasttext_id, score_pred
+
+
+
+class LangDetection:
+    #adapted from https://github.com/bigcode-project/bigcode-analysis/blob/main/data_analysis/python_data_analysis/nl_language_identification/language_identifier.py
+    def __init__(self,model_dump_path:str) -> None:
+        self.lang_model_path : str = model_dump_path
+        self.model = load_fasttext_model(self.lang_model_path)
+
+
+    def detect(self, text: str) -> str:
+        """
+        Detects the language of the text
+        args:
+            text (str) : Text to detect the language
+
+        returns:
+            language (str) : Predicted Language of the text
+            score_pred (str) : confidence of the prediction
+
+        """
+        text = text.lower()
+
+        fasttext_pred = get_fasttext_info(
+            text, self.model
+        )
+        return fasttext_pred[0], fasttext_pred[1]
+
+
+import re
+
+
+class LicensePattern:
+    cc_pattern = re.compile("http[s]?://creativecommons\\.org/licenses/(by|by-sa|by-nd|by-nc|by-nc-sa|by-nc-nd|publicdomain)[\"/ >]")
+
+def detect_licence(html:str):
+    """
+    Given a HTML string, this function detects the licence of the page.
+    It returns a string with the licence name, or NO-LICENCE-FOUND if no licence is found.
+    """
+    license_attribute_pattern = re.compile(LicensePattern.cc_pattern)
+
+    # storing counts of all difference occurrences of link to CC
+    multiple_occurrences_map = {}
+
+    # add all of them to the list
+    for match in license_attribute_pattern.finditer(html):
+        licence = match.group(1)
+
+        # add entry
+        if licence not in multiple_occurrences_map:
+            multiple_occurrences_map[licence] = 0
+
+        # and increase count
+        multiple_occurrences_map[licence] += 1
+
+    # no licence found
+    if not multiple_occurrences_map:
+        return "no-licence-found"
+
+    # only one link found or if multiple links found but the same type
+    if len(multiple_occurrences_map) == 1:
+        return list(multiple_occurrences_map.keys())[0]
+
+    # if multiple different links found, we return a general CC-UNSPECIFIED
+    return "cc-unspecified"
+
+
+if __name__ == "__main__":
+    lang_detector = LangDetection("lid_model_dump/lid.176.bin")
+    print(lang_detector.detect("Das ist ein Test."))
+
+#output : ('de', 1.000038981437683)