Support Embeddings in mltransform #29564

AnandInguva · 2023-11-29T19:04:04Z

Support embeddings in MLTransform using config based structure.

Each config will define its own model handler so that model handler can be compatible with the inputs passed to the MLTransform.
This PR supports vertex_ai, hugging face(sentence-transformers).

Changes made:

Make ProcessHandler a PTransform.
Create _TextEmbeddingHandler which takes a ModelHandler but it is responsible to work on Dict[str, Any] inputs.
Create a helper class that would create and saves the PTransform list required to process data processing steps or embedding configs. MLTransform would be the container that holds the list of transforms and these list of transforms are passed to the _MLTransformToPTransformMapper which maps each data processing transforms to its PTransform
- For example, tft based transforms are mapped to TFTProcessHandler.
- Embedding configs are mapped to RunInference
Use jsonpickle to store and load the ptransforms instances. This is done to seal the gap between training and inference. jsonpickle is cross compatible across python versions and has backward compatibility with older versions of its versions.

This PR will throw if input is a Dict[str, List[str]]. Works for Dict[str, str]

Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:

Mention the appropriate issue in your description (for example: addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, comment fixes #<ISSUE NUMBER> instead.
Update CHANGES.md with noteworthy changes.
If this contribution is large, please file an Apache Individual Contributor License Agreement.

See the Contributor Guide for more tips on how to make review process smoother.

To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md

GitHub Actions Tests Status (on master branch)

See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.

Fix tox.ini Fix pydoc Fix indent in pydoc

AnandInguva · 2023-11-29T19:58:06Z

R: @damccorm This is ready for review.

github-actions · 2023-11-29T19:59:21Z

Stopping reviewer notifications for this pull request: review requested by someone other than the bot, ceding control

damccorm

Chunking the review a bit. Just reviewed base.py thus far, though I imagine that's the place I'll have the most comments

sdks/python/apache_beam/ml/transforms/base.py

sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py

sdks/python/apache_beam/ml/transforms/embeddings/sentence_transformer.py

sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub.py

sdks/python/apache_beam/ml/transforms/embeddings/sentence_transformer.py

sdks/python/apache_beam/ml/transforms/embeddings/sentence_transformer_test.py

sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py

sdks/python/apache_beam/ml/transforms/utils.py

… on different machines

This reverts commit cfb1883.

damccorm

This mostly looks good, had a few more comments though

sdks/python/apache_beam/ml/transforms/base.py

sdks/python/apache_beam/ml/transforms/tft.py

sdks/python/apache_beam/ml/transforms/base.py

sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py

sdks/python/apache_beam/ml/transforms/utils.py

sdks/python/tox.ini

…va/beam into embeddings_mltransform

damccorm · 2023-12-08T20:04:11Z

sdks/python/apache_beam/ml/transforms/base.py

+      raise FileExistsError(
+          "The artifact location %s already exists and contains %s. Please "
+          "specify a different location." %
+          (artifact_location, _ATTRIBUTE_FILE_NAME))


Good call - one possible future enhancement would be to support an overwrite argument that allows users to do this

sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai_test.py

sdks/python/apache_beam/ml/transforms/base.py

sdks/python/apache_beam/ml/transforms/embeddings/vertex_ai.py

sdks/python/apache_beam/ml/transforms/utils.py

damccorm

One last nit, otherwise this LGTM. Feel free to make the change and merge once checks pass

CHANGES.md

Co-authored-by: Danny McCormick <dannymccormick@google.com>

* Make base.py framework agnostic and add helper transforms * Add tests for base.py * Add sentence-transformers * Add tensorflow hub * Add vertex_ai * Make TFTProcessHandler a PTransform * raise RuntimeError in ArtifactsFetcher when it is used for embeddings * Add JsonPickle to requirements * Add tox tests * Mock frameworks in pydocs Fix tox.ini Fix pydoc Fix indent in pydoc * Add Row type check * Remove requires_chaining * change name of PTransformProvider to MLTransformProvider * remove batch_len in utility fun * Change type annotation and redundant comments * Remove get_transforms method * remove requires_chaining from tft * add tests to sentence-transformers * Pass inference_args to RunInference * Add TODO GH issue * refactor variables in vertex_ai embeddings * remove try/catch and throw error if options is empty for GCS artifact location * Refactor NotImplementedError message * remove tensorflow hub from this PR * Add _validate_transform method * add more tests * fix test * Fix test * Add more tests in sentence-transformer * use np.max instead of max * round to 2 decimals * Remove gradle command action * Refactor throwing dataflow client exception * skip the test if gcp is not installed * remove toxTests for hub * remove toxTests for hub * Fix values in assert for sentence_transformer_test * rename sentence_transformers to huggingface * fix pydocs * Change the model name for tests since it is getting different results on different machines * Fix pydoc in vertexai * add suffix to artifact_location * Revert "add suffix to artifact_location" This reverts commit cfb1883. * add no_xdist * Try fixing pydoc for vertexai * change tox.ini to use pytest directly * raise FileExistError if Attribute file is already present * modify build.gradle to match tox task names * Add note to CHANGES.md * change gcs bucket to gs://temp-storage-for-perf-tests * Add TODO GH links * Update CHANGES.md Co-authored-by: Danny McCormick <dannymccormick@google.com> --------- Co-authored-by: Danny McCormick <dannymccormick@google.com>

AnandInguva added 10 commits November 29, 2023 13:54

Make base.py framework agnostic and add helper transforms

80e5c4a

Add tests for base.py

0d34847

Add sentence-transformers

58b24f6

Add tensorflow hub

88f9ceb

Add vertex_ai

23f7027

Make TFTProcessHandler a PTransform

04ebdb0

raise RuntimeError in ArtifactsFetcher when it is used for embeddings

f86c259

Add JsonPickle to requirements

fc4ec00

Add tox tests

3da5ce8

Mock frameworks in pydocs

4b4ee58

Fix tox.ini Fix pydoc Fix indent in pydoc

github-actions bot added python docker labels Nov 29, 2023

AnandInguva requested a review from damccorm November 29, 2023 19:57

damccorm reviewed Dec 1, 2023

View reviewed changes

AnandInguva added 2 commits December 4, 2023 18:25

Add Row type check

01ba217

Remove requires_chaining

f080c25

AnandInguva commented Dec 4, 2023

View reviewed changes

sdks/python/apache_beam/ml/transforms/embeddings/tensorflow_hub_test.py Outdated Show resolved Hide resolved

AnandInguva added 5 commits December 4, 2023 18:39

change name of PTransformProvider to MLTransformProvider

6111c31

remove batch_len in utility fun

ba24e81

Change type annotation and redundant comments

d690aec

Remove get_transforms method

af7496b

remove requires_chaining from tft

d713555

damccorm reviewed Dec 4, 2023

View reviewed changes

AnandInguva and others added 2 commits December 4, 2023 22:42

add tests to sentence-transformers

50450f3

Merge remote-tracking branch 'origin/master' into embeddings_mltransform

c2b691f

github-actions bot added the build label Dec 5, 2023

AnandInguva and others added 2 commits December 5, 2023 15:24

Pass inference_args to RunInference

8823a75

Add TODO GH issue

a7e2bd3

AnandInguva and others added 9 commits December 7, 2023 18:34

fix pydocs

5cae04b

Change the model name for tests since it is getting different results…

489200f

… on different machines

Fix pydoc in vertexai

816174a

add suffix to artifact_location

cfb1883

Revert "add suffix to artifact_location"

2cb6f03

This reverts commit cfb1883.

add no_xdist

cd7050e

Try fixing pydoc for vertexai

98cd949

change tox.ini to use pytest directly

8ea0906

Merge remote-tracking branch 'origin/master' into embeddings_mltransform

5187b0e

damccorm reviewed Dec 8, 2023

View reviewed changes

AnandInguva mentioned this pull request Dec 8, 2023

[Task]: Support embeddings using ML models in MLTransform #29356

Closed

16 tasks

AnandInguva added 4 commits December 8, 2023 20:00

raise FileExistError if Attribute file is already present

6f83d3c

Merge branch 'embeddings_mltransform' of https://github.com/AnandIngu…

c9ddb25

…va/beam into embeddings_mltransform

modify build.gradle to match tox task names

9dce3cf

Add note to CHANGES.md

539c9ad

damccorm reviewed Dec 8, 2023

View reviewed changes

change gcs bucket to gs://temp-storage-for-perf-tests

b967cd8

damccorm reviewed Dec 11, 2023

View reviewed changes

sdks/python/apache_beam/ml/transforms/utils.py Outdated Show resolved Hide resolved

sdks/python/apache_beam/ml/transforms/utils.py Outdated Show resolved Hide resolved

AnandInguva added 2 commits December 11, 2023 09:55

Add TODO GH links

f1bb42c

Merge remote-tracking branch 'origin/master' into embeddings_mltransform

8d0b47d

damccorm approved these changes Dec 11, 2023

View reviewed changes

CHANGES.md Outdated Show resolved Hide resolved

Update CHANGES.md

c173d6a

Co-authored-by: Danny McCormick <dannymccormick@google.com>

AnandInguva merged commit aacf1ee into apache:master Dec 11, 2023
78 of 79 checks passed

AnandInguva linked an issue Jan 5, 2024 that may be closed by this pull request

[Task]: Support Vertex AI model garden for Vertex AI embeddings #29805

Open

16 tasks

AnandInguva removed a link to an issue Jan 5, 2024

[Task]: Support Vertex AI model garden for Vertex AI embeddings #29805

Open

16 tasks

AnandInguva linked an issue Jan 5, 2024 that may be closed by this pull request

[Task]: Support embeddings using ML models in MLTransform #29356

Closed

16 tasks

tvalentyn mentioned this pull request Feb 1, 2024

[Failing Test]: MLTransform IT test is failing #30191

Closed

16 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Support Embeddings in mltransform #29564

Support Embeddings in mltransform #29564

AnandInguva commented Nov 29, 2023 •

edited

AnandInguva commented Nov 29, 2023

github-actions bot commented Nov 29, 2023

damccorm left a comment

damccorm left a comment

damccorm Dec 8, 2023

damccorm left a comment

Support Embeddings in mltransform #29564

Support Embeddings in mltransform #29564

Conversation

AnandInguva commented Nov 29, 2023 • edited

GitHub Actions Tests Status (on master branch)

AnandInguva commented Nov 29, 2023

github-actions bot commented Nov 29, 2023

damccorm left a comment

Choose a reason for hiding this comment

damccorm left a comment

Choose a reason for hiding this comment

damccorm Dec 8, 2023

Choose a reason for hiding this comment

damccorm left a comment

Choose a reason for hiding this comment

AnandInguva commented Nov 29, 2023 •

edited