fix: LR schedule, analogy examples, remove old pretrain.py

app.py

@@ -215,7 +215,7 @@ with gr.Blocks(title="microembeddings", theme=gr.themes.Soft()) as demo:
             analogy_btn = gr.Button("Solve", variant="primary")
             gr.Examples(
                 [["man", "king", "woman"], ["france", "paris", "germany"],
-                 ["bigger", "big", "small"]],
+                 ["big", "bigger", "small"]],
                 inputs=[a_input, b_input, c_input]
             )
             analogy_text = gr.Textbox(label="Results", interactive=False, lines=6)

microembeddings.py

@@ -80,9 +80,9 @@ def train(corpus, vocab_size, neg_dist, epochs=EPOCHS, embed_dim=EMBED_DIM,
     W = (np.random.randn(vocab_size, embed_dim) * scale).astype(np.float32)
     C = np.zeros((vocab_size, embed_dim), dtype=np.float32)
 
-    # Each corpus position generates ~window context pairs on average
-    # (random window from 1..window, mean = (window+1)/2, times 2 sides)
-    total_steps = epochs * len(corpus) * window
+    # Each position draws actual_window ~ uniform(1..window), generating
+    # 2*actual_window context pairs. Expected pairs = 2 * E[uniform(1..w)] = w+1
+    total_steps = epochs * len(corpus) * (window + 1)
     step = 0
     losses = []
 
@@ -186,7 +186,7 @@ if __name__ == "__main__":
         print(f"\n{word}: {', '.join(f'{w} ({s:.3f})' for w, s in neighbors[:5])}")
 
     print("\n--- Analogies ---")
-    for a, b, c in [("king", "man", "woman"), ("paris", "france", "germany"),
+    for a, b, c in [("man", "king", "woman"), ("france", "paris", "germany"),
                      ("big", "bigger", "small")]:
         results = analogy(a, b, c, W_norm, word2idx, idx2word)
         ans = results[0][0] if results else "?"

pretrain.py

@@ -1,24 +0,0 @@
-"""Pre-train embeddings and save for the Gradio app."""
-import json
-import numpy as np
-from microembeddings import (
-    load_text8, build_vocab, prepare_corpus, build_neg_table, train
-)
-
-words = load_text8(3000000)
-word2idx, idx2word, freqs = build_vocab(words)
-corpus = prepare_corpus(words, word2idx, freqs)
-neg_dist = build_neg_table(freqs)
-
-def progress(epoch, i, total, loss):
-    pct = i / total * 100
-    print(f"  Epoch {epoch+1}: {pct:.0f}%, loss={loss:.4f}", end="\r")
-
-W, losses = train(corpus, len(word2idx), neg_dist, epochs=3, callback=progress)
-
-# Save embeddings as raw binary .npy (no pickle), vocab/losses as JSON
-np.save("pretrained_W.npy", W)
-vocab_list = [idx2word[i] for i in range(len(idx2word))]
-with open("pretrained_vocab.json", "w") as f:
-    json.dump({"vocab": vocab_list, "losses": [float(x) for x in losses]}, f)
-print(f"\nSaved: {W.shape[0]} words x {W.shape[1]} dims")

pretrain_gensim.py

@@ -0,0 +1,60 @@
+"""Pre-train embeddings using gensim's optimized Word2Vec on full text8.
+
+This is the canonical script for generating pretrained_W.npy and
+pretrained_vocab.json. Run locally (not on HF Space):
+
+    pip install gensim numpy
+    python pretrain_gensim.py
+
+gensim is a dev-only dependency — the Space runtime uses only numpy,
+gradio, plotly, and scikit-learn (see requirements.txt).
+"""
+import json
+import numpy as np
+from gensim.models import Word2Vec
+from microembeddings import load_text8
+
+# Load full text8 corpus (17M words)
+words = load_text8(max_words=17_000_000)
+# gensim expects list of sentences — split into ~1000-word chunks
+sentences = [words[i:i+1000] for i in range(0, len(words), 1000)]
+
+print(f"Training on {len(words)} words, {len(sentences)} sentences...")
+model = Word2Vec(
+    sentences,
+    vector_size=50,
+    window=5,
+    min_count=5,
+    sg=1,           # skip-gram
+    negative=5,
+    epochs=5,
+    workers=4,
+    max_final_vocab=10000,
+)
+
+# Export to our format
+vocab_list = list(model.wv.index_to_key)
+W = np.array([model.wv[w] for w in vocab_list], dtype=np.float32)
+
+np.save("pretrained_W.npy", W)
+with open("pretrained_vocab.json", "w") as f:
+    json.dump({"vocab": vocab_list, "losses": []}, f)
+
+print(f"Saved: {W.shape[0]} words x {W.shape[1]} dims")
+
+# Quick quality check
+from microembeddings import normalize, most_similar, analogy
+W_norm = normalize(W)
+word2idx = {w: i for i, w in enumerate(vocab_list)}
+idx2word = {i: w for i, w in enumerate(vocab_list)}
+
+print("\n--- Nearest Neighbors ---")
+for word in ["king", "france", "dog", "computer"]:
+    neighbors = most_similar(word, W_norm, word2idx, idx2word, topn=5)
+    print(f"{word}: {', '.join(f'{w} ({s:.3f})' for w, s in neighbors)}")
+
+print("\n--- Analogies ---")
+for a, b, c in [("man", "king", "woman"), ("france", "paris", "germany"), ("big", "bigger", "small")]:
+    results = analogy(a, b, c, W_norm, word2idx, idx2word)
+    ans = results[0][0] if results else "?"
+    print(f"{a} : {b} :: {c} : {ans}")