refactor: move script to file

README.md

@@ -17,63 +17,18 @@ datasets:
 
 StableVicuna-13B is a [Vicuna-13B](https://vicuna.lmsys.org/) model fine-tuned using reinforcement learning from human feedback (RLHF) via Proximal Policy Optimization (PPO) on various conversational and instructional datasets.
 
-### Apply Delta weights
+### Apply Delta Weights
 
-```python
-"""
-Usage:
-python3 apply_delta.py --base /path/to/model_weights/llama-13b --target stable-vicuna-13b --delta pvduy/stable-vicuna-13b-delta
-"""
-import argparse
-
-import torch
-from tqdm import tqdm
-from transformers import AutoTokenizer, AutoModelForCausalLM
-
-
-def apply_delta(base_model_path, target_model_path, delta_path):
-    print("Loading base model")
-    base = AutoModelForCausalLM.from_pretrained(
-        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-
-    print("Loading delta")
-    delta = AutoModelForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
-    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
+StableVicuna-13B cannot be used from the `stability/stable-vicuna-13b-delta` weights alone. To obtain the correct model, one must add back the difference between LLaMA 13B and `stability/stable-vicuna-13b-delta` weights. We provide the [`apply_delta.py`](https://huggingface.co/CarperAI/stable-vicuna-13b-delta/raw/main/apply_delta.py) script to automate the conversion, which you can run as:
 
-    DEFAULT_PAD_TOKEN = "[PAD]"
-    base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=False)
-    num_new_tokens = base_tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))
-
-    base.resize_token_embeddings(len(base_tokenizer))
-    input_embeddings = base.get_input_embeddings().weight.data
-    output_embeddings = base.get_output_embeddings().weight.data
-    input_embeddings[-num_new_tokens:] = 0
-    output_embeddings[-num_new_tokens:] = 0
-
-    print("Applying delta")
-    for name, param in tqdm(base.state_dict().items(), desc="Applying delta"):
-        assert name in delta.state_dict()
-        param.data += delta.state_dict()[name]
-
-    print("Saving target model")
-    base.save_pretrained(target_model_path)
-    delta_tokenizer.save_pretrained(target_model_path)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--base-model-path", type=str, required=True)
-    parser.add_argument("--target-model-path", type=str, required=True)
-    parser.add_argument("--delta-path", type=str, required=True)
-
-    args = parser.parse_args()
-
-    apply_delta(args.base_model_path, args.target_model_path, args.delta_path)
+```sh
+python3 apply_delta.py --base /path/to/model_weights/llama-13b --target stable-vicuna-13b --delta stabilityai/stable-vicuna-13b-delta
 ```
 
+
 ## Usage
 
-Quickly get started chatting with the model by using the [`transformers`](https://huggingface.co/docs/transformers) library:
+Once the delta weights are applied, get started chatting with the model by using the [`transformers`](https://huggingface.co/docs/transformers) library:
 
 ```python
 from transformers import AutoTokenizer, AutoModelForCausalLM
@@ -118,14 +73,14 @@ print(tokenizer.decode(tokens[0], skip_special_tokens=True))
 
 ### Training Dataset
 
-`stabilityai/stable-vicuna-13b` is fine-tuned on a mix of three datasets. [OpenAssistant Conversations Dataset (OASST1)](https://huggingface.co/datasets/OpenAssistant/oasst1), a human-generated, human-annotated assistant-style conversation corpus consisting of 161,443 messages distributed across 66,497 conversation trees, in 35 different languages;
+StableVicuna-13B is fine-tuned on a mix of three datasets. [OpenAssistant Conversations Dataset (OASST1)](https://huggingface.co/datasets/OpenAssistant/oasst1), a human-generated, human-annotated assistant-style conversation corpus consisting of 161,443 messages distributed across 66,497 conversation trees, in 35 different languages;
 [GPT4All Prompt Generations](https://huggingface.co/datasets/nomic-ai/gpt4all_prompt_generations), a dataset of 400k prompts and responses generated by GPT-4; and [Alpaca](https://huggingface.co/datasets/tatsu-lab/alpaca),  a dataset of 52,000 instructions and demonstrations generated by OpenAI's text-davinci-003 engine.
 
 The reward model used during RLHF was also trained on [OpenAssistant Conversations Dataset (OASST1)](https://huggingface.co/datasets/OpenAssistant/oasst1) along with two other datasets: [Anthropic HH-RLHF](https://huggingface.co/datasets/Anthropic/hh-rlhf), a dataset of preferences about AI assistant helpfulness and harmlessness; and [Stanford Human Preferences Dataset](https://huggingface.co/datasets/stanfordnlp/SHP) a dataset of 385K collective human preferences over responses to questions/instructions in 18 different subject areas, from cooking to legal advice.
 
 ### Training Procedure
 
-`stabilityai/sstable-vicuna-13b` was trained using PPO as implemented in [`trlX`](https://github.com/CarperAI/trlx/blob/main/trlx/trainer/accelerate_ppo_trainer.py) with the following configuration:
+`stabilityai/sstable-vicuna-13b-delta` was trained using PPO as implemented in [`trlX`](https://github.com/CarperAI/trlx/blob/main/trlx/trainer/accelerate_ppo_trainer.py) with the following configuration:
 
 |  Hyperparameter   |  Value  |
 |-------------------|---------|

apply_delta.py

@@ -0,0 +1,49 @@
+"""
+Usage:
+python3 apply_delta.py --base /path/to/model_weights/llama-13b --target stable-vicuna-13b --delta pvduy/stable-vicuna-13b-delta
+"""
+import argparse
+
+import torch
+from tqdm import tqdm
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+
+def apply_delta(base_model_path, target_model_path, delta_path):
+    print("Loading base model")
+    base = AutoModelForCausalLM.from_pretrained(
+        base_model_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+
+    print("Loading delta")
+    delta = AutoModelForCausalLM.from_pretrained(delta_path, torch_dtype=torch.float16, low_cpu_mem_usage=True)
+    delta_tokenizer = AutoTokenizer.from_pretrained(delta_path)
+
+    DEFAULT_PAD_TOKEN = "[PAD]"
+    base_tokenizer = AutoTokenizer.from_pretrained(base_model_path, use_fast=False)
+    num_new_tokens = base_tokenizer.add_special_tokens(dict(pad_token=DEFAULT_PAD_TOKEN))
+
+    base.resize_token_embeddings(len(base_tokenizer))
+    input_embeddings = base.get_input_embeddings().weight.data
+    output_embeddings = base.get_output_embeddings().weight.data
+    input_embeddings[-num_new_tokens:] = 0
+    output_embeddings[-num_new_tokens:] = 0
+
+    print("Applying delta")
+    for name, param in tqdm(base.state_dict().items(), desc="Applying delta"):
+        assert name in delta.state_dict()
+        param.data += delta.state_dict()[name]
+
+    print("Saving target model")
+    base.save_pretrained(target_model_path)
+    delta_tokenizer.save_pretrained(target_model_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--base-model-path", type=str, required=True)
+    parser.add_argument("--target-model-path", type=str, required=True)
+    parser.add_argument("--delta-path", type=str, required=True)
+
+    args = parser.parse_args()
+
+    apply_delta(args.base_model_path, args.target_model_path, args.delta_path)