diff --git a/.gitignore b/.gitignore
index 18ce86f..f99b29d 100644
--- a/.gitignore
+++ b/.gitignore
@@ -11,4 +11,7 @@ wandb
**/*.parquet
**/*.tar.gz
.mypy_cache
-.pytest_cache
\ No newline at end of file
+.pytest_cache
+assets/*.bin
+build/*
+assets/*.out
\ No newline at end of file
diff --git a/assets/spark_slurm.sh b/assets/spark_slurm.sh
new file mode 100644
index 0000000..ec15aac
--- /dev/null
+++ b/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
diff --git a/assets/worker_spark_on_slurm.sh b/assets/worker_spark_on_slurm.sh
new file mode 100644
index 0000000..0dc4681
--- /dev/null
+++ b/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
diff --git a/cc2dataset/__init__.py b/cc2dataset/__init__.py
index 9c1bcc8..6b238a1 100644
--- a/cc2dataset/__init__.py
+++ b/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
+
diff --git a/cc2dataset/index_utils.py b/cc2dataset/index_utils.py
new file mode 100644
index 0000000..4359ed3
--- /dev/null
+++ b/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 "s3://commoncrawl/crawl-data/" in e]
+ l = [
+ e.split(" ")[0].replace("s3://commoncrawl/", "https://data.commoncrawl.org/").replace("", "")
+ 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
\ No newline at end of file
diff --git a/cc2dataset/lang_utils.py b/cc2dataset/lang_utils.py
new file mode 100644
index 0000000..be1bc93
--- /dev/null
+++ b/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)
\ No newline at end of file
diff --git a/cc2dataset/main.py b/cc2dataset/main.py
index efbd6bb..5821384 100644
--- a/cc2dataset/main.py
+++ b/cc2dataset/main.py
@@ -8,7 +8,7 @@
from loguru import logger
import hashlib
import datetime
-from multiprocessing.pool import ThreadPool
+
from pyspark import SparkContext
from pyspark.sql.functions import rand
from pyspark.sql import SparkSession
@@ -16,200 +16,11 @@
import math
import time
from .spark_session_builder import build_spark_session
-from io import BytesIO
-
-
-def valid_video_link(link):
- valid_http = link.get("url", "").startswith("http")
- valid_video = any(
- link.get("url", "").endswith(ext) for ext in [".avi", ".mp4", ".mkv", ".webm", ".mov", ".mpg", ".mpeg", ".m4v"]
- )
- return valid_http and valid_video
-
-
-def extract_video_from_links(links):
- filtered_links = [{"url": link["url"], "alt": link.get("text", "")} for link in links if valid_video_link(link)]
- return filtered_links
-
-
-text_extensions = set(
- [
- "pdf",
- "epub",
- "djvu",
- "mobi",
- "doc",
- "docx",
- "rtf",
- "txt",
- "odt",
- "ppt",
- "pptx",
- "pages",
- "keynote",
- "wps",
- "md",
- ]
-)
-
-
-def valid_text_link(link):
- if not link.get("url", "").startswith("http"):
- return False
- splits = link.get("url", "").split(".")
- if len(splits) < 2:
- return False
- if splits[-1] not in text_extensions:
- return False
- return True
-
-
-def extract_text_from_links(links):
- filtered_links = [{"url": link["url"], "alt": link.get("text", "")} for link in links if valid_text_link(link)]
- return filtered_links
-
-
-def valid_audio_link(link):
- valid_http = link.get("url", "").startswith("http")
- valid_audio = any(link.get("url", "").endswith(ext) for ext in [".ogg", ".wav", ".mp3", ".flac", ".m4a"])
- return valid_http and valid_audio
-
-
-def extract_audio_from_links(links):
- """Extract image from links"""
- filtered_links = [{"url": link["url"], "alt": link.get("text", "")} for link in links if valid_audio_link(link)]
- return filtered_links
-
-
-def valid_image_link(link):
- valid_path = link.get("path", "") == "IMG@/src"
- valid_alt = len(link.get("alt", "")) > 0
- valid_http = link.get("url", "").startswith("http")
- return valid_path and valid_http and valid_alt
-
-
-def extract_image_from_links(links):
- """Extract image from links"""
- filtered_links = [{"url": link["url"], "alt": link["alt"]} for link in links if valid_image_link(link)]
- return filtered_links
-
-
-def extract_documents_from_links(links, document_type):
- """Extract documents from links ; this function returns a list of dict {"alt": ..., "url": ...}"""
-
- if document_type == "image":
- return extract_image_from_links(links)
- elif document_type == "audio":
- return extract_audio_from_links(links)
- elif document_type == "text":
- return extract_text_from_links(links)
- elif document_type == "video":
- return extract_video_from_links(links)
- else:
- raise ValueError(f"Unknown document type {document_type}")
-
-
-def extract_documents_from_wat(stream, document_type):
- """Extract document from stream"""
- all_links = []
- try:
- for record in ArchiveIterator(stream, record_types=WarcRecordType.metadata, parse_http=False):
- try:
- record_data = simdjson.load(record.reader) # type: ignore
- except: # pylint: disable=bare-except
- logger.info("A shard record failed")
- continue
- envelope = record_data["Envelope"]
- payload = envelope["Payload-Metadata"]
- if "HTTP-Response-Metadata" not in payload:
- continue
- http_resp = payload["HTTP-Response-Metadata"]
- if "HTML-Metadata" not in http_resp:
- continue
- metadata = http_resp["HTML-Metadata"]
- if "Links" not in metadata:
- continue
-
- links = metadata["Links"]
-
- filtered_links = extract_documents_from_links(links, document_type)
- for link in filtered_links:
- link["uid"] = str(hashlib.md5((link["alt"] + link["url"]).encode()).hexdigest())
- all_links.extend(filtered_links)
- except Exception as e: # pylint: disable=broad-except
- logger.info(e)
- logger.info("A shard failed to parse")
- return []
-
- return all_links
-
-
-def process_wat(path, document_type):
- """Process a single wat file"""
- begin_read = timer()
- with fsspec.open(path, "rb") as f:
- for i in range(10):
- try:
- tf = BytesIO(f.read())
- break
- except Exception as ex: # pylint: disable=broad-except
- if i == 9:
- logger.info("failed 10 times, skipping ", path)
- return
- logger.info(ex)
- logger.info(f"retrying reading {i}/10")
- time.sleep(1)
-
- for e in extract_documents_from_wat(tf, document_type):
- yield (e["uid"], e["url"], e["alt"])
- end_read = timer()
- tot_read_time = end_read - begin_read
- logger.info(f"Took {tot_read_time} to parse")
-
-
-def get_cc_wat_links(source_cc_protocol):
- """Get cc wat links"""
- if source_cc_protocol == "s3":
- fs, p = fsspec.core.url_to_fs("s3://commoncrawl/crawl-data/")
- links = ["s3://" + e for e in fs.glob(p + "/*/wat.paths.gz")]
- 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 "s3://commoncrawl/crawl-data/" in e]
- l = [
- e.split(" ")[0].replace("s3://commoncrawl/", "https://data.commoncrawl.org/").replace("", "")
- for e in l
- ]
- l = [(e + "/wat.paths.gz").replace("//wat", "/wat") for e in l]
- return l
- else:
- raise ValueError(f"Unknown protocol {source_cc_protocol}")
-
-
-def read_wat_index_file(wat_index):
- with fsspec.open(wat_index, "rb", compression="gzip") as f:
- wats = [a.decode("utf8").strip() for a in f.readlines()]
- return wats
-
-
-def read_wat_index_files(shard_count, wat_count, source_cc_protocol):
- """Read all wat index files"""
- cc_wat_links = get_cc_wat_links(source_cc_protocol)
- if shard_count is not None:
- cc_wat_links = cc_wat_links[-shard_count:] # pylint: disable=invalid-unary-operand-type
- all_wats = []
- with ThreadPool(16) as pool:
- for wats in pool.imap_unordered(read_wat_index_file, cc_wat_links):
- all_wats.extend(wats)
- if wat_count is not None:
- all_wats = random.choices(all_wats, k=wat_count)
- else:
- # shuffle to increase duplication over each part hence reduce size of each part after duplication
- random.shuffle(all_wats)
- return all_wats
+from .wat_utils import process_wat
+from .warc_utils import process_warc
+from .index_utils import read_index_files
+
+
def deduplicate_repartition_count(df, output_path, wat_count, spark, shuffle=False):
@@ -227,25 +38,41 @@ def deduplicate_repartition_count(df, output_path, wat_count, spark, shuffle=Fal
logger.info(f"Size: {df.count()}")
-def process_one_part(output_path, wat_index_files, build_spark, shuffle, document_type, source_cc_protocol):
+def process_one_part(output_path, cc_index_files, build_spark, shuffle, document_type, source_cc_protocol,ccfile):
"""Process one part"""
spark = build_spark()
sc = SparkContext.getOrCreate()
- wat_count = len(wat_index_files)
- wat_rdd = sc.parallelize(wat_index_files, wat_count)
+ ccfile_count = len(cc_index_files)
+ wat_rdd = sc.parallelize(cc_index_files, ccfile_count)
+
if source_cc_protocol == "s3":
prefix = "s3://commoncrawl/"
elif source_cc_protocol == "http":
prefix = "https://data.commoncrawl.org/"
- def extract(x):
- x = list(x)
- yield from process_wat(prefix + x[0], document_type)
+ if ccfile == "warc":
+ def extract(x):
+ x = list(x)
+ yield from process_warc(prefix + x[0])
+
+ elif ccfile == "wat":
+ def extract(x):
+ x = list(x)
+ yield from process_wat(prefix + x[0], document_type)
+
+ elif ccfile == "wet":
+ def extract(x):
+ x = list(x)
+ yield from process_wet(prefix + x[0])
+ else:
+ raise ValueError(f"Unknown ccfile: {ccfile}")
+
output = wat_rdd.mapPartitions(extract)
- df = output.toDF(["uid", "url", "alt"])
+ # e["uid"], e["url"], e["text"],e['lang'],e['license'],e['perplexity']
+ df = output.toDF(["uid", "url", "text","lang","license","fwords","cratio","crep","swords","wnum","perplexity"])
- deduplicate_repartition_count(df, output_path, wat_count, spark, shuffle)
+ deduplicate_repartition_count(df, output_path, ccfile_count, spark, shuffle)
def get_last_successful_part(output_path):
@@ -258,7 +85,7 @@ def get_last_successful_part(output_path):
def process_multi_part(
- output_path, wat_index_files, build_spark, multipart, shuffle, resume, document_type, source_cc_protocol
+ output_path, cc_index_files, build_spark, multipart, shuffle, resume, document_type, source_cc_protocol,ccfile
):
"""Process multi part"""
if resume:
@@ -266,16 +93,16 @@ def process_multi_part(
else:
start_part = 0
- wat_count = len(wat_index_files)
- wat_per_part = math.ceil(wat_count / multipart)
+ ccfile_count = len(cc_index_files)
+ ccfile_per_part = math.ceil(ccfile_count / multipart)
part_paths = []
for i in range(start_part, multipart):
- start = i * wat_per_part
- end = (i + 1) * wat_per_part
+ start = i * ccfile_per_part
+ end = (i + 1) * ccfile_per_part
part_path = f"{output_path}/part_{i}"
part_paths.append(part_path)
logger.info(f"Processing part {i} from {start} to {end} into {part_path}")
- process_one_part(part_path, wat_index_files[start:end], build_spark, False, document_type, source_cc_protocol)
+ process_one_part(part_path, cc_index_files[start:end], build_spark, False, document_type, source_cc_protocol,ccfile)
spark = build_spark()
logger.info("Merging parts")
@@ -287,7 +114,7 @@ def process_multi_part(
else:
df = df.union(spark.read.parquet(part_path))
- deduplicate_repartition_count(df, output_path + "/merged", wat_count, spark, shuffle)
+ deduplicate_repartition_count(df, output_path + "/merged", ccfile_count, spark, shuffle)
def get_date_str():
@@ -296,8 +123,8 @@ def get_date_str():
def cc2dataset(
output_path,
- wat_index_count=1,
- wat_count=100,
+ crawl_index_count=1,
+ files_count=100,
master="local",
num_cores=128,
mem_gb=256,
@@ -307,6 +134,8 @@ def cc2dataset(
spark_builder=None,
document_type="image",
source_cc_protocol="s3",
+ ccfile="wat",
+ crawl_index_list=None,
):
"""Convert common crawl to image caption set"""
@@ -332,20 +161,19 @@ def build_spark():
return spark_builder()
if resume is None:
- wat_index_files = read_wat_index_files(wat_index_count, wat_count, source_cc_protocol)
- # write wat index files to disk in output_path with fsspec
- with fsspec.open(f"{output_path}/wat_index_files.txt", "w", encoding="utf8") as f:
- f.write("\n".join(wat_index_files))
+ cc_index_files = read_index_files(crawl_index_count, files_count, source_cc_protocol, ccfile, crawl_index_list, shuffle )
+ # write ccfile index files to disk in output_path with fsspec
+ with fsspec.open(f"{output_path}/crawl_index_files.txt", "w", encoding="utf8") as f:
+ f.write("\n".join(cc_index_files))
else:
- with fsspec.open(f"{output_path}/wat_index_files.txt", "r", encoding="utf8") as f:
- wat_index_files = f.read().splitlines()
+ with fsspec.open(f"{output_path}/crawl_index_files.txt", "r", encoding="utf8") as f:
+ cc_index_files = f.read().splitlines()
if multipart is None:
- process_one_part(output_path, wat_index_files, build_spark, shuffle, document_type, source_cc_protocol)
+ process_one_part(output_path, cc_index_files, build_spark, shuffle, document_type, source_cc_protocol,ccfile)
else:
process_multi_part(
- output_path, wat_index_files, build_spark, multipart, shuffle, resume, document_type, source_cc_protocol
- )
+ output_path, cc_index_files, build_spark, multipart, shuffle, resume, document_type, source_cc_protocol, ccfile)
def main():
diff --git a/cc2dataset/warc_utils.py b/cc2dataset/warc_utils.py
new file mode 100644
index 0000000..5e0240e
--- /dev/null
+++ b/cc2dataset/warc_utils.py
@@ -0,0 +1,113 @@
+from fastwarc import ArchiveIterator
+import hashlib
+import fsspec
+from resiliparse.parse import detect_encoding
+from resiliparse.parse.html import HTMLTree
+from resiliparse.extract.html2text import extract_plain_text
+from io import BytesIO
+from timeit import default_timer as timer
+from .lang_utils import LangDetection,detect_licence
+from loguru import logger
+from squeakily.filter import (
+ check_char_repetition,
+ check_flagged_words,
+ check_stop_word_ratio,
+ check_compression_ratio,
+ check_word_number,
+ check_perplexity,
+)
+
+from squeakily.helpers import KenlmModel
+
+model = KenlmModel.from_pretrained(
+ model_dataset="wikipedia",
+ language="en",
+ lower_case=True,
+ remove_accents=True,
+ normalize_numbers=True,
+ punctuation=1,
+)
+
+#check_perplexity(low_test_str, model=model,dry_run=True)
+
+def extract_documents_from_warc(stream):
+ """Extract document from stream"""
+ all_extend = []
+ # download https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin and store to a path
+ lang_model_path = "/home/harrysaini/cc2dataset/assets/lid.176.bin"
+ detector = LangDetection(lang_model_path)
+ try:
+ for idx, record in enumerate(ArchiveIterator(stream, max_content_length=4 * 1024**2)):
+ try:
+ if record.headers is None:
+ continue
+ if record.http_headers is None:
+ continue
+ if (
+ record.headers["WARC-Type"] == "response"
+ and record.content_length >= 128
+ ):
+ content_type = str(record.http_content_type).lower()
+ if content_type.startswith("text/html"):
+ url = str(record.headers["WARC-Target-URI"])
+ html_bytes = record.reader.read()
+ encoding = detect_encoding(html_bytes)
+ licence = detect_licence(str(html_bytes))
+ tree = HTMLTree.parse_from_bytes(html_bytes, encoding)
+
+ for ele in tree.body.get_elements_by_tag_name("nav"):
+ ele.parent.remove_child(ele)
+
+
+ text = extract_plain_text(tree, preserve_formatting=False,
+ main_content=False, list_bullets=False,
+ alt_texts=True, links=False,
+ form_fields=False, noscript=False)
+ text = text.replace("\n", " ").replace("\t", " ").replace("\r", " ")
+
+ cre=dict()
+ cre["text"] = text
+ cre["url"] = url
+ cre["uid"] = str(hashlib.md5((text+url).encode()).hexdigest())
+ cre["lang"],_ = detector.detect(text)
+ cre['license'] = licence
+ cre['fwords'] = check_flagged_words(text,dry_run=True)
+ cre['cratio'] = check_compression_ratio(text,dry_run=True)
+ cre['crep'] = check_char_repetition(text,dry_run=True)
+ cre['swords'] = check_stop_word_ratio(text,dry_run=True)
+ cre['wnum'] = check_word_number(text,dry_run=True)
+ cre['perplexity'] = check_perplexity(text, model=model,dry_run=True)
+
+ #cre.update(permodel(text, prefix="perplexity/"))
+ all_extend.append(cre)
+ except:
+ continue
+
+ except Exception as e: # pylint: disable=broad-except
+ logger.info(e)
+ logger.info("A shard failed to parse")
+ return []
+
+ return all_extend
+
+def process_warc(path):
+ """Process a single warc file"""
+ begin_read = timer()
+ with fsspec.open(path, "rb") as f:
+ for i in range(10):
+ try:
+ tf = BytesIO(f.read())
+ break
+ except Exception as ex: # pylint: disable=broad-except
+ if i == 9:
+ logger.info("failed 10 times, skipping ", path)
+ return
+ logger.info(ex)
+ logger.info(f"retrying reading {i}/10")
+ time.sleep(1)
+
+ for e in extract_documents_from_warc(tf):
+ yield (e["uid"], e["url"], e["text"], e["lang"], e["license"])
+ end_read = timer()
+ tot_read_time = end_read - begin_read
+ logger.info(f"Took {tot_read_time} to parse")
diff --git a/cc2dataset/wat_utils.py b/cc2dataset/wat_utils.py
new file mode 100644
index 0000000..10dd34d
--- /dev/null
+++ b/cc2dataset/wat_utils.py
@@ -0,0 +1,154 @@
+import hashlib
+from fastwarc.warc import ArchiveIterator, WarcRecordType
+import simdjson
+import fsspec
+from timeit import default_timer as timer
+from loguru import logger
+from io import BytesIO
+
+def valid_video_link(link):
+ valid_http = link.get("url", "").startswith("http")
+ valid_video = any(
+ link.get("url", "").endswith(ext) for ext in [".avi", ".mp4", ".mkv", ".webm", ".mov", ".mpg", ".mpeg", ".m4v"]
+ )
+ return valid_http and valid_video
+
+
+def extract_video_from_links(links):
+ filtered_links = [{"url": link["url"], "alt": link.get("text", "")} for link in links if valid_video_link(link)]
+ return filtered_links
+
+
+text_extensions = set(
+ [
+ "pdf",
+ "epub",
+ "djvu",
+ "mobi",
+ "doc",
+ "docx",
+ "rtf",
+ "txt",
+ "odt",
+ "ppt",
+ "pptx",
+ "pages",
+ "keynote",
+ "wps",
+ "md",
+ ]
+)
+
+
+def valid_text_link(link):
+ if not link.get("url", "").startswith("http"):
+ return False
+ splits = link.get("url", "").split(".")
+ if len(splits) < 2:
+ return False
+ if splits[-1] not in text_extensions:
+ return False
+ return True
+
+
+def extract_text_from_links(links):
+ filtered_links = [{"url": link["url"], "alt": link.get("text", "")} for link in links if valid_text_link(link)]
+ return filtered_links
+
+
+def valid_audio_link(link):
+ valid_http = link.get("url", "").startswith("http")
+ valid_audio = any(link.get("url", "").endswith(ext) for ext in [".ogg", ".wav", ".mp3", ".flac", ".m4a"])
+ return valid_http and valid_audio
+
+
+def extract_audio_from_links(links):
+ """Extract image from links"""
+ filtered_links = [{"url": link["url"], "alt": link.get("text", "")} for link in links if valid_audio_link(link)]
+ return filtered_links
+
+
+def valid_image_link(link):
+ valid_path = link.get("path", "") == "IMG@/src"
+ valid_alt = len(link.get("alt", "")) > 0
+ valid_http = link.get("url", "").startswith("http")
+ return valid_path and valid_http and valid_alt
+
+
+def extract_image_from_links(links):
+ """Extract image from links"""
+ filtered_links = [{"url": link["url"], "alt": link["alt"]} for link in links if valid_image_link(link)]
+ return filtered_links
+
+
+def extract_documents_from_links(links, document_type):
+ """Extract documents from links ; this function returns a list of dict {"alt": ..., "url": ...}"""
+
+ if document_type == "image":
+ return extract_image_from_links(links)
+ elif document_type == "audio":
+ return extract_audio_from_links(links)
+ elif document_type == "text":
+ return extract_text_from_links(links)
+ elif document_type == "video":
+ return extract_video_from_links(links)
+ else:
+ raise ValueError(f"Unknown document type {document_type}")
+
+
+def extract_documents_from_wat(stream, document_type):
+ """Extract document from stream"""
+ all_links = []
+ try:
+ for record in ArchiveIterator(stream, record_types=WarcRecordType.metadata, parse_http=False):
+ try:
+ record_data = simdjson.load(record.reader) # type: ignore
+ except: # pylint: disable=bare-except
+ logger.info("A shard record failed")
+ continue
+ envelope = record_data["Envelope"]
+ payload = envelope["Payload-Metadata"]
+ if "HTTP-Response-Metadata" not in payload:
+ continue
+ http_resp = payload["HTTP-Response-Metadata"]
+ if "HTML-Metadata" not in http_resp:
+ continue
+ metadata = http_resp["HTML-Metadata"]
+ if "Links" not in metadata:
+ continue
+
+ links = metadata["Links"]
+
+ filtered_links = extract_documents_from_links(links, document_type)
+ for link in filtered_links:
+ link["uid"] = str(hashlib.md5((link["alt"] + link["url"]).encode()).hexdigest())
+ all_links.extend(filtered_links)
+ except Exception as e: # pylint: disable=broad-except
+ logger.info(e)
+ logger.info("A shard failed to parse")
+ return []
+
+ return all_links
+
+def process_wat(path, document_type):
+ """Process a single wat file"""
+ begin_read = timer()
+ with fsspec.open(path, "rb") as f:
+ for i in range(10):
+ try:
+ tf = BytesIO(f.read())
+ break
+ except Exception as ex: # pylint: disable=broad-except
+ if i == 9:
+ logger.info("failed 10 times, skipping ", path)
+ return
+ logger.info(ex)
+ logger.info(f"retrying reading {i}/10")
+ time.sleep(1)
+
+ for e in extract_documents_from_wat(tf, document_type):
+ yield (e["uid"], e["url"], e["alt"])
+ end_read = timer()
+ tot_read_time = end_read - begin_read
+ logger.info(f"Took {tot_read_time} to parse")
+
diff --git a/examples/single_warc.py b/examples/single_warc.py
new file mode 100644
index 0000000..5ded9a8
--- /dev/null
+++ b/examples/single_warc.py
@@ -0,0 +1,16 @@
+from cc2dataset import process_warc
+import os
+import pandas as pd
+
+if __name__ == "__main__":
+ from_s3 = True
+ wurl = "crawl-data/CC-MAIN-2022-33/segments/1659882570651.49/warc/CC-MAIN-20220807150925-20220807180925-00000.warc.gz"
+ if from_s3:
+ url = "s3://commoncrawl/" + wurl
+ else:
+ url = "https://data.commoncrawl.org/" + wurl
+
+ results = process_warc(url)
+ df = pd.DataFrame(results, columns=["uid", "url", "text",'lang','license'])
+ df.to_parquet(os.getcwd() + "/output.parquet")
+ print(df)
diff --git a/examples/single_warc_example.py b/examples/single_wat_example.py
similarity index 100%
rename from examples/single_warc_example.py
rename to examples/single_wat_example.py
diff --git a/requirements.txt b/requirements.txt
index 6f00d62..c9dd112 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -7,4 +7,10 @@ pyarrow
fastwarc
s3fs
fire
-requests
\ No newline at end of file
+requests
+squeakily
+spacy
+spacy_fastlang
+resiliparse
+pip install https://github.com/kpu/kenlm/archive/master.zip sentencepiece
+pip install git+https://github.com/CarperAI/squeakily.git
\ No newline at end of file
diff --git a/tests/test_kenlm.py b/tests/test_kenlm.py
new file mode 100644
index 0000000..6980bb9
--- /dev/null
+++ b/tests/test_kenlm.py
@@ -0,0 +1,17 @@
+
+from os import environ
+import pytest
+
+from cc2dataset import kenlm as mod
+
+DEFAULT_CACHEDIR = '/tmp'
+
+
+def test_perplexity_scorer():
+ cache_dir = environ.get("TMPDIR", DEFAULT_CACHEDIR)
+ model = mod.PerplexityScorer(cache_dir=cache_dir)
+ result = model("this is an example", prefix="perp/")
+ print(result)
+ assert list(result) == ['perp/ccnet/wikipedia', 'perp/ontocord/riverbed_kenlm']
+ assert result['perp/ccnet/wikipedia'] == pytest.approx(5394.3)
+ assert result['perp/ontocord/riverbed_kenlm'] == pytest.approx(30.2)
diff --git a/tests/test_newmain.py b/tests/test_newmain.py
new file mode 100644
index 0000000..c3c73d5
--- /dev/null
+++ b/tests/test_newmain.py
@@ -0,0 +1,26 @@
+import pytest
+from cc2dataset import cc2dataset
+from cc2dataset.index_utils import get_cc_links
+import os
+import pandas as pd
+import tempfile
+from glob import glob
+
+
+def test_main():
+ with tempfile.TemporaryDirectory() as tmpdir:
+ cc2dataset(
+ tmpdir,
+ crawl_index_count=None,
+ files_count=1,
+ master="local",
+ num_cores=1,
+ mem_gb=2,
+ multipart=None,
+ source_cc_protocol="s3",
+ shuffle=False,ccfile='warc',crawl_index_list=[-3]
+ )
+ files = list(glob(os.path.join(tmpdir, "*/*.parquet")))
+ assert len(files) == 256
+ df = pd.read_parquet(files)
+ assert len(df) > 100