Files
Yi-1.5-9B-Chat-16K-abliterated/abliterate-yi-1-5-9b-chat-16k.ipynb

1 line
1.1 MiB
Plaintext
Raw Normal View History

{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[],"dockerImageVersionId":30762,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# LLM Abliteration v1.3 script, adapted for 01-ai/Yi-1.5-9B-Chat-16K\n\nAuthor: byroneverson\n\nThis script ran at kaggle.com, accelerator: P100, persistence: Files only","metadata":{}},{"cell_type":"markdown","source":"# Download original abliterator script for harmful and harmless instructions txt files\nCredit: https://github.com/Sumandora/remove-refusals-with-transformers","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/working\n!git clone https://github.com/Sumandora/remove-refusals-with-transformers.git","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Install requirements for the entire process\n\nAssumes your environment already has all the usual libs installed (transformers, etc). This script was run with kaggle so only a few extras are needed.","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/working\n!pip install jaxtyping\n!pip install transformers\n!pip install tqdm\n!pip install einops\n!pip install torch\n!pip install bitsandbytes\n!pip install accelerate","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Download 01-ai/Yi-1.5-9B-Chat-16K model locally","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/working\n\nfrom huggingface_hub import snapshot_download\nsnapshot_download(repo_id=\"01-ai/Yi-1.5-9B-Chat-16K\", local_dir=\"./Yi-1.5-9B-Chat-16K\")\n","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"# Generate layer embedding outputs for all sample instructions and save locally","metadata":{}},{"cell_type":"code","source":"%cd /kaggle/working\n\nimport random\nimport os\nimport gc\nimport torch\nfrom transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig\nfrom datasets import load_dataset\nimport jaxtyping\nimport einops\nfrom tqdm import tqdm\n\n# Clear memory of past model usage\nmodel = None\ntokenizer = None\ngc.collect()\ntorch.cuda.empty_cache()\n\ntorch.inference_mode()\n\nlocal_repo_dir = \"/kaggle/working/Yi-1.5-9B-Chat-16K\"\nworking_dir = \"/kaggle/working\"\n\nmodel = AutoModelForCausalLM.from_pretrained(local_repo_dir, local_files_only=True, trust_remote_code=True, \n torch_dtype=torch.float16, \n quantization_config=BitsAndBytesConfig(load_in_4bit=True,\n bnb_4bit_compute_dtype=torch.float16))\ntokenizer = AutoTokenizer.from_pretrained(local_repo_dir, local_files_only=True, trust_remote_code=True)\ntokenizer.pad_toke_id = 0\n\n# Settings\n# I have used 128 and 256 with success but may as well use the max for a better estimation\ninstructions = 512\n#layer_idx = int(len(model.model.layers) * 0.5) #6)\n\nprint(\"Instruction count: \" + str(instructions))\n\ndataset = load_dataset(\"byroneverson/abliterate-refusal\", split=\"train\")\n\n# Filter the dataset based on 'target'\nharmful_dataset = dataset.filter(lambda x: x['target'] == True)\nharmless_dataset = dataset.filter(lambda x: x['target'] == False)\n\n# Randomly select 512 entries from each filtered dataset\nharmful_instructions = random.sample(harmful_dataset['prompt'], instructions)\nharmless_instructions = random.sample(harmless_dataset['prompt'], instructions)\n\ndataset = None\nharmful_dataset = None\nharmless_dataset = No