llm-app

1
from llm_app.model_wrappers.api_clients.clients import HuggingFaceClient
2
from llm_app.model_wrappers.base import BaseModel
3

4

5
class HuggingFaceAPIModel(BaseModel):
6
    def __init__(self, **kwargs):
7
        super().__init__()
8
        self.api_client = self.get_client(**kwargs)
9

10
    def get_client(self, **kwargs) -> HuggingFaceClient:
11
        return HuggingFaceClient(**kwargs)
12

13
    def call_api(self, **kwargs):
14
        """
15
        Makes a request to the Hugging Face API and returns the result.
16

17
        The method accepts arguments as keyword arguments (**kwargs).
18
        The expected arguments are 'model' and others that depend on the specific task.
19
        Please check [HuggingFace Inference API](https://huggingface.co/docs/api-inference/detailed_parameters)
20
        'model' is a string representing the pre-trained model to use for the call.
21

22
        Examples:
23

24
        1) Question-Answering Task usage:
25

26
        model = HuggingFaceModel(api_key=token)
27
        result = model.call(
28
            inputs=dict(
29
                context="Pathway is a realtime stream data processing framework. It has a python api",
30
                question="Does Pathway have a python API ?"
31
            ),
32
            model='deepset/roberta-base-squad2'
33
        )
34
        The expected output for the example above is:
35
        {
36
            'score': 0.41046786308288574,
37
            'start': 56,
38
            'end': 75,
39
            'answer': 'It has a python api'
40
        }
41

42
        2) Sentiment analysis model usage:
43

44
        result = model.call(
45
            inputs="What a performance that was!",
46
            model='distilbert-base-uncased-finetuned-sst-2-english'
47
        )
48
        The expected output for the example above is:
49
        [
50
            [
51
                {'label': 'POSITIVE', 'score': 0.9988183379173279},
52
                {'label': 'NEGATIVE', 'score': 0.0011816363548859954}
53
            ]
54
        ]
55

56
        Args:
57
            **kwargs: The arguments for the model call. 'inputs' and 'model' keys are expected.
58

59
        Returns:
60
            The response from the Hugging Face API, the format depends on the model being used.
61
        """
62
        return self.api_client.make_request(**kwargs)
63

64

65
class HFApiFeatureExtractionTask(HuggingFaceAPIModel):
66
    def __call__(
67
        self, text: str, locator="sentence-transformers/all-MiniLM-L6-v2", **kwargs
68
    ):
69
        response = self.call_api(inputs=[text], model=locator, **kwargs)
70
        return response
71

72

73
class HFApiTextGenerationTask(HuggingFaceAPIModel):
74
    """
75
    A class that represents a text generation task using the Hugging Face API.
76

77
    It inherits from the HuggingFaceModel class and overrides the call method to
78
    specifically work with text generation models.
79

80
    This class allows users to simply pass a text string and get a generated
81
    text in return.
82

83
    Args:
84
        api_key (str): The API key to access the Hugging Face API.
85

86
    Example:
87

88
    # >>> model = HFTextGenerationTask(api_key=token)
89
    # >>> text = "Once upon a time"
90
    # >>> generated_text = model(text, locator="gpt2")
91
    # >>> print(generated_text)
92
    'Once upon a time in a land far away...'
93

94
    """
95

96
    def __call__(self, text: str, locator="gpt2", **kwargs):
97
        response = self.call_api(inputs=text, model=locator, **kwargs)
98
        return response[0]["generated_text"]
99