| import os |
| from abc import ABC, abstractmethod |
| from typing import Dict, List, Optional, Union |
|
|
| import torch.distributed as dist |
|
|
| from lmms_eval.api.registry import ALL_TASKS |
| from lmms_eval.tasks import ( |
| ConfigurableTask, |
| get_task_dict, |
| include_path, |
| initialize_tasks, |
| ) |
|
|
|
|
| def rank0_print(*args): |
| if dist.is_initialized(): |
| if dist.get_rank() == 0: |
| print(f"Rank {dist.get_rank()}: ", *args) |
| else: |
| print(*args) |
|
|
|
|
| class BaseEmbedder(ABC): |
| def __init__(self, name: str, output_path: str) -> None: |
| super().__init__() |
| self.name = name |
| self.output_path = output_path |
| os.makedirs(self.output_path, exist_ok=True) |
| initialize_tasks() |
|
|
| def flatten(self, input): |
| new_list = [] |
| for i in input: |
| for j in i: |
| new_list.append(j) |
| return new_list |
|
|
| |
| |
| @staticmethod |
| def init_task(task: str, ignored_ids: Union[set, List] = None): |
| task_dict = get_task_dict([task], model_name="llava") |
| task_obj = task_dict[task] |
| if type(task_obj) == tuple: |
| group, task_obj = task_obj |
| DATASET_PATH = task_obj.DATASET_PATH |
| DATASET_NAME = None |
| if task_obj.DATASET_NAME is not None: |
| DATASET_NAME = task_obj.DATASET_NAME |
|
|
| docs = task_obj.test_docs() if task_obj.has_test_docs() else task_obj.validation_docs() |
| split = task_obj.config.test_split if task_obj.has_test_docs() else task_obj.config.validation_split |
| rank0_print(f"\nTask : {task_obj.config.task}\n - #num : {len(task_obj.test_docs()) if task_obj.has_test_docs() else task_obj.validation_docs()}") |
| task_obj.build_all_requests() |
| requests = [] |
| for instance in task_obj.instances: |
| reqtype = instance.request_type |
| contexts, all_gen_kwargs, doc_to_visual, doc_id, task, split = instance.args |
| if ignored_ids is not None and doc_id in ignored_ids: |
| continue |
| requests.append(instance) |
| return DATASET_PATH, DATASET_NAME, split, requests, task_obj, docs |
|
|
| @abstractmethod |
| def embed_task(self, task: str): |
| return |
|
|
