feat(embed): EmbeddingGemma 2 as a local, text-only embedding provider

[laithalsaadoon]

Embed and search on this machine with EmbeddingGemma 2

atif-sql embed and atif-sql search send step text and queries to Cohere Embed v4 on Bedrock, which bills an AWS account and takes the text off the machine. This adds Google's EmbeddingGemma 2 (google/embeddinggemma-2, Apache 2.0, released 2026-10-06) as a second embedding provider, run locally in its text-only setup. ATIF_SQL_EMBED_PROVIDER=gemma selects it; Cohere stays the default and nothing changes for an existing store.

uv tool install 'atif-sql[local]' --torch-backend cpu   # or, in a checkout: uv sync --all-packages --extra local
export ATIF_SQL_EMBED_PROVIDER=gemma
atif-sql embed --limit 200
atif-sql search 'the flaky watermark test'

What changed

• One optional extra, local (build(deps) commit), on the embed package and the published root: sentence-transformers>=6.1.0, transformers>=5.19.0, torch>=2.14, torchvision>=0.29, pillow>=12.3.0. A plain uv sync --all-packages, CI included, stays torch-free.
  • transformers 5.19.0 is the first release with EmbeddingGemma 2's model class and was uploaded 2026-10-06T16:38:56Z, inside the seven-day cooldown, so it has an exclude-newer-package exemption to 2026-10-06T16:39:00Z, to delete after 2026-10-13T16:39Z.
  • pillow and torchvision are needed even for text: sentence-transformers builds the model's AutoProcessor, and a text-only load without them fails with "requires the PIL library", then "No module named 'torchvision'".
  • torch and torchvision on Linux resolve from PyTorch's CPU index (explicit = true, marker sys_platform == 'linux'), so uv.lock holds 2.14.1+cpu and 0.29.1+cpu and no nvidia-*, cuda-* or triton package (grep -c gives 0). The index publishes PEP 700 upload times (data-upload-time), so the cooldown applies to it as to PyPI. macOS and Windows take PyPI's CPU builds.
  • test_distribution.py holds the root's extras equal to the members' extras, as it does for dependencies.
• The adapter (feat(embed) commit), atif_embed/infrastructure/gemma_local.py:
  • loads the model once, lazily, through sentence-transformers with config_kwargs={"vision_config": None, "audio_config": None} (270M of 740M parameters) at pinned revision 914f7f89142e33e77833254d9c9b90c3cef7303b;
  • bfloat16 on a CUDA device that supports it, float32 elsewhere, never float16 (the model card: float16 returns NaN or degraded vectors without an error); a non-finite vector is dropped as a failed slot;
  • documents encode under the model's Document prompt and queries under SearchQuery, with truncate_dim and normalize_embeddings=True; widths 768 (default), 512, 256, 128;
  • a failing batch leaves None slots, the port's bounded-loss contract;
  • batches form shortest first under a padded-character budget, and the loader caps max_seq_length at the model's 8,192-token window. Both came from the smoke run: the tokenizer declares no maximum, and 32 long texts padded into one batch took 26 GB of RSS on CPU.
• A store per provider. Cohere keeps <corpus_root>/embeddings_lance; Gemma writes <corpus_root>/embeddings_lance_gemma. Switching never meets the (model, dim) guard and both stores coexist; ATIF_SQL_LANCE_URI still names one store for either. status reports embed_provider and embeddings_store, and scripts/atif-sql-refresh.sh reads that path for its terminal-marker check, falling back to Cohere's path for an older CLI.
• Wiring and errors. build_embedder(settings) is the one place that names an adapter, for both backfill and search. The dry-run plan adds provider, dim and store. Gemma without the extra raises a terminal domain error naming the install command (exit 78); a width the provider cannot emit exits 64 instead of a pydantic traceback; search classifies a failed query embedding (70 or 78) instead of raising a traceback.
• A fix found on the way. ATIF_SQL_OUTPUT_DIMENSION never parsed from the environment: as an integer Literal, pydantic refuses the string an environment variable carries (Input should be 256, 512, 1024 or 1536 [type=literal_error, input_value='512', input_type=str] on main). It is an int checked against the selected provider's widths.
• Lean import. torch, transformers and sentence_transformers join the forbidden list of test_lean_import.py, and a new probe checks the Cohere path imports none of them. The adapter imports them by name, so ty and pyright pass without the extra.
• Docs (docs commit): README install line, a "What you can do" row, a "Search without sending text off the machine" section; CLI reference for embed and search; the embed package README, AGENTS.md, CONTRIBUTING.md, SECURITY.md, the site index and CONTRACT-V2.

Verification

• mise run check green on each of the three commits from a fresh worktree, with mise run lock:check after it:
  • build(deps): test 2495 passed, 1 skipped, 1 xfailed; pyright 0 errors; security:vex:check green; lock:check "Resolved 156 packages".
  • feat(embed): test 2541 passed, 1 skipped, 1 xfailed; pyright 0 errors; lock:check green.
  • docs: test 2541 passed, 1 skipped, 1 xfailed; docs:prose 0 errors at the 273 baseline; lock:check green.
  • mise run test:cov on the head: 90.47% against the 89% floor; gemma_local.py 95%, providers.py and settings.py 100%.
• mise run docs:install && mise run docs:gate: 0 broken internal links over 20 pages, astro check 0 errors, 73 vitest tests passed. docs:prose: 0 errors, 273 warnings, at the baseline.
• scripts/atif-sql-refresh-selftest.sh: 0 failures.
• osv-scanner over the new uv.lock (torch 2.14.1, transformers 5.19.0, sentence-transformers 6.1.0, pillow 12.3.0, torchvision 0.29.1 and the rest): "No issues found".
• Tests use a fake SentenceTransformer and fake torch/sentence_transformers modules, so they need no download or network and pass with or without the extra; 62 of them also ran in a venv with the extra installed. Planted defects: returning torch.float16 on CUDA turns test_device_and_dtype_never_float16 red; encoding queries under the Document prompt turns three query tests red; dropping torch from the root extra turns test_root_extras_equal_the_members_extras red.
• Real smoke on CPU (the lock's 2.14.1+cpu, float32, HF_HUB_OFFLINE=1 after one download of 1,488,915,288 bytes), 20 sessions copied from a Claude Code corpus, --no-lake:
  • embed --dry-run: provider gemma, model google/embeddinggemma-2, dim 768, 919 candidates, batch 32, concurrency 1, store .../embeddings_lance_gemma.
  • embed --limit 200: 200 rows in 156 s of encoding (1.3 vectors/s; 170.6 s wall with the load), peak RSS 3.9 GB, on a shared 16-core host with load average 12 to 20. Short steps run at about 30 vectors/s (32 texts of 150 characters in 1.04 s); the 8 steps over 2,000 characters take most of the time. The store holds 200 rows stamped google/embeddinggemma-2, dim 768, fixed_size_list<float>[768], unit-norm vectors.
  • search with a paraphrase of one embedded step (it shares only "forecast bands" and "instrument" with the step, which says CUDA where the query says GPU) ranks that step first at sim 0.7794, next 0.7247; 9.8 s wall and 2.3 GB RSS, most of it the model load.
  • uv tool install --torch-backend cpu "atif-sql[local] @ file://<checkout>" installs torch 2.14.1+cpu and torchvision 0.29.1+cpu and no nvidia-* or cuda-* package.

Risks

• CPU speed on long steps. A long step near the 8,192-token window takes tens of seconds on CPU, so a full backfill of a large corpus takes hours there. The lockfile installs CPU torch; a CUDA user installs a CUDA torch build (--torch-backend auto) to get bfloat16 on the GPU, which this PR did not measure.
• The transformers exemption must be deleted in the first uv lock --upgrade after 2026-10-13T16:39Z.
• Install weight. The extra adds torch (187 MiB wheel), transformers, scikit-learn and scipy to an install that opts in; the default install is unchanged.
• Status JSON gains two keys under the vector block; nothing is removed or renamed.

What to look at

• plan_batches and BATCH_CHAR_BUDGET in gemma_local.py: the budget (32,000 padded characters per forward pass) is a measured CPU choice, not a model limit.
• Exit 78 for "gemma selected, extra missing": it is terminal so an unattended refresh lane stops retrying, the same contract as a provider mismatch.
• The refresh script's store-path change in embed_store_dir.

Sources read 2026-10-09: the model card (huggingface.co/google/embeddinggemma-2), uv's PyTorch guide (docs.astral.sh/uv/guides/integration/pytorch) and resolution docs on exclude-newer and upload times (docs.astral.sh/uv/concepts/resolution), PyPI JSON for transformers 5.19.0, sentence-transformers 6.1.0, torch 2.14.1, torchvision 0.29.1 and pillow 12.3.0, and the PyTorch CPU index's data-upload-time attributes.