Added space in prompt for better infills and proper stop token

[Kevsnz]

Hi!

I've been experimenting with CodeLlama FIM last couple of days. What I discovered is that CodeLlama gives more robust results when sentinel tokens in prompt is surrounded by spaces. It's especially noticeable at the beginning of the file when 'system' part of prefix dominates the prompt.

Typical failure is shown below:

After the fix model fills the code properly:

I also changed stop tokens. I added <EOT> and <EOD> according to the paper and removed <PRE>, <SUF> and <MID> tokens since they stopped showing up in model responses after I fixed the prompt.

[Kevsnz]

So, after some more digging I think I've found the culprit.

I noticed that this happens only in Windows. Windows uses \r\n as newlines. For some reason CodeLlama (or it's tokenizer) fails to recognize <SUF> token if it's followed by \r.

That's why adding space after <SUF> stops those infilling fails. Removing \r after the token also fixes the problem.

Experiments

Failure: <SUF>\r\n

OK: <SUF> \r\n

OK: <SUF>\n (funny thing is that the model has put missing \r in generated completion 🤣):

Other possible way to fix it could be to replace all occurrences of \r\n' with \nin the prompt and put\r`s back into model output, and do this only in Windows. I'm not sure how to check which OS is running and/or if it's possible in VSCode at all.

[ex3ndr]

Wow, that's crazy! did they changed FIM tokens? Which paper you are referring to?

[Kevsnz]

It doesn't look like those tokens changed.

In the Code Llama paper (part 2.3) they refer to Bavarian et al. (2022) (part 3) where prompt format for fill-in-the-middle is described. <EOT> token finishes the whole sequence there.

However I can't find where I saw <EOD> token, probably in the source files somewhere on HugginFace. I haven't seen it in model's output, so maybe it can be removed.

Also, today I encountered same FIM failure on my Mac (using original extension version) where there is no \rs in prompts. Adding space after <SUF> token seem to have fixed the problem there.

[ex3ndr]

Perfect, thanks!