Code-prose-composition tagger.

python/dolma/taggers/__init__.py

@@ -1,6 +1,7 @@
 from . import (
     c4,
     code,
+    code_composition,
     gopher,
     jigsaw,
     language,

python/dolma/taggers/code_composition.py

@@ -0,0 +1,98 @@
+"""
+
+Code Prose Composition Classifier.
+
+This tagger classifies the composition of code and prose in a given text slice
+at the document level. It uses a FastText model trained on code and prose
+composition data.
+
+Tags include information about the number of code-prose boundaries, the
+composition of code and prose in the text, and the entropy of the predicted
+labels.
+
+@robertb
+
+"""
+
+import math
+from typing import Dict, Iterable, List, Tuple
+
+from ..core.data_types import TextSlice
+from ..core.ft_tagger import BaseFastTextTagger, Prediction
+from ..core.registry import TaggerRegistry
+
+
+@TaggerRegistry.add("code-prose-composition")
+class CodeProseCompositionClassifier(BaseFastTextTagger):
+    MODEL_PATH = "hf://techarb/code-prose-composition/code-comment-prose-model.bin"  # noqa: E501
+
+    def __init__(self):
+        super().__init__(model_path=self.MODEL_PATH, model_mode=self.DOCUMENT_LEVEL_TAGGER)
+
+    def calculate_entropy(self, distribution: List[float]) -> float:
+        entropy = 0.0
+        for p in distribution:
+            if p > 0:
+                entropy -= p * math.log2(p)
+        return entropy
+
+    def mean_entropy(self, list_of_distributions: List[List[float]]) -> float:
+        if not list_of_distributions:
+            return 0
+
+        total_entropy = 0.0
+        for dist in list_of_distributions:
+            total_entropy += self.calculate_entropy(dist)
+        return total_entropy / len(list_of_distributions)
+
+    def line_label(self, line: str) -> Tuple[str, List[float]]:
+        label = "other"
+        probabilities = []
+        if len(line) > 3:
+            labels, probabilities = self.classifier.predict(line, k=-1)
+
+            label = labels[0].lstrip("__label__")
+        return label, probabilities
+
+    def predictions(
+        self,
+        code_prose_boundaries: int,
+        class_counts: Dict[str, int],
+        prediction_distributions: Dict[str, List[List[float]]],
+    ) -> Iterable[Prediction]:
+        composition = {}
+        for label, count in class_counts.items():
+            composition[label] = round((count / sum(class_counts.values())), 2)
+
+        out = [Prediction(label="code_prose_boundaries", score=code_prose_boundaries)]
+
+        for label in composition.keys():
+            out.append(Prediction(label=f"{label}_composition", score=composition[label]))
+            out.append(Prediction(label=f"{label}_count", score=class_counts.get(label, 0)))
+            out.append(
+                Prediction(
+                    label=f"{label}_mean_entropy", score=self.mean_entropy(prediction_distributions.get(label, []))
+                )
+            )
+
+        return out
+
+    def predict_slice(self, text_slice: TextSlice) -> Iterable[Prediction]:
+        class_counts: Dict[str, int] = {}
+        prediction_distributions: Dict[str, List[List[float]]] = {}
+        active_class, code_prose_boundaries = None, 0
+
+        for line in [line.strip() for line in text_slice.text.splitlines()]:
+            if not line:
+                continue
+
+            label, probabilities = self.line_label(line)
+
+            prediction_distributions.setdefault(label, []).append(probabilities)
+            class_counts[label] = class_counts.get(label, 0) + 1
+
+            if active_class in ["code", "prose"] and label in ["code", "prose"] and label != active_class:
+                code_prose_boundaries += 1
+            active_class = label
+
+        return self.predictions(code_prose_boundaries, class_counts, prediction_distributions)

tests/python/conftest.py

@@ -0,0 +1,15 @@
+import multiprocessing
+
+import pytest
+
+
+# The default multiprocessing start method is "fork" which is not compatible with
+# with runtime assertions that it is set to "spawn". When running unit tests, it's
+# possible to call an external library that sets the start method to the default.
+# Here we set the start method to be "spawn" for all tests before executing.
+@pytest.fixture(scope="session", autouse=True)
+def initialize_multiprocessing_start_method():
+    try:
+        multiprocessing.set_start_method("spawn")
+    except Exception:
+        pass

tests/python/test_code_composition.py

@@ -0,0 +1,98 @@
+from unittest import TestCase
+
+from dolma.core.data_types import Document
+from dolma.taggers.code_composition import CodeProseCompositionClassifier
+
+PROSE_TEXT = """
+The Allen Institute for AI (abbreviated AI2) is a 501(c)(3) non-profit research institute founded by late Microsoft co-founder and philanthropist Paul Allen in 2014. The institute seeks to conduct high-impact AI research and engineering in service of the common good. Oren Etzioni was appointed by Paul Allenin September 2013 to direct the research at the institute. After leading the organization for nine years, Oren Etzioni stepped down from his role as CEO on September 30, 2022. He was replaced in an interim capacity by the leading researcher of the company's Aristo project, Peter Clark. On June 20, 2023, AI2 announced Ali Farhadi as its next CEO starting July 31, 2023. The company's board formed a search committee for a new CEO. AI2 also has an active office in Tel Aviv, Israel.
+"""
+
+CODE_TEXT = """
+def foo():
+    if True:
+        print("Hello, world!")
+"""
+
+CODE_PROSE_TEXT = """
+The following function adds two numbers together.
+Then it returns the result.
+
+def foo():
+    x = 1 + 1
+    return x
+
+Next we demonstrate multiplying two numbers together.
+Note that these are floats.
+We return the result rounded to 2 decimal places.
+
+def bar():
+    x = 1.1 * 2.2
+    return x
+
+Finally, we show how to divide two numbers.
+
+def baz():
+    x = 1 / 2
+    return x
+"""
+
+
+class TestDolmaCodeProseCompositionClassifier(TestCase):
+    def setUp(self) -> None:
+        self.code_composition_tagger = CodeProseCompositionClassifier()
+
+    def test_prose_text(self):
+        doc = Document(source="fixtures", id="1", text=PROSE_TEXT, version="v0")
+        pred = self.code_composition_tagger.predict(doc)
+
+        self.assertEqual(len(pred.spans), 4)
+        self.assertEqual(
+            {s.type for s in pred.spans},
+            {"prose_mean_entropy", "code_prose_boundaries", "prose_composition", "prose_count"},
+        )
+
+        scores = {s.type: s.score for s in pred.spans}
+        self.assertEqual(scores["code_prose_boundaries"], 0)
+        self.assertEqual(scores["prose_composition"], 1)
+        self.assertEqual(scores["prose_count"], 1)
+        self.assertLess(scores["prose_mean_entropy"], 0.5)
+
+    def test_code_text(self):
+        doc = Document(source="fixtures", id="1", text=CODE_TEXT, version="v0")
+        pred = self.code_composition_tagger.predict(doc)
+
+        self.assertEqual(len(pred.spans), 4)
+        self.assertEqual(
+            {s.type for s in pred.spans},
+            {"code_mean_entropy", "code_composition", "code_count", "code_prose_boundaries"},
+        )
+
+        scores = {s.type: s.score for s in pred.spans}
+        self.assertEqual(scores["code_prose_boundaries"], 0)
+        self.assertEqual(scores["code_composition"], 1)
+        self.assertEqual(scores["code_count"], 3)
+        self.assertLess(scores["code_mean_entropy"], 0.5)
+
+    def test_code_prose_text(self):
+        doc = Document(source="fixtures", id="1", text=CODE_PROSE_TEXT, version="v0")
+        pred = self.code_composition_tagger.predict(doc)
+
+        self.assertEqual(len(pred.spans), 7)
+        self.assertEqual(
+            {s.type for s in pred.spans},
+            {
+                "code_count",
+                "prose_count",
+                "prose_mean_entropy",
+                "code_composition",
+                "prose_composition",
+                "code_prose_boundaries",
+                "code_mean_entropy",
+            },
+        )
+
+        scores = {s.type: s.score for s in pred.spans}
+        self.assertEqual(scores["code_prose_boundaries"], 5)
+        self.assertGreater(scores["code_composition"], 0.5)
+        self.assertEqual(scores["code_count"], 9)
+        self.assertLess(scores["code_mean_entropy"], 0.3)