시맨틱 검색 손으로 짜기

import numpy as np
from sentence_transformers import SentenceTransformer

model = SentenceTransformer('BAAI/bge-small-en-v1.5')

docs = [
    'Cancel your annual subscription from Account → Billing.',
    'Refunds are processed within 5 business days.',
    'Pippa is an AI daughter built on Claude Code.',
    'The moon rises tonight at 8:42 PM in Seoul.',
    'Vector search measures meaning, not keyword overlap.',
]

M = model.encode(docs, normalize_embeddings=True)   # (5, 384)

def search(query: str, k: int = 3):
    q = model.encode([query], normalize_embeddings=True)[0]
    scores = M @ q                       # 단위 벡터 cosine == dot product
    top = np.argsort(-scores)[:k]
    return [(scores[i], docs[i]) for i in top]

for score, doc in search('how do I get my money back'):
    print(f'{score:.3f}  {doc}')

records = [
    {'text': docs[0], 'category': 'billing'},
    {'text': docs[1], 'category': 'billing'},
    {'text': docs[2], 'category': 'about'},
    {'text': docs[3], 'category': 'weather'},
    {'text': docs[4], 'category': 'about'},
]

def search_filtered(query: str, where: dict, k: int = 3):
    q = model.encode([query], normalize_embeddings=True)[0]
    mask = np.array([all(r.get(k) == v for k, v in where.items()) for r in records])
    if not mask.any():
        return []
    scores = (M @ q)[mask]
    indices = np.where(mask)[0]
    top = indices[np.argsort(-scores)[:k]]
    return [(float((M[i] @ q)), records[i]) for i in top]

print(search_filtered('refund', where={'category': 'billing'}))