Skip to content

feat: add support for the embedding batch task - #516

Merged
cswaney merged 2 commits into
mainfrom
feat/embed-batch-task
Aug 27, 2026
Merged

feat: add support for the embedding batch task#516
cswaney merged 2 commits into
mainfrom
feat/embed-batch-task

Conversation

@cswaney

@cswaney cswaney commented Aug 27, 2026

Copy link
Copy Markdown
Member

Summary

  • Registers the tigerflow-ml `embed` task (added in 0.2.0, PR Release v0.5.0 - Vite Migration #178/Slurm profile using localhost still needs to use profile home directory #180) so batch jobs can produce vector embeddings for text, images, and PDFs.
  • Maps `embed` -> `tigerflow_ml.multimodal.embed.local` in `SUPPORTED_TASKS`, defaults input to `.txt` and output to `.npy` (the task rejects any other suffix).
  • Wires the launcher: `COMPATIBLE_PIPELINES["embedding"]` -> `sentence-similarity` + `feature-extraction`; `NewJobModal` gets an "embed" entry, `multiPageOnly` param gate, and `applyEmbedEncodeMode` helper.
  • Adds `embed` (and pre-existing gap `chat`) to `EXPECTED_TASKS` in `test_tigerflow.py` so the contract test guards them against future image bumps.

Test plan

  • `uv run just lint`
  • `uv run just test` (972 pass, 7 skipped)
  • `npm test` (552 pass)
  • End-to-end on Della: run an embed job over a directory of `.txt` files, confirm `.npy` outputs resolve in the results table.

Notes

  • Depends on tigerflow-ml 0.2.0, pinned in chore(lib): bump tigerflow-ml to 0.2.0 #510.
  • Rebased onto today's `main`; the branch's original chat/embed under `text/` was updated to the 0.2.0 module layout (`multimodal/`) as part of the rebase.
  • Dropped three tautological tests (`test_embed_task_is_supported`, `test_embed_default_input_ext_is_text`, `test_embed_default_output_ext_is_npy`) that just re-read `SUPPORTED_TASKS` / `DEFAULT_*_EXT`. Kept the two load-bearing ones (`test_build_pipeline_config_builds_embed_task`, `test_embed_pipeline_yaml_sets_npy_output_ext`) that exercise real assembly / yaml rendering.
  • Removed the "Until the pinned image is updated" note from the CLI docs (moot now).

Closes #463

Register the tigerflow-ml embed task and add a launcher entry (text and
image/PDF inputs, encode mode, normalize, truncate dimension).

Default embed output to .npy: the task rejects any other suffix, so an
unset default would make tigerflow write .out and fail every file.

Expose encode mode as one select rather than two booleans, since the task
raises when use_encode_document and use_encode_query are both true.

Closes #463.
@claude

claude Bot commented Aug 27, 2026

Copy link
Copy Markdown

Review

Nicely scoped change — task registry, launcher UI, and docs are all updated together, and the applyEmbedEncodeMode trick for collapsing two mutually-exclusive booleans into one select is a good pattern (mirrors applyTranslateSourceLang). Test coverage for the new pure helpers (applyEmbedEncodeMode, isParamVisible w/ multiPageOnly) is thorough.

Possible bug: encode_kwargs is collected as a raw string but never parsed to JSON

NewJobModal.jsx adds an encode_kwargs textarea param whose help text explicitly promises JSON input (Additional encode() kwargs as JSON, e.g. {"prompt": "query: "}), but:

  • It has no valueType (unlike normalize/truncate_dim/batch_size), so coerceParamValue leaves it as a plain string (NewJobModal.jsx around line 627-640).
  • There's no client-side JSON.parse/validation anywhere in the submit path — a malformed value (trailing comma, unquoted keys) is silently accepted and submitted as-is.
  • On the backend, build_pipeline_config (tasks.py:125-126) just drops params straight into the pipeline dict with no parsing either.

So a user-entered {"prompt": "query: "} ends up serialized into the job's YAML as the string '{"prompt": "query: "}', not a mapping. If tigerflow_ml.multimodal.embed.local does something like model.encode(**encode_kwargs), a string there would blow up at runtime (can't **-unpack a str) rather than failing at submission with a clear error.

This may be fine if the tigerflow-ml task itself does its own json.loads on string input (unverifiable from this repo), but there's no test anywhere — frontend or backend — that exercises encode_kwargs with an actual value, and the PR's own test plan only checks a plain .txt embed run, not one using this field. Worth either confirming tigerflow-ml parses encode_kwargs when it arrives as a string, or parsing/validating the JSON client-side (e.g., JSON.parse on submit with an inline error) before sending it through, consistent with how normalize/truncate_dim are coerced to real types.

Minor / non-blocking

  • lib/docs/usage/cli.md's "Supported tasks" table still doesn't include chat, even though this PR is already touching that exact table to add embed and separately fixes the chat gap in the contract test's EXPECTED_TASKS. Since you're already here, might be worth adding the chat row too (pre-existing gap, not introduced by this PR).
  • The new embed row in that same markdown table isn't column-aligned with the others (cosmetic only — GitHub/mkdocs render pipe tables fine regardless of alignment).

Everything else looked solid

  • SUPPORTED_TASKS / DEFAULT_INPUT_EXT / DEFAULT_OUTPUT_EXT wiring in tasks.py is consistent with the other tasks, and the comment explaining why .npy must be a hard default (task rejects any other suffix) is genuinely useful context.
  • COMPATIBLE_PIPELINES["embedding"]sentence-similarity/feature-extraction reuses the existing generic lookup in asgi.py (both get_models query sites), so model filtering by service should work correctly.
  • multiPageOnly gating (MULTI_PAGE_INPUT_EXTS = new Set([".pdf"])) correctly scopes batch_size visibility to PDF-only inputs, matching tigerflow's count_images behavior per the comment.
  • Dropping the three tautological tests while keeping the two load-bearing ones (per the PR notes) is a reasonable call — they weren't testing behavior, just restating the dict literals.

@cswaney

cswaney commented Aug 27, 2026

Copy link
Copy Markdown
Member Author

`encode_kwargs` — false positive on the parsing concern.

Verified against `tigerflow_ml.multimodal.embed._base.Params` at v0.2.0:

```python
encode_kwargs: Annotated[
str,
typer.Option(help="Additional kwargs for SentenceTransformer's encode() (e.g., {'prompt':'query: '})..."),
] = "{}"
```

The field is declared as `str` on purpose. Line 105 hands it to `parse_kwargs`, which tries `json.loads` first, then Python literal syntax. So the string form is expected, and — notably — the help text example (`{'prompt':'query: '}`) is Python-literal, not JSON. Adding client-side `JSON.parse` in Blackfish would be stricter than what tigerflow-ml accepts and would reject the very syntax the help text advertises. Leaving it as-is; malformed input surfaces as a clear tigerflow-ml error at task setup.

`chat` row in CLI docs table — fixed. Added in ac42a00 while we were already editing that table.

embed row misalignment — skipping. Cosmetic; pipe tables render fine.

@cswaney
cswaney merged commit d7a7d0b into main Aug 27, 2026
11 checks passed
@cswaney
cswaney deleted the feat/embed-batch-task branch August 27, 2026 17:33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Support embedding batch jobs

1 participant