The base class of each model worker.
| 18 | |
| 19 | @dataclass |
| 20 | class BaseWorker(): |
| 21 | """ |
| 22 | The base class of each model worker. |
| 23 | """ |
| 24 | cfg: dict |
| 25 | input_pth: str |
| 26 | output_pth: str |
| 27 | batch_size: int |
| 28 | use_cot: bool = False |
| 29 | use_qa: bool = False |
| 30 | generate_fewshot_examples_only: bool = False |
| 31 | use_fewshot: bool = False, |
| 32 | |
| 33 | def __post_init__(self): |
| 34 | if self.generate_fewshot_examples_only: # no need to do post_init if we only need to generate fewshot examples |
| 35 | return |
| 36 | self.print_in_main(f'loading config: {self.cfg.load}') |
| 37 | self.model, self.tokenizer = self.load_model_and_tokenizer(self.cfg.load) |
| 38 | self.device = self.cfg.load.device |
| 39 | self.accelerator = Accelerator() |
| 40 | self.prompt_wrapper = PromptWrapper( |
| 41 | self.tokenizer, |
| 42 | self.instruction_template_with_fewshot if self.use_fewshot else self.instruction_template, |
| 43 | conv_collater=self.collate_conv, |
| 44 | use_cot=self.use_cot, |
| 45 | ) |
| 46 | self.wrap_model() |
| 47 | self.init_generation_config(self.cfg) |
| 48 | self.init_dataloader(self.input_pth, self.batch_size) |
| 49 | self.init_writer(self.output_pth) |
| 50 | |
| 51 | |
| 52 | @classmethod |
| 53 | def from_config( |
| 54 | cls, |
| 55 | cfg, |
| 56 | input_pth: str = '', |
| 57 | output_pth: str = '', |
| 58 | batch_size = 1, |
| 59 | use_qa = False, |
| 60 | use_cot = False, |
| 61 | generate_fewshot_examples_only = False, |
| 62 | use_fewshot = False, |
| 63 | ): |
| 64 | assert cfg.get('load', None) is not None |
| 65 | |
| 66 | return cls( |
| 67 | cfg, |
| 68 | input_pth, |
| 69 | output_pth, |
| 70 | batch_size, |
| 71 | use_cot = use_cot, |
| 72 | use_qa = use_qa, |
| 73 | generate_fewshot_examples_only = generate_fewshot_examples_only, |
| 74 | use_fewshot = use_fewshot, |
| 75 | ) |
| 76 | |
| 77 |
nothing calls this directly
no outgoing calls
no test coverage detected