nuonuo/experiments/exp04_real_embeddings.py

"""Experiment 4: End-to-end with real sentence embeddings.

All previous experiments used random vectors. Now test with actual semantic
embeddings from a sentence transformer model. Key questions:

1. Does pattern separation preserve semantic neighborhoods?
   (Similar sentences → similar/related codes?)
2. Can we retrieve memories using paraphrased/related queries?
3. Does the multi-hop chaining work with semantic embeddings?
4. Noise tolerance: does embedding-space noise behave differently?
5. Does a learned separator trained on real data improve noise tolerance?
"""

import sys
import time
import json
from pathlib import Path

import torch
import torch.nn as nn
import numpy as np

DEVICE = "cuda"
RESULTS_DIR = Path(__file__).parent.parent / "doc"


def cosine(a, b):
    if a.norm() == 0 or b.norm() == 0:
        return 0.0
    return nn.functional.cosine_similarity(a.unsqueeze(0), b.unsqueeze(0)).item()


def winner_take_all(x, k):
    _, idx = x.topk(k, dim=-1)
    out = torch.zeros_like(x)
    out.scatter_(-1, idx, 1.0)
    return out


# --- Test data: conversation-like memory pairs ---
MEMORY_PAIRS = [
    # (context/cue, memory/fact to recall)
    ("What's the weather like today?", "User prefers to check weather every morning"),
    ("Let's deploy the new version", "The deployment pipeline uses GitHub Actions with k3s"),
    ("The database is slow again", "Last time DB was slow it was because of missing index on users table"),
    ("Can you review my pull request?", "User prefers small PRs with clear commit messages"),
    ("I need to fix the authentication bug", "Auth service uses JWT tokens with 24h expiry stored in Redis"),
    ("Let's set up monitoring", "Prometheus + Grafana stack is already running on the OCI cluster"),
    ("The API is returning 500 errors", "Last 500 error was caused by OOM in the Python worker"),
    ("I want to learn Rust", "User has strong Python and Go background, new to systems programming"),
    ("Schedule a meeting with the team", "Team standup is at 10am London time, Mon-Fri"),
    ("How do I configure nginx?", "The project uses Traefik as reverse proxy, not nginx"),
    ("The tests are failing in CI", "CI runs on Gitea Actions, tests need postgres service container"),
    ("Let's optimize the search function", "Search uses Elasticsearch, recently upgraded to v8"),
    ("I need to backup the database", "Backups run daily at 3am UTC via cron job to S3"),
    ("The memory usage is too high", "Python service has a known memory leak in the websocket handler"),
    ("Can you help with the Docker setup?", "Project uses docker-compose for local dev, k3s for production"),
    ("I want to add caching", "Redis is already available at redis.internal:6379"),
    ("The frontend is loading slowly", "CDN is CloudFlare, assets should be cached with 1h TTL"),
    ("Let's refactor the payment module", "Payment uses Stripe API, webhook handler is in payments/webhook.py"),
    ("I need to set up a new server", "Standard setup: Ubuntu 22.04, Docker, Tailscale, monitoring agent"),
    ("The log files are too large", "Logs rotate daily, kept for 30 days, shipped to Loki"),
]

# Paraphrased queries (semantically similar to cues but different wording)
PARAPHRASED_QUERIES = [
    "How's the weather outside?",
    "We should push the new release",
    "The DB performance is terrible",
    "Please look at my code changes",
    "There's a login bug I need to fix",
    "We need better observability",
    "Getting internal server errors from the API",
    "I'm interested in learning a new language like Rust",
    "Need to organize a team meeting",
    "How to set up nginx as a web server?",
    "CI tests keep breaking",
    "The search feature needs to be faster",
    "How do I create a database backup?",
    "The service is using too much RAM",
    "Help me with Docker configuration",
    "I want to implement caching for the API",
    "The website is really slow",
    "The payment system needs restructuring",
    "Setting up a fresh Linux server",
    "Logs are eating up disk space",
]


def load_model():
    """Load a small, fast sentence transformer."""
    from sentence_transformers import SentenceTransformer
    print("Loading sentence-transformers model...")
    model = SentenceTransformer("all-MiniLM-L6-v2", device=DEVICE)
    print(f"  Model loaded. Embedding dim: {model.get_sentence_embedding_dimension()}")
    return model


def embed_texts(model, texts):
    """Encode texts to normalized embeddings on GPU."""
    embeddings = model.encode(texts, convert_to_tensor=True,
                               normalize_embeddings=True, device=DEVICE)
    return embeddings


class HebbianMemory:
    def __init__(self, input_dim, code_dim=16384, k=20):
        self.k = k
        self.code_dim = code_dim
        self.input_dim = input_dim
        self.proj = (torch.randn(input_dim, code_dim, device=DEVICE)
                     * (1.0 / input_dim**0.5))
        self.target_proj = (torch.randn(input_dim, code_dim, device=DEVICE)
                            * (1.0 / input_dim**0.5))
        self.W = torch.zeros(code_dim, code_dim, device=DEVICE)
        self.cue_store = []  # For coarse retrieval
        self.target_store = []
        self.metadata = []  # Store original text for debugging

    def sep(self, x):
        return winner_take_all(x @ self.proj, self.k)

    def sep_target(self, x):
        return winner_take_all(x @ self.target_proj, self.k)

    def learn(self, cue_emb, target_emb, cue_text="", target_text=""):
        cc = self.sep(cue_emb)
        tc = self.sep_target(target_emb)
        self.W += torch.outer(tc, cc)
        self.cue_store.append(cue_emb.detach().clone())
        self.target_store.append(target_emb.detach().clone())
        self.metadata.append({"cue": cue_text, "target": target_text})

    def recall_direct(self, query_emb):
        """Direct WTA recall (no coarse retrieval)."""
        cc = self.sep(query_emb)
        raw = self.W @ cc
        return winner_take_all(raw, self.k)

    def recall_coarse_to_fine(self, query_emb, top_n=3):
        """Coarse: NN in embedding space. Fine: Hebbian recall from best match."""
        if not self.cue_store:
            return torch.zeros(self.code_dim, device=DEVICE)

        cue_matrix = torch.stack(self.cue_store)
        sims = nn.functional.cosine_similarity(
            query_emb.unsqueeze(0), cue_matrix, dim=-1)
        best_idx = sims.argmax()
        best_cue = self.cue_store[best_idx]

        cc = self.sep(best_cue)
        raw = self.W @ cc
        return winner_take_all(raw, self.k), best_idx.item()

    def find_nearest_target(self, recalled_code, top_n=3):
        """Given a recalled code, find which stored targets it matches."""
        target_codes = [self.sep_target(t) for t in self.target_store]
        sims = [cosine(recalled_code, tc) for tc in target_codes]
        sorted_idx = np.argsort(sims)[::-1]
        return [(int(i), sims[i], self.metadata[i]) for i in sorted_idx[:top_n]]


def test_basic_recall(model, mem):
    """Test: can we recall the correct memory for each cue?"""
    print("\n=== Test 1: Direct Recall (exact cues) ===")

    cue_texts = [p[0] for p in MEMORY_PAIRS]
    target_texts = [p[1] for p in MEMORY_PAIRS]

    correct_count = 0
    for i in range(len(MEMORY_PAIRS)):
        cue_emb = embed_texts(model, [cue_texts[i]])[0]
        recalled = mem.recall_direct(cue_emb)
        matches = mem.find_nearest_target(recalled, top_n=3)

        is_correct = matches[0][0] == i
        correct_count += is_correct

        if not is_correct and i < 5:  # Show first few errors
            print(f"  ✗ Cue: '{cue_texts[i][:40]}...'")
            print(f"    Expected: [{i}] '{target_texts[i][:50]}...'")
            print(f"    Got:      [{matches[0][0]}] '{matches[0][2]['target'][:50]}...' "
                  f"(sim={matches[0][1]:.3f})")

    print(f"  Direct recall: {correct_count}/{len(MEMORY_PAIRS)} "
          f"({correct_count/len(MEMORY_PAIRS):.0%})")
    return correct_count / len(MEMORY_PAIRS)


def test_paraphrase_recall(model, mem):
    """Test: can we recall memories using paraphrased queries?"""
    print("\n=== Test 2: Paraphrase Recall ===")

    target_texts = [p[1] for p in MEMORY_PAIRS]

    # Direct recall (WTA)
    direct_correct = 0
    coarse_correct = 0

    for i, query in enumerate(PARAPHRASED_QUERIES):
        query_emb = embed_texts(model, [query])[0]

        # Direct
        recalled = mem.recall_direct(query_emb)
        matches = mem.find_nearest_target(recalled, top_n=3)
        is_direct = matches[0][0] == i
        direct_correct += is_direct

        # Coarse-to-fine
        recalled_cf, best_idx = mem.recall_coarse_to_fine(query_emb)
        matches_cf = mem.find_nearest_target(recalled_cf, top_n=3)
        is_coarse = matches_cf[0][0] == i
        coarse_correct += is_coarse

        if i < 5:
            status_d = "✓" if is_direct else "✗"
            status_c = "✓" if is_coarse else "✗"
            print(f"  [{status_d}/{status_c}] Q: '{query[:50]}...'")
            if not is_direct:
                print(f"       Direct got: [{matches[0][0]}] "
                      f"'{matches[0][2]['target'][:50]}...'")
            if is_coarse and not is_direct:
                print(f"       Coarse-fine got it right! (via cue #{best_idx})")

    n = len(PARAPHRASED_QUERIES)
    print(f"\n  Direct recall: {direct_correct}/{n} ({direct_correct/n:.0%})")
    print(f"  Coarse-to-fine: {coarse_correct}/{n} ({coarse_correct/n:.0%})")
    return direct_correct / n, coarse_correct / n


def test_semantic_neighborhood(model, mem):
    """Test: do semantically related cues retrieve related memories?"""
    print("\n=== Test 3: Semantic Neighborhood ===")

    test_queries = [
        "server is down",          # Should relate to: API 500, deployment, monitoring
        "performance problem",      # Should relate to: DB slow, memory, search
        "security issue",           # Should relate to: auth bug, JWT tokens
        "infrastructure setup",     # Should relate to: server, Docker, k3s
    ]

    for query in test_queries:
        query_emb = embed_texts(model, [query])[0]
        recalled = mem.recall_direct(query_emb)
        matches = mem.find_nearest_target(recalled, top_n=3)

        print(f"\n  Query: '{query}'")
        for rank, (idx, sim, meta) in enumerate(matches):
            print(f"    #{rank+1} (sim={sim:.3f}): {meta['target'][:60]}...")


def test_multihop_semantic(model, mem):
    """Test: multi-hop with semantic embeddings.
    Learn: "weather" → "morning routine" → "coffee shop"
    Can we go from "weather" to "coffee shop" in 2 hops?
    """
    print("\n=== Test 4: Multi-hop with Semantic Chains ===")

    chains = [
        ["What's the weather?", "I usually check weather before going out",
         "My favorite coffee shop is around the corner", "They have great latte art"],
        ["Let's review the code", "The code review found a memory leak",
         "Memory leaks often cause OOM kills", "We need to add memory limits to k8s pods"],
        ["Deploy to production", "Production uses blue-green deployment",
         "The blue environment is currently active", "Switch DNS to green when ready"],
    ]

    for chain_idx, chain in enumerate(chains):
        print(f"\n  Chain {chain_idx+1}: {' → '.join([c[:20]+'...' for c in chain])}")

        # Create a separate small memory for this chain
        chain_mem = HebbianMemory(384, code_dim=8192, k=20)

        chain_embs = [embed_texts(model, [text])[0] for text in chain]

        # Learn consecutive pairs
        for i in range(len(chain) - 1):
            chain_mem.learn(chain_embs[i], chain_embs[i+1],
                           chain[i], chain[i+1])

        # Test recall at each hop distance
        for hops in range(1, len(chain)):
            start_emb = chain_embs[0]
            target_code = chain_mem.sep_target(chain_embs[hops])

            # Multi-hop
            code = chain_mem.sep(start_emb)
            for _ in range(hops):
                raw = chain_mem.W @ code
                code = winner_take_all(raw, chain_mem.k)

            sim = cosine(code, target_code)
            print(f"    {hops} hop(s): '{chain[0][:25]}...' → "
                  f"'{chain[hops][:25]}...' sim={sim:.4f}")


def test_embedding_distances(model):
    """Analyze: how far apart are original and paraphrased embeddings?"""
    print("\n=== Test 5: Embedding Distance Analysis ===")

    cue_texts = [p[0] for p in MEMORY_PAIRS]
    cue_embs = embed_texts(model, cue_texts)
    para_embs = embed_texts(model, PARAPHRASED_QUERIES)

    # Same-pair distances
    same_pair_sims = []
    for i in range(len(cue_texts)):
        s = cosine(cue_embs[i], para_embs[i])
        same_pair_sims.append(s)

    # Different-pair distances
    diff_pair_sims = []
    for i in range(len(cue_texts)):
        for j in range(len(cue_texts)):
            if i != j:
                diff_pair_sims.append(cosine(cue_embs[i], para_embs[j]))

    print(f"  Same-pair cosine sim:  mean={np.mean(same_pair_sims):.4f}, "
          f"min={np.min(same_pair_sims):.4f}, max={np.max(same_pair_sims):.4f}")
    print(f"  Diff-pair cosine sim:  mean={np.mean(diff_pair_sims):.4f}, "
          f"min={np.min(diff_pair_sims):.4f}, max={np.max(diff_pair_sims):.4f}")
    print(f"  Gap: {np.mean(same_pair_sims) - np.mean(diff_pair_sims):.4f}")

    # Show some examples
    print("\n  Sample distances:")
    for i in range(5):
        print(f"    '{cue_texts[i][:35]}...' ↔ '{PARAPHRASED_QUERIES[i][:35]}...' "
              f"sim={same_pair_sims[i]:.4f}")


def main():
    print("=" * 60)
    print("Experiment 4: Real Sentence Embeddings")
    print("=" * 60)

    model = load_model()

    # Analyze embedding space first
    test_embedding_distances(model)

    # Build memory
    print("\n--- Building memory ---")
    embed_dim = model.get_sentence_embedding_dimension()
    mem = HebbianMemory(embed_dim, code_dim=16384, k=20)

    cue_texts = [p[0] for p in MEMORY_PAIRS]
    target_texts = [p[1] for p in MEMORY_PAIRS]

    cue_embs = embed_texts(model, cue_texts)
    target_embs = embed_texts(model, target_texts)

    for i in range(len(MEMORY_PAIRS)):
        mem.learn(cue_embs[i], target_embs[i], cue_texts[i], target_texts[i])

    print(f"  Stored {len(MEMORY_PAIRS)} memory pairs")

    # Run tests
    test_basic_recall(model, mem)
    test_paraphrase_recall(model, mem)
    test_semantic_neighborhood(model, mem)
    test_multihop_semantic(model, mem)


if __name__ == "__main__":
    main()