Knowledge Base Search

Inputs

Outputs

Source

python

inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]

_MODEL_CACHE = globals().setdefault("_ivy_embed_models", {})


def embed(texts, model_name):
    """fastembed (ONNX, no torch). The model is loaded once per kernel."""
    from fastembed import TextEmbedding
    model = _MODEL_CACHE.get(model_name)
    if model is None:
        model = _MODEL_CACHE[model_name] = TextEmbedding(model_name)
    return [[float(x) for x in v] for v in model.embed(list(texts))]

import json

question = inp["question"]
index_path = inp["index_path"]
top_k = int(inp.get("top_k", 3))
with open(index_path) as handle:
    index = json.load(handle)
items = index.get("items", [])
if not items:
    raise ValueError(f"{index_path} holds no chunks.")
query = embed([question], index["model"])[0]


def cosine(a, b):
    dot = sum(x * y for x, y in zip(a, b))
    na = sum(x * x for x in a) ** 0.5
    nb = sum(y * y for y in b) ** 0.5
    return dot / (na * nb) if na and nb else 0.0


scored = sorted(((cosine(query, it["vector"]), it["text"]) for it in items), key=lambda p: -p[0])[:top_k]

out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["hits"] = [{"text": t, "score": round(s, 4)} for s, t in scored]
out["context"] = "\n\n".join(t for _, t in scored)
out["question"] = question

Tests

Requires: python:3.9, fastembed, model download on first use

• finds-refunds

  A refund question finds the refund chunk first.

• empty-index

  A file with no chunks is refused.