LMQL and Pandas 🐼

When used from within Python, LMQL queries can be treated as simple python functions. This means building pipelines with LMQL is as easy as chaining together functions. However, next to easy integration, LMQL queries also offer a guaranteed output format when it comes to the data types and structure of the returned values. This makes it easy to process the output of LMQL queries in a structured way, e.g. with tabular/array-based data processing libraries such as Pandas.

For example, the following code snippet processes the output of an LMQL query with pandas:

import lmql
import pandas as pd

@lmql.query
async def generate_dogs(n: int):
    '''lmql
    sample(n=n)
        """Generate a dog with the following characteristics:
        Name:[NAME]
        Age: [AGE]
        Breed:[BREED]
        Quirky Move:[MOVE]
        """
    from
        "openai/text-davinci-003"
    where
        STOPS_BEFORE(NAME, "\n") and STOPS_BEFORE(BREED, "\n") and
        STOPS_BEFORE(MOVE, "\n") and INT(AGE) and len(AGE) < 3
    '''

result = await generate_dogs(8)
df = pd.DataFrame([r.variables for r in result])
df

	NAME	AGE	BREED	MOVE
0	Houston	8	Beagle	Houston loves to play hide and seek, but inst...
1	Mya	3	Golden Retriever	The Mya shake, a playful shoulder shimmy deve...
2	Sandy	8	Golden Retriever	Sandy loves to wiggle her front legs when she...
3	Ranger	7	Bulldog	Special scratches his ear with his hind leg l...
4	Spot	4	German Shepherd	A somersault during a jump when running
5	Fido	7	Beagle	A funny shuffle dance when playing fetch
6	Tippy	4	Bernese Mountain Dog	Doing the shadow dance
7	Max	8	Golden Retriever	The Wiggle Jump

Note how we sample multiple sequences (i.e. dog instances) using the sample(n=n) decoder configuration. The returned result is a list of lmql.LMQLResult objects (see Python integration), which we can easily convert to a pandas.DataFrame by accessing r.variables on each item.

In the query, we use scripted prompting and constraints to make sure the generated dog samples are valid and each provides the necessary attributes for name, age, breed, and move.

Having the resulting values for LMQL template variables NAME, AGE, BREED, and MOVE in a pandas.DataFrame, makes it easy to apply further processing and work with the generated data. For instance, we can easily determine the average age of the generated dogs:

# determine average age
df["AGE"].mean()

6.0

Note how the INT(AGE) constraints automatically converted the AGE values to integers, which makes the generated AGE values automatically amendable to arithmetic operations such as mean().

Based on this tabular representation, it is now also easy to filter and aggregate the data. For instance, we can easily determine the average age of the generated dogs per breed:

# group df by BREED and compute the average age for each breed
df.groupby("BREED")["AGE"].mean()

BREED
 Alaskan Malamute               8.0
 Dachshund                      6.0
 Golden Retriever               5.0
 Jack Russell Terrier           6.0
 Labrador Retriever             5.0
 Shetland Sheepdog              3.0
 Wirehaired Pointing Griffon    7.0
Name: AGE, dtype: float64