Token Alignment via Character Matching for Subword Completion

This repo releasing the code and benchmark datasets for paper "Token Alignment via Character Matching for Subword Completion" in ACL Findings 2024. In our paper, we noticed LLMs usually generate sub-optimal responses when their input prompts ending with partial tokens. The issues are generically associated to tokenization artifacts, and are especially important when applying LLMs in practical applications like code completions. This paper proposes a novel technique, namely token alignment, to well mitigate such issues. We optimized the inference efficiency with pre-built character trie and mask cache. We conducted extensive experiments on code completions and general NL tasks to validate the improvements of our proposed token alignment.

Evaluation of public models on execution datasets

Code to evaluate public models on MBXP datasets.

Usage :

For Falcon models we need pytorch 2.0 and transformers>=4.30

```python CUDAVISIBLEDEVICES=0,1,2,3,4,5,6,7 python3 -m torch.distributed.run --nprocpernode 8 --masterport 20000 evaluatemodel.py \ --modelnameorpath $modelpath/starcoder/falconmodels/models/falcon7binstruct/ \ --dosample \ --temperature 0.4 \ --modelcontextlength 8192 \ --maxgenerationlength 256 \ --numsamplesperexample 5 \ --batchsize 1 \ --taskname mxeval/mbxp \ --programminglang python \ --outputdir ~/results/falcon7binstruction \ --usestoppingcriteria \ --bf16 \ --overrideprevious_results

```

python3 evaluate_model.py \ --model_name_or_path facebook/opt-6.7b \ --tokenizer_name facebook/opt-6.7b \ --model_context_length 512 \ --output_dir ./eval_results/humaneval/opt-6.7b \ --do_sample \ --bf16 \ --debug

python3 -m torch.distributed.launch --nproc_per_node 8 --master_port 20000 evaluate_model.py \ --model_name_or_path bigscience/bloom-6b3 \ --tokenizer_name bigscience/bloom-6b3 \ --model_context_length 512 \ --output_dir ./eval_results/humaneval/bloom-6b3 \ --debug

Starcoder Notes

• Requires
  • transformers>=4.28.1
  • tokenizers>=0.13
• BigCode models tend to generate garbage when prompts have traling newline chars

Other dependencies

``` pip install datasets accelerate pygtrie pip install nvidia-pyindex pip install pytorch-quantization

```

Need to install the latest mxeval as well https://github.com/amazon-science/mxeval

Token Align

Given a prompt below to the model (StarCoder in this case)

```

write a function to get three maximum numbers from a list

def three_max(l): re ```

The model is not likely to generate return due to the fact that return due to the artifact of tokenization. Instead, the model will generate the following:

```

write a function to get three maximum numbers from a list

def three_max(l): re = [] for i in range(len(l)): if i == 0: re.append(l[0]) else: if l[i] > re[0]: re[0] = l[i] elif l[i] > re[1]: re[1] = l[i] elif l[i] > re[2]: re[2] = l[i] return re ```

However, with TokenAlign, the model aligns with the existing tokens and are able to generate correctly.

```

write a function to get three maximum numbers from a list

def three_max(l): return sorted(l, reverse=True)[:3] ```

Below are the code to replicate

With token align CUDA_VISIBLE_DEVICES=0 python3 evaluate_model.py \ --model_name_or_path bigcode/starcoder \ --tokenizer_name bigcode/starcoder \ --model_context_length 2048 \ --output_dir ./eval_results/humaneval/starcoder \ --do_sample \ --bf16 \ --debug \ --use_custom_generate 1 \ --use_token_align 1 \ --single_context_batch_sampling 1 \ --custom_generate_verbose 1 \

Without token align CUDA_VISIBLE_DEVICES=0 python3 evaluate_model.py \ --model_name_or_path bigcode/starcoder \ --tokenizer_name bigcode/starcoder \ --model_context_length 2048 \ --output_dir ./eval_results/humaneval/starcoder \ --do_sample \ --bf16 \ --debug \ --use_custom_generate 1 \ --use_token_align 0 \ --single_context_batch_sampling 1 \ --custom_generate_verbose 1 \

To cite our work

``` @article{tokenalign_2024, title = {Token Alignment via Character Matching for Subword Completion}, author = {Athiwaratkun, Ben and Wang, Shiqi and Shang, Mingyue and Tian, Yuchen and Wang, Zijian and Gonugondla, Sujan Kumar and Gouda, Sanjay Krishna and Kwiatowski, Rob and Nallapati, Ramesh and Xiang, Bing }, url = {https://arxiv.org/pdf/2403.08688}, publisher = {ACL Findings}, year = {2024}, copyright = {Creative Commons Attribution 4.0 International} }

```