# 评估 k 最近邻语言模型

11from typing import Optional, List
12
13import faiss
14import numpy as np
15import torch
16
17from labml import monit, lab
18from labml.logger import inspect
19from labml_nn.transformers.knn.train_model import Configs

## -NN 要得到

22def knn(queries: torch.Tensor, index: faiss.IndexFlatL2, keys_store: np.ndarray, vals_store: np.ndarray, n_tokens: int):

31    queries_shape = queries.shape

34    queries = queries.view(-1, queries_shape[-1])

38    distance, idx = index.search(queries.numpy(), 10)

41    keys_found = queries.new_tensor(keys_store[idx])

43    vals_found = torch.tensor(vals_store[idx]).squeeze(-1)

48    keys_found_n = keys_found / torch.sqrt((keys_found ** 2).sum(-1, keepdims=True) + 1e-10)

50    queries_n = queries / torch.sqrt((queries ** 2).sum(-1, keepdims=True) + 1e-10)

53    dot_prod = (keys_found_n * queries_n.unsqueeze(1)).sum(-1)

56    logits_token = dot_prod.new_zeros(queries.shape[0], n_tokens)

58    _ = logits_token.scatter_(dim=1, index=vals_found, src=dot_prod, reduce='add')

61    logits_token = logits_token.reshape(queries_shape[0], queries_shape[1], -1)
62
63    return logits_token

## 计算验证损失

66def validation_loss(knn_weights: List[float], last_n: Optional[int], conf: Configs, index: faiss.IndexFlatL2,
67                    keys_store: np.ndarray, vals_store: np.ndarray):

77    losses = [[] for _ in knn_weights]

79    n_samples = []
80    with torch.no_grad():

82        for i, batch in monit.enum("Validation", conf.validator.data_loader, is_children_silent=True):

84            data, target = batch[0].to(conf.device), batch[1].to(conf.device)

86            res = conf.model(data)

88            res_knn = knn(conf.model.ff_input.cpu(), index, keys_store, vals_store, conf.n_tokens)
89            res_knn = res_knn.to(conf.device)

94            if last_n:
95                res = res[-last_n:]
96                res_knn = res_knn[-last_n:]
97                target = target[-last_n:]

100            n_s = res.shape[0] * data.shape[1]
101            n_samples.append(n_s)

104            for i, c in enumerate(knn_weights):

106                loss = conf.loss_func(res_knn * c + (1 - c) * res, target)
107                losses[i].append(loss * n_s)
108
109    return losses, n_samples

## 加载索引

112def load_index(conf: Configs, n_probe: int = 8):

117    d_model = conf.transformer.d_model

119    data_loader = conf.trainer.data_loader

122    n_keys = data_loader.data.shape[0] * data_loader.data.shape[1] - 1

125    with monit.section('Load index'):
126        index = faiss.read_index(str(lab.get_data_path() / 'faiss.index'))

128    index.nprobe = n_probe

131    keys_store = np.memmap(str(lab.get_data_path() / 'keys.npy'), dtype=np.float32, mode='r', shape=(n_keys, d_model))
132    vals_store = np.memmap(str(lab.get_data_path() / 'vals.npy'), dtype=np.int, mode='r', shape=(n_keys, 1))
133
134    return index, keys_store, vals_store
137def main():
138    from labml_nn.transformers.knn.build_index import load_experiment

141    conf = load_experiment('4984b85c20bf11eb877a69c1a03717cd')

143    conf.model.eval()

146    index, keys_store, vals_store = load_index(conf)

-NN 预测赋予的权重列表。我们将评估每个权重的验证损失

149    knn_weights = [i / 20 for i in range(10)]

151    losses, n_samples = validation_loss(knn_weights, None, conf, index, keys_store, vals_store)

153    inspect({c: np.sum(losses[i]) / np.sum(n_samples) for i, c in enumerate(knn_weights)})
154
155
156if __name__ == '__main__':
157    main()