embed.py

"""
One-time script to pre-compute embeddings for the knowledge base.
Run this locally, then commit embeddings.json to the repo.
"""

import json
from sentence_transformers import SentenceTransformer

def chunk_text(text: str, chunk_size: int = 500, overlap: int = 50) -> list[dict]:
    """Split text into overlapping chunks."""
    lines = text.strip().split('\n')
    chunks = []
    current_chunk = []
    current_length = 0

    for line in lines:
        line_length = len(line)

        # If adding this line exceeds chunk size, save current chunk
        if current_length + line_length > chunk_size and current_chunk:
            chunk_text = '\n'.join(current_chunk)
            chunks.append({"text": chunk_text})

            # Keep some overlap
            overlap_lines = []
            overlap_length = 0
            for l in reversed(current_chunk):
                if overlap_length + len(l) <= overlap:
                    overlap_lines.insert(0, l)
                    overlap_length += len(l)
                else:
                    break
            current_chunk = overlap_lines
            current_length = overlap_length

        current_chunk.append(line)
        current_length += line_length

    # Don't forget the last chunk
    if current_chunk:
        chunk_text = '\n'.join(current_chunk)
        chunks.append({"text": chunk_text})

    return chunks


def main():
    # Load knowledge base
    with open("knowledge.md", "r", encoding="utf-8") as f:
        content = f.read()

    # Chunk the content
    chunks = chunk_text(content, chunk_size=500, overlap=50)
    print(f"Created {len(chunks)} chunks")

    # Load embedding model
    print("Loading embedding model...")
    model = SentenceTransformer('all-MiniLM-L6-v2')

    # Generate embeddings
    print("Generating embeddings...")
    texts = [chunk["text"] for chunk in chunks]
    embeddings = model.encode(texts, show_progress_bar=True)

    # Save to JSON
    data = {
        "chunks": [
            {
                "text": chunk["text"],
                "embedding": embedding.tolist()
            }
            for chunk, embedding in zip(chunks, embeddings)
        ]
    }

    with open("embeddings.json", "w", encoding="utf-8") as f:
        json.dump(data, f, indent=2)

    print(f"Saved embeddings to embeddings.json")


if __name__ == "__main__":
    main()