GloVe

58import argparse
59import json
60import os
61import re
62import time
63from collections import Counter
64from pathlib import Path
65
66import numpy as np
67from sklearn.decomposition import PCA
68
69
70# ── stopwords ───────────────────────────────────────────────────────
71STOPWORDS = frozenset({

89def load_corpus(filepath: str):
90    """Load corpus with preprocessing: lowercase, strip punctuation,
91    split hyphens, remove numeric tokens, normalize whitespace."""
92    raw, cleaned = [], []
93    with open(filepath) as f:
94        for line in f:
95            line = line.strip()
96            if not line or line.startswith("#"):
97                continue
98            line = line.lower()
99            line = re.sub(r"-", " ", line)
100            line = re.sub(r"[^a-z\s]", "", line)
101            line = re.sub(r"\s+", " ", line).strip()
102            tokens = line.split()
103            raw.append(tokens)
104            content = [t for t in tokens if t not in STOPWORDS]
105            if len(content) >= 2:
106                cleaned.append(content)
107    return raw, cleaned

110def build_vocab(sentences, min_count=3):
111    counts = Counter()
112    for s in sentences:
113        counts.update(s)
114    vocab = sorted(w for w, c in counts.items() if c >= min_count)
115    w2i = {w: i for i, w in enumerate(vocab)}
116    i2w = {i: w for w, i in w2i.items()}
117    return vocab, w2i, i2w, counts
118
119
120# ── co-occurrence matrix ────────────────────────────────────────────

122def build_cooccurrence(sentences, w2i, window, symmetric):
123    """Build dense co-occurrence matrix with 1/d distance weighting.
124
125    For each word pair within the context window, the contribution
126    is 1/distance (matching the GloVe paper). With --symmetric (default),
127    both X[i][j] and X[j][i] are incremented. With --no-symmetric,
128    only the left context contributes (asymmetric window).
129    """
130    V = len(w2i)
131    X = np.zeros((V, V), dtype=np.float64)
132
133    for sent in sentences:
134        idxs = [w2i[w] for w in sent if w in w2i]
135        n = len(idxs)
136        for i in range(n):
137            for j in range(max(0, i - window), i):
138                dist = i - j
139                weight = 1.0 / dist
140                X[idxs[i], idxs[j]] += weight
141                if symmetric:
142                    X[idxs[j], idxs[i]] += weight
143            if not symmetric:
144                continue
145            for j in range(i + 1, min(n, i + window + 1)):
146                dist = j - i
147                weight = 1.0 / dist
148                X[idxs[i], idxs[j]] += weight
149                X[idxs[j], idxs[i]] += weight
150
151    return X

154def compute_cooccurrence_stats(X, vocab):
155    """Compute sparsity, value distribution, and Zipf fit for X."""
156    nonzero_mask = X > 0
157    nonzero_vals = X[nonzero_mask]
158    total_entries = X.shape[0] * X.shape[1]
159    num_nonzero = int(nonzero_mask.sum())
160    sparsity = 1.0 - num_nonzero / total_entries
161
162    distribution = {}
163    if len(nonzero_vals) > 0:
164        distribution = {
165            "min": float(nonzero_vals.min()),
166            "max": float(nonzero_vals.max()),
167            "mean": float(nonzero_vals.mean()),
168            "median": float(np.median(nonzero_vals)),
169            "percentiles": {
170                str(p): float(np.percentile(nonzero_vals, p))
171                for p in [25, 50, 75, 90, 95, 99]
172            },
173        }
174
175    zipf_data = {}
176    if len(nonzero_vals) > 10:
177        sorted_vals = np.sort(nonzero_vals)[::-1]
178        ranks = np.arange(1, len(sorted_vals) + 1, dtype=np.float64)
179        log_ranks = np.log(ranks)
180        log_vals = np.log(sorted_vals)
181        valid = np.isfinite(log_vals) & np.isfinite(log_ranks)
182        if valid.sum() > 2:
183            coeffs = np.polyfit(log_ranks[valid], log_vals[valid], 1)
184            alpha = -coeffs[0]
185            predicted = coeffs[0] * log_ranks[valid] + coeffs[1]
186            ss_res = np.sum((log_vals[valid] - predicted) ** 2)
187            ss_tot = np.sum((log_vals[valid] - log_vals[valid].mean()) ** 2)
188            r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0
189            zipf_data = {
190                "alpha": float(alpha),
191                "r_squared": float(r_squared),
192                "intercept": float(coeffs[1]),
193            }
194
195    return {
196        "total_entries": total_entries,
197        "num_nonzero": num_nonzero,
198        "sparsity": float(sparsity),
199        "value_distribution": distribution,
200        "zipf": zipf_data,
201    }
202
203
204# ── weighting function ─────────────────────────────────────────────

206def f_weight(x, x_max, alpha):
207    """GloVe weighting: (x/x_max)^alpha if x < x_max, else 1."""
208    return (x / x_max) ** alpha if x < x_max else 1.0

211def f_weight_array(x_vals, x_max, alpha):
212    """Vectorized weighting for an array of co-occurrence values."""
213    weights = np.where(x_vals < x_max, (x_vals / x_max) ** alpha, 1.0)
214    return weights

229KEY_PAIRS = [
230    ("ice", "solid"), ("steam", "gas"), ("ice", "water"), ("steam", "water"),
231    ("ice", "cold"), ("steam", "hot"), ("ice", "freeze"), ("steam", "boil"),
232    ("king", "queen"), ("prince", "princess"), ("man", "woman"), ("boy", "girl"),
233    ("father", "mother"), ("son", "daughter"), ("brother", "sister"),
234    ("husband", "wife"), ("cat", "dog"), ("lion", "wolf"),
235]
236
237ANALOGY_TESTS = [
238    ("king", "man", "woman", "queen"),
239    ("queen", "woman", "man", "king"),
240    ("ice", "solid", "gas", "steam"),
241    ("steam", "gas", "solid", "ice"),
242    ("prince", "boy", "girl", "princess"),
243    ("father", "man", "woman", "mother"),
244]

309def train(X, vocab, w2i, i2w, dim, lr, x_max, alpha, epochs, seed):
310    """GloVe weighted least squares with manual AdaGrad.
311
312    Following the Stanford C implementation:
313    - W, W_tilde initialized uniform(-0.5/d, 0.5/d)
314    - b, b_tilde initialized uniform(-0.5/d, 0.5/d)
315    - AdaGrad accumulators initialized to 1.0 (not 0)
316    - Gradient clipping at [-100, 100]
317    """
318    rng = np.random.RandomState(seed)
319    V = len(vocab)
320
321    W = rng.uniform(-0.5 / dim, 0.5 / dim, (V, dim))
322    W_tilde = rng.uniform(-0.5 / dim, 0.5 / dim, (V, dim))
323    b = rng.uniform(-0.5 / dim, 0.5 / dim, V)
324    b_tilde = rng.uniform(-0.5 / dim, 0.5 / dim, V)
325
326    G_W = np.ones((V, dim))
327    G_Wt = np.ones((V, dim))
328    G_b = np.ones(V)
329    G_bt = np.ones(V)
330
331    nonzero_pairs = []
332    for i in range(V):
333        for j in range(V):
334            if X[i, j] > 0:
335                nonzero_pairs.append((i, j, X[i, j]))
336
337    log_X = np.zeros_like(X)
338    nz_mask = X > 0
339    log_X[nz_mask] = np.log(X[nz_mask])
340
341    f_weights_arr = np.array([f_weight(x, x_max, alpha)
342                              for _, _, x in nonzero_pairs])
343
344    print(f"  Non-zero pairs: {len(nonzero_pairs)}")
345    print(f"  Parameters: {2 * V * dim + 2 * V} "
346          f"({V}×{dim}×2 vectors + {V}×2 biases)")
347
348    records = []
349    all_losses = []
350    t0 = time.time()
351
352    rec = capture_epoch(0, W, W_tilde, b, b_tilde, 0.0, X, log_X,
353                        nonzero_pairs, w2i, i2w, vocab, f_weights_arr)
354    records.append(rec)
355
356    for epoch in range(1, epochs + 1):
357        indices = rng.permutation(len(nonzero_pairs))
358        total_cost = 0.0
359        total_weight = 0.0
360
361        for idx in indices:
362            i, j, x_ij = nonzero_pairs[idx]
363
364            diff = np.dot(W[i], W_tilde[j]) + b[i] + b_tilde[j] - log_X[i, j]
365            fw = f_weights_arr[idx]
366            fdiff = fw * diff
367
368            total_cost += fw * diff * diff
369            total_weight += 1.0
370
371            grad_w = fdiff * W_tilde[j]
372            grad_wt = fdiff * W[i]
373            grad_b = fdiff
374            grad_bt = fdiff
375
376            grad_w = np.clip(grad_w, -100, 100)
377            grad_wt = np.clip(grad_wt, -100, 100)
378            grad_b = np.clip(grad_b, -100, 100)
379            grad_bt = np.clip(grad_bt, -100, 100)
380
381            G_W[i] += grad_w ** 2
382            W[i] -= lr * grad_w / np.sqrt(G_W[i])
383
384            G_Wt[j] += grad_wt ** 2
385            W_tilde[j] -= lr * grad_wt / np.sqrt(G_Wt[j])
386
387            G_b[i] += grad_b ** 2
388            b[i] -= lr * grad_b / np.sqrt(G_b[i])
389
390            G_bt[j] += grad_bt ** 2
391            b_tilde[j] -= lr * grad_bt / np.sqrt(G_bt[j])
392
393        avg_loss = total_cost / total_weight if total_weight > 0 else 0.0
394        all_losses.append({"epoch": epoch, "loss": round(float(avg_loss), 8)})
395
396        record = should_record(epoch, epochs)
397        if record:
398            rec = capture_epoch(epoch, W, W_tilde, b, b_tilde, avg_loss,
399                                X, log_X, nonzero_pairs, w2i, i2w, vocab,
400                                f_weights_arr)
401            records.append(rec)
402
403        if epoch % 10 == 0 or epoch == epochs:
404            elapsed = time.time() - t0
405            print(f"  epoch {epoch:>4}/{epochs}  "
406                  f"loss={avg_loss:.6f}  [{elapsed:.1f}s]  "
407                  f"({len(records)} recorded)")
408
409    final_state = {
410        "W": W.copy(), "W_tilde": W_tilde.copy(),
411        "b": b.copy(), "b_tilde": b_tilde.copy(),
412        "G_W": G_W.copy(), "G_Wt": G_Wt.copy(),
413        "G_b": G_b.copy(), "G_bt": G_bt.copy(),
414    }
415
416    return records, all_losses, final_state
417
418
419# ── projections ─────────────────────────────────────────────────────

219def should_record(epoch, total_epochs):
220    if epoch <= 10:
221        return True
222    if epoch % 5 == 0:
223        return True
224    if epoch == total_epochs:
225        return True
226    return False

247def capture_epoch(epoch, W, W_tilde, b, b_tilde, avg_loss, X, log_X,
248                  nonzero_pairs, w2i, i2w, vocab, f_weights):
249    """Capture full state for a single recorded epoch."""
250    V, d = W.shape
251    combined = W + W_tilde
252
253    cosines = {}
254    for w1, w2 in KEY_PAIRS:
255        if w1 in w2i and w2 in w2i:
256            v1 = combined[w2i[w1]]
257            v2 = combined[w2i[w2]]
258            n1, n2 = np.linalg.norm(v1), np.linalg.norm(v2)
259            cos = float(np.dot(v1, v2) / (n1 * n2)) if n1 > 0 and n2 > 0 else 0.0
260            cosines[f"{w1}_{w2}"] = round(cos, 6)
261
262    errors = []
263    for idx, (i, j, x_ij) in enumerate(nonzero_pairs):
264        pred = float(np.dot(W[i], W_tilde[j]) + b[i] + b_tilde[j])
265        target = float(log_X[i, j])
266        err = abs(pred - target)
267        errors.append({
268            "i": int(i), "j": int(j),
269            "word_i": i2w[i], "word_j": i2w[j],
270            "x_ij": float(x_ij),
271            "log_x": round(target, 6),
272            "pred": round(pred, 6),
273            "error": round(err, 6),
274            "f_weight": round(float(f_weights[idx]), 6),
275        })
276
277    errors_sorted = sorted(errors, key=lambda e: e["error"], reverse=True)
278    top_loss_pairs = errors_sorted[:10]
279
280    overall_mse = np.mean([e["error"] ** 2 for e in errors])
281    overall_mae = np.mean([e["error"] for e in errors])
282    targets = np.array([e["log_x"] for e in errors])
283    preds = np.array([e["pred"] for e in errors])
284    ss_res = np.sum((targets - preds) ** 2)
285    ss_tot = np.sum((targets - targets.mean()) ** 2)
286    r_squared = 1.0 - ss_res / ss_tot if ss_tot > 0 else 0.0
287
288    return {
289        "epoch": epoch,
290        "avg_loss": round(float(avg_loss), 8),
291        "W": W.tolist(),
292        "W_tilde": W_tilde.tolist(),
293        "b": b.tolist(),
294        "b_tilde": b_tilde.tolist(),
295        "cosine_similarities": cosines,
296        "top_loss_pairs": top_loss_pairs,
297        "reconstruction": {
298            "mse": round(float(overall_mse), 8),
299            "mae": round(float(overall_mae), 8),
300            "r_squared": round(float(r_squared), 6),
301            "num_pairs": len(errors),
302        },
303        "all_errors": errors,
304    }
305
306
307# ── training loop ──────────────────────────────────────────────────

421def add_projections(records, i2w):
422    """Add 2D PCA projections of W+W_tilde to each recorded epoch."""
423    V = len(i2w)
424    pca = PCA(n_components=2, random_state=42)
425
426    last_combined = np.array(records[-1]["W"]) + np.array(records[-1]["W_tilde"])
427    pca.fit(last_combined)
428    var = pca.explained_variance_ratio_.tolist()
429
430    combined_snapshots = []
431    for rec in records:
432        combined = np.array(rec["W"]) + np.array(rec["W_tilde"])
433        combined_snapshots.append(combined)
434
435    for i, rec in enumerate(records):
436        proj = pca.transform(combined_snapshots[i])
437        rec["embeddings_2d"] = {i2w[j]: proj[j].tolist() for j in range(V)}
438        rec["pca_variance_explained"] = var
439
440
441# ── storage ─────────────────────────────────────────────────────────

443def save_local(records, all_losses, X, vocab, w2i, counts, config,
444               raw_sentences, cooc_stats, out_dir):
445    """Write all training data as structured local files."""
446    steps_dir = os.path.join(out_dir, "steps")
447    os.makedirs(steps_dir, exist_ok=True)
448
449    V = len(vocab)
450    nz_mask = X > 0
451    log_X = np.zeros_like(X)
452    log_X[nz_mask] = np.log(X[nz_mask])
453    f_X = np.zeros_like(X)
454    for i in range(V):
455        for j in range(V):
456            if X[i, j] > 0:
457                f_X[i, j] = f_weight(X[i, j], config["x_max"], config["alpha"])
458
459    # ── run_metadata.json ────────────────────────────────────
460    meta = {
461        "topic": "glove",
462        "corpus_name": "curated_v1",
463        "vocab": vocab,
464        "vocab_size": V,
465        "embed_dim": config["embed_dim"],
466        "config": config,
467        "total_epochs": config["epochs"],
468        "recorded_epochs": len(records),
469        "word_frequencies": {w: counts[w] for w in vocab},
470        "raw_sentences": [" ".join(s) for s in raw_sentences],
471        "cooccurrence_stats": cooc_stats,
472    }
473    _write(os.path.join(out_dir, "run_metadata.json"), meta, indent=2)
474
475    # ── cooccurrence_matrix.json ─────────────────────────────
476    cooc_data = {
477        "matrix": X.tolist(),
478        "log_matrix": [[round(float(log_X[i, j]), 6) if X[i, j] > 0 else None
479                         for j in range(V)] for i in range(V)],
480        "weights": [[round(float(f_X[i, j]), 6) if X[i, j] > 0 else None
481                      for j in range(V)] for i in range(V)],
482        "vocab": vocab,
483        "stats": cooc_stats,
484    }
485    _write(os.path.join(out_dir, "cooccurrence_matrix.json"), cooc_data)
486
487    # ── loss_curve.json ──────────────────────────────────────
488    _write(os.path.join(out_dir, "loss_curve.json"), all_losses)
489
490    # ── step_index.json ──────────────────────────────────────
491    step_index = [rec["epoch"] for rec in records]
492    _write(os.path.join(out_dir, "step_index.json"), step_index)
493
494    # ── embeddings_timeline.json ─────────────────────────────
495    timeline = []
496    for rec in records:
497        entry = {
498            "epoch": rec["epoch"],
499            "positions": rec.get("embeddings_2d", {}),
500            "loss": rec["avg_loss"],
501        }
502        if rec.get("pca_variance_explained"):
503            entry["pca_variance_explained"] = rec["pca_variance_explained"]
504        timeline.append(entry)
505    _write(os.path.join(out_dir, "embeddings_timeline.json"), timeline)
506
507    # ── steps/{N}.json ───────────────────────────────────────
508    total_step_bytes = 0
509    for rec in records:
510        step_path = os.path.join(steps_dir, f"{rec['epoch']}.json")
511        step_data = {
512            "epoch": rec["epoch"],
513            "avg_loss": rec["avg_loss"],
514            "W": rec["W"],
515            "W_tilde": rec["W_tilde"],
516            "b": rec["b"],
517            "b_tilde": rec["b_tilde"],
518            "cosine_similarities": rec["cosine_similarities"],
519            "top_loss_pairs": rec["top_loss_pairs"],
520            "reconstruction": rec["reconstruction"],
521            "all_errors": rec["all_errors"],
522            "embeddings_2d": rec.get("embeddings_2d"),
523            "pca_variance_explained": rec.get("pca_variance_explained"),
524        }
525        _write(step_path, step_data)
526        total_step_bytes += os.path.getsize(step_path)
527
528    # ── report ───────────────────────────────────────────────
529    sizes = {}
530    for name in ["run_metadata.json", "cooccurrence_matrix.json",
531                 "loss_curve.json", "step_index.json",
532                 "embeddings_timeline.json"]:
533        p = os.path.join(out_dir, name)
534        sizes[name] = os.path.getsize(p)
535
536    print(f"\n  Output directory: {out_dir}/")
537    for name, sz in sizes.items():
538        print(f"    {name:<35} {sz / 1024:>8.1f} KB")
539    print(f"    steps/ ({len(records)} files){' ' * 19}"
540          f"{total_step_bytes / 1e6:>7.1f} MB")
541    print(f"    {'TOTAL':<35} "
542          f"{(sum(sizes.values()) + total_step_bytes) / 1e6:>7.1f} MB")