Dataclass type constraint issues (Infinite arrays & LogitBias Warnings)

[0xc1c4da]

There is a number of issues when using a Dataclass type constraint with a List[str] property.

Using this example:

@lmql.query(max_len=200)
async def example():
    '''lmql
        from dataclasses import dataclass
        from typing import List

        @dataclass
        class State:
            tasks: List[str]
            role: str
        
        argmax
            """
            Goal: Produce a list of tasks to make a ham sandwich

            The best role to solve this goal is [ROLE]

            Give your answer in the following JSON format, the JSON must be parseable by Python json.loads(<json string>):
                {{ "tasks" : [["task 1", "task 2", ... "task N" ]] }}
            
            Answer:
            [STATE]
            """
        where 
            type(STATE) is State
            and STOPS_BEFORE(ROLE, '\n')
            # and STOPS_BEFORE(STATE, '\n')
    '''

lmql will appear to hang until maximum tokens is reached, printing context.prompt in the types.py stack pop while loop shows an output that looks like this:

Goal: Produce a list of tasks to make a ham sandwich
The best role to solve this goal is 
Give your answer in the following JSON format, the JSON must be parseable by Python json.loads(<json string>):
    { "tasks" : ["task 1", "task 2", ... "task N" ] }

then

(continue in JSON) {"tasks":"Gather the ingredients: bread, ham, lettuce, tomato, mayonnaise, butter, knife.\nSpread butter on one side of each slice of bread.\nPlace the ham on one slice of bread.\nTop the ham with lettuce and tomato.\nSpread mayonnaise on the other slice of bread.\nPut the two slices of bread together.\nCut the sandwich in half.\nEnjoy!","}","","","","","","",

or

(continue in JSON) {"tasks":"Gather the ingredients: bread, ham, cheese, lettuce, tomato, mayonnaise, butter, knife.","Spread butter on one side of each slice of bread.","Lay one slice of bread, butter side down, on a plate.","Layer ham, cheese, lettuce, and tomato on the bread.","Spread mayonnaise on the other slice of bread.","Place the other slice of bread, mayonnaise side down, on top of the sandwich.","Cut the sandwich in half.","Enjoy!","","","","","","","","",

At this point I halt the program because:

• "", repeats indefinitely, until the max token length is reached.
• the output misses the the square brackets for the array.
• sometimes the list is not separated by "," and is instead encapsulated in a string and separated by \n

Other projects that solve JSON structure that lqml's dataclass type constraint requires are

• JSONFormer
• Outlines

[0xc1c4da]

Looks like a culprit is "[" on line 173 in types.py it should be "[["

[lbeurerkellner]

Thanks for reporting this. @dataclass constraints are still in preview and a work-in-progress, so this is valuable feedback. Thanks also for the links.

[0xc1c4da]

No worries! changing line 173 is substantially more reliable

I am now getting consistent correct output with this reduced query

    '''lmql
        from dataclasses import dataclass
        from typing import List

        @dataclass
        class State:
            tasks: List[str]
        
        argmax
            """
            Goal: Produce a list of tasks to make a ham sandwich

            The best role to solve this goal is [ROLE]
            
            Answer:
            [STATE]
            """
        where 
            type(STATE) is State
            and STOPS_BEFORE(ROLE, '\n')
    '''

results in

LMQLResult(prompt='\n'
                   'Goal: Produce a list of tasks to make a ham sandwich\n'
                   'The best role to solve this goal is \n'
                   'Answer:\n'
                   "State(tasks=['Gather ingredients: bread, ham, condiments "
                   "of choice, knife', 'Spread condiments of choice on one "
                   "slice of bread', 'Place ham on top of condiments', 'Place "
                   "second slice of bread on top of ham', 'Cut sandwich in "
                   "half'])\n",
            variables={'ROLE': '',
                       'STATE': State(tasks=['Gather ingredients: bread, ham, '
                                             'condiments of choice, knife',
                                             'Spread condiments of choice on '
                                             'one slice of bread',
                                             'Place ham on top of condiments',
                                             'Place second slice of bread on '
                                             'top of ham',
                                             'Cut sandwich in half'])},
            distribution_variable=None,
            distribution_values=None)]

Not sure why role is empty here (Edit: fixed it with a tokens len constraint), the dataclass with role field works however.

[0xc1c4da]

Hijacking the same issue, I also notice it's very easy to hit the OpenAILogitBiasLimitationWarning: the required logit_bias is too large to be handled by the OpenAI API and will be limited to the first 300 tokens.

Just adding one more string property (objective) to the previous dataclass causes the issue.

Either of these examples results in the warning

@dataclass
class State:
    objective: str
    tasks: List[str]
    role: str

or

@dataclass
class Task:
    id: int
    task: str

@dataclass
class State:
    tasks: List[Task]

[lbeurerkellner]

On the first one I don't get warnings on logit bias limitations, but the second one is due to the integer constraining. This is an API limitation from the OpenAI side.

Check out some experimental changes I tried out in 376736a (not on main), where I tweaked the type query code a bit, by adding some more precise constraints, to avoid unnecessary extra tokens.

[0xc1c4da]

Yup it was definitely the integer constraint, and types-fixes branch seems fine for List[str] the output is:

[LMQLResult(prompt='\n'
                   'Goal: Produce a list of tasks to make a ham sandwich\n'
                   'The best person to solve this goal would be: \n'
                   'A Chef or Cook\n'
                   'Answer:\n'
                   "State(tasks=['Gather ingredients: ham, bread, condiments, "
                   "and other desired toppings', 'Spread condiments on one "
                   "side of the bread', 'Place ham on the bread', 'Add desired "
                   "toppings', 'Place the other slice of bread on top'], "
                   "role='Chef or Cook')\n",
            variables={'ROLE': '\nA Chef or Cook',
                       'STATE': State(tasks=['Gather ingredients: ham, bread, '
                                             'condiments, and other desired '
                                             'toppings',
                                             'Spread condiments on one side of '
                                             'the bread',
                                             'Place ham on the bread',
                                             'Add desired toppings',
                                             'Place the other slice of bread '
                                             'on top'],
                                      role='Chef or Cook')},
            distribution_variable=None,
            distribution_values=None)]

for lists of dataclasses ie List[Task], it will only return 1 value in the list (I tried modifying the query to force a list of tasks first)

    '''lmql
        from dataclasses import dataclass
        from typing import List

        @dataclass
        class Task:
            text: str

        @dataclass
        class State:
            tasks: List[Task]
            role: str
        
        argmax
            """
            Goal: Produce a list of tasks to make a ham sandwich

            The best person to solve this goal would be: [ROLE]
            
            Tasks:
            - [TASKS]

            Tasks as JSON:
            [STATE]
            """
        where 
            type(STATE) is State
            and STOPS_BEFORE(ROLE, '\n')
            and len(TOKENS(ROLE)) > 1
            # and STOPS_BEFORE(STATE, '\n')
    '''

outputs

[LMQLResult(prompt='\n'
                   'Goal: Produce a list of tasks to make a ham sandwich\n'
                   'The best person to solve this goal would be: \n'
                   'A Chef or Cook\n'
                   'Tasks:\n'
                   '-  Gather ingredients: ham, bread, condiments (mayonnaise, '
                   'mustard, etc.), lettuce, tomato, cheese (optional)\n'
                   '-  Slice the ham\n'
                   '-  Toast the bread\n'
                   '-  Spread condiments on the bread\n'
                   '-  Place the ham on the bread\n'
                   '-  Add lettuce, tomato, and cheese (optional)\n'
                   '-  Place the other slice of bread on top\n'
                   '-  Cut the sandwich in half\n'
                   'Tasks as JSON:\n'
                   "State(tasks=[Task(text='Gather ingredients: ham, bread, "
                   'condiments (mayonnaise, mustard, etc.), lettuce, tomato, '
                   "cheese (optional)')], role='Chef or Cook')\n",
            variables={'ROLE': '\nA Chef or Cook',
                       'STATE': State(tasks=[Task(text='Gather ingredients: '
                                                       'ham, bread, condiments '
                                                       '(mayonnaise, mustard, '
                                                       'etc.), lettuce, '
                                                       'tomato, cheese '
                                                       '(optional)')],
                                      role='Chef or Cook'),
                       'TASKS': ' Gather ingredients: ham, bread, condiments '
                                '(mayonnaise, mustard, etc.), lettuce, tomato, '
                                'cheese (optional)\n'
                                '-  Slice the ham\n'
                                '-  Toast the bread\n'
                                '-  Spread condiments on the bread\n'
                                '-  Place the ham on the bread\n'
                                '-  Add lettuce, tomato, and cheese '
                                '(optional)\n'
                                '-  Place the other slice of bread on top\n'
                                '-  Cut the sandwich in half'},
            distribution_variable=None,
            distribution_values=None)]

[younisshah]

I just started out with LMQL and I am getting a similar error with a simple LMQL program (example reproduced from the docs):

@lmql.query
def sentiment(review):
    """lmql
    "Review: {review}"
    "Q: What is the underlying sentiment of this review and why?"
    "A: [ANALYSIS]" where not "\n" in ANALYSIS

    "Based on this, the overall sentiment of the message\
     can be considered to be [CLS]" where CLS in ["positive", "neutral", "negative"]

    return ANALYSIS, CLS
    """

I am using Azure OpenAI - GPT-4. This is my model:

gpt_azure = lmql.model(
    "openai/gpt-4",
    api_type="azure-chat",
    api_base="https://<MY_BASE>.openai.azure.com/",
    api_key="MY_KEY",
    api_version="2023-03-15-preview",
    verbose=True
)

And this is how I call the sentiment function:

print(sentiment(model=gpt_azure, review="We had a great stay. Hiking in the mountains was fabulous and the food is "
                                        "really good"))

Not sure what the issue is.

[kentjhall]

Looks like a culprit is "[" on line 173 in types.py it should be "[["

Is there any movement on getting this fix patched in? I ran into the same issue and this definitely seems to work. Seems pretty trivial to fix, and this was reported back in August.

[lbeurerkellner]

Dataclass support is a feature in preview. The core team is more focused on the next major version of LMQL right now which will also help with this particular feature, but PRs are welcome for the current revision as well :)