Spaces:
Running
Running
| """ | |
| Generic loader/cache for Hugging Face pipelines pulled from the Hub at | |
| request time. | |
| This lets a caller pick *any* HF model id - for sentiment analysis, named | |
| entity recognition, or zero-shot topic classification - instead of being | |
| limited to the built-in default models. Pipelines are expensive to | |
| instantiate (they download weights on first use), so loaded pipelines are | |
| kept in a small in-memory cache, namespaced by task so the same model id | |
| used for two different tasks never collides. | |
| Security note: ``trust_remote_code`` is never enabled. Enabling it would let | |
| an arbitrary Hub repository execute Python code inside this process, which | |
| is not acceptable for a model id supplied by an API caller. | |
| """ | |
| from __future__ import annotations | |
| import threading | |
| from typing import Any | |
| # Maximum number of distinct (task, model) pipelines kept warm in memory at | |
| # once. Oldest (first loaded) is evicted when the cache is full - simple | |
| # FIFO, adequate for a demo/thesis system that isn't serving many concurrent | |
| # model ids. | |
| _MAX_CACHED_PIPELINES = 6 | |
| _cache: dict[str, Any] = {} | |
| _cache_order: list[str] = [] | |
| _load_errors: dict[str, str] = {} | |
| _lock = threading.Lock() | |
| def _cache_key(model_id: str, task: str) -> str: | |
| return f"{task}::{model_id}" | |
| def get_pipeline(model_id: str, task: str = "sentiment-analysis", **pipeline_kwargs: Any): | |
| """ | |
| Return a cached (or newly loaded) Hugging Face pipeline for | |
| ``(task, model_id)``. Extra ``pipeline_kwargs`` (e.g. | |
| ``aggregation_strategy="simple"`` for NER) are forwarded to | |
| ``transformers.pipeline`` on first load only. | |
| Raises ``RuntimeError`` if the model cannot be loaded (unknown repo, | |
| incompatible with the task, network error, ...) so callers can turn it | |
| into a clean HTTP error instead of crashing. | |
| """ | |
| key = _cache_key(model_id, task) | |
| with _lock: | |
| if key in _cache: | |
| return _cache[key] | |
| if key in _load_errors: | |
| raise RuntimeError(_load_errors[key]) | |
| try: | |
| from transformers import pipeline | |
| clf = pipeline( | |
| task, | |
| model=model_id, | |
| tokenizer=model_id, | |
| trust_remote_code=False, | |
| **pipeline_kwargs, | |
| ) | |
| except Exception as e: # pragma: no cover - depends on network/model | |
| with _lock: | |
| _load_errors[key] = str(e) | |
| raise RuntimeError(f"Could not load model '{model_id}' for task '{task}': {e}") from e | |
| with _lock: | |
| _cache[key] = clf | |
| _cache_order.append(key) | |
| while len(_cache_order) > _MAX_CACHED_PIPELINES: | |
| oldest = _cache_order.pop(0) | |
| _cache.pop(oldest, None) | |
| return clf | |
| def cached_models(task: str | None = None) -> dict[str, list[str]]: | |
| """ | |
| Model ids currently kept warm in memory, grouped by task. | |
| If ``task`` is given, returns just that task's list under the same key | |
| (still as a dict, for a consistent return type). | |
| """ | |
| with _lock: | |
| keys = list(_cache_order) | |
| grouped: dict[str, list[str]] = {} | |
| for key in keys: | |
| key_task, _, model_id = key.partition("::") | |
| if task is not None and key_task != task: | |
| continue | |
| grouped.setdefault(key_task, []).append(model_id) | |
| return grouped | |