feat(embed): add support for audio inputs

README.md

@@ -28,6 +28,7 @@ pip install tigerflow-ml[vllm]
 | Chat             | Apply a chat prompt to images or text | `chat` / `chat-local`             |
 | Transcription    | Transcribe audio to text              | `transcribe` / `transcribe-local` |
 | Object Detection | Detect objects in images and videos   | `detect` / `detect-local`         |
+| Embed            | Embed text, images, or audio          | `embed` / `embed-local`           |
 
 Each task provides both a Slurm variant (for HPC) and a Local variant (for development).
 
@@ -47,6 +48,7 @@ python -m tigerflow_ml.text.translate.slurm --help
 python -m tigerflow_ml.text.chat.slurm --help
 python -m tigerflow_ml.audio.transcribe.slurm --help
 python -m tigerflow_ml.image.detect.slurm --help
+python -m tigerflow_ml.text.embed.slurm --help
 ```
 
 ## Container

docs/mkdocs.yml

@@ -65,3 +65,4 @@ nav:
       - Transcription: tasks/transcribe.md
       - Object Detection: tasks/detect.md
       - Chat: tasks/chat.md
+      - Embed: tasks/embed.md

docs/mkdocs/index.md

@@ -40,6 +40,7 @@ hide:
 - [Transcription](tasks/transcribe.md) — Transcribe audio to text
 - [Object Detection](tasks/detect.md) — Detect objects in images and videos
 - [Chat](tasks/chat.md) - Apply a chat prompt to text or image documents
+- [Embed](tasks/embed.md) — Embed text, images, or audio
 
 ## Installation
 

docs/mkdocs/tasks/embed.md

@@ -0,0 +1,180 @@
+# Embed
+
+Embed text, images, or audio using HuggingFace sentence-transformers models.
+
+## Parameters
+
+| Parameter          | Default            | Description                                                                                                      |
+|---------------------|--------------------|--------------------------------------------------------------------------------------------------------------------|
+| `--model`           |                    | HuggingFace model repo ID                                                                                          |
+| `--revision`        | `main`             | Model revision (branch, tag, or commit hash)                                                                       |
+| `--cache-dir`       |                    | HuggingFace cache directory for model files                                                                        |
+| `--device`          | `auto`             | Device to use (`cuda`, `cpu`, or `auto`)                                                                            |
+| `--allow-fetch`     | `--no-allow-fetch` | Allow downloads from HuggingFace Hub (network access required)                                                     |
+| `--seed`            | `42`               | The seed to set for more reproducible behavior                                                                     |
+| `--per-line`        | `--no-per-line`    | Embed each non-empty line of the input file independently, producing one vector per line instead of a single vector for the whole file (text input only) |
+| `--batch-size`      | `32`               | Number of lines encoded per batch when `--per-line` is set, or number of pages per batch when embedding a multi-page PDF |
+| `--prompt`          |                    | Raw text prepended to each input before encoding (e.g. `query: `). Mutually exclusive with `--prompt-name`         |
+| `--prompt-name`     |                    | Name of a prompt predefined in the model's config (e.g. `query` or `passage` for e5/bge models). Mutually exclusive with `--prompt` |
+| `--normalize`       | `--no-normalize`   | Whether to normalize returned vectors to have length 1                                                             |
+| `--truncate-dim`    |                    | The dimension to truncate sentence embeddings to                                                                   |
+
+## Supported Input Formats
+
+- Text files (`.txt`, `.text`, `.md`, `.log`, `.rtf`)
+- Image files (`.jpg`, `.jpeg`, `.png`, `.tiff`, `.tif`, `.bmp`, `.heic`, `.heif`)
+- PDF files (`.pdf`) — each page is rendered to an image and embedded
+- Audio files (`.wav`, `.flac`, `.ogg`, `.aiff`, `.aif`, `.mp3`) — decoded, averaged to
+  mono, and resampled to the rate the model expects
+
+## Output Format
+
+NumPy binary (`.npy`).
+
+- By default, a text file is embedded as one document, producing a 1-D array of shape `(dim,)`.
+- With `--per-line`, each non-empty line of a text file is embedded independently, producing a 2-D array of shape `(n_lines, dim)`.
+- A single image (or single-page PDF) produces a 1-D array of shape `(dim,)`.
+- A multi-page PDF produces a 2-D array of shape `(n_pages, dim)`, one row per page.
+- A single audio file produces a 1-D array of shape `(dim,)`.
+
+## Models
+
+Any HuggingFace model compatible with the [sentence-transformers](https://sbert.net) library, including plain text encoder models. Embedding image or PDF input requires a multi-modal model (e.g. CLIP-style) that supports image encoding. Embedding audio requires an audio-capable model (e.g. `wav2vec2`, `HuBERT`, `WavLM`, a Whisper encoder, or CLAP) — the audio is automatically resampled to whatever rate that model's feature extractor expects.
+
+## Examples
+
+### Embed a document
+
+=== "Config"
+
+    ```yaml title="config.yaml"
+    tasks:
+      - name: embed
+        kind: local
+        module: tigerflow_ml.text.embed.local
+        input_ext: .txt
+        output_ext: .npy
+        params:
+          model: sentence-transformers/all-MiniLM-L6-v2
+          allow-fetch: True
+    ```
+
+=== "Input (.txt)"
+
+    ```text title="Raven.txt"
+    "The Raven" by Edgar Allan Poe
+    ```
+
+=== "Output (.npy)"
+
+    A single vector of shape `(384,)`.
+
+### Embed each line independently
+
+Use `--per-line` to embed a corpus file with one record per line, producing one vector per line instead of a single document vector.
+
+=== "Config"
+
+    ```yaml title="config.yaml"
+    tasks:
+      - name: embed
+        kind: local
+        module: tigerflow_ml.text.embed.local
+        input_ext: .txt
+        output_ext: .npy
+        params:
+          model: sentence-transformers/all-MiniLM-L6-v2
+          per-line: True
+          batch-size: 3
+          allow-fetch: True
+    ```
+
+=== "Input (.txt)"
+
+    ```text title="corpus.txt"
+    The quick brown fox jumps over the lazy dog.
+    Princeton University is in New Jersey.
+    Embeddings map text to dense vectors.
+    ```
+
+=== "Output (.npy)"
+
+    An array of shape `(3, 384)` — one row per line.
+
+### Embed an image
+
+Use a multi-modal (CLIP-style) model to embed images. PDFs are supported the same
+way, with one row of output per page.
+
+=== "Config"
+
+    ```yaml title="config.yaml"
+    tasks:
+      - name: embed
+        kind: local
+        module: tigerflow_ml.text.embed.local
+        input_ext: .jpg
+        output_ext: .npy
+        params:
+          model: sentence-transformers/clip-ViT-B-32
+          allow-fetch: True
+    ```
+
+=== "Input"
+
+    An image file, e.g. `photo.jpg`.
+
+=== "Output (.npy)"
+
+    A single vector of shape `(512,)`.
+
+### Embed audio
+
+Use an audio-capable model to embed a sound file. The file is decoded, averaged to
+mono, and resampled to the model's expected sampling rate before encoding.
+
+=== "Config"
+
+    ```yaml title="config.yaml"
+    tasks:
+      - name: embed
+        kind: local
+        module: tigerflow_ml.text.embed.local
+        input_ext: .mp3
+        output_ext: .npy
+        params:
+          model: openai/whisper-tiny
+          allow-fetch: True
+    ```
+
+=== "Input"
+
+    An audio recording, e.g. `clip.mp3`.
+
+=== "Output (.npy)"
+
+    A single vector of shape `(384,)`.
+
+### Run on HPC with Slurm
+
+For bulk embedding across large text collections, use the Slurm variant to distribute work across compute nodes:
+
+```yaml title="config.yaml"
+tasks:
+  - name: embed
+    kind: slurm
+    module: tigerflow_ml.text.embed.slurm
+    input_ext: .txt
+    output_ext: .npy
+    max_workers: 4
+    worker_resources:
+      cpus: 2
+      gpus: 1
+      memory: 16G
+      time: 04:00:00
+    params:
+      model: BAAI/bge-base-en-v1.5
+      per_line: True
+      batch_size: 64
+      cache_dir: ~/path/to/model/hub
+```

pyproject.toml

@@ -39,6 +39,7 @@ dependencies = [
     "soundfile>=0.12",
     "soxr>=0.3",
     "pillow-heif>=1.4.0",
+    "sentence-transformers[image]>=5.6.0",
 ]
 
 [project.optional-dependencies]
@@ -78,6 +79,8 @@ detect = "tigerflow_ml.image.detect.slurm:Detect"
 detect-local = "tigerflow_ml.image.detect.local:Detect"
 chat = "tigerflow_ml.text.chat.slurm:Chat"
 chat-local = "tigerflow_ml.text.chat.local:Chat"
+embed = "tigerflow_ml.text.embed.slurm:Embed"
+embed-local = "tigerflow_ml.text.embed.local:Embed"
 
 [build-system]
 requires = ["hatchling"]

src/tigerflow_ml/__init__.py

@@ -10,6 +10,7 @@
     from tigerflow_ml.audio.transcribe.slurm import Transcribe
     from tigerflow_ml.image.detect.slurm import Detect
     from tigerflow_ml.text.chat.slurm import Chat
+    from tigerflow_ml.text.embed.slurm import Embed
     from tigerflow_ml.text.ocr.slurm import OCR
     from tigerflow_ml.text.translate.slurm import Translate
 
@@ -19,9 +20,10 @@
     "Transcribe": "tigerflow_ml.audio.transcribe.slurm",
     "Detect": "tigerflow_ml.image.detect.slurm",
     "Chat": "tigerflow_ml.text.chat.slurm",
+    "Embed": "tigerflow_ml.text.embed.slurm",
 }
 
-__all__ = ["OCR", "Translate", "Transcribe", "Detect", "Chat"]
+__all__ = ["OCR", "Translate", "Transcribe", "Detect", "Chat", "Embed"]
 
 
 def __getattr__(name: str):

src/tigerflow_ml/text/embed/__init__.py

@@ -0,0 +1,14 @@
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from tigerflow_ml.text.embed.slurm import Embed
+
+__all__ = ["Embed"]
+
+
+def __getattr__(name: str):
+    if name == "Embed":
+        from tigerflow_ml.text.embed.slurm import Embed
+
+        return Embed
+    raise AttributeError(f"module 'tigerflow_ml.text.embed' has no attribute {name!r}")