import torch
import torch.nn as nn
import torch.optim as optim
import numpy as np
from typing import List, Tuple, Dict

class LexicalTopology(nn.Module):
    """Learn dense word representations via center-context prediction."""
    
    def __init__(self, vocab_size: int, embedding_dim: int):
        super().__init__()
        self.target_proj = nn.Embedding(vocab_size, embedding_dim)
        self.context_proj = nn.Embedding(vocab_size, embedding_dim)
        self._init_weights()
        
    def _init_weights(self):
        for module in self.modules():
            if isinstance(module, nn.Embedding):
                nn.init.xavier_uniform_(module.weight)
                
    def forward(self, target_ids: torch.Tensor, context_ids: torch.Tensor) -> torch.Tensor:
        t_vec = self.target_proj(target_ids)
        c_vec = self.context_proj(context_ids)
        return (t_vec * c_vec).sum(dim=1)
    
    def extract_lexical_vectors(self) -> torch.Tensor:
        return self.target_proj.weight.detach().clone()

def build_cooccurrence_pairs(tokens: List[str], window_size: int, vocab_map: Dict[str, int]) -> List[Tuple[int, int]]:
    pairs = []
    for idx, center in enumerate(tokens):
        if center not in vocab_map:
            continue
        start = max(0, idx - window_size)
        end = min(len(tokens), idx + window_size + 1)
        for ctx_idx in range(start, end):
            if ctx_idx != idx and tokens[ctx_idx] in vocab_map:
                pairs.append((vocab_map[center], vocab_map[tokens[ctx_idx]]))
    return pairs

def train_semantic_space(model: LexicalTopology, pairs: List[Tuple[int, int]], epochs: int, lr: float = 0.02):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    
    target_tensor = torch.tensor([p[0] for p in pairs], dtype=torch.long, device=device)
    context_tensor = torch.tensor([p[1] for p in pairs], dtype=torch.long, device=device)
    positive_labels = torch.ones(len(pairs), device=device)
    
    optimizer = optim.Adam(model.parameters(), lr=lr)
    loss_fn = nn.BCEWithLogitsLoss()
    
    for epoch in range(epochs):
        optimizer.zero_grad()
        logits = model(target_tensor, context_tensor)
        loss = loss_fn(logits, positive_label

s) loss.backward() optimizer.step()

    if (epoch + 1) % 50 == 0:
        print(f"Epoch {epoch+1}/{epochs} | Loss: {loss.item():.4f}")
        
return model

**Why this works:** The model learns to maximize the dot product between center and context vectors while minimizing it for non-occurring pairs. Over hundreds of epochs, the optimization landscape forces semantically related tokens into neighboring regions of the embedding space.

### Step 3: Query the Semantic Manifold

Cosine similarity is the standard metric for lexical proximity. Unlike Euclidean distance, it measures angular alignment, which is invariant to vector magnitude and better captures directional semantic relationships.

```python
def angular_similarity(vec_a: np.ndarray, vec_b: np.ndarray) -> float:
    dot_product = np.dot(vec_a, vec_b)
    norm_a = np.linalg.norm(vec_a)
    norm_b = np.linalg.norm(vec_b)
    return dot_product / (norm_a * norm_b + 1e-9)

def retrieve_nearest_neighbors(query_word: str, vocab_map: Dict[str, int], 
                               vectors: torch.Tensor, top_k: int = 5) -> List[Tuple[str, float]]:
    if query_word not in vocab_map:
        raise ValueError(f"Unknown token: {query_word}")
    
    query_vec = vectors[vocab_map[query_word]].cpu().numpy()
    all_vecs = vectors.cpu().numpy()
    
    similarities = [angular_similarity(query_vec, v) for v in all_vecs]
    ranked_indices = np.argsort(similarities)[::-1]
    
    results = []
    for idx in ranked_indices:
        word = [w for w, i in vocab_map.items() if i == idx][0]
        if word != query_word:
            results.append((word, similarities[idx]))
            if len(results) == top_k:
                break
    return results

Step 4: Integrate Pretrained and Contextual Models

Training embeddings from scratch requires massive corpora and careful hyperparameter tuning. Production systems typically load pretrained static vectors or switch to contextual transformers.

Static Embeddings (Gensim):

import gensim.downloader as api

# Load pre-optimized GloVe vectors trained on 6B tokens
glove_vectors = api.load("glove-wiki-gigaword-100")

# Semantic lookup
similar_terms = glove_vectors.most_similar("algorithm", topn=5)
print("Static similarity:", similar_terms)

# Lexical arithmetic
analogy_result = glove_vectors.most_similar(positive=["doctor", "woman"], negative=["man"])
print("Vector arithmetic:", analogy_result)

Contextual Embeddings (HuggingFace):

from transformers import AutoTokenizer, AutoModel
import torch

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
context_model = AutoModel.from_pretrained("bert-base-uncased")
context_model.eval()

sample_sentences = [
    "The engineer deployed the server cluster.",
    "The restaurant server cleared the table."
]

contextual_vectors = []
for sentence in sample_sentences:
    inputs = tokenizer(sentence, return_tensors="pt", truncation=True)
    with torch.no_grad():
        outputs = context_model(**inputs)
    
    token_ids = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
    target_pos = token_ids.index("server")
    vec = outputs.last_hidden_state[0, target_pos, :].numpy()
    contextual_vectors.append(vec)
    print(f"Context: '{sentence}' | Vector norm: {np.linalg.norm(vec):.4f}")

cross_context_sim = angular_similarity(contextual_vectors[0], contextual_vectors[1])
print(f"Same token, different context similarity: {cross_context_sim:.4f}")

Contextual models generate token-specific representations conditioned on the entire sequence. The word "server" receives distinct vectors in technical vs. hospitality contexts, resolving polysemy that static embeddings cannot handle.

Pitfall Guide

1. The Polysemy Blind Spot

Explanation: Static embeddings assign one vector per token type. Words with multiple meanings ("lead", "crane", "bank") collapse into a single averaged representation, degrading performance in domain-specific tasks. Fix: Switch to contextual embeddings (BERT, RoBERTa, or LLM encoders) when polysemy impacts downstream accuracy. Use static vectors only for monosemous domains or high-throughput similarity search.

2. Distance Metric Mismatch

Explanation: Using Euclidean or Manhattan distance on raw embedding vectors penalizes magnitude differences that are semantically irrelevant. This causes nearest-neighbor queries to return outliers. Fix: Always normalize vectors before distance computation, or use cosine similarity directly. Cosine similarity measures directional alignment, which correlates with semantic relatedness.

3. Dimensionality Overprovisioning

Explanation: Defaulting to 300D or 768D embeddings increases memory bandwidth, cache misses, and inference latency without proportional accuracy gains for simple classification or clustering tasks. Fix: Benchmark 50D, 100D, and 200D variants on your validation set. Most tabular NLP and retrieval tasks saturate at 100D. Use higher dimensions only when capturing fine-grained syntactic or cross-lingual relationships.

4. Unhandled Out-of-Vocabulary (OOV) Tokens

Explanation: Pretrained vocabularies rarely cover domain-specific jargon, neologisms, or misspellings. Unhandled OOV tokens cause index errors or fallback to zero vectors, breaking similarity calculations. Fix: Implement subword tokenization (Byte-Pair Encoding or WordPiece) or map unknown tokens to a learned <UNK> vector. For production, maintain a dynamic vocabulary extension layer that assigns random projections to new tokens and fine-tunes them on domain data.

5. Context Window Misconfiguration

Explanation: Training Skip-gram with a window size of 1 captures only immediate adjacency, missing syntactic dependencies. A window of 10+ introduces noise from unrelated clauses, diluting semantic signals. Fix: Use a window size of 3–5 for general text. Adjust based on domain: legal/medical text benefits from larger windows due to long-range dependencies; code or log data performs better with smaller windows.

6. Ignoring Text Normalization

Explanation: Feeding raw text with mixed casing, punctuation, and HTML artifacts inflates vocabulary size and fragments co-occurrence statistics. "API", "api", and "Api" become three separate tokens. Fix: Apply consistent lowercasing, strip punctuation, and normalize whitespace before vocabulary construction. Preserve case only when it carries semantic weight (e.g., proper nouns, code identifiers).

7. Fine-Tuning Without Freezing Strategy

Explanation: Updating embedding weights during downstream task training without proper regularization causes catastrophic forgetting. The model overfits to task-specific labels and loses general semantic structure. Fix: Freeze static embeddings for small datasets. For contextual models, use layer-wise learning rate decay or LoRA adapters. Always validate embedding stability by monitoring cosine similarity drift on a held-out lexical benchmark.

Production Bundle

Action Checklist

• Normalize text casing and punctuation before vocabulary construction to prevent token fragmentation
• Select cosine similarity as the default distance metric; avoid Euclidean distance on raw vectors
• Benchmark embedding dimensionality (50D vs 100D vs 200D) against your specific latency and accuracy SLAs
• Implement OOV handling via subword tokenization or dynamic vocabulary extension
• Choose static embeddings for high-throughput similarity search; switch to contextual models for polysemous or syntactically complex domains
• Freeze embedding weights during initial downstream training; enable fine-tuning only after baseline convergence
• Validate semantic topology using lexical analogy benchmarks before deploying to production

Decision Matrix

Scenario	Recommended Approach	Why	Cost Impact
Real-time semantic search (<50ms)	Static embeddings (GloVe/FastText)	Cacheable, O(1) lookup, minimal compute	Low memory, negligible CPU
Financial/medical NLP with polysemy	Contextual embeddings (BERT/RoBERTa)	Resolves context-dependent meanings	Higher GPU memory, 5-20x latency
Low-resource language or domain	Subword static embeddings + domain fine-tuning	Handles unseen tokens, adapts to jargon	Moderate training cost, high recall
High-throughput log parsing	Static 50D embeddings + approximate nearest neighbor	Balances speed and semantic grouping	Minimal infrastructure, scales horizontally
Cross-lingual retrieval	Multilingual contextual model (XLM-R)	Aligns semantic spaces across languages	High initial load, unified pipeline

Configuration Template

# production_embedding_config.py
import torch
from transformers import AutoTokenizer, AutoModel
import numpy as np

class EmbeddingPipeline:
    def __init__(self, model_name: str = "bert-base-uncased", device: str = "cpu"):
        self.device = torch.device(device)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(self.device).eval()
        self.model.requires_grad_(False)  # Freeze for inference
        
    @torch.no_grad()
    def encode(self, texts: list[str]) -> np.ndarray:
        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt").to(self.device)
        outputs = self.model(**inputs)
        # Mean pooling over token dimension
        mask = inputs["attention_mask"].unsqueeze(-1)
        pooled = (outputs.last_hidden_state * mask).sum(dim=1) / mask.sum(dim=1)
        return pooled.cpu().numpy()
    
    @staticmethod
    def cosine_similarity_matrix(vectors: np.ndarray) -> np.ndarray:
        norms = np.linalg.norm(vectors, axis=1, keepdims=True)
        normalized = vectors / (norms + 1e-9)
        return normalized @ normalized.T

Quick Start Guide

1. Install dependencies: pip install torch transformers numpy
2. Initialize the pipeline: Instantiate EmbeddingPipeline(device="cpu") or "cuda" for GPU acceleration.
3. Encode text: Pass a list of strings to .encode(). Returns a normalized matrix ready for similarity computation.
4. Query semantics: Use .cosine_similarity_matrix() to compute pairwise scores, or extract row vectors for nearest-neighbor search with FAISS or Annoy.
5. Validate: Run a lexical sanity check ("doctor" vs "physician" vs "hammer") to confirm semantic alignment before integrating into your application layer.