From db8b744c2bc101a6ace0f28bacea5d59cc3bc4da Mon Sep 17 00:00:00 2001 From: robert berry Date: Fri, 28 Feb 2025 11:07:22 -0800 Subject: [PATCH 1/3] Code-prose-composition tagger. Add a tagger that adds attributes for code-prose-other composition of files based on line classifications. Coverage for code-prose-composition tagger. Improve error messages for spawn method checks. Set multiprocessing start method in test setup Set multiprocessing in test case with error handling. Add before suite hook to set mp start method. 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 "fork". Here we enforce the start method to be "spawn" for all tests before executing. linting. Remove error log messages. --- python/dolma/taggers/code_composition.py | 98 ++++++++++++++++++++++++ tests/python/conftest.py | 15 ++++ tests/python/test_code_composition.py | 98 ++++++++++++++++++++++++ 3 files changed, 211 insertions(+) create mode 100644 python/dolma/taggers/code_composition.py create mode 100644 tests/python/conftest.py create mode 100644 tests/python/test_code_composition.py diff --git a/python/dolma/taggers/code_composition.py b/python/dolma/taggers/code_composition.py new file mode 100644 index 00000000..ced61dc0 --- /dev/null +++ b/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) diff --git a/tests/python/conftest.py b/tests/python/conftest.py new file mode 100644 index 00000000..39b3839a --- /dev/null +++ b/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 "fork". +# Here we enforce the start method to be "spawn" for all tests before executing. +@pytest.fixture(scope="session", autouse=True) +def initialize_data_environment(): + try: + multiprocessing.set_start_method("spawn") + except Exception: + pass diff --git a/tests/python/test_code_composition.py b/tests/python/test_code_composition.py new file mode 100644 index 00000000..5b7b5308 --- /dev/null +++ b/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) From e71559103501d6c9484762cd120113b4b21c83e6 Mon Sep 17 00:00:00 2001 From: robert berry Date: Fri, 28 Feb 2025 14:01:19 -0800 Subject: [PATCH 2/3] Add code-prose tagger to init import. Additionally update commentary and word choice. --- python/dolma/taggers/__init__.py | 1 + tests/python/conftest.py | 8 ++++---- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/python/dolma/taggers/__init__.py b/python/dolma/taggers/__init__.py index ece1ae3b..52a9f996 100644 --- a/python/dolma/taggers/__init__.py +++ b/python/dolma/taggers/__init__.py @@ -1,6 +1,7 @@ from . import ( c4, code, + code_composition, gopher, jigsaw, language, diff --git a/tests/python/conftest.py b/tests/python/conftest.py index 39b3839a..86147674 100644 --- a/tests/python/conftest.py +++ b/tests/python/conftest.py @@ -4,11 +4,11 @@ # 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 "fork". -# Here we enforce the start method to be "spawn" for all tests before executing. +# 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_data_environment(): +def initialize_multiprocessing_start_method(): try: multiprocessing.set_start_method("spawn") except Exception: From 845fe2f92766b638823e96df117982f5c363447e Mon Sep 17 00:00:00 2001 From: robert berry Date: Sat, 1 Mar 2025 09:38:40 -0800 Subject: [PATCH 3/3] Brevity in prediction labels. Shorten prediction labels for readability and type-ability. --- python/dolma/taggers/code_composition.py | 10 +++--- tests/python/test_code_composition.py | 42 ++++++++++++------------ 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/python/dolma/taggers/code_composition.py b/python/dolma/taggers/code_composition.py index ced61dc0..c8198b41 100644 --- a/python/dolma/taggers/code_composition.py +++ b/python/dolma/taggers/code_composition.py @@ -22,7 +22,7 @@ from ..core.registry import TaggerRegistry -@TaggerRegistry.add("code-prose-composition") +@TaggerRegistry.add("code_composition") class CodeProseCompositionClassifier(BaseFastTextTagger): MODEL_PATH = "hf://techarb/code-prose-composition/code-comment-prose-model.bin" # noqa: E501 @@ -64,14 +64,14 @@ def predictions( 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)] + out = [Prediction(label="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}_pct", score=composition[label])) + out.append(Prediction(label=f"{label}", score=class_counts.get(label, 0))) out.append( Prediction( - label=f"{label}_mean_entropy", score=self.mean_entropy(prediction_distributions.get(label, [])) + label=f"{label}_entropy", score=self.mean_entropy(prediction_distributions.get(label, [])) ) ) diff --git a/tests/python/test_code_composition.py b/tests/python/test_code_composition.py index 5b7b5308..0b1f060d 100644 --- a/tests/python/test_code_composition.py +++ b/tests/python/test_code_composition.py @@ -48,14 +48,14 @@ def test_prose_text(self): 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"}, + {"prose_entropy", "boundaries", "prose_pct", "prose"}, ) 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) + self.assertEqual(scores["boundaries"], 0) + self.assertEqual(scores["prose_pct"], 1) + self.assertEqual(scores["prose"], 1) + self.assertLess(scores["prose_entropy"], 0.5) def test_code_text(self): doc = Document(source="fixtures", id="1", text=CODE_TEXT, version="v0") @@ -64,14 +64,14 @@ def test_code_text(self): 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"}, + {"code_entropy", "code_pct", "code", "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) + self.assertEqual(scores["boundaries"], 0) + self.assertEqual(scores["code_pct"], 1) + self.assertEqual(scores["code"], 3) + self.assertLess(scores["code_entropy"], 0.5) def test_code_prose_text(self): doc = Document(source="fixtures", id="1", text=CODE_PROSE_TEXT, version="v0") @@ -81,18 +81,18 @@ def test_code_prose_text(self): 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", + "code", + "prose", + "prose_entropy", + "code_pct", + "prose_pct", + "boundaries", + "code_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) + self.assertEqual(scores["boundaries"], 5) + self.assertGreater(scores["code_pct"], 0.5) + self.assertEqual(scores["code"], 9) + self.assertLess(scores["code_entropy"], 0.3)