embed.py

#!/usr/bin/env python3
"""Embed memory entries in SurrealDB using Qwen3-Embedding-4B."""

import argparse
import gc
import json
import sys
import time
import requests
import torch
import numpy as np
from sentence_transformers import SentenceTransformer

SURREAL_URL = "http://localhost:8000"
SURREAL_NS = "kit"
SURREAL_DB = "memory"
SURREAL_USER = "root"
SURREAL_PASS = "root"

MODEL_NAME = "Qwen/Qwen3-Embedding-4B"
EMBEDDING_DIM = 1024  # MRL truncation from native 2560


def surreal_query(sql: str, vars: dict | None = None) -> list:
    """Execute a SurrealQL query via HTTP."""
    headers = {
        "Accept": "application/json",
        "surreal-ns": SURREAL_NS,
        "surreal-db": SURREAL_DB,
    }
    payload = {"query": sql}
    if vars:
        payload["variables"] = vars
    # Use the SQL endpoint
    resp = requests.post(
        f"{SURREAL_URL}/sql",
        headers=headers,
        data=sql,
        auth=(SURREAL_USER, SURREAL_PASS),
    )
    resp.raise_for_status()
    return resp.json()


def load_model() -> SentenceTransformer:
    """Load Qwen3-Embedding-4B onto GPU."""
    print(f"Loading {MODEL_NAME} onto GPU...")
    t0 = time.time()
    model = SentenceTransformer(MODEL_NAME, trust_remote_code=True, model_kwargs={"torch_dtype": torch.float16})
    model.max_seq_length = 8192
    # Truncate to target dim
    model.truncate_dim = EMBEDDING_DIM
    print(f"Model loaded in {time.time() - t0:.1f}s")
    return model


def unload_model(model):
    """Free GPU memory."""
    del model
    gc.collect()
    torch.cuda.empty_cache()
    print("Model unloaded, VRAM freed.")


def get_memories(only_missing: bool = False) -> list[dict]:
    """Fetch memory entries from SurrealDB."""
    sql = "SELECT id, title, content, embedding FROM memory;"
    results = surreal_query(sql)
    rows = results[0].get("result", [])
    if only_missing:
        rows = [r for r in rows if not r.get("embedding")]
    return rows


def embed_and_update(model: SentenceTransformer, memories: list[dict], batch_size: int = 8):
    """Generate embeddings and update SurrealDB records."""
    if not memories:
        print("No memories to embed.")
        return

    contents = [m["content"] for m in memories]
    ids = [m["id"] for m in memories]

    print(f"Embedding {len(contents)} memories...")
    t0 = time.time()
    embeddings = model.encode(contents, batch_size=batch_size, show_progress_bar=True, normalize_embeddings=True)
    elapsed = time.time() - t0
    print(f"Embedded {len(contents)} entries in {elapsed:.1f}s ({len(contents)/elapsed:.1f} entries/s)")

    # Update each record
    for i, (mid, emb) in enumerate(zip(ids, embeddings)):
        vec = emb.tolist()
        # SurrealDB expects the id as-is (e.g. memory:identity)
        sql = f"UPDATE {mid} SET embedding = {json.dumps(vec)}, updated_at = time::now();"
        surreal_query(sql)

    print(f"Updated {len(ids)} records in SurrealDB.")


def main():
    parser = argparse.ArgumentParser(description="Embed memory entries with Qwen3-Embedding-4B")
    parser.add_argument("--all", action="store_true", help="Re-embed all entries (default: only missing)")
    parser.add_argument("--batch-size", type=int, default=8, help="Batch size for encoding")
    args = parser.parse_args()

    memories = get_memories(only_missing=not args.all)
    if not memories:
        print("All memories already have embeddings. Use --all to re-embed.")
        return

    model = load_model()
    try:
        embed_and_update(model, memories, batch_size=args.batch_size)
    finally:
        unload_model(model)


if __name__ == "__main__":
    main()