tgi_client_inference_workflow - Infernet ML
Description: Infernet ML Documentation

Skip to content

# Module: `tgi_client_inference_workflow`

## TGI Client Inference Workflow

A class that uses the HuggingFace Text Generation Inference client to run LLM inference on any TGI-compliant inference server.

### Additional Installations

Since this workflow uses some additional libraries, you'll need to install `infernet-ml[tgi_inference]`. Alternatively, you can install those packages directly. The optional dependencies `"[tgi_inference]"` are provided for your convenience.

uv pip

    uv pip install "infernet-ml[tgi_inference]"

    pip install "infernet-ml[tgi_inference]"

### Example Usage

In the example below we use an API key from Hugging Face to access the `Mixtral-8x7B-Instruct-v0.1` model. You can obtain an API key by signing up on the HuggingFace website.

    import os
    from infernet_ml.workflows.inference.tgi_client_inference_workflow import (
        TGIClientInferenceWorkflow,
        TgiInferenceRequest,
    )

    def main():
        server_url = "https://api-inference.huggingface.co/models/mistralai/Mixtral-8x7B-Instruct-v0.1"

        # Instantiate the workflow
        workflow: TGIClientInferenceWorkflow = TGIClientInferenceWorkflow(
            server_url,
            timeout=10,
            headers={"Authorization": f"Bearer {os.environ['HF_TOKEN']}"},
        )

        # Setup the workflow
        workflow.setup()

        # Run the inference
        res = workflow.inference(TgiInferenceRequest(text="Is the sky blue during a clear day?"))
        print(f"response: {res}")

        # Stream the inference
        collected_res = ""
        for r in workflow.stream(TgiInferenceRequest(text="Is the sky blue during a clear day?")):
            collected_res += r.token.text
        print(f"streaming: {collected_res}")

    if __name__ == "__main__":
        main()

Outputs:

    response:

    Yes, the sky is blue during a clear day.

    streaming:

    Yes, the sky is blue during a clear day.

### More Information

For more info, check out the reference docs below.

## `TGIClientInferenceWorkflow`

Bases: `BaseInferenceWorkflow`

Inference workflow for requesting LLM inference on TGI-compliant inference servers.

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    116
    117
    118
    119
    120
    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    152
    153
    154
    155
    156
    157
    158
    159
    160
    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    174
    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    190
    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    211
    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    223
    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    235
    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    265
    266
    267
    268
    269
    270
    class TGIClientInferenceWorkflow(BaseInferenceWorkflow):
        """
        Inference workflow for requesting LLM inference on TGI-compliant inference servers.
        """

        def __init__(
            self,
            server_url: str,
            timeout: int = 30,
            headers: dict[str, str] | None = None,
            cookies: dict[str, str] | None = None,
            retry_params: Optional[RetryParams] = None,
            **inference_params: dict[str, Any],
        ) -> None:
            """
            Constructor. Any named arguments passed to LLM during inference.

            Args:
                server_url (str): url of inference server
            """
            super().__init__()
            self.client: Client = Client(
                server_url, timeout=timeout, headers=headers, cookies=cookies
            )
            self.inference_params: dict[str, Any] = inference_params

            self.retry_params = {
                **DEFAULT_RETRY_PARAMS.model_dump(),
                "exceptions": (
                    ShardNotReadyError,
                    ShardTimeoutError,
                    RateLimitExceededError,
                    OverloadedError,
                ),
                **({} if retry_params is None else retry_params.model_dump()),
            }

        def do_setup(self) -> bool:
            """
            no specific setup needed
            """
            # dummy call to fail fast if client is misconfigured
            self.client.generate("hello", **self.inference_params)
            return True

        def do_preprocessing(self, input_data: TgiInferenceRequest) -> str:
            """
            Implement any preprocessing of the raw input. For example, you may want to
            append additional context. By default, returns the value associated with the
            text key in a dictionary.

            Args:
                input_data (TgiInferenceRequest): user input

            Returns:
                str: transformed user input prompt
            """
            return input_data.text

        def do_postprocessing(self, input_data: TgiInferenceRequest, gen_text: str) -> str:
            """
            Implement any postprocessing here. For example, you may need to return
            additional data. By default returns a dictionary with a single
            output key.

            Args:
                input_data (TgiInferenceRequest): user input
                gen_text (str): generated text from the model.

            Returns:
                str: transformed llm output
            """

            return gen_text

        def generate_inference(self, preprocessed_data: str) -> str:
            """Use tgi client to generate inference.

            Args:
                preprocessed_data (str): input to tgi

            Returns:
                str: output of tgi inference
            """

            @retry(**self.retry_params)
            def _run() -> str:
                return cast(
                    str,
                    self.client.generate(
                        preprocessed_data, **self.inference_params
                    ).generated_text,
                )

            return _run()

        def stream(self, input_data: TgiInferenceRequest) -> Iterator[StreamResponse]:
            """
            Stream results from the model.

            Args:
                input_data (TgiInferenceRequest): user input

            Returns:
                Iterator[StreamResponse]: stream of results
            """
            yield from super().stream(input_data)

        def do_stream(self, _input: str) -> Iterator[StreamResponse]:
            """
            Stream results from the model.

            Args:
                _input (str): user input

            Returns:
                Iterator[StreamResponse]: stream of results
            """
            yield from self.client.generate_stream(_input, **self.inference_params)

        def do_run_model(self, prompt: str) -> str:
            """
            Run the model with the given prompt.

            Args:
                prompt (str): user prompt

            Returns:
                Any: result of inference
            """
            try:
                return self.generate_inference(prompt)
            except (
                BadRequestError,
                GenerationError,
                IncompleteGenerationError,
                NotFoundError,
                NotSupportedError,
                OverloadedError,
                RateLimitExceededError,
                ShardNotReadyError,
                ShardTimeoutError,
                UnknownError,
                ValidationError,
            ) as e:
                # we catch expected service exceptions and return ServiceException
                # this is so we can handle unexpected vs. expected exceptions
                # downstream
                raise InfernetMLException(e) from e

        def do_generate_proof(self) -> Any:
            """
            raise error by default
            """
            raise NotImplementedError

### `__init__(server_url, timeout=30, headers=None, cookies=None, retry_params=None, **inference_params)`

Constructor. Any named arguments passed to LLM during inference.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`server_url` | `str` | url of inference server | required

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    121
    122
    123
    124
    125
    126
    127
    128
    129
    130
    131
    132
    133
    134
    135
    136
    137
    138
    139
    140
    141
    142
    143
    144
    145
    146
    147
    148
    149
    150
    151
    def __init__(
        self,
        server_url: str,
        timeout: int = 30,
        headers: dict[str, str] | None = None,
        cookies: dict[str, str] | None = None,
        retry_params: Optional[RetryParams] = None,
        **inference_params: dict[str, Any],
    ) -> None:
        """
        Constructor. Any named arguments passed to LLM during inference.

        Args:
            server_url (str): url of inference server
        """
        super().__init__()
        self.client: Client = Client(
            server_url, timeout=timeout, headers=headers, cookies=cookies
        )
        self.inference_params: dict[str, Any] = inference_params

        self.retry_params = {
            **DEFAULT_RETRY_PARAMS.model_dump(),
            "exceptions": (
                ShardNotReadyError,
                ShardTimeoutError,
                RateLimitExceededError,
                OverloadedError,
            ),
            **({} if retry_params is None else retry_params.model_dump()),
        }

### `do_generate_proof()`

raise error by default

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    266
    267
    268
    269
    270
    def do_generate_proof(self) -> Any:
        """
        raise error by default
        """
        raise NotImplementedError

### `do_postprocessing(input_data, gen_text)`

Implement any postprocessing here. For example, you may need to return additional data. By default returns a dictionary with a single output key.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data` | `TgiInferenceRequest` | user input | required
`gen_text` | `str` | generated text from the model. | required

Returns:

Name | Type | Description
--- | --- | ---
`str` | `str` | transformed llm output

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    175
    176
    177
    178
    179
    180
    181
    182
    183
    184
    185
    186
    187
    188
    189
    def do_postprocessing(self, input_data: TgiInferenceRequest, gen_text: str) -> str:
        """
        Implement any postprocessing here. For example, you may need to return
        additional data. By default returns a dictionary with a single
        output key.

        Args:
            input_data (TgiInferenceRequest): user input
            gen_text (str): generated text from the model.

        Returns:
            str: transformed llm output
        """

        return gen_text

### `do_preprocessing(input_data)`

Implement any preprocessing of the raw input. For example, you may want to append additional context. By default, returns the value associated with the text key in a dictionary.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data` | `TgiInferenceRequest` | user input | required

Returns:

Name | Type | Description
--- | --- | ---
`str` | `str` | transformed user input prompt

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    161
    162
    163
    164
    165
    166
    167
    168
    169
    170
    171
    172
    173
    def do_preprocessing(self, input_data: TgiInferenceRequest) -> str:
        """
        Implement any preprocessing of the raw input. For example, you may want to
        append additional context. By default, returns the value associated with the
        text key in a dictionary.

        Args:
            input_data (TgiInferenceRequest): user input

        Returns:
            str: transformed user input prompt
        """
        return input_data.text

### `do_run_model(prompt)`

Run the model with the given prompt.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`prompt` | `str` | user prompt | required

Returns:

Name | Type | Description
--- | --- | ---
`Any` | `str` | result of inference

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    236
    237
    238
    239
    240
    241
    242
    243
    244
    245
    246
    247
    248
    249
    250
    251
    252
    253
    254
    255
    256
    257
    258
    259
    260
    261
    262
    263
    264
    def do_run_model(self, prompt: str) -> str:
        """
        Run the model with the given prompt.

        Args:
            prompt (str): user prompt

        Returns:
            Any: result of inference
        """
        try:
            return self.generate_inference(prompt)
        except (
            BadRequestError,
            GenerationError,
            IncompleteGenerationError,
            NotFoundError,
            NotSupportedError,
            OverloadedError,
            RateLimitExceededError,
            ShardNotReadyError,
            ShardTimeoutError,
            UnknownError,
            ValidationError,
        ) as e:
            # we catch expected service exceptions and return ServiceException
            # this is so we can handle unexpected vs. expected exceptions
            # downstream
            raise InfernetMLException(e) from e

### `do_setup()`

no specific setup needed

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    153
    154
    155
    156
    157
    158
    159
    def do_setup(self) -> bool:
        """
        no specific setup needed
        """
        # dummy call to fail fast if client is misconfigured
        self.client.generate("hello", **self.inference_params)
        return True

### `do_stream(_input)`

Stream results from the model.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`_input` | `str` | user input | required

Returns:

Type | Description
--- | ---
`Iterator[StreamResponse]` | Iterator[StreamResponse]: stream of results

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    224
    225
    226
    227
    228
    229
    230
    231
    232
    233
    234
    def do_stream(self, _input: str) -> Iterator[StreamResponse]:
        """
        Stream results from the model.

        Args:
            _input (str): user input

        Returns:
            Iterator[StreamResponse]: stream of results
        """
        yield from self.client.generate_stream(_input, **self.inference_params)

### `generate_inference(preprocessed_data)`

Use tgi client to generate inference.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`preprocessed_data` | `str` | input to tgi | required

Returns:

Name | Type | Description
--- | --- | ---
`str` | `str` | output of tgi inference

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    191
    192
    193
    194
    195
    196
    197
    198
    199
    200
    201
    202
    203
    204
    205
    206
    207
    208
    209
    210
    def generate_inference(self, preprocessed_data: str) -> str:
        """Use tgi client to generate inference.

        Args:
            preprocessed_data (str): input to tgi

        Returns:
            str: output of tgi inference
        """

        @retry(**self.retry_params)
        def _run() -> str:
            return cast(
                str,
                self.client.generate(
                    preprocessed_data, **self.inference_params
                ).generated_text,
            )

        return _run()

### `stream(input_data)`

Stream results from the model.

Parameters:

Name | Type | Description | Default
--- | --- | --- | ---
`input_data` | `TgiInferenceRequest` | user input | required

Returns:

Type | Description
--- | ---
`Iterator[StreamResponse]` | Iterator[StreamResponse]: stream of results

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    212
    213
    214
    215
    216
    217
    218
    219
    220
    221
    222
    def stream(self, input_data: TgiInferenceRequest) -> Iterator[StreamResponse]:
        """
        Stream results from the model.

        Args:
            input_data (TgiInferenceRequest): user input

        Returns:
            Iterator[StreamResponse]: stream of results
        """
        yield from super().stream(input_data)

## `TgiInferenceRequest`

Bases: `BaseModel`

Represents an TGI Inference Request

Source code in `src/infernet_ml/workflows/inference/tgi_client_inference_workflow.py`

    108
    109
    110
    111
    112
    113
    class TgiInferenceRequest(BaseModel):
        """
        Represents an TGI Inference Request
        """

        text: str  # query to the LLM backend