Fix LLaMA tokenization issue

[gakada]

As per the regression script:

model	flan-t5-small	gpt2	opt-125m	gpt-neo-125m	pythia-160m	llama-7b
boolq (master)	65.60 ± 1.50	49.90 ± 1.58	54.80 ± 1.57	60.90 ± 1.54	54.00 ± 1.58	71.10 ± 1.43
boolq (llama)	65.60 ± 1.50	49.90 ± 1.58	54.80 ± 1.57	60.90 ± 1.54	54.00 ± 1.58	73.90 ± 1.39
boolq (diff)	0	0	0	0	0	+2.8
lambada_openai (master)	8.80 ± 0.90	31.30 ± 1.47	38.10 ± 1.54	34.70 ± 1.51	32.20 ± 1.48	0
lambada_openai (llama)	8.80 ± 0.90	31.30 ± 1.47	38.10 ± 1.54	34.70 ± 1.51	32.20 ± 1.48	72.20 ± 1.42
lambada_openai (diff)	0	0	0	0	0	+72.2
winogrande (master)	49.30 ± 1.58	53.10 ± 1.58	51.20 ± 1.58	51.40 ± 1.58	52.30 ± 1.58	67.30 ± 1.48
winogrande (llama)	49.30 ± 1.58	53.10 ± 1.58	51.20 ± 1.58	51.40 ± 1.58	52.30 ± 1.58	70.20 ± 1.45
winogrande (diff)	0	0	0	0	0	+2.9
hellaswag (master)	28.70 ± 1.43	31.00 ± 1.46	30.50 ± 1.46	29.60 ± 1.44	29.40 ± 1.44	58.00 ± 1.56
hellaswag (llama)	28.70 ± 1.43	31.00 ± 1.46	30.50 ± 1.46	29.60 ± 1.44	29.40 ± 1.44	59.30 ± 1.55
hellaswag (diff)	0	0	0	0	0	+1.3
openbookqa (master)	12.00 ± 1.45	16.20 ± 1.65	17.20 ± 1.69	15.80 ± 1.63	15.00 ± 1.60	29.00 ± 2.03
openbookqa (llama)	12.00 ± 1.45	16.20 ± 1.65	17.20 ± 1.69	15.80 ± 1.63	15.00 ± 1.60	34.00 ± 2.12
openbookqa (diff)	0	0	0	0	0	+5
piqa (master)	59.30 ± 1.55	62.20 ± 1.53	62.00 ± 1.54	62.10 ± 1.53	60.80 ± 1.54	77.00 ± 1.33
piqa (llama)	59.30 ± 1.55	62.20 ± 1.53	62.00 ± 1.54	62.10 ± 1.53	60.80 ± 1.54	77.80 ± 1.31
piqa (diff)	0	0	0	0	0	+0.8
wikitext (master)	0	3738.14	3206.76	3430.35	3374.91	954.64
wikitext (llama)	0	3738.14	3206.76	3430.35	3374.91	954.64
wikitext (diff)	0	0	0	0	0	0

[HaniItani]

Hello @gakada and @haileyschoelkopf ,

I would really appreciate it if you can explain the fix to me and why it only affects LLaMA models.

[gakada]

@HaniItani it is because LLaMA tokenizer doesn't satisfy enc(a + b) = enc(a) + enc(b) (which was used in the prediction code, we want logits corresponding to b from model(a + b), so there should be a way to get an encoding for a + b and to separate the b part), so a weaker property should be used, like enc(a + b) = enc(a) + enc(a + b)[len(enc(a)):] (otherwise it is the same, so a + b is encoded and b is separated), e.g.

import torch
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained('huggyllama/llama-7b')
model = AutoModelForCausalLM.from_pretrained('huggyllama/llama-7b', device_map='auto')

enc = lambda s: tokenizer.encode(s, add_special_tokens=False)

def test(question_enc, answer_enc):
    # get b logprobs out of model(a + b)
    logprobs = F.log_softmax(model(torch.tensor(question_enc + answer_enc).unsqueeze(0).cuda())['logits'], dim=-1).cpu().squeeze(0)[len(question_enc) - 1 : -1]
    guess = logprobs.argmax(dim=-1)
    print(f'|{tokenizer.decode(guess)}| |{tokenizer.decode(answer_enc)}|')

# assume enc(a + b) = enc(a) + enc(b)

test(enc('Paris is the capital of'), enc(' France.'))
# |France2,| | France.|

# assume enc(a + b) = enc(a) + enc(a + b)[len(enc(a)):]

sentence_enc = enc('Paris is the capital of France.')
question_enc = enc('Paris is the capital of')
answer_enc = sentence_enc[len(question_enc):]

test(question_enc, answer_enc)
# |France.| |France.|

This doesn't affect other tokenizers, though in all cases it is also assumed that a doesn't end with spaces.

[HaniItani]

@gakada thank you very much for your elaborate response. It's all clear to me now.

[PkuRainBow]

+1