| """Hugging Face components (install ``fgao[huggingface]``; ``transformers`` for local models). |
| |
| * ``HFInferenceLLM`` chat model behind the Hugging Face Inference API / Providers |
| * ``TransformersLLM`` local ``transformers`` text-generation model |
| * ``hf_embedder`` embedding function for ``DenseRetriever`` (Inference API) |
| * ``st_embedder`` embedding function from a local ``sentence-transformers`` model |
| * ``load_hf_corpus`` documents from a Hugging Face dataset |
| """ |
| from __future__ import annotations |
|
|
| import os |
| from typing import Any, Callable, List, Optional |
|
|
| from ..rag.documents import Document |
| from ..rag.prompts import SYSTEM_PROMPT |
|
|
|
|
| class HFInferenceLLM: |
| """``huggingface_hub`` chat completion (Inference API / Providers / a TGI endpoint) as an FGAO LLM. |
| |
| Sync calls use ``InferenceClient``, async calls ``AsyncInferenceClient``. |
| The token defaults to ``HF_TOKEN``. ``temperature=None`` leaves the |
| provider default; set ``temperature=0`` where the provider accepts it for |
| deterministic calls. ``base_url`` targets a dedicated endpoint. |
| """ |
|
|
| def __init__(self, model: str, *, token: Optional[str] = None, provider: Optional[str] = None, |
| base_url: Optional[str] = None, max_tokens: int = 128, temperature: Optional[float] = None, |
| system: Optional[str] = SYSTEM_PROMPT, timeout: Optional[float] = None, client: Any = None, |
| aclient: Any = None, **kwargs: Any) -> None: |
| self.model = model |
| self._ckw: dict = {"token": token or os.environ.get("HF_TOKEN")} |
| for k, v in (("provider", provider), ("base_url", base_url), ("timeout", timeout)): |
| if v is not None: |
| self._ckw[k] = v |
| self._client, self._aclient = client, aclient |
| self.max_tokens, self.temperature, self.system, self.kwargs = max_tokens, temperature, system, kwargs |
|
|
| @property |
| def client(self): |
| if self._client is None: |
| from huggingface_hub import InferenceClient |
|
|
| self._client = InferenceClient(**self._ckw) |
| return self._client |
|
|
| @property |
| def aclient(self): |
| if self._aclient is None: |
| from huggingface_hub import AsyncInferenceClient |
|
|
| self._aclient = AsyncInferenceClient(**self._ckw) |
| return self._aclient |
|
|
| def _args(self, prompt: str) -> dict: |
| messages = ([{"role": "system", "content": self.system}] if self.system else []) + [ |
| {"role": "user", "content": prompt}] |
| kw = dict(self.kwargs) |
| if self.temperature is not None: |
| kw["temperature"] = self.temperature |
| kw.update(messages=messages, max_tokens=self.max_tokens) |
| if "base_url" not in self._ckw: |
| kw["model"] = self.model |
| return kw |
|
|
| def complete(self, prompt: str) -> str: |
| out = self.client.chat_completion(**self._args(prompt)) |
| return (out.choices[0].message.content or "").strip() |
|
|
| async def acomplete(self, prompt: str) -> str: |
| out = await self.aclient.chat_completion(**self._args(prompt)) |
| return (out.choices[0].message.content or "").strip() |
|
|
|
|
| class TransformersLLM: |
| """A local ``transformers`` model (greedy decoding, chat template when available).""" |
|
|
| def __init__(self, model: Any, *, max_new_tokens: int = 64, system: Optional[str] = SYSTEM_PROMPT, |
| **pipeline_kwargs: Any) -> None: |
| if isinstance(model, str): |
| from transformers import pipeline |
|
|
| model = pipeline("text-generation", model=model, **pipeline_kwargs) |
| self.pipe, self.max_new_tokens, self.system = model, max_new_tokens, system |
|
|
| def complete(self, prompt: str) -> str: |
| tok = getattr(self.pipe, "tokenizer", None) |
| if tok is not None and getattr(tok, "chat_template", None): |
| msgs = ([{"role": "system", "content": self.system}] if self.system else []) + [ |
| {"role": "user", "content": prompt}] |
| out = self.pipe(msgs, max_new_tokens=self.max_new_tokens, do_sample=False) |
| gen = out[0]["generated_text"] |
| return (gen[-1]["content"] if isinstance(gen, list) else str(gen)).strip() |
| out = self.pipe(prompt, max_new_tokens=self.max_new_tokens, do_sample=False, return_full_text=False) |
| return out[0]["generated_text"].strip() |
|
|
|
|
| def hf_embedder(model: str, *, token: Optional[str] = None, client: Any = None) -> Callable[[List[str]], List[List[float]]]: |
| """Embedding function backed by ``InferenceClient.feature_extraction``.""" |
| if client is None: |
| from huggingface_hub import InferenceClient |
|
|
| client = InferenceClient(token=token or os.environ.get("HF_TOKEN")) |
|
|
| def embed(texts: List[str]) -> List[List[float]]: |
| rows = [] |
| for t in texts: |
| v = client.feature_extraction(t, model=model) |
| v = v.tolist() if hasattr(v, "tolist") else v |
| while isinstance(v, list) and v and isinstance(v[0], list) and len(v) == 1: |
| v = v[0] |
| if v and isinstance(v[0], list): |
| v = [sum(col) / len(v) for col in zip(*v)] |
| rows.append([float(x) for x in v]) |
| return rows |
|
|
| return embed |
|
|
|
|
| def st_embedder(model: str, **kwargs: Any) -> Callable[[List[str]], Any]: |
| """Embedding function from a local ``sentence-transformers`` model.""" |
| from sentence_transformers import SentenceTransformer |
|
|
| m = SentenceTransformer(model, **kwargs) |
| return lambda texts: m.encode(list(texts), normalize_embeddings=True) |
|
|
|
|
| def load_hf_corpus(path: str, *, split: str = "train", text_field: str = "text", title_field: Optional[str] = None, |
| id_field: Optional[str] = None, limit: Optional[int] = None, **kwargs: Any) -> List[Document]: |
| """Documents from a Hugging Face dataset (``datasets.load_dataset``).""" |
| from datasets import load_dataset |
|
|
| ds = load_dataset(path, split=split, **kwargs) |
| out: List[Document] = [] |
| for i, row in enumerate(ds): |
| if limit is not None and i >= limit: |
| break |
| out.append(Document(text=str(row[text_field]), title=str(row[title_field]) if title_field else "", |
| id=str(row[id_field]) if id_field else str(i))) |
| return out |
|
|
|
|
| __all__ = ["HFInferenceLLM", "TransformersLLM", "hf_embedder", "st_embedder", "load_hf_corpus"] |
|
|
