[Prefactor] Return keyed BatchedEmbeddingResult from image embedding generators#1563
[Prefactor] Return keyed BatchedEmbeddingResult from image embedding generators#1563JonasWurst wants to merge 2 commits into
Conversation
Prefactor for per-item broken-file tolerance (LIG-10034). The image and image-crop embedding paths now take (sample_id, filepath) / (sample_id, crop) pairs and return BatchedEmbeddingResult(embeddings, keys): each embedding is tagged with the sample id that produced it. This lets callers map results back to their samples by identity instead of by position, so the manager no longer keeps parallel sample_id/filepath lists in lock-step. No behavior change: every input is embedded and keys covers them all. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Repository UI Review profile: CHILL Plan: Pro Run ID: 📒 Files selected for processing (2)
🚧 Files skipped from review as they are similar to previous changes (1)
📝 WalkthroughWalkthroughThis PR changes image embedding APIs to accept UUID-keyed inputs and return ChangesKeyed embedding API refactor
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 27cc280107
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| import logging | ||
| from dataclasses import dataclass | ||
| from uuid import UUID | ||
| from uuid import UUID, uuid4 |
There was a problem hiding this comment.
The Python guideline in ai_guidelines/python.md asks function imports to use the containing module with dot notation, except for typing helpers. This new direct uuid4 import should be import uuid with uuid.uuid4() at the call site to keep production imports consistent.
Useful? React with 👍 / 👎.
| @@ -1,5 +1,7 @@ | |||
| from __future__ import annotations | |||
|
|
|||
| from uuid import uuid4 | |||
There was a problem hiding this comment.
The Python guideline in ai_guidelines/python.md asks function imports to use the containing module with dot notation, except for typing helpers. This added test import should use import uuid and call uuid.uuid4() instead of importing uuid4 directly.
Useful? React with 👍 / 👎.
| @@ -1,4 +1,5 @@ | |||
| from pathlib import Path | |||
| from uuid import uuid4 | |||
There was a problem hiding this comment.
The Python guideline in ai_guidelines/python.md asks function imports to use the containing module with dot notation, except for typing helpers. This added test import should use import uuid and call uuid.uuid4() instead of importing uuid4 directly.
Useful? React with 👍 / 👎.
| @@ -1,4 +1,5 @@ | |||
| from pathlib import Path | |||
| from uuid import uuid4 | |||
There was a problem hiding this comment.
The Python guideline in ai_guidelines/python.md asks function imports to use the containing module with dot notation, except for typing helpers. This added test import should use import uuid and call uuid.uuid4() instead of importing uuid4 directly.
Useful? React with 👍 / 👎.
| if kept.batches | ||
| else np.empty((0, context.embedding_dimension), dtype=np.float32) | ||
| ) | ||
| return BatchedEmbeddingResult(embeddings=embeddings, keys=kept.keys) |
There was a problem hiding this comment.
This now returns kept.keys in grouped encoding order rather than the original input order, but the function docstring still says it preserves input order. Please update that wording so callers do not assume result.embeddings is input-aligned when the row order must be read from result.keys.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
lightly_studio/src/lightly_studio/dataset/image_embedding.py (1)
86-95: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicate empty-input/validation logic with
image_crop_embedding.py.The "if not total_images: return empty BatchedEmbeddingResult" + "if context.max_batch_size <= 0: raise ValueError" block is duplicated verbatim in
embed_image_crops_batched(image_crop_embedding.py, lines 47-56). Since both helpers are slated for coordinated changes to support per-item broken-file skipping (per PR objectives), extracting a shared helper now would reduce the risk of the two implementations drifting apart.♻️ Proposed shared helper
+def _validate_batch_config( + total_items: int, embedding_dimension: int, max_batch_size: int +) -> BatchedEmbeddingResult | None: + """Return an empty result if there's nothing to embed, else validate config.""" + if not total_items: + return BatchedEmbeddingResult( + embeddings=np.empty((0, embedding_dimension), dtype=np.float32), + keys=[], + ) + if max_batch_size <= 0: + raise ValueError("max_batch_size must be positive.") + return NoneThen in both
embed_image_files_batchedandembed_image_crops_batched:- total_images = len(keyed_filepaths) - if not total_images: - return BatchedEmbeddingResult( - embeddings=np.empty((0, context.embedding_dimension), dtype=np.float32), - keys=[], - ) - - if context.max_batch_size <= 0: - raise ValueError("max_batch_size must be positive.") + total_images = len(keyed_filepaths) + if empty_result := _validate_batch_config( + total_images, context.embedding_dimension, context.max_batch_size + ): + return empty_result🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lightly_studio/src/lightly_studio/dataset/image_embedding.py` around lines 86 - 95, The empty-input and max-batch-size validation in embed_image_files_batched is duplicated in embed_image_crops_batched, so factor this shared setup into a common helper used by both image_embedding and image_crop_embedding. Move the “return empty BatchedEmbeddingResult when there are no items” and “raise ValueError when context.max_batch_size <= 0” logic into the shared helper, then call it from both embed_image_files_batched and embed_image_crops_batched to keep the coordinated per-item skipping changes consistent.lightly_studio/src/lightly_studio/dataset/embedding_manager.py (1)
292-298: 🩺 Stability & Availability | 🔵 Trivial | 💤 Low valueNote:
result.embeddings[0]will need a length/result.keyscheck once broken-file skipping lands.Currently safe since this PR guarantees exactly one embedding is returned for one input. Once the planned per-item skip logic (LIG-10034) is wired into
embed_image_files_batched, a broken file here would yield an emptyresult.embeddings, causing anIndexErrorinstead of a clear error message.Not a blocker now, but worth tracking for the follow-up PR.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@lightly_studio/src/lightly_studio/dataset/embedding_manager.py` around lines 292 - 298, The single-item embedding path in `embed_image` currently assumes `result.embeddings[0]` always exists; when broken-file skipping is added to `embed_image_files_batched`, a failed file could produce an empty result and raise `IndexError`. Update `embed_image` to validate `result.keys`/`result.embeddings` after `model.embed_images(...)`, and raise a clear error (or otherwise handle the missing embedding) before indexing into `result.embeddings[0]`.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@lightly_studio/tests/dataset/test_embedding_generator.py`:
- Around line 11-32: Add coverage in TestRandomEmbeddingGeneratorCrops for
RandomEmbeddingGenerator.embed_images, since only embed_image_crops currently
checks the keyed contract. Create a test that passes keyed filepaths through
embed_images and assert the returned result.keys preserves the same input keys,
using the existing RandomEmbeddingGenerator and result structure as the
reference point.
---
Nitpick comments:
In `@lightly_studio/src/lightly_studio/dataset/embedding_manager.py`:
- Around line 292-298: The single-item embedding path in `embed_image` currently
assumes `result.embeddings[0]` always exists; when broken-file skipping is added
to `embed_image_files_batched`, a failed file could produce an empty result and
raise `IndexError`. Update `embed_image` to validate
`result.keys`/`result.embeddings` after `model.embed_images(...)`, and raise a
clear error (or otherwise handle the missing embedding) before indexing into
`result.embeddings[0]`.
In `@lightly_studio/src/lightly_studio/dataset/image_embedding.py`:
- Around line 86-95: The empty-input and max-batch-size validation in
embed_image_files_batched is duplicated in embed_image_crops_batched, so factor
this shared setup into a common helper used by both image_embedding and
image_crop_embedding. Move the “return empty BatchedEmbeddingResult when there
are no items” and “raise ValueError when context.max_batch_size <= 0” logic into
the shared helper, then call it from both embed_image_files_batched and
embed_image_crops_batched to keep the coordinated per-item skipping changes
consistent.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 91a9b065-fa11-4b4c-a94b-166e6d17bcad
📒 Files selected for processing (12)
lightly_studio/src/lightly_studio/dataset/embedding_generator.pylightly_studio/src/lightly_studio/dataset/embedding_manager.pylightly_studio/src/lightly_studio/dataset/image_crop_embedding.pylightly_studio/src/lightly_studio/dataset/image_embedding.pylightly_studio/src/lightly_studio/dataset/mobileclip_embedding_generator.pylightly_studio/src/lightly_studio/dataset/perception_encoder_embedding_generator.pylightly_studio/tests/dataset/test_embedding_generator.pylightly_studio/tests/dataset/test_embedding_manager.pylightly_studio/tests/dataset/test_image_crop_embedding.pylightly_studio/tests/dataset/test_image_embedding.pylightly_studio/tests/dataset/test_mobileclip_embedding_generator.pylightly_studio/tests/dataset/test_perception_encoder_embedding_generator.py
| class TestRandomEmbeddingGeneratorCrops: | ||
| def test_embed_image_crops__returns_one_embedding_per_crop(self) -> None: | ||
| generator = RandomEmbeddingGenerator(dimension=4) | ||
| image_crops = [ | ||
| ImageCrop(filepath="a.jpg", x=0, y=0, width=10, height=10), | ||
| ImageCrop(filepath="a.jpg", x=5, y=5, width=20, height=20), | ||
| ImageCrop(filepath="b.jpg", x=0, y=0, width=30, height=30), | ||
| keys = [uuid4(), uuid4(), uuid4()] | ||
| keyed_crops = [ | ||
| (keys[0], ImageCrop(filepath="a.jpg", x=0, y=0, width=10, height=10)), | ||
| (keys[1], ImageCrop(filepath="a.jpg", x=5, y=5, width=20, height=20)), | ||
| (keys[2], ImageCrop(filepath="b.jpg", x=0, y=0, width=30, height=30)), | ||
| ] | ||
|
|
||
| embeddings = generator.embed_image_crops(image_crops) | ||
| result = generator.embed_image_crops(keyed_crops) | ||
|
|
||
| assert embeddings.shape == (3, 4) | ||
| assert result.embeddings.shape == (3, 4) | ||
| assert result.keys == keys | ||
|
|
||
| def test_embed_image_crops__empty_input_returns_empty_array(self) -> None: | ||
| generator = RandomEmbeddingGenerator(dimension=4) | ||
|
|
||
| embeddings = generator.embed_image_crops([]) | ||
| result = generator.embed_image_crops([]) | ||
|
|
||
| assert embeddings.shape == (0, 4) | ||
| assert result.embeddings.shape == (0, 4) | ||
| assert result.keys == [] |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check for tests exercising RandomEmbeddingGenerator.embed_images with the new keyed API.
rg -nP -C5 '\bembed_images\s*\(' lightly_studio/tests/dataset/test_embedding_generator.py
rg -nP 'class Test' lightly_studio/tests/dataset/test_embedding_generator.pyRepository: lightly-ai/lightly-studio
Length of output: 207
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== File contents =="
cat -n lightly_studio/tests/dataset/test_embedding_generator.py
echo
echo "== Search for embed_images coverage =="
rg -n -C 3 'embed_images\s*\(' lightly_studio/tests
echo
echo "== Search for RandomEmbeddingGenerator references =="
rg -n -C 3 'RandomEmbeddingGenerator' lightly_studio/testsRepository: lightly-ai/lightly-studio
Length of output: 27487
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== Direct RandomEmbeddingGenerator embed_images test references =="
rg -n -C 3 'RandomEmbeddingGenerator.*embed_images|embed_images.*RandomEmbeddingGenerator' lightly_studio/tests lightly_studio | sed -n '1,200p'
echo
echo "== Any test names or assertions mentioning result.keys for embed_images =="
rg -n -C 2 'result\.keys|keys ==|keys\)' lightly_studio/tests/dataset/test_*embedding_generator.py lightly_studio/tests/dataset/test_embedding_manager.py | sed -n '1,240p'Repository: lightly-ai/lightly-studio
Length of output: 230
Add embed_images coverage for RandomEmbeddingGenerator The keyed contract is only exercised for embed_image_crops here; add a test for embed_images that passes keyed filepaths and asserts result.keys matches the input keys.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@lightly_studio/tests/dataset/test_embedding_generator.py` around lines 11 -
32, Add coverage in TestRandomEmbeddingGeneratorCrops for
RandomEmbeddingGenerator.embed_images, since only embed_image_crops currently
checks the keyed contract. Create a test that passes keyed filepaths through
embed_images and assert the returned result.keys preserves the same input keys,
using the existing RandomEmbeddingGenerator and result structure as the
reference point.
1 participant
What
Prefactor for LIG-10034 (per-item broken-file tolerance). The image and image-crop embedding paths now take
(sample_id, filepath)/(sample_id, crop)pairs and returnBatchedEmbeddingResult(embeddings, keys: list[UUID]): each embedding is tagged with the sample id that produced it.Why
Today the embedding helpers return a bare array aligned by position, so the manager keeps parallel
sample_ids/filepathslists in lock-step and re-associates by integer index. That's brittle — any reordering silently mismaps embeddings to samples. Carrying the sample id with each input lets callers map results back by identity, which is what the upcoming broken-file skipping (LIG-10034) needs: a skipped file just omits its id fromkeys.Behavior
No behavior change. Every input is embedded and
keyscovers them all; the manager mapsresult.keysstraight to storage.🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes
Tests