For the detailed repliation package, please go to: https://github.com/invisiblehead/Sparse_Attention_on_Transformer-based_model
For the processed dataset, please go to: https://zenodo.org/deposit/7606184 Also, the original dataset without preprocessing can be found here: https://osf.io/d45bw/
$ conda create -n ltp python=3.8
$ conda activate ltp
$ pip install -r requirements.txt
Replace the standard Transformer implementation files to support the functionality of token pruning based one sparsity attention. Here are files we implement and change for the implementation of Learned Token Pruning (LTP) based on the code base of Hugging Face (https://huggingface.co/) Transformer. All required files are under transformer_revised/:
You may also use revise_local_transformer_package.sh to replace the files under /anaconda/envs/ltp/lib/python3.8/site-packages/transformers/
. Please take care of the src
and des
.
-
trainer_ltp.py
: replaced the trainer.py with trainer_ltp.py in my implementation -
trainer.py
: fix minor deprecated error -
modeling_utils.py
: minor change to fix error -
__init__.py
: update ltp -
training_args.py
: update the arguments of ltp(lr_threshold, lambda_threshold, weight_decay_threshold, masking_mode, temperature, etc.) to pass it to longformer -
models/
-
__init__.py
: update ltp -
ltp/
: major change on most of the files within this folder__init__.py
: replace the deprecated modules like _BaseLazyModule with _LazyModuleconfiguration_ltp.py
ltp_model_output.py
: add global attention and attention_maskmodeling_ltp.py
: replace the IBERT with longformer, given LTP implemented based on IBERT, update the attention mechanism only for sequence classification taskprune_modules.py
: implement on how to get average score of local and global attention for pruning in LTP
-
auto/
: updated the ltp configurations given LTP is not clarifiedconfiguration_auto.py
: updated the ltp configurationsmodeling_auto.py
: updated the ltp configurations
-
Once the conda virtual env: ltp is set up on your VM (Please don't update the default transformer in this virtual env, since it takes time and effort to set up ltp in huggingface transformer.):
$ conda activate ltp
You can also go to https://github.com/kssteven418/LTP/blob/ltp/main/README.md for the meaning of each argument.
This is just a demo about how to play with the code and run the whole process. Please use your own train_file and val_file, and update with your own output_dir. And update these arguments based the your task and the computing power of your GPU. For security defect detection datasets, please specify the task name here:
$ python ./run_glue_code_metrics.py --model_name_or_path allenai/longformer-base-4096 --do_train --do_eval --train_file ./Data/train_subset.csv --validation_file ./Data/test.csv --max_seq_length 1024 --per_device_train_batch_size 4 --learning_rate 2e-5 --num_train_epochs 3 --overwrite_output_dir --output_dir code_test
After generate the checkpoint, the final model will be checkpointed in cmd_test.
- Remove cmd_test/trainer_state.json
- In the configuration file cmd_test/config.json, change (1) "architectures" to ["LTPForSequenceClassification"] and (2) "model_type" to "ltp"
- Add the following lines in the configuration file cmd_test/config.json: "prune_mode": "absolute_threshold", "final_token_threshold": 0.01,
Final_token_threshold determines the token threshold of the last layer, and the thresholds of the remaining layers will be linearly scaled. For instance, the thresholds for the 3rd, 6th, and 9th layers will be 0.0025, 0.005, and 0.0075, respectively, when setting the final_token_threshold , i.e., the threshold for the last (12th) layer, to 0.01. This number is a hyperparameter, and we found that 0.01 works well in many cases.
3. Run the soft pruning (remember to delete the previously generated checkpoint in the checkpoints folder before you rerun the soft pruning for the same task with the same parameter in case that it won’t overwrite):
$ python ./run_code_metrics.py --arch ltp-base --task CODE --restore code_test --train_file ./Data/train_subset.csv --validation_file ./Data/test.csv --lr 2e-5 --temperature 1e-5 --lambda 0.1 --weight_decay 0 --bs 4 --masking_mode soft --epoch 3 --save_step 100 --no_load
The final model will be checkpointed in {CKPT_soft} = checkpoints/base/CODE/absolute_threshold/rate_{final_token_threshold}/temperature_{T}/lambda_{lambda}/lr_{lr}. Remove trainer_state.json from the checkpoint file in {CKPT_soft}.
$ python ./run_code_metrics.py --arch ltp-base --task CODE --restore checkpoints/base/CODE/absolute_threshold/rate_0.03/temperautre_1e-05/lambda_0.1/lr_2e-05 --train_file ./Data/train_subset.csv --validation_file ./Data/test.csv --lr 2e-5 --bs 4 --masking_mode hard --epoch 1 --save_step 100 --eval
The final model will be checkpointed in {CKPT_soft}/hard/lr_{LR}
5. Run the visualization(demo for text classification task). (Please make sure the pathes of train_file, validation file and model_name_or_path are correct):
$ bash test.sh
Please see ltp_test.py for more information
6. save the dataset for powershell after token pruning. (Please make sure the pathes of train_file, validation file, model_name_or_path and the output dir are correct):
$ bash test.sh
please see : ltp_extract_and_save.py for more information