Package: pastalib
🍝 PASTA: Post-hoc Attention Steering for LLMs 🍝Tell Your Model Where to Attend: Post-hoc Attention Steering for LLMs(Zhang et al. 2023).PASTA allows a user to improve LLM controllability by simply emphasizing part of the prompt (e.g. the instruction) that the LLM should focus on. It requires no changes to LLM weights and no increase in inference time.Quickstart -- use PASTA for improved inferenceInstallpastalib:pipinstall-e.# Alternatively,# pip install pastalib# pip install git+https://github.com/QingruZhang/PASTAInitialize PASTA along with a pre-trained LLM.frompastalib.pastaimportPASTAfromtransformersimportAutoModelForCausalLM,AutoTokenizername="huggyllama/llama-7b"model=AutoModelForCausalLM.from_pretrained(name)tokenizer=AutoTokenizer.from_pretrained(name)# Select the attention heads to be steered,# following the format of {'layer_id': [head_ids]}:head_config={'16':[16,21],'12':[23,26,0],'10':[23,7,22],}# Initialize the PASTA steererpasta=PASTA(model=model,tokenizer=tokenizer,head_config=head_config,alpha=0.01,# scaling coefficientscale_position="exclude",# downweighting unselected tokens)Select specific input spans to emphasize, and then run inference as normal.# Model Inputtexts=["Mary is a doctor. She obtains her bachelor degree from ... Return her occupation in JSON format."]# ===== Before =====# inputs = tokenizer(texts, return_tensors="pt")# outputs = model.generate(**inputs)# ---------------------# ["Mary is a doctor"]  # returns answer in the wrong format# ===== After =====inputs,offset_mapping=pasta.inputs_from_batch(texts)# User highlighs specific input spansemphasized_texts=["Return her occupation in JSON format"]# PASTA registers the pre_forward_hook to edit attentionwithpasta.apply_steering(model=model,strings=texts,substrings=emphasized_texts,model_input=inputs,offsets_mapping=offset_mapping)assteered_model:outputs=steered_model.generate(**inputs)# -------------------------------# ["{'occupation': 'doctor'}"]  # returns answer in the correct formatAdditional Notepastalibworks with any models that apply causal attention by summing up query-key inner product with attention mask and can be applied to LLMs in a plug-and-play manner. For example,LlamaForCausalLMandGPTJForCausalLMfromtransformers, whose attention moduels apply attention masks followingtorch.matmul(query, key) + attention_mask. However,pastalibcurrently only supports LLAMA, LLAMA-2 and GPT-J (more models in progress!).We provide different options ofhead_configfor LLAMA-7B and GPT-J in the folder ofconfig/head_config, including multi-task, task-agnostic and task-specific settings. Please see detailed discussion in our paper.OverviewThe overview of this repo is as follows:pastalib: contains the source code of PASTA libary, which can be applied to models from huggingfacetransformers.evaluation: consists of evaluation pipelines for different tasks, including data/model preprocessing and task evaluators/metrices.config: includes thehead_configfor steering attention modules of LLAMA-7B and GPT-J with PASTA.scripts: consists of running scripts of four tasks: JSON Formatting, Pronouns Changing, Bias in Bios, and CounterFact.EvaluationThe evaluation pipeline are mainly refactored fromREMEDI repo. Please see more details there.Environment SetupSet up the environment with the following commands:condacreate-npastapython=3.10pipinstall-rrequirements.txtpipinstall-e.python-mspacydownloaden_core_web_sm
python-Wignore-mnltk.downloaderpunktcmudictBy default, the preprocessed datasets, models, and results are saved in the local directory of./data,./models, and./results. You can change the directory of their by setting the environment variables:exportCM_DATA_DIR=<datapath>exportCM_MODELS_DIR=<modelspath>exportCM_RESULTS_DIR=<resultspath>Dataset SetupFor CounterFact, our scripts can automatically download the dataset.For Bias in Bios, we cannot release the dataset without the authorization. The dataset must be downloaded withthe offical release. After downloading the data examples into theBIOS.pklfile, you can run the following scripts:pythonreformat_dataset.py\--biasbios_raw_path<pathofBIOS.pkl>\--biasbios_save_filebiasbios.jsonThen, the preprocessed biasbios dataset file will be saved inCM_DATA_DIR/biasbios.json.EvaluationChoose anyhead_configfiles fromconfig/head_configand evaluation the performance of PASTA with the following command.JSON Formattingpython-mscripts.eval_biasbios_instruction\--taskjson\--apply_pasta\--emphasized_textinstruct\--alpha0.01\--scale_positionexclude\--pasta_head_config<head_config_path>\--modelhuggyllama/llama-7b\--prompt_idx0\--batch_size16\--max_new_tokens128\--experiment_namellama_evaluation\--devicecudaPronouns Changingpython-mscripts.eval_biasbios_instruction\--taskpronchange\--apply_pasta\--emphasized_textinstruct\--alpha0.01\--scale_positionexclude\--pasta_head_config<head_config_path>\--prompt_idx0\--modelhuggyllama/llama-7b\--max_new_tokens128\--batch_size16\--experiment_namellama_evaluation\--devicecudaBias in Biospython-mscripts.eval_bias_gen\--modelhuggyllama/llama-7b\--apply_pasta\--alpha0.01\--scale_positionexclude\--pasta_head_config<head_config_path>\--max_length256\--batch_size16\--experiment_namellama_evaluation\--devicecudaCounterFactpython-mscripts.eval_fact_gen\--modelhuggyllama/llama-7b\--apply_pasta\--alpha0.01\--scale_positionexclude\--pasta_head_config<head_config_path>\--add_unmediated_factTrue\--benchmarksefficacyparaphrasegeneration\--experiment_namellama_evaluationContactPlease contact us or post an issue if you have any questions:Qingru Zhang (qingru.zhang@gatech.edu)Chandan Singh (chansingh@microsoft.com)Liyuan Liu (lucliu@microsoft.com)Xiaodong Liu (xiaodl@microsoft.com)@misc{zhang2023tell,title={TellYourModelWheretoAttend:Post-hocAttentionSteeringforLLMs},author={QingruZhangandChandanSinghandLiyuanLiuandXiaodongLiuandBinYuandJianfengGaoandTuoZhao},year={2023},archivePrefix={arXiv},primaryClass={cs.CL}}