| | from typing import Dict, List, Optional, Union |
| |
|
| | import torch |
| |
|
| | from opencompass.models.base import BaseModel |
| | from opencompass.models.base_api import APITemplateParser |
| | from opencompass.utils.logging import get_logger |
| | from opencompass.utils.prompt import PromptList |
| |
|
| | PromptType = Union[PromptList, str] |
| |
|
| |
|
| | class Mixtral(BaseModel): |
| | """Mixtral model wrapper https://github.com/open-compass/MixtralKit. |
| | |
| | Args: |
| | path (str): path to the model directory |
| | max_seq_len (int): max sequence length |
| | max_batch_size (int): max batch size |
| | tokenizer_only (bool): whether to load tokenizer only |
| | tokenizer_path (str): path to the tokenizer directory |
| | meta_template (dict): meta template for the model |
| | """ |
| |
|
| | def __init__( |
| | self, |
| | path: str, |
| | max_seq_len: int = 2048, |
| | max_batch_size: int = 8, |
| | tokenizer_only: bool = False, |
| | tokenizer_path: Optional[str] = None, |
| | meta_template: Optional[Dict] = None, |
| | num_gpus: int = 2, |
| | ): |
| | if tokenizer_only: |
| | self._load_tokenizer(tokenizer_path=tokenizer_path) |
| | else: |
| | self._load_model(path=path, |
| | max_seq_len=max_seq_len, |
| | max_batch_size=max_batch_size, |
| | tokenizer_path=tokenizer_path, |
| | num_gpus=num_gpus) |
| | self.max_seq_len = max_seq_len |
| | self.template_parser = APITemplateParser(meta_template) |
| | self.logger = get_logger() |
| |
|
| | def _load_model(self, |
| | path: str, |
| | max_seq_len: int, |
| | max_batch_size: int, |
| | tokenizer_path: Optional[str] = None, |
| | num_gpus: int = 2): |
| | from mixtralkit.mixtral import Mixtral |
| | self.generator = Mixtral.build(ckpt_dir=path, |
| | tokenizer_path=tokenizer_path, |
| | max_seq_len=max_seq_len, |
| | max_batch_size=max_batch_size, |
| | num_gpus=num_gpus) |
| | self.tokenizer = self.generator.tokenizer |
| | self.model = self.generator.model |
| |
|
| | def _load_tokenizer(self, tokenizer_path: str): |
| | from mixtralkit.layers import Tokenizer |
| | self.tokenizer = Tokenizer(tokenizer_path) |
| |
|
| | def generate(self, inputs: List[str], max_out_len: int) -> List[str]: |
| | prompt_tokens = [] |
| | for input in inputs: |
| | tokens = self.tokenizer.encode(input, True, False) |
| | num_token = min(self.model.params.max_seq_len, len(tokens)) |
| | prompt_tokens.append(tokens[-num_token:]) |
| | generation_tokens, _ = self.generator.generate( |
| | prompt_tokens=prompt_tokens, |
| | max_gen_len=max_out_len, |
| | temperature=0, |
| | ) |
| | results = [self.tokenizer.decode(t) for t in generation_tokens] |
| | return results |
| |
|
| | def get_ppl(self, |
| | inputs: List[str], |
| | mask_length: Optional[List[int]] = None) -> List[float]: |
| | assert mask_length is None, 'mask_length is not supported' |
| | bsz = len(inputs) |
| | params = self.model.params |
| | assert bsz <= params.max_batch_size, (bsz, params.max_batch_size) |
| | |
| | prompt_tokens = [self.tokenizer.encode(x, True, False) for x in inputs] |
| | max_prompt_size = max([len(t) for t in prompt_tokens]) |
| | total_len = min(params.max_seq_len, max_prompt_size) |
| | tokens = torch.zeros((bsz, total_len)).cuda().long() |
| | for k, t in enumerate(prompt_tokens): |
| | num_token = min(total_len, len(t)) |
| | tokens[k, :num_token] = torch.tensor(t[-num_token:]).long() |
| | |
| | outputs = self.model.forward(tokens, 0) |
| | |
| | shift_logits = outputs[..., :-1, :].contiguous().float() |
| | shift_labels = tokens[..., 1:].contiguous() |
| | shift_logits = shift_logits.view(-1, shift_logits.size(-1)) |
| | shift_labels = shift_labels.view(-1) |
| | loss_fct = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=0) |
| | loss = loss_fct(shift_logits, shift_labels).view(bsz, -1) |
| | lens = (tokens != 0).sum(-1).cpu().numpy() |
| | ce_loss = loss.sum(-1).cpu().detach().numpy() / lens |
| | return ce_loss |
| |
|
| | def get_token_len(self, prompt: str) -> int: |
| | return len(self.tokenizer.encode(prompt, True, True)) |
| |
|
