Lable Mapping not aligned when prefixes are set to be an empty string

[Punchwes]

Hi @ericwtodd

I am trying to play with different prefixes. The default one is: prefixes = {"input":"Q:", "output":"A:","instructions":""}. When I change it to empty string prefixes = {"input":"", "output":"","instructions":""}, I get a error saying that:

activation_storage[n] = stack_filtered
RuntimeError: The expanded size of the tensor (54) must match the existing size (53) at non-singleton dimension 2.  Target sizes: [40, 40, 54, 128].  Tensor sizes: [40, 40, 53, 128]

I have further checked it and found that this is caused by different label size (the size of idx_map is different from dummy_labels).
When the prefix is made of empty strings, the function compute_duplicated_labels() will an index_map that is of different size of dummy label. I try to debug where went wrong and found it too hard to dive in...

The error only pops up when I have empty strings for prefixes. All other prefixes do not have this problem....

It would be great if you could share some insights on this. I wonder have you experienced similar problems when using empty string prefixes...

[ericwtodd]

Hey @Punchwes, thanks for raising the issue. I previously spent a while trying to get this auto-labeling scheme to work generally, but it looks like it still has some bugs, so thanks! I had never tried using completely empty prefixes, but I can indeed replicate the issue with the llama-2 family of models. Other models (gpt-j & llama-3) seem to work okay with empty prefixes so I wonder if it's a tokenization quirk of how llama-2 deals with the space in front of the first token.
The goal of this code is to extend labels across tokenization of words, but sometimes the tokenization doesn't work as I'd expect - which is where the issue is coming from.

It looks like the issue happens when the dummy labels are created - as they conflict with the actual labels. A position that is normally labeled as "structural" suddenly becomes re-labeled as a second "demonstration_1_token" when no prefixes are present.

A temporary fix for this makes two changes in the utils.prompt_utils file. The first is to add the following code snippet at the end of the get_token_meta_labels function right before the return statement.:

if prompt_data['prefixes']['input'] == '' and prompt_data['prefixes']['instructions'] == '' and prompt_data['instructions'] == '':
    token_labels[1] = (1, '', 'structural_token')

The second is to include a similar snippet to the end of the get_dummy_token_labels function right before the return statement:

if prefixes is not None and prefixes["input"] == "" and 'llama' in model_config['name_or_path'].lower():
    final_token_labels[1] = (1, 'structural_token')

I will look into this issue a bit more to see if there's a more clean way to get around this issue a little later. But hopefully this will unblock you!