File size: 1,528 Bytes
c3babff
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Hugging Face client - Inference (HTTP router + Inference Endpoints) and hub.

Uses huggingface_hub when present (pip install huggingface_hub), else raw HTTP.
"""
import os
import requests

ROUTER = os.getenv("HF_INFERENCE_API", "https://router.huggingface.co")
TOKEN = os.getenv("HF_TOKEN")


def _h():
    if not TOKEN:
        raise RuntimeError("HF_TOKEN not set")
    return {"Authorization": f"Bearer {TOKEN}", "Content-Type": "application/json"}


def chat_completion(model, messages, **kwargs):
    return requests.post(f"{ROUTER}/v1/chat/completions", headers=_h(),
                         json={"model": model, "messages": messages, **kwargs}, timeout=120).json()


def text_generation(model, prompt, **kwargs):
    return requests.post(f"{ROUTER}/models/{model}", headers=_h(),
                         json={"inputs": prompt, **kwargs}, timeout=120).json()


def list_models():
    return requests.get(f"{ROUTER}/v1/models", headers=_h(), timeout=30).json()


def infer_endpoint(endpoint_id: str, payload: dict):
    """Query a dedicated Inference Endpoint (set HF_INFERENCE_ENDPOINT base)."""
    base = os.getenv("HF_INFERENCE_ENDPOINT", ROUTER)
    return requests.post(f"{base}/{endpoint_id}", headers=_h(), json=payload, timeout=120).json()


def whoami():
    """Hub identity via huggingface_hub (optional dep)."""
    try:
        from huggingface_hub import whoami as _whoami
        return _whoami()
    except ImportError:
        return {"note": "install huggingface_hub for hub identity/upload APIs"}