[Feature] Add icl_sliding_k_retriever.py and update __init__.py (#1305)

* Add icl_sliding_k_retriever.py and update __init__.py

* Fix flake8, isort, and yapf issues for Sliding Window Retriever

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

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