Vector Embeddings from Scratch with PyTorch and Sentence-Transformers

#pgvector #Python AI #PyTorch #rag #semantic search #sentence-transformers #vector embeddings

Vector Embeddings are the fundamental mathematical unit of modern semantic search engines, Retrieval-Augmented Generation (RAG) pipelines, and AI recommendation systems. An embedding model maps raw text strings (words, sentences, or paragraphs) into dense high-dimensional floating-point vectors (e.g., 384 or 1,536 numbers) where geometric distance between vectors correlates directly with semantic meaning similarity.

This technical guide teaches software engineers how vector embeddings work under the hood in PyTorch, how to use the sentence-transformers library, how mean pooling mathematics convert token matrices into sentence vectors, and how to store embeddings efficiently inside PostgreSQL using pgvector.

Building production-ready software systems around artificial intelligence demands a rigorous understanding of system design and execution parameters. Unlike traditional software modules where deterministic input produces predictable output, probabilistic models require strict boundary controls, context management, and defensive error handling. Operating at scale requires software engineers to establish reliable operational contracts between application backends and language model endpoints.

Furthermore, engineering teams must evaluate the latency, throughput, and financial costs of model invocation. Every request introduces network I/O overhead and computational processing costs. Implementing structured input validation, semantic caching, and dynamic model routing allows developers to optimize resource utilization while maintaining high service reliability across diverse production workloads.

1. Mathematical Principles of Semantic Vector Representations

Vector embeddings convert unstructured text documents into dense numerical vector representations that capture semantic relationships in high-dimensional vector spaces. Traditional keyword search algorithms rely on exact character matching, failing to identify conceptually relevant documents that use alternative terminology. Embedding models resolve keyword limitations by mapping semantically similar phrases to close geometric coordinates.

Passing text through transformer neural networks generates token embedding matrices across attention layers. Applying mean pooling across non-padded token vectors yields a single normalized vector representing document content. Normalizing output vectors enables fast cosine similarity calculations using dot product matrix operations in vector databases.

1. The Mathematics of Text Embeddings

Traditional search engines (like Elasticsearch or SQL LIKE queries) rely on exact keyword matching. If a user searches for "automobile repair," a keyword engine misses documents containing "car fix" because the character strings do not match. Vector embedding models solve this by encoding the underlying **semantic concept** into high-dimensional space.

When passing text through a Transformer model (like BERT or RoBERTa), the output is a 3D matrix of token embeddings: [batch_size, sequence_length, hidden_dimension]. To collapse this matrix into a single vector representing the entire sentence, we apply **Mean Pooling** across the non-padded tokens.

import torch
import torch.nn.functional as F

# Manual PyTorch Mean Pooling Implementation
def calculate_mean_pooling(model_output: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
    # Extract token embeddings (last hidden state)
    token_embeddings = model_output[0] # Shape: [batch_size, seq_len, hidden_dim]
    
    # Expand attention mask to match token embedding dimensions
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    
    # Sum embeddings along sequence axis (ignoring padding tokens)
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, dim=1)
    
    # Divide by total non-padded token count
    sum_mask = torch.clamp(input_mask_expanded.sum(dim=1), min=1e-9)
    
    sentence_vectors = sum_embeddings / sum_mask
    
    # Normalize vectors to unit length (L2 norm) for fast Cosine Similarity calculation
    return F.normalize(sentence_vectors, p=2, dim=1)

Defining system instructions as software contracts ensures that language models operate within predefined operational boundaries. In enterprise applications, system prompts are version-controlled, modularized, and tested using automated evaluation pipelines. Decoupling system instructions into distinct role specifications, behavioral constraints, and output format requirements minimizes generation variance and improves system predictability.

Additionally, system prompt architecture must account for unexpected user inputs and edge cases. Explicitly defining fallback rules and error responses prevents application crashes and provides end users with clear, actionable feedback. Treating prompt configurations with the same discipline as traditional source code ensures long-term system maintainability.

2. Streamlined Vector Generation with Sentence-Transformers

While writing manual PyTorch pooling code provides complete mathematical control, the open-source sentence-transformers library wraps these operations into a clean, high-performance API:

from sentence_transformers import SentenceTransformer
import numpy as np

# Load Pre-Trained Embedding Model (384 Dimensions)
model = SentenceTransformer("all-MiniLM-L6-v2")

sentences = [
    "How do I reset my user account password?",
    "Where is the password reset link located?",
    "What is your refund policy for annual software subscriptions?"
]

# Generate Vector Embeddings in Batches
embeddings = model.encode(sentences, batch_size=32, show_progress_bar=False)

print(f"Generated Embeddings Shape: {embeddings.shape}") # (3, 384)

# Calculate Cosine Similarity Matrix (Dot Product of Normalized Vectors)
similarity_matrix = np.dot(embeddings, embeddings.T)

print(f"Similarity (Query 1 vs Query 2): {similarity_matrix[0][1]:.4f}") # High Similarity (~0.88)
print(f"Similarity (Query 1 vs Query 3): {similarity_matrix[0][2]:.4f}") # Low Similarity (~0.12)

Structured output generation is fundamental to integrating language models into automated microservice architectures. Backend APIs require strictly formatted JSON payloads matching predefined data schemas to execute downstream database operations and external service calls. Relying on natural language formatting instructions often leads to broken parsers and system failures.

Leveraging native JSON schema enforcement and function calling guarantees schema compliance at the API SDK level. By constraining output decoding logits, developers ensure that returned payloads adhere strictly to expected field types, enabling seamless ingestion into backend application pipelines.

3. High-Throughput Batch Embedding Pipelines and Parallel Vectorization

Generating vector embeddings across enterprise document repositories demands efficient batch processing pipelines. Executing model inference across batch arrays of text passages maximizes GPU parallel matrix computation speed, preventing system bottlenecks during document ingestion. Managing batch memory allocation prevents server VRAM exhaustion under high data volumes.

Incorporating worker queues and asynchronous pipeline processing keeps document ingestion workflows responsive. Normalizing output vectors enables rapid cosine similarity search using native vector operations in PostgreSQL databases.

3. High-Throughput Batch Embedding Pipelines

When embedding millions of document chunks during initial RAG document ingestion, calling model inference one sentence at a time creates severe CPU/GPU bottleneck overhead. Process inputs in optimized **batch arrays of 128 to 512 items**:

def process_large_document_corpus(documents: list[str], batch_size: int = 128) -> list[list[float]]:
    all_vectors = []
    
    for i in range(0, len(documents), batch_size):
        batch_texts = documents[i:i + batch_size]
        batch_vectors = model.encode(batch_texts, batch_size=batch_size, normalize_embeddings=True)
        all_vectors.extend(batch_vectors.tolist())
        
    return all_vectors

Security engineering for AI-powered backends requires protecting applications against prompt injection attacks, data leakage, and unauthorized function execution. Adversarial users often attempt to bypass system constraints by embedding malicious commands within user input fields. Establishing strict architectural isolation between system instructions and user parameters is essential to maintaining system security.

Implementing multi-layered defense mechanisms โ€” including input sanitization, character length limits, pre-execution content moderation, and fine-grained authorization checks โ€” mitigates security risks. Developers must treat all model-generated parameters as untrusted input, verifying permissions before executing database updates or external API calls.

4. Storing & Querying Embeddings in PostgreSQL (pgvector)

Once PyTorch computes vector embeddings in your Python service, persist them directly inside PostgreSQL using the pgvector extension for fast nearest-neighbor SQL search:

-- SQL Schema: Store 384-dimensional sentence-transformer vectors
CREATE EXTENSION IF NOT EXISTS vector;

CREATE TABLE document_embeddings (
    id SERIAL PRIMARY KEY,
    document_title VARCHAR(255),
    chunk_content TEXT,
    embedding vector(384) -- Matches all-MiniLM-L6-v2 dimensions!
);

-- Cosine Distance Query in PostgreSQL
SELECT chunk_content, 
       1 - (embedding <=> '[0.012, -0.045, ...]'::vector) AS similarity_score
FROM document_embeddings
ORDER BY embedding <=> '[0.012, -0.045, ...]'::vector ASC
LIMIT 5;

Architecting robust AI integrations combines modern backend development patterns with specialized machine learning engineering discipline. By prioritizing deterministic schema enforcement, multi-layered security controls, efficient prompt design, and continuous evaluation, software developers create scalable AI solutions that deliver consistent enterprise value.

Conclusion

Understanding vector embeddings from mathematical principles in PyTorch to high-level execution with Sentence-Transformers gives software engineers complete mastery over modern semantic search and RAG architecture. Storing normalized vectors in pgvector enables scalable AI retrieval pipelines.

Need assistance building vector embedding engines or pgvector search pipelines? Visit my AI integration development page or Hire AI Developer.


Prakash Tank

Prakash Tank

Full-Stack Architect & Tech Enthusiast. Passionate about building scalable applications and sharing knowledge with the community.