diff --git a/AGENTS.md b/AGENTS.md index 7b8fa803e..c34c4f4be 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -42,7 +42,7 @@ - `__init__.py` — must define `DATASOURCE` (internal ID) and `NAME` (display name); optionally imports `init_datasource`. - A search worker file (e.g. `search_bsky.py`) whose `type` must follow the `{DATASOURCE}-search` or `{DATASOURCE}-import` naming convention. - Optional: `DESCRIPTION.md` (shown in the UI), `database.sql` (datasource-specific tables), Explorer CSS/HTML files. - - `processors/`: Modular Python scripts that manipulate datasets in some way. This is very diverse; processors include machine learning analyses as well as 'Download images' and simple metric calculations. Organized into subdirectories by category: `audio/`, `conversion/`, `filtering/`, `machine-learning/`, `machine_learning/`, `metrics/`, `networks/`, `presets/`, `statistics/`, `text-analysis/`, `twitter/`, `visualisation/`. + - `processors/`: Modular Python scripts that manipulate datasets in some way. This is very diverse; processors include machine learning analyses as well as 'Download images' and simple metric calculations. Organized into subdirectories by category: `audio/`, `conversion/`, `filtering/`, `machine_learning/`, `machine_learning/`, `metrics/`, `networks/`, `presets/`, `statistics/`, `text-analysis/`, `twitter/`, `visualisation/`. - `helper-scripts/`: Shared helper scripts. This importantly also contains various migration scripts to upgrade to new 4CAT versions. This often has to do with database manipulation. See the *Versioning and Migrations* section below. - `config/`: Configuration directory. Contains `config.ini` (INI-based runtime config), `module_config.bin` (pickled module-defined config options), `.current-version` (last-migrated version). `config/extensions/` holds installed extensions. Extensions are modular add-ons to 4CAT that are loaded in on startup. These can be data sources, processors, or workers. Folders here may be managed externally and concern symlinks. Each can have its own `requirements.txt`). @@ -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/common/lib/config_definition.py b/common/lib/config_definition.py index 252e844d3..794bb0276 100644 --- a/common/lib/config_definition.py +++ b/common/lib/config_definition.py @@ -652,13 +652,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.", }, } }, @@ -762,6 +762,15 @@ "help": "Maximum images to show", "tooltip": "Maximum number of images to show in the image visualization tab of a dataset. This is to prevent " "issues with large datasets.", + }, + "embeddings.max_embeddings": { + "type": UserInput.OPTION_TEXT, + "default": 100000, + "coerce_type": int, + "help": "Maximum embeddings to analyse", + "tooltip": "Maximum number of embeddings a single processor may read into memory at once, e.g. for clustering, " + "or reducing dimensions, which need to hold every vector in memory at the same time. " + "'0' means no limit.", } } @@ -782,6 +791,7 @@ "ui": "User interface", "proxies": "Proxied HTTP requests", "image-visuals": "Image visualization", + "embeddings": "Embeddings", "extensions": "Extensions", "llm": "LLM servers" } diff --git a/common/lib/embeddings.py b/common/lib/embeddings.py new file mode 100644 index 000000000..c2dcb731d --- /dev/null +++ b/common/lib/embeddings.py @@ -0,0 +1,228 @@ +""" +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 +from common.lib.user_input import UserInput + +# How often to report progress while reading, in items +STATUS_EVERY = 1000 + +# Read embeddings in batches to avoid memory problems +READ_BATCH = 8192 + +MAX_EMBEDDINGS = "embeddings.max_embeddings" + +# processor queue +EMBEDDING_QUEUE = "embeddings" + + +def get_max_items(config) -> int: + """ + The most embeddings a processor may read. + + :param config: Configuration reader, or `None` if there is none to read + :return int: Most embeddings to read, or 0 for no limit + """ + if config is None: + return 0 + + try: + return max(0, int(config.get(MAX_EMBEDDINGS, 0) or 0)) + except (TypeError, ValueError): + return 0 + + +def amount_option(config, tooltip: str, default: int = 0) -> dict: + """ + The 'No. of items' option capped by admin setting. Follows the image visualisation processors. + + :param config: Configuration reader, or `None` + :param str tooltip: Tooltip saying what a larger number costs + :param int default: Number to start the option at + :return dict: Option definition + """ + maximum = get_max_items(config) + + option = { + "type": UserInput.OPTION_TEXT, + "help": "No. of items", + "default": min(default, maximum) if (default and maximum) else default, + "min": 0, + "coerce_type": int, + "tooltip": tooltip + (f" '0' uses all embeddings, up to the {maximum:,} this 4CAT allows." + if maximum else " '0' uses all embeddings."), + } + + if maximum: + option["max"] = maximum + + return option + + +def items_to_read(parameters: dict, config, available: int) -> int: + """ + How many embeddings to read. + + :param dict parameters: Processor parameters + :param config: Configuration reader, or `None` + :param int available: Embeddings in the parent dataset + :return int: Number of embeddings to read + """ + wanted = parameters.get("amount", 0) or available + maximum = get_max_items(config) + + if maximum: + wanted = min(wanted, maximum) + + return min(wanted, available) + + +def get_vector(item: dict) -> list: + """ + Get the embedding of an item as a list of numbers. + + :param dict item: Embedding item + :return list: The vector, or an empty list if the item has none + """ + vector = item.get("embedding") + if isinstance(vector, str): + vector = [float(value) for value in vector.split() if value] + return vector or [] + + +def merge_batches(batches: list, rows: int, dimensions: int): + """ + Join the batches read into one array of exactly `rows` rows. Used to make + sure embeddings are not held in memory twice. + + :param list batches: Batches of rows, in the order they were read + :param int rows: Rows actually filled across all batches + :param int dimensions: Width of each row + :return: One float32 array of `rows` rows + """ + if len(batches) == 1: + # one batch is already the right width; copy only if it has a tail of + # rows that were never filled, which the caller would otherwise hold on + # to for as long as it keeps the matrix + return batches[0] if len(batches[0]) == rows else batches[0][:rows].copy() + + matrix = np.empty((rows, dimensions), dtype=np.float32) + offset = 0 + while batches: + # each batch is released as soon as it is copied over, so the two never + # hold the embeddings at the same time + batch = batches.pop(0) + taken = min(len(batch), rows - offset) + matrix[offset:offset + taken] = batch[:taken] + offset += taken + del batch + + return matrix + + +def read_embedding_matrix(dataset, limit: int, processor=None, accept=None, keep=None) -> tuple: + """ + Read up to `limit` embeddings into one float32 array. + + Read in batch rather than into one array of `limit` rows: `limit` is a + ceiling, not a count, and an array allocated at its full size keeps the + difference for as long as the caller holds the matrix. With an `accept` + that leaves most of the dataset out, that difference is nearly all of it. + + :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 + """ + batches, rows, dimensions = [], 0, 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 rows >= limit: + break + + record = item.original + + if accept is not None and not accept(record): + continue + + vector = get_vector(record) + if not vector: + continue + + if dimensions is None: + dimensions = len(vector) + elif len(vector) != dimensions: + # embeddings shouldn't have different lenghts, something is up! + raise ValueError("The embeddings are not all the same size.") + + if rows % READ_BATCH == 0: + # never more rows than are still wanted, so a nearly reached limit + # does not allocate a whole batch for a handful of vectors + batches.append(np.empty((min(READ_BATCH, limit - rows), dimensions), dtype=np.float32)) + + batches[-1][rows % READ_BATCH] = vector + rows += 1 + positions.append(position) + if keep is not None: + kept.append(keep(record)) + + if processor is not None and rows % STATUS_EVERY == 0: + processor.dataset.update_status(f"Read {rows:,} embeddings") + + if not rows: + return np.empty((0, 0), dtype=np.float32), [], [] + + return merge_batches(batches, rows, dimensions), positions, kept + + +def find_embedding(dataset, match, processor=None): + """ + Find the first item that matches, and its vector. + + :param dataset: Embeddings dataset + :param match: `function(item) -> bool` + :param processor: Processor reading it, to stop when it is interrupted + :return tuple: `(vector as float32 array, item)`, 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") + + item = item.original + if match(item): + vector = get_vector(item) + if vector: + return np.asarray(vector, dtype=np.float32), item + 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 209d76db6..54d888d4f 100644 --- a/common/lib/exceptions.py +++ b/common/lib/exceptions.py @@ -168,3 +168,9 @@ class MetadataException(FourcatException): Raise when there is an issue with metadata """ 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..174c9406d 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"` **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", # a multimodal embedder also shows vision and audio "vision": "image", - "embedding": "embedding", + "audio": "audio", } for cap in capabilities: mapped = _cap_map.get(cap) @@ -88,6 +91,151 @@ 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 shows capabilities in two places: under `metadata` for + the `/api/show` response that `list_models()` attaches, or inline on the + `/api/tags` entry itself. + + :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") + + # If we can't tell, assume the model generates. + 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 we can batch calls! + + 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). todo: update when this changes + * `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..cfc5461b5 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,197 @@ 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 + (and hasn't released a multimodal embedding model yet) + todo: keep track of when this changes + + * 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 is like a system prompt, other text should be part of the embedded content + 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, `audio_url` for + audio, `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": "video_url", "audio": "audio_url"}.get(media_type, "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 +248,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/clients/thirdparty_client.py b/common/lib/llm/clients/thirdparty_client.py index 30ed8e2e4..1e0a1e15c 100644 --- a/common/lib/llm/clients/thirdparty_client.py +++ b/common/lib/llm/clients/thirdparty_client.py @@ -30,7 +30,7 @@ def parse_supported_media_types(self, meta: dict) -> list[str]: :returns: Ordered list of supported media type strings. Returns `[]` when `meta` is `None` """ - return meta.get("supported_media_types", ["text"]) + return meta.get("supported_media_types", []) def format_display_name(self, meta: dict) -> str: """ diff --git a/common/lib/llm/llm_client.py b/common/lib/llm/llm_client.py index 60783294a..1cc5179ee 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,24 @@ 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 + # Accept both "Authorization: Bearer" and just the bearer + 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 @@ -77,7 +89,6 @@ def get_status(self) -> bool | int: :return: `False` if the server is not responding, or an HTTP status code. """ try: - print(f"{self.base_url}{self._models_info_path}") r = self._session.get( f"{self.base_url}{self._models_info_path}", headers=self._headers, @@ -129,9 +140,82 @@ 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. So use the + model name as an imperfect fallback: if it contains "embed", it's an embedder; + otherwise, it's 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 so assume generative + 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. + + Not done through LangChain since it doesn't support multimodal embeddings. + # todo: keep tabs on if/when this changes + + 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 +280,84 @@ 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_available_models(config, task: str = "generate") -> dict: + """ + Return which models are available to a particular user + + :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: `{model ID: model entry}` + """ + available_models = config.get("llm.available_models", {}) + enabled_model_ids = config.get("llm.enabled_models", []) + if not config.get("llm.access"): + enabled_model_ids = [_ for _ in enabled_model_ids if _.startswith("thirdparty")] + + return {model_id: model for model_id, model in available_models.items() + if model_id in enabled_model_ids and supports_task(model, task)} + + +def resolve_model(config, model_id: str, task: str = "generate") -> tuple: + """ + Look up a model and its server, by the same rules the form offers them. + + :param config: 4CAT config reader (context-aware) + :param str model_id: ID of the model to look up + :param str task: Task the model must support - `"generate"` or `"embed"`. + :return tuple: `(model, server, None)`, or `(None, None, error message)` + """ + models = get_available_models(config, task) + + if model_id not in models: + return None, None, (f"The '{model_id}' model is not available. Ask an admin to enable it, or pick another " + f"model.") + + model = models[model_id] + server = config.get("llm.servers", {}).get(model["server"]) + if not server: + return None, None, f"The LLM server for model '{model_id}' is no longer configured." + + return model, server, None + + +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. + """ + servers = config.get("llm.servers", {}) + + models_option = {} + for key, value in get_available_models(config, task).items(): + # a model whose server was removed cannot be offered; skipping it keeps + # the processor's options page working instead of raising a KeyError + server = servers.get(value["server"]) + if not server: + continue + + 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 6ba43e2e0..dab66b444 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/processors/machine_learning/cluster_embeddings.py b/processors/machine_learning/cluster_embeddings.py new file mode 100644 index 000000000..79c36be8f --- /dev/null +++ b/processors/machine_learning/cluster_embeddings.py @@ -0,0 +1,540 @@ +""" +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 (EMBEDDING_QUEUE, amount_option, items_to_read, 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 + +__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` - carrying `(labels, membership strengths)` - + or `error` + """ + try: + labels, probabilities = ClusterEmbeddings.reduce_and_cluster( + matrix, parameters, lambda kind, message: messages.put((kind, message))) + messages.put(("labels", (labels, probabilities))) + 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 result 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" + + annotation_labels = { + "cluster": annotation_label, + "cluster_probability": "cluster membership strength", + "cluster_concentration": "cluster concentration", + } + + # 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_queue_id(cls, remote_id, details, dataset) -> str: + """ + Shared queue for the processors that hold a full embedding matrix. + + Each of these reads every vector of its dataset into memory at once, so + running several at the same time multiplies the largest thing 4CAT + allocates. One at a time instead. + + :param str remote_id: Job item ID + :param dict details: Job details + :param DataSet dataset: Dataset to run job for + :return str: Queue ID + """ + return EMBEDDING_QUEUE + + @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": amount_option(config, "Clustering on a reduction slows down on very large datasets."), + "save_annotations": { + "type": UserInput.OPTION_ANNOTATIONS, + "options": { + "cluster": "Cluster number", + "cluster_probability": "Cluster membership strength (HDBSCAN only)", + "cluster_concentration": "Cluster concentration", + }, + "default": [] + }, + } + + 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. + """ + max_processed = items_to_read(self.parameters, self.config, 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.") + + clustered = self.reduce_and_cluster_in_child(matrix) + if clustered is None: + del matrix + return + + labels, probabilities = clustered + clusters = self.renumber_by_size(labels) + concentration = self.cluster_concentration(matrix, clusters) + del matrix + + 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", []) + 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]) + + if probabilities is not None: + record["cluster_probability"] = round(float(probabilities[row]), 4) + if record["cluster"] in concentration: + record["cluster_concentration"] = concentration[record["cluster"]] + + 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() + metadata = self.get_media_metadata(archive) if archive else None + post_id_map = metadata.filename_to_post_ids() if metadata else {} + post_ids = post_id_map.get(record.get("filename"), []) + if not post_ids: + unmatched += 1 + + annotations.extend({ + "item_id": str(post_id), + "label": self.annotation_labels[field], + "value": str(record[field]), + "type": "text", + } for post_id in post_ids for field in save_annotations if field in record) + + # 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. + + :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. + + Similar to colink_urls process spawning. + + :param matrix: 2D numpy array, one normalised embedding per row + :return tuple|None: `(cluster labels, membership strengths or None)`, + 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() + + result, error = None, None + try: + while result 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": + result = value + elif kind == "error": + error = f"Clustering failed: {value}" + else: + report(kind, value) + finally: + if result is None and child.is_alive(): + child.terminate() + child.join(timeout=5) + + if error: + self.dataset.finish_with_error(error) + return None + return result + + @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 tuple: `(cluster label per row, membership strength per row or + None)` + """ + 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 tuple: `(cluster label per row, membership strength per row or + None)`. Labels are -1 for HDBSCAN's noise. Only HDBSCAN gives cluster strength. + """ + 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), None + + 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, None + + 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 + clusterer = HDBSCAN(min_cluster_size=min_cluster_size, min_samples=min_samples, + algorithm="kd_tree", copy=False).fit(matrix) + + return clusterer.labels_, clusterer.probabilities_ + + @staticmethod + def cluster_concentration(matrix, clusters: list) -> dict: + """ + How tightly each cluster sits together, from 0 to 1. + + The mean cosine similarity between a cluster's members and its centre. + Calculated via full embeddings since UMAP does not preserve density. + + :param matrix: 2D numpy array, one normalised embedding per row + :param list clusters: Cluster number per row + :return dict: Concentration per cluster number + """ + numbers = np.asarray(clusters) + + return {cluster: round(float(np.linalg.norm(matrix[numbers == cluster].mean(axis=0))), 4) + for cluster in sorted({int(number) for number in numbers if number != -1})} + + @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"), + "cluster_probability": value_or_missing(item, "cluster_probability", -1), + "cluster_concentration": value_or_missing(item, "cluster_concentration", -1), + # 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..2053017a1 --- /dev/null +++ b/processors/machine_learning/embed_audio.py @@ -0,0 +1,147 @@ +""" +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): + """ + Transform audio files into embeddings. + """ + 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..8806c8675 --- /dev/null +++ b/processors/machine_learning/embed_images.py @@ -0,0 +1,130 @@ +""" +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.") + 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: + """ + Compress an image and return it as a base64 media descriptor. + + :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..879b9f506 --- /dev/null +++ b/processors/machine_learning/embed_media.py @@ -0,0 +1,472 @@ +""" +Shared behaviour for processors that embed media files via a multimodal model. + +""" +import abc +import base64 +import json +import time + +from datetime import datetime + +from backend.lib.processor import BasicProcessor +from common.lib.embeddings import amount_option, items_to_read +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, resolve_model as resolve_llm_model +from common.lib.user_input import UserInput + +__author__ = "Sal Hagen" +__credits__ = ["Sal Hagen"] +__maintainer__ = "Sal Hagen" +__email__ = "4cat@oilab.eu" + + +class EmbedMedia(BasicProcessor): + """ + Tranform media files into embeddings via a multimodal embedding model. + """ + + 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"] = { + **amount_option(config, f"Each {cls.media_label} is a separate request, so this is slow for large " + f"datasets.", default=100), + "help": f"No. of {cls.media_label_plural}", + } + + 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") + # make sure users can't run models they don't have access to + model, server, error = resolve_llm_model(self.config, model_id, task="embed") + if error: + self.dataset.finish_with_error(error) + 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_media_metadata(self): + """ + The archive's metadata, which says which posts each file came from. + + :return MediaArchiveMetadata|None: The metadata, or `None` if the + archive has none that can be used + """ + metadata = self.get_media_metadata() + if metadata: + self.dataset.log(f"Found post IDs for {len(metadata):,} {self.media_label_plural} in the archive " + f"metadata") + + return metadata + + 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() + + max_processed = items_to_read(self.parameters, self.config, 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 + media_metadata = self.load_media_metadata() + 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 + + # get file names either as full path or just file name (parent processors varyingly store media in + # subfolders) + try: + member = media_path.relative_to(staging_area).as_posix() + except ValueError: + member = media_path.name + + 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 = media_metadata.post_ids_for(member) if media_metadata else [] + + time_created = int(time.time()) + outfile.write(json.dumps({ + "id": media_path.stem, + "text": media_path.name, + "filename": member, + "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 media_metadata: + 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..287fe52f2 --- /dev/null +++ b/processors/machine_learning/embed_text.py @@ -0,0 +1,348 @@ +""" +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.embeddings import amount_option, items_to_read +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, resolve_model +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 reachable on this server.") + extension = "ndjson" # extension of result file, used internally and in UI + + compatibility = Compatibility(extensions={"csv", "ndjson"}) + + annotation_label = "text embedding" + + 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": amount_option(config, "Embedding is slower than it looks for large datasets; test on a sample " + "first.", default=100), + "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