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.
This commit is contained in:
2026-04-22 21:18:02 +07:00
commit 2e3e61dcbb
9 changed files with 406 additions and 0 deletions
+17
View File
@@ -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/
+18
View File
@@ -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/
+23
View File
@@ -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"]
+99
View File
@@ -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.
View File
+138
View File
@@ -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])
+89
View File
@@ -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
+17
View File
@@ -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:
+5
View File
@@ -0,0 +1,5 @@
fastapi==0.115.0
uvicorn[standard]==0.32.0
gensim==4.3.3
numpy<2.0
pydantic==2.9.2