S104 · Embeddings — word2vec, GloVe, Contrastive Learning

# mini_w2v.py — a from-scratch skip-gram with negative sampling.
# Usage: uv run mini_w2v.py path/to/corpus.txt
import sys
import random
import re
from collections import Counter
 
import torch
import torch.nn as nn
import torch.nn.functional as F
 
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
EMB_DIM = 100
WINDOW = 5
NEG_SAMPLES = 5
EPOCHS = 5
BATCH = 512
LR = 5e-3
MIN_COUNT = 5
 
 
def tokenise(text: str) -> list[str]:
    return re.findall(r"[a-z']+", text.lower())
 
 
def build_vocab(tokens: list[str], min_count: int) -> tuple[dict[str, int], list[str], list[float]]:
    counts = Counter(tokens)
    vocab = [w for w, c in counts.most_common() if c >= min_count]
    stoi = {w: i for i, w in enumerate(vocab)}
    # Unigram distribution raised to 0.75 (Mikolov's trick) for negative sampling.
    freqs = torch.tensor([counts[w] for w in vocab], dtype=torch.float)
    weights = (freqs ** 0.75)
    weights /= weights.sum()
    return stoi, vocab, weights.tolist()
 
 
def make_pairs(tokens: list[str], stoi: dict[str, int], window: int) -> list[tuple[int, int]]:
    ids = [stoi[t] for t in tokens if t in stoi]
    pairs = []
    for i, center in enumerate(ids):
        w = random.randint(1, window)
        for j in range(max(0, i - w), min(len(ids), i + w + 1)):
            if i != j:
                pairs.append((center, ids[j]))
    return pairs
 
 
class SkipGram(nn.Module):
    def __init__(self, vocab_size: int, dim: int):
        super().__init__()
        self.center = nn.Embedding(vocab_size, dim)
        self.context = nn.Embedding(vocab_size, dim)
        nn.init.uniform_(self.center.weight, -0.5 / dim, 0.5 / dim)
        nn.init.zeros_(self.context.weight)
 
    def forward(
        self, center: torch.Tensor, pos: torch.Tensor, neg: torch.Tensor
    ) -> torch.Tensor:
        v_c = self.center(center)         # [B, D]
        v_pos = self.context(pos)         # [B, D]
        v_neg = self.context(neg)         # [B, K, D]
        pos_score = (v_c * v_pos).sum(dim=1)              # [B]
        neg_score = torch.bmm(v_neg, v_c.unsqueeze(2)).squeeze(2)  # [B, K]
        loss = -F.logsigmoid(pos_score).mean() - F.logsigmoid(-neg_score).mean()
        return loss
 
 
def train(path: str) -> None:
    text = open(path, "r", encoding="utf-8").read()
    tokens = tokenise(text)
    stoi, vocab, neg_weights = build_vocab(tokens, MIN_COUNT)
    print(f"vocab={len(vocab):,}  tokens={len(tokens):,}")
    pairs = make_pairs(tokens, stoi, WINDOW)
    print(f"training pairs: {len(pairs):,}")
 
    model = SkipGram(len(vocab), EMB_DIM).to(DEVICE)
    opt = torch.optim.Adam(model.parameters(), lr=LR)
    weights_t = torch.tensor(neg_weights)
 
    for epoch in range(1, EPOCHS + 1):
        random.shuffle(pairs)
        total = 0.0
        steps = 0
        for i in range(0, len(pairs), BATCH):
            batch = pairs[i : i + BATCH]
            if len(batch) < BATCH:
                continue
            c, p = zip(*batch)
            c = torch.tensor(c, device=DEVICE)
            p = torch.tensor(p, device=DEVICE)
            neg = torch.multinomial(weights_t, BATCH * NEG_SAMPLES, replacement=True)
            neg = neg.view(BATCH, NEG_SAMPLES).to(DEVICE)
            opt.zero_grad()
            loss = model(c, p, neg)
            loss.backward()
            opt.step()
            total += loss.item()
            steps += 1
        print(f"epoch {epoch}  loss={total/steps:.4f}")
 
    # Save + demo queries.
    emb = F.normalize(model.center.weight.detach(), dim=1)  # normalise for cosine
    itos = {i: w for w, i in stoi.items()}
    for word in ["king", "computer", "love", "war"]:
        if word not in stoi:
            continue
        v = emb[stoi[word]]
        sims = emb @ v
        top = sims.topk(6).indices.tolist()
        print(f"nearest to {word!r}: {[itos[i] for i in top if i != stoi[word]][:5]}")
 
 
if __name__ == "__main__":
    train(sys.argv[1])