Make training data of bert configurable

Author: jvamvas

data/synthetic.py

@@ -0,0 +1,31 @@
+from pathlib import Path
+from typing import List
+
+import jsonlines
+
+from rsd.data import DifferenceDataset
+from rsd.recognizers.utils import DifferenceSample
+
+class SyntheticSTSDataset(DifferenceDataset):
+
+    def __init__(self,
+                 path: Path,
+                 ):
+        super().__init__()
+        self.path = path
+        with jsonlines.open(path) as reader:
+            self.dataset = list(reader)
+
+    def get_samples(self) -> List[DifferenceSample]:
+        samples = []
+        for sample in self.dataset:
+            samples.append(DifferenceSample(
+                tokens_a=sample["text_a"].split(),
+                tokens_b=sample["text_b"].split(),
+                labels_a=sample["labels_a"],
+                labels_b=sample["labels_b"],
+            ))
+        return samples
+
+    def __str__(self):
+        return f"SyntheticDataset({self.path.name})"

experiments/benchmark.py

@@ -189,3 +189,35 @@ def to_dataset(self, both_directions: bool = False) -> datasets.Dataset:
                 data["labels_b"].append(list(doc.labels_a))
 
         return datasets.Dataset.from_dict(data)
+
+
+class MultiLengthDifferenceRecognitionBenchmark:
+
+    def __init__(self,
+                 positive_dataset: DifferenceDataset,
+                 negative_dataset: DifferenceDataset = None,
+                 positive_ratio: float = 1.0,
+                 max_sentences_per_document: int = 1,
+                 max_inversions: int = 0,
+                 seed: int = 42,
+                 ):
+        assert max_sentences_per_document >= 1
+        assert max_inversions <= max_sentences_per_document
+        self.num_sentences_range = list(range(1, max_sentences_per_document + 1))
+        self.num_inversions_range = [0] * (max_sentences_per_document - max_inversions) + list(range(1, max_inversions + 1))
+        assert len(self.num_inversions_range) == len(self.num_sentences_range)
+        self.benchmarks = []
+        for num_sentences, num_inversions in zip(self.num_sentences_range, self.num_inversions_range):
+            benchmark = DifferenceRecognitionBenchmark(
+                positive_dataset=positive_dataset,
+                negative_dataset=negative_dataset,
+                positive_ratio=positive_ratio,
+                num_sentences_per_document=num_sentences,
+                num_inversions=num_inversions,
+                seed=seed,
+            )
+            self.benchmarks.append(benchmark)
+
+    def to_dataset(self, both_directions: bool = False) -> datasets.Dataset:
+        ds = [benchmark.to_dataset(both_directions=both_directions) for benchmark in self.benchmarks]
+        return datasets.concatenate_datasets(ds)

tests/test_benchmark.py

@@ -1,4 +1,5 @@
 import random
+from pathlib import Path
 from typing import Dict, Tuple
 from unittest import TestCase
 
@@ -7,15 +8,19 @@
 
 from rsd.data.ists import ISTSDataset
 from rsd.data.pawsx import PAWSXDataset
-from rsd.experiments.benchmark import DifferenceRecognitionBenchmark
+from rsd.data.synthetic import SyntheticSTSDataset
+from rsd.experiments.benchmark import DifferenceRecognitionBenchmark, MultiLengthDifferenceRecognitionBenchmark
 from rsd.recognizers.base import DifferenceRecognizer
 from rsd.recognizers.utils import DifferenceSample
 
 
 class DifferenceRecognitionBenchmarkTestCase(TestCase):
 
     def setUp(self) -> None:
-        self.positive_dataset = ISTSDataset()
+        # self.positive_dataset = ISTSDataset()
+        self.positive_dataset = SyntheticSTSDataset(
+            path=Path(__file__).parent.parent.parent / "ists_finetuning" / "data" / "synthetic_rsd" / "ft:gpt-4o-mini-2024-07-18:cl-uzh:rsd-test-en-v3:AxfIBckt_train_v2.jsonl",
+        )
         self.negative_dataset = PAWSXDataset()
         self.benchmark = DifferenceRecognitionBenchmark(
             positive_dataset=self.positive_dataset,
@@ -109,7 +114,7 @@ def predict(self, a: str, b: str, *args, **kwargs):
 
         recognizer = OracleRecognizer(self.benchmark)
         result = self.benchmark.evaluate(recognizer)
-        self.assertEqual(1, result.spearman)
+        self.assertAlmostEqual(1, result.spearman)
 
     def test_to_dataset(self):
         # Test conversion to HuggingFace dataset
@@ -153,3 +158,39 @@ def test_to_dataset_both_directions(self):
         self.assertEqual(second_example["labels_a"], list(first_doc.labels_b))
         self.assertEqual(second_example["labels_b"], list(first_doc.labels_a))
 
+
+class TestMultiLengthDifferenceRecognitionBenchmark(TestCase):
+    def setUp(self):
+        self.positive_dataset = SyntheticSTSDataset(
+            path=Path(__file__).parent.parent.parent / "ists_finetuning" / "data" / "synthetic_rsd" / "ft:gpt-4o-mini-2024-07-18:cl-uzh:rsd-test-en-v3:AxfIBckt_train_v2.jsonl",
+        )
+        self.negative_dataset = PAWSXDataset()
+
+    def test_basic_functionality(self):
+        benchmark = MultiLengthDifferenceRecognitionBenchmark(
+            positive_dataset=self.positive_dataset,
+            negative_dataset=self.negative_dataset,
+            positive_ratio=0.5,
+            max_sentences_per_document=3,
+            max_inversions=2,
+            seed=42
+        )
+
+        # Check that we have the expected number of benchmarks
+        self.assertEqual(len(benchmark.benchmarks), 3)
+
+        # Check that num_sentences_range is correct
+        self.assertEqual(benchmark.num_sentences_range, [1, 2, 3])
+
+        # Check that num_inversions_range is correct (1 for first entry, then 1, 2)
+        self.assertEqual(benchmark.num_inversions_range, [0, 1, 2])
+
+        # Test dataset conversion
+        dataset = benchmark.to_dataset()
+        self.assertIsInstance(dataset, datasets.Dataset)
+        self.assertTrue({"text_a", "text_b", "labels_a", "labels_b"}.issubset(dataset.features))
+
+        # Test both directions
+        dataset_both = benchmark.to_dataset(both_directions=True)
+        self.assertEqual(len(dataset_both), len(dataset) * 2)
+

tests/test_data.py

@@ -1,7 +1,9 @@
+from pathlib import Path
 from unittest import TestCase
 
 from rsd.data.ists import ISTSDataset
 from rsd.data.pawsx import PAWSXDataset, CrosslingualPAWSXDataset
+from rsd.data.synthetic import SyntheticSTSDataset
 from rsd.recognizers.utils import DifferenceSample
 
 
@@ -59,3 +61,17 @@ def test_get_samples(self):
         self.assertSetEqual(set(sample.labels_b), {-1})
         print(samples[0])
         print(samples[1])
+
+
+class SyntheticSTSDatasetTestCase(TestCase):
+
+    def setUp(self) -> None:
+        self.dataset = SyntheticSTSDataset(
+            path=Path(__file__).parent.parent.parent / "ists_finetuning" / "data" / "synthetic_rsd" / "ft:gpt-4o-mini-2024-07-18:cl-uzh:rsd-test-en-v3:AxfIBckt_train_v2.jsonl",
+        )
+
+    def test_get_samples(self):
+        sample: DifferenceSample = self.dataset.get_samples()[0]
+        self.assertEqual(len(sample.tokens_a), len(sample.labels_a))
+        self.assertEqual(len(sample.tokens_b), len(sample.labels_b))
+        print(sample)