| """Votek Retriever.""" | |
| import json | |
| import os | |
| import random | |
| from collections import defaultdict | |
| from typing import Optional | |
| import numpy as np | |
| from sklearn.metrics.pairwise import cosine_similarity | |
| from opencompass.openicl.icl_retriever.icl_topk_retriever import TopkRetriever | |
| class VotekRetriever(TopkRetriever): | |
| """Vote-k In-context Learning Retriever, subclass of `TopkRetriever`. | |
| **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, | |
| sentence_transformers_model_name: Optional[ | |
| str] = 'all-mpnet-base-v2', | |
| tokenizer_name: Optional[str] = 'gpt2-xl', | |
| batch_size: Optional[int] = 1, | |
| votek_k: Optional[int] = 3) -> None: | |
| super().__init__(dataset, ice_separator, ice_eos_token, ice_num, | |
| sentence_transformers_model_name, tokenizer_name, | |
| batch_size) | |
| self.votek_k = votek_k | |
| def votek_select(self, | |
| embeddings=None, | |
| select_num=None, | |
| k=None, | |
| overlap_threshold=None, | |
| vote_file=None): | |
| n = len(embeddings) | |
| if vote_file is not None and os.path.isfile(vote_file): | |
| with open(vote_file, encoding='utf-8') as f: | |
| vote_stat = json.load(f) | |
| else: | |
| vote_stat = defaultdict(list) | |
| for i in range(n): | |
| cur_emb = embeddings[i].reshape(1, -1) | |
| cur_scores = np.sum(cosine_similarity(embeddings, cur_emb), | |
| axis=1) | |
| sorted_indices = np.argsort(cur_scores).tolist()[-k - 1:-1] | |
| for idx in sorted_indices: | |
| if idx != i: | |
| vote_stat[idx].append(i) | |
| if vote_file is not None: | |
| with open(vote_file, 'w', encoding='utf-8') as f: | |
| json.dump(vote_stat, f) | |
| votes = sorted(vote_stat.items(), | |
| key=lambda x: len(x[1]), | |
| reverse=True) | |
| j = 0 | |
| selected_indices = [] | |
| while len(selected_indices) < select_num and j < len(votes): | |
| candidate_set = set(votes[j][1]) | |
| flag = True | |
| for pre in range(j): | |
| cur_set = set(votes[pre][1]) | |
| if len(candidate_set.intersection( | |
| cur_set)) >= overlap_threshold * len(candidate_set): | |
| flag = False | |
| break | |
| if not flag: | |
| j += 1 | |
| continue | |
| selected_indices.append(int(votes[j][0])) | |
| j += 1 | |
| if len(selected_indices) < select_num: | |
| unselected_indices = [] | |
| cur_num = len(selected_indices) | |
| for i in range(n): | |
| if i not in selected_indices: | |
| unselected_indices.append(i) | |
| selected_indices += random.sample(unselected_indices, | |
| select_num - cur_num) | |
| return selected_indices | |
| def vote_k_search(self): | |
| vote_k_idxs = self.votek_select(embeddings=self.embed_list, | |
| select_num=self.ice_num, | |
| k=self.votek_k, | |
| overlap_threshold=1) | |
| return [vote_k_idxs[:] for _ in range(len(self.test_ds))] | |
| def retrieve(self): | |
| return self.vote_k_search() | |