| import os
|
| import collections
|
| import json
|
| import logging
|
| import argparse
|
| import numpy as np
|
| import pandas as pd
|
| import torch
|
| from time import time
|
| from torch import optim
|
| from tqdm import tqdm
|
| from torch.utils.data import DataLoader
|
|
|
| from rq_llama import *
|
| from index.datasets import EmbDataset
|
|
|
| def if_collided(all_indices_str):
|
| tot_item = len(all_indices_str)
|
| tot_indice = len(set(all_indices_str.tolist()))
|
| return tot_item == tot_indice
|
|
|
| def get_indices_count(all_indices_str):
|
| indices_count = collections.defaultdict(int)
|
| for index in all_indices_str:
|
| indices_count[index] += 1
|
| return indices_count
|
|
|
| def get_collision_item(all_indices_str):
|
| index2id = {}
|
| for i, index in enumerate(all_indices_str):
|
| if index not in index2id:
|
| index2id[index] = []
|
| index2id[index].append(i)
|
| collision_item_groups = []
|
| for index in index2id:
|
| if len(index2id[index]) > 1:
|
| collision_item_groups.append(index2id[index])
|
| return collision_item_groups
|
|
|
| def parse_args():
|
| parser = argparse.ArgumentParser(description = "Index")
|
| parser.add_argument("--ckpt_path", type = str, default = "", help = "")
|
| parser.add_argument("--data_path", type = str, default = "", help = "")
|
| parser.add_argument("--save_path", type = str, default = "", help = "")
|
| parser.add_argument("--device_map", type = str, default = "1", help = "gpu or cpu")
|
| return parser.parse_args()
|
|
|
| args = parse_args()
|
| print(args)
|
|
|
| data = EmbDataset(args.data_path)
|
| data_loader = DataLoader(data, num_workers = 4, batch_size = 64, shuffle = False, pin_memory = True)
|
| device_map = {'': int(args.device_map)}
|
| MODEL = LlamaWithRQ.from_pretrained(args.ckpt_path, torch_dtype = torch.float16, low_cpu_mem_usage = True, device_map = device_map)
|
| MODEL.eval()
|
| device = MODEL.device
|
| rqvae = MODEL.rqvae
|
| prefix = MODEL.prefix
|
|
|
| all_indices = []
|
| all_indices_str = []
|
| with torch.no_grad():
|
| for x in tqdm(data_loader):
|
| indices = rqvae.get_indices(x.to(device), False)
|
| indices = indices.view(-1, indices.shape[-1]).cpu().numpy()
|
| for index in indices:
|
| code = []
|
| for i, ind in enumerate(index):
|
| code.append(prefix[i].format(int(ind)))
|
|
|
| all_indices.append(code)
|
| all_indices_str.append(str(code))
|
|
|
| all_indices = np.array(all_indices)
|
| all_indices_str = np.array(all_indices_str)
|
|
|
| for vq in rqvae.rq.vq_layers[:-1]:
|
| vq.sk_epsilon=0.0
|
| if rqvae.rq.vq_layers[-1].sk_epsilon == 0.0:
|
| rqvae.rq.vq_layers[-1].sk_epsilon = 0.003
|
|
|
| tt = 0
|
| while True:
|
| if tt >= 20 or if_collided(all_indices_str):
|
| break
|
|
|
| collision_item_groups = get_collision_item(all_indices_str)
|
|
|
| print(len(collision_item_groups))
|
| with torch.no_grad():
|
| for collision_items in collision_item_groups:
|
| indices = rqvae.get_indices(data[collision_items].to(device), True)
|
| indices = indices.view(-1, indices.shape[-1]).cpu().numpy()
|
| for item, index in zip(collision_items, indices):
|
| code = []
|
| for i, ind in enumerate(index):
|
| code.append(prefix[i].format(int(ind)))
|
|
|
| all_indices[item] = code
|
| all_indices_str[item] = str(code)
|
| tt += 1
|
|
|
| print("All indices number: ",len(all_indices))
|
| print("Max number of conflicts: ", max(get_indices_count(all_indices_str).values()))
|
|
|
| tot_item = len(all_indices_str)
|
| tot_indice = len(set(all_indices_str.tolist()))
|
| print("Collision Rate",(tot_item - tot_indice) / tot_item)
|
|
|
| all_indices_dict = {}
|
| for item, indices in enumerate(all_indices.tolist()):
|
| all_indices_dict[item] = list(indices)
|
|
|
| with open(args.save_path, 'w',encoding = 'utf-8') as f:
|
| json.dump(all_indices_dict, f) |
