From 2e3e61dcbb5378378a38db4c3602a722f9252be8 Mon Sep 17 00:00:00 2001
From: tiennm99 <tiennm99@outlook.com>
Date: Wed, 22 Apr 2026 21:18:02 +0700
Subject: [PATCH] feat: initial scaffold for word2sim similarity API

Stateless FastAPI service exposing word2vec cosine similarity,
nearest neighbors, vocab lookup, and random-word picker.
Dockerized with gensim GoogleNews pretrained model support.
---
 .dockerignore      |  17 ++++++
 .gitignore         |  18 ++++++
 Dockerfile         |  23 ++++++++
 README.md          |  99 ++++++++++++++++++++++++++++++++
 app/__init__.py    |   0
 app/main.py        | 138 +++++++++++++++++++++++++++++++++++++++++++++
 app/vectors.py     |  89 +++++++++++++++++++++++++++++
 docker-compose.yml |  17 ++++++
 requirements.txt   |   5 ++
 9 files changed, 406 insertions(+)
 create mode 100644 .dockerignore
 create mode 100644 .gitignore
 create mode 100644 Dockerfile
 create mode 100644 README.md
 create mode 100644 app/__init__.py
 create mode 100644 app/main.py
 create mode 100644 app/vectors.py
 create mode 100644 docker-compose.yml
 create mode 100644 requirements.txt

diff --git a/.dockerignore b/.dockerignore
new file mode 100644
index 0000000..18a0e61
--- /dev/null
+++ b/.dockerignore
@@ -0,0 +1,17 @@
+__pycache__/
+*.pyc
+.venv/
+.git/
+.gitignore
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+.env
+.env.*
+*.egg-info/
+dist/
+build/
+README.md
+vectors.bin
+*.bin
+gensim-cache/
diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..426aa4e
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,18 @@
+__pycache__/
+*.pyc
+*.pyo
+.venv/
+.env
+.env.*
+!.env.example
+.pytest_cache/
+.mypy_cache/
+.ruff_cache/
+*.egg-info/
+dist/
+build/
+
+# models + caches — never commit
+vectors.bin
+*.bin
+gensim-cache/
diff --git a/Dockerfile b/Dockerfile
new file mode 100644
index 0000000..57d1e99
--- /dev/null
+++ b/Dockerfile
@@ -0,0 +1,23 @@
+FROM python:3.11-slim
+
+WORKDIR /app
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+      curl \
+    && rm -rf /var/lib/apt/lists/*
+
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+COPY app ./app
+
+ENV GENSIM_DATA_DIR=/data/gensim-cache \
+    MODEL_NAME=word2vec-google-news-300 \
+    PORT=8000
+
+EXPOSE 8000
+
+HEALTHCHECK --interval=30s --timeout=5s --start-period=300s --retries=3 \
+  CMD curl -fsS http://localhost:8000/health || exit 1
+
+CMD ["uvicorn", "app.main:app", "--host", "0.0.0.0", "--port", "8000"]
diff --git a/README.md b/README.md
new file mode 100644
index 0000000..85bbda2
--- /dev/null
+++ b/README.md
@@ -0,0 +1,99 @@
+# word2sim
+
+Tiny HTTP service that returns word2vec cosine similarity and nearest neighbors.
+Stateless. No sessions. Just the math.
+
+Designed as a backend building block — a Semantle-style game, a search re-ranker,
+or a writing-assistance tool can all sit on top.
+
+## Stack
+- FastAPI + uvicorn
+- gensim (loads pretrained `word2vec-google-news-300` by default: 3M tokens × 300 dims, ~3.4GB RAM)
+
+## Endpoints
+
+| Method | Path | Purpose |
+|---|---|---|
+| GET | `/health` | liveness probe |
+| GET | `/similarity?a=X&b=Y` | cosine similarity between two words |
+| GET | `/neighbors?word=X&topn=10` | nearest-neighbor words with scores |
+| GET | `/vocab?word=X` | check if a word is in vocab; return canonical form |
+| GET | `/random` | random vocab word, filtered for game-friendliness |
+
+### Examples
+
+```bash
+curl 'http://localhost:8000/similarity?a=king&b=queen'
+# {"a":"king","b":"queen","canonical_a":"king","canonical_b":"queen",
+#  "in_vocab_a":true,"in_vocab_b":true,"similarity":0.6510957}
+
+curl 'http://localhost:8000/neighbors?word=ocean&topn=5'
+# {"word":"ocean","canonical":"ocean","in_vocab":true,
+#  "neighbors":[{"word":"oceans","similarity":0.78},{"word":"sea","similarity":0.75}, ...]}
+
+curl 'http://localhost:8000/vocab?word=Paris'
+# {"word":"Paris","canonical":"Paris","in_vocab":true}
+
+curl 'http://localhost:8000/random?min_rank=500&max_rank=20000&min_len=4&max_len=8'
+# {"word":"harbor","rank":8421}
+```
+
+### `/random` query params
+| Param | Default | Meaning |
+|---|---|---|
+| `min_rank` | 100 | skip the top-N most frequent tokens (common function words) |
+| `max_rank` | 50000 | cap at top-N most frequent (avoids rare/noisy tail) |
+| `alpha_only` | true | reject phrases (`new_york`), digits, punctuation |
+| `min_len` | 3 | |
+| `max_len` | 12 | |
+
+Uses rejection sampling over the frequency-sorted vocab; returns 503 if no word matches within 1000 attempts (loosen the filters).
+
+Out-of-vocab words return `in_vocab:false` and `similarity:null`. Case-insensitive lookup tries exact → lower → capitalized.
+
+## Quick start
+
+```bash
+docker compose up --build
+# first boot downloads ~1.6GB model into the gensim-cache volume; later boots are instant
+```
+
+## Using your own vectors
+Skip the download by mounting a locally trained `vectors.bin`:
+```yaml
+# docker-compose.yml
+services:
+  word2sim:
+    environment:
+      MODEL_PATH: /models/vectors.bin
+    volumes:
+      - ./vectors.bin:/models/vectors.bin:ro
+```
+(Train one with `bash demo-word.sh` from the upstream word2vec repo.)
+
+## Config (env vars)
+| Var | Default | Meaning |
+|---|---|---|
+| `MODEL_NAME` | `word2vec-google-news-300` | gensim downloader id |
+| `MODEL_PATH` | _(unset)_ | if set + file exists, load this `.bin` instead (skips download) |
+| `GENSIM_DATA_DIR` | `/data/gensim-cache` | where gensim caches downloaded models |
+
+## Project layout
+```
+word2sim/
+├── app/
+│   ├── main.py       # FastAPI routes
+│   └── vectors.py    # model loader + similarity/neighbors
+├── Dockerfile
+├── docker-compose.yml
+└── requirements.txt
+```
+
+## Building a Semantle-style game on top
+The game server keeps state (session, secret, guess log); it calls word2sim per guess:
+```
+new game:  GET  /random?min_rank=500&max_rank=20000&min_len=4&max_len=10
+           GET  /neighbors?word={secret}&topn=1000     → cache ranks locally
+on guess:  GET  /similarity?a={secret}&b={guess}
+```
+word2sim stays stateless and cache-friendly.
diff --git a/app/__init__.py b/app/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/app/main.py b/app/main.py
new file mode 100644
index 0000000..ba444f6
--- /dev/null
+++ b/app/main.py
@@ -0,0 +1,138 @@
+"""FastAPI entry point for word2sim.
+
+Exposes pure read-only endpoints over a word2vec model: pairwise similarity,
+nearest neighbors, vocab lookup. Stateless — no sessions, no persistence.
+"""
+
+from __future__ import annotations
+
+from contextlib import asynccontextmanager
+
+from fastapi import FastAPI, HTTPException, Query
+from pydantic import BaseModel
+
+from app.vectors import canonicalize, load_model, neighbors, random_word, similarity
+
+
+@asynccontextmanager
+async def lifespan(_app: FastAPI):
+    # Force model load before accepting traffic (GoogleNews cold start ~30s).
+    load_model()
+    yield
+
+
+app = FastAPI(title="word2sim", version="0.1.0", lifespan=lifespan)
+
+
+class SimilarityResponse(BaseModel):
+    a: str
+    b: str
+    canonical_a: str | None
+    canonical_b: str | None
+    in_vocab_a: bool
+    in_vocab_b: bool
+    similarity: float | None  # null iff either side is out-of-vocab
+
+
+class NeighborEntry(BaseModel):
+    word: str
+    similarity: float
+
+
+class NeighborsResponse(BaseModel):
+    word: str
+    canonical: str | None
+    in_vocab: bool
+    neighbors: list[NeighborEntry]
+
+
+class VocabResponse(BaseModel):
+    word: str
+    canonical: str | None
+    in_vocab: bool
+
+
+class RandomWordResponse(BaseModel):
+    word: str
+    rank: int  # position in frequency-sorted vocab (0 = most frequent)
+
+
+@app.get("/health")
+def health() -> dict[str, str]:
+    return {"status": "ok"}
+
+
+@app.get("/similarity", response_model=SimilarityResponse)
+def get_similarity(
+    a: str = Query(min_length=1, max_length=64),
+    b: str = Query(min_length=1, max_length=64),
+) -> SimilarityResponse:
+    kv = load_model()
+    canonical_a = canonicalize(kv, a)
+    canonical_b = canonicalize(kv, b)
+    score = (
+        similarity(kv, canonical_a, canonical_b)
+        if canonical_a and canonical_b
+        else None
+    )
+    return SimilarityResponse(
+        a=a,
+        b=b,
+        canonical_a=canonical_a,
+        canonical_b=canonical_b,
+        in_vocab_a=canonical_a is not None,
+        in_vocab_b=canonical_b is not None,
+        similarity=score,
+    )
+
+
+@app.get("/neighbors", response_model=NeighborsResponse)
+def get_neighbors(
+    word: str = Query(min_length=1, max_length=64),
+    topn: int = Query(default=10, ge=1, le=1000),
+) -> NeighborsResponse:
+    kv = load_model()
+    canonical = canonicalize(kv, word)
+    if canonical is None:
+        return NeighborsResponse(
+            word=word, canonical=None, in_vocab=False, neighbors=[]
+        )
+    entries = [
+        NeighborEntry(word=w, similarity=s) for w, s in neighbors(kv, canonical, topn)
+    ]
+    return NeighborsResponse(
+        word=word, canonical=canonical, in_vocab=True, neighbors=entries
+    )
+
+
+@app.get("/vocab", response_model=VocabResponse)
+def get_vocab(word: str = Query(min_length=1, max_length=64)) -> VocabResponse:
+    kv = load_model()
+    canonical = canonicalize(kv, word)
+    return VocabResponse(word=word, canonical=canonical, in_vocab=canonical is not None)
+
+
+@app.get("/random", response_model=RandomWordResponse)
+def get_random(
+    min_rank: int = Query(default=100, ge=0),
+    max_rank: int = Query(default=50000, ge=1),
+    alpha_only: bool = Query(default=True),
+    min_len: int = Query(default=3, ge=1, le=64),
+    max_len: int = Query(default=12, ge=1, le=64),
+) -> RandomWordResponse:
+    if min_rank >= max_rank:
+        raise HTTPException(status_code=400, detail="min_rank must be < max_rank")
+    if min_len > max_len:
+        raise HTTPException(status_code=400, detail="min_len must be <= max_len")
+    kv = load_model()
+    word = random_word(
+        kv,
+        min_rank=min_rank,
+        max_rank=max_rank,
+        alpha_only=alpha_only,
+        min_len=min_len,
+        max_len=max_len,
+    )
+    if word is None:
+        raise HTTPException(status_code=503, detail="no word matched filter; loosen the constraints")
+    return RandomWordResponse(word=word, rank=kv.key_to_index[word])
diff --git a/app/vectors.py b/app/vectors.py
new file mode 100644
index 0000000..5f7b445
--- /dev/null
+++ b/app/vectors.py
@@ -0,0 +1,89 @@
+"""Word2vec model loader and similarity primitives.
+
+Process-wide KeyedVectors singleton; loaded lazily on first use.
+Supports a gensim-downloader id (MODEL_NAME) or a local .bin file (MODEL_PATH).
+"""
+
+from __future__ import annotations
+
+import os
+import random as _random
+import string
+from typing import Optional
+
+from gensim.models import KeyedVectors
+
+_MODEL: Optional[KeyedVectors] = None
+
+
+def load_model() -> KeyedVectors:
+    """Return the singleton KeyedVectors, loading it on first call."""
+    global _MODEL
+    if _MODEL is not None:
+        return _MODEL
+
+    local_path = os.environ.get("MODEL_PATH")
+    if local_path and os.path.exists(local_path):
+        _MODEL = KeyedVectors.load_word2vec_format(local_path, binary=True)
+        return _MODEL
+
+    # Defer gensim.downloader import so MODEL_PATH users avoid the network path.
+    import gensim.downloader as api
+
+    model_name = os.environ.get("MODEL_NAME", "word2vec-google-news-300")
+    _MODEL = api.load(model_name)
+    return _MODEL
+
+
+def canonicalize(kv: KeyedVectors, word: str) -> Optional[str]:
+    """Resolve `word` to its in-vocab form, trying exact → lower → capitalized.
+
+    GoogleNews vectors are case-sensitive; this matches most user expectations
+    without forcing callers to know the casing conventions of the training corpus.
+    """
+    for candidate in (word, word.lower(), word.capitalize()):
+        if candidate in kv:
+            return candidate
+    return None
+
+
+def similarity(kv: KeyedVectors, a: str, b: str) -> float:
+    """Cosine similarity between two in-vocab words. Caller must canonicalize first."""
+    return float(kv.similarity(a, b))
+
+
+def neighbors(kv: KeyedVectors, word: str, topn: int) -> list[tuple[str, float]]:
+    """Top-N nearest-neighbor words with cosine scores. Caller must canonicalize first."""
+    return [(w, float(s)) for w, s in kv.most_similar(word, topn=topn)]
+
+
+def random_word(
+    kv: KeyedVectors,
+    *,
+    min_rank: int = 0,
+    max_rank: Optional[int] = None,
+    alpha_only: bool = True,
+    min_len: int = 1,
+    max_len: int = 64,
+    max_attempts: int = 1000,
+) -> Optional[str]:
+    """Return a random vocab word matching the filters, or None if no match within budget.
+
+    `index_to_key` is frequency-ordered for word2vec .bin files, so `min_rank`/`max_rank`
+    act as a frequency window — e.g. min_rank=100 skips the most common function words,
+    max_rank=50000 avoids the rare/noisy tail. `alpha_only=True` rejects phrases
+    (`new_york` has `_`), digits, and punctuation.
+    """
+    vocab = kv.index_to_key
+    upper = min(max_rank, len(vocab)) if max_rank is not None else len(vocab)
+    if min_rank >= upper:
+        return None
+    allowed = set(string.ascii_letters) if alpha_only else None
+    for _ in range(max_attempts):
+        word = vocab[_random.randrange(min_rank, upper)]
+        if not (min_len <= len(word) <= max_len):
+            continue
+        if allowed is not None and not all(c in allowed for c in word):
+            continue
+        return word
+    return None
diff --git a/docker-compose.yml b/docker-compose.yml
new file mode 100644
index 0000000..86c1d55
--- /dev/null
+++ b/docker-compose.yml
@@ -0,0 +1,17 @@
+services:
+  word2sim:
+    build: .
+    container_name: word2sim
+    ports:
+      - "8000:8000"
+    environment:
+      MODEL_NAME: word2vec-google-news-300
+    volumes:
+      - gensim-cache:/data/gensim-cache
+      # mount a locally trained vectors.bin instead of downloading:
+      # - ./vectors.bin:/models/vectors.bin:ro
+      # and set MODEL_PATH: /models/vectors.bin above
+    restart: unless-stopped
+
+volumes:
+  gensim-cache:
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..19986c2
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1,5 @@
+fastapi==0.115.0
+uvicorn[standard]==0.32.0
+gensim==4.3.3
+numpy<2.0
+pydantic==2.9.2