Skip to content

Latest commit

 

History

History
104 lines (76 loc) · 4.37 KB

File metadata and controls

104 lines (76 loc) · 4.37 KB
description Native support for LLMs in your ML workflows

LLMs

Aqueduct supports a number of built-in large language models (LLMs) that you can use either standalone as a part of Aqueduct workflows.

Invoking an LLM

The Aqueduct SDK comes prepackaged with an llm_op API:

import aqueduct as aq

client = aq.Client()

# use `aq.supported_llms` to see a list of LLMs we currently support
vicuna = aq.llm_op(
    # the name of the LLM you want to use 
    name="vicuna_7b",
    # the name of the compute engine to run this LLM operator on
    engine="my_k8s_engine",
)

response = vicuna("How do I make my business successful?")

response.get()

In this case, the vicuna function takes in a prompt and returns the text generated by the LLM. The response object return by vicuna() is an Aqueduct artifact, which can be used as an input to other operators in your workflow.

Aqueduct currently supports the following LLMs:

  • LLaMa-7B: llama_7b
  • Vicuna-7B: vicuna_7b
  • Dolly v2-3B: dolly_v2_3b
  • Dolly v2-7B: dolly_v2_7b
  • ... more coming soon!

Processing a dataset

You can also feed a Pandas DataFrame to an LLM operator. When you pass in your dataset, you specify an input column (the name of the column in your dataset) and an output column where the results of the LLM generation will be stored. Aqueduct will automatically invoke the LLM on the column specified and return your dataset augmented with the new column.

When invoking an LLM operator, you can optionally specify a configuration map as the second argument to the invocation. This map takes in model-specific configuration parameters (see aqueduct.llm_op.md for more details) and also allows you to specify a prompt that will be injected before each invocation.

vicuna = aq.llm_op(
    name="vicuna_7b",
    engine="my_k8s_engine",
    column_name="review", # the column to send to the LLM
    output_column_name="response", # the column name for the response
)

snowflake = client.resource("snowflake")

reviews_table = snowflake.sql("select * from hotel_reviews;")

updated_table = vicuna(
    # the dataset we're processing
    reviews_table,
    # an optional list of parameters for the LLM; see below for a link 
    # the full documentation
    {
        "prompt": "Generate a customer service response for this review: ",
        "max_new_tokens": 512,
    },)
updated_table.get() # the `response` column is appended to the original table

Specifying compute engine

As in the example above, you can specify the compute resources to run LLM operators by using the engine argument to the llm_op function. LLM operators are currently optimized for running on Kubernetes. When you specify the engine to be a Kubernetes resource, we automatically generate the appropriate resource requests (CPU, RAM, GPU, etc.) for the LLM you are using.

If you want to run LLM operators on other type of compute engines, you need to make sure the engine has sufficient hardware resources to run the LLM. Please refer to the LLM package documentation for the minimum requirements of each LLM.

Using LLMs via Aqueduct decorators

Although llm_op is the easiest way to invoke LLMs, you can also use the Aqueduct decorators (op, metric, check) to combine LLMs with arbitrary Python code, which gives you more flexibility. All you need to do is to include aqueduct-llm in the requirements list when decorating your function:

from aqueduct import op

@op(
    engine="my_k8s_engine",
    requirements=["aqueduct-llm"],
    resources={
        'memory': '16GB',
        'gpu_resource_name': 'nvidia.com/gpu',
    }
) 
def vicuna(message):
    from aqueduct_llm import vicuna_7b
    # Your custom code here
    # ...
    # ...
    return vicuna_7b.generate(message, max_new_tokens=512)

response = vicuna("How to generate revenue for my business?")

{% hint style="info" %} Since the operator can contain arbitrary Python code, you need to explicitly specify the resources required by the operator when running on Kubernetes. Please refer to the LLM package documentation for the minimum requirements of each LLM. {% endhint %}