├── .gitignore ├── LICENSE.txt ├── README.md ├── cog.yaml ├── predict.py ├── scripts ├── download_weights.py └── tensorize_model.py ├── subclass.py └── tokenizer ├── special_tokens_map.json ├── tokenizer.json └── tokenizer_config.json /.gitignore: -------------------------------------------------------------------------------- 1 | .cog 2 | __pycache__ -------------------------------------------------------------------------------- /LICENSE.txt: -------------------------------------------------------------------------------- 1 | Apache License 2 | Version 2.0, January 2004 3 | http://www.apache.org/licenses/ 4 | 5 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 6 | 7 | 1. Definitions. 8 | 9 | "License" shall mean the terms and conditions for use, reproduction, 10 | and distribution as defined by Sections 1 through 9 of this document. 11 | 12 | "Licensor" shall mean the copyright owner or entity authorized by 13 | the copyright owner that is granting the License. 14 | 15 | "Legal Entity" shall mean the union of the acting entity and all 16 | other entities that control, are controlled by, or are under common 17 | control with that entity. For the purposes of this definition, 18 | "control" means (i) the power, direct or indirect, to cause the 19 | direction or management of such entity, whether by contract or 20 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 21 | outstanding shares, or (iii) beneficial ownership of such entity. 22 | 23 | "You" (or "Your") shall mean an individual or Legal Entity 24 | exercising permissions granted by this License. 25 | 26 | "Source" form shall mean the preferred form for making modifications, 27 | including but not limited to software source code, documentation 28 | source, and configuration files. 29 | 30 | "Object" form shall mean any form resulting from mechanical 31 | transformation or translation of a Source form, including but 32 | not limited to compiled object code, generated documentation, 33 | and conversions to other media types. 34 | 35 | "Work" shall mean the work of authorship, whether in Source or 36 | Object form, made available under the License, as indicated by a 37 | copyright notice that is included in or attached to the work 38 | (an example is provided in the Appendix below). 39 | 40 | "Derivative Works" shall mean any work, whether in Source or Object 41 | form, that is based on (or derived from) the Work and for which the 42 | editorial revisions, annotations, elaborations, or other modifications 43 | represent, as a whole, an original work of authorship. For the purposes 44 | of this License, Derivative Works shall not include works that remain 45 | separable from, or merely link (or bind by name) to the interfaces of, 46 | the Work and Derivative Works thereof. 47 | 48 | "Contribution" shall mean any work of authorship, including 49 | the original version of the Work and any modifications or additions 50 | to that Work or Derivative Works thereof, that is intentionally 51 | submitted to Licensor for inclusion in the Work by the copyright owner 52 | or by an individual or Legal Entity authorized to submit on behalf of 53 | the copyright owner. For the purposes of this definition, "submitted" 54 | means any form of electronic, verbal, or written communication sent 55 | to the Licensor or its representatives, including but not limited to 56 | communication on electronic mailing lists, source code control systems, 57 | and issue tracking systems that are managed by, or on behalf of, the 58 | Licensor for the purpose of discussing and improving the Work, but 59 | excluding communication that is conspicuously marked or otherwise 60 | designated in writing by the copyright owner as "Not a Contribution." 61 | 62 | "Contributor" shall mean Licensor and any individual or Legal Entity 63 | on behalf of whom a Contribution has been received by Licensor and 64 | subsequently incorporated within the Work. 65 | 66 | 2. Grant of Copyright License. Subject to the terms and conditions of 67 | this License, each Contributor hereby grants to You a perpetual, 68 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 69 | copyright license to reproduce, prepare Derivative Works of, 70 | publicly display, publicly perform, sublicense, and distribute the 71 | Work and such Derivative Works in Source or Object form. 72 | 73 | 3. Grant of Patent License. Subject to the terms and conditions of 74 | this License, each Contributor hereby grants to You a perpetual, 75 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 76 | (except as stated in this section) patent license to make, have made, 77 | use, offer to sell, sell, import, and otherwise transfer the Work, 78 | where such license applies only to those patent claims licensable 79 | by such Contributor that are necessarily infringed by their 80 | Contribution(s) alone or by combination of their Contribution(s) 81 | with the Work to which such Contribution(s) was submitted. If You 82 | institute patent litigation against any entity (including a 83 | cross-claim or counterclaim in a lawsuit) alleging that the Work 84 | or a Contribution incorporated within the Work constitutes direct 85 | or contributory patent infringement, then any patent licenses 86 | granted to You under this License for that Work shall terminate 87 | as of the date such litigation is filed. 88 | 89 | 4. Redistribution. You may reproduce and distribute copies of the 90 | Work or Derivative Works thereof in any medium, with or without 91 | modifications, and in Source or Object form, provided that You 92 | meet the following conditions: 93 | 94 | (a) You must give any other recipients of the Work or 95 | Derivative Works a copy of this License; and 96 | 97 | (b) You must cause any modified files to carry prominent notices 98 | stating that You changed the files; and 99 | 100 | (c) You must retain, in the Source form of any Derivative Works 101 | that You distribute, all copyright, patent, trademark, and 102 | attribution notices from the Source form of the Work, 103 | excluding those notices that do not pertain to any part of 104 | the Derivative Works; and 105 | 106 | (d) If the Work includes a "NOTICE" text file as part of its 107 | distribution, then any Derivative Works that You distribute must 108 | include a readable copy of the attribution notices contained 109 | within such NOTICE file, excluding those notices that do not 110 | pertain to any part of the Derivative Works, in at least one 111 | of the following places: within a NOTICE text file distributed 112 | as part of the Derivative Works; within the Source form or 113 | documentation, if provided along with the Derivative Works; or, 114 | within a display generated by the Derivative Works, if and 115 | wherever such third-party notices normally appear. The contents 116 | of the NOTICE file are for informational purposes only and 117 | do not modify the License. You may add Your own attribution 118 | notices within Derivative Works that You distribute, alongside 119 | or as an addendum to the NOTICE text from the Work, provided 120 | that such additional attribution notices cannot be construed 121 | as modifying the License. 122 | 123 | You may add Your own copyright statement to Your modifications and 124 | may provide additional or different license terms and conditions 125 | for use, reproduction, or distribution of Your modifications, or 126 | for any such Derivative Works as a whole, provided Your use, 127 | reproduction, and distribution of the Work otherwise complies with 128 | the conditions stated in this License. 129 | 130 | 5. Submission of Contributions. Unless You explicitly state otherwise, 131 | any Contribution intentionally submitted for inclusion in the Work 132 | by You to the Licensor shall be under the terms and conditions of 133 | this License, without any additional terms or conditions. 134 | Notwithstanding the above, nothing herein shall supersede or modify 135 | the terms of any separate license agreement you may have executed 136 | with Licensor regarding such Contributions. 137 | 138 | 6. Trademarks. This License does not grant permission to use the trade 139 | names, trademarks, service marks, or product names of the Licensor, 140 | except as required for reasonable and customary use in describing the 141 | origin of the Work and reproducing the content of the NOTICE file. 142 | 143 | 7. Disclaimer of Warranty. Unless required by applicable law or 144 | agreed to in writing, Licensor provides the Work (and each 145 | Contributor provides its Contributions) on an "AS IS" BASIS, 146 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 147 | implied, including, without limitation, any warranties or conditions 148 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 149 | PARTICULAR PURPOSE. You are solely responsible for determining the 150 | appropriateness of using or redistributing the Work and assume any 151 | risks associated with Your exercise of permissions under this License. 152 | 153 | 8. Limitation of Liability. In no event and under no legal theory, 154 | whether in tort (including negligence), contract, or otherwise, 155 | unless required by applicable law (such as deliberate and grossly 156 | negligent acts) or agreed to in writing, shall any Contributor be 157 | liable to You for damages, including any direct, indirect, special, 158 | incidental, or consequential damages of any character arising as a 159 | result of this License or out of the use or inability to use the 160 | Work (including but not limited to damages for loss of goodwill, 161 | work stoppage, computer failure or malfunction, or any and all 162 | other commercial damages or losses), even if such Contributor 163 | has been advised of the possibility of such damages. 164 | 165 | 9. Accepting Warranty or Additional Liability. While redistributing 166 | the Work or Derivative Works thereof, You may choose to offer, 167 | and charge a fee for, acceptance of support, warranty, indemnity, 168 | or other liability obligations and/or rights consistent with this 169 | License. However, in accepting such obligations, You may act only 170 | on Your own behalf and on Your sole responsibility, not on behalf 171 | of any other Contributor, and only if You agree to indemnify, 172 | defend, and hold each Contributor harmless for any liability 173 | incurred by, or claims asserted against, such Contributor by reason 174 | of your accepting any such warranty or additional liability. 175 | 176 | END OF TERMS AND CONDITIONS 177 | 178 | APPENDIX: How to apply the Apache License to your work. 179 | 180 | To apply the Apache License to your work, attach the following 181 | boilerplate notice, with the fields enclosed by brackets "[]" 182 | replaced with your own identifying information. (Don't include 183 | the brackets!) The text should be enclosed in the appropriate 184 | comment syntax for the file format. We also recommend that a 185 | file or class name and description of purpose be included on the 186 | same "printed page" as the copyright notice for easier 187 | identification within third-party archives. 188 | 189 | Copyright 2022, Replicate, Inc. 190 | 191 | Licensed under the Apache License, Version 2.0 (the "License"); 192 | you may not use this file except in compliance with the License. 193 | You may obtain a copy of the License at 194 | 195 | http://www.apache.org/licenses/LICENSE-2.0 196 | 197 | Unless required by applicable law or agreed to in writing, software 198 | distributed under the License is distributed on an "AS IS" BASIS, 199 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 200 | See the License for the specific language governing permissions and 201 | limitations under the License. 202 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # Cog StableLM 2 | 3 | [![Replicate](https://replicate.com/stability-ai/stablelm-tuned-alpha-7b/badge)](https://replicate.com/stability-ai/stablelm-tuned-alpha-7b) 4 | 5 | This repository is an implementation of [StableLM](https://github.com/Stability-AI/StableLM) as a Cog model. [Cog packages machine learning models as standard containers.](https://github.com/replicate/cog). 6 | 7 | ## Prerequisites 8 | 9 | * Model weights. 10 | 11 | * GPU machine. You'll need a Linux machine with an NVIDIA GPU attached, and the NVIDIA Container Toolkit installed. If you don't already have access to a machine with a GPU, check out our guide to getting a GPU machine. This codebase currently assumes a single device with sufficient VRAM (>24GB) is available. If, instead, you have access to a multi-device environment, you can modify the code to distribute your model across devices. 12 | 13 | * Docker. You'll be using the Cog command-line tool to build and push a model. Cog uses Docker to create containers for models. 14 | 15 | ## Step 0: Install Cog 16 | 17 | First, [install Cog](https://github.com/replicate/cog#install): 18 | 19 | ```bash 20 | curl -o /usr/local/bin/cog -L "https://github.com/replicate/cog/releases/download/v0.7.0-beta17/cog_Linux_x86_64" 21 | sudo chmod +x /usr/local/bin/cog 22 | ``` 23 | 24 | ## Step 1: Run the model 25 | 26 | 27 | You can run the model locally to test it: 28 | 29 | ```bash 30 | cog predict -i prompt="What is a meme, and what's the history behind this word?" 31 | ``` 32 | 33 | ## Step 2: Create a model on Replicate 34 | 35 | Go to [replicate.com/create](https://replicate.com/create) to create a Replicate model. 36 | 37 | Make sure to specify "private" to keep the model private. 38 | 39 | ## Step 3: Configure the model to run on A100 GPUs 40 | 41 | Replicate supports running models on a variety of GPUs. The default GPU type is a T4, but for best performance you'll want to configure your model to run on an A100. 42 | 43 | Click on the "Settings" tab on your model page, scroll down to "GPU hardware", and select "A100". Then click "Save". 44 | 45 | ## Step 4: Push the model to Replicate 46 | 47 | Log in to Replicate: 48 | 49 | ```bash 50 | cog login 51 | ``` 52 | 53 | Push the contents of your current directory to Replicate, using the model name you specified in step 3: 54 | 55 | ```bash 56 | cog push r8.im/username/modelname 57 | ``` 58 | 59 | [Learn more about pushing models to Replicate.](https://replicate.com/docs/guides/push-a-model) 60 | 61 | 62 | ## Step 5: Run the model on Replicate 63 | 64 | Now that you've pushed the model to Replicate, you can run it from the website or with an API. 65 | 66 | To use your model in the browser, go to your model page. 67 | 68 | To use your model with an API, click on the "API" tab on your model page. You'll see commands to run the model with cURL, Python, etc. 69 | 70 | To learn more about how to use Replicate, [check out our documentation](https://replicate.com/docs). 71 | -------------------------------------------------------------------------------- /cog.yaml: -------------------------------------------------------------------------------- 1 | # Configuration for Cog ⚙️ 2 | # Reference: https://github.com/replicate/cog/blob/main/docs/yaml.md 3 | 4 | build: 5 | # set to true if your model requires a GPU 6 | gpu: true 7 | cuda: "11.7" 8 | 9 | # python version in the form '3.8' or '3.8.12' 10 | python_version: "3.8" 11 | 12 | system_packages: 13 | - "apt-transport-https" 14 | - "ca-certificates" 15 | - "gnupg" 16 | - "curl" 17 | 18 | # a list of packages in the format == 19 | python_packages: 20 | - "torch==1.13.1" 21 | - "transformers==4.27.4" 22 | - "accelerate==0.18.0" 23 | - "bitsandbytes==0.38.1" 24 | - "tensorizer==1.0.1" 25 | - "jinja2==3.1.2" 26 | - "huggingface_hub==0.13.4" 27 | 28 | run: 29 | # speed up downloading weights with gcloud storage cp 30 | - "echo 'deb [signed-by=/usr/share/keyrings/cloud.google.gpg] https://packages.cloud.google.com/apt cloud-sdk main' | tee -a /etc/apt/sources.list.d/google-cloud-sdk.list" 31 | - "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key --keyring /usr/share/keyrings/cloud.google.gpg add -" 32 | - "apt-get update && apt-get install google-cloud-cli" 33 | 34 | # predict.py defines how predictions are run on your model 35 | predict: "predict.py:Predictor" 36 | -------------------------------------------------------------------------------- /predict.py: -------------------------------------------------------------------------------- 1 | import subprocess 2 | import time 3 | from collections import OrderedDict 4 | from typing import Optional 5 | 6 | import torch 7 | from cog import BasePredictor, Input, Path, ConcatenateIterator 8 | from tensorizer import TensorDeserializer 9 | from tensorizer.utils import no_init_or_tensor 10 | from transformers import ( 11 | AutoConfig, 12 | AutoTokenizer, 13 | StoppingCriteria, 14 | StoppingCriteriaList, 15 | ) 16 | 17 | from subclass import YieldingCausalLM 18 | 19 | DEFAULT_MODEL = "stabilityai/stablelm-tuned-alpha-7b" 20 | CACHE_DIR = "pretrained_weights" 21 | TOKENIZER_PATH = "./tokenizer" 22 | 23 | TENSORIZER_WEIGHTS_PATH = "gs://replicate-weights/stablelm-tuned-alpha-7b.tensors" 24 | 25 | SYSTEM_PROMPT = """<|SYSTEM|># StableLM Tuned (Alpha version) 26 | - StableLM is a helpful and harmless open-source AI language model developed by StabilityAI. 27 | - StableLM is excited to be able to help the user, but will refuse to do anything that could be considered harmful to the user. 28 | - StableLM is more than just an information source, StableLM is also able to write poetry, short stories, and make jokes. 29 | - StableLM will refuse to participate in anything that could harm a human. 30 | """ 31 | 32 | 33 | class StopOnTokens(StoppingCriteria): 34 | def __call__( 35 | self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs 36 | ) -> bool: 37 | stop_ids = [50278, 50279, 50277, 1, 0] 38 | for stop_id in stop_ids: 39 | if input_ids[0][-1] == stop_id: 40 | return True 41 | return False 42 | 43 | 44 | def maybe_download(path): 45 | if path.startswith("gs://"): 46 | output_path = "/tmp/weights.tensors" 47 | subprocess.check_call(["gcloud", "storage", "cp", path, output_path]) 48 | return output_path 49 | return path 50 | 51 | 52 | class Predictor(BasePredictor): 53 | def setup(self, weights: Optional[Path] = None): 54 | if weights is not None and weights.name == "weights": 55 | weights = None 56 | if weights is None and TENSORIZER_WEIGHTS_PATH: 57 | print("Loading tensorized weights from public path") 58 | self.model = self.load_tensorizer( 59 | weights=maybe_download(TENSORIZER_WEIGHTS_PATH) 60 | ) 61 | 62 | elif (hasattr(weights, "filename") and "tensors" in weights.filename) or str( 63 | weights 64 | ).endswith(".tensors"): 65 | self.model = self.load_tensorizer(weights) 66 | else: 67 | self.model = self.load_huggingface_model(weights=weights) 68 | 69 | self.tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_PATH) 70 | self.stop = StopOnTokens() 71 | 72 | def load_huggingface_model(self, weights=None): 73 | st = time.time() 74 | print(f"loading weights from {weights} w/o tensorizer") 75 | model = YieldingCausalLM.from_pretrained(weights, cache_dir=CACHE_DIR).to( 76 | "cuda:0" 77 | ) 78 | print(f"weights loaded in {time.time() - st}") 79 | return model 80 | 81 | def load_tensorizer(self, weights): 82 | st = time.time() 83 | print(f"deserializing weights from {weights}") 84 | config = AutoConfig.from_pretrained(DEFAULT_MODEL) 85 | 86 | model = no_init_or_tensor( 87 | lambda: YieldingCausalLM.from_pretrained( 88 | None, config=config, state_dict=OrderedDict() 89 | ) 90 | ) 91 | des = TensorDeserializer(weights, plaid_mode=True) 92 | des.load_into_module(model) 93 | print(f"weights loaded in {time.time() - st}") 94 | return model 95 | 96 | def predict( 97 | self, 98 | prompt: str = Input( 99 | description=f"Input Prompt.", default="What's your mood today?" 100 | ), 101 | max_tokens: int = Input( 102 | description="Maximum number of tokens to generate. A word is generally 2-3 tokens", 103 | ge=1, 104 | default=100, 105 | ), 106 | top_p: float = Input( 107 | description="Valid if you choose top_p decoding. When decoding text, samples from the top p percentage of most likely tokens; lower to ignore less likely tokens", 108 | ge=0.01, 109 | le=1.0, 110 | default=1.0, 111 | ), 112 | temperature: float = Input( 113 | description="Adjusts randomness of outputs, greater than 1 is random and 0 is deterministic, 0.75 is a good starting value.", 114 | ge=0.01, 115 | le=5, 116 | default=0.75, 117 | ), 118 | repetition_penalty: float = Input( 119 | description="Penalty for repeated words in generated text; 1 is no penalty, values greater than 1 discourage repetition, less than 1 encourage it.", 120 | ge=0.01, 121 | le=5, 122 | default=1.2, 123 | ), 124 | ) -> ConcatenateIterator[str]: 125 | 126 | prompt_text = f"{SYSTEM_PROMPT}<|USER|>{prompt}<|ASSISTANT|>" 127 | 128 | input_ids = self.tokenizer(prompt_text, return_tensors="pt").input_ids.to( 129 | "cuda:0" 130 | ) 131 | with torch.inference_mode(): 132 | first_token_yielded = False 133 | prev_ids = [] 134 | for output in self.model.generate( 135 | input_ids, 136 | max_new_tokens=max_tokens, 137 | do_sample=True, 138 | num_return_sequences=1, 139 | num_beams=1, 140 | temperature=temperature, 141 | top_p=top_p, 142 | repetition_penalty=repetition_penalty, 143 | stopping_criteria=StoppingCriteriaList([self.stop]), 144 | ): 145 | cur_id = output.item() 146 | # in order to properly handle spaces, we need to do our own tokenizing. Fun! 147 | # we're building up a buffer of sub-word / punctuation tokens until we hit a space, and then yielding whole words + punctuation. 148 | cur_token = self.tokenizer.convert_ids_to_tokens(cur_id) 149 | 150 | # skip initial newline, which this almost always yields. hack - newline id = 187. 151 | if not first_token_yielded and not prev_ids and cur_id == 187: 152 | continue 153 | 154 | # Space is represented as "Ġ". 155 | # Yield previous IDs if we hit a space 156 | # or if the current token includes a space 157 | # at its start (e.g. ' is' -> 'Ġis') 158 | if cur_token.startswith("Ġ"): 159 | if prev_ids: 160 | yield self.tokenizer.decode(prev_ids) 161 | prev_ids = [] 162 | 163 | prev_ids = [cur_id] 164 | continue 165 | 166 | # End token 167 | if cur_token == "<|endoftext|>": 168 | break 169 | 170 | prev_ids.append(cur_id) 171 | 172 | if prev_ids: 173 | yield self.tokenizer.decode(prev_ids) 174 | prev_ids = [] 175 | -------------------------------------------------------------------------------- /scripts/download_weights.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | 3 | import os 4 | import shutil 5 | from distutils.dir_util import copy_tree 6 | from pathlib import Path 7 | from tempfile import TemporaryDirectory 8 | from huggingface_hub import snapshot_download 9 | from transformers import AutoTokenizer 10 | 11 | MODEL_REPO = "stabilityai/stablelm-tuned-alpha-7b" 12 | LOCAL_MODEL_DIR = "./pretrained_weights" 13 | LOCAL_TOKENIZER_DIR = "./tokenizer" 14 | 15 | if os.path.exists(LOCAL_MODEL_DIR): 16 | shutil.rmtree(LOCAL_MODEL_DIR) 17 | 18 | # setup temporary directory 19 | with TemporaryDirectory() as tmpdir: 20 | # download snapshot 21 | snapshot_dir = snapshot_download( 22 | repo_id=MODEL_REPO, 23 | cache_dir=tmpdir, 24 | allow_patterns=["*.bin", "*.json", "*.md"], 25 | ) 26 | # copy snapshot to model dir 27 | copy_tree(snapshot_dir, str(LOCAL_MODEL_DIR)) 28 | 29 | # Move tokenizer to separate location 30 | tokenizer = AutoTokenizer.from_pretrained(LOCAL_MODEL_DIR, padding_side="left") 31 | tokenizer.save_pretrained(LOCAL_TOKENIZER_DIR) 32 | -------------------------------------------------------------------------------- /scripts/tensorize_model.py: -------------------------------------------------------------------------------- 1 | #!/usr/bin/env python 2 | import torch 3 | 4 | from tensorizer import TensorSerializer 5 | from transformers import AutoModelForCausalLM 6 | 7 | model = AutoModelForCausalLM.from_pretrained( 8 | "./pretrained_weights/", torch_dtype=torch.float16 9 | ).to("cuda:0") 10 | 11 | path = "./tensorized_models/stabilityai/stablelm-tuned-alpha-7b.tensors" 12 | serializer = TensorSerializer(path) 13 | serializer.write_module(model) 14 | serializer.close() 15 | -------------------------------------------------------------------------------- /subclass.py: -------------------------------------------------------------------------------- 1 | """sampling code pulled from Transformers & slightly modified to stream tokens""" 2 | import warnings 3 | from typing import List, Optional, Union 4 | 5 | import torch 6 | import torch.distributed as dist 7 | from torch import nn 8 | 9 | from transformers.generation.logits_process import LogitsProcessorList 10 | from transformers.generation.stopping_criteria import ( 11 | StoppingCriteriaList, 12 | validate_stopping_criteria, 13 | ) 14 | from transformers.generation.utils import ( 15 | SampleOutput, 16 | SampleDecoderOnlyOutput, 17 | SampleEncoderDecoderOutput, 18 | ) 19 | 20 | from transformers import GPTNeoXForCausalLM 21 | 22 | 23 | class YieldingCausalLM(GPTNeoXForCausalLM): 24 | """Overriding sample to yield tokens""" 25 | 26 | def sample( 27 | self, 28 | input_ids: torch.LongTensor, 29 | logits_processor: Optional[LogitsProcessorList] = None, 30 | stopping_criteria: Optional[StoppingCriteriaList] = None, 31 | logits_warper: Optional[LogitsProcessorList] = None, 32 | max_length: Optional[int] = None, 33 | pad_token_id: Optional[int] = None, 34 | eos_token_id: Optional[Union[int, List[int]]] = None, 35 | output_attentions: Optional[bool] = None, 36 | output_hidden_states: Optional[bool] = None, 37 | output_scores: Optional[bool] = None, 38 | return_dict_in_generate: Optional[bool] = None, 39 | synced_gpus: Optional[bool] = False, 40 | **model_kwargs, 41 | ) -> Union[SampleOutput, torch.LongTensor]: 42 | # init values 43 | logits_processor = ( 44 | logits_processor if logits_processor is not None else LogitsProcessorList() 45 | ) 46 | stopping_criteria = ( 47 | stopping_criteria 48 | if stopping_criteria is not None 49 | else StoppingCriteriaList() 50 | ) 51 | if max_length is not None: 52 | warnings.warn( 53 | "`max_length` is deprecated in this function, use" 54 | " `stopping_criteria=StoppingCriteriaList(MaxLengthCriteria(max_length=max_length))` instead.", 55 | UserWarning, 56 | ) 57 | stopping_criteria = validate_stopping_criteria( 58 | stopping_criteria, max_length 59 | ) 60 | logits_warper = ( 61 | logits_warper if logits_warper is not None else LogitsProcessorList() 62 | ) 63 | pad_token_id = ( 64 | pad_token_id 65 | if pad_token_id is not None 66 | else self.generation_config.pad_token_id 67 | ) 68 | eos_token_id = ( 69 | eos_token_id 70 | if eos_token_id is not None 71 | else self.generation_config.eos_token_id 72 | ) 73 | if isinstance(eos_token_id, int): 74 | eos_token_id = [eos_token_id] 75 | output_scores = ( 76 | output_scores 77 | if output_scores is not None 78 | else self.generation_config.output_scores 79 | ) 80 | output_attentions = ( 81 | output_attentions 82 | if output_attentions is not None 83 | else self.generation_config.output_attentions 84 | ) 85 | output_hidden_states = ( 86 | output_hidden_states 87 | if output_hidden_states is not None 88 | else self.generation_config.output_hidden_states 89 | ) 90 | return_dict_in_generate = ( 91 | return_dict_in_generate 92 | if return_dict_in_generate is not None 93 | else self.generation_config.return_dict_in_generate 94 | ) 95 | 96 | # init attention / hidden states / scores tuples 97 | scores = () if (return_dict_in_generate and output_scores) else None 98 | decoder_attentions = ( 99 | () if (return_dict_in_generate and output_attentions) else None 100 | ) 101 | cross_attentions = ( 102 | () if (return_dict_in_generate and output_attentions) else None 103 | ) 104 | decoder_hidden_states = ( 105 | () if (return_dict_in_generate and output_hidden_states) else None 106 | ) 107 | 108 | # if model is an encoder-decoder, retrieve encoder attention weights and hidden states 109 | if return_dict_in_generate and self.config.is_encoder_decoder: 110 | encoder_attentions = ( 111 | model_kwargs["encoder_outputs"].get("attentions") 112 | if output_attentions 113 | else None 114 | ) 115 | encoder_hidden_states = ( 116 | model_kwargs["encoder_outputs"].get("hidden_states") 117 | if output_hidden_states 118 | else None 119 | ) 120 | 121 | # keep track of which sequences are already finished 122 | unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) 123 | 124 | this_peer_finished = False # used by synced_gpus only 125 | 126 | # auto-regressive generation 127 | while True: 128 | if synced_gpus: 129 | # Under synced_gpus the `forward` call must continue until all gpus complete their sequence. 130 | # The following logic allows an early break if all peers finished generating their sequence 131 | this_peer_finished_flag = torch.tensor( 132 | 0.0 if this_peer_finished else 1.0 133 | ).to(input_ids.device) 134 | # send 0.0 if we finished, 1.0 otherwise 135 | dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM) 136 | # did all peers finish? the reduced sum will be 0.0 then 137 | if this_peer_finished_flag.item() == 0.0: 138 | break 139 | 140 | # prepare model inputs 141 | model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs) 142 | 143 | # forward pass to get next token 144 | outputs = self( 145 | **model_inputs, 146 | return_dict=True, 147 | output_attentions=output_attentions, 148 | output_hidden_states=output_hidden_states, 149 | ) 150 | 151 | if synced_gpus and this_peer_finished: 152 | continue # don't waste resources running the code we don't need 153 | 154 | next_token_logits = outputs.logits[:, -1, :] 155 | 156 | # pre-process distribution 157 | next_token_scores = logits_processor(input_ids, next_token_logits) 158 | next_token_scores = logits_warper(input_ids, next_token_scores) 159 | 160 | # Store scores, attentions and hidden_states when required 161 | if return_dict_in_generate: 162 | if output_scores: 163 | scores += (next_token_scores,) 164 | if output_attentions: 165 | decoder_attentions += ( 166 | (outputs.decoder_attentions,) 167 | if self.config.is_encoder_decoder 168 | else (outputs.attentions,) 169 | ) 170 | if self.config.is_encoder_decoder: 171 | cross_attentions += (outputs.cross_attentions,) 172 | 173 | if output_hidden_states: 174 | decoder_hidden_states += ( 175 | (outputs.decoder_hidden_states,) 176 | if self.config.is_encoder_decoder 177 | else (outputs.hidden_states,) 178 | ) 179 | 180 | # sample 181 | probs = nn.functional.softmax(next_token_scores, dim=-1) 182 | next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) 183 | 184 | # finished sentences should have their next token be a padding token 185 | if eos_token_id is not None: 186 | if pad_token_id is None: 187 | raise ValueError( 188 | "If `eos_token_id` is defined, make sure that `pad_token_id` is defined." 189 | ) 190 | next_tokens = next_tokens * unfinished_sequences + pad_token_id * ( 191 | 1 - unfinished_sequences 192 | ) 193 | 194 | # update generated ids, model inputs, and length for next step 195 | input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) 196 | model_kwargs = self._update_model_kwargs_for_generation( 197 | outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder 198 | ) 199 | 200 | # if eos_token was found in one sentence, set sentence to finished 201 | if eos_token_id is not None: 202 | unfinished_sequences = unfinished_sequences.mul( 203 | (sum(next_tokens != i for i in eos_token_id)).long() 204 | ) 205 | 206 | # stop when each sentence is finished, or if we exceed the maximum length 207 | if unfinished_sequences.max() == 0 or stopping_criteria(input_ids, scores): 208 | if not synced_gpus: 209 | break 210 | else: 211 | this_peer_finished = True 212 | else: 213 | yield next_tokens 214 | 215 | if return_dict_in_generate: 216 | if self.config.is_encoder_decoder: 217 | yield SampleEncoderDecoderOutput( 218 | sequences=input_ids, 219 | scores=scores, 220 | encoder_attentions=encoder_attentions, 221 | encoder_hidden_states=encoder_hidden_states, 222 | decoder_attentions=decoder_attentions, 223 | cross_attentions=cross_attentions, 224 | decoder_hidden_states=decoder_hidden_states, 225 | ) 226 | else: 227 | yield SampleDecoderOnlyOutput( 228 | sequences=input_ids, 229 | scores=scores, 230 | attentions=decoder_attentions, 231 | hidden_states=decoder_hidden_states, 232 | ) 233 | else: 234 | yield next_tokens 235 | -------------------------------------------------------------------------------- /tokenizer/special_tokens_map.json: -------------------------------------------------------------------------------- 1 | { 2 | "bos_token": "<|endoftext|>", 3 | "eos_token": "<|endoftext|>", 4 | "unk_token": "<|endoftext|>" 5 | } 6 | -------------------------------------------------------------------------------- /tokenizer/tokenizer_config.json: -------------------------------------------------------------------------------- 1 | { 2 | "add_prefix_space": false, 3 | "bos_token": "<|endoftext|>", 4 | "clean_up_tokenization_spaces": true, 5 | "eos_token": "<|endoftext|>", 6 | "model_max_length": 1000000000000000019884624838656, 7 | "padding_side": "left", 8 | "special_tokens_map_file": "./pretrained_weights/special_tokens_map.json", 9 | "tokenizer_class": "GPTNeoXTokenizer", 10 | "unk_token": "<|endoftext|>" 11 | } 12 | --------------------------------------------------------------------------------