import torch import itertools import threading import numpy as np from sklearn.metrics.pairwise import cosine_similarity from collections import Counter, defaultdict from loguru import logger from abc import ABCMeta, abstractmethod from .paper_client import PaperClient from .paper_crawling import PaperCrawling from .llms_api import APIHelper from .hash import get_embedding_model class UnionFind: def __init__(self, n): self.parent = list(range(n)) self.rank = [1] * n def find(self, x): if self.parent[x] != x: self.parent[x] = self.find(self.parent[x]) return self.parent[x] def union(self, x, y): rootX = self.find(x) rootY = self.find(y) if rootX != rootY: if self.rank[rootX] > self.rank[rootY]: self.parent[rootY] = rootX elif self.rank[rootX] < self.rank[rootY]: self.parent[rootX] = rootY else: self.parent[rootY] = rootX self.rank[rootX] += 1 def can_merge(uf, similarity_matrix, i, j, threshold): """Condition of i and j can be merged: After merging, the similarity of any two nodes from root_i and root_j are larger than threshold """ root_i = uf.find(i) root_j = uf.find(j) for k in range(len(similarity_matrix)): if uf.find(k) == root_i or uf.find(k) == root_j: if ( similarity_matrix[i][k] < threshold or similarity_matrix[j][k] < threshold ): return False return True class CoCite: _instance = None _initialized = False def __new__(cls, *args, **kwargs): if cls._instance is None: cls._instance = super(CoCite, cls).__new__(cls) return cls._instance def __init__(self) -> None: if not self._initialized: logger.debug("init co-cite map begin...") self.paper_client = PaperClient() citemap = self.paper_client.build_citemap() self.comap = defaultdict(lambda: defaultdict(int)) for paper_id, cited_id in citemap.items(): for id0, id1 in itertools.combinations(cited_id, 2): # ensure comap[id0][id1] == comap[id1][id0] self.comap[id0][id1] += 1 self.comap[id1][id0] += 1 logger.debug("init co-cite map success") CoCite._initialized = True def get_cocite_ids(self, id_, k=1): """ """ sorted_items = sorted(self.comap[id_].items(), key=lambda x: x[1], reverse=True) top_k = sorted_items[:k] paper_ids = [] for item in top_k: paper_ids.append(item[0]) paper_ids = self.paper_client.filter_paper_id_list(paper_ids) return paper_ids class Retriever(object): """The superclass of all retrievers Args: config: Returns: A Retriever instance """ __metaclass__ = ABCMeta retriever_name = "BASE" def __init__(self, config): self.config = config self.use_cocite = config.RETRIEVE.use_cocite self.use_cluster_to_filter = config.RETRIEVE.use_cluster_to_filter self.paper_client = PaperClient() self.cocite = CoCite() self.api_helper = APIHelper(config=config) self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") self.embedding_model = get_embedding_model(config) self.paper_crawling = PaperCrawling(config=config) if self.config.DEFAULT.embedding == "sentence-transformers/all-MiniLM-L6-v2": self.embedding_postfix = "" elif self.config.DEFAULT.embedding == "BAAI/llm-embedder": self.embedding_postfix = "_llm_embedder" elif self.config.DEFAULT.embedding == "jinaai/jina-embeddings-v3": self.embedding_postfix = "_jina_v3" if self.config.DEFAULT.embedding_database == "text-matching": self.embedding_postfix += "_text_matching" elif self.config.DEFAULT.embedding_database == "retrieval.query": self.embedding_postfix += "_query" elif self.config.DEFAULT.embedding_database == "retrieval.passage": self.embedding_postfix += "_passage" @abstractmethod def retrieve(self, bg, entities, use_evaluate): """Retrieve papers, should be implemented by the sub-class Args: None Returns: None """ pass def retrieve_entities_by_enties(self, entities): """The method do three things: 1. Expand entities according to entities co-occurence 2. Count the number of papers related to each expanded entity. Sort entities in terms of their occurence times in ascending order 3. Initial new entities. Retrieve entities one by one until the number of related papers reach a threshold Args: entities: A List of entities, e.g., [str, str, ...] Returns: new_entities: A List of entities after expansion, e.g., [str, str, ...] """ # TODO: KG expand_entities = self.paper_client.find_related_entities_by_entity_list( entities, n=self.config.RETRIEVE.kg_jump_num, k=self.config.RETRIEVE.kg_cover_num, relation_name=self.config.RETRIEVE.relation_name, ) expand_entities = list(set(entities + expand_entities)) entity_paper_num_dict = self.paper_client.get_entities_related_paper_num( expand_entities ) new_entities = [] entity_paper_num_dict = { k: v for k, v in entity_paper_num_dict.items() if v != 0 } entity_paper_num_dict = dict( sorted(entity_paper_num_dict.items(), key=lambda item: item[1]) ) sum_paper_num = 0 for key, value in entity_paper_num_dict.items(): if sum_paper_num <= self.config.RETRIEVE.sum_paper_num: sum_paper_num += value new_entities.append(key) elif ( value < self.config.RETRIEVE.limit_num and sum_paper_num < self.config.RETRIEVE.sum_paper_num ): sum_paper_num += value new_entities.append(key) return new_entities def update_related_paper(self, paper_id_list): """ Args: paper_id_list (List of hash_id): e.g., [1231214, 46345] Return: related_paper (List of dict): """ related_paper = self.paper_client.update_papers_from_client(paper_id_list) return related_paper def calculate_similarity(self, entities, related_entities_list, use_weight=False): """[Deprecated] Calculate the similarities between two lists of entities """ if use_weight: vec1 = self.vectorizer.transform([" ".join(entities)]).toarray()[0] weighted_vec1 = np.array( [ vec1[i] * self.log_inverse_freq.get(word, 1) for i, word in enumerate(self.vectorizer.get_feature_names_out()) ] ) vecs2 = self.vectorizer.transform( [ " ".join(related_entities) for related_entities in related_entities_list ] ).toarray() weighted_vecs2 = np.array( [ [ vec2[i] * self.log_inverse_freq.get(word, 1) for i, word in enumerate( self.vectorizer.get_feature_names_out() ) ] for vec2 in vecs2 ] ) similarity = cosine_similarity([weighted_vec1], weighted_vecs2)[0] else: vec1 = self.vectorizer.transform([" ".join(entities)]) vecs2 = self.vectorizer.transform( [ " ".join(related_entities) for related_entities in related_entities_list ] ) similarity = cosine_similarity(vec1, vecs2)[0] return similarity def cal_related_score( self, embedding, related_paper_id_list, type_name="background_embedding" ): """Calculate the cosine similarity between the input background's embedding and given list of papers Args: embedding: the embedding of the input background related_paper_id_list (List of int): the paper ids in the database Returns: Empty dict: {} Empty dict: {} score_all_dict: paper_id1: score1, paper_id2: score2, ... """ score_1 = np.zeros((len(related_paper_id_list))) # score_2 = np.zeros((len(related_paper_id_list))) origin_vector = torch.tensor(embedding).to(self.device).unsqueeze(0) context_embeddings = self.paper_client.get_papers_attribute( related_paper_id_list, type_name ) if len(context_embeddings) > 0: context_embeddings = torch.tensor(context_embeddings).to(self.device) score_1 = torch.nn.functional.cosine_similarity( origin_vector, context_embeddings ) score_1 = score_1.cpu().numpy() if self.config.RETRIEVE.need_normalize: score_1 = score_1 / np.max(score_1) score_all_dict = dict(zip(related_paper_id_list, score_1)) # score_en_dict = dict(zip(related_paper_id_list, score_2)) """ score_all_dict = dict( zip( related_paper_id_list, score_1 * self.config.RETRIEVE.alpha + score_2 * self.config.RETRIEVE.beta, ) ) """ return {}, {}, score_all_dict def filter_related_paper(self, score_dict, top_k): """Pick top_k papers from all retrieved papers in terms of score_dict. If clustering is not used, top_k papers with highest scores will be picked. If clustering is used, we will pick papers from each cluster in turn util top_k papers are chosen. Args: score_dict (dict): dict of (paper_id, similarity with user input background) top_k (int): pick top_k papers Returns: """ if len(score_dict) <= top_k: return list(score_dict.keys()) if not self.use_cluster_to_filter: paper_id_list = ( list(score_dict.keys())[:top_k] if len(score_dict) >= top_k else list(score_dict.keys()) ) return paper_id_list else: ## Calculate the final embedding for each paper, which is the weighted average ## background_embedding (embedding), contribution_embedding, and summary_embedding. # clustering filter, ensure that each category the highest score save first # background embedding paper_id_list = list(score_dict.keys()) paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"background_embedding{self.embedding_postfix}") for paper_id in paper_id_list ] paper_embedding = np.array(paper_embedding_list) # contribution embedding paper_embedding_list = [ self.paper_client.get_paper_attribute( paper_id, f"contribution_embedding{self.embedding_postfix}" ) for paper_id in paper_id_list ] paper_contribution_embedding = np.array(paper_embedding_list) # summary embedding paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"summary_embedding{self.embedding_postfix}") for paper_id in paper_id_list ] paper_summary_embedding = np.array(paper_embedding_list) # abstract embedding paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"abstract_embedding{self.embedding_postfix}") for paper_id in paper_id_list ] paper_abstract_embedding = np.array(paper_embedding_list) weight_background = self.config.RETRIEVE.s_bg weight_contribution = self.config.RETRIEVE.s_contribution weight_summary = self.config.RETRIEVE.s_summary weight_abstract = self.config.RETRIEVE.s_abstract paper_embedding = ( weight_background * paper_embedding + weight_contribution * paper_contribution_embedding + weight_summary * paper_summary_embedding + weight_abstract * paper_abstract_embedding ) ## similarity_matrix of all retrieved papers similarity_matrix = np.dot(paper_embedding, paper_embedding.T) related_labels = self.cluster_algorithm(paper_id_list, similarity_matrix) related_paper_label_dict = dict(zip(paper_id_list, related_labels)) label_group = {} for paper_id, label in related_paper_label_dict.items(): if label not in label_group: label_group[label] = [] label_group[label].append(paper_id) paper_id_list = [] # randomly pick a paper from each cluster in turn until top_k papers are chosen while len(paper_id_list) < top_k: for label, papers in label_group.items(): if papers: paper_id_list.append(papers.pop(0)) if len(paper_id_list) >= top_k: break return paper_id_list def cosine_similarity_search(self, embedding, k=1, type_name="background_embedding"): """Retrieve papers through embedding Args: embedding: the input embedding Returns: result (List of Papers): return related papers with the least embedding distance """ result = self.paper_client.cosine_similarity_search( embedding, k, type_name=type_name ) # backtrack: first is itself result = result[1:] return result def cluster_algorithm(self, paper_id_list, similarity_matrix): """ """ threshold = self.config.RETRIEVE.similarity_threshold uf = UnionFind(len(paper_id_list)) # merge for i in range(len(similarity_matrix)): for j in range(i + 1, len(similarity_matrix)): if similarity_matrix[i][j] >= threshold: if can_merge(uf, similarity_matrix, i, j, threshold): uf.union(i, j) cluster_labels = [uf.find(i) for i in range(len(similarity_matrix))] return cluster_labels def eval_related_paper_in_all(self, score_all_dict, target_paper_id_list): score_all_dict = dict( sorted(score_all_dict.items(), key=lambda item: item[1], reverse=True) ) result = {} related_paper_id_list = list(score_all_dict.keys()) if len(related_paper_id_list) == 0: for k in self.config.RETRIEVE.top_k_list: result[k] = {"recall": 0, "precision": 0} return result, 0, 0, 0 ## merge retrieved papers and target papers and clustering ## clustering according to the combination of background, contribution, and summary_embedding all_paper_id_set = set(related_paper_id_list) all_paper_id_set.update(target_paper_id_list) all_paper_id_list = list(all_paper_id_set) # get all target papers' background_embedding paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"background_embedding{self.embedding_postfix}") for paper_id in target_paper_id_list ] paper_embedding = np.array(paper_embedding_list) # get all target papers' contribution_embedding paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"contribution_embedding{self.embedding_postfix}") for paper_id in target_paper_id_list ] paper_contribution_embedding = np.array(paper_embedding_list) # get all target papers' summary_embedding paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"summary_embedding{self.embedding_postfix}") for paper_id in target_paper_id_list ] # abstract embedding paper_embedding_list = [ self.paper_client.get_paper_attribute(paper_id, f"abstract_embedding{self.embedding_postfix}") for paper_id in target_paper_id_list ] paper_abstract_embedding = np.array(paper_embedding_list) paper_summary_embedding = np.array(paper_embedding_list) weight_background = self.config.RETRIEVE.s_bg weight_contribution = self.config.RETRIEVE.s_contribution weight_summary = self.config.RETRIEVE.s_summary weight_abstract = self.config.RETRIEVE.s_abstract # 2D matrix of size [# of target papers, embedding dimension] target_paper_embedding = ( weight_background * paper_embedding + weight_contribution * paper_contribution_embedding + weight_summary * paper_summary_embedding + weight_abstract * paper_abstract_embedding ) similarity_threshold = self.config.RETRIEVE.similarity_threshold similarity_matrix = np.dot(target_paper_embedding, target_paper_embedding.T) # return each target_paper's cluster label target_labels = self.cluster_algorithm(target_paper_id_list, similarity_matrix) target_paper_label_dict = dict(zip(target_paper_id_list, target_labels)) logger.debug("Target paper cluster result: {}".format(target_paper_label_dict)) logger.debug( { paper_id: self.paper_client.get_paper_attribute(paper_id, "title") for paper_id in target_paper_label_dict.keys() } ) ## calculate the similarity between each two papers all_labels = [] for paper_id in all_paper_id_list: # for each paper, get its background_embedding paper_bg_embedding = [ self.paper_client.get_paper_attribute(paper_id, f"background_embedding{self.embedding_postfix}") ] paper_bg_embedding = np.array(paper_bg_embedding) # for each paper, get its contribution_embedding paper_contribution_embedding = [ self.paper_client.get_paper_attribute( paper_id, f"contribution_embedding{self.embedding_postfix}" ) ] paper_contribution_embedding = np.array(paper_contribution_embedding) # for each paper, get its summary_embedding paper_summary_embedding = [ self.paper_client.get_paper_attribute(paper_id, f"summary_embedding{self.embedding_postfix}") ] paper_summary_embedding = np.array(paper_summary_embedding) # for each paper, get its abstract_embedding paper_abstract_embedding = [ self.paper_client.get_paper_attribute(paper_id, f"abstract_embedding{self.embedding_postfix}") ] paper_abstract_embedding = np.array(paper_abstract_embedding) paper_embedding = ( weight_background * paper_bg_embedding + weight_contribution * paper_contribution_embedding + weight_summary * paper_summary_embedding + weight_abstract * paper_abstract_embedding ) # vector of size embedding dimension similarities = cosine_similarity(paper_embedding, target_paper_embedding)[0] if np.any(similarities >= similarity_threshold): all_labels.append(target_labels[np.argmax(similarities)]) else: all_labels.append(-1) # other class: -1 all_paper_label_dict = dict(zip(all_paper_id_list, all_labels)) all_label_counts = Counter(all_paper_label_dict.values()) logger.debug(f"All labels and the number of papers of each label: {all_label_counts}") target_label_counts = Counter(target_paper_label_dict.values()) logger.debug(f"All labels and the number of target papers of each label : {target_label_counts}") target_label_list = list(target_label_counts.keys()) max_k = max(self.config.RETRIEVE.top_k_list) logger.info("=== Begin filter related paper ===") max_k_paper_id_list = self.filter_related_paper(score_all_dict, top_k=max_k) logger.info("=== End filter related paper ===") ## calculate recall and precision of first {10, 20, 30, ...} papers for k in self.config.RETRIEVE.top_k_list: # 前top k 的文章 top_k = min(k, len(max_k_paper_id_list)) top_k_paper_id_list = max_k_paper_id_list[:top_k] top_k_paper_label_dict = {} for paper_id in top_k_paper_id_list: top_k_paper_label_dict[paper_id] = all_paper_label_dict[paper_id] logger.debug( "=== ideal top {}, real top {} paper id list : {}".format(k, top_k, top_k_paper_label_dict) ) logger.debug( { paper_id: self.paper_client.get_paper_attribute(paper_id, "title") for paper_id in top_k_paper_label_dict.keys() } ) top_k_label_counts = Counter(top_k_paper_label_dict.values()) logger.debug(f"Retrieved {top_k} papers have K different label: {top_k_label_counts}") top_k_label_list = list(top_k_label_counts.keys()) match_label_list = list(set(target_label_list) & set(top_k_label_list)) logger.debug(f"match label list : {match_label_list}") recall = 0 precision = 0 for label in match_label_list: recall += target_label_counts[label] for label in match_label_list: precision += top_k_label_counts[label] recall /= len(target_paper_id_list) precision /= len(top_k_paper_id_list) result[k] = {"recall": recall, "precision": precision} ## calculate recall and precision of all retrieved papers related_paper_id_list = list(score_all_dict.keys()) related_paper_label_dict = {} for paper_id in related_paper_id_list: related_paper_label_dict[paper_id] = all_paper_label_dict[paper_id] related_label_counts = Counter(related_paper_label_dict.values()) logger.debug(f"top K label counts : {related_label_counts}") related_label_list = list(related_label_counts.keys()) match_label_list = list(set(target_label_list) & set(related_label_list)) recall = 0 precision = 0 for label in match_label_list: recall += target_label_counts[label] for label in match_label_list: precision += related_label_counts[label] recall /= len(target_paper_id_list) precision /= len(related_paper_id_list) logger.debug(result) return result, len(target_label_counts), recall, precision class RetrieverFactory(object): """RetrieverFactory is a singleton class, which will return cls._instance if it has been created, it saves all Retriever instances. Args: None Returns: The singleton instance of the RetrieverFactory """ _instance = None _lock = threading.Lock() def __new__(cls, *args, **kwargs): with cls._lock: if cls._instance is None: cls._instance = super(RetrieverFactory, cls).__new__( cls, *args, **kwargs ) cls._instance.init_factory() return cls._instance def init_factory(self): self.retriever_classes = {} @staticmethod def get_retriever_factory(): """The method can also return the singleton instance of the RetrieverFactory Args: None Returns: The singleton instance of the RetrieverFactory """ if RetrieverFactory._instance is None: RetrieverFactory._instance = RetrieverFactory() return RetrieverFactory._instance def register_retriever(self, retriever_name, retriever_class) -> bool: """Register a new retriever class (not instance) to the RetrieverFactory Args: retriever_name: str retriever_class: a class object (not instance) Returns: True if add successfully, False otherwise """ if retriever_name not in self.retriever_classes: self.retriever_classes[retriever_name] = retriever_class return True else: return False def delete_retriever(self, retriever_name) -> bool: if retriever_name in self.retriever_classes: self.retriever_classes[retriever_name] = None del self.retriever_classes[retriever_name] return True else: return False def __getitem__(self, key): return self.retriever_classes[key] def __len__(self): return len(self.retriever_classes) def create_retriever(self, retriever_name, *args, **kwargs) -> Retriever: """Return a retriever instance Args: retriever_name: str Returns: The retriever """ if retriever_name not in self.retriever_classes: raise ValueError(f"Unknown retriever type: {retriever_name}. retriever_name should be one of {self.retriever_classes.keys()}") else: return self.retriever_classes[retriever_name](*args, **kwargs) class autoregister: def __init__(self, retriever_name, *args, **kwds): self.retriever_name = retriever_name def __call__(self, cls, *args, **kwds): if RetrieverFactory.get_retriever_factory().register_retriever( self.retriever_name, cls ): cls.retriever_name = self.retriever_name return cls else: raise KeyError() @autoregister("SN") class SNRetriever(Retriever): def __init__(self, config): super().__init__(config) def retrieve_paper(self, bg): """Retrieve papers P (a set) according to embeddings' similarity between the input background and the backgrounds from the database. Optionally, you can also retrieve papers co-cited with P. Args: bg (str): the input background Returns: result (dict): "background_embedding": embedding of the input background, "paper" (List of int): all retrieved related_papers' ids, "entities" (List): An empty list (TODO: remove), "cocite_paper" (List of int): all papers cocited with embedding-retrieved papers """ entities = [] embedding = self.embedding_model.encode(bg, device=self.device) sn_paper_id_list = self.cosine_similarity_search( embedding=embedding, k=self.config.RETRIEVE.sn_retrieve_paper_num, type_name=f"{self.config.RETRIEVE.SN_field_name}_embedding{self.embedding_postfix}" ) related_paper = set() related_paper.update(sn_paper_id_list) cocite_id_set = set() if self.use_cocite: for paper_id in related_paper: cocite_id_set.update( self.cocite.get_cocite_ids( paper_id, k=self.config.RETRIEVE.cocite_top_k ) ) related_paper = related_paper.union(cocite_id_set) related_paper = list(related_paper) logger.debug(f"paper num before filter: {len(related_paper)}") result = { f"background_embedding{self.embedding_postfix}": embedding, "paper": related_paper, "entities": entities, "cocite_paper": list(cocite_id_set), } return result def retrieve(self, bg, entities, need_evaluate=True, target_paper_id_list=[]): """ Args: bg (str): The user input background Return: result (dict): "recall": recall of paper retrieval, 0 if need_evaluate==False, "precision": precision of paper retrieval, 0 if need_evaluate==False,, "filtered_recall": recall of paper retrieval after filtering, 0 if need_evaluate==False,, "filtered_precision": precision of paper retrieval after filtering, 0 if need_evaluate==False,, "related_paper": all retrieved related_papers. !!! [ The most important item ] "related_paper_id_list": all retrieved related_papers' ids. !!! [ The most important item ] "cocite_paper_id_list": retrieve_result["cocite_paper"], "entities": retrieve_result["entities"], always empty "top_k_matrix": top_k_matrix, 0 if need_evaluate==False "gt_reference_num": len(target_paper_id_list) "retrieve_paper_num": len(related_paper_id_list), "label_num": TODO, """ if need_evaluate: if target_paper_id_list is None or len(target_paper_id_list) == 0: logger.error( "If you need evaluate retriever, please input target paper is list..." ) else: target_paper_id_list = list(set(target_paper_id_list)) retrieve_result = self.retrieve_paper(bg) related_paper_id_list = retrieve_result["paper"] retrieve_paper_num = len(related_paper_id_list) # scores between the input background and all retrieved papers _, _, score_all_dict = self.cal_related_score( retrieve_result[f"background_embedding{self.embedding_postfix}"], related_paper_id_list=related_paper_id_list, type_name=f"{self.config.RETRIEVE.SN_field_name}_embedding{self.embedding_postfix}" ) top_k_matrix = {} recall = 0 precision = 0 filtered_recall = 0 label_num = 0 filtered_precision = 0 if need_evaluate: top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all( score_all_dict, target_paper_id_list ) logger.debug("Top K matrix:{}".format(top_k_matrix)) logger.debug("before filter:") logger.debug(f"Recall: {recall:.3f}") logger.debug(f"Precision: {precision:.3f}") ## For idea generation, only top 10 papers will be used, which has no relations with retriveal evaluation related_paper = self.filter_related_paper(score_all_dict, top_k=self.config.RETRIEVE.all_retrieve_paper_num) related_paper = self.update_related_paper(related_paper) result = { "recall": recall, "precision": precision, "filtered_recall": filtered_recall, "filtered_precision": filtered_precision, "related_paper": related_paper, "related_paper_id_list": related_paper_id_list, "cocite_paper_id_list": retrieve_result["cocite_paper"], "entities": retrieve_result["entities"], "top_k_matrix": top_k_matrix, "gt_reference_num": len(target_paper_id_list), "retrieve_paper_num": retrieve_paper_num, "label_num": label_num, } return result @autoregister("KG") class KGRetriever(Retriever): def __init__(self, config): super().__init__(config) def retrieve_paper(self, entities): """Retrieve according to entities """ new_entities = self.retrieve_entities_by_enties(entities) logger.debug("KG entities for retriever: {}".format(new_entities)) related_paper = set() for entity in new_entities: paper_id_set = set(self.paper_client.find_paper_by_entity(entity)) related_paper = related_paper.union(paper_id_set) cocite_id_set = set() if self.use_cocite: for paper_id in related_paper: cocite_id_set.update(self.cocite.get_cocite_ids(paper_id)) related_paper = related_paper.union(cocite_id_set) related_paper = list(related_paper) logger.debug(f"paper num before filter: {len(related_paper)}") result = { "paper": related_paper, "entities": entities, "cocite_paper": list(cocite_id_set), } return result def retrieve(self, bg, entities, need_evaluate=False, target_paper_id_list=[]): """ Args: context: string Return: list(dict) """ if need_evaluate: if target_paper_id_list is None or len(target_paper_id_list) == 0: logger.error( "If you need evaluate retriever, please input target paper is list..." ) else: target_paper_id_list = list(set(target_paper_id_list)) logger.debug(f"target paper id list: {target_paper_id_list}") retrieve_result = self.retrieve_paper(entities) related_paper_id_list = retrieve_result["paper"] retrieve_paper_num = len(related_paper_id_list) embedding = self.embedding_model.encode(bg, device=self.device) _, _, score_all_dict = self.cal_related_score( embedding, related_paper_id_list=related_paper_id_list, type_name=f"background_embedding{self.embedding_postfix}" ) top_k_matrix = {} recall = 0 precision = 0 filtered_recall = 0 label_num = 0 filtered_precision = 0 if need_evaluate: top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all( score_all_dict, target_paper_id_list ) logger.debug("Top P ACC:{}".format(top_k_matrix)) logger.debug("before filter:") logger.debug(f"Recall: {recall:.3f}") logger.debug(f"Precision: {precision:.3f}") related_paper = self.filter_related_paper(score_all_dict, top_k=self.config.RETRIEVE.all_retrieve_paper_num) related_paper = self.update_related_paper(related_paper) result = { "recall": recall, "precision": precision, "filtered_recall": filtered_recall, "filtered_precision": filtered_precision, "related_paper": related_paper, "related_paper_id_list": related_paper_id_list, "cocite_paper_id_list": retrieve_result["cocite_paper"], "entities": retrieve_result["entities"], "top_k_matrix": top_k_matrix, "gt_reference_num": len(target_paper_id_list), "retrieve_paper_num": retrieve_paper_num, "label_num": label_num, } return result @autoregister("SNKG") class SNKGRetriever(Retriever): def __init__(self, config): super().__init__(config) def retrieve_paper(self, bg, entities): sn_entities = [] ## 1. Retrieve papers according to the embeddings of input background embedding = self.embedding_model.encode(bg, device=self.device) sn_paper_id_list = self.cosine_similarity_search( embedding, k=self.config.RETRIEVE.sn_num_for_entity, type_name=f"{self.config.RETRIEVE.SN_field_name}_embedding{self.embedding_postfix}" ) related_paper = set() related_paper.update(sn_paper_id_list) logger.debug(f"SN retrieve {len(related_paper)} papers") ## 2. Retrieve papers according to entites # Fetch all entities from embedding-retrieved papers sn_entities += self.paper_client.find_entities_by_paper_list(sn_paper_id_list) logger.debug("SN entities for retriever: {}".format(sn_entities)) entities = list(set(entities + sn_entities)) # Expand entity list through synonyms new_entities = self.retrieve_entities_by_enties(entities) logger.debug("SNKG entities for retriever: {}".format(new_entities)) paper_id_set = set() for entity in new_entities: paper_id_set.update(self.paper_client.find_paper_by_entity(entity)) related_paper = related_paper.union(paper_id_set) logger.debug(f"Entity retrieve {len(paper_id_set)} papers") logger.debug(f"SN+entity retrieve {len(related_paper)} papers") ## 3. Retrieve papers according to citation co-occurrence cocite_id_set = set() if self.use_cocite: for paper_id in related_paper: cocite_id_set.update(self.cocite.get_cocite_ids(paper_id)) related_paper = related_paper.union(cocite_id_set) logger.debug(f"Cocite retrieve {len(cocite_id_set)} papers") logger.debug(f"SN+entity+cocite retrieve {len(related_paper)} papers") ## 4. Return retrieval results related_paper = list(related_paper) result = { f"background_embedding{self.embedding_postfix}": embedding, "paper": related_paper, "entities": entities, "cocite_paper": list(cocite_id_set), } return result def retrieve( self, bg, entities, need_evaluate=True, target_paper_id_list=[] ): """ Args: context: string Return: list(dict) """ if need_evaluate: if target_paper_id_list is None or len(target_paper_id_list) == 0: logger.error( "If you need evaluate retriever, please input target paper is list..." ) else: target_paper_id_list = list(set(target_paper_id_list)) logger.debug(f"target paper id list: {target_paper_id_list}") retrieve_result = self.retrieve_paper(bg, entities) related_paper_id_list = retrieve_result["paper"] retrieve_paper_num = len(related_paper_id_list) logger.info("=== Begin cal related paper score ===") _, _, score_all_dict = self.cal_related_score( retrieve_result[f"background_embedding{self.embedding_postfix}"], related_paper_id_list=related_paper_id_list, type_name=f"background_embedding{self.embedding_postfix}" ) logger.info("=== End cal related paper score ===") top_k_matrix = {} recall = 0 precision = 0 filtered_recall = 0 filtered_precision = 0 label_num = 0 if need_evaluate: top_k_matrix, label_num, recall, precision = self.eval_related_paper_in_all( score_all_dict, target_paper_id_list ) logger.debug("Top K matrix:{}".format(top_k_matrix)) logger.debug("before filter:") logger.debug(f"Recall: {recall:.3f}") logger.debug(f"Precision: {precision:.3f}") logger.info("=== Begin filter related paper score ===") related_paper = self.filter_related_paper(score_all_dict, self.config.RETRIEVE.all_retrieve_paper_num) logger.info("=== End filter related paper score ===") related_paper = self.update_related_paper(related_paper) result = { "recall": recall, "precision": precision, "filtered_recall": filtered_recall, "filtered_precision": filtered_precision, "related_paper": related_paper, "cocite_paper_id_list": retrieve_result["cocite_paper"], "related_paper_id_list": related_paper_id_list, "entities": retrieve_result["entities"], "top_k_matrix": top_k_matrix, "gt_reference_num": ( len(target_paper_id_list) if target_paper_id_list is not None else 0 ), "retrieve_paper_num": retrieve_paper_num, "label_num": label_num, } return result