Coverage for code-prose-composition tagger.

tests/python/test_code.py

@@ -1,238 +1,98 @@
-"""
-
-Unit tests for code taggers.
+from unittest import TestCase
 
-@soldni
+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.
 """
 
-import re
-import unittest
-
-from bs4 import BeautifulSoup
-
-from dolma.core.data_types import Document, DocumentWithMetadata
-from dolma.taggers.code import (
-    CodeCopyrightTagger,
-    CodeRedPajamaTaggers,
-    CodeSecretsTagger,
-    CodeStarCoderTaggers2,
-)
-
-DOC_WITH_SECRETS_AND_COPYRIGHT = """
-/* copyright: Test 2023 **/
-
-This is a document.
-
-This line contains a secret: https://username:[email protected]
-
-This is a line with just text.
+CODE_TEXT = """
+def foo():
+    if True:
+        print("Hello, world!")
 """
 
+CODE_PROSE_TEXT = """
+The following function adds two numbers together.
+Then it returns the result.
 
-class TestCodeTaggers(unittest.TestCase):
-    def setUp(self) -> None:
-        self.doc = Document(id="0", text=DOC_WITH_SECRETS_AND_COPYRIGHT.strip(), source=__file__)
-        return super().setUp()
-
-    def test_code_secrets_tagger(self) -> None:
-        tagger = CodeSecretsTagger()
-        result = tagger.predict(self.doc)
-
-        self.assertEqual(len(result.spans), 3)
-
-        self.assertEqual(result.spans[0].type, "SECRET_Secret_Keyword")
-        self.assertEqual(result.spans[0].select(self.doc), "https://username:[email protected]")
-
-        self.assertEqual(result.spans[1].type, "SECRET_Basic_Auth_Credentials")
-        self.assertEqual(result.spans[1].select(self.doc), "password")
-
-        self.assertEqual(result.spans[2].type, "doc")
-        self.assertEqual(result.spans[2].select(self.doc), self.doc.text)
-
-    def test_copyright_notice(self):
-        tagger = CodeCopyrightTagger()
-        result = tagger.predict(self.doc)
-
-        self.assertEqual(len(result.spans), 2)
-
-        self.assertEqual(result.spans[0].type, "copyright_notice")
-        self.assertEqual(result.spans[0].select(self.doc), "/* copyright: Test 2023 **/")
-
-        self.assertEqual(result.spans[1].type, "doc")
-        self.assertEqual(result.spans[1].select(self.doc), self.doc.text)
-
-
-class TestRedPajamaTaggers(unittest.TestCase):
-    def setUp(self) -> None:
-        self.doc = Document(id="0", text=DOC_WITH_SECRETS_AND_COPYRIGHT.strip(), source=__file__)
-        self.whitespace_regex = re.compile(r"\w+|[^\w\s]+")
-        return super().setUp()
-
-    def test_code_red_pajama_taggers(self) -> None:
-        tagger = CodeRedPajamaTaggers()
-        result = tagger.predict(self.doc)
-
-        # handy precomputed values
-        line_lengths = list(map(len, self.doc.text.splitlines()))
-        words = self.whitespace_regex.findall(self.doc.text)
-        self.assertGreater(len(line_lengths), 0)
-        self.assertGreater(len(words), 0)
-
-        self.assertEqual(len(result.spans), 4)
-        self.assertEqual(result.spans[0].type, "max_line_length_doc")
-        self.assertEqual(result.spans[0].score, max(line_lengths))
-
-        self.assertEqual(result.spans[1].type, "avg_line_length_doc")
-        self.assertEqual(result.spans[1].score, sum(line_lengths) / len(line_lengths))
-
-        self.assertEqual(result.spans[2].type, "alnum_prop_doc")
-        self.assertEqual(result.spans[2].score, len(list(filter(str.isalnum, self.doc.text))) / len(self.doc.text))
-
-        # TODO: This test fail; check with Akshita if this is expected
-        # self.assertEqual(result.spans[3].type, "alpha_token_prop_doc")
-        # self.assertEqual(result.spans[3].score, len(list(filter(str.isalpha, words))) / len(words))
-
-
-DOC_WITH_METADATA = """
-An XML file begins as follows:
-
-```
-<?xml version="1.0" encoding="UTF-8"?>
-...
-</xml>
-```
-
-An HTML file begins as follows:
+def foo():
+    x = 1 + 1
+    return x
 
-```
-<!DOCTYPE html>
-<html>
-...
-</html>
-```
+Next we demonstrate multiplying two numbers together.
+Note that these are floats.
+We return the result rounded to 2 decimal places.
 
-These are different.
-"""
+def bar():
+    x = 1.1 * 2.2
+    return x
 
-DOC_WITH_PYTHON_CONTENT = """
-def foo():
-    # prints hello world
-    print("Hello, World!")
-"""
+Finally, we show how to divide two numbers.
 
-DOC_WITH_HTML_CONTENT = """
-<!DOCTYPE html>
-<html>
-<head>
-    <title>Page Title</title>
-</head>
-<body>
-    <h1>This is a Heading</h1>
-    <p>This is a paragraph.</p>
-</body>
-<javascript>
-    console.log("Hello, World!")
-    for (let i = 0; i < 10; i++) {
-        console.log(i)
-    }
-</javascript>
-</html>
+def baz():
+    x = 1 / 2
+    return x
 """
 
 
-class TestStarCoderTaggers(unittest.TestCase):
+class TestDolmaCodeProseCompositionClassifier(TestCase):
     def setUp(self) -> None:
-        self.md_doc = DocumentWithMetadata(
-            id="0",
-            text=DOC_WITH_METADATA.strip(),
-            source=__file__,
-            metadata={"ext": "md", "max_stars_count": 10},
-        )
-        self.python_doc = DocumentWithMetadata(
-            id="1",
-            text=DOC_WITH_PYTHON_CONTENT.strip(),
-            source=__file__,
-            metadata={"ext": "py", "max_stars_count": 1},
-        )
-        self.html_doc = DocumentWithMetadata(
-            id="2",
-            text=DOC_WITH_HTML_CONTENT.strip(),
-            source=__file__,
-            metadata={"ext": "html", "max_stars_count": 5},
-        )
-        self.tagger = CodeStarCoderTaggers2()
-        return super().setUp()
-
-    def test_metadata_tagger(self):
-        result = self.tagger.predict(self.md_doc)
-        self.assertEqual(len(result.spans), 4)
-
-        self.assertEqual(result.spans[0].type, "has_xml_template_doc")
-        self.assertEqual(result.spans[0].score, 1.0)
-
-        self.assertEqual(result.spans[1].type, "num_github_stars_doc")
-        self.assertEqual(result.spans[1].score, 10.0)
+        self.code_composition_tagger = CodeProseCompositionClassifier()
 
-        # not a python, js, or java, so this is pinned to 0.5
-        self.assertEqual(result.spans[2].type, "code_to_comment_ratio_doc")
-        self.assertEqual(result.spans[2].score, 0.5)
+    def test_prose_text(self):
+        doc = Document(source="fixtures", id="1", text=PROSE_TEXT, version="v0")
+        pred = self.code_composition_tagger.predict(doc)
 
-        # not html, so this is pinned to 1.0
-        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
-        self.assertEqual(result.spans[3].score, 1.0)
-
-    def test_python_tagger(self):
-        result = self.tagger.predict(self.python_doc)
-
-        comment_lines = [
-            lns.split("#")[1].strip()
-            for ln in self.python_doc.text.split("\n")
-            if (lns := ln.strip()).startswith("#")
-        ]
-
-        self.assertEqual(len(result.spans), 4)
-
-        self.assertEqual(result.spans[0].type, "has_xml_template_doc")
-        self.assertEqual(result.spans[0].score, 0.0)
-
-        self.assertEqual(result.spans[1].type, "num_github_stars_doc")
-        self.assertEqual(result.spans[1].score, 1.0)
-
-        self.assertEqual(result.spans[2].type, "code_to_comment_ratio_doc")
-        self.assertEqual(result.spans[2].score, sum(map(len, comment_lines)) / len(self.python_doc.text))
-
-        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
-        self.assertEqual(result.spans[3].score, 1.0)
-
-    def test_html_tagger(self):
-        result = self.tagger.predict(self.html_doc)
-
-        soup = BeautifulSoup(self.html_doc.text, features="html.parser")
-
-        self.assertEqual(len(result.spans), 4)
-
-        self.assertEqual(result.spans[0].type, "has_xml_template_doc")
-        self.assertEqual(result.spans[0].score, 0.0)
+        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"},
+        )
 
-        self.assertEqual(result.spans[1].type, "num_github_stars_doc")
-        self.assertEqual(result.spans[1].score, 5.0)
+        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(result.spans[2].type, "code_to_comment_ratio_doc")
-        self.assertEqual(result.spans[2].score, 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(result.spans[3].type, "code_to_text_ratio_html_doc")
-        self.assertEqual(result.spans[3].score, len(soup.get_text()) / len(self.html_doc.text))
+        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"},
+        )
 
-    def test_html_tagger_doc_too_short(self):
-        doc = DocumentWithMetadata(
-            id="3",
-            text="<html><head></head><body></body></html>",
-            source=__file__,
-            metadata={"ext": "html", "max_stars_count": 5},
+        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",
+            },
         )
-        doc.text = doc.text[:100]
-        result = self.tagger.predict(doc)
 
-        self.assertEqual(result.spans[3].type, "code_to_text_ratio_html_doc")
-        self.assertEqual(result.spans[3].score, 0.0)
+        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)