OpenICL Inference Task. This task is used to run the inference process.
| 19 | |
| 20 | @TASKS.register_module() |
| 21 | class OpenICLInferTask(BaseTask): |
| 22 | """OpenICL Inference Task. |
| 23 | |
| 24 | This task is used to run the inference process. |
| 25 | """ |
| 26 | |
| 27 | name_prefix = 'OpenICLInfer' |
| 28 | log_subdir = 'logs/infer' |
| 29 | output_subdir = 'predictions' |
| 30 | |
| 31 | def __init__(self, cfg: ConfigDict): |
| 32 | super().__init__(cfg) |
| 33 | run_cfg = self.model_cfgs[0].get('run_cfg', {}) |
| 34 | self.num_gpus = run_cfg.get('num_gpus', 0) |
| 35 | self.num_procs = run_cfg.get('num_procs', 1) |
| 36 | self.logger = get_logger() |
| 37 | |
| 38 | def get_command(self, cfg_path, template): |
| 39 | """Get the command template for the task. |
| 40 | |
| 41 | Args: |
| 42 | cfg_path (str): The path to the config file of the task. |
| 43 | template (str): The template which have '{task_cmd}' to format |
| 44 | the command. |
| 45 | """ |
| 46 | sys.path.append(os.getcwd()) |
| 47 | script_path = __file__ |
| 48 | backend_keys = ['VLLM', 'Lmdeploy'] |
| 49 | use_backend = any( |
| 50 | key in str(self.model_cfgs[0].get('type', '')) |
| 51 | or key in str(self.model_cfgs[0].get('llm', {}).get('type', '')) |
| 52 | for key in backend_keys) |
| 53 | if self.num_gpus > 1 and not use_backend: |
| 54 | port = random.randint(12000, 32000) |
| 55 | command = (f'torchrun --master_port={port} ' |
| 56 | f'--nproc_per_node {self.num_procs} ' |
| 57 | f'{script_path} {cfg_path}') |
| 58 | else: |
| 59 | python = sys.executable |
| 60 | command = f'{python} {script_path} {cfg_path}' |
| 61 | |
| 62 | return template.format(task_cmd=command) |
| 63 | |
| 64 | def run(self, cur_model=None, cur_model_abbr=None): |
| 65 | self.logger.info(f'Task {task_abbr_from_cfg(self.cfg)}') |
| 66 | for model_cfg, dataset_cfgs in zip(self.model_cfgs, self.dataset_cfgs): |
| 67 | self.max_out_len = model_cfg.get('max_out_len', None) |
| 68 | self.batch_size = model_cfg.get('batch_size', None) |
| 69 | self.min_out_len = model_cfg.get('min_out_len', None) |
| 70 | if cur_model and cur_model_abbr == model_abbr_from_cfg(model_cfg): |
| 71 | self.model = cur_model |
| 72 | else: |
| 73 | self.model = build_model_from_cfg(model_cfg) |
| 74 | |
| 75 | for dataset_cfg in dataset_cfgs: |
| 76 | self.model_cfg = model_cfg |
| 77 | self.dataset_cfg = dataset_cfg |
| 78 | self.infer_cfg = self.dataset_cfg['infer_cfg'] |
no outgoing calls
no test coverage detected