AI/LLM Cheatsheet 04 — Embeddings

response = client.embeddings.create(
    model="text-embedding-3-small",      # 1536 dim, cheap
    # model="text-embedding-3-large",    # 3072 dim
    input=["text 1", "text 2"],
)

vectors = [d.embedding for d in response.data]

import cohere
co = cohere.Client(...)

response = co.embed(
    texts=["text 1", "text 2"],
    model="embed-english-v3.0",
    input_type="search_document",       # or search_query
)

import voyageai
vo = voyageai.Client()
result = vo.embed(["text"], model="voyage-3", input_type="document")

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-base-en-v1.5")
embeddings = model.encode(["text 1", "text 2"], normalize_embeddings=True)

import numpy as np

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

# If normalized: dot product alone

CREATE EXTENSION vector;

CREATE TABLE docs (
    id SERIAL PRIMARY KEY,
    text TEXT,
    embedding vector(1536)
);

CREATE INDEX ON docs USING hnsw (embedding vector_cosine_ops);

-- Insert
INSERT INTO docs (text, embedding) VALUES ('...', '[0.1, 0.2, ...]');

-- Query
SELECT text, embedding <=> '[...]' AS distance
FROM docs ORDER BY distance LIMIT 5;

from qdrant_client import QdrantClient

client = QdrantClient(":memory:")
client.create_collection("docs", vectors_config={"size": 1536, "distance": "Cosine"})

client.upsert(collection_name="docs", points=[
    {"id": 1, "vector": [...], "payload": {"text": "..."}},
])

results = client.search(collection_name="docs", query_vector=[...], limit=5)

import chromadb
client = chromadb.PersistentClient(path="./db")
coll = client.get_or_create_collection("docs")

coll.add(ids=["1"], embeddings=[[...]], metadatas=[{"src": "..."}], documents=["text"])
results = coll.query(query_embeddings=[[...]], n_results=5)

def chunk(text, chunk_size=500, overlap=50):
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        chunk = " ".join(words[i:i+chunk_size])
        chunks.append(chunk)
        i += chunk_size - overlap
    return chunks

client.embeddings.create(
    model="text-embedding-3-large",
    input=["..."],
    dimensions=512,                     # cut from 3072
)

from rank_bm25 import BM25Okapi

bm25_scores = bm25.get_scores(tokenized_query)
vector_scores = ...

# Reciprocal Rank Fusion
def rrf(rankings, k=60):
    scores = {}
    for ranking in rankings:
        for rank, doc_id in enumerate(ranking):
            scores[doc_id] = scores.get(doc_id, 0) + 1 / (k + rank + 1)
    return sorted(scores.items(), key=lambda x: -x[1])

from sentence_transformers import CrossEncoder

reranker = CrossEncoder("BAAI/bge-reranker-large")
pairs = [(query, doc) for doc in candidates]
scores = reranker.predict(pairs)
ranked = sorted(zip(candidates, scores), key=lambda x: -x[1])