| 45 | self.logger = get_logger() |
| 46 | |
| 47 | def _load_model(self, |
| 48 | path: str, |
| 49 | max_seq_len: int, |
| 50 | max_batch_size: int, |
| 51 | tokenizer_path: Optional[str] = None, |
| 52 | num_gpus: int = 2): |
| 53 | from mixtralkit.mixtral import Mixtral |
| 54 | self.generator = Mixtral.build(ckpt_dir=path, |
| 55 | tokenizer_path=tokenizer_path, |
| 56 | max_seq_len=max_seq_len, |
| 57 | max_batch_size=max_batch_size, |
| 58 | num_gpus=num_gpus) |
| 59 | self.tokenizer = self.generator.tokenizer |
| 60 | self.model = self.generator.model |
| 61 | |
| 62 | def _load_tokenizer(self, tokenizer_path: str): |
| 63 | from mixtralkit.layers import Tokenizer |