| from typing import Dict, Iterable, List, Optional, Union |
|
|
| import numpy as np |
| import torch.distributed as dist |
|
|
| 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 LLaMA2AccessoryModel(BaseModel): |
| """LLaMA2-Accessory model wrapper. |
| |
| Project: https://github.com/Alpha-VLLM/LLaMA2-Accessory |
| |
| Args: |
| tokenizer_only (bool): whether to load tokenizer only |
| meta_template (dict): meta template for the model |
| additional_stop_symbols: (Iterable[str]): additional symbols that mark |
| the end of generation, e.g. the "###" symbol for separating turns |
| in the chat template. |
| from_pretrained_kwargs: kwargs that will be passed to |
| `accessory.MetaModel.from_pretrained` for model instantiation. |
| """ |
|
|
| def __init__(self, |
| tokenizer_only: bool = False, |
| meta_template: Optional[Dict] = None, |
| additional_stop_symbols: Iterable[str] = (), |
| **from_pretrained_kwargs): |
| if tokenizer_only: |
| self._load_tokenizer(from_pretrained_kwargs) |
| else: |
| self._load_model(from_pretrained_kwargs) |
|
|
| self.additional_stop_symbols = additional_stop_symbols |
| self.max_seq_len = from_pretrained_kwargs.get('max_seq_len', 4096) |
| self.template_parser = APITemplateParser(meta_template) |
| self.logger = get_logger() |
|
|
| def _load_model(self, from_pretrained_kwargs): |
| from accessory.model.meta import MetaModel |
| from accessory.util.misc import init_distributed_mode |
| if not dist.is_initialized(): |
| init_distributed_mode() |
|
|
| model_parallel_group = dist.GroupMember.WORLD |
| from_pretrained_kwargs['mp_group'] = model_parallel_group |
|
|
| self.model = MetaModel.from_pretrained(**from_pretrained_kwargs) |
| self.tokenizer = self.model.tokenizer |
| self.logger = get_logger() |
|
|
| def _load_tokenizer(self, from_pretrained_kwargs): |
| from accessory.model.tokenizer import ( |
| Tokenizer, probe_tokenizer_path_from_pretrained) |
| if 'tokenizer_path' in from_pretrained_kwargs: |
| tokenizer_path = from_pretrained_kwargs['tokenizer_path'] |
| else: |
| pretrained_path = from_pretrained_kwargs['pretrained_path'] |
| if isinstance(pretrained_path, str): |
| pretrained_path = [pretrained_path] |
| tokenizer_path = probe_tokenizer_path_from_pretrained( |
| pretrained_path[-1]) |
|
|
| self.tokenizer = Tokenizer(tokenizer_path) |
|
|
| def generate(self, inputs: List[str], max_out_len: int) -> List[str]: |
| results = self.model.generate( |
| prompts=inputs, |
| max_gen_len=max_out_len, |
| temperature=0., |
| additional_stop_symbols=self.additional_stop_symbols) |
| return results |
|
|
| def get_ppl(self, |
| inputs: List[str], |
| mask_length: Optional[List[int]] = None): |
| assert mask_length is None, 'mask_length is not supported' |
| evaluation_results = self.model.evaluate_examples(examples=inputs) |
| ppl = evaluation_results['ppl'] |
| return np.array(ppl, dtype=np.float32) |
|
|
| def get_token_len(self, prompt: str) -> int: |
| return len(self.tokenizer.encode(prompt, True, True)) |
|
|
