Embeddings: Meaning as Coordinates

import numpy as np

# A small corpus where the structure is visible by eye. Function
# words are left out: "the" appears beside everything, so it
# carries no information about anything and only adds noise.
documents = [
    'cat drank milk', 'kitten drank milk',
    'cat chased mouse', 'kitten chased mouse',
    'cat purred loudly', 'kitten purred loudly',
    'dog ate bone', 'puppy ate bone',
    'dog barked loudly', 'puppy barked loudly',
    'dog fetched stick', 'puppy fetched stick',
    'engineer deployed service', 'developer deployed service',
    'engineer debugged service', 'developer debugged service',
    'engineer reviewed code', 'developer reviewed code',
    'server crashed during deployment',
    'service restarted after crash',
]

words = sorted({word for document in documents
                for word in document.split()})
index = {word: i for i, word in enumerate(words)}

# Count how often each pair of words appears in the same document.
counts = np.zeros((len(words), len(words)))
for document in documents:
    present = sorted(set(document.split()))
    for a in present:
        for b in present:
            if a != b:
                counts[index[a], index[b]] += 1

print(f'{len(words)} distinct words, {len(documents)} documents')
print()
print('the co-occurrence row for "cat":')
row = counts[index['cat']]
for word, value in sorted(zip(words, row), key=lambda p: -p[1])[:6]:
    print(f'  {word:<12} {value:.0f}')
print()
print(f'That row IS an embedding - a {len(words)}-dimensional')
print('vector describing "cat" by the company it keeps. Crude,')
print('mostly zeroes, and the principle is already complete.')
print()


def cosine(a, b):
    denominator = np.linalg.norm(a) * np.linalg.norm(b)
    return 0.0 if denominator == 0 else float(a @ b / denominator)


pairs = [('cat', 'kitten'), ('dog', 'puppy'), ('cat', 'dog'),
         ('cat', 'engineer'), ('engineer', 'developer'),
         ('milk', 'bone'), ('deployed', 'debugged')]
print('cosine similarity of the raw co-occurrence rows:')
for a, b in pairs:
    print(f'  {a:<10} {b:<12} {cosine(counts[index[a]], counts[index[b]]):+.3f}')
print()
print('Nobody supplied a definition, a dictionary or a label.')
print('Twenty short sentences produced all of this:')
print()
print('cat/kitten and dog/puppy are 1.000 - exactly identical,')
print('because in this corpus they occur in precisely the same')
print('contexts. That is the distributional hypothesis at its')
print('purest: two words with the same neighbours are, as far as')
print('this method can tell, the same word.')
print()
print('cat/dog is only 0.167, because they share "loudly" and')
print('nothing else. Same category to a human, different')
print('neighbourhoods in the text.')
print()
print('cat/engineer and milk/bone are 0.000 - no shared context')
print('at all. The animal half and the software half of the corpus')
print('have separated themselves with no help from anyone.')

import numpy as np

documents = [
    'cat drank milk', 'kitten drank milk',
    'cat chased mouse', 'kitten chased mouse',
    'cat purred loudly', 'kitten purred loudly',
    'dog ate bone', 'puppy ate bone',
    'dog barked loudly', 'puppy barked loudly',
    'dog fetched stick', 'puppy fetched stick',
    'engineer deployed service', 'developer deployed service',
    'engineer debugged service', 'developer debugged service',
    'engineer reviewed code', 'developer reviewed code',
    'server crashed during deployment',
    'service restarted after crash',
]
words = sorted({w for d in documents for w in d.split()})
index = {w: i for i, w in enumerate(words)}

counts = np.zeros((len(words), len(words)))
for document in documents:
    present = sorted(set(document.split()))
    for a in present:
        for b in present:
            if a != b:
                counts[index[a], index[b]] += 1

# Positive pointwise mutual information, then a truncated SVD.
# This is latent semantic analysis, and it is a real embedding
# method - the same linear algebra that word2vec turned out to be
# approximating.
total = counts.sum()
row_sums = counts.sum(axis=1, keepdims=True)
column_sums = counts.sum(axis=0, keepdims=True)
with np.errstate(divide='ignore', invalid='ignore'):
    expected = row_sums @ column_sums / total
    ratio = np.where(expected > 0, counts / np.where(expected > 0,
                                                     expected, 1), 1)
    pmi = np.log(ratio)
pmi = np.nan_to_num(np.maximum(pmi, 0.0))

DIMENSIONS = 6
u, singular, _ = np.linalg.svd(pmi, full_matrices=False)
embeddings = u[:, :DIMENSIONS] * singular[:DIMENSIONS]

print(f'{pmi.shape[1]} sparse dimensions -> {DIMENSIONS} dense ones')
print(f'storage per word: {pmi.shape[1] * 4} bytes before, '
      f'{DIMENSIONS * 4} bytes after')
print()
print('the embedding of "cat", in six numbers:')
print(' ', embeddings[index['cat']].round(3))
print()


def cosine(a, b):
    denominator = np.linalg.norm(a) * np.linalg.norm(b)
    return 0.0 if denominator == 0 else float(a @ b / denominator)


print('nearest neighbours in the dense space:')
for target in ['cat', 'dog', 'engineer', 'deployed', 'milk']:
    scores = [(cosine(embeddings[index[target]], embeddings[index[w]]), w)
              for w in words if w != target]
    best = sorted(scores, reverse=True)[:3]
    shown = ', '.join(f'{w} {s:.2f}' for s, w in best)
    print(f'  {target:<10} -> {shown}')
print()
print('Six numbers a word instead of twenty-nine, and every')
print('neighbourhood is still right: cat finds kitten, dog finds')
print('puppy, engineer finds developer, deployed finds debugged.')
print()
print('That compression is the whole value of an embedding. A')
print('sparse count matrix grows with the square of the')
print('vocabulary; a dense embedding is a fixed width however')
print('much text you have, which is why a real model can give you')
print('1,536 numbers for any sentence in any language.')

import numpy as np

# Three documents. Two are about the same thing at different
# lengths; the third is about something else.
short_about_cats = np.array([3.0, 1.0, 0.0, 0.0])
long_about_cats = np.array([30.0, 10.0, 0.0, 0.0])
short_about_servers = np.array([0.0, 0.0, 3.0, 1.0])


def cosine(a, b):
    return float(a @ b / (np.linalg.norm(a) * np.linalg.norm(b)))


def euclidean(a, b):
    return float(np.linalg.norm(a - b))


print('vectors (counts of four terms):')
for name, vector in [('short, about cats', short_about_cats),
                     ('long, about cats', long_about_cats),
                     ('short, about servers', short_about_servers)]:
    print(f'  {name:<22} {vector}  length {np.linalg.norm(vector):.1f}')
print()
print(f'{"pair":<44} {"cosine":>8} {"euclidean":>10}')
comparisons = [
    ('short cats vs long cats (same topic)',
     short_about_cats, long_about_cats),
    ('short cats vs short servers (different topic)',
     short_about_cats, short_about_servers),
    ('long cats vs short servers (different topic)',
     long_about_cats, short_about_servers),
]
for label, a, b in comparisons:
    print(f'{label:<44} {cosine(a, b):>8.3f} {euclidean(a, b):>10.3f}')
print()
print('Cosine gets it exactly right: 1.000 for the two cat')
print('documents whatever their length, 0.000 for the unrelated')
print('pair.')
print()
print('Euclidean distance gets it backwards. The two documents')
print('about cats are 28.5 apart, and the short cat document is')
print('only 4.5 from the server document - so a nearest-neighbour')
print('search by Euclidean distance would return the WRONG topic.')
print()
print('The reason is that length encodes how much text there is,')
print('not what it is about. Cosine measures the angle and')
print('discards the length, which is what you want. And once')
print('vectors are normalised to unit length the two measures')
print('agree, because then the only thing left is the angle:')
print()
unit_short = short_about_cats / np.linalg.norm(short_about_cats)
unit_long = long_about_cats / np.linalg.norm(long_about_cats)
print(f'  normalised, cosine    {cosine(unit_short, unit_long):.6f}')
print(f'  normalised, euclidean {euclidean(unit_short, unit_long):.6f}')
print()
print('Which is the practical rule: NORMALISE ON THE WAY IN. Then')
print('a dot product is a cosine, every vector database agrees')
print('with every other, and the question of which metric to use')
print('stops arising.')

import numpy as np

DOCUMENT = (
    'Refund policy. Customers may return physical goods within 30 '
    'days of delivery for a full refund, provided the packaging is '
    'intact. Digital goods are non-refundable once downloaded. '
    'Shipping costs are not refunded unless the item was faulty. '
    'For faulty items we also cover return postage. '
    'Gift cards cannot be refunded or exchanged. '
    'Refunds are processed within 5 working days of receipt. '
    'Contact support with your order number to begin a return.'
)

sentences = [s.strip() + '.' for s in DOCUMENT.split('. ') if s.strip()]


def chunk_by_sentences(sentences, per_chunk, overlap=0):
    chunks = []
    step = max(1, per_chunk - overlap)
    for start in range(0, len(sentences), step):
        piece = sentences[start:start + per_chunk]
        if piece:
            chunks.append(' '.join(piece))
        if start + per_chunk >= len(sentences):
            break
    return chunks


print(f'{len(sentences)} sentences, {len(DOCUMENT.split())} words')
print()
for per_chunk, overlap in [(1, 0), (3, 0), (3, 1), (8, 0)]:
    chunks = chunk_by_sentences(sentences, per_chunk, overlap)
    sizes = [len(c.split()) for c in chunks]
    print(f'{per_chunk} sentences, overlap {overlap}: {len(chunks)} chunks, '
          f'{min(sizes)}-{max(sizes)} words each')
print()

# The question that decides the chunk size: can a chunk answer a
# question on its own?
QUESTIONS = {
    'How long do I have to return a physical item?':
        ['30 days', 'physical'],
    'Can I get a refund on a downloaded ebook?':
        ['Digital', 'non-refundable'],
    'Who pays return postage on a faulty item?':
        ['faulty', 'postage'],
    'How long until the money arrives?':
        ['5 working days'],
}

print('can a chunk answer each question on its own?')
print()
for per_chunk, overlap in [(1, 0), (3, 1), (8, 0)]:
    chunks = chunk_by_sentences(sentences, per_chunk, overlap)
    answerable = 0
    for question, needed in QUESTIONS.items():
        if any(all(term.lower() in chunk.lower() for term in needed)
               for chunk in chunks):
            answerable += 1
    average = sum(len(c.split()) for c in chunks) / len(chunks)
    print(f'  {per_chunk} sentences (overlap {overlap}): '
          f'{answerable}/{len(QUESTIONS)} answerable, '
          f'{len(chunks)} chunks, {average:.0f} words each')
print()
print('One sentence per chunk is precise and loses context: the')
print('sentence about return postage does not say which items it')
print('applies to, so a chunk containing it alone cannot answer')
print('the question.')
print()
print('Eight sentences is one chunk containing everything, which')
print('answers every question and retrieves nothing useful - at')
print('that size every query matches the same single chunk and')
print('ranking has nothing to rank.')
print()
print('Three sentences with one of overlap is the usual answer,')
print('and the overlap is what stops a fact being split across a')
print('boundary. The rule of thumb that follows: a chunk should')
print('be the smallest passage that can answer a question WITHOUT')
print('the sentences around it.')

import numpy as np

# A knowledge base, embedded with TF-IDF and reduced by SVD.
# This is a real, usable retrieval system in fifty lines.
DOCUMENTS = [
    'Refunds for physical goods are available within 30 days of delivery.',
    'Digital downloads and ebooks cannot be refunded once the file has '
    'been accessed.',
    'Software licences are non-refundable after activation.',
    'Shipping charges are refunded only when the item arrived damaged '
    'or faulty.',
    'For faulty items we pay the return postage ourselves.',
    'Gift cards are non-refundable and cannot be exchanged for cash.',
    'Refunded money is paid back to the original card within five '
    'working days.',
    'To start a return, contact support with your order number.',
    'Returned items must have their original packaging intact.',
    'Our warehouse dispatches orders on the next working day.',
    'Standard delivery takes two to four working days within the UK.',
    'Express delivery arrives the next working day if ordered before noon.',
    'International delivery takes seven to fourteen working days.',
    'Delivery to the Channel Islands may take longer and incurs duty.',
    'You can change your delivery address before the order is dispatched.',
    'Track your parcel using the link in your dispatch email.',
    'A missing parcel should be reported after ten working days.',
    'Orders over fifty pounds qualify for free standard delivery.',
    'We do not deliver to PO boxes or forwarding addresses.',
    'Signature on delivery is required for orders over two hundred pounds.',
    'Your account password can be reset from the sign-in page.',
    'Two-factor authentication can be enabled in account settings.',
    'Invoices are available to download from your order history.',
    'VAT receipts are issued automatically for business accounts.',
]

STOP = set(('the a an is are of to in for and or be been can cannot '
            'your you with only once has if before within on our from '
            'do does').split())


def tokenise(text):
    cleaned = ''.join(c.lower() if c.isalnum() else ' ' for c in text)
    return [w for w in cleaned.split() if w not in STOP and len(w) > 2]


vocabulary = sorted({w for d in DOCUMENTS for w in tokenise(d)})
position = {w: i for i, w in enumerate(vocabulary)}

# Term frequency.
tf = np.zeros((len(DOCUMENTS), len(vocabulary)))
for row, document in enumerate(DOCUMENTS):
    for word in tokenise(document):
        tf[row, position[word]] += 1

# Inverse document frequency: a word in every document carries no
# information, so it gets no weight.
document_frequency = (tf > 0).sum(axis=0)
idf = np.log((1 + len(DOCUMENTS)) / (1 + document_frequency)) + 1.0
tfidf = tf * idf

# Reduce, then normalise. Normalising last is what makes a dot
# product a cosine.
u, singular, vt = np.linalg.svd(tfidf, full_matrices=False)
DIMENSIONS = 10
reduced = u[:, :DIMENSIONS] * singular[:DIMENSIONS]
norms = np.linalg.norm(reduced, axis=1, keepdims=True)
vectors = reduced / np.where(norms > 0, norms, 1)

print(f'{len(DOCUMENTS)} documents, {len(vocabulary)} terms, '
      f'reduced to {DIMENSIONS} dimensions')
print(f'index size: {vectors.nbytes:,} bytes '
      f'({vectors.nbytes / len(DOCUMENTS):.0f} bytes a document)')
print()


def embed_query(text):
    row = np.zeros(len(vocabulary))
    unknown = []
    for word in tokenise(text):
        if word in position:
            row[position[word]] += 1
        else:
            unknown.append(word)
    # Project into the same space using the right singular vectors.
    projected = (row * idf) @ vt[:DIMENSIONS].T
    norm = np.linalg.norm(projected)
    return projected / (norm if norm > 0 else 1), unknown


def search(query, k=2):
    vector, unknown = embed_query(query)
    scores = vectors @ vector
    order = np.argsort(scores)[::-1][:k]
    return [(float(scores[i]), DOCUMENTS[i]) for i in order], unknown


QUERIES = [
    'who pays postage if it is broken',
    'when will my parcel arrive',
    'reset my password',
    'can I get money back for an ebook',
    'what is the capital of France',
]
for query in QUERIES:
    results, unknown = search(query)
    print(f'Q: {query}')
    if unknown:
        print(f'   not in the index: {", ".join(unknown)}')
    for score, document in results:
        print(f'   {score:+.3f}  {document[:68]}')
    print()

print('Four different outcomes in five queries, and all four are')
print('worth recognising.')
print()
print('IT WORKS. "Who pays postage if it is broken" scores 0.993')
print('against "For faulty items we pay the return postage')
print('ourselves." The query does not contain "faulty" and the')
print('document does not contain "broken"; the match came through')
print('the co-occurrence structure the SVD learned. That is what')
print('"search by meaning" amounts to, and it is genuinely')
print('useful.')
print()
print('IT GUESSES. "Reset my password" finds the right document')
print('at 1.000 - and the SECOND result, about VAT receipts,')
print('scores 0.983. Nothing about VAT answers the question. A')
print('system that retrieves the top two just put an irrelevant')
print('document in front of the model at near-perfect confidence.')
print()
print('A MISSING WORD IS SILENT. "Ebook" is not in the index, so')
print('it contributes nothing, and the query quietly becomes')
print('"money back" - which finds the refund TIMING rule at 0.979')
print('instead of the rule saying ebooks are not refundable at')
print('all. Note where that right answer went: SECOND, at 0.709.')
print('It was retrieved - just not first. Take the top result')
print('only and you get the wrong rule; take the top two and the')
print('model sees both. That is the practical argument for')
print('retrieving several chunks rather than one, and it is worth')
print('more than any amount of tuning the embedding.')
print()
print('ALL-UNKNOWN IS DETECTABLE. "What is the capital of France"')
print('has no in-vocabulary terms, so the query vector is all')
print('zeroes and every score is exactly 0.000. That case you can')
print('catch.')
print()
print('Which is the warning worth taking away: the dangerous')
print('query is not the one with NO known words - that one scores')
print('zero and can be rejected. It is the one with SOME known')
print('words, which scores 0.98 on the wrong document and looks')
print('exactly like a success.')

import numpy as np

rng = np.random.default_rng(0)

# What similarity scores actually look like: relevant matches and
# irrelevant ones, as two overlapping distributions.
relevant = np.clip(rng.normal(0.62, 0.13, 400), -1, 1)
irrelevant = np.clip(rng.normal(0.28, 0.11, 4000), -1, 1)

print('score distributions:')
print(f'  relevant   mean {relevant.mean():.3f}, '
      f'5th-95th [{np.percentile(relevant, 5):.3f}, '
      f'{np.percentile(relevant, 95):.3f}]')
print(f'  irrelevant mean {irrelevant.mean():.3f}, '
      f'5th-95th [{np.percentile(irrelevant, 5):.3f}, '
      f'{np.percentile(irrelevant, 95):.3f}]')
print()
print(f'{"threshold":>10} {"recall":>8} {"precision":>10} '
      f'{"missed":>8} {"false":>7}')
for threshold in [0.20, 0.30, 0.40, 0.45, 0.50, 0.60, 0.70]:
    kept_relevant = int((relevant >= threshold).sum())
    kept_irrelevant = int((irrelevant >= threshold).sum())
    recall = kept_relevant / len(relevant)
    precision = (kept_relevant / (kept_relevant + kept_irrelevant)
                 if kept_relevant + kept_irrelevant else 0.0)
    print(f'{threshold:>10.2f} {recall:>8.1%} {precision:>10.1%} '
          f'{len(relevant) - kept_relevant:>8} {kept_irrelevant:>7}')
print()
print('There is no threshold that keeps everything relevant and')
print('nothing else, because the two distributions overlap. Every')
print('choice is a trade, and which way to trade depends entirely')
print('on what happens next:')
print()
print(' - feeding a model that must answer only from the')
print('   documents: favour PRECISION. An irrelevant chunk')
print('   invites a wrong answer.')
print(' - feeding a reranker that will filter again: favour')
print('   RECALL. Anything you drop here is gone for good.')
print(' - showing results to a person: favour precision. People')
print('   judge a search by its first wrong answer.')
print()
print('And the number itself does not transfer. A cosine of 0.45')
print('means different things in different embedding spaces, so')
print('the threshold has to be measured on YOUR data with YOUR')
print('model - and re-measured the day you change either.')

import numpy as np

cases = [
    ('Exact identifiers',
     'order 10294-B', 'order 10295-B',
     'Nearly identical vectors, completely different orders. '
     'Embeddings capture meaning, and an order number has none.'),
    ('Negation',
     'refunds are available for digital goods',
     'refunds are not available for digital goods',
     'One word apart, opposite meanings, and cosine barely moves.'),
    ('Rare proper nouns',
     'the Thamesford depot closes at six',
     'the Hartwell depot closes at six',
     'The names are the only difference and the only thing that '
     'matters.'),
    ('Numbers and quantities',
     'orders over 50 pounds ship free',
     'orders over 500 pounds ship free',
     'A factor of ten, nearly the same vector.'),
]
for title, a, b, note in cases:
    print(f'{title}')
    print(f'  "{a}"')
    print(f'  "{b}"')
    print(f'  -> {note}')
    print()

# Demonstrated, on character overlap as a proxy for what a
# bag-of-terms embedding sees.
def term_overlap(a, b):
    first, second = set(a.lower().split()), set(b.lower().split())
    return len(first & second) / len(first | second)


print('shared-term overlap for each pair above:')
for title, a, b, _ in cases:
    print(f'  {title:<22} {term_overlap(a, b):.2f}')
print()
print('Between 0.6 and 0.8 shared terms in every case - which is')
print('what any embedding of these sentences will see, because the')
print('distinguishing token is one word out of seven.')
print()
print('The fixes are all the same shape: do not ask the embedding')
print('to do the part it cannot.')
print()
print('  identifiers  -> exact or keyword match, not vectors')
print('  negation     -> a reranker, or a model that reads the')
print('                  chunk rather than its vector')
print('  proper nouns -> hybrid search, with the keyword half')
print('                  carrying the name')
print('  numbers      -> extract them into metadata and filter')
print()
print('This is why production retrieval is almost always HYBRID:')
print('a vector search for meaning, a keyword search for exact')
print('terms, and the two result sets combined. The next lesson')
print('builds one.')

import numpy as np

DOCUMENTS = [
    'Refunds for physical goods are available within 30 days of delivery.',
    'Digital downloads and ebooks cannot be refunded once the file has '
    'been accessed.',
    'Software licences are non-refundable after activation.',
    'Shipping charges are refunded only when the item arrived damaged '
    'or faulty.',
    'For faulty items we pay the return postage ourselves.',
    'Gift cards are non-refundable and cannot be exchanged for cash.',
    'Refunded money is paid back to the original card within five '
    'working days.',
    'To start a return, contact support with your order number.',
    'Returned items must have their original packaging intact.',
    'Our warehouse dispatches orders on the next working day.',
    'Standard delivery takes two to four working days within the UK.',
    'Express delivery arrives the next working day if ordered before noon.',
    'International delivery takes seven to fourteen working days.',
    'Delivery to the Channel Islands may take longer and incurs duty.',
    'You can change your delivery address before the order is dispatched.',
    'Track your parcel using the link in your dispatch email.',
    'A missing parcel should be reported after ten working days.',
    'Orders over fifty pounds qualify for free standard delivery.',
    'We do not deliver to PO boxes or forwarding addresses.',
    'Signature on delivery is required for orders over two hundred pounds.',
    'Your account password can be reset from the sign-in page.',
    'Two-factor authentication can be enabled in account settings.',
    'Invoices are available to download from your order history.',
    'VAT receipts are issued automatically for business accounts.',
]

# The labelled set: a query, and which document indices answer it.
# Twenty-one of these took about ten minutes to write down.
LABELLED = [
    ('return a physical item', {0}),
    ('refund a downloaded file', {1}),
    ('licence refund after activation', {2}),
    ('postage refunded damaged item', {3}),
    ('return postage faulty', {4}),
    ('gift card refund', {5}),
    ('how soon is money paid back', {6}),
    ('begin a return', {7}),
    ('packaging needed for returns', {8}),
    ('express next working day', {11}),
    ('international delivery time', {12, 13}),
    ('change delivery address', {14}),
    ('track my parcel', {15}),
    ('parcel never arrived', {16}),
    ('free delivery threshold', {17}),
    ('signature required', {19}),
    ('password reset', {20}),
    ('two factor authentication', {21}),
    ('download an invoice', {22}),
    ('vat receipt', {23}),
    ('refund rules', {0, 1, 2, 3, 5}),
]

STOP = set(('the a an is are of to in for and or be been can cannot '
            'your you with only once has if before within on our from '
            'do does how my').split())


def tokenise(text):
    cleaned = ''.join(c.lower() if c.isalnum() else ' ' for c in text)
    return [w for w in cleaned.split() if w not in STOP and len(w) > 2]


vocabulary = sorted({w for d in DOCUMENTS for w in tokenise(d)})
position = {w: i for i, w in enumerate(vocabulary)}
tf = np.zeros((len(DOCUMENTS), len(vocabulary)))
for row, document in enumerate(DOCUMENTS):
    for word in tokenise(document):
        tf[row, position[word]] += 1
idf = np.log((1 + len(DOCUMENTS)) / (1 + (tf > 0).sum(axis=0))) + 1.0
tfidf = tf * idf
u, singular, vt = np.linalg.svd(tfidf, full_matrices=False)


def build(dimensions):
    reduced = u[:, :dimensions] * singular[:dimensions]
    norms = np.linalg.norm(reduced, axis=1, keepdims=True)
    return reduced / np.where(norms > 0, norms, 1), vt[:dimensions].T


def rank(query, index, projection):
    row = np.zeros(len(vocabulary))
    for word in tokenise(query):
        if word in position:
            row[position[word]] += 1
    vector = (row * idf) @ projection
    norm = np.linalg.norm(vector)
    if norm > 0:
        vector = vector / norm
    return list(np.argsort(index @ vector)[::-1])


def evaluate(dimensions, k_values=(1, 3, 5)):
    index, projection = build(dimensions)
    results = {}
    for k in k_values:
        found = 0
        wanted = 0
        for query, relevant in LABELLED:
            found += len(relevant & set(rank(query, index, projection)[:k]))
            # The denominator is how many relevant documents EXIST,
            # not how many could fit in k. Dividing by min(len, k)
            # makes recall fall as k rises, which is not a metric.
            wanted += len(relevant)
        results[f'recall@{k}'] = found / wanted

    # Mean reciprocal rank: how far down the first correct answer is.
    reciprocals = []
    for query, relevant in LABELLED:
        ranked = rank(query, index, projection)
        hit = next((i for i, d in enumerate(ranked, 1) if d in relevant),
                   None)
        reciprocals.append(1 / hit if hit else 0.0)
    results['MRR'] = float(np.mean(reciprocals))
    return results


print(f'{len(LABELLED)} labelled queries over {len(DOCUMENTS)} documents, '
      f'{len(vocabulary)} terms')
print()
print(f'{"dimensions":>11} {"recall@1":>9} {"recall@3":>9} '
      f'{"recall@5":>9} {"MRR":>7}')
for dimensions in [2, 4, 8, 12, 16, 20, 24]:
    scores = evaluate(dimensions)
    print(f'{dimensions:>11} {scores["recall@1"]:>9.2f} '
          f'{scores["recall@3"]:>9.2f} {scores["recall@5"]:>9.2f} '
          f'{scores["MRR"]:>7.3f}')
print()
print('Four things to read off that table, and one of them is')
print('uncomfortable.')
print()
print('recall@3 beats recall@1 everywhere, by 11 points at the')
print('best setting and by 35 at the worst. The right document is')
print('usually in the top few even when it is not first, which is')
print('the whole argument for retrieving several chunks and')
print('letting the model or a reranker choose between them.')
print()
print('Two dimensions is not enough. recall@1 of 0.15 means the')
print('right document came first for three queries out of')
print('twenty-one, because twenty-four documents cannot be told')
print('apart in a plane. Note that recall@5 is already 0.62 at')
print('two dimensions: even a badly compressed space gets the')
print('answer into the top five most of the time, which is how a')
print('broken index passes a casual review.')
print()
print('The gain flattens. 8, 12 and 16 dimensions score')
print('identically on recall; only MRR distinguishes them, and')
print('barely. Every extra dimension costs memory and query time')
print('forever, so the flat part of this curve is where to ship.')
print()
print('And the uncomfortable one: 24 dimensions - which for')
print('24 documents is NO REDUCTION AT ALL - scores best, at 0.69')
print('and 0.899. On a corpus this small the SVD only loses')
print('information. Its value is compression, and compression')
print('only pays when there is something to compress. If your')
print('corpus is a few thousand chunks, skip the reduction and')
print('search the full vectors.')
print()
print('MRR is the number to watch over time. It moves when a')
print('change pushes correct answers up or down the list, which')
print('recall@5 cannot see at all - and a regression that shows')
print('up only in the ordering is exactly the kind that ships.')
print()
print('Twenty-one labelled queries took ten minutes to write and')
print('they turn every later change from an opinion into a')
print('measurement. Do this before tuning anything.')

# 1. Changing the embedding model means reindexing EVERYTHING.
#    Vectors from two models are not comparable, and mixing them
#    does not error - it silently returns nonsense.
import numpy as np

rng = np.random.default_rng(0)
model_a = rng.normal(0, 1, 8)
model_a = model_a / np.linalg.norm(model_a)
model_b = rng.normal(0, 1, 8)
model_b = model_b / np.linalg.norm(model_b)

print('the same sentence, embedded by two different models:')
print('  model A:', model_a.round(3))
print('  model B:', model_b.round(3))
print(f'  cosine between them: {float(model_a @ model_b):+.3f}')
print()
print('Both are unit vectors of the same width, so every vector')
print('database will accept both and compare them happily. The')
print('similarity is meaningless.')
print()
print('So: store the model name and version alongside every')
print('vector, refuse to query an index with a mismatched model,')
print('and budget for a full reindex whenever you upgrade. A')
print('million documents is a real job, not a configuration')
print('change.')
print()

# 2. The cost of an index, which decides the architecture.
print('index size for various corpus sizes, at 1,536 dimensions:')
for documents in [1000, 100000, 10000000, 1000000000]:
    float32 = documents * 1536 * 4
    float16 = documents * 1536 * 2
    int8 = documents * 1536
    print(f'  {documents:>13,} chunks: '
          f'{float32 / 1024 ** 3:8.1f} GiB float32, '
          f'{float16 / 1024 ** 3:8.1f} GiB float16, '
          f'{int8 / 1024 ** 3:8.1f} GiB int8')
print()
print('A million chunks in float32 is 5.7 GiB, which fits in')
print('memory on one machine and needs no vector database at all -')
print('a numpy array and a dot product will do it in milliseconds.')
print()
print('A billion chunks is 5.7 TiB and needs real infrastructure:')
print('quantisation, approximate nearest neighbours, sharding.')
print('Know which of those two you are before choosing a tool,')
print('because most projects are the first one and buy the')
print('second.')

import numpy as np

rng = np.random.default_rng(1)

# Two candidate spaces. One separates topics; one has collapsed,
# which happens when embeddings are trained or averaged badly.
def make_space(separation):
    centres = rng.normal(0, 1, (3, 16)) * separation
    points, labels = [], []
    for topic, centre in enumerate(centres):
        for _ in range(40):
            vector = centre + rng.normal(0, 1, 16)
            points.append(vector / np.linalg.norm(vector))
            labels.append(topic)
    return np.array(points), np.array(labels)


print(f'{"space":<22} {"same-topic":>11} {"other-topic":>12} '
      f'{"gap":>7}  verdict')
for label, separation in [('well separated', 2.5), ('weak', 0.8),
                          ('collapsed', 0.15)]:
    points, labels = make_space(separation)
    similarity = points @ points.T
    same = np.equal.outer(labels, labels)
    np.fill_diagonal(same, False)
    off_diagonal = ~np.eye(len(points), dtype=bool)
    within = similarity[same].mean()
    between = similarity[off_diagonal & ~same].mean()
    gap = within - between
    verdict = ('usable' if gap > 0.25 else
               'marginal' if gap > 0.10 else
               'USELESS - everything is similar to everything')
    print(f'{label:<22} {within:>11.3f} {between:>12.3f} '
          f'{gap:>7.3f}  {verdict}')
print()
print('Look at the "weak" row. Its same-topic similarity is 0.354,')
print('which is a respectable-looking number and the one people')
print('quote. On its own it says nothing: what makes that space')
print('usable is that its other-topic similarity is -0.012, so')
print('the GAP is 0.366.')
print()
print('Then look at "collapsed". Same-topic 0.024, other-topic')
print('-0.002, gap 0.026. Every document is equally similar to')
print('every other, so the ranking is noise and the top result is')
print('arbitrary. An index in that state returns confident')
print('nonsense and no threshold saves it.')
print()
print('Run this on your own index with twenty queries whose')
print('answers you know:')
print()
print('  1. compute the mean similarity of each query to its')
print('     correct document;')
print('  2. compute the mean similarity of each query to a random')
print('     document;')
print('  3. subtract.')
print()
print('If the gap is under 0.1, no amount of threshold tuning will')
print('help, because there is no threshold that separates the two')
print('populations. The problem is upstream: the chunks are too')
print('long, the model is wrong for the domain, or the text was')
print('normalised away. That diagnosis takes three lines of numpy')
print('and saves a week of tuning retrieval parameters that')
print('cannot work.')