torch-rechub 0.0.1__py3-none-any.whl → 0.0.3__py3-none-any.whl

torch_rechub/utils/data.py

@@ -0,0 +1,360 @@
+import random
+import torch
+import numpy as np
+import pandas as pd
+import tqdm
+from sklearn.preprocessing import LabelEncoder
+from sklearn.metrics import roc_auc_score, mean_squared_error
+from torch.utils.data import Dataset, DataLoader, random_split
+
+
+class TorchDataset(Dataset):
+
+    def __init__(self, x, y):
+        super().__init__()
+        self.x = x
+        self.y = y
+
+    def __getitem__(self, index):
+        return {k: v[index] for k, v in self.x.items()}, self.y[index]
+
+    def __len__(self):
+        return len(self.y)
+
+
+class PredictDataset(Dataset):
+
+    def __init__(self, x):
+        super().__init__()
+        self.x = x
+
+    def __getitem__(self, index):
+        return {k: v[index] for k, v in self.x.items()}
+
+    def __len__(self):
+        return len(self.x[list(self.x.keys())[0]])
+
+
+class MatchDataGenerator(object):
+
+    def __init__(self, x, y=[]):
+        super().__init__()
+        if len(y) != 0:
+            self.dataset = TorchDataset(x, y)
+        else:  # For pair-wise model, trained without given label
+            self.dataset = PredictDataset(x)
+
+    def generate_dataloader(self, x_test_user, x_all_item, batch_size, num_workers=8):
+        train_dataloader = DataLoader(self.dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+        test_dataset = PredictDataset(x_test_user)
+
+        # shuffle = False to keep same order as ground truth
+        test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+        item_dataset = PredictDataset(x_all_item)
+        item_dataloader = DataLoader(item_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+        return train_dataloader, test_dataloader, item_dataloader
+
+
+class DataGenerator(object):
+
+    def __init__(self, x, y):
+        super().__init__()
+        self.dataset = TorchDataset(x, y)
+        self.length = len(self.dataset)
+
+    def generate_dataloader(self, x_val=None, y_val=None, x_test=None, y_test=None, split_ratio=None, batch_size=16,
+                            num_workers=0):
+        if split_ratio != None:
+            train_length = int(self.length * split_ratio[0])
+            val_length = int(self.length * split_ratio[1])
+            test_length = self.length - train_length - val_length
+            print("the samples of train : val : test are  %d : %d : %d" % (train_length, val_length, test_length))
+            train_dataset, val_dataset, test_dataset = random_split(self.dataset,
+                                                                    (train_length, val_length, test_length))
+        else:
+            train_dataset = self.dataset
+            val_dataset = TorchDataset(x_val, y_val)
+            test_dataset = TorchDataset(x_test, y_test)
+
+        train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=num_workers)
+        val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+        test_dataloader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers)
+        return train_dataloader, val_dataloader, test_dataloader
+
+
+def get_auto_embedding_dim(num_classes):
+    """ Calculate the dim of embedding vector according to number of classes in the category
+    emb_dim = [6 * (num_classes)^(1/4)]
+    reference: Deep & Cross Network for Ad Click Predictions.(ADKDD'17)
+    Args:
+        num_classes: number of classes in the category
+    
+    Returns:
+        the dim of embedding vector
+    """
+    return int(np.floor(6 * np.pow(num_classes, 0.25)))
+
+
+def get_loss_func(task_type="classification"):
+    if task_type == "classification":
+        return torch.nn.BCELoss()
+    elif task_type == "regression":
+        return torch.nn.MSELoss()
+    else:
+        raise ValueError("task_type must be classification or regression")
+
+
+def get_metric_func(task_type="classification"):
+    if task_type == "classification":
+        return roc_auc_score
+    elif task_type == "regression":
+        return mean_squared_error
+    else:
+        raise ValueError("task_type must be classification or regression")
+
+
+def generate_seq_feature(data,
+                         user_col,
+                         item_col,
+                         time_col,
+                         item_attribute_cols=[],
+                         min_item=0,
+                         shuffle=True,
+                         max_len=50):
+    """generate sequence feature and negative sample for ranking.
+
+    Args:
+        data (pd.DataFrame): the raw data.
+        user_col (str): the col name of user_id
+        item_col (str): the col name of item_id
+        time_col (str): the col name of timestamp
+        item_attribute_cols (list[str], optional): the other attribute cols of item which you want to generate sequence feature. Defaults to `[]`.
+        sample_method (int, optional): the negative sample method `{
+            0: "random sampling",
+            1: "popularity sampling method used in word2vec",
+            2: "popularity sampling method by `log(count+1)+1e-6`",
+            3: "tencent RALM sampling"}`.
+            Defaults to 0.
+        min_item (int, optional): the min item each user must have. Defaults to 0.
+        shuffle (bool, optional): shulle if True
+        max_len (int, optional): the max length of a user history sequence.
+
+    Returns:
+        pd.DataFrame: split train, val and test data with sequence features by time.
+    """
+    for feat in data:
+        le = LabelEncoder()
+        data[feat] = le.fit_transform(data[feat])
+        data[feat] = data[feat].apply(lambda x: x + 1)  # 0 to be used as the symbol for padding
+    data = data.astype('int32')
+
+    # generate item to attribute mapping
+    n_items = data[item_col].max()
+    item2attr = {}
+    if len(item_attribute_cols) > 0:
+        for col in item_attribute_cols:
+            map = data[[item_col, col]]
+            item2attr[col] = map.set_index([item_col])[col].to_dict()
+
+    train_data, val_data, test_data = [], [], []
+    data.sort_values(time_col, inplace=True)
+    # Sliding window to construct negative samples
+    for uid, hist in tqdm.tqdm(data.groupby(user_col), desc='generate sequence features'):
+        pos_list = hist[item_col].tolist()
+        len_pos_list = len(pos_list)
+        if len_pos_list < min_item:  # drop this user when his pos items < min_item
+            continue
+
+        neg_list = [neg_sample(pos_list, n_items) for _ in range(len_pos_list)]
+        for i in range(1, min(len_pos_list, max_len)):
+            hist_item = pos_list[:i]
+            hist_item = hist_item + [0] * (max_len - len(hist_item))
+            pos_item = pos_list[i]
+            neg_item = neg_list[i]
+            pos_seq = [1, pos_item, uid, hist_item]
+            neg_seq = [0, neg_item, uid, hist_item]
+            if len(item_attribute_cols) > 0:
+                for attr_col in item_attribute_cols:  # the history of item attribute features
+                    hist_attr = hist[attr_col].tolist()[:i]
+                    hist_attr = hist_attr + [0] * (max_len - len(hist_attr))
+                    pos2attr = [hist_attr, item2attr[attr_col][pos_item]]
+                    neg2attr = [hist_attr, item2attr[attr_col][neg_item]]
+                    pos_seq += pos2attr
+                    neg_seq += neg2attr
+            if i == len_pos_list - 1:
+                test_data.append(pos_seq)
+                test_data.append(neg_seq)
+            elif i == len_pos_list - 2:
+                val_data.append(pos_seq)
+                val_data.append(neg_seq)
+            else:
+                train_data.append(pos_seq)
+                train_data.append(neg_seq)
+
+    col_name = ['label', 'target_item_id', user_col, 'hist_item_id']
+    if len(item_attribute_cols) > 0:
+        for attr_col in item_attribute_cols:  # the history of item attribute features
+            name = ['hist_'+attr_col, 'target_'+attr_col]
+            col_name += name
+
+    # shuffle
+    if shuffle:
+        random.shuffle(train_data)
+        random.shuffle(val_data)
+        random.shuffle(test_data)
+
+    train = pd.DataFrame(train_data, columns=col_name)
+    val = pd.DataFrame(val_data, columns=col_name)
+    test = pd.DataFrame(test_data, columns=col_name)
+
+    return train, val, test
+
+
+def df_to_dict(data):
+    """
+    Convert the DataFrame to a dict type input that the network can accept
+    Args:
+        data (pd.DataFrame): datasets of type DataFrame
+    Returns:
+        The converted dict, which can be used directly into the input network
+    """
+    data_dict = data.to_dict('list')
+    for key in data.keys():
+        data_dict[key] = np.array(data_dict[key])
+    return data_dict
+
+
+def neg_sample(click_hist, item_size):
+    neg = random.randint(1, item_size)
+    while neg in click_hist:
+        neg = random.randint(1, item_size)
+    return neg
+
+
+def pad_sequences(sequences, maxlen=None, dtype='int32', padding='pre', truncating='pre', value=0.):
+    """ Pads sequences (list of list) to the ndarray of same length.
+        This is an equivalent implementation of tf.keras.preprocessing.sequence.pad_sequences
+        reference: https://github.com/huawei-noah/benchmark/tree/main/FuxiCTR/fuxictr
+    
+    Args:
+        sequences (pd.DataFrame): data that needs to pad or truncate
+        maxlen (int): maximum sequence length. Defaults to None.
+        dtype (str, optional): Defaults to 'int32'.
+        padding (str, optional): if len(sequences) less than maxlen, padding style, {'pre', 'post'}. Defaults to 'pre'.
+        truncating (str, optional): if len(sequences) more than maxlen, truncate style, {'pre', 'post'}. Defaults to 'pre'.
+        value (_type_, optional): Defaults to 0..
+
+    Returns:
+        _type_: _description_
+    """
+    
+
+    assert padding in ["pre", "post"], "Invalid padding={}.".format(padding)
+    assert truncating in ["pre", "post"], "Invalid truncating={}.".format(truncating)
+
+    if maxlen is None:
+        maxlen = max(len(x) for x in sequences)
+    arr = np.full((len(sequences), maxlen), value, dtype=dtype)
+    for idx, x in enumerate(sequences):
+        if len(x) == 0:
+            continue  # empty list
+        if truncating == 'pre':
+            trunc = x[-maxlen:]
+        else:
+            trunc = x[:maxlen]
+        trunc = np.asarray(trunc, dtype=dtype)
+
+        if padding == 'pre':
+            arr[idx, -len(trunc):] = trunc
+        else:
+            arr[idx, :len(trunc)] = trunc
+    return arr
+
+
+def array_replace_with_dict(array, dic):
+    """Replace values in NumPy array based on dictionary.
+    Args:
+        array (np.array): a numpy array
+        dic (dict): a map dict
+
+    Returns:
+        np.array: array with replace
+    """
+    # Extract out keys and values
+    k = np.array(list(dic.keys()))
+    v = np.array(list(dic.values()))
+
+    # Get argsort indices
+    idx = k.argsort()
+    return v[idx[np.searchsorted(k, array, sorter=idx)]]
+
+
+# Temporarily reserved for testing purposes(1985312383@qq.com)
+def create_seq_features(data, seq_feature_col=['item_id', 'cate_id'], max_len=50, drop_short=3, shuffle=True):
+    """Build a sequence of user's history by time.
+
+    Args:
+        data (pd.DataFrame): must contain keys: `user_id, item_id, cate_id, time`.
+        seq_feature_col (list): specify the column name that needs to generate sequence features, and its sequence features will be generated according to userid.
+        max_len (int): the max length of a user history sequence.
+        drop_short (int): remove some inactive user who's sequence length < drop_short.
+        shuffle (bool): shuffle data if true.
+
+    Returns:
+        train (pd.DataFrame): target item will be each item before last two items.
+        val (pd.DataFrame): target item is the second to last item of user's history sequence.
+        test (pd.DataFrame): target item is the last item of user's history sequence.
+    """
+    for feat in data:
+        le = LabelEncoder()
+        data[feat] = le.fit_transform(data[feat])
+        data[feat] = data[feat].apply(lambda x: x + 1)  # 0 to be used as the symbol for padding
+    data = data.astype('int32')
+
+    n_items = data["item_id"].max()
+
+    item_cate_map = data[['item_id', 'cate_id']]
+    item2cate_dict = item_cate_map.set_index(['item_id'])['cate_id'].to_dict()
+
+    data = data.sort_values(['user_id', 'time']).groupby('user_id').agg(click_hist_list=('item_id', list), cate_hist_hist=('cate_id', list)).reset_index()
+
+    # Sliding window to construct negative samples
+    train_data, val_data, test_data = [], [], []
+    for item in data.itertuples():
+        if len(item[2]) < drop_short:
+            continue
+        user_id = item[1]
+        click_hist_list = item[2][:max_len]
+        cate_hist_list = item[3][:max_len]
+
+        neg_list = [neg_sample(click_hist_list, n_items) for _ in range(len(click_hist_list))]
+        hist_list = []
+        cate_list = []
+        for i in range(1, len(click_hist_list)):
+            hist_list.append(click_hist_list[i - 1])
+            cate_list.append(cate_hist_list[i - 1])
+            hist_list_pad = hist_list + [0] * (max_len - len(hist_list))
+            cate_list_pad = cate_list + [0] * (max_len - len(cate_list))
+            if i == len(click_hist_list) - 1:
+                test_data.append([user_id, hist_list_pad, cate_list_pad, click_hist_list[i], cate_hist_list[i], 1])
+                test_data.append([user_id, hist_list_pad, cate_list_pad, neg_list[i], item2cate_dict[neg_list[i]], 0])
+            if i == len(click_hist_list) - 2:
+                val_data.append([user_id, hist_list_pad, cate_list_pad, click_hist_list[i], cate_hist_list[i], 1])
+                val_data.append([user_id, hist_list_pad, cate_list_pad, neg_list[i], item2cate_dict[neg_list[i]], 0])
+            else:
+                train_data.append([user_id, hist_list_pad, cate_list_pad, click_hist_list[i], cate_hist_list[i], 1])
+                train_data.append([user_id, hist_list_pad, cate_list_pad, neg_list[i], item2cate_dict[neg_list[i]], 0])
+
+    # shuffle
+    if shuffle:
+        random.shuffle(train_data)
+        random.shuffle(val_data)
+        random.shuffle(test_data)
+
+    col_name = ['user_id', 'history_item', 'history_cate', 'target_item', 'target_cate', 'label']
+    train = pd.DataFrame(train_data, columns=col_name)
+    val = pd.DataFrame(val_data, columns=col_name)
+    test = pd.DataFrame(test_data, columns=col_name)
+
+    return train, val, test
+

torch_rechub/utils/match.py

@@ -0,0 +1,274 @@
+import tqdm
+import pandas as pd
+import numpy as np
+import copy
+import random
+from collections import OrderedDict, Counter
+from annoy import AnnoyIndex
+from .data import pad_sequences, df_to_dict
+from pymilvus import Collection,CollectionSchema,DataType,FieldSchema,connections,utility
+
+def gen_model_input(df, user_profile, user_col, item_profile, item_col, seq_max_len, padding='pre', truncating='pre'):
+    """Merge user_profile and item_profile to df, pad and truncate history sequence feature
+
+    Args:
+        df (pd.DataFrame): data with history sequence feature
+        user_profile (pd.DataFrame): user data
+        user_col (str): user column name
+        item_profile (pd.DataFrame): item data
+        item_col (str): item column name
+        seq_max_len (int): sequence length of every data
+        padding (str, optional): padding style, {'pre', 'post'}. Defaults to 'pre'.
+        truncating (str, optional): truncate style, {'pre', 'post'}. Defaults to 'pre'.
+
+    Returns:
+        dict: The converted dict, which can be used directly into the input network
+    """
+    df = pd.merge(df, user_profile, on=user_col, how='left')  # how=left to keep samples order same as the input
+    df = pd.merge(df, item_profile, on=item_col, how='left')
+    for col in df.columns.to_list():
+        if col.startswith("hist_"):
+            df[col] = pad_sequences(df[col], maxlen=seq_max_len, value=0, padding=padding, truncating=truncating).tolist()
+    for col in df.columns.to_list():
+        if col.startswith("tag_"):
+            df[col] = pad_sequences(df[col], maxlen=seq_max_len, value=0, padding=padding, truncating=truncating).tolist()
+
+    input_dict = df_to_dict(df)
+    return input_dict
+
+
+def negative_sample(items_cnt_order, ratio, method_id=0):
+    """Negative Sample method for matching model
+    reference: https://github.com/wangzhegeek/DSSM-Lookalike/blob/master/utils.py
+    update more method and redesign this function.
+
+    Args:
+        items_cnt_order (dict): the item count dict, the keys(item) sorted by value(count) in reverse order.
+        ratio (int): negative sample ratio, >= 1
+        method_id (int, optional): 
+        `{
+            0: "random sampling", 
+            1: "popularity sampling method used in word2vec", 
+            2: "popularity sampling method by `log(count+1)+1e-6`",
+            3: "tencent RALM sampling"}`. 
+            Defaults to 0.
+            
+    Returns:
+        list: sampled negative item list
+    """
+    items_set = [item for item, count in items_cnt_order.items()]
+    if method_id == 0:
+        neg_items = np.random.choice(items_set, size=ratio, replace=True)
+    elif method_id == 1:
+        #items_cnt_freq = {item: count/len(items_cnt) for item, count in items_cnt_order.items()}
+        #p_sel = {item: np.sqrt(1e-5/items_cnt_freq[item]) for item in items_cnt_order}
+        #The most popular paramter is item_cnt**0.75:
+        p_sel = {item: count**0.75 for item, count in items_cnt_order.items()}
+        p_value = np.array(list(p_sel.values())) / sum(p_sel.values())
+        neg_items = np.random.choice(items_set, size=ratio, replace=True, p=p_value)
+    elif method_id == 2:
+        p_sel = {item: np.log(count + 1) + 1e-6 for item, count in items_cnt_order.items()}
+        p_value = np.array(list(p_sel.values())) / sum(p_sel.values())
+        neg_items = np.random.choice(items_set, size=ratio, replace=True, p=p_value)
+    elif method_id == 3:
+        p_sel = {item: (np.log(k + 2) - np.log(k + 1)) / np.log(len(items_cnt_order) + 1) for item, k in items_cnt_order.items()}
+        p_value = np.array(list(p_sel.values())) / sum(p_sel.values())
+        neg_items = np.random.choice(items_set, size=ratio, replace=False, p=p_value)
+    else:
+        raise ValueError("method id should in (0,1,2,3)")
+    return neg_items
+
+
+def generate_seq_feature_match(data,
+                               user_col,
+                               item_col,
+                               time_col,
+                               item_attribute_cols=None,
+                               sample_method=0,
+                               mode=0,
+                               neg_ratio=0,
+                               min_item=0):
+    """generate sequence feature and negative sample for match.
+
+    Args:
+        data (pd.DataFrame): the raw data.
+        user_col (str): the col name of user_id 
+        item_col (str): the col name of item_id 
+        time_col (str): the col name of timestamp
+        item_attribute_cols (list[str], optional): the other attribute cols of item which you want to generate sequence feature. Defaults to `[]`.
+        sample_method (int, optional): the negative sample method `{
+            0: "random sampling", 
+            1: "popularity sampling method used in word2vec", 
+            2: "popularity sampling method by `log(count+1)+1e-6`",
+            3: "tencent RALM sampling"}`. 
+            Defaults to 0.
+        mode (int, optional): the training mode, `{0:point-wise, 1:pair-wise, 2:list-wise}`. Defaults to 0.
+        neg_ratio (int, optional): negative sample ratio, >= 1. Defaults to 0.
+        min_item (int, optional): the min item each user must have. Defaults to 0.
+
+    Returns:
+        pd.DataFrame: split train and test data with sequence features.
+    """
+    if item_attribute_cols is None:
+        item_attribute_cols = []
+    if mode == 2:  # list wise learning
+        assert neg_ratio > 0, 'neg_ratio must be greater than 0 when list-wise learning'
+    elif mode == 1:  # pair wise learning
+        neg_ratio = 1
+    print("preprocess data")
+    data.sort_values(time_col, inplace=True)  #sort by time from old to new
+    train_set, test_set = [], []
+    n_cold_user = 0
+
+    items_cnt = Counter(data[item_col].tolist())
+    items_cnt_order = OrderedDict(sorted((items_cnt.items()), key=lambda x: x[1], reverse=True))  #item_id:item count
+    neg_list = negative_sample(items_cnt_order, ratio=data.shape[0] * neg_ratio, method_id=sample_method)
+    neg_idx = 0
+    for uid, hist in tqdm.tqdm(data.groupby(user_col), desc='generate sequence features'):
+        pos_list = hist[item_col].tolist()
+        if len(pos_list) < min_item:  #drop this user when his pos items < min_item
+            n_cold_user += 1
+            continue
+
+        for i in range(1, len(pos_list)):
+            hist_item = pos_list[:i]
+            sample = [uid, pos_list[i], hist_item, len(hist_item)]
+            if len(item_attribute_cols) > 0:
+                for attr_col in item_attribute_cols:  #the history of item attribute features
+                    sample.append(hist[attr_col].tolist()[:i])
+            if i != len(pos_list) - 1:
+                if mode == 0:  #point-wise, the last col is label_col, include label 0 and 1
+                    last_col = "label"
+                    train_set.append(sample + [1])
+                    for _ in range(neg_ratio):
+                        sample[1] = neg_list[neg_idx]
+                        neg_idx += 1
+                        train_set.append(sample + [0])
+                elif mode == 1:  #pair-wise, the last col is neg_col, include one negative item
+                    last_col = "neg_items"
+                    for _ in range(neg_ratio):
+                        sample_copy = copy.deepcopy(sample)
+                        sample_copy.append(neg_list[neg_idx])
+                        neg_idx += 1
+                        train_set.append(sample_copy)
+                elif mode == 2:  #list-wise, the last col is neg_col, include neg_ratio negative items
+                    last_col = "neg_items"
+                    sample.append(neg_list[neg_idx: neg_idx + neg_ratio])
+                    neg_idx += neg_ratio
+                    train_set.append(sample)
+                else:
+                    raise ValueError("mode should in (0,1,2)")
+            else:
+                test_set.append(sample + [1])  #Note: if mode=1 or 2, the label col is useless.
+
+    random.shuffle(train_set)
+    random.shuffle(test_set)
+
+    print("n_train: %d, n_test: %d" % (len(train_set), len(test_set)))
+    print("%d cold start user dropped " % n_cold_user)
+
+    attr_hist_col = ["hist_" + col for col in item_attribute_cols]
+    df_train = pd.DataFrame(train_set,
+                            columns=[user_col, item_col, "hist_" + item_col, "histlen_" + item_col] + attr_hist_col + [last_col])
+    df_test = pd.DataFrame(test_set,
+                           columns=[user_col, item_col, "hist_" + item_col, "histlen_" + item_col] + attr_hist_col + [last_col])
+
+    return df_train, df_test
+
+
+class Annoy(object):
+    """Vector matching by Annoy
+
+    Args:
+        metric (str): distance metric
+        n_trees (int): n_trees
+        search_k (int): search_k
+    """
+
+    def __init__(self, metric='angular', n_trees=10, search_k=-1):
+        self._n_trees = n_trees
+        self._search_k = search_k
+        self._metric = metric
+
+    def fit(self, X):
+        self._annoy = AnnoyIndex(X.shape[1], metric=self._metric)
+        for i, x in enumerate(X):
+            self._annoy.add_item(i, x.tolist())
+        self._annoy.build(self._n_trees)
+
+    def set_query_arguments(self, search_k):
+        self._search_k = search_k
+
+    def query(self, v, n):
+        return self._annoy.get_nns_by_vector(v.tolist(), n, self._search_k, include_distances=True)  #
+
+    def __str__(self):
+        return 'Annoy(n_trees=%d, search_k=%d)' % (self._n_trees, self._search_k)
+
+    
+class Milvus(object):
+    """Vector matching by Milvus.
+
+    Args:
+        dim (int): embedding dim
+        host (str): host address of Milvus
+        port (str): port of Milvus
+    """
+
+    def __init__(self, dim=64, host="localhost", port="19530"):
+        print("Start connecting to Milvus")
+        connections.connect("default", host=host, port=port)
+        self.dim = dim
+        has = utility.has_collection("rechub")
+        #print(f"Does collection rechub exist? {has}")
+        if has:
+            utility.drop_collection("rechub")
+        # Create collection
+        fields = [
+            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True),
+            FieldSchema(name="embeddings", dtype=DataType.FLOAT_VECTOR, dim=dim),
+        ]
+        schema = CollectionSchema(fields=fields)
+        self.milvus = Collection("rechub", schema=schema)
+
+    def fit(self, X):
+        if torch.is_tensor(X):
+            X = X.cpu().numpy()
+        self.milvus.release()
+        entities = [[i for i in range(len(X))], X]
+        self.milvus.insert(entities)
+        print(
+            f"Number of entities in Milvus: {self.milvus.num_entities}"
+        )  # check the num_entites
+
+        index = {
+            "index_type": "IVF_FLAT",
+            "metric_type": "L2",
+            "params": {"nlist": 128},
+        }
+        self.milvus.create_index("embeddings", index)
+
+    @staticmethod
+    def process_result(results):
+        idx_list = []
+        score_list = []
+        for r in results:
+            temp_idx_list = []
+            temp_score_list = []
+            for i in range(len(r)):
+                temp_idx_list.append(r[i].id)
+                temp_score_list.append(r[i].distance)
+            idx_list.append(temp_idx_list)
+            score_list.append(temp_score_list)
+        return idx_list, score_list
+
+    def query(self, v, n):
+        if torch.is_tensor(v):
+            v = v.cpu().numpy().reshape(-1, self.dim)
+        self.milvus.load()
+        search_params = {"metric_type": "L2", "params": {"nprobe": 16}}
+        results = self.milvus.search(v, "embeddings", search_params, n)
+        return self.process_result(results)
+
+#annoy = Annoy(n_trees=10)
+#annoy.fit(item_embs)