Three implementations for https://nlp100.github.io/2025/ja/ch06.html#id3

similar_word.py

# -*- python -*-
# Three implementations for https://nlp100.github.io/2025/ja/ch06.html#id3

from gensim.models import KeyedVectors
from tap import Tap
import gdown
import os
import numpy as np
import torch

class Args(Tap):
    modelid: str = "0B7XkCwpI5KDYNlNUTTlSS21pQmM"
    modelfile: str = "GoogleNews-vectors-negative300.bin.gz"
    target: str = "United_States"
    topn: int = 10
    usetorch: bool = False
    usenumpy: bool = False

def load_model(args):
    if not os.path.exists(args.modelfile):
        gdown.download(id=args.modelid, output=args.modelfile, quiet=False)
    return KeyedVectors.load_word2vec_format(args.modelfile, binary=True)

def most_similar_torch(model, target, topn=10):
    device = torch.device("cuda")
    vectors = torch.tensor(model.vectors, device=device)
    vectors = torch.nn.functional.normalize(vectors, dim=1)

    query = torch.tensor(model[target], device=device)
    query = torch.nn.functional.normalize(query, dim=0)

    sims = vectors @ query
    top_scores, top_indices = torch.topk(sims, topn + 1)

    results = [
        (model.index_to_key[idx.item()], top_scores[i].item())
        for i, idx in enumerate(top_indices)
        if model.index_to_key[idx.item()] != target
    ]
    return results[:topn]

def most_similar_numpy(model, target, topn=10):
    vectors = model.vectors
    vectors = vectors / np.linalg.norm(vectors, axis=1, keepdims=True)

    query = model[target]
    query = query / np.linalg.norm(query)

    sims = vectors @ query
    top_indices = np.argpartition(sims, -(topn + 1))[-(topn + 1):]
    top_indices = top_indices[np.argsort(sims[top_indices])[::-1]]

    results = [
        (model.index_to_key[idx], float(sims[idx]))
        for idx in top_indices
        if model.index_to_key[idx] != target
    ]
    return results[:topn]

def main(args):
    model = load_model(args)
    if args.usetorch:
        similar_words = most_similar_torch(model, args.target, topn=args.topn)
    elif args.usenumpy:
        similar_words = most_similar_numpy(model, args.target, topn=args.topn)
    else:
        similar_words = model.similar_by_word(args.target, topn=args.topn)
    for w, sim in similar_words:
        print(f"{w:20} : {sim:.4f}")

if __name__ == "__main__":
    args = Args().parse_args()
    main(args)