File size: 6,397 Bytes
0703e2b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 | """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): # token embeddings: mean-pool
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"]
|