opencompass-my-api/opencompass/models/glm.py

import re
from functools import partial
from typing import Dict, List, Optional, Union

import numpy as np
import torch

from opencompass.models.base import BaseModel, LMTemplateParser
from opencompass.registry import MODELS
from opencompass.utils.prompt import PromptList

PromptType = Union[PromptList, str]


@MODELS.register_module(name=['GLM-130B'])
class GLM130B(BaseModel):

    def __init__(self,
                 pkg_root: str,
                 ckpt_path: str,
                 tokenizer_only: bool = False,
                 meta_template: Optional[Dict] = None,
                 **kwargs):
        assert not tokenizer_only, 'LLama does not support tokenizer only mode'
        self.pkg_root = pkg_root
        self.ckpt_path = ckpt_path
        self._load_model(**kwargs)

        self.template_parser = LMTemplateParser(meta_template)
        self.eos_token_id = None
        if meta_template and 'eos_token_id' in meta_template:
            self.eos_token_id = meta_template['eos_token_id']

    def _load_model(self, **kwargs):
        import sys
        sys.path.insert(0, self.pkg_root)
        from argparse import Namespace

        from evaluation.model import ModelForEvaluation, batch_filling_sequence
        from generate import get_masks_and_position_ids
        from generation import BaseStrategy, BeamSearchStrategy
        from initialize import initialize_model_and_tokenizer
        from SwissArmyTransformer import get_args

        self.get_masks_and_position_ids = get_masks_and_position_ids
        self.batch_filling_sequence = batch_filling_sequence

        kwargs = {
            'bminf': False,
            'bminf_memory_limit': 20,
            'quantization_bit_width': None,
            'from_quantized_checkpoint': False,
            'sequential_initialization': False,
            'sampling_strategy': 'BaseStrategy',
            'min_gen_length': 0,
            'print_all_beams': False,
            **kwargs,
        }

        args_list = [
            ['--seed', '1234'],
            ['--mode', 'inference'],
            ['--out-seq-length', '256'],
            ['--num-beams', '4'],
            ['--length-penalty', '1.0'],
            ['--no-repeat-ngram-size', '3'],
            ['--temperature', '1.0'],
            ['--top_k', '0'],
            ['--top_p', '0'],
            ['--output-path', 'samples'],
            ['--model-parallel-size', '8'],
            ['--num-layers', '70'],
            ['--hidden-size', '12288'],
            ['--inner-hidden-size', '32768'],
            ['--vocab-size', '150528'],
            ['--num-attention-heads', '96'],
            ['--max-sequence-length', '2048'],
            ['--tokenizer-type', 'icetk-glm-130B'],
            ['--layernorm-order', 'post'],
            ['--load', self.ckpt_path],
            ['--skip-init'],
            ['--fp16'],
            ['--input-source', 'interactive'],
        ]  # Come from the default initialize arguments of official repo
        args = get_args(sum(args_list, []))
        args = Namespace(**vars(args), **kwargs)
        args.do_train = False
        self.args = args

        model, tokenizer = initialize_model_and_tokenizer(args)
        self.model = model
        self.model_for_eval = ModelForEvaluation(model)
        self.tokenizer = tokenizer
        self.device = args.device

        end_tokens = [
            tokenizer.get_command('eop'),
            tokenizer.get_command('eos')
        ]
        if args.sampling_strategy == 'BaseStrategy':
            self.strategy = BaseStrategy(batch_size=1,
                                         temperature=args.temperature,
                                         top_k=args.top_k,
                                         top_p=args.top_p,
                                         end_tokens=end_tokens)
        elif args.sampling_strategy == 'BeamSearchStrategy':
            self.strategy = BeamSearchStrategy(
                1,
                args.num_beams,
                length_penalty=args.length_penalty,
                consider_end=True,
                end_tokens=end_tokens,
                no_repeat_ngram_size=args.no_repeat_ngram_size,
                min_gen_length=args.min_gen_length,
            )
        else:
            raise ValueError(f'unknown strategy {args.sampling_strategy}')

        sys.path.pop(0)

    def get_token_len(self, prompt: str) -> int:
        """Get lengths of the tokenized strings.

        Args:
            prompt (str): Input string.

        Returns:
            int: Length of the input tokens
        """
        return len(self.tokenizer.tokenize(prompt))

    def choice(self, inputs, choices):
        import sys
        sys.path.insert(0, self.pkg_root)
        from unittest.mock import MagicMock

        from evaluation.dataset import MultiChoiceTaskDataset
        sys.path.pop(0)

        choice_tokens = [self.tokenizer.tokenize(item) for item in choices]
        is_single_token = all(len(token) == 1 for token in choice_tokens)

        data_items = []
        mock_dataset = MagicMock(is_single_token=is_single_token)
        from mmengine.dist import is_main_process
        for text in inputs:
            if is_main_process():
                print(f"\033[92m'text'\033[0m: {text}")
            data_item = MultiChoiceTaskDataset.build_multiple_choice_sample(
                text=self.tokenizer.tokenize(text),
                #  text=self.tokenizer.tokenize(text) + [20019],
                choices=[self.tokenizer.tokenize(item) for item in choices],
                is_single_token=is_single_token,
            )
            data_items.append(data_item)
        batch = MultiChoiceTaskDataset.collate_fn(mock_dataset, data_items)

        log_probs = self.model_for_eval.cond_log_prob(batch)

        answers = []
        for log_prob in zip(log_probs):
            answers.append(choices[np.argmax(log_prob).item()])

        return answers

    def generate(self, inputs: List[str], max_out_len: int) -> List[str]:
        """Generate results given a list of inputs.

        Args:
            inputs (List[str]): A list of strings.
            max_out_len (int): The maximum length of the output.

        Returns:
            List[str]: A list of generated strings.
        """
        if isinstance(inputs, list):
            return sum((self.generate(raw_text, max_out_len)
                        for raw_text in inputs), [])
        else:
            raw_text = inputs

        from mmengine.dist import is_main_process
        if is_main_process():
            print(f"\033[92m'raw_text'\033[0m: \n{raw_text}")

        # add MASK
        generation_mask = '[gMASK]'
        if '[MASK]' in raw_text:
            generation_mask = '[MASK]'
        elif '[sMASK]' in raw_text:
            generation_mask = '[sMASK]'
        use_gmask = '[MASK]' not in raw_text and '[sMASK]' not in raw_text

        mask_pattern = r'\[[sg]?MASK\]'
        text_list = re.split(mask_pattern, raw_text)
        pattern_list = re.compile(mask_pattern).findall(raw_text)
        seq = []
        for i in range(len(pattern_list)):
            pattern = pattern_list[i]
            sub_text = text_list[i]
            seq.extend(self.tokenizer.tokenize(sub_text))
            seq.append(self.tokenizer.get_command(pattern))

        seq.extend(self.tokenizer.tokenize(text_list[-1]))
        prompt_token_length = len(seq)

        if 'MASK]' not in raw_text:
            seq += [self.tokenizer.get_command(generation_mask)]
            raw_text += ' ' + generation_mask
        if not raw_text.endswith('MASK]'):
            seq = seq + [self.tokenizer.get_command('eos')]
        if len(seq) > self.args.max_sequence_length:
            raise ValueError('text too long.')

        # generation
        output_list = [seq]
        if self.args.sampling_strategy == 'BeamSearchStrategy':
            num_output = self.args.num_beams
        else:
            num_output = 1
        last_pos = [0] * num_output

        # continually detect the first mark position
        while True:
            seq = output_list[0]
            # detect mask position
            mask_token = self.tokenizer.get_command(generation_mask)
            if mask_token not in seq:
                break
            mask_position = seq.index(mask_token)

            output_list = []

            input_seq = torch.cuda.LongTensor(
                [seq + [self.tokenizer.get_command('sop')]],
                device=self.device,
            )
            output, _ = self.batch_filling_sequence(
                self.model,
                input_seq,
                torch.cuda.LongTensor([input_seq.shape[-1]],
                                      device=self.device),
                strategy=self.strategy,
                get_masks_and_position_ids=partial(
                    self.get_masks_and_position_ids,
                    mask_position=mask_position,
                    max_gen_length=max_out_len,
                    gmask=use_gmask,
                ),
            )
            if isinstance(output, torch.Tensor):  # different strategies
                output = output.tolist()
            output = output[0]  # batch_size = 1
            output_list.extend(output)

            # clip -1s and fill back generated things into seq
            for i in range(len(output_list)):
                output = output_list[i].tolist() if isinstance(
                    output_list[i], torch.Tensor) else output_list[i]
                try:
                    unfinished = output.index(-1)
                except ValueError:
                    unfinished = len(output)
                if output[unfinished - 1] in self.strategy.end_tokens:
                    unfinished -= 1
                bog = output.index(self.tokenizer.get_command('sop'))

                last_pos[i] = mask_position + unfinished - (bog + 1)
                output_list[i] = output[:mask_position] + output[
                    bog + 1:unfinished] + output[mask_position + 1:bog]

        # Select the best answer
        output = output_list[0]
        if output[-1] == self.tokenizer.get_command('eos'):
            output = output[:-1]

        # Avoid generate out-of-range id, replace to unk
        output = np.array(output)
        output[output < 20000] = 20000
        output = output.tolist()
        answer = self.tokenizer.detokenize(output[prompt_token_length:])
        if is_main_process():
            print(f"\033[92m'answer'\033[0m: \n{answer}")

        return [answer]

    def get_logits(self, inputs: List[str]):
        mask_id = self.tokenizer.get_command('[MASK]')
        sop_id = self.tokenizer.get_command('sop')

        tokens = []
        targets = []
        position_ids = []
        attn_masks = []
        from mmengine.dist import is_main_process
        for raw_text in inputs:
            mask_pattern = r'\[MASK\]'
            text_list = re.split(mask_pattern, raw_text, 1)

            token = sum([
                self.tokenizer.tokenize(text_list[0]),
                [mask_id, sop_id],
                self.tokenizer.tokenize(text_list[1]),
            ], [])[:-1]
            target = sum([
                self.tokenizer.tokenize(text_list[0]),
                [mask_id],
                self.tokenizer.tokenize(text_list[1]),
            ], [])
            if is_main_process():
                print(f"\033[92m'raw_text'\033[0m: {raw_text}")
                print(f"\033[92m'token'\033[0m: {token}")

            seq_length = len(token)

            attn_mask = np.ones((seq_length, seq_length), dtype=np.int64)

            tokens.append(np.array(token, dtype=np.int64))
            targets.append(np.array(target, dtype=np.int64))
            position_ids.append(np.arange(0, seq_length, dtype=np.int64))
            attn_masks.append(attn_mask)

        TILE = 32
        length_to_pad = (max(map(len, tokens)) + TILE - 1) // TILE * TILE
        token_batch, target_batch, position_id_batch, attention_mask_batch = [], [], [], []  # noqa: E501
        for token, target, position_id, attn_mask in zip(
                tokens, targets, position_ids, attn_masks):
            attn_mask = np.pad(
                attn_mask,
                pad_width=((0, length_to_pad - len(token)), ),
                mode='constant',
                constant_values=0,
            )
            token = np.concatenate(
                (token, np.zeros(length_to_pad - len(token), dtype=np.int64)))
            target = np.concatenate((target,
                                     np.full(length_to_pad - len(target),
                                             -1,
                                             dtype=np.int64)))
            position_id = np.concatenate(
                (position_id,
                 np.zeros(length_to_pad - len(position_id), dtype=np.int64)))

            token_batch.append(token)
            target_batch.append(target)
            position_id_batch.append(position_id)
            attention_mask_batch.append(attn_mask)

        token_batch = torch.tensor(np.array(token_batch),
                                   dtype=torch.int64).to(self.device)
        target_batch = torch.tensor(np.array(target_batch),
                                    dtype=torch.int64).to(self.device)
        position_id_batch = torch.tensor(np.array(position_id_batch),
                                         dtype=torch.int64).to(self.device)
        attention_mask_batch = (torch.tensor(
            np.array(attention_mask_batch), dtype=torch.int64) < 0.5).to(
                self.device).bool().unsqueeze(1)

        logits, *out_per_layers = self.model(token_batch,
                                             position_id_batch,
                                             attention_mask_batch,
                                             log_attention_weights=None)
        if is_main_process():
            print(f"\033[92m'target_batch'\033[0m: {target_batch}")

        return logits, target_batch

    def get_ppl(self,
                inputs: List[str],
                mask_length: Optional[List[int]] = None) -> List[float]:
        """Get perplexity scores given a list of inputs.

        Args:
            inputs (List[str]): A list of strings.
            mask_length (Optional[List[int]]): A list of mask lengths. If
                provided, the perplexity scores will be calculated with the
                first mask_length[i] tokens masked out. It's okay to skip
                its implementation if advanced features in PPLInfernecer is
                not needed.

        Returns:
            List[float]: A list of perplexity scores.
        """
        logits, targets = self.get_logits(inputs)

        loss_fn = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=-1)
        loss = loss_fn(logits.view(-1, logits.size(-1)),
                       targets.view(-1)).view(targets.size())
        from mmengine.dist import is_main_process
        if is_main_process():
            print(f"\033[92m'loss'\033[0m: {loss}")

        if mask_length is not None:
            mask = torch.zeros_like(targets)  # [batch,seqlen]
            for i in range(len(mask)):
                for j in range(mask_length[i] - 1, len(mask[i])):
                    mask[i][j] = 1
            loss = loss * mask

        lens = (targets != -1).sum(-1).cpu().numpy()
        if mask_length is not None:
            lens -= np.array(mask_length)
        ce_loss = loss.sum(-1).cpu().detach().numpy() / lens
        if is_main_process():
            print(f"\033[92m'lens'\033[0m: {lens}")
            print(f"\033[92m'ce_loss'\033[0m: {ce_loss}")
        return ce_loss