IVYXSTUDIO · IVY NODE

Knowledge Base Ingest

ivy.node.kb-ingest · v0.0.0

ivyx✓

Ingests text documents into a Milvus collection for retrieval. Embeds each document with a Hugging Face model (CLS pooling + L2 normalize — identical to ivy.node.rag-kb-answer so query and index vectors match), (re)creates the target collection sized to the embedding dimension, and inserts {vector, text} rows. Run this once to populate the knowledge base before querying with rag-kb-answer. Demonstrates the multi-cell program layout (input / implementation / run cells share one kernel session).

#rag#kb#ingest#milvus

Inputs

Field	Type	Description
docsrequired	array	Text documents/passages to index.
model_namerequired	string	Hugging Face embedding model id (must match the query node).
hugging_face_token	string	Hugging Face token; empty string for public models.
milvus_urirequired	string	Milvus connection URI (e.g. /tmp/kb.db).
collection_namerequired	string	Target Milvus collection name.
drop_existing	boolean	Drop the collection first if it already exists.
features	object	Features
labels	object	Labels

Outputs

Field	Type	Description
inserted_countrequired	integer	Number of documents inserted.
dim	integer	Embedding dimension of the created collection.
collection_name	string	The collection that was populated.
features	object	Features
labels	object	Labels

Source

python

# Cell 1 — input preparation
inp = __ivy_ctx__["nodes"][__ivy_node_id__]["input"]
docs = inp["docs"]
model_name = inp["model_name"]
hugging_face_token = inp.get("hugging_face_token", "") or None
milvus_uri = inp["milvus_uri"]
collection_name = inp["collection_name"]
drop_existing = bool(inp.get("drop_existing", True))

# Cell 2 — implementation (shares the session with cell 1)
from typing import List
import torch
from transformers import AutoTokenizer, AutoModel
from pymilvus import MilvusClient

class KbIngest:
    """Embeds documents and inserts them into a Milvus collection."""

    def __init__(self, model_name: str, hugging_face_token, milvus_uri: str):
        self.tokenizer = AutoTokenizer.from_pretrained(model_name, token=hugging_face_token)
        self.model = AutoModel.from_pretrained(model_name, token=hugging_face_token)
        self.milvus = MilvusClient(uri=milvus_uri)

    def embed(self, text: str) -> List[float]:
        inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True)
        with torch.no_grad():
            outputs = self.model(**inputs)
        emb = outputs.last_hidden_state[:, 0, :]
        return torch.nn.functional.normalize(emb, p=2, dim=1).squeeze().tolist()

    def ingest(self, docs: List[str], collection_name: str, drop_existing: bool):
        if not docs:
            raise ValueError("docs is empty")
        vectors = [self.embed(d) for d in docs]
        dim = len(vectors[0])
        if drop_existing and self.milvus.has_collection(collection_name):
            self.milvus.drop_collection(collection_name)
        if not self.milvus.has_collection(collection_name):
            self.milvus.create_collection(collection_name=collection_name, dimension=dim, auto_id=True)
        rows = [{"vector": v, "text": d} for v, d in zip(vectors, docs)]
        self.milvus.insert(collection_name, rows)
        return len(rows), dim

# Cell 3 — run + output collection
ingestor = KbIngest(model_name=model_name, hugging_face_token=hugging_face_token, milvus_uri=milvus_uri)
inserted_count, dim = ingestor.ingest(docs, collection_name, drop_existing)

out = __ivy_ctx__["nodes"][__ivy_node_id__]["output"]
out["inserted_count"] = inserted_count
out["dim"] = dim
out["collection_name"] = collection_name

Tests

Requires: python:3.11

basic-ingest
Embed and insert a few docs (network; tolerate success or environment error).
ragingestbasic
empty-docs
Empty docs list should fail.
ragingesterror