mirror of
https://github.com/tiennm99/word2sim.git
synced 2026-08-14 08:26:40 +00:00
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:
@@ -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
@@ -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
@@ -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"]
|
||||
@@ -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.
|
||||
+138
@@ -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])
|
||||
@@ -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
|
||||
@@ -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:
|
||||
@@ -0,0 +1,5 @@
|
||||
fastapi==0.115.0
|
||||
uvicorn[standard]==0.32.0
|
||||
gensim==4.3.3
|
||||
numpy<2.0
|
||||
pydantic==2.9.2
|
||||
Reference in New Issue
Block a user