Skip to content

Hf

HuggingFaceAgent

Bases: BaseAgent

Source code in Docs2KG/agents/hf.py 
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
class HuggingFaceAgent(BaseAgent):
    def __init__(self, name: str):
        """
        Initialize HuggingFaceAgent with model name.

        Args:
            name: Name of the model to use (e.g., 'gpt2')
        """
        super().__init__(name)
        self.client = self._init_huggingface_client()

    def _init_huggingface_client(self) -> InferenceClient:
        """
        Initialize HuggingFace client with API token from either:
        1. Environment variable HF_API_TOKEN
        2. .env file
        """
        try:
            # Initialize client with config
            client = InferenceClient(
                model=self.name,
                token=PROJECT_CONFIG.huggingface.api_token.get_secret_value(),
            )

            logger.info(
                f"Successfully initialized HuggingFace client for model {self.name}"
            )
            return client

        except Exception as e:
            logger.error(f"Failed to initialize HuggingFace client: {str(e)}")
            raise

    def process(self, input_data: Any) -> Any:
        """
        Process input using the HuggingFace client.

        Args:
            input_data: The input to be processed by the model

        Returns:
            Dict containing the model response and metadata
        """
        logger.info(f"Processing input with HuggingFace: {input_data}")

        try:
            # Query the model with proper error handling
            response = self.client.text_generation(
                prompt=str(input_data),
                details=True,  # Get detailed response including token counts
                return_full_text=False,  # Only return generated text, not the prompt
            )

            return {
                "model": self.name,
                "input": input_data,
                "status": "processed",
                "response": response.generated_text,
                "usage": {
                    "prompt_tokens": len(str(input_data).split()),
                    "completion_tokens": len(response.generated_text.split()),
                    "total_tokens": len(str(input_data).split())
                    + len(response.generated_text.split()),
                },
            }

        except Exception as e:
            logger.error(f"Error processing input with HuggingFace: {str(e)}")
            raise

__init__(name)

Initialize HuggingFaceAgent with model name.

Parameters:

Name Type Description Default
name str

Name of the model to use (e.g., 'gpt2')

 required
Source code in Docs2KG/agents/hf.py 
11
12
13
14
15
16
17
18
19
def __init__(self, name: str):
    """
    Initialize HuggingFaceAgent with model name.

    Args:
        name: Name of the model to use (e.g., 'gpt2')
    """
    super().__init__(name)
    self.client = self._init_huggingface_client()

process(input_data)

Process input using the HuggingFace client.

Parameters:

Name Type Description Default
input_data Any

The input to be processed by the model

 required

Returns:

Type Description
Any

Dict containing the model response and metadata
Source code in Docs2KG/agents/hf.py 
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
def process(self, input_data: Any) -> Any:
    """
    Process input using the HuggingFace client.

    Args:
        input_data: The input to be processed by the model

    Returns:
        Dict containing the model response and metadata
    """
    logger.info(f"Processing input with HuggingFace: {input_data}")

    try:
        # Query the model with proper error handling
        response = self.client.text_generation(
            prompt=str(input_data),
            details=True,  # Get detailed response including token counts
            return_full_text=False,  # Only return generated text, not the prompt
        )

        return {
            "model": self.name,
            "input": input_data,
            "status": "processed",
            "response": response.generated_text,
            "usage": {
                "prompt_tokens": len(str(input_data).split()),
                "completion_tokens": len(response.generated_text.split()),
                "total_tokens": len(str(input_data).split())
                + len(response.generated_text.split()),
            },
        }

    except Exception as e:
        logger.error(f"Error processing input with HuggingFace: {str(e)}")
        raise