USE 131 - Framework for embedding input strategies

[ghukill]

Purpose and background context

This PR establishes a framework for "strategies" used to extract and transform data from TIMDEX records into single strings used to create an embedding from.

As noted in the commit message, while we only have a single strategy at the moment of using the full record, we want to ensure we can support other strategies in the future and potentially multiple strategies for a record, in a single run (resulting in the cartesian product of records * strategies requested).

This establishes a new embeddings/strategies module, with a base class BaseStrategy, and our first implementation class FullRecordStrategy.

Overall, it's quite simple: a bit of boilerplate to get down to the strategy's extract_text() method which is opinionated for how text is extracted to create an embedding from.

How can a reviewer manually see the effects of these changes?

1- Set Dev1 AWS credentials

2- Set env vars:

TIMDEX_DATASET_LOCATION=s3://timdex-extract-dev-222053980223/dataset

3- Start an ipython shell:

uv run --env-file .env ipython

4- Prepare an iterator of TIMDEX records:

import os
from timdex_dataset_api import TIMDEXDataset

# init TIMDEXDataset and get an iterator or records
td = TIMDEXDataset(os.environ["TIMDEX_DATASET_LOCATION"])
records = td.read_dicts_iter(source="libguides", limit=3)

5- Create EmbeddingInput instances:

from embeddings.strategies.processor import create_embedding_inputs

# create EmbeddingInputs
embedding_inputs = list(create_embedding_inputs(records, ["full_record"]))

The following is an example EmbeddingInput returned:

EmbeddingInput(timdex_record_id='libguides:guides-175846', run_id='370abbb9-fcfc-4356-bff9-7b000dff109e', run_record_offset=0, embedding_strategy='full_record', text='<full record JSON dump here...>')

Includes new or updated dependencies?

NO

Changes expectations for external applications?

NO

What are the relevant tickets?

• https://mitlibraries.atlassian.net/browse/USE-131
• https://mitlibraries.atlassian.net/browse/USE-132

Code review

• Code review best practices are documented here and you are encouraged to have a constructive dialogue with your reviewers about their preferences and expectations.

[ehanson8]

ipython testing worked as expected, a few minor comments and questions but nothing blocking. Great work, this is clean and tight!

[ehanson8]

Is this check necessary since the attribute exists in the base class even if it's not defined?

[ghukill]

Yes, I think it is. Without this, you could leave out the class level STRATEGY_NAME attribute on a real strategy class.

Another pattern could be something like this:

class BaseStrategy(ABC):
  @absractmethod
  @classmethod
  def strategy_name(self) -> str:
    # return strategy name...

Which would allow getting the strategy name for an uninstantiated class, which is important. We do something similar in other projects I think.

But this pattern, with a little logic in the base class, enforces that child classes define it.

[jonavellecuerdo]

Hmm, I'm curious if you had an idea of other methods that could be added to a child class of BaseStrategy? 🤔 When I think of this class, I understand it as a "collection of methods that could be applied to TIMDEX records", and while hard to explain, I was surprised by the number of required variables for the __init__ method:

self.timdex_record_id = timdex_record_id
self.run_id = run_id
self.run_record_offset = run_record_offset
self.transformed_record = transformed_record

It didn't feel clear to me as to why all these fields were required to instantiate a "strategy", if that makes sense. I do understand why it is important information for the resulting EmbeddingInput returned by the BaseStrategy.to_embedding_input() method, which is why I find it fits more as parameters for that method specifically, instead of the __init__ method. However, perhaps there is something I'm not thinking of re: other methods that you expect to define for BaseStrategy classes that will heavily rely on these fields. 🤔

[ghukill]

Hmm, I'm curious if you had an idea of other methods that could be added to a child class of BaseStrategy? 🤔 When I think of this class, I understand it as a "collection of methods that could be applied to TIMDEX records", and while hard to explain, I was surprised by the number of required variables for the __init__ method:

self.timdex_record_id = timdex_record_id
self.run_id = run_id
self.run_record_offset = run_record_offset
self.transformed_record = transformed_record

It didn't feel clear to me as to why all these fields were required to instantiate a "strategy", if that makes sense. I do understand why it is important information for the resulting EmbeddingInput returned by the BaseStrategy.to_embedding_input() method, which is why I find it fits more as parameters for that method specifically, instead of the __init__ method. However, perhaps there is something I'm not thinking of re: other methods that you expect to define for BaseStrategy classes that will heavily rely on these fields. 🤔

It's an excellent question, and one I don't have an excellent answer for. I'm glad you caught and surfaced this, as this was something I had intended to return to but did not. I think what we have here is a pretty classic major reshuffling during work on this, and this is an awkward leftover that needs attention.

One thing is known: we want strategies to emit fully formed EmbeddingInput's that can be passed directly to an embedding class's create_embedding() method. What I think is wrong with the current implementation -- and I think you touch on this -- is instantiating the transformer strategy class in a per record fashion, hence those record level details.

I'm going to take another pass at this with this discussion in mind, and will re-tag everyone for a review. Thanks for raising this @jonavellecuerdo!

[ghukill]

@jonavellecuerdo, @ehanson8 - I've added a new commit per the discussion above: 421ba71.

Thanks again @jonavellecuerdo, that rough edge completely slipped my mind and/or I couldn't see it anymore when working through this. This further simplifies the transformer strategies as just that: accept a TIMDEX JSON record, be opinionated about how to return a string representation based on the strategy, and then get out of the way.

Because we are utilizing iterators most everywhere in the CLI, there was a need for the embeddings.strategies.processor.create_embedding_inputs() function to yield EmbeddingInput's. I think this change makes that function "do" a bit more, and that also feels correct. It actually is creating EmbeddingInput's objects now, by combining the record details with the text-to-embed from the transformer strategy.

[ehanson8]

It got even cleaner and tighter! Great comment @jonavellecuerdo and great code in response @ghukill !

[jonavellecuerdo]

Changes look great and code flows better! Great work. :D