codebook / potato /ai /huggingface_endpoint.py
davidjurgens's picture
Deploy: Potato — Codebook Annotation
aceb1b2 verified
Raw
History Blame Contribute Delete
2.09 kB
"""
Hugging Face AI endpoint implementation.
This module provides integration with Hugging Face's Inference API for LLM inference.
"""
from huggingface_hub import InferenceClient
from .ai_endpoint import BaseAIEndpoint, AIEndpointRequestError
DEFAULT_MODEL = "meta-llama/Llama-3.2-3B-Instruct"
class HuggingfaceEndpoint(BaseAIEndpoint):
"""Hugging Face endpoint for cloud-based LLM inference."""
def _initialize_client(self) -> None:
"""Initialize the Hugging Face client."""
api_key = self.ai_config.get("api_key", "")
if not api_key:
raise AIEndpointRequestError("Hugging Face API key is required")
# Default timeout of 30 seconds, configurable via ai_config
timeout = self.ai_config.get("timeout", 30)
self.client = InferenceClient(
model=self.model,
token=api_key,
timeout=timeout
)
def _get_default_model(self) -> str:
"""Get the default Hugging Face model."""
return DEFAULT_MODEL
def query(self, prompt: str, output_format: dict) -> str:
"""
Send a query to Hugging Face and return the response.
Args:
prompt: The prompt to send to the model
Returns:
The model's response as a string
Raises:
AIEndpointRequestError: If the request fails
"""
try:
response = self.client.chat_completion(
messages=[{"role": "user", "content": prompt}],
max_tokens=self.max_tokens,
temperature=self.temperature,
response_format= {
"type": "json_schema",
"json_schema": {
"name": "output_format",
"schema": output_format.model_json_schema(),
"strict": True,
}
}
)
return response.choices[0].message.content
except Exception as e:
raise AIEndpointRequestError(f"Hugging Face request failed: {e}")