Add rolling normalized repetitions to the documents

pyproject.toml

@@ -1,4 +1,3 @@
-
 [project]
 name = "dactory"
 description = "The data factory"
@@ -28,6 +27,11 @@ classifiers = [
 ]
 dynamic = ["version"]
 
+[dependency-groups]
+dev = [
+    "pytest>=8.4.2",
+]
+
 [project.scripts]
 dactory = "dactory.main:main"
 
@@ -54,3 +58,6 @@ line-length = 95
 
 [tool.ruff.format]
 skip-magic-trailing-comma = true
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

python/dactory/create.py

@@ -13,7 +13,7 @@
 from retry import retry
 from tqdm import tqdm
 
-from dactory import dedup_document
+from dactory import compute_long_words, compute_repetitions_rolling, dedup_document
 from dactory.bloom_filter import load_bloom_filter
 from dactory.scoring import ScoringModels
 from dactory.zstd_writer import zstd_writer
@@ -81,6 +81,8 @@ def get_record_dict(
         group_idx=group_idx,
         warc_file=warc_file,
         record_idx=record_idx,
+        repetitions=None,  # will be filled later
+        long_words=None,  # will be filled later
     )
 
 
@@ -220,11 +222,14 @@ def download_warcs_for_group(args: LoadedArgs, group_idx: int, warc_paths: list[
                 )
                 if len(document.text) < args.min_length:
                     continue
+            document.repetitions = compute_repetitions_rolling(document.text, 20)
+            document.long_words = compute_long_words(document.text, min_length=15)
 
             if args.scoring_models is not None:
                 scores = args.scoring_models.get_doc_scores(document.text, document.language)
                 if scores["rand"] > args.max_rand_score:
                     continue
+                document.scores = {k: round(v, 3) for k, v in scores.items()}
 
             progress_bar_bytes.update(len(document.text))
             work_already_done[document.warc_file].last_record_seen = document.record_idx

python/dactory/document.py

@@ -14,6 +14,8 @@ class Document(BaseModel):
     group_idx: int
     warc_file: str
     record_idx: int
+    repetitions: float | None
+    long_words: float | None
 
     class Config:
         validate_by_name = True

python/dactory/download_models.py

@@ -6,6 +6,7 @@
 CACHE_DIRECTORY = Path.home() / ".cache" / "dactory"
 HF_PREFIX = "hf://"
 
+
 def download_if_necessary(path_or_url: str) -> Path:
     if path_or_url.startswith(HF_PREFIX):
         path_or_url = path_or_url.removeprefix(HF_PREFIX)

python/dactory/main.py

@@ -1,9 +1,12 @@
+import io
+from collections import Counter
 from pathlib import Path
 from typing import Annotated
 
 import fasttext
 import pydantic
 import typer
+import zstandard as zstd
 from typer import Argument, Option
 
 import dactory.create
@@ -14,13 +17,14 @@
 from dactory.profiling import profile
 from dactory.scoring import get_scoring_models
 from dactory.warc_groups import get_warc_groups
-from .download_models import HF_PREFIX
 
+from .document import Document
+from .download_models import HF_PREFIX
 
 KYUTAI_HF_REPOSITORY = HF_PREFIX + "kyutai/dactory-models"
 
 DEFAULT_LANGUAGE_DETECTOR_MODEL = (
-    f"https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin"
+    "https://dl.fbaipublicfiles.com/fasttext/supervised-models/lid.176.bin"
 )
 
 # fmt: off
@@ -35,11 +39,6 @@
 
 app = typer.Typer()
 
-from .document import Document
-import zstandard as zstd
-from io import TextIOWrapper
-import io
-from collections import Counter
 
 @app.command()
 def stats():
@@ -61,11 +60,6 @@ def stats():
         print(f"{lang}: {count / total * 100:.2f}%")
 
 
-
-
-
-
-
 @app.command()
 def list_languages(
     lang_detection_model: Annotated[
@@ -107,7 +101,8 @@ class CreateArgs(pydantic.BaseModel):
         "CC-MAIN-2024-51"
     )
     load_models_early: Annotated[
-        bool, Option(help="Load scoring models before downloading, disable for faster iteration.")
+        bool,
+        Option(help="Load scoring models before downloading, disable for faster iteration."),
     ] = True
     workers: Annotated[
         int,

python/dactory/profiling.py

@@ -2,5 +2,5 @@
 
 profile = line_profiler.LineProfiler()
 
-# Add @profile to any function in the codebase, and run any command, 
-# you'll see a profiling report once the command is finished (even if it is stopped).
+# Add @profile to any function in the codebase, and run any command,
+# you'll see a profiling report once the command is finished (even if it is stopped).

src/lib.rs

@@ -3,11 +3,14 @@ use pyo3::prelude::*;
 mod bloom;
 mod code;
 mod entry;
+mod repetitions;
 
 /// A Python module implemented in Rust.
 #[pymodule]
 fn dactory(m: &Bound<'_, PyModule>) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(entry::dedup_document, m)?)?;
+    m.add_function(wrap_pyfunction!(repetitions::compute_repetitions_rolling, m)?)?;
+    m.add_function(wrap_pyfunction!(repetitions::compute_long_words, m)?)?;
     m.add_class::<bloom::BloomFilter>()?;
     m.add_class::<entry::FastTextPyWrapper>()?;
     Ok(())

src/repetitions.rs

@@ -1,3 +1,7 @@
+use std::collections::HashMap;
+use pyo3::prelude::*;
+
+
 fn fnv1a_64(data: &[u8]) -> u64 {
     let mut h = 14695981039346656037u64;
     for b in data {
@@ -17,14 +21,15 @@ fn init_adler_32(data: &[u8], n: usize) -> (u32, u32) {
     (a, b)
 }
 
-fn compute_repetitions_rolling(text: &String, n: usize) -> u32 {
+#[pyfunction]
+pub fn compute_repetitions_rolling(text: &str, n: usize) -> f32 {
     let mut n_repetitions = 0;
     let text = text.as_bytes();
     let mut ngrams = vec![0 as u32; 4 * text.len()];
     //let mut ngrams = vec![0 as u8; 16 * text.len()];
 
     if text.len() <= n {
-	return 0;
+	return 0.0;
     }
 
     //let (mut a, mut b) = init_adler_32(&text, n);
@@ -68,7 +73,21 @@ fn compute_repetitions_rolling(text: &String, n: usize) -> u32 {
 	    ngrams[p] = h32;
 	}
     }
-    n_repetitions
+    n_repetitions as f32 / text.len() as f32
+}
+
+#[pyfunction]
+pub fn compute_long_words(text: &str, min_length: usize) -> f32 {
+    if text.len() == 0 {
+        return 0.0;
+    }
+    let mut n_long_words = 0;
+    for word in text.split_whitespace() {
+        if word.len() >= min_length {
+            n_long_words += word.len();
+        }
+    }
+    n_long_words as f32 / text.len() as f32
 }
 
 fn compute_repetitions(text: &String, n: usize) -> u32 {

tests/rust/test_repetitions.py

@@ -0,0 +1,31 @@
+from dactory import compute_long_words
+
+
+class TestComputeLongWords:
+    """Tests for `compute_long_words` function"""
+
+    error_margin = 1e-6
+
+    def test_empty_string(self):
+        text = ""
+        res = compute_long_words(text, min_length=0)
+        expected = 0.0
+        assert res == expected
+
+    def test_no_long_words(self):
+        text = "shrot words only"
+        res = compute_long_words(text, min_length=6)
+        expected = 0.0
+        assert res == expected
+
+    def test_all_long_words(self):
+        text = "these are some longwords"
+        res = compute_long_words("these are some longwords", min_length=3)
+        expected = 21 / len(text)
+        assert abs(res - expected) < self.error_margin
+
+    def test_mixed_words(self):
+        text = "this sentence contains some longwords"
+        res = compute_long_words(text, min_length=8)
+        expected = 25 / len(text)
+        assert abs(res - expected) < self.error_margin