ramixpe/1.8testing

1# Copyright (c) Alibaba Cloud.2#3# This source code is licensed under the license found in the4# LICENSE file in the root directory of this source tree.5 6"""Generation support."""7 8from typing import Tuple, List, Union, Iterable9 10import numpy as np11import torch12import torch.nn.functional as F13from transformers import PreTrainedTokenizer14from transformers import logging15from transformers.generation import LogitsProcessor16 17logger = logging.get_logger(__name__)18 19# Types.20HistoryType = List[Tuple[str, str]]21TokensType = List[int]22BatchTokensType = List[List[int]]23 24 25def pad_batch(batch: BatchTokensType, pad_id: int, seq_length: int) -> BatchTokensType:26    for tokens in batch:27        context_length = len(tokens)28        if context_length < seq_length:29            tokens.extend([pad_id] * (seq_length - context_length))30    return batch31 32 33def get_ltor_masks_and_position_ids(34    data,35    eod_token,36    reset_position_ids,37    reset_attention_mask,38    eod_mask_loss,39):40    """Build masks and position id for left to right model."""41 42    # Extract batch size and sequence length.43    micro_batch_size, seq_length = data.size()44 45    # Attention mask (lower triangular).46    if reset_attention_mask:47        att_mask_batch = micro_batch_size48    else:49        att_mask_batch = 150    attention_mask = torch.tril(51        torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)52    ).view(att_mask_batch, 1, seq_length, seq_length)53 54    # Loss mask.55    loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)56    if eod_mask_loss:57        loss_mask[data == eod_token] = 0.058 59    # Position ids.60    position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)61    position_ids = position_ids.unsqueeze(0).expand_as(data)62    # We need to clone as the ids will be modifed based on batch index.63    if reset_position_ids:64        position_ids = position_ids.clone()65 66    if reset_position_ids or reset_attention_mask:67        # Loop through the batches:68        for b in range(micro_batch_size):69 70            # Find indecies where EOD token is.71            eod_index = position_ids[b, data[b] == eod_token]72            # Detach indecies from positions if going to modify positions.73            if reset_position_ids:74                eod_index = eod_index.clone()75 76            # Loop through EOD indecies:77            prev_index = 078            for j in range(eod_index.size()[0]):79                i = eod_index[j]80                # Mask attention loss.81                if reset_attention_mask:82                    attention_mask[b, 0, (i + 1) :, : (i + 1)] = 083                # Reset positions.84                if reset_position_ids:85                    position_ids[b, (i + 1) :] -= i + 1 - prev_index86                    prev_index = i + 187 88    # Convert attention mask to binary:89    attention_mask = attention_mask < 0.590 91    return attention_mask, loss_mask, position_ids92 93 94def get_batch(context_tokens: torch.LongTensor, eod_id: int):95    """Generate batch from context tokens."""96    # Move to GPU.97    tokens = context_tokens.contiguous().to(context_tokens.device)98    # Get the attention mask and postition ids.99    attention_mask, _, position_ids = get_ltor_masks_and_position_ids(100        tokens,101        eod_id,102        reset_position_ids=False,103        reset_attention_mask=False,104        eod_mask_loss=False,105    )106    return tokens, attention_mask, position_ids107 108 109def get_stop_words_ids(chat_format, tokenizer):110    if chat_format == "raw":111        stop_words_ids = [tokenizer.encode("Human:"), [tokenizer.eod_id]]112    elif chat_format == "chatml":113        stop_words_ids = [[tokenizer.im_end_id], [tokenizer.im_start_id]]114    else:115        raise NotImplementedError(f"Unknown chat format {chat_format!r}")116    return stop_words_ids117 118 119def make_context(120    tokenizer: PreTrainedTokenizer,121    query: str,122    history: List[Tuple[str, str]] = None,123    system: str = "",124    max_window_size: int = 6144,125    chat_format: str = "chatml",126):127    if history is None:128        history = []129 130    if chat_format == "chatml":131        im_start, im_end = "<|im_start|>", "<|im_end|>"132        im_start_tokens = [tokenizer.im_start_id]133        im_end_tokens = [tokenizer.im_end_id]134        nl_tokens = tokenizer.encode("\n")135 136        def _tokenize_str(role, content):137            return f"{role}\n{content}", tokenizer.encode(138                role, allowed_special=set()139            ) + nl_tokens + tokenizer.encode(content, allowed_special=set())140 141        system_text, system_tokens_part = _tokenize_str("system", system)142        system_tokens = im_start_tokens + system_tokens_part + im_end_tokens143 144        raw_text = ""145        context_tokens = []146 147        for turn_query, turn_response in reversed(history):148            query_text, query_tokens_part = _tokenize_str("user", turn_query)149            query_tokens = im_start_tokens + query_tokens_part + im_end_tokens150            response_text, response_tokens_part = _tokenize_str(151                "assistant", turn_response152            )153            response_tokens = im_start_tokens + response_tokens_part + im_end_tokens154 155            next_context_tokens = nl_tokens + query_tokens + nl_tokens + response_tokens156            prev_chat = (157                f"\n{im_start}{query_text}{im_end}\n{im_start}{response_text}{im_end}"158            )159 160            current_context_size = (161                len(system_tokens) + len(next_context_tokens) + len(context_tokens)162            )163            if current_context_size < max_window_size:164                context_tokens = next_context_tokens + context_tokens165                raw_text = prev_chat + raw_text166            else:167                break168 169        context_tokens = system_tokens + context_tokens170        raw_text = f"{im_start}{system_text}{im_end}" + raw_text171        context_tokens += (172            nl_tokens173            + im_start_tokens174            + _tokenize_str("user", query)[1]175            + im_end_tokens176            + nl_tokens177            + im_start_tokens178            + tokenizer.encode("assistant")179            + nl_tokens180        )181        raw_text += f"\n{im_start}user\n{query}{im_end}\n{im_start}assistant\n"182 183    elif chat_format == "raw":184        raw_text = query185        context_tokens = tokenizer.encode(raw_text)186    else:187        raise NotImplementedError(f"Unknown chat format {chat_format!r}")188 189    return raw_text, context_tokens190 191 192def _decode_default(193    tokens: List[int],194    *,195    stop_words: List[str],196    eod_words: List[str],197    tokenizer: PreTrainedTokenizer,198    raw_text_len: int,199    verbose: bool = False,200    return_end_reason: bool = False,201    errors: str='replace',202):203    trim_decode_tokens = tokenizer.decode(tokens, errors=errors)[raw_text_len:]204    if verbose:205        print("\nRaw Generate: ", trim_decode_tokens)206 207    end_reason = f"Gen length {len(tokens)}"208    for stop_word in stop_words:209        trim_decode_tokens = trim_decode_tokens.replace(stop_word, "").strip()210    for eod_word in eod_words:211        if eod_word in trim_decode_tokens:212            end_reason = f"Gen {eod_word!r}"213        trim_decode_tokens = trim_decode_tokens.split(eod_word)[0]214    trim_decode_tokens = trim_decode_tokens.strip()215    if verbose:216        print("\nEnd Reason:", end_reason)217        print("\nGenerate: ", trim_decode_tokens)218 219    if return_end_reason:220        return trim_decode_tokens, end_reason221    else:222        return trim_decode_tokens223 224 225def _decode_chatml(226    tokens: List[int],227    *,228    stop_words: List[str],229    eod_token_ids: List[int],230    tokenizer: PreTrainedTokenizer,231    raw_text_len: int,232    context_length: int,233    verbose: bool = False,234    return_end_reason: bool = False,235    errors: str='replace'236):237    end_reason = f"Gen length {len(tokens)}"238    eod_token_idx = context_length239    for eod_token_idx in range(context_length, len(tokens)):240        if tokens[eod_token_idx] in eod_token_ids:241            end_reason = f"Gen {tokenizer.decode([tokens[eod_token_idx]])!r}"242            break243 244    trim_decode_tokens = tokenizer.decode(tokens[:eod_token_idx], errors=errors)[raw_text_len:]245    if verbose:246        print("\nRaw Generate w/o EOD:", tokenizer.decode(tokens, errors=errors)[raw_text_len:])247        print("\nRaw Generate:", trim_decode_tokens)248        print("\nEnd Reason:", end_reason)249    for stop_word in stop_words:250        trim_decode_tokens = trim_decode_tokens.replace(stop_word, "").strip()251    trim_decode_tokens = trim_decode_tokens.strip()252    if verbose:253        print("\nGenerate:", trim_decode_tokens)254 255    if return_end_reason:256        return trim_decode_tokens, end_reason257    else:258        return trim_decode_tokens259 260 261def decode_tokens(262    tokens: Union[torch.LongTensor, TokensType],263    tokenizer: PreTrainedTokenizer,264    raw_text_len: int,265    context_length: int,266    chat_format: str,267    verbose: bool = False,268    return_end_reason: bool = False,269    errors: str="replace",270) -> str:271    if torch.is_tensor(tokens):272        tokens = tokens.cpu().numpy().tolist()273 274    if chat_format == "chatml":275        return _decode_chatml(276            tokens,277            stop_words=[],278            eod_token_ids=[tokenizer.im_start_id, tokenizer.im_end_id],279            tokenizer=tokenizer,280            raw_text_len=raw_text_len,281            context_length=context_length,282            verbose=verbose,283            return_end_reason=return_end_reason,284            errors=errors,285        )286    elif chat_format == "raw":287        return _decode_default(288            tokens,289            stop_words=["<|endoftext|>"],290            eod_words=["<|endoftext|>"],291            tokenizer=tokenizer,292            raw_text_len=raw_text_len,293            verbose=verbose,294            return_end_reason=return_end_reason,295            errors=errors,296        )297    else:298        raise NotImplementedError(f"Unknown chat format {chat_format!r}")299 300 301class StopWordsLogitsProcessor(LogitsProcessor):302    """303    :class:`transformers.LogitsProcessor` that enforces that when specified sequences appear, stop geration.304 305    Args:306        stop_words_ids (:obj:`List[List[int]]`):307            List of list of token ids of stop ids. In order to get the tokens of the words308            that should not appear in the generated text, use :obj:`tokenizer(bad_word,309            add_prefix_space=True).input_ids`.310        eos_token_id (:obj:`int`):311            The id of the `end-of-sequence` token.312    """313 314    def __init__(self, stop_words_ids: Iterable[Iterable[int]], eos_token_id: int):315 316        if not isinstance(stop_words_ids, List) or len(stop_words_ids) == 0:317            raise ValueError(318                f"`stop_words_ids` has to be a non-emtpy list, but is {stop_words_ids}."319            )320        if any(not isinstance(bad_word_ids, list) for bad_word_ids in stop_words_ids):321            raise ValueError(322                f"`stop_words_ids` has to be a list of lists, but is {stop_words_ids}."323            )324        if any(325            any(326                (not isinstance(token_id, (int, np.integer)) or token_id < 0)327                for token_id in stop_word_ids328            )329            for stop_word_ids in stop_words_ids330        ):331            raise ValueError(332                f"Each list in `stop_words_ids` has to be a list of positive integers, but is {stop_words_ids}."333            )334 335        self.stop_words_ids = list(336            filter(337                lambda bad_token_seq: bad_token_seq != [eos_token_id], stop_words_ids338            )339        )340        self.eos_token_id = eos_token_id341        for stop_token_seq in self.stop_words_ids:342            assert (343                len(stop_token_seq) > 0344            ), "Stop words token sequences {} cannot have an empty list".format(345                stop_words_ids346            )347 348    def __call__(349        self, input_ids: torch.LongTensor, scores: torch.FloatTensor350    ) -> torch.FloatTensor:351        stopped_samples = self._calc_stopped_samples(input_ids)352        for i, should_stop in enumerate(stopped_samples):353            if should_stop:354                scores[i, self.eos_token_id] = float(2**15)355        return scores356 357    def _tokens_match(self, prev_tokens: torch.LongTensor, tokens: List[int]) -> bool:358        if len(tokens) == 0:359            # if bad word tokens is just one token always ban it360            return True361        elif len(tokens) > len(prev_tokens):362            # if bad word tokens are longer then prev input_ids they can't be equal363            return False364        elif prev_tokens[-len(tokens) :].tolist() == tokens:365            # if tokens match366            return True367        else:368            return False369 370    def _calc_stopped_samples(self, prev_input_ids: Iterable[int]) -> Iterable[int]:371        stopped_samples = []372        for prev_input_ids_slice in prev_input_ids:373            match = False374            for stop_token_seq in self.stop_words_ids:375                if self._tokens_match(prev_input_ids_slice, stop_token_seq):376                    # if tokens do not match continue377                    match = True378                    break379            stopped_samples.append(match)380 381        return stopped_samples382 383 384def top_k_logits(logits, top_k=0, top_p=0.0, filter_value=-float("Inf")):385    """This function has been mostly taken from huggingface conversational386    ai code at387        https://medium.com/huggingface/how-to-build-a-state-of-the-art-388             conversational-ai-with-transfer-learning-2d818ac26313"""389 390    if top_k > 0:391        # Remove all tokens with a probability less than the392        # last token of the top-k393        indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]394        logits[indices_to_remove] = filter_value395 396    if top_p > 0.0:397        # Cconvert to 1D398        sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)399        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)400 401        # Remove tokens with cumulative probability above the threshold402        sorted_indices_to_remove = cumulative_probs > top_p403        # Shift the indices to the right to keep also the first token404        # above the threshold405        sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()406        sorted_indices_to_remove[..., 0] = 0407        for i in range(sorted_indices.size(0)):408            indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]]409            logits[i][indices_to_remove] = filter_value410 411    return logits412 413 414def switch(val1, val2, boolean):415    boolean = boolean.type_as(val1)416    return (1 - boolean) * val1 + boolean * val2417