Small corpus → weak on analogies (esp. capital-country, morphology). For better analogy
accuracy, train on a larger corpus (enwik9+). Lowercased English only; drops OOV.
1from huggingface_hub import hf_hub_download
2import torch
3import torch.nn.functional as F
4
5path = hf_hub_download(repo_id="ocdbytes/embeddings", filename="embeddings_200.pt")
6# weights_only=False because the checkpoint bundles Python dicts (word2idx/idx2word),
7# which the default restricted loader (torch>=2.6) may reject.
8ck = torch.load(path, map_location="cpu", weights_only=False)
9
10syn0 = ck["syn0"]
11word2idx, idx2word = ck["word2idx"], ck["idx2word"]
12emb = F.normalize(syn0, dim=1)
13
14def neighbours(word, n=10):
15 i = word2idx[word]
16 sims = emb @ emb[i]
17 top = sims.topk(n + 1).indices.tolist()
18 return [idx2word[j] for j in top if j != i][:n]
19
20def analogy(a, b, c, n=5):
21 t = F.normalize(emb[word2idx[b]] - emb[word2idx[a]] + emb[word2idx[c]], dim=0)
22 sims = emb @ t
23 ban = {word2idx[a], word2idx[b], word2idx[c]}
24 top = sims.topk(n + len(ban)).indices.tolist()
25 return [idx2word[j] for j in top if j not in ban][:n]
26
27print(neighbours("king")) # -> ['viii', 'elizabeth', 'queen', ...]
28print(analogy("france", "paris", "germany")) # -> ['berlin', ...]
