Init transformer strategy once for all records

Why these changes are being introduced:

Formerly, a transformer strategy class was instantiated in a per-record fashion,
where things like the timdex_record_id and other record-level values were passed.
This ultimately felt awkward, when we could just as easily instantiate it once in
a more generic fashion, then build EmbeddingInput instances with the *result* of
the strategy extracting text from the TIMDEX JSON record.

How this addresses that need:

All record-level details are removed as arguments for initializing a transformer
strategy.  Instead, the helper function create_embedding_inputs() is responsible
for passing the TIMDEX JSON record to the transformer strategies, and then
building an EmbeddingInput object before yielding.  This keeps the init of those
strategies much simpler, and preventing properties in the class they don't really
need.

Side effects of this change:
* None

Relevant ticket(s):
* https://mitlibraries.atlassian.net/browse/USE-131
* https://mitlibraries.atlassian.net/browse/USE-132

Author: ghukill

embeddings/strategies/base.py

@@ -1,7 +1,5 @@
 from abc import ABC, abstractmethod
 
-from embeddings.embedding import EmbeddingInput
-
 
 class BaseStrategy(ABC):
     """Base class for embedding input strategies.
@@ -11,26 +9,6 @@ class BaseStrategy(ABC):
 
     STRATEGY_NAME: str  # type hint to document the requirement
 
-    def __init__(
-        self,
-        timdex_record_id: str,
-        run_id: str,
-        run_record_offset: int,
-        transformed_record: dict,
-    ) -> None:
-        """Initialize strategy with TIMDEX record metadata.
-
-        Args:
-            timdex_record_id: TIMDEX record ID
-            run_id: TIMDEX ETL run ID
-            run_record_offset: record offset within the run
-            transformed_record: parsed TIMDEX record JSON
-        """
-        self.timdex_record_id = timdex_record_id
-        self.run_id = run_id
-        self.run_record_offset = run_record_offset
-        self.transformed_record = transformed_record
-
     def __init_subclass__(cls, **kwargs: dict) -> None:  # noqa: D105
         super().__init_subclass__(**kwargs)
 
@@ -43,15 +21,9 @@ def __init_subclass__(cls, **kwargs: dict) -> None:  # noqa: D105
             )
 
     @abstractmethod
-    def extract_text(self) -> str:
-        """Extract text to be embedded from transformed_record."""
-
-    def to_embedding_input(self) -> EmbeddingInput:
-        """Create EmbeddingInput instance with strategy-specific extracted text."""
-        return EmbeddingInput(
-            timdex_record_id=self.timdex_record_id,
-            run_id=self.run_id,
-            run_record_offset=self.run_record_offset,
-            embedding_strategy=self.STRATEGY_NAME,
-            text=self.extract_text(),
-        )
+    def extract_text(self, timdex_record: dict) -> str:
+        """Extract text to be embedded from transformed_record.
+
+        Args:
+            timdex_record: TIMDEX JSON record ("transformed_record" in TIMDEX dataset)
+        """

embeddings/strategies/full_record.py

@@ -8,6 +8,6 @@ class FullRecordStrategy(BaseStrategy):
 
     STRATEGY_NAME = "full_record"
 
-    def extract_text(self) -> str:
+    def extract_text(self, timdex_record: dict) -> str:
         """Serialize the entire transformed_record as JSON."""
-        return json.dumps(self.transformed_record)
+        return json.dumps(timdex_record)

embeddings/strategies/processor.py

@@ -6,7 +6,7 @@
 
 
 def create_embedding_inputs(
-    timdex_records: Iterator[dict],
+    timdex_dataset_records: Iterator[dict],
     strategies: list[str],
 ) -> Iterator[EmbeddingInput]:
     """Yield EmbeddingInput instances for all records x all strategies.
@@ -15,7 +15,7 @@ def create_embedding_inputs(
     yielding one EmbeddingInput per combination.
 
     Args:
-        timdex_records: Iterator of TIMDEX records.
+        timdex_dataset_records: Iterator of TIMDEX records.
             Expected keys: timdex_record_id, run_id, run_record_offset,
             transformed_record (bytes)
         strategies: List of strategy names to apply
@@ -26,17 +26,24 @@ def create_embedding_inputs(
     Example:
         100 records x 3 strategies = 300 EmbeddingInput instances
     """
-    for timdex_record in timdex_records:
-        # decode and parse the TIMDEX JSON record
-        transformed_record = json.loads(timdex_record["transformed_record"].decode())
-
-        # apply all strategies to the record and yield
-        for strategy_name in strategies:
-            strategy_class = get_strategy_class(strategy_name)
-            strategy_instance = strategy_class(
-                timdex_record_id=timdex_record["timdex_record_id"],
-                run_id=timdex_record["run_id"],
-                run_record_offset=timdex_record["run_record_offset"],
-                transformed_record=transformed_record,
+    # instantiate strategy transformers
+    transformers = [get_strategy_class(strategy)() for strategy in strategies]
+
+    # loop through records and apply all strategies, yielding an EmbeddingInput for each
+    for timdex_dataset_record in timdex_dataset_records:
+
+        # decode and parse the TIMDEX JSON record once for all requested strategies
+        timdex_record = json.loads(timdex_dataset_record["transformed_record"].decode())
+
+        for transformer in transformers:
+            # prepare text for embedding from transformer strategy
+            text = transformer.extract_text(timdex_record)
+
+            # emit an EmbeddingInput instance
+            yield EmbeddingInput(
+                timdex_record_id=timdex_dataset_record["timdex_record_id"],
+                run_id=timdex_dataset_record["run_id"],
+                run_record_offset=timdex_dataset_record["run_record_offset"],
+                embedding_strategy=transformer.STRATEGY_NAME,
+                text=text,
             )
-            yield strategy_instance.to_embedding_input()

tests/test_strategies.py

@@ -8,22 +8,14 @@
 from embeddings.strategies.registry import get_strategy_class
 
 
-def test_full_record_strategy_creates_embedding_input():
-    transformed_record = {"timdex_record_id": "test-123", "title": ["Test Title"]}
-    strategy = FullRecordStrategy(
-        timdex_record_id="test-123",
-        run_id="run-456",
-        run_record_offset=42,
-        transformed_record=transformed_record,
-    )
+def test_full_record_strategy_extracts_text():
+    timdex_record = {"timdex_record_id": "test-123", "title": ["Test Title"]}
+    strategy = FullRecordStrategy()
 
-    embedding_input = strategy.to_embedding_input()
+    text = strategy.extract_text(timdex_record)
 
-    assert embedding_input.timdex_record_id == "test-123"
-    assert embedding_input.run_id == "run-456"
-    assert embedding_input.run_record_offset == 42
-    assert embedding_input.embedding_strategy == "full_record"
-    assert embedding_input.text == json.dumps(transformed_record)
+    assert text == json.dumps(timdex_record)
+    assert strategy.STRATEGY_NAME == "full_record"
 
 
 def test_create_embedding_inputs_yields_cartesian_product():