diff --git a/AGENTS.md b/AGENTS.md index 1cf2cf2c3..13e7ce951 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -118,7 +118,7 @@ Invocation (the plain `docker compose` command uses `docker-compose.yml`): - **Extensions**: Installed under `config/extensions/`. Each extension can include its own `requirements.txt` (auto-installed by `setup.py`). Enabled/disabled via `extensions.enabled` setting. ## Versioning and Migrations -- The current version is stored in the `VERSION` file (currently `1.53`). **Do not edit this file casually.** +- The current version is stored in the `VERSION` file. **Do not edit this file casually.** - `config/.current-version` tracks the last-migrated version for the running instance. - `helper-scripts/migrate.py` compares `VERSION` to `.current-version` and runs the appropriate `migrate-X.XX-X.XX.py` scripts from `helper-scripts/migrate/` in sequence. - Migration scripts handle database schema changes, data transformations, and config updates between versions. diff --git a/VERSION b/VERSION index abf39ad47..ac7c5e2a2 100644 --- a/VERSION +++ b/VERSION @@ -1,4 +1,4 @@ -1.57 +1.58 This file should not be modified. It is used by 4CAT to determine whether it needs to run migration scripts to e.g. update the database structure to a more diff --git a/common/assets/llms.json b/common/assets/llms.json index 75fbf4d35..015c7078e 100644 --- a/common/assets/llms.json +++ b/common/assets/llms.json @@ -1,4 +1,25 @@ [ + { + "model": "gpt-5.6-sol", + "name": "[OpenAI] GPT-5.6 Sol", + "model_card": "https://developers.openai.com/api/docs/models/gpt-5.6-sol", + "provider": "openai" + },{ + "model": "gpt-5.6-terra", + "name": "[OpenAI] GPT-5.6 Terra", + "model_card": "https://platform.openai.com/docs/models/gpt-5.6-terra", + "provider": "openai" + },{ + "model": "gpt-5.6-luna", + "name": "[OpenAI] GPT-5.6 Luna", + "model_card": "https://platform.openai.com/docs/models/gpt-5.6-luna", + "provider": "openai" + },{ + "model": "gpt-5.5", + "name": "[OpenAI] GPT-5.5", + "model_card": "https://platform.openai.com/docs/models/gpt-5.5", + "provider": "openai" + }, { "model": "gpt-5.5", "name": "[OpenAI] GPT-5.5", @@ -48,44 +69,32 @@ "provider": "google" }, { - "model": "gemini-3.5-flash", - "name": "[Google] Gemini 3.5 Flash", - "model_card": "https://ai.google.dev/gemini-api/docs/models/gemini-3.5-flash", + "model": "gemini-3.8-flash", + "name": "[Google] Gemini 3.8 Flash", + "model_card": "https://ai.google.dev/gemini-api/docs/models/gemini-3.8-flash", "provider": "google" }, { - "model": "gemini-3-flash-preview", - "name": "[Google] Gemini 3 Flash", - "model_card": "https://cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-flash", + "model": "gemini-3.5-flash-lite-preview", + "name": "[Google] Gemini 3.5 Flash Lite", + "model_card": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-5-flash-lite", "provider": "google" }, { - "model": "gemini-3.1-flash-lite-preview", - "name": "[Google] Gemini 3.1 Flash Lite", - "provider": "google", - "model_card": "https://docs.cloud.google.com/vertex-ai/generative-ai/docs/models/gemini/3-1-flash-lite" - }, - { - "model": "claude-fable-5", - "name": "[Anthropic] Claude Fable 5", + "model": "claude-fable-5-1", + "name": "[Anthropic] Claude Fable 5.1", "model_card": "https://www.anthropic.com/claude/fable", "provider": "anthropic" }, { - "model": "claude-opus-4-8", - "name": "[Anthropic] Claude Opus 4.8", - "model_card": "https://www.anthropic.com/claude/opus", - "provider": "anthropic" - }, - { - "model": "claude-opus-4-6", - "name": "[Anthropic] Claude Opus 4.6", + "model": "claude-opus-5", + "name": "[Anthropic] Claude Opus 5", "model_card": "https://www.anthropic.com/claude/opus", "provider": "anthropic" }, { - "model": "claude-sonnet-4-6", - "name": "[Anthropic] Claude Sonnet 4.6 (latest)", + "model": "claude-sonnet-5", + "name": "[Anthropic] Claude Sonnet 5 (latest)", "model_card": "https://www.anthropic.com/claude/sonnet", "provider": "anthropic" }, diff --git a/common/lib/config_definition.py b/common/lib/config_definition.py index d913a2bdc..2ccf0ccbf 100644 --- a/common/lib/config_definition.py +++ b/common/lib/config_definition.py @@ -636,13 +636,13 @@ "type": UserInput.OPTION_TEXT, "help": "Authentication Header", "default": "", - "tooltip": "The HTTP header used to authenticate with the server (e.g. 'X-API-KEY', 'Authorization'). Passed with the Authentication Key as value.", + "tooltip": "The HTTP header used to authenticate with the server (e.g. 'X-API-KEY', 'Authorization'). Passed with the Authentication Key as value. For a bearer token you can write either 'Authorization' with 'Bearer ' as the value, or 'Authorization: Bearer' with just the key.", }, "auth_key": { "type": UserInput.OPTION_TEXT, "default": "", "help": "Authentication Key", - "tooltip": "The API key to access the LLM server, if required.", + "tooltip": "The API key to access the LLM server, if required. 'Bearer ' for bearer tokens.", }, } }, diff --git a/common/lib/embeddings.py b/common/lib/embeddings.py new file mode 100644 index 000000000..caeff6cd9 --- /dev/null +++ b/common/lib/embeddings.py @@ -0,0 +1,121 @@ +""" +Helpers for processors that need all vectors of an embeddings dataset at once. +Mostly to reduce memory errors + +""" +import numpy as np + +from common.lib.exceptions import ProcessorInterruptedException + +# How often to report progress while reading, in items +STATUS_EVERY = 1000 + + +def get_vector(record: dict) -> list: + """ + Get the embedding of a record as a list of numbers. + + :param dict record: Embedding record + :return list: The vector, or an empty list if the record has none + """ + vector = record.get("embedding") + if isinstance(vector, str): + vector = [float(value) for value in vector.split() if value] + return vector or [] + + +def read_embedding_matrix(dataset, limit: int, processor=None, accept=None, keep=None) -> tuple: + """ + Read up to `limit` embeddings into one float32 array. + + The array is allocated once, at its full size, instead of collecting rows + and stacking them, which would hold every vector twice for a moment. + + :param dataset: Embeddings dataset + :param int limit: Most embeddings to read + :param processor: Processor reading it, for progress and interruption + :param accept: Optional `function(record) -> bool`; records it rejects + are skipped + :param keep: Optional `function(record) -> anything`, to keep a small + summary of each record read (never the vector itself) + :return tuple: `(matrix, positions, kept)`: one row per embedding read; + the position of each among the items `iterate_items()` yields, to find + them again in a second pass; and what `keep` returned for each (empty + without `keep`) + :raises ValueError: If the embeddings are not all the same size + """ + matrix = None + positions, kept = [], [] + + for position, item in enumerate(dataset.iterate_items(processor, get_annotations=False)): + if processor is not None and processor.interrupted: + raise ProcessorInterruptedException("Interrupted while reading embeddings") + + if len(positions) >= limit: + break + + record = item.original + + if accept is not None and not accept(record): + continue + + vector = get_vector(record) + if not vector: + continue + + if matrix is None: + matrix = np.empty((limit, len(vector)), dtype=np.float32) + elif len(vector) != matrix.shape[1]: + # a ragged matrix cannot be compared; this means the dataset mixes + # models, which should not happen but is worth saying out loud + raise ValueError("The embeddings are not all the same size.") + + matrix[len(positions)] = vector + positions.append(position) + if keep is not None: + kept.append(keep(record)) + + if processor is not None and len(positions) % STATUS_EVERY == 0: + processor.dataset.update_status(f"Read {len(positions):,} embeddings") + + if matrix is None: + return np.empty((0, 0), dtype=np.float32), [], [] + + return matrix[:len(positions)], positions, kept + + +def find_embedding(dataset, match, processor=None): + """ + Find the first record that matches, and its vector. + + :param dataset: Embeddings dataset + :param match: `function(record) -> bool` + :param processor: Processor reading it, to stop when it is interrupted + :return tuple: `(vector as float32 array, record)`, or `(None, None)` + """ + for item in dataset.iterate_items(processor, get_annotations=False): + if processor is not None and processor.interrupted: + raise ProcessorInterruptedException("Interrupted while reading embeddings") + + record = item.original + if match(record): + vector = get_vector(record) + if vector: + return np.asarray(vector, dtype=np.float32), record + return None, None + + +def normalise_in_place(matrix): + """ + Scale each row to length 1, without making a copy. + + Embeddings are compared by angle (cosine similarity), not length. On unit + vectors, Euclidean distances and dot products follow the angles. + + :param matrix: 2D float array, one embedding per row; changed in place + :return: The same array + """ + norms = np.linalg.norm(matrix, axis=1, keepdims=True) + norms[norms == 0] = 1 + matrix /= norms + return matrix diff --git a/common/lib/exceptions.py b/common/lib/exceptions.py index b0b798a4c..17b5208ae 100644 --- a/common/lib/exceptions.py +++ b/common/lib/exceptions.py @@ -162,3 +162,9 @@ class MediaSignatureException(FourcatException): Raise in media processors when the media cannot be read """ pass + +class LLMServerException(FourcatException): + """ + Raise when an LLM server request fails or returns something unusable + """ + pass diff --git a/common/lib/llm/clients/ollama_client.py b/common/lib/llm/clients/ollama_client.py index 6484e1095..6eff2a653 100644 --- a/common/lib/llm/clients/ollama_client.py +++ b/common/lib/llm/clients/ollama_client.py @@ -3,6 +3,7 @@ """ import requests +from common.lib.exceptions import LLMServerException from common.lib.llm.llm_client import LLMServerClient @@ -46,7 +47,8 @@ def parse_supported_media_types(self, meta: dict) -> list[str]: **Primary path**: reads `meta["capabilities"]`: - `"completion"` → `"text"` - `"vision"` → `"image"` - - `"embedding"` → `"embedding"` + - `"audio"` → `"audio"` + - `"embedding"` → `"text"` (an embedder still consumes text) **Fallback path** (used when capabilities are absent or only yield `"text"`): inspects GGUF `model_info` / `details` for vision signals and adds @@ -57,16 +59,17 @@ def parse_supported_media_types(self, meta: dict) -> list[str]: Returns `[]` when `meta` is `None` (unknown — callers should include the model, not block it). """ - if meta is None or not meta.get("metadata"): + if meta is None or not (meta.get("metadata") or meta.get("capabilities")): return [] - capabilities = meta["metadata"].get("capabilities", []) + capabilities = self._get_capabilities(meta) media_types: list[str] = [] _cap_map = { "completion": "text", + "embedding": "text", "vision": "image", - "embedding": "embedding", + "audio": "audio", } for cap in capabilities: mapped = _cap_map.get(cap) @@ -88,6 +91,156 @@ def parse_supported_media_types(self, meta: dict) -> list[str]: return media_types + @staticmethod + def _get_capabilities(meta: dict) -> list[str]: + """ + Read a model's capability list. + + Ollama reports capabilities in two places: nested under `metadata` for + the `/api/show` response `list_models()` attaches, and inline on the + `/api/tags` entry itself. Prefer the former and fall back to the latter, + so a model whose `/api/show` request failed is still classified. + + :param dict meta: Model metadata, or `None`. + :returns list[str]: Ollama capability strings, or `[]` if unknown. + """ + if not meta: + return [] + + metadata = meta.get("metadata") or {} + return metadata.get("capabilities") or meta.get("capabilities") or [] + + def parse_supported_tasks(self, meta: dict) -> list[str]: + """ + Derive the tasks a model supports from its Ollama capabilities. + + :param dict meta: `/api/show` response dict, or `None`. + :returns list[str]: Supported tasks - `"generate"`, `"embed"`, or both. + """ + capabilities = self._get_capabilities(meta) + + tasks = [] + if "completion" in capabilities: + tasks.append("generate") + if "embedding" in capabilities: + tasks.append("embed") + + # An unreachable or partial /api/show leaves nothing to go on. Assume + # generation, matching the base class, rather than hiding the model + # from every processor at once. + return tasks if tasks else ["generate"] + + def embed(self, model_id: str, inputs: list, media: list | None = None, timeout: int = 300, + text_as_instruction: bool = False, dimensions: int | None = None) -> list[list[float]]: + """ + Embed inputs via Ollama's `/api/embed` endpoint. + + The endpoint takes a list and returns one vector per entry, so a caller + that batches its inputs gets one round-trip per batch rather than one + per item. + + An entry is either a plain string, or an object carrying any of `text`, + `image` and `audio`. Two limits: + + * video is not yet supported (can be approximated with video frame + audio). + * `image` and `audio` are single values, so one entry holds at most one + of each. Several frames cannot be fused into one vector this way. + + :param str model_id: Ollama model name. + :param list inputs: Text per item. With media, `inputs[i]` accompanies + `media[i]` and may be an empty string. + :param list media: Optional media per item, a descriptor or a list of + them, see `LLMServerClient.as_descriptors`. + :param int timeout: Request timeout in seconds. + :param bool text_as_instruction: Unused here since Ollama has no 'system prompt' + you can add to the embedding as e.g. Qwen3 embedding does. + :param int dimensions: Truncate vectors to this many dimensions. + :returns list[list[float]]: One vector per input, in input order. + :raises LLMServerException + """ + if not inputs and not media: + return [] + + if media and len(media) != len(inputs): + raise LLMServerException( + f"Got media for {len(media)} items but {len(inputs)} inputs; each input needs its own entry " + f"(use an empty string as the text if there is none)") + + payload = {"model": model_id, "input": self.build_input(inputs, media)} + if dimensions: + payload["dimensions"] = dimensions + + try: + response = self._session.post( + f"{self.base_url}/api/embed", + headers=self._headers, + json=payload, + timeout=timeout, + ) + except requests.RequestException as e: + raise LLMServerException(f"Could not reach Ollama server at {self.base_url}: {e}") + + if response.status_code != 200: + raise LLMServerException( + f"Ollama server returned status {response.status_code} while embedding with model " + f"'{model_id}': {response.text}") + + try: + embeddings = response.json().get("embeddings") + except ValueError as e: + raise LLMServerException(f"Ollama server returned invalid JSON while embedding: {e}") + + if not embeddings: + raise LLMServerException(f"Ollama server returned no embeddings for model '{model_id}'") + + if len(embeddings) != len(inputs): + raise LLMServerException( + f"Ollama server returned {len(embeddings)} embeddings for {len(inputs)} inputs; cannot map " + f"vectors back to items") + + return embeddings + + def build_input(self, inputs: list, media: list | None) -> list: + """ + Build the `input` list, as plain strings or as media objects. + + :param list inputs: Text per item + :param list media: Media per item, or None + :return list: Entries for the request body + :raises LLMServerException: For media this endpoint cannot accept + """ + if not media: + return list(inputs) + + entries = [] + for text, item_media in zip(inputs, media): + entry = {} + if text: + entry["text"] = text + + for descriptor in self.as_descriptors(item_media): + kind = descriptor.get("type") + if kind == "video": + raise LLMServerException( + "Ollama's /embed endpoint does not (yet) accept video. Send a frame and the audio track " + "instead or use an OpenAI-like API server that takes video.") + + if kind not in ("image", "audio"): + raise LLMServerException(f"Ollama cannot embed media of type '{kind}'") + + if kind in entry: + raise LLMServerException( + f"Ollama accepts one {kind} per item, and this item has several.") + + entry[kind] = descriptor.get("data") + + if not entry: + raise LLMServerException("An item has neither text nor usable media to embed") + + entries.append(entry) + + return entries + def format_display_name(self, meta: dict) -> str: """ Build a human-readable display name for a model. diff --git a/common/lib/llm/clients/openai_client.py b/common/lib/llm/clients/openai_client.py index 7cdf92050..382b7be8e 100644 --- a/common/lib/llm/clients/openai_client.py +++ b/common/lib/llm/clients/openai_client.py @@ -5,6 +5,9 @@ some useful API endpoints exclusive to it that we can benefit from, so use the dedicated class for tht instead. """ +import requests + +from common.lib.exceptions import LLMServerException from common.lib.llm.llm_client import LLMServerClient @@ -15,6 +18,194 @@ class OpenAICompatibleClient(LLMServerClient): _models_info_key = "models" _model_id_key = "key" + # vLLM and any strict OpenAI-compatible server expose the standard + # /v1/models, returning {"object": "list", "data": [{"id": ...}]}. The + # /api/v1/models shape below is kept as a fallback for servers that only + # speak that dialect, so existing connections keep working. + _openai_models_path = "/v1/models" + + def get_status(self) -> bool | int: + """ + Check whether the server is reachable. + + Probes the same two model-listing dialects `list_models()` does. Without + this a vLLM server reads as down: the base class only knows + `_models_info_path`, which vLLM answers with a 404 even while healthy. + + :return: `200` if either path answers, else the last status code seen, + or `False` if the server could not be reached at all. + """ + last_status = False + for path in (self._openai_models_path, self._models_info_path): + try: + response = self._session.get( + f"{self.base_url}{path}", headers=self._headers, timeout=self.timeout) + if response.status_code == 200: + return 200 + last_status = response.status_code + except requests.RequestException as e: + if self.log: + self.log.warning(f"{self.__class__.__name__}: server not available at {self.base_url}{path}: {e}") + + return last_status + + def list_models(self) -> list[dict]: + """ + List available models, trying the standard OpenAI path first. + + Tries `/v1/models` and falls back to this class's `_models_info_path`. + vLLM often does not prepend 'api/'... + + :return list[dict]: List of models available, or `[]` on failure. + """ + try: + response = self._session.get( + f"{self.base_url}{self._openai_models_path}", + headers=self._headers, + timeout=self.timeout, + ) + if response.status_code == 200: + models = response.json().get("data", []) + if models: + return [{**model, self._model_id_key: model.get("id")} for model in models] + except (requests.RequestException, ValueError) as e: + if self.log: + self.log.debug( + f"{self.__class__.__name__}: {self._openai_models_path} not usable at {self.base_url} ({e}), " + f"falling back to {self._models_info_path}") + + return super().list_models() + + def embed(self, model_id: str, inputs: list, media: list | None = None, timeout: int = 300, + text_as_instruction: bool = False, dimensions: int | None = None) -> list[list[float]]: + """ + Embed via the server's OpenAI-compatible `/v1/embeddings` endpoint. + + Two request shapes, because the OpenAI only handles text embeddings for now: + + * text only -> `{"input": [...]}`, the standard batched form, one vector + per input. + * with media -> `{"messages": [...]}`, vLLM's format, where the + content parts carry the media. A conversation is a single input, so + this embeds exactly one item per request; `inputs` and `media` are + paired by index and each pair becomes its own call. The conversation + ends in an empty assistant turn left open with + `continue_final_message`, as in vLLM's Qwen3-VL-Embedding example, + so the prompt ends where such models take their vector. + + A server that is not running a pooling/embedding model has no + `/v1/embeddings` route at all and answers 404. + + :param str model_id: Model ID within this server's context. + :param list inputs: Texts to embed. With media, `inputs[i]` is the text + accompanying `media[i]` and may be an empty string. + :param list media: Optional media descriptors, one per input - see + `LLMServerClient.embed`. + :param int timeout: Request timeout in seconds. + :param bool text_as_instruction: With media, send `inputs[i]` as a + system message that steers the embedding (e.g. Qwen3-VL-Embedding's + instruction), rather than as text embedded along with the media. + :returns list[list[float]]: One vector per input, in input order. + :raises LLMServerException: On transport failure, an error response, or + a vector count that does not match the input count. + """ + if not inputs and not media: + return [] + + if media and len(media) != len(inputs): + raise LLMServerException( + f"Got {len(media)} media items for {len(inputs)} inputs; each input needs exactly one media entry " + f"(use an empty string as the text if there is none)") + + if not media: + payload = {"model": model_id, "input": list(inputs)} + if dimensions: + payload["dimensions"] = dimensions + return self._post_embeddings(payload, timeout, len(inputs)) + + vectors = [] + for text, item_media in zip(inputs, media): + messages = [] + # several parts in one message is the natural shape here: unlike + # Ollama, this API fuses a whole content array into a single vector + user_content = [self._media_content_part(part) for part in self.as_descriptors(item_media)] + + # an instruction steers the embedding; any other text is embedded along with the media + if text and text_as_instruction: + messages.append({"role": "system", "content": [{"type": "text", "text": text}]}) + elif text: + user_content.append({"type": "text", "text": text}) + + messages.append({"role": "user", "content": user_content}) + messages.append({"role": "assistant", "content": [{"type": "text", "text": ""}]}) # necessary for vLLM + + payload = { + "model": model_id, + "messages": messages, + "continue_final_message": True, + "add_special_tokens": True, + } + if dimensions: + payload["dimensions"] = dimensions + vectors += self._post_embeddings(payload, timeout, 1) + + return vectors + + @staticmethod + def _media_content_part(item_media: dict) -> dict: + """ + Build one OpenAI-style content part for a media descriptor. + + vLLM keys the part by modality - `video_url` for video, `image_url` for + images - and accepts a `data:` URI in place of an http URL. + + :param dict item_media: Media descriptor (`type`, `mime`, `data`). + :return dict: Content part for the `messages` payload. + """ + media_type = item_media.get("type", "image") + mime = item_media.get("mime") or f"{media_type}/*" + url = item_media.get("url") or f"data:{mime};base64,{item_media.get('data', '')}" + + key = "video_url" if media_type == "video" else "image_url" + return {"type": key, key: {"url": url}} + + def _post_embeddings(self, payload: dict, timeout: int, expected: int) -> list[list[float]]: + """ + POST one embeddings request and get the vectors from the response. + + :param dict payload: Request body. + :param int timeout: Request timeout in seconds. + :param int expected: How many vectors the response must contain. + :return list[list[float]]: Vectors, ordered by the response `index` + field - the API does not promise to return them in input order. + """ + url = f"{self.base_url}/v1/embeddings" + try: + response = self._session.post(url, headers=self._headers, json=payload, timeout=timeout) + except requests.RequestException as e: + raise LLMServerException(f"Could not reach {url}: {e}") + + if response.status_code == 404: + raise LLMServerException( + f"{url} does not exist on this server.") + + if response.status_code != 200: + raise LLMServerException( + f"Server returned status {response.status_code} while embedding with model " + f"'{payload.get('model')}': {response.text[:500]}") + + try: + data = response.json().get("data", []) + except ValueError as e: + raise LLMServerException(f"Server returned invalid JSON while embedding: {e}") + + if len(data) != expected: + raise LLMServerException( + f"Server returned {len(data)} embeddings for {expected} input(s); cannot map vectors back to items") + + ordered = sorted(data, key=lambda row: row.get("index", 0)) + return [row["embedding"] for row in ordered] + def parse_supported_media_types(self, meta: dict) -> list[str]: """ Derive the media types a model supports from its LiteLLM metadata. @@ -54,6 +245,9 @@ def format_display_name(self, meta: dict) -> str: if meta.get("params_string"): extra_bits.append(meta["params_string"]) - model_name += f" ({', '.join(extra_bits)})" + # only append the parenthetical when there is something to put in it - + # a plain /v1/models entry has neither publisher nor size + if extra_bits: + model_name += f" ({', '.join(extra_bits)})" return model_name diff --git a/common/lib/llm/llm_client.py b/common/lib/llm/llm_client.py index 60783294a..d0794ff4c 100644 --- a/common/lib/llm/llm_client.py +++ b/common/lib/llm/llm_client.py @@ -12,6 +12,8 @@ import requests import re +from common.lib.exceptions import LLMServerException + class LLMServerClient: _headers = {} server_config = {} @@ -59,14 +61,29 @@ def __init__(self, config, server_config: dict, timeout: int = 10, log=None) -> self.base_url = server_config["url"].rstrip("/") if self.base_url.endswith("v1"): - # get rid of the 'v1' - we'll add this in the path - self.base_url = f"{self.base_url[:-2]}" + # get rid of the 'v1' - we'll add this in the path. + self.base_url = self.base_url[:-2].rstrip("/") self._session = requests.Session() self._headers = {"Content-Type": "application/json"} if self.auth_type: - self._headers[self.auth_type] = self.auth_key + # The setting asks for a header *name*, but "Authorization: Bearer" + # is the natural thing to type and is what people paste from curl + # examples. As a header name that is malformed - requests raises + # InvalidHeader, which surfaces as "server unavailable" and sends + # people hunting for a network problem. So accept both spellings: + # anything after a colon is treated as a prefix on the value. + header_name, _, value_prefix = self.auth_type.partition(":") + header_name = header_name.strip() + value_prefix = value_prefix.strip() + + auth_value = self.auth_key or "" + if value_prefix and not auth_value.lower().startswith(value_prefix.lower()): + auth_value = f"{value_prefix} {auth_value}".strip() + + if header_name: + self._headers[header_name] = auth_value self.log = log @@ -129,9 +146,85 @@ def build_model_entry(self, meta: dict) -> dict: "server": self.server_config["_id"], "wrapper": self.server_config["type"], "supported_media_types": self.parse_supported_media_types(meta), + "supported_tasks": self.parse_supported_tasks(meta), "metadata": meta, } + def parse_supported_tasks(self, meta: dict) -> list[str]: + """ + Derive the tasks a model can be used for from its metadata. + + Most servers do not report this. An OpenAI-compatible `/v1/models` lists + an id and nothing else, so there is no capability field to read. The + fallback is therefore the model's own name: an id containing "embed" + (Qwen3-VL-Embedding, mxbai-embed-large, nomic-embed-text) reads as an + embedder, anything else as generative. + + # todo: check if there's a more rigid way of separating later. + + :param dict meta: Model metadata, or `None` if unavailable. + :returns list[str]: Supported tasks - `"generate"`, `"embed"`, or both. + """ + try: + model_id = str(self.get_model_id(meta) or "") + except (KeyError, TypeError): + # nothing to go on; assume generative, which is how every server + # behaved before embedding support existed + return ["generate"] + + return ["embed"] if "embed" in model_id.lower() else ["generate"] + + @staticmethod + def as_descriptors(item_media) -> list: + """ + Normalise one item's media to a list of descriptors. + + :param item_media: A descriptor, a list of them, or None + :return list: Descriptors, possibly empty + """ + if not item_media: + return [] + + return list(item_media) if isinstance(item_media, (list, tuple)) else [item_media] + + def embed(self, model_id: str, inputs: list, media: list | None = None, timeout: int = 300, + text_as_instruction: bool = False, dimensions: int | None = None) -> list[list[float]]: + """ + Embed one or more inputs, returning one vector per input. + + Deliberately *not* routed through `LLMAdapter`/LangChain. LangChain's + embedding interface is text-only, so it + cannot handle multimodal embeddings. + + A client that cannot embed media must raise an exception when `media` is given + rather than embedding the text alone. + + :param str model_id: Model ID *within this server's context* (i.e. a + `local_id`, not a global model ID). + :param list inputs: Inputs to embed. + :param list media: Optional media to embed, one entry per input and + paired by index, for servers that support multimodal embedding. Each + entry is a descriptor dict: + + {"type": "video"|"image", "mime": "video/mp4", "data": ""} + + or, in place of `data`, a `"url"` the server can fetch itself. Each + client translates this to whatever its own API expects. + :param int timeout: Request timeout in seconds. + :param int dimensions: Truncate each vector to this many dimensions, + for models trained with Matryoshka representation learning. `None` + keeps the model's full output. + :param bool text_as_instruction: With media, treat `inputs[i]` as an + + instruction that steers the embedding rather than as text to embed + along with the media. Clients without such a distinction ignore it. + :returns list[list[float]]: One vector per input, in input order. + """ + raise LLMServerException( + f"Embedding is not supported for {self.server_config.get('type', 'this')} connections, so the model " + f"'{model_id}' cannot be used to generate embeddings. Use a model on an Ollama or OpenAI-compatible " + f"connection instead.") + def get_model_card_url(self, meta: dict) -> str: """ Get a URL for a model card for a given model @@ -196,4 +289,45 @@ def get_global_model_id(self, meta: dict) -> str: """ domain = re.sub(r"^https?://", "", self.server_config["url"]) domain = domain.rstrip("/") - return f"{self.server_config['type']}://{domain}/{self.get_model_id(meta)}" \ No newline at end of file + return f"{self.server_config['type']}://{domain}/{self.get_model_id(meta)}" + + +def supports_task(model: dict, task: str = "generate") -> bool: + """ + Check whether a model entry can be used for a given task. + + :param dict model: A single `llm.available_models` entry. + :param str task: Task to check for - `"generate"` or `"embed"`. + :return bool: Whether the model supports the task. + """ + return task in model.get("supported_tasks", ["generate"]) + + +def get_model_library(config, task: str = "generate") -> dict: + """ + Get the LLM models available for a given task, grouped by server. + + :param config: 4CAT config reader (context-aware, so per-user `llm.access` + is respected) + :param str task: Task the model must support - `"generate"` or `"embed"`. + :return dict: `{server name: {model ID: model display name}}`, shaped for + a `UserInput.OPTION_CHOICE` option. + """ + available_models = config.get("llm.available_models", {}) + enabled_model_ids = config.get("llm.enabled_models", []) + servers = config.get("llm.servers", {}) + if not config.get("llm.access"): + enabled_model_ids = [_ for _ in enabled_model_ids if _.startswith("thirdparty")] + + models_option = {} + for key, value in {k: v for k, v in available_models.items() if k in enabled_model_ids}.items(): + if not supports_task(value, task): + continue + + server = servers[value["server"]] + if server["name"] not in models_option: + models_option[server["name"]] = {} + + models_option[server["name"]][key] = value["name"] + + return models_option diff --git a/common/lib/module_loader.py b/common/lib/module_loader.py index 3dc3f3ccf..a050c1261 100644 --- a/common/lib/module_loader.py +++ b/common/lib/module_loader.py @@ -421,7 +421,8 @@ def _load_datasource(subdirectory, expiration): "id": subdirectory.parts[-1], "init": datasource.init_datasource, "config": {} if not hasattr(datasource, "config") else datasource.config, - "explorer-templates": self.load_datasource_explorer_templates(datasource_id, subdirectory) + "explorer-templates": self.load_datasource_explorer_templates(datasource_id, subdirectory), + "explorer-templates-from": getattr(datasource, "EXPLORER_TEMPLATE", None) } # Load 4CAT core datasources @@ -444,6 +445,22 @@ def _load_datasource(subdirectory, expiration): if subdirectory.is_dir(): _load_datasource(subdirectory, expiration) + # a data source can reuse another one's Explorer templates by setting + # EXPLORER_TEMPLATE to that data source's ID in its __init__.py. + own_templates = {datasource_id: datasource.get("explorer-templates") for datasource_id, datasource in + self.datasources.items()} + for datasource_id, datasource in self.datasources.items(): + template_source = datasource.get("explorer-templates-from") + if not template_source: + continue + + if template_source not in own_templates: + self.log_buffer += ("Data source '%s' uses the Explorer templates of '%s', but that data source is not " + "loaded; using its own templates instead.\n" % (datasource_id, template_source)) + continue + + datasource["explorer-templates"] = own_templates[template_source] + sorted_datasources = {datasource_id: self.datasources[datasource_id] for datasource_id in sorted(self.datasources, key=lambda id: self.datasources[id]["name"])} self.datasources = sorted_datasources diff --git a/datasources/tiktok/tiktok-explorer.css b/datasources/tiktok/tiktok-explorer.css index 8b07e38ef..6dfe52529 100644 --- a/datasources/tiktok/tiktok-explorer.css +++ b/datasources/tiktok/tiktok-explorer.css @@ -96,6 +96,37 @@ color: rgb(104, 119, 130); } +.items .item { + scroll-margin-top: 20px; +} + +.metrics .item-nav button { + display: block; + width: 48px; + height: 48px; + margin-bottom: 16px; + padding: 0; + border: none; + border-radius: 100%; + background-color: rgb(241, 241, 242); + color: rgb(22, 24, 35); + font-size: 20px; + cursor: pointer; +} + +.metrics .item-nav button:hover { + background-color: rgb(227, 227, 228); +} + +.metrics .item-nav button:disabled { + opacity: 0.5; + cursor: default; +} + +.metrics .item-nav button:disabled:hover { + background-color: rgb(241, 241, 242); +} + span.hashtag { color: rgb(29, 155, 240); } diff --git a/datasources/tiktok/tiktok-explorer.html b/datasources/tiktok/tiktok-explorer.html index 1f1b84dd1..7b31d3299 100644 --- a/datasources/tiktok/tiktok-explorer.html +++ b/datasources/tiktok/tiktok-explorer.html @@ -78,8 +78,69 @@ {% elif item.get("tiktok_url") and not pseudonymised %} {% endif %} + + + + + +{# This template is included once per item, but the script only needs to be on the page once #} +{% if item_index == 0 %} + +{% endif %} + diff --git a/datasources/tiktok_urls/__init__.py b/datasources/tiktok_urls/__init__.py index e4a989b84..daeb656fd 100644 --- a/datasources/tiktok_urls/__init__.py +++ b/datasources/tiktok_urls/__init__.py @@ -9,4 +9,7 @@ # Internal identifier for this data source DATASOURCE = "tiktok-urls" -NAME = "TikTok (via URLs)" \ No newline at end of file +NAME = "TikTok (via URLs)" + +# Items are mapped like TikTok imports, so reuse those Explorer templates +EXPLORER_TEMPLATE = "tiktok" diff --git a/datasources/tiktok_urls/tiktok_urls-explorer.css b/datasources/tiktok_urls/tiktok_urls-explorer.css deleted file mode 100644 index ef3ef9e3a..000000000 --- a/datasources/tiktok_urls/tiktok_urls-explorer.css +++ /dev/null @@ -1,19 +0,0 @@ -body { - color: black; -} - -.items li.item { - border-radius: 10px; - padding: 20px; - max-width: 200px; - background-color: white; - border-bottom: 1px grey; -} - - -.items li.item .preview { - display: block; - margin: 0 auto; - max-width: 300px; -} - diff --git a/helper-scripts/migrate/migrate-1.57-1.58.py b/helper-scripts/migrate/migrate-1.57-1.58.py new file mode 100644 index 000000000..88feb2ba0 --- /dev/null +++ b/helper-scripts/migrate/migrate-1.57-1.58.py @@ -0,0 +1,40 @@ +# Model entries in `llm.available_models` gained a `supported_tasks` field, so +# that processors can tell generative models and embedding models apart. +# +# Clearing the inventory forces it to be rebuilt. +import sys +import os + +from pathlib import Path + +sys.path.insert(0, os.path.join(os.path.abspath(os.path.dirname(__file__)), "../..")) +from common.lib.database import Database +from common.lib.logger import Logger + +import configparser # noqa: E402 + +log = Logger(output=True) +ini = configparser.ConfigParser() +ini.read(Path(__file__).parent.parent.parent.resolve().joinpath("config/config.ini")) +db_config = ini["DATABASE"] + +db = Database( + logger=log, + dbname=db_config["db_name"], + user=db_config["db_user"], + password=db_config["db_password"], + host=db_config["db_host"], + port=db_config["db_port"], + appname="4cat-migrate", +) + +print(" Clearing known LLM models so they are re-indexed with task support...") +has_setting = db.fetchone("SELECT COUNT(*) AS num FROM settings WHERE name = 'llm.available_models'") + +if has_setting["num"] > 0: + db.upsert("settings", {"name": "llm.available_models", "value": "{}", "tag": ""}, constraints=["name", "tag"]) + print(" ...cleared; models will be re-indexed on 4CAT restart") +else: + print(" ...no models indexed yet, nothing to do") + +print(" - done!") diff --git a/processors/machine_learning/cluster_embeddings.py b/processors/machine_learning/cluster_embeddings.py new file mode 100644 index 000000000..9d9c866bc --- /dev/null +++ b/processors/machine_learning/cluster_embeddings.py @@ -0,0 +1,500 @@ +""" +Assign each embedding to a cluster of similar items. +""" +import json +import multiprocessing +from collections import Counter +from queue import Empty + +import numpy as np +import psutil + +from backend.lib.processor import BasicProcessor +from common.lib.compatibility import Compatibility +from common.lib.embeddings import normalise_in_place, read_embedding_matrix +from common.lib.exceptions import ProcessorInterruptedException +from common.lib.item_mapping import MappedItem, value_or_missing +from common.lib.user_input import UserInput +from processors.machine_learning.embed_media import read_post_id_map + +__author__ = "Sal Hagen" +__credits__ = ["Sal Hagen"] +__maintainer__ = "Sal Hagen" +__email__ = "4cat@oilab.eu" + +# Silhouette scores compare every point with every other, so they are +# computed on a sample of at most this many points +SILHOUETTE_SAMPLE = 5000 + + +def _cluster_in_child(matrix, parameters, messages): + """ + Reduce and cluster in a child process, reporting back through a queue. + + :param matrix: 2D numpy array, one normalised embedding per row + :param dict parameters: Processor parameters + :param messages: Queue for `(kind, value)` tuples: `log` and `status` + messages, then `labels` or `error` + """ + try: + labels = ClusterEmbeddings.reduce_and_cluster(matrix, parameters, + lambda kind, message: messages.put((kind, message))) + messages.put(("labels", labels)) + except Exception as e: + messages.put(("error", f"{type(e).__name__}: {e}")) + raise + + +class ClusterEmbeddings(BasicProcessor): + """ + Group embeddings into clusters with k-means or HDBSCAN. + + Writes the embeddings unchanged plus a `cluster` field, so that anything + that runs on embeddings can also run on the reulst here. Clusters are numbered by size, the + largest first; HDBSCAN's outliers are cluster -1. + """ + type = "cluster-embeddings" # job type ID + category = "Machine learning" # category + title = "Cluster embeddings" # title displayed in UI + description = ("Assign clusters to embeddings using HDBSCAN or k-means.") + extension = "ndjson" # extension of result file, used internally and in UI + + compatibility = Compatibility( + types={"text-embeddings", "video-embeddings", "image-embeddings", "audio-embeddings"}, + preferred_followups=["reduce-embeddings"], + ) + + # Label the clusters are written under when saved as annotations + annotation_label = "embedding cluster" + + # Annotations are saved in batches of this many while writing + annotation_batch_size = 1000 + + references = [ + "[Campello, Ricardo J. G. B., Davoud Moulavi, and Jörg Sander. 2013. 'Density-Based Clustering Based on " + "Hierarchical Density Estimates.' PAKDD 2013.](https://doi.org/10.1007/978-3-642-37456-2_14)", + "[How HDBSCAN works](https://hdbscan.readthedocs.io/en/latest/how_hdbscan_works.html)", + "[Basic UMAP parameters](https://umap-learn.readthedocs.io/en/latest/parameters.html)", + "[scikit-learn: k-means](https://scikit-learn.org/stable/modules/clustering.html#k-means)", + "[scikit-learn: silhouette score](https://scikit-learn.org/stable/modules/clustering.html#silhouette-coefficient)", + ] + + @classmethod + def get_options(cls, parent_dataset=None, config=None) -> dict: + """ + Get processor options + + :param parent_dataset DataSet: An object representing the dataset that + the processor would be or was run on. + :param config ConfigManager|None config: Configuration reader (context-aware) + :return dict: Options for this processor + """ + return { + "algorithm": { + "type": UserInput.OPTION_CHOICE, + "help": "Clustering algorithm", + "options": { + "hdbscan": "HDBSCAN (finds the number of clusters)", + "kmeans": "K-means (a set number of clusters)", + }, + "default": "hdbscan", + "tooltip": "HDBSCAN finds clusters of any shape, decides how many there are, and assigns items that " + "fit nowhere a 'noise' cluster -1. K-means puts every item in one of a set " + "number of clusters, and assumes they are round and of similar size.", + }, + "min_cluster_size": { + "type": UserInput.OPTION_TEXT, + "help": "Minimum cluster size", + "default": 10, + "min": 2, + "max": 10000, + "coerce_type": int, + "tooltip": "The smallest group HDBSCAN will call a cluster. Higher values give fewer but larger " + "clusters as well as more noise.", + "requires": "algorithm==hdbscan", + }, + "min_samples": { + "type": UserInput.OPTION_TEXT, + "help": "Minimum samples", + "default": 10, + "min": 1, + "max": 500, + "coerce_type": int, + "tooltip": "How many close neighbours an item needs to count as part of a dense area. Higher values " + "make clustering more conservative: more items become noise and clusters are denser. " + "Memory use grows with this number times the number of items.", + "requires": "algorithm==hdbscan", + }, + "n_clusters": { + "type": UserInput.OPTION_TEXT, + "help": "K-means: number of clusters", + "default": 8, + "min": 2, + "max": 100, + "coerce_type": int, + "tooltip": "How many clusters to make. With automatic choice on, use this number as the max clusters to try.", + "requires": "algorithm==kmeans", + }, + "auto_clusters": { + "type": UserInput.OPTION_TOGGLE, + "help": "Find optimal number of clusters (silhouette)", + "default": False, + "tooltip": "Try every number of clusters from 2 up to the number above, and keep the one with the " + "best silhouette score (how much closer items are to their own cluster than to the next). " + "All scores are written to the log.", + "requires": "algorithm==kmeans", + }, + "reduce_first": { + "type": UserInput.OPTION_TOGGLE, + "help": "Reduce with UMAP first (recommended for HDBSCAN)", + "default": True, + "tooltip": "Cluster on a UMAP reduction of the embeddings instead of the full vectors. HDBSCAN finds " + "little in vectors with hundreds or thousands of dimensions, so for HDBSCAN this is " + "strongly recommended and k-means usually benefits too.", + }, + "reduce_dimensions": { + "type": UserInput.OPTION_TEXT, + "help": "Dimensions to reduce to", + "default": 25, + "min": 2, + "max": 100, + "coerce_type": int, + "tooltip": "Check what is recommended for k-means/HBDSCAN and test with your data.", + "requires": "reduce_first==true", + }, + "amount": { + "type": UserInput.OPTION_TEXT, + "help": "No. of items", + "default": 0, + "min": 0, + "coerce_type": int, + "tooltip": "Use '0' for all items. Clustering on a reduction slows down on very large datasets.", + }, + "save_annotations": { + "type": UserInput.OPTION_ANNOTATION, + "label": cls.annotation_label, + "tooltip": "Add each item's cluster number as an annotation to the original dataset. Media files " + "annotate every post they appeared in.", + "default": False, + }, + } + + def process(self): + """ + Read the embeddings, cluster them, and write them back with a cluster. + + Two passes over the embeddings: the first reads only the vectors, + into one array; the second reads each record again to write it out + with its cluster. Holding the records themselves in between would keep + every vector in memory a second time, as Python numbers. + """ + limit = self.parameters.get("amount", 0) + max_processed = min(limit, self.source_dataset.num_rows) if limit else self.source_dataset.num_rows + + self.dataset.update_status("Reading embeddings") + try: + matrix, positions, _ = read_embedding_matrix(self.source_dataset, max_processed, self) + except ValueError: + self.dataset.finish_with_error("The parent dataset contains embeddings of different sizes, so they " + "cannot be clustered together.") + return + + point_count = len(positions) + if point_count < 3: + self.dataset.finish_with_error(f"Not enough embeddings to cluster (found {point_count}, need at least " + f"3).") + return + + normalise_in_place(matrix) + + if not self.parameters.get("reduce_first", True) and self.parameters.get("algorithm") == "hdbscan": + self.dataset.log("Clustering the full vectors with HDBSCAN. This often finds few clusters and much noise; " + "reducing with UMAP first usually works better.") + + labels = self.reduce_and_cluster_in_child(matrix) + del matrix + if labels is None: + return + clusters = self.renumber_by_size(labels) + + sizes = Counter(clusters) + noise = sizes.pop(-1, 0) + summary = f"Found {len(sizes):,} clusters" + if noise: + summary += f" and {noise:,} items that fit no cluster" + self.dataset.log(summary + ". Cluster sizes: " + + ", ".join(f"{cluster}: {size:,}" for cluster, size in sorted(sizes.items()))) + + save_annotations = self.parameters.get("save_annotations", False) + annotations = [] + post_id_map = None + unmatched = 0 + + self.dataset.update_status("Writing clustered embeddings") + with self.dataset.get_results_path().open("w", encoding="utf-8") as outfile: + # the second pass: the same items in the same order, so the + # positions from the first pass say which row each one is + row = 0 + for position, item in enumerate(self.source_dataset.iterate_items(self, get_annotations=False)): + if self.interrupted: + raise ProcessorInterruptedException("Interrupted while writing clustered embeddings") + if row == point_count: + break + if position != positions[row]: + continue + + # a copy, so adding the cluster leaves the item itself alone + record = dict(item.original) + record["cluster"] = int(clusters[row]) + row += 1 + outfile.write(json.dumps(record, ensure_ascii=False) + "\n") + + if not save_annotations: + continue + + # media files are named after a hash that matches no post; the + # posts they appeared in come from the record or the archive + post_ids = self.get_post_ids(record) + if post_ids is None: + if post_id_map is None: + archive = self.source_dataset.get_parent() + post_id_map = read_post_id_map(archive.get_results_path() if archive else None, + log=self.dataset.log) + post_ids = post_id_map.get(str(record.get("id")), []) + if not post_ids: + unmatched += 1 + + annotations.extend({ + "item_id": str(post_id), + "label": self.annotation_label, + "value": str(record["cluster"]), + "type": "text", + } for post_id in post_ids) + + # saved in batches as we go: progress shows, and a run that + # fails halfway keeps what it saved + if len(annotations) >= self.annotation_batch_size: + self.save_annotations(annotations) + annotations = [] + self.dataset.update_status(f"Writing clustered embeddings and annotations ({row:,}/" + f"{point_count:,})") + + if annotations: + self.save_annotations(annotations) + if unmatched: + self.dataset.log(f"{unmatched:,} files could not be linked to a post in the original dataset, so their " + f"clusters were not added as annotations.") + + self.dataset.update_status(summary, is_final=True) + self.dataset.finish(point_count) + + @staticmethod + def get_post_ids(record: dict) -> list | None: + """ + Get the posts in the original dataset an embedding stands for. + + A text embedding's ID is its post's ID. A media file's ID is a hash of + the file, which matches no post: it stands for the posts it appeared + in. Embeddings made before those were always recorded lack them. + + :param dict record: Embedding record + :return list|None: Post IDs, or `None` for a media file whose posts + were not recorded and have to be looked up in the archive + """ + if "filename" not in record: + return [record.get("id")] + return record.get("post_ids") or None + + def reduce_and_cluster_in_child(self, matrix) -> np.ndarray | None: + """ + Reduce and cluster the embeddings in a separate process. + + UMAP and HDBSCAN run for a long time in compiled code that holds + Python's global interpreter lock, which would freeze every other thread + in the backend, the API included, until they finish. A child process + also gives its memory back to the system when it ends. + + The child is forked, not spawned: spawning re-runs the daemon's main + script. Where forking is not possible (Windows), this runs in-process. + + :param matrix: 2D numpy array, one normalised embedding per row + :return: Cluster label per row, or `None` if the dataset was finished + with an error + """ + def report(kind, message): + (self.dataset.log if kind == "log" else self.dataset.update_status)(message) + + if "fork" not in multiprocessing.get_all_start_methods(): + return self.reduce_and_cluster(matrix, self.parameters, report) + + context = multiprocessing.get_context("fork") + messages = context.Queue() + child = context.Process(target=_cluster_in_child, args=(matrix, self.parameters, messages)) + child.start() + + labels, error = None, None + try: + while labels is None and error is None: + if self.interrupted: + raise ProcessorInterruptedException("Interrupted while clustering embeddings") + if psutil.virtual_memory().percent >= 90: + self.log.warning(f"Memory usage above 90% while clustering; terminating clustering for dataset " + f"{self.dataset.key}") + error = "Clustering was stopped because memory usage exceeded 90%. Try fewer items." + break + + try: + kind, value = messages.get(timeout=0.5) + except Empty: + if not child.is_alive(): + error = f"Clustering stopped unexpectedly (exit code {child.exitcode})." + continue + + if kind == "labels": + labels = value + elif kind == "error": + error = f"Clustering failed: {value}" + else: + report(kind, value) + finally: + if labels is None and child.is_alive(): + child.terminate() + child.join(timeout=5) + + if error: + self.dataset.finish_with_error(error) + return None + return labels + + @staticmethod + def reduce_and_cluster(matrix, parameters: dict, report) -> np.ndarray: + """ + Reduce the embeddings if asked to, and cluster them. + + :param matrix: 2D numpy array, one normalised embedding per row + :param dict parameters: Processor parameters + :param report: `function(kind, message)`, with kind `log` or `status` + :return: Cluster label per row; -1 for HDBSCAN's noise + """ + if parameters.get("reduce_first", True): + matrix = ClusterEmbeddings.reduce(matrix, parameters, report) + + report("status", f"Clustering {len(matrix):,} embeddings") + return ClusterEmbeddings.cluster(matrix, parameters, report) + + @staticmethod + def reduce(matrix, parameters: dict, report): + """ + Reduce the embeddings with UMAP, only to find clusters in. + + Settings follow BERTopic's: `min_dist` 0 packs similar items tightly, + which is what a clustering algorithm needs, rather than what a readable + map needs. + + :param matrix: 2D numpy array, one normalised embedding per row + :param dict parameters: Processor parameters + :param report: `function(kind, message)`, with kind `log` or `status` + :return: 2D numpy array with the reduced vectors + """ + import umap + + point_count, dimensions = matrix.shape + target = min(parameters.get("reduce_dimensions", 5), dimensions - 1, point_count - 2) + report("status", f"Reducing {point_count:,} embeddings to {target} dimensions with UMAP before clustering; " + f"this may take a while") + report("log", f"Reducing to {target} dimensions with UMAP (cosine, n_neighbors=15, min_dist=0.0) before " + f"clustering") + + # fixed random_state: the same input should give the same clusters + return umap.UMAP(n_components=target, n_neighbors=min(15, point_count - 1), min_dist=0.0, metric="cosine", + random_state=42).fit_transform(matrix) + + @staticmethod + def cluster(matrix, parameters: dict, report) -> np.ndarray: + """ + Run the chosen clustering algorithm. + + :param matrix: 2D numpy array, one (normalised or reduced) vector per row + :param dict parameters: Processor parameters + :param report: `function(kind, message)`, with kind `log` or `status` + :return: Cluster label per row; -1 for HDBSCAN's noise + """ + point_count = len(matrix) + + if parameters.get("algorithm", "hdbscan") == "kmeans": + from sklearn.cluster import KMeans + + most = min(parameters.get("n_clusters", 8), point_count - 1) + if not parameters.get("auto_clusters"): + report("log", f"Running k-means with {most} clusters") + return KMeans(n_clusters=most, n_init="auto", random_state=42).fit_predict(matrix) + + from sklearn.metrics import silhouette_score + + best_score, best_labels = None, None + for count in range(2, most + 1): + report("status", f"Trying {count} of up to {most} clusters") + labels = KMeans(n_clusters=count, n_init="auto", random_state=42).fit_predict(matrix) + score = silhouette_score(matrix, labels, sample_size=min(point_count, SILHOUETTE_SAMPLE), + random_state=42) + report("log", f"{count} clusters: silhouette score {score:.3f}") + if best_score is None or score > best_score: + best_score, best_labels = score, labels + + report("log", f"Chose {len(set(best_labels))} clusters (silhouette score {best_score:.3f})") + return best_labels + + from sklearn.cluster import HDBSCAN + + min_cluster_size = min(parameters.get("min_cluster_size", 10), point_count) + # left unset, min_samples follows min_cluster_size, and finding each + # item's neighbours then holds items x min_samples distances and + # indices: over 10 GB for 130,000 items and a cluster size of 5,000 + min_samples = min(parameters.get("min_samples", 10), point_count) + report("log", f"Running HDBSCAN with a minimum cluster size of {min_cluster_size} and {min_samples} " + f"minimum samples") + # kd_tree: never fall back to brute force, which holds all distances + # between all items at once + return HDBSCAN(min_cluster_size=min_cluster_size, min_samples=min_samples, + algorithm="kd_tree", copy=False).fit_predict(matrix) + + @staticmethod + def renumber_by_size(labels) -> list: + """ + Number clusters by size, the largest 0, keeping noise at -1. + + The algorithms number clusters arbitrarily. By size, the numbers mean + something, and plots can give the largest clusters the clearest colours. + + :param labels: Cluster label per item + :return list: New cluster number per item + """ + sizes = Counter(int(label) for label in labels if label != -1) + order = {label: rank for rank, (label, _) in enumerate(sorted(sizes.items(), key=lambda kv: (-kv[1], kv[0])))} + return [order.get(int(label), -1) for label in labels] + + @staticmethod + def map_item(item): + """ + Map a clustered embedding to a flat row. + + The same shape the embedding datasets map to, plus the cluster. + + :param item: Item to map + :return MappedItem: Mapped item + """ + embedding = item.get("embedding", []) + + return MappedItem({ + "id": item.get("id"), + "text": item.get("text"), + "cluster": item.get("cluster"), + # only media embeddings have these + "filename": value_or_missing(item, "filename", ""), + "post_ids": ", ".join([str(post_id) for post_id in item.get("post_ids", [])]) + if "post_ids" in item else value_or_missing(item, "post_ids", ""), + "model": item.get("model"), + "dimensions": item.get("dimensions"), + "embedding": " ".join([str(value) for value in embedding]), + "timestamp": item.get("time_created"), + }) diff --git a/processors/machine_learning/embed_audio.py b/processors/machine_learning/embed_audio.py new file mode 100644 index 000000000..75fc3a0e2 --- /dev/null +++ b/processors/machine_learning/embed_audio.py @@ -0,0 +1,149 @@ +""" +Generate an embedding per audio file via a multimodal embedding model. +""" +import shutil + +from common.lib.compatibility import Compatibility, is_executable +from common.lib.user_input import UserInput +from processors.machine_learning.embed_media import EmbedMedia + +__author__ = "Sal Hagen" +__credits__ = ["Sal Hagen"] +__maintainer__ = "Sal Hagen" +__email__ = "4cat@oilab.eu" + + +class EmbedAudio(EmbedMedia): + """ + Send each audio file to a multimodal embedding model and store the vector it + returns, so recordings can be compared, mapped and searched by how they + sound. + """ + type = "audio-embeddings" # job type ID + category = "Audio" # category + title = "Generate audio embeddings" # title displayed in UI + description = ("Generate a numerical representation (an 'embedding') of each audio file using an " + "embedding model. This allows recordings to be compared and clustered by how they sound. Requires " + "an embedding model that accepts audio.") + extension = "ndjson" # extension of result file, used internally and in UI + + media_label = "audio file" + media_label_plural = "audio files" + model_requirement_note = "not only text or images" + working_filename = "prepared.wav" + annotation_label = "audio embedding" + + compatibility = Compatibility( + media_types={"audio"}, + type_prefixes={"audio-extractor"}, + required_settings={("video-downloader.ffmpeg_path", is_executable)}, + preferred_followups=["reduce-embeddings", "cluster-embeddings", "embedding-similarity"], + ) + + references = [ + "[EmbeddingGemma 2](https://huggingface.co/google/embeddinggemma-2)", + ] + + @classmethod + def get_media_options(cls, config=None) -> dict: + """ + Options controlling which part of the audio is sent. + + :param ConfigManager|None config: Configuration reader (context-aware) + :return dict: Options for this processor + """ + return { + "window_info": { + "type": UserInput.OPTION_INFO, + "help": "Audio embedding models usually only accept a short stretch of audio (often around 80 seconds). " + "Longer recordings are cut down to the window below.", + }, + "window_position": { + "type": UserInput.OPTION_TEXT, + "help": "Audio start %", + "default": 0, + "min": 0, + "max": 100, + "coerce_type": int, + "tooltip": "The percentage point of where the audio starts.", + }, + "window_seconds": { + "type": UserInput.OPTION_TEXT, + "help": "Seconds of audio", + "default": 75, + "min": 5, + "max": 80, + "coerce_type": int, + "tooltip": "Length of the audio fragments to embed. Max: 80s. Min: 5s.", + }, + } + + def check_requirements(self) -> str | None: + """ + Resolve ffmpeg, which trims and resamples the audio. + + :return str|None: An error message, or `None` when ffmpeg is present. + """ + self.ffmpeg = shutil.which(self.config.get("video-downloader.ffmpeg_path")) + if not self.ffmpeg: + return "ffmpeg is not available, so audio cannot be prepared for embedding." + + return None + + def prepare_media(self, media_path, output_path) -> tuple: + """ + Trim a window out of the audio and return it as a base64 descriptor. + + Resampled to 16 kHz mono, which is what these models are trained on, and + what makes the clip small enough to travel inside the request. + + :param Path media_path: Audio file to prepare + :param Path output_path: Where to write the trimmed clip + :return tuple: `(media descriptor, original size, sent size)` + """ + output_path.unlink(missing_ok=True) + + position = self.parameters.get("window_position", 10) / 100 + offset = self.probe_duration(media_path) * position + + result = self.run_interruptable_process([ + self.ffmpeg, "-y", "-ss", f"{offset:.3f}", "-i", str(media_path), + "-t", str(self.parameters.get("window_seconds", 75)), + "-vn", "-ac", "1", "-ar", "16000", str(output_path), + ], cleanup_paths=(output_path,)) + + if result.returncode != 0 or not output_path.exists() or not output_path.stat().st_size: + error = result.stderr.decode("utf-8", errors="replace") if result.stderr else "" + raise RuntimeError(f"ffmpeg exited {result.returncode}: {' '.join(error.split(chr(10))[-2:])}") + + original_size = media_path.stat().st_size + sent_size = output_path.stat().st_size + self.dataset.log(f"{media_path.name}: {self.parameters.get('window_seconds', 75)}s from {offset:.1f}s, " + f"{original_size / 1024:.0f} KB -> {sent_size / 1024:.0f} KB") + + payload = {"type": "audio", "mime": "audio/wav", "data": self.encode(output_path)} + output_path.unlink(missing_ok=True) + + return payload, original_size, sent_size + + def probe_duration(self, media_path) -> float: + """ + Read a recording's length, so a percentage becomes a timestamp. + + :param Path media_path: File to measure + :return float: Length in seconds, or 0.0 when it cannot be read + """ + ffprobe = shutil.which("ffprobe".join(self.ffmpeg.rsplit("ffmpeg", 1))) if self.ffmpeg else None + if not ffprobe: + return 0.0 + + probe = self.run_interruptable_process([ + ffprobe, "-v", "error", "-show_entries", "format=duration", + "-of", "default=noprint_wrappers=1:nokey=1", str(media_path), + ]) + + try: + return float(probe.stdout.decode("utf-8", errors="replace").strip()) + except (ValueError, AttributeError): + # an unreadable duration is not worth failing over; start at 0 + return 0.0 diff --git a/processors/machine_learning/embed_images.py b/processors/machine_learning/embed_images.py new file mode 100644 index 000000000..e2de40be8 --- /dev/null +++ b/processors/machine_learning/embed_images.py @@ -0,0 +1,135 @@ +""" +Generate an embedding per image via a multimodal embedding model. +""" +from PIL import Image, ImageOps, UnidentifiedImageError + +from common.lib.compatibility import Compatibility +from common.lib.user_input import UserInput +from processors.machine_learning.embed_media import EmbedMedia + +__author__ = "Sal Hagen" +__credits__ = ["Sal Hagen"] +__maintainer__ = "Sal Hagen" +__email__ = "4cat@oilab.eu" + + +class EmbedImages(EmbedMedia): + """ + Send each image to a multimodal embedding model and store the vector it + returns, so images can be compared, mapped and searched by content. + """ + type = "image-embeddings" # job type ID + category = "Visual" # category + title = "Generate image embeddings" # title displayed in UI + description = ("Generate a numerical representation (an 'embedding') of each image using a multimodal embedding " + "model. This allows images to be compared and clustered. Images are compressed before sending. " + "Requires an embedding model that accepts images, such as Qwen3-VL-Embedding.") + extension = "ndjson" # extension of result file, used internally and in UI + + media_label = "image" + media_label_plural = "images" + model_requirement_note = "not only text" + working_filename = "compressed.jpg" + annotation_label = "image embedding" + + compatibility = Compatibility( + media_types={"image"}, + type_prefixes={"image-downloader"}, + preferred_followups=["reduce-embeddings", "cluster-embeddings", "embedding-similarity"], + ) + + @classmethod + def get_media_options(cls, config=None) -> dict: + """ + Options controlling how an image is made smaller before it is sent. + + :param ConfigManager|None config: Configuration reader (context-aware) + :return dict: Options for this processor + """ + return { + "compression_info": { + "type": UserInput.OPTION_INFO, + "help": "Compression speeds up requests. Reducing the size helps the most, and most models scale " + "images down to a few hundred pixels anyway. Very large images can also exceed the model's " + "context window and be rejected.", + }, + "max_size": { + "type": UserInput.OPTION_TEXT, + "help": "Max width/height in pixels", + "default": 448, + "min": 64, + "max": 1920, + "coerce_type": int, + "tooltip": "Images larger than this are scaled down, keeping their aspect ratio.", + }, + "quality": { + "type": UserInput.OPTION_TEXT, + "help": "Compression level", + "default": 85, + "min": 1, + "max": 95, + "coerce_type": int, + "tooltip": "Min: 1, max: 95. JPEG quality: lower means smaller files and lower quality. 85 is a " + "reasonable trade-off, below 50 makes the image grainy.", + }, + } + + def skip_reason(self, media_path) -> str | None: + """ + Leave out files Pillow cannot open as a bitmap. + + SVGs are vector files and are common in image archives, so they are + named explicitly rather than reported as a generic failure. + + :param Path media_path: File that would be embedded + :return str|None: Why the file is skipped, or `None` to embed it. + """ + if media_path.suffix.lower() == ".svg": + return "SVG files cannot be embedded" + + return None + + def prepare_media(self, media_path, output_path) -> tuple: + """ + Re-encode an image smaller and return it as a base64 media descriptor. + + Pixel size matters most: models scale images to a fixed resolution + anyway, so anything larger is bytes spent on nothing. Everything is + written out as JPEG - the model reads a bitmap, so a PNG's lossless + compression only costs transfer size. + + :param Path media_path: Image to prepare + :param Path output_path: Where to write the re-encoded image + :return tuple: `(media descriptor, original size, sent size)` + """ + output_path.unlink(missing_ok=True) + max_size = self.parameters.get("max_size", 448) + + try: + with Image.open(media_path) as image: + # lets the JPEG decoder skip straight to roughly the size we + # want instead of decoding at full resolution first + image.draft("RGB", (max_size, max_size)) + + # photos carry their rotation in EXIF rather than in the pixels; + # without this a sideways photo is embedded sideways + image = ImageOps.exif_transpose(image) + + # JPEG has no alpha channel, and paletted or 16-bit images need + # converting before they can be saved as one + if image.mode != "RGB": + image = image.convert("RGB") + + image.thumbnail((max_size, max_size), Image.LANCZOS) + image.save(output_path, "JPEG", quality=self.parameters.get("quality", 85), optimize=True) + except (UnidentifiedImageError, OSError, ValueError) as e: + raise RuntimeError(f"could not read image: {e}") + + original_size = media_path.stat().st_size + sent_size = output_path.stat().st_size + self.dataset.log(f"{media_path.name}: {original_size / 1024:.0f} KB -> {sent_size / 1024:.0f} KB") + + payload = {"type": "image", "mime": "image/jpeg", "data": self.encode(output_path)} + output_path.unlink(missing_ok=True) + + return payload, original_size, sent_size diff --git a/processors/machine_learning/embed_media.py b/processors/machine_learning/embed_media.py new file mode 100644 index 000000000..bb3275269 --- /dev/null +++ b/processors/machine_learning/embed_media.py @@ -0,0 +1,532 @@ +""" +Shared behaviour for processors that embed media files via a multimodal model. + +`EmbedVideos` and `EmbedImages` differ only in how a file is made smaller before +it is sent; everything around that - resolving the model, walking the archive, +batching, writing results - is identical, and lives here. + +This class is abstract (`prepare_media` has no implementation), so +`ModuleCollector` skips it: `is_4cat_class()` ignores anything +`inspect.isabstract()` reports, and also ignores classes whose `__module__` is +not the file being scanned, which covers the subclass imports. +""" +import abc +import base64 +import json +import time +import zipfile + +from datetime import datetime + +from backend.lib.processor import BasicProcessor +from common.lib.exceptions import LLMServerException, ProcessorInterruptedException, QueryParametersException +from common.lib.item_mapping import MappedItem +from common.lib.llm.llm_client import LLMServerClient, get_model_library +from common.lib.user_input import UserInput + +__author__ = "Sal Hagen" +__credits__ = ["Sal Hagen"] +__maintainer__ = "Sal Hagen" +__email__ = "4cat@oilab.eu" + + +def read_post_id_map(archive_path, log=None) -> dict: + """ + Map each media filename in an archive to the posts it was downloaded for. + + A media archive's `.metadata.json` is keyed by source URL, and each entry + carries the `post_ids` it came from plus the `files` written for it. One + media file can belong to several posts, and its own filename is a hash + rather than a post ID, so anything annotating posts has to go through this. + + Module-level rather than a method, because processors that merely *read* a + media embedding - and so never subclass `EmbedMedia` - need it too. + + :param Path archive_path: Archive to read the metadata from + :param callable log: Optional logger for explaining a failed read + :return dict: `{filename without extension: [post ID, ...]}` + """ + post_ids = {} + if not archive_path or not zipfile.is_zipfile(archive_path): + return post_ids + + try: + with zipfile.ZipFile(archive_path) as archive: + if ".metadata.json" not in archive.namelist(): + return post_ids + + metadata = json.loads(archive.read(".metadata.json")) + except (zipfile.BadZipFile, json.JSONDecodeError, KeyError, OSError) as e: + if log: + log(f"Could not read .metadata.json, post IDs are unavailable: {e}") + return post_ids + + for entry in metadata.values(): + if not isinstance(entry, dict): + continue + for file in entry.get("files", []): + filename = file.get("filename") if isinstance(file, dict) else None + if filename: + post_ids[".".join(filename.split(".")[:-1])] = entry.get("post_ids", []) + + return post_ids + + +class EmbedMedia(BasicProcessor): + """ + Send each media file to a multimodal embedding model and store the vector it + returns, so the files can be compared, mapped and searched by content. + """ + + media_label = "media file" + media_label_plural = "media files" + model_requirement_note = "not only text" + working_filename = "compressed" + + annotation_label = "media_embedding" + annotation_batch_size = 100 + + # Stop after this many files in a row could not be embedded. + max_consecutive_failures = 10 + + @classmethod + def get_queue_id(cls, remote_id, details, dataset) -> str: + """ + Shared queue for locally hosted models + + :param str remote_id: Job item ID + :param dict details: Job details + :param DataSet dataset: Dataset to run job for + :return str: Queue ID + """ + # Unique queue for locally hosted models; used by other local model processors as well + return "local_models" + + @classmethod + def is_compatible_with(cls, module=None, config=None): + """ + Media datasets, but only where an embedding model is available. + + :param module: Dataset or processor to check against + :param ConfigManager|None config: Context-aware configuration reader + :return bool: + """ + if config is None or not cls.compatibility.is_compatible_with(module, config=config): + return False + + return bool(get_model_library(config, task="embed")) + + @classmethod + def get_media_options(cls, config=None) -> dict: + """ + Options controlling how a file is made smaller before it is sent. + + Subclasses override this; the surrounding model and amount options are + the same everywhere and are added by `get_options()`. + + :param ConfigManager|None config: Configuration reader (context-aware) + :return dict: Options for this processor + """ + return {} + + @classmethod + def get_options(cls, parent_dataset=None, config=None) -> dict: + """ + Get processor options + + :param parent_dataset DataSet: An object representing the dataset that + the processor would be or was run on. + :param config ConfigManager|None config: Configuration reader (context-aware) + :return dict: Options for this processor + """ + options = { + "model_info": { + "type": UserInput.OPTION_INFO, + "help": f"The model must accept **{cls.media_label}**, {cls.model_requirement_note}. 4CAT cannot tell " + f"which models do, so every embedding model is listed here; one that cannot read " + f"{cls.media_label} will return an error.", + }, + "model": { + "type": UserInput.OPTION_CHOICE, + "help": "Embedding model", + "options": get_model_library(config, task="embed") if config else {}, + "default": "", + }, + "instruction": { + "type": UserInput.OPTION_TEXT, + "help": "[optional] Prompt", + "default": "", + "tooltip": f"Requires an instruction-aware embedding model. This is a small system prompt sent with " + f"each {cls.media_label}. The text is not embedded itself but used by instruction-aware " + f"models like Qwen3-VL-Embedding to steer the embedding. Leave empty for model's own default.", + }, + } + + options.update(cls.get_media_options(config)) + + options["amount"] = { + "type": UserInput.OPTION_TEXT, + "help": f"No. of {cls.media_label_plural}", + "default": 100, + "min": 0, + "coerce_type": int, + "tooltip": f"Use '0' for all {cls.media_label_plural}. Each {cls.media_label} is a separate request, so " + f"this is slow for large datasets.", + } + + options["dimensions"] = { + "type": UserInput.OPTION_CHOICE, + "help": "Embedding size", + "options": { + "": "Full (model default)", + "768": "768", + "512": "512", + "256": "256", + "128": "128", + }, + "default": "", + "tooltip": "Increases efficiency but only works well for models trained with Matryoshka representation " + "learning, which keep most of their accuracy when truncated.", + } + + options["save_annotations"] = { + "type": UserInput.OPTION_ANNOTATION, + "label": cls.annotation_label, + "tooltip": "Add embeddings scores as annotations to the parent dataset.", + "default": False, + "hide_in_explorer": True + } + + return options + + @classmethod + def validate_query(cls, query, request, config): + """ + Validate input + + :param dict query: Query parameters, as submitted by the user via the web interface + :param request: Flask request + :param ConfigManager config: Configuration reader + :return dict: Validated query parameters + """ + allowed_models = { + model_id + for server_models in get_model_library(config, task="embed").values() + for model_id in server_models + } + + if query.get("model") not in allowed_models: + raise QueryParametersException(f"The '{query.get('model')}' embedding model is not currently available.") + + return query + + def check_requirements(self) -> str | None: + """ + Check anything this processor needs beyond a model, before starting. + + Runs once, before any file is read, so a missing dependency is reported + immediately rather than after the first item fails. Subclasses that need + an external tool override this and can store what they resolved on + `self`. + + :return str|None: An error message, or `None` when everything is present. + """ + return None + + def skip_reason(self, media_path) -> str | None: + """ + Decide whether to leave a file out, before any work is done on it. + + :param Path media_path: File that would be embedded + :return str|None: Why the file is skipped, or `None` to embed it. + """ + return None + + @abc.abstractmethod + def prepare_media(self, media_path, output_path) -> tuple: + """ + Make one file smaller and return it as a base64 media descriptor. + + The file travels inside the request body, so the point is to send as few + bytes as possible without losing what the model looks at. + + :param Path media_path: File to prepare + :param Path output_path: Where to write the smaller version + :return tuple: `(media descriptor, original size, sent size)` + """ + pass + + def embedding_dimensions(self) -> int | None: + """ + The vector size the user asked for, if they asked for one. + + The option is a dropdown, so the value arrives as a string; an empty + one means "whatever the model produces". + + :return int|None: Dimensions to truncate to, or `None` for the full + vector. + """ + chosen = self.parameters.get("dimensions") + return int(chosen) if chosen else None + + def resolve_model(self): + """ + Resolve the chosen model to a server and a client for it. + + :return tuple|None: `(model metadata, client)`, or `None` when + something is missing - in which case the dataset has been finished + with an error explaining what. + """ + model_id = self.parameters.get("model") + available_models = { + k: v for k, v in self.config.get("llm.available_models", {}).items() + if k in self.config.get("llm.enabled_models", []) + } + + if model_id not in available_models: + self.dataset.finish_with_error(f"Embedding model '{model_id}' is no longer available. Ask an admin to " + f"enable it, or pick another model.") + return None + + model = available_models[model_id] + server = self.config.get("llm.servers", {}).get(model["server"]) + if not server: + self.dataset.finish_with_error(f"The LLM server for model '{model_id}' is no longer configured.") + return None + + try: + client = LLMServerClient.get_client(self.config, server, self.log) + except ValueError as e: + self.dataset.finish_with_error(str(e)) + return None + + return model, client + + @staticmethod + def encode(path) -> str: + """ + Read a file and base64-encode it for transport in a request body. + + :param Path path: File to encode + :return str: Base64-encoded contents + """ + with path.open("rb") as infile: + return base64.b64encode(infile.read()).decode("utf-8") + + def load_post_id_map(self) -> dict: + """ + Map each media filename to the posts it was downloaded for. + + Read straight from the source archive rather than the staging area: the + staging area is filled lazily as items are iterated, so the metadata is + not reliably there before the first file is handled - and annotations + are written as we go. + + :return dict: `{filename without extension: [post ID, ...]}` + """ + post_ids = read_post_id_map(self.source_file, log=self.dataset.log) + self.dataset.log(f"Found post IDs for {len(post_ids):,} {self.media_label_plural} in .metadata.json") + return post_ids + + def process(self): + """ + Compress each file, embed it, and write one vector per file. + """ + resolved = self.resolve_model() + if not resolved: + return + + model, client = resolved + + requirement_error = self.check_requirements() + if requirement_error: + self.dataset.finish_with_error(requirement_error) + return + + instruction = self.parameters.get("instruction", "").strip() + + limit = self.parameters.get("amount", 100) + max_processed = min(limit, self.source_dataset.num_rows) if limit else self.source_dataset.num_rows + + staging_area = self.dataset.get_staging_area() + output_path = staging_area.joinpath(self.working_filename) + + save_annotations = self.parameters.get("save_annotations", False) + # always, not only for annotations: the posts behind each file are what + # later steps (clustering annotations, plotting by a post's date) need + post_id_map = self.load_post_id_map() + annotations = [] + + embedded = 0 + failed_embeddings = 0 + consecutive_failures = 0 + stopped = "" + skipped = 0 + processed = 0 + total_sent = 0 + + embedding_error = "" + + self.dataset.update_status(f"Embedding {self.media_label_plural} with {model['local_id']}") + + with self.dataset.get_results_path().open("w", encoding="utf-8") as outfile: + for item in self.source_dataset.iterate_items(self, staging_area=staging_area, + immediately_delete=False): + if self.interrupted: + raise ProcessorInterruptedException(f"Interrupted while embedding {self.media_label_plural}") + + if processed >= max_processed: + break + + media_path = item.file if hasattr(item, "file") else None + if not media_path or media_path.name == ".metadata.json": + # the archive carries its metadata alongside the media files + continue + + reason = self.skip_reason(media_path) + if reason: + self.dataset.log(f"Skipping {media_path.name}: {reason}") + skipped += 1 + continue + + processed += 1 + self.dataset.update_status(f"Embedding {self.media_label} {processed:,}/{max_processed:,} " + f"({media_path.name})") + + try: + payload, original_size, sent_size = self.prepare_media(media_path, output_path) + except ProcessorInterruptedException: + raise + except Exception as e: + self.dataset.log(f"Skipping {media_path.name}: could not compress ({e})") + skipped += 1 + continue + + if not payload: + self.dataset.log(f"Skipping {media_path.name}: compression produced no output") + skipped += 1 + continue + + # do the actual embedding + max_retries = 3 + retries = 0 + vector = None + while retries < max_retries: + try: + vector = client.embed(model["local_id"], [instruction], media=[payload], + dimensions=self.embedding_dimensions(), + text_as_instruction=True)[0] + break # success! + except LLMServerException as e: + retries += 1 + self.dataset.log(f"Error embedding {media_path.name}: {e}. Retrying ({retries}/{max_retries})...") + if retries >= max_retries: + embedding_error = f"Failed to embed {media_path.name} after {max_retries} attempts: {e}" + failed_embeddings += 1 + consecutive_failures += 1 + break # not a success... + time.sleep(2) # wait a bit before retrying + + if vector: + # the filename is a hash, so carry the posts it came from into + # the result: makes it easier for child processors to reach .metadata.json + post_ids = post_id_map.get(media_path.stem, []) + + time_created = int(time.time()) + outfile.write(json.dumps({ + "id": media_path.stem, + "text": media_path.name, + "filename": media_path.name, + "post_ids": post_ids, + "embedding": vector, + "dimensions": len(vector), + "model": model["local_id"], + "original_bytes": original_size, + "sent_bytes": sent_size, + "time_created": datetime.fromtimestamp(time_created).strftime("%Y-%m-%d %H:%M:%S"), + "time_created_utc": time_created, + }) + "\n") + + if save_annotations: + # one vector can belong to several posts, and each gets its + # own annotation + for post_id in post_ids: + annotations.append({ + "item_id": post_id, + "label": self.annotation_label, + "value": " ".join([str(value) for value in vector]), + "type": "text", + }) + + embedded += 1 + total_sent += sent_size + consecutive_failures = 0 + + if embedded % self.annotation_batch_size == 0: + outfile.flush() + if annotations: + self.save_annotations(annotations, hide_in_explorer=True) + annotations = [] + + if consecutive_failures >= self.max_consecutive_failures: + stopped = (f"Stopped after {consecutive_failures} {self.media_label_plural} in a row could not " + f"be embedded, which suggests a problem with the server rather than the files. " + f"{embedding_error}") + break + + self.dataset.update_progress(processed / max_processed) + + if annotations: + self.save_annotations(annotations, hide_in_explorer=True) + + if not embedded: + self.dataset.finish_with_error(stopped or f"No {self.media_label_plural} could be embedded.") + return + + if save_annotations and not post_id_map: + self.dataset.log("No annotations were written: the archive has no .metadata.json linking files to posts.") + + status = f"Embedded {embedded:,} {self.media_label_plural} ({total_sent / 1024 / 1024:.1f} MB sent in total)" + if skipped: + status += f", skipped {skipped:,}" + + if stopped: + self.dataset.finish_with_warning( + embedded, f"{status}. {stopped} {failed_embeddings:,} {self.media_label_plural} could not be " + f"embedded in total; see the log for details.") + return + + if failed_embeddings: + self.dataset.finish_with_warning( + embedded, f"{status}. {failed_embeddings:,} {self.media_label_plural} could not be embedded; see " + f"the log for details.") + return + + self.dataset.update_status(status, is_final=True) + self.dataset.finish(embedded) + + @staticmethod + def map_item(item): + """ + Map a media embedding to a flat row. + + Deliberately the same shape `text-embeddings` produces, so that the + processors chaining off embeddings work on these unchanged. + + :param item: Item to map + :return MappedItem: Mapped item + """ + embedding = item.get("embedding", []) + + return MappedItem({ + "id": item.get("id"), + "text": item.get("text"), + "filename": item.get("filename"), + "post_ids": ", ".join([str(post_id) for post_id in item.get("post_ids", [])]), + "model": item.get("model"), + "dimensions": item.get("dimensions"), + "embedding": " ".join([str(value) for value in embedding]), + "original_bytes": item.get("original_bytes"), + "sent_bytes": item.get("sent_bytes"), + "timestamp": item.get("time_created"), + }) diff --git a/processors/machine_learning/embed_text.py b/processors/machine_learning/embed_text.py new file mode 100644 index 000000000..52e208bf9 --- /dev/null +++ b/processors/machine_learning/embed_text.py @@ -0,0 +1,368 @@ +""" +Generate text embeddings via a configured LLM server. +""" +import json +import time + +from datetime import datetime + +from backend.lib.processor import BasicProcessor +from common.lib.compatibility import Compatibility +from common.lib.exceptions import LLMServerException, ProcessorInterruptedException, QueryParametersException +from common.lib.item_mapping import MappedItem +from common.lib.llm.llm_client import LLMServerClient, get_model_library +from common.lib.user_input import UserInput + +__author__ = "Sal Hagen" +__credits__ = ["Sal Hagen"] +__maintainer__ = "Sal Hagen" +__email__ = "4cat@oilab.eu" + + +class GenerateTextEmbeddings(BasicProcessor): + """ + Generate an embedding vector per dataset item using an embedding model + hosted on a configured LLM server. + """ + type = "text-embeddings" # job type ID + category = "Text analysis" # category + title = "Generate text embeddings" # title displayed in UI + description = ("Generate a numerical representation (an 'embedding') of the text of each item using an LLM embedding " + "model. Embeddings can be used to e.g. map, cluster, and retrieve similar kinds of texts. Requires " + "an embedding model to be installed on this server.") + extension = "ndjson" # extension of result file, used internally and in UI + + compatibility = Compatibility(extensions={"csv", "ndjson"}) + + #: Hard-coded label the embeddings are written back under, when saved as + #: annotations. Not user-editable: a vector is only meaningful next to the + #: model that produced it, so a free-form label invites mismatched columns. + annotation_label = "text embedding" + #: How many annotations to buffer before writing them to the database + annotation_batch_size = 100 + + references = [ + "[Ollama embedding models](https://ollama.com/blog/embedding-models)", + ] + + @classmethod + def get_queue_id(cls, remote_id, details, dataset) -> str: + """ + Shared queue for locally hosted models + + :param str remote_id: Job item ID + :param dict details: Job details + :param DataSet dataset: Dataset to run job for + :return str: Queue ID + """ + # Unique queue for locally hosted models; used by other local model processors as well + return "local_models" + + @classmethod + def is_compatible_with(cls, module=None, config=None): + """ + Text datasets, but only where an embedding model is actually available. + + :param module: Dataset or processor to check against + :param ConfigManager|None config: Context-aware configuration reader + :return bool: + """ + if config is None or not cls.compatibility.is_compatible_with(module, config=config): + return False + + return bool(get_model_library(config, task="embed")) + + @classmethod + def get_options(cls, parent_dataset=None, config=None) -> dict: + """ + Get processor options + + :param parent_dataset DataSet: An object representing the dataset that + the processor would be or was run on. + :param config ConfigManager|None config: Configuration reader (context-aware) + :return dict: Options for this processor + """ + options = { + "model": { + "type": UserInput.OPTION_CHOICE, + "help": "Embedding model", + # embedding models only - a generative model cannot return a vector + "options": get_model_library(config, task="embed") if config else {}, + "default": "", + "tooltip": "Only models that can produce embeddings are listed. Ask an admin to enable more via the " + "control panel.", + }, + "columns": { + "type": UserInput.OPTION_TEXT, + "help": "Column(s) to embed", + "default": "body", + "inline": True, + "tooltip": "Values of multiple columns are joined with a newline before being embedded.", + }, + "amount": { + "type": UserInput.OPTION_TEXT, + "help": "No. of items", + "default": 100, + "min": 0, + "coerce_type": int, + "tooltip": "Use '0' to embed all items. Embedding is slower than it looks for large datasets; test on " + "a sample first.", + }, + "batch_size": { + "type": UserInput.OPTION_TEXT, + "help": "Items per request", + "default": 50, + "min": 1, + "max": 500, + "coerce_type": int, + "tooltip": "How many items to send to the server at once. Higher is faster but uses more memory on the " + "server; lower this if requests time out.", + }, + "dimensions": { + "type": UserInput.OPTION_CHOICE, + "help": "Embedding size", + "options": { + "": "Full (model default)", + "768": "768", + "512": "512", + "256": "256", + "128": "128", + }, + "default": "", + "tooltip": "Increases efficiency but only works well for models trained with Matryoshka representation " + "learning, which keep most of their accuracy when truncated.", + }, + "save_annotations": { + "type": UserInput.OPTION_ANNOTATION, + # must match the label the annotations are written under, below: + # the form renders this as "Add