| 7 | |
| 8 | |
| 9 | class ToolRetriever: |
| 10 | def __init__(self, corpus_tsv_path = "", model_path=""): |
| 11 | self.corpus_tsv_path = corpus_tsv_path |
| 12 | self.model_path = model_path |
| 13 | self.corpus, self.corpus2tool = self.build_retrieval_corpus() |
| 14 | self.embedder = self.build_retrieval_embedder() |
| 15 | self.corpus_embeddings = self.build_corpus_embeddings() |
| 16 | |
| 17 | def build_retrieval_corpus(self): |
| 18 | print("Building corpus...") |
| 19 | documents_df = pd.read_csv(self.corpus_tsv_path, sep='\t') |
| 20 | corpus, corpus2tool = process_retrieval_ducoment(documents_df) |
| 21 | corpus_ids = list(corpus.keys()) |
| 22 | corpus = [corpus[cid] for cid in corpus_ids] |
| 23 | return corpus, corpus2tool |
| 24 | |
| 25 | def build_retrieval_embedder(self): |
| 26 | print("Building embedder...") |
| 27 | embedder = SentenceTransformer(self.model_path) |
| 28 | return embedder |
| 29 | |
| 30 | def build_corpus_embeddings(self): |
| 31 | print("Building corpus embeddings with embedder...") |
| 32 | corpus_embeddings = self.embedder.encode(self.corpus, convert_to_tensor=True) |
| 33 | return corpus_embeddings |
| 34 | |
| 35 | def retrieving(self, query, top_k=5, excluded_tools={}): |
| 36 | print("Retrieving...") |
| 37 | start = time.time() |
| 38 | query_embedding = self.embedder.encode(query, convert_to_tensor=True) |
| 39 | hits = util.semantic_search(query_embedding, self.corpus_embeddings, top_k=10*top_k, score_function=util.cos_sim) |
| 40 | retrieved_tools = [] |
| 41 | for rank, hit in enumerate(hits[0]): |
| 42 | category, tool_name, api_name = self.corpus2tool[self.corpus[hit['corpus_id']]].split('\t') |
| 43 | category = standardize_category(category) |
| 44 | tool_name = standardize(tool_name) # standardizing |
| 45 | api_name = change_name(standardize(api_name)) # standardizing |
| 46 | if category in excluded_tools: |
| 47 | if tool_name in excluded_tools[category]: |
| 48 | top_k += 1 |
| 49 | continue |
| 50 | tmp_dict = { |
| 51 | "category": category, |
| 52 | "tool_name": tool_name, |
| 53 | "api_name": api_name |
| 54 | } |
| 55 | retrieved_tools.append(tmp_dict) |
| 56 | return retrieved_tools |