Sentence Similarity with BERT and Flask

How a sentence becomes a vector

Tokenisation. BERT's tokenizer splits text into WordPiece sub-words, adds a [CLS] token at the start and [SEP] at the end, and maps each piece to an ID. BERT-base has 512 position embeddings, so longer input must be truncated. When you batch sentences of different lengths, the tokenizer pads shorter ones and returns an attention_mask with 1 for real tokens and 0 for padding.

Encoding. BERT-base runs the IDs through 12 transformer layers with 768-dimensional hidden states and about 110 million parameters. Self-attention makes each output vector depend on the whole sentence, so 'bank' in 'river bank' and 'bank loan' get different vectors.

Pooling. The encoder returns one vector per token; similarity needs one per sentence. The options are the [CLS] vector, the mean of token vectors, or the max over them. Mean pooling must ignore padding, otherwise a short sentence batched next to a long one is averaged with dozens of padding vectors and its embedding changes depending on what it was batched with.

Comparison. Cosine similarity is the dot product of two vectors divided by the product of their lengths. If you L2-normalise embeddings once, cosine becomes a plain dot product, and comparing one query against a million stored vectors becomes a single matrix-vector multiply.

import torch
from transformers import AutoTokenizer, AutoModel

def mean_pool(last_hidden, mask):
    mask = mask.unsqueeze(-1).to(last_hidden.dtype)      # [B, T, 1]
    summed = (last_hidden * mask).sum(dim=1)              # ignore padding
    counts = mask.sum(dim=1).clamp(min=1e-9)
    return summed / counts

tok = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME).eval()

@torch.inference_mode()
def embed(texts):
    batch = tok(texts, padding=True, truncation=True, max_length=256,
                return_tensors="pt")
    out = model(**batch)
    vecs = mean_pool(out.last_hidden_state, batch["attention_mask"])
    return torch.nn.functional.normalize(vecs, dim=1)    # unit length

How a sentence becomes a vector

Tokenisation. BERT's tokenizer splits text into WordPiece sub-words, adds a [CLS] token at the start and [SEP] at the end, and maps each piece to an ID. BERT-base has 512 position embeddings, so longer input must be truncated. When you batch sentences of different lengths, the tokenizer pads shorter ones and returns an attention_mask with 1 for real tokens and 0 for padding.

Encoding. BERT-base runs the IDs through 12 transformer layers with 768-dimensional hidden states and about 110 million parameters. Self-attention makes each output vector depend on the whole sentence, so 'bank' in 'river bank' and 'bank loan' get different vectors.

Pooling. The encoder returns one vector per token; similarity needs one per sentence. The options are the [CLS] vector, the mean of token vectors, or the max over them. Mean pooling must ignore padding, otherwise a short sentence batched next to a long one is averaged with dozens of padding vectors and its embedding changes depending on what it was batched with.

Comparison. Cosine similarity is the dot product of two vectors divided by the product of their lengths. If you L2-normalise embeddings once, cosine becomes a plain dot product, and comparing one query against a million stored vectors becomes a single matrix-vector multiply.

import torch
from transformers import AutoTokenizer, AutoModel

def mean_pool(last_hidden, mask):
    mask = mask.unsqueeze(-1).to(last_hidden.dtype)      # [B, T, 1]
    summed = (last_hidden * mask).sum(dim=1)              # ignore padding
    counts = mask.sum(dim=1).clamp(min=1e-9)
    return summed / counts

tok = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModel.from_pretrained(MODEL_NAME).eval()

@torch.inference_mode()
def embed(texts):
    batch = tok(texts, padding=True, truncation=True, max_length=256,
                return_tensors="pt")
    out = model(**batch)
    vecs = mean_pool(out.last_hidden_state, batch["attention_mask"])
    return torch.nn.functional.normalize(vecs, dim=1)    # unit length

Why raw BERT embeddings disappoint

BERT was pre-trained to predict masked words, not to place sentences with similar meaning near each other. Its raw output vectors occupy a narrow cone of the space, so almost any two sentences score a high cosine and the useful signal is squeezed into a small range. Reimers and Gurevych showed in the Sentence-BERT paper (2019) that averaged raw BERT embeddings did worse on semantic textual similarity benchmarks than simply averaging GloVe word vectors, and that the [CLS] vector was worse still.

The fix is a model fine-tuned for the task. Sentence-BERT trains a siamese network on sentence pairs, natural language inference data and later large collections of question-answer and paraphrase pairs, with a contrastive objective that pulls paraphrases together and pushes unrelated sentences apart. The sentence-transformers library packages these models. Two common starting points are all-MiniLM-L6-v2, a 6-layer model with 384-dimensional output that truncates input at 256 word pieces, and all-mpnet-base-v2, larger and slower with 768-dimensional output and better quality on most benchmarks. Check the model card for the pooling method and maximum length of whichever model you choose, because they differ between models.

There is a second family worth knowing. A bi-encoder embeds each sentence separately, so you can precompute and index a corpus. A cross-encoder feeds both sentences into one transformer pass and outputs a score directly; it is more accurate but must run once per pair, so it cannot search a million documents. Production systems retrieve candidates with a bi-encoder and re-rank the top few with a cross-encoder.

Building the Flask service

The service below loads the model once at import, exposes a pairwise endpoint and a corpus search endpoint, enforces input limits and reports the model identity with every answer, so clients can tell when scores came from a different model.

import hashlib, os, time
import numpy as np
from flask import Flask, jsonify, request
from sentence_transformers import SentenceTransformer

MODEL_ID = os.environ.get("MODEL_ID", "sentence-transformers/all-MiniLM-L6-v2")
MAX_CHARS, MAX_BATCH = 2000, 64

app = Flask(__name__)
model = SentenceTransformer(MODEL_ID, device="cpu")   # once per process
corpus = [line.strip() for line in open("faq.txt", encoding="utf-8") if line.strip()]
corpus_vecs = model.encode(corpus, batch_size=64, normalize_embeddings=True)
cache = {}

def encode(texts):
    keys = [hashlib.sha256(t.encode()).hexdigest() for t in texts]
    missing = [t for t, k in zip(texts, keys) if k not in cache]
    if missing:
        vecs = model.encode(missing, batch_size=MAX_BATCH, normalize_embeddings=True)
        for t, v in zip(missing, vecs):
            cache[hashlib.sha256(t.encode()).hexdigest()] = v
    return np.stack([cache[k] for k in keys])

def bad_input(texts):
    return (not texts or len(texts) > MAX_BATCH
            or any(not isinstance(t, str) or not t.strip() or len(t) > MAX_CHARS for t in texts))

@app.post("/similarity")
def similarity():
    body = request.get_json(silent=True) or {}
    pair = [body.get("a"), body.get("b")]
    if bad_input(pair):
        return jsonify(error="a and b must be non-empty strings"), 400
    t0 = time.perf_counter()
    a, b = encode(pair)
    return jsonify(score=float(a @ b), model=MODEL_ID,
                   ms=round((time.perf_counter() - t0) * 1000, 1))

@app.post("/search")
def search():
    q = (request.get_json(silent=True) or {}).get("query")
    if bad_input([q]):
        return jsonify(error="query must be a non-empty string"), 400
    scores = corpus_vecs @ encode([q])[0]
    top = np.argsort(-scores)[:5]
    return jsonify(model=MODEL_ID,
                   results=[{"text": corpus[i], "score": float(scores[i])} for i in top])

The cache here is a plain dictionary to keep the example short; in production bound it with an LRU or move it to Redis keyed by model identifier plus text hash, so a model upgrade never serves stale vectors. SentenceTransformer.encode already disables gradient tracking, which matters because autograd would otherwise keep activations alive for a backward pass that never happens.

Serving it in production

Flask's development server is single-process and not for production. Run the app under Gunicorn, and remember that each Gunicorn worker is a separate process with its own copy of the model and its own PyTorch thread pool.

# 2 workers x 4 intra-op threads = 8 cores; never let the product exceed the core count
OMP_NUM_THREADS=4 gunicorn -w 2 -b 0.0.0.0:8000 --timeout 30 app:app

• Memory. MiniLM's weights are under 100 MB in float32, but each worker also holds the corpus matrix, the cache and PyTorch's runtime. A 768-dimensional model with eight workers can reach several gigabytes. Measure resident memory per worker and divide the container's memory by it.
• Threads. PyTorch uses all cores per process by default. Eight workers each spawning eight threads on an 8-core machine thrash the CPU and latency gets worse as you add workers. Set OMP_NUM_THREADS or torch.set_num_threads explicitly.
• Batching. Encoding 32 sentences in one call costs far less than 32 calls, because the matrix multiplies amortise. Accept lists in the API, and for high traffic put a micro-batching queue in front of the model that waits a few milliseconds to collect requests.
• Warm-up. The first call after start-up is slow. Encode a dummy sentence at import and only report ready from the health check after that.
• Scaling further. Export to ONNX Runtime or use a GPU and a dedicated inference server when CPU latency or throughput is not enough; keep Flask as a thin API in front.

Worked example: matching questions to an FAQ

A support team has 1,200 FAQ answers and wants incoming questions matched to an answer automatically when the match is confident, and routed to a human otherwise. The corpus matrix is 1,200 x 384 floats, under 2 MB, so brute-force dot products are fast and no vector database is needed yet.

The key decision is the threshold. Cosine scores are not probabilities and their scale differs between models, so a threshold copied from a blog post or from another model is meaningless. Label a few hundred real pairs as match or no match, then pick the threshold that meets your precision target:

import numpy as np

def pick_threshold(scores, labels, min_precision=0.95):
    best = None
    for t in np.unique(scores):
        pred = scores >= t
        tp = np.sum(pred & labels); fp = np.sum(pred & ~labels)
        if tp + fp == 0:
            continue
        precision, recall = tp / (tp + fp), tp / labels.sum()
        if precision >= min_precision and (best is None or recall > best[2]):
            best = (float(t), float(precision), float(recall))
    return best   # (threshold, precision, recall) or None

scores = np.array([...])           # model scores on labelled pairs
labels = np.array([...], dtype=bool)
print(pick_threshold(scores, labels))

Questions above the threshold get the answer automatically; the rest go to an agent, together with the top three candidates as suggestions. Re-run the calibration whenever the model or the FAQ changes.

Failure modes

• Negation and numbers. 'The refund was approved' and 'The refund was not approved' share almost every token and often score high. Bi-encoders are weak at negation, quantities and dates; re-rank with a cross-encoder or add rules where they matter.
• Silent truncation. Text beyond the model's maximum length is dropped without an error. Chunk long documents and embed the chunks.
• Mixed embedding spaces. Vectors from two different models, or two revisions of one model, are not comparable. Store the model identifier with every vector and re-embed on upgrade.
• Batch-dependent results. Pooling that ignores the attention mask makes a sentence's vector depend on its neighbours in the batch.
• Domain shift. General models underperform on legal, medical or internal jargon. Evaluate on your own labelled pairs and fine-tune if the gap is large.
• Unbounded input. A 5 MB request body ties up a worker. Enforce size and batch limits before encoding.

Testing and monitoring

Treat the model like any other dependency with behaviour you must pin and observe.

• Pin the revision. Load a specific model revision or a copy in your own artifact store, not whatever the hub serves today, and bake it into the container image so start-up never downloads.
• Golden tests. Keep twenty or so sentence pairs with expected orderings, for example a paraphrase must outscore an unrelated sentence, and run them in CI. Test orderings, not exact scores, since floating-point results shift slightly across hardware and library versions.
• Metrics. Export request latency split into encode time and total time, batch size, cache hit rate, input length in tokens and the share of inputs that were truncated. A rising truncation share means users are sending documents to an endpoint built for sentences.
• Score drift. Track the distribution of top-1 scores per day. A sudden drop in the median often means new topics the FAQ does not cover, which is useful product feedback, or a broken deployment that loaded the wrong model.

Trade-offs

Start with the smallest model that meets your accuracy target on labelled data from your own domain; accuracy on public benchmarks does not transfer reliably to a narrow corpus.

What to do next

Keep learning with the mathematics of embeddings, Matryoshka embeddings, rerankers, semantic caching, vector search at scale and quantizing embeddings. Then:

1. Replace raw BERT averaging with a sentence-embedding model and mask-aware pooling.
2. Normalise embeddings once and compare with dot products.
3. Run the Flask app under Gunicorn with an explicit worker and thread budget, and load the model once per worker.
4. Label 200 to 500 pairs from your own traffic and calibrate a threshold for your precision target.
5. Store the model identifier with every vector and cache embeddings by model and text hash.
6. Add a cross-encoder re-ranker if negation or near-miss errors show up in review.