Build semantic search for text

%pip install -qU pixeltable sentence-transformers

import pixeltable as pxt
from pixeltable.functions.huggingface import sentence_transformer

# Create a fresh directory
pxt.drop_dir('search_demo', force=True)
pxt.create_dir('search_demo')

# Create table with content and metadata
kb = pxt.create_table(
    'search_demo/articles',
    {'title': pxt.String, 'content': pxt.String, 'category': pxt.String},
)

# Insert sample content
kb.insert(
    [
        {
            'title': 'Debugging best practices',
            'content': 'Use logging, breakpoints, and unit tests to identify and fix issues in your code.',
            'category': 'engineering',
        },
        {
            'title': 'Machine learning model optimization',
            'content': 'Improve training efficiency with batch normalization, learning rate schedules, and early stopping.',
            'category': 'ml',
        },
        {
            'title': 'Production infrastructure setup',
            'content': 'Deploy applications using containers, load balancers, and automated scaling.',
            'category': 'devops',
        },
        {
            'title': 'API design principles',
            'content': 'Create RESTful endpoints with proper versioning, authentication, and error handling.',
            'category': 'engineering',
        },
    ]
)

# Add embedding index
kb.add_embedding_index(
    column='content',
    string_embed=sentence_transformer.using(model_id='all-MiniLM-L6-v2'),
)

# Search by meaning
query = 'how to fix bugs'
sim = kb.content.similarity(string=query)

results = (
    kb.order_by(sim, asc=False)
    .select(kb.title, kb.content, score=sim)
    .limit(2)
)
results.collect()

# Search within a specific category
query = 'best practices'
sim = kb.content.similarity(string=query)

results = (
    kb.where(kb.category == 'engineering')  # Filter first
    .order_by(sim, asc=False)
    .select(kb.title, kb.category, score=sim)
    .limit(2)
)
results.collect()