diff --git a/opencompass/openicl/icl_retriever/__init__.py b/opencompass/openicl/icl_retriever/__init__.py index c1b48c2b..1f39b3f8 100644 --- a/opencompass/openicl/icl_retriever/__init__.py +++ b/opencompass/openicl/icl_retriever/__init__.py @@ -7,3 +7,4 @@ from .icl_random_retriever import RandomRetriever # noqa from .icl_topk_retriever import TopkRetriever # noqa from .icl_votek_retriever import VotekRetriever # noqa from .icl_zero_retriever import ZeroRetriever # noqa +from .icl_sliding_k_retriever import SlidingWindowRetriever # noqa diff --git a/opencompass/openicl/icl_retriever/icl_sliding_k_retriever.py b/opencompass/openicl/icl_retriever/icl_sliding_k_retriever.py new file mode 100644 index 00000000..f9cfe8f0 --- /dev/null +++ b/opencompass/openicl/icl_retriever/icl_sliding_k_retriever.py @@ -0,0 +1,67 @@ +"""Sliding Window Retriever.""" + +from typing import Optional + +from tqdm import trange + +from opencompass.openicl.icl_retriever import BaseRetriever +from opencompass.openicl.utils.logging import get_logger +from opencompass.registry import ICL_RETRIEVERS + +logger = get_logger(__name__) + + +@ICL_RETRIEVERS.register_module() +class SlidingWindowRetriever(BaseRetriever): + """Sliding Window Retriever. Each in-context example of the test prompts is + retrieved based on a sliding window from the index set. + + Args: + dataset (`BaseDataset`): + Any BaseDataset instances. + Attributes of ``reader``, ``train`` and ``test`` will be used. + k (int): + The number of in-context examples to retrieve for each test prompt. + ice_separator (`Optional[str]`): + The separator between each in-context + example template when origin `PromptTemplate` is provided. Defaults + to '\n'. + ice_eos_token (`Optional[str]`): + The end of sentence token for + in-context example template when origin `PromptTemplate` is + provided. Defaults to '\n'. + ice_num (`Optional[int]`): + The number of in-context example template + when origin `PromptTemplate` is provided. Defaults to 1. + """ + + def __init__(self, + dataset, + k: int, + ice_separator: Optional[str] = '\n', + ice_eos_token: Optional[str] = '\n', + ice_num: Optional[int] = 1) -> None: + super().__init__(dataset, ice_separator, ice_eos_token, ice_num) + self.k = k + + def retrieve(self): + """Retrieve the in-context example index for each test example.""" + num_idx = len(self.index_ds) + rtr_idx_list = [] + for current_index in trange(len(self.test_ds), + disable=not self.is_main_process): + if current_index < self.k: + """For the first few examples, + get the previous ones and pad with the last ones""" + start_index = max(0, current_index - self.k) + previous_shots = list(range(start_index, current_index)) + if len(previous_shots) < self.k: + pad_needed = self.k - len(previous_shots) + previous_shots = list(range(num_idx - pad_needed, + num_idx)) + previous_shots + else: + # For other examples, retrieve the previous k examples + previous_shots = list( + range(current_index - self.k, current_index)) + rtr_idx_list.append(previous_shots) + return rtr_idx_list