Delete retriever

retriever/retriever.py

@@ -1,129 +0,0 @@
-from transformers import AutoTokenizer, AutoModel
-from sklearn.metrics.pairwise import cosine_similarity
-import json
-import torch
-from tqdm import tqdm
-import os
-import pandas as pd
-import numpy as np
-from datasets import load_dataset
-from utils.utils import read_yaml_file
-
-
-def generate_topic_level_embeddings(model, tokenizer, paper_list, tmp_id_2_abs):
-    id2topics = {
-        entry["paper_id"]: [entry["Level 1"], entry["Level 2"], entry["Level 3"]]
-        for entry in tmp_id_2_abs['train']
-    }
-
-    for topic_level in ['Level 1', 'Level 2', 'Level 3']:
-        i = 0
-        batch_size = 2048
-        candidate_emb_list = []
-        pbar = tqdm(total=len(paper_list))
-        while i < len(paper_list):
-            yield i / len(paper_list) / 3 if topic_level == 'Level 1' else 0.33 + i / len(paper_list) / 3 if topic_level == 'Level 2' else 0.66 + i / len(paper_list) / 3
-            paper_batch = paper_list[i:i+batch_size]
-            paper_text_batch = []
-            for paper_id in paper_batch:
-                topics = id2topics[paper_id][int(topic_level[6])-1]
-                topic_text = ''
-                for t in topics:
-                    topic_text += t + ','
-                paper_text_batch.append(topic_text)
-            inputs = tokenizer(paper_text_batch, return_tensors='pt', padding=True, truncation=True)
-            with torch.no_grad():
-                outputs = model(**inputs.to('cuda'))
-                candidate_embeddings = outputs.last_hidden_state[:, 0, :].cpu()
-                candidate_embeddings = candidate_embeddings.reshape(-1, 1024)
-                candidate_emb_list.append(candidate_embeddings)
-
-                i += len(candidate_embeddings)
-                pbar.update(len(candidate_embeddings))
-
-        all_candidate_embs = torch.cat(candidate_emb_list, 0)
-        
-        df = pd.DataFrame({
-            "paper_id": paper_list,
-            "embedding": list(all_candidate_embs.numpy())
-        })
-        
-        if not os.path.exists('datasets/topic_level_embeds'):
-            os.makedirs('datasets/topic_level_embeds')
-
-        df.to_parquet(f'datasets/topic_level_embeds/{topic_level}_emb.parquet', engine='pyarrow', compression='snappy')
-        
-    all_candidate_embs_L1 = torch.tensor(np.array(pd.read_parquet('datasets/topic_level_embeds/Level 1_emb.parquet')['embedding'].tolist()))
-    all_candidate_embs_L2 = torch.tensor(np.array(pd.read_parquet('datasets/topic_level_embeds/Level 2_emb.parquet')['embedding'].tolist()))
-    all_candidate_embs_L3 = torch.tensor(np.array(pd.read_parquet('datasets/topic_level_embeds/Level 3_emb.parquet')['embedding'].tolist()))
-    all_candidate_embs = all_candidate_embs_L1 + all_candidate_embs_L2 + all_candidate_embs_L3
-    
-    df = pd.DataFrame({
-        "paper_id": paper_list,
-        "embedding": list(all_candidate_embs.numpy())
-    })
-    
-    df.to_parquet('datasets/topic_level_embeds/arxiv_papers_embeds.parquet', engine='pyarrow', compression='snappy')
-
-
-
-def retriever(query, retrieval_nodes_path):    
-    yield 0
-    config = read_yaml_file('configs/config.yaml')
-
-    # Load the model and tokenizer to generate the embeddings
-    embedder_name = config['retriever']['embedder']
-    tokenizer = AutoTokenizer.from_pretrained(embedder_name)
-    model = AutoModel.from_pretrained(embedder_name).to(device='cuda', dtype=torch.float16)
-
-
-    # Load the arXiv dataset
-    tmp_id_2_abs = load_dataset("AliMaatouk/arXiv_Topics", cache_dir="datasets/arxiv_topics")
-    paper_list = list(tmp_id_2_abs['train']['paper_id'])
-
-
-    # Generate the query embeddings
-    inputs = tokenizer([query], return_tensors='pt', padding=True, truncation=True)
-    with torch.no_grad():
-        outputs = model(**inputs.to('cuda'))
-        query_embeddings = outputs.last_hidden_state[:, 0, :].cpu()
-
-    # Generate the candidate embeddings
-    # Load the embeddings from the dataset, otherwise generate the embeddings and save them
-    if config['retriever']['load_arxiv_embeds']:
-        dataset = load_dataset("AliMaatouk/arXiv-Topics-Embeddings", cache_dir="datasets/topic_level_embeds")
-        table = dataset["train"].data  # Get PyArrow Table
-        all_candidate_embs = table.column("embedding").to_numpy()
-    else:        
-        # If the file does not exist, generate the embeddings, otherwise, load the embeddings
-        if not os.path.exists('datasets/topic_level_embeds/arxiv_papers_embeds.parquet'):
-            yield from generate_topic_level_embeddings(model, tokenizer, paper_list, tmp_id_2_abs)
-        
-        all_candidate_embs = torch.tensor(np.array(pd.read_parquet('datasets/topic_level_embeds/arxiv_papers_embeds.parquet')['embedding'].tolist()))
-        all_candidate_embs = all_candidate_embs.cpu().numpy()
-
-    all_candidate_embs = np.stack(all_candidate_embs)
-
-
-    # Calculate the cosine similarity between the query and all candidate embeddings
-    query_embeddings = np.array(query_embeddings)
-    similarity_scores = cosine_similarity(query_embeddings, all_candidate_embs)[0]
-
-
-    # Sort the papers by similarity scores and select the top K papers
-    id_score_list = []
-    for i in range(len(paper_list)):
-        id_score_list.append([paper_list[i], similarity_scores[i]])
-
-    sorted_scores = sorted(id_score_list, key=lambda i: i[-1], reverse = True)
-    top_K_paper = [sample[0] for sample in sorted_scores[:config['retriever']['num_retrievals']]]
-
-    papers_results = {
-        paper: True
-        for paper in top_K_paper
-    }
-
-    with open(retrieval_nodes_path, 'w') as f:
-        json.dump(papers_results, f)
-
-    yield 1.0