Skip to main content

A Python wrapper around TGI and TEI servers

Project description

Py-TXI (previously Py-TGI)

PyPI version PyPI - Python Version PyPI - Format Downloads PyPI - License Tests

Py-TXI is a Python wrapper around Text-Generation-Inference and Text-Embedding-Inference that enables creating and running TGI/TEI instances through the awesome docker-py in a similar style to Transformers API.

Installation

pip install py-txi

Py-TXI is designed to be used in a similar way to Transformers API. We use docker-py (instead of a dirty subprocess solution) so that the containers you run are linked to the main process and are stopped automatically when your code finishes or fails.

Usage

Here's an example of how to use it:

from py_txi import TGI, is_nvidia_system, is_rocm_system

llm = TGI(config=TGIConfig(sharded="false"))
output = llm.generate(["Hi, I'm a language model", "I'm fine, how are you?"])
print("LLM:", output)
llm.close()

Output: LLM: ["er. I'm a language modeler. I'm a language modeler. I'm a language", " I'm fine, how are you? I'm fine, how are you? I'm fine,"]

from py_txi import TEI, is_nvidia_system

embed = TEI(config=TEIConfig(pooling="cls"))
output = embed.encode(["Hi, I'm an embedding model", "I'm fine, how are you?"])
print("Embed:", output)
embed.close()

Output: [array([[ 0.01058742, -0.01588806, -0.03487622, ..., -0.01613717, 0.01772875, -0.02237891]], dtype=float32), array([[ 0.02815401, -0.02892136, -0.0536355 , ..., 0.01225784, -0.00241452, -0.02836569]], dtype=float32)]

That's it! Now you can write your Python scripts using the power of TGI and TEI without having to worry about the underlying Docker containers.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

py-txi-0.4.0.tar.gz (6.8 kB view hashes)

Uploaded Source

Built Distribution

py_txi-0.4.0-py3-none-any.whl (7.4 kB view hashes)

Uploaded Python 3

Supported by

AWS AWS Cloud computing and Security Sponsor Datadog Datadog Monitoring Fastly Fastly CDN Google Google Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Sentry Sentry Error logging StatusPage StatusPage Status page