(self,
path: str,
hf_cache_dir: Optional[str] = None,
max_seq_len: int = 2048,
tokenizer_path: Optional[str] = None,
tokenizer_kwargs: dict = dict(),
peft_path: Optional[str] = None,
tokenizer_only: bool = False,
model_kwargs: dict = dict(device_map='auto'),
generation_kwargs: dict = dict(),
meta_template: Optional[Dict] = None,
extract_pred_after_decode: bool = False,
batch_padding: bool = False,
pad_token_id: Optional[int] = None,
mode: str = 'none',
use_fastchat_template: bool = False,
end_str: Optional[str] = None)
| 89 | """ |
| 90 | |
| 91 | def __init__(self, |
| 92 | path: str, |
| 93 | hf_cache_dir: Optional[str] = None, |
| 94 | max_seq_len: int = 2048, |
| 95 | tokenizer_path: Optional[str] = None, |
| 96 | tokenizer_kwargs: dict = dict(), |
| 97 | peft_path: Optional[str] = None, |
| 98 | tokenizer_only: bool = False, |
| 99 | model_kwargs: dict = dict(device_map='auto'), |
| 100 | generation_kwargs: dict = dict(), |
| 101 | meta_template: Optional[Dict] = None, |
| 102 | extract_pred_after_decode: bool = False, |
| 103 | batch_padding: bool = False, |
| 104 | pad_token_id: Optional[int] = None, |
| 105 | mode: str = 'none', |
| 106 | use_fastchat_template: bool = False, |
| 107 | end_str: Optional[str] = None): |
| 108 | super().__init__(path=path, |
| 109 | max_seq_len=max_seq_len, |
| 110 | tokenizer_only=tokenizer_only, |
| 111 | meta_template=meta_template) |
| 112 | if hf_cache_dir is None: |
| 113 | hf_cache_dir = os.getenv('HF_MODEL_HUB', None) |
| 114 | self.logger = get_logger() |
| 115 | self.pad_token_id = pad_token_id |
| 116 | assert mode in ['none', 'mid'] |
| 117 | self.mode = mode |
| 118 | self._load_tokenizer(path=path, |
| 119 | tokenizer_path=tokenizer_path, |
| 120 | tokenizer_kwargs=tokenizer_kwargs) |
| 121 | self.batch_padding = batch_padding |
| 122 | self.extract_pred_after_decode = extract_pred_after_decode |
| 123 | if not tokenizer_only: |
| 124 | self._load_model(path=path, |
| 125 | model_kwargs=model_kwargs, |
| 126 | peft_path=peft_path) |
| 127 | self.generation_kwargs = generation_kwargs |
| 128 | self.use_fastchat_template = use_fastchat_template |
| 129 | self.end_str = end_str |
| 130 | |
| 131 | def _load_tokenizer(self, path: str, tokenizer_path: Optional[str], |
| 132 | tokenizer_kwargs: dict): |
no test coverage detected