text embed cache batching for big speed-up

[bghira]

This pull request adds support for batching text encoder operations when precomputing uncached text embeddings, and updates the documentation and configuration options across multiple languages to describe the new text_encoder_batch_size parameter. It also refactors and improves the batching logic in text_embeds.py to ensure correct handling of multi-prompt records and efficient caching.

Major features and improvements:

Support for batched text encoding and caching:

• Added the text_encoder_batch_size parameter to allow batching multiple captions in a single forward pass when precomputing text embeddings, improving throughput and VRAM utilization. (F9357724L77R76, [1] [2] [3]
• Implemented new methods _slice_batch_output_for_cache and _encode_and_cache_prompt_batch in simpletuner/helpers/caching/text_embeds.py to handle splitting batched encoder outputs and saving them to cache per sample.
• Improved _normalize_prompt_records to correctly handle records with multiple prompts and keys, ensuring data consistency and error checking for mismatches.

Documentation and configuration updates:

• Updated dataloader and options documentation in English, Spanish, Hindi, Japanese, Portuguese, and Chinese to describe the new text_encoder_batch_size parameter, its usage, and default behavior. [1] [2] [3] [4] [5] [6] [7] [8] [9] [10] [11] [12] [13] [14] [15] [16] [17] [18]

Default behavior and parameter handling:

• Changed the default value for text_encoder_batch_size to 1 (was previously 4), ensuring backward compatibility and predictable resource usage. (F9357724L77R76)
• Ensured that text_encoder_batch_size is always at least 1 and properly type-cast, preventing invalid configurations.

These changes collectively improve the performance and flexibility of text embedding precomputation, especially for large datasets, and clarify usage for users in all supported languages.

[Copilot]

Pull request overview

This pull request introduces configurable batching for text-encoder forward passes when precomputing uncached text embeddings, aiming to improve cache precompute throughput while keeping default behavior unchanged (batch size = 1). It also wires the new text_encoder_batch_size option through the SDK/WebUI config surface and updates user-facing documentation across all supported languages.

Changes:

• Add text_encoder_batch_size end-to-end (CLI/env mapping, backend config, SDK/WebUI field registry + dataset blueprint).
• Refactor TextEmbeddingCache to support batched encode+cache for precompute runs, including per-sample slicing of batched outputs.
• Add unit tests and update multilingual OPTIONS/DATALOADER docs to document the new parameter.

Reviewed changes

Copilot reviewed 20 out of 20 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/test_text_embeds.py	Adds coverage for multi-prompt normalization and batched encode-to-cache slicing behavior.
simpletuner/helpers/caching/text_embeds.py	Implements batched encode+cache flow and prompt-record normalization updates; clamps batch size to >= 1.
simpletuner/helpers/models/ernie/model.py	Enables tokenizer padding suitable for batched text encoding (padding="longest").
simpletuner/helpers/data_backend/factory.py	Passes text_encoder_batch_size into TextEmbeddingCache construction for backend configs.
simpletuner/helpers/configuration/env_file.py	Adds env var mapping for TEXT_ENCODER_BATCH_SIZE.
simpletuner/simpletuner_sdk/server/services/cache_job_service.py	Includes text_encoder_batch_size in job defaults and forwards it into the cache executor.
simpletuner/simpletuner_sdk/server/services/field_registry/sections/advanced.py	Registers the new option as a configurable advanced caching field in the SDK/WebUI.
simpletuner/simpletuner_sdk/server/data/dataset_blueprints.py	Exposes the new option in dataset blueprint schema for text_embeds datasets.
documentation/OPTIONS.md	Documents --text_encoder_batch_size in English CLI options output.
documentation/OPTIONS.es.md	Documents --text_encoder_batch_size in Spanish CLI options output.
documentation/OPTIONS.hi.md	Documents --text_encoder_batch_size in Hindi CLI options output.
documentation/OPTIONS.ja.md	Documents --text_encoder_batch_size in Japanese CLI options output.
documentation/OPTIONS.pt-BR.md	Documents --text_encoder_batch_size in Portuguese (Brazil) CLI options output.
documentation/OPTIONS.zh.md	Documents --text_encoder_batch_size in Chinese CLI options output.
documentation/DATALOADER.md	Adds text_encoder_batch_size to English dataloader docs for dataset_type=text_embeds.
documentation/DATALOADER.es.md	Adds text_encoder_batch_size to Spanish dataloader docs for dataset_type=text_embeds.
documentation/DATALOADER.hi.md	Adds text_encoder_batch_size to Hindi dataloader docs for dataset_type=text_embeds.
documentation/DATALOADER.ja.md	Adds text_encoder_batch_size to Japanese dataloader docs for dataset_type=text_embeds.
documentation/DATALOADER.pt-BR.md	Adds text_encoder_batch_size to Portuguese (Brazil) dataloader docs for dataset_type=text_embeds.
documentation/DATALOADER.zh.md	Adds text_encoder_batch_size to Chinese dataloader docs for dataset_type=text_embeds.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.