File size: 3,749 Bytes
256a159 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
"""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()
|