Hugging Face LLM Client

Inputs

Outputs

Source

python

# Input preparation
inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]
model_name = inp["model_name"]
prompt = inp["prompt"]
token = inp["token"]
timeout = inp["timeout"]
max_new_tokens = inp["max_new_tokens"]

# Compute
from huggingface_hub import InferenceClient
from typing import Optional
import os

class HuggingFaceLLMClient:
    """
    A class to interact with a language model hosted on Hugging Face Hub.
    """

    def __init__(self, model_name: str, token: str, timeout: int):
        """
        Initialize the LLMClient with a model name, token, and timeout.

        :param model_name: The name of the model to use.
        :param token: The Hugging Face API token.
        :param timeout: The timeout for the inference request in seconds.
        """
        self.model_name = model_name
        self.token = token
        self.timeout = timeout
        self.client = InferenceClient(model=model_name, token=token, timeout=timeout)

    def generate_response(self, prompt: str, max_new_tokens: int) -> Optional[str]:
        """
        Generate a response from the language model for the given prompt.

        :param prompt: The input prompt for the model (str).
        :param max_new_tokens: The maximum number of tokens to generate.
        :return: The generated response as a string, or None if an error occurs.
        """
        try:
            response = self.client.text_generation(prompt, max_new_tokens=max_new_tokens)
            return response.strip()
        except Exception as e:
            print(f"An unexpected error occurred while generating response: {e}")
        return None

os.environ["HF_TOKEN"] = token

hugging_face_llm_client = HuggingFaceLLMClient(model_name=model_name, token=token, timeout=timeout)

response = hugging_face_llm_client.generate_response(prompt=prompt, max_new_tokens=max_new_tokens)

print(response)

# Output collection (runner reads __ivy_ctx__)
out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["response"] = response

Tests

Requires: python:3.11

• basic-generation

  Generate text response with valid inputs (should succeed).

  huggingfacellmgenerationbasic
• error-invalid-token

  Invalid token should handle error gracefully.

  error-handlingauthentication