| """Random Retriever.""" | |
| from typing import Optional | |
| import numpy as np | |
| from tqdm import trange | |
| from opencompass.openicl.icl_retriever import BaseRetriever | |
| from opencompass.openicl.utils.logging import get_logger | |
| logger = get_logger(__name__) | |
| class RandomRetriever(BaseRetriever): | |
| """Random Retriever. Each in-context example of the test prompts is | |
| retrieved in a random way. | |
| **WARNING**: This class has not been tested thoroughly. Please use it with | |
| caution. | |
| """ | |
| def __init__(self, | |
| dataset, | |
| ice_separator: Optional[str] = '\n', | |
| ice_eos_token: Optional[str] = '\n', | |
| ice_num: Optional[int] = 1, | |
| seed: Optional[int] = 43) -> None: | |
| super().__init__(dataset, ice_separator, ice_eos_token, ice_num) | |
| self.seed = seed | |
| def retrieve(self): | |
| np.random.seed(self.seed) | |
| num_idx = len(self.index_ds) | |
| rtr_idx_list = [] | |
| logger.info('Retrieving data for test set...') | |
| for _ in trange(len(self.test_ds), disable=not self.is_main_process): | |
| idx_list = np.random.choice(num_idx, self.ice_num, | |
| replace=False).tolist() | |
| rtr_idx_list.append(idx_list) | |
| return rtr_idx_list | |