Build Entity-Aware Semantic Search with Pinecone Metadata Filters

Entity-aware semantic search combines two retrieval signals: dense embeddings that capture meaning, and named entities that anchor a query to specific people, places, and organizations. This guide extracts named entities from each document, stores them as Pinecone metadata, and applies a metadata filter at query time so vectors are ranked only when they share at least one entity with the query. The result is a retrieval system that can distinguish a general travel article from Greece-specific material without reducing the recall of embedding-based matching.

This guide is adapted from Pinecone's ner-powered-search.ipynb, available at https://github.com/pinecone-io/examples under the MIT license. The code cells below match that notebook's order and contents. The notebook does not publish timing or cost benchmarks, so benchmark on the hardware you intend to use before planning capacity.

Prerequisites

You need Python 3.8+ and a Jupyter notebook or script environment. The notebook reads PINECONE_API_KEY, PINECONE_CLOUD, and PINECONE_REGION from the environment; you must create a Pinecone key before running. Required packages are sentence_transformers, pinecone-client==3.1.0, datasets, torch, transformers, and tqdm. A CUDA GPU is optional. The source checks torch.cuda.is_available() and uses CPU when no GPU is present.

Step 1: Load and prepare the dataset

The dataset is fabiochiu/medium-articles, loaded as a CSV from Hugging Face. After loading, the source drops empty rows, samples 20,000 rows, truncates each article to its first 1,000 characters, and creates a single title_text column from the title and article text. That column is the input to both the NER pipeline and the retriever.

!pip install sentence_transformers pinecone-client==3.1.0 datasets -qU

The install command appears as a notebook cell. If you are running a plain script, use the same packages in your terminal or requirements file.

from datasets import load_dataset

# load the dataset and convert to pandas dataframe
df = load_dataset(
    "fabiochiu/medium-articles",
    data_files="medium_articles.csv",
    split="train"
).to_pandas()

The loader returns a dataset and converts it to a pandas DataFrame. The next cell keeps the experiment small.

# drop empty rows and select 20k articles
df = df.dropna().sample(20000, random_state=32)
df.head()

The random state fixes the sample so the index contents are reproducible. Dropping empty rows prevents the NER and embedding steps from wasting calls on empty strings.

# select first 1000 characters
df["text"] = df["text"].str[:1000]
# join article title and the text
df["title_text"] = df["title"] + ". " + df["text"]

Truncation keeps the token counts and metadata size manageable. It is a deliberate source choice: the filtered title-and-text prefix is what gets embedded and passed through NER.

Step 2: Initialize the NER model

The notebook uses dslim/bert-base-NER, a BERT-base token-classification model. The device selector uses the current CUDA device when available; otherwise it passes None to let Hugging Face place the model on CPU.

import torch

# set device to GPU if available
device = torch.cuda.current_device() if torch.cuda.is_available() else None

from transformers import AutoTokenizer, AutoModelForTokenClassification
from transformers import pipeline

model_id = "dslim/bert-base-NER"

# load the tokenizer from huggingface
tokenizer = AutoTokenizer.from_pretrained(
    model_id
)
# load the NER model from huggingface
model = AutoModelForTokenClassification.from_pretrained(
    model_id
)
# load the tokenizer and model into a NER pipeline
nlp = pipeline(
    "ner",
    model=model,
    tokenizer=tokenizer,
    aggregation_strategy="max",
    device=device
)

The aggregation_strategy="max" argument is important: it groups subword tokens into full entity words instead of leaving fragments such as "United" and "Kingdom" separate.

text = "London is the capital of England and the United Kingdom"
# use the NER pipeline to extract named entities from the text
nlp(text)

This sanity-check sentence should produce London, England, and United Kingdom as separate entity predictions. At this point nlp accepts a string or a list of strings and returns per-text entity lists.

Step 3: Initialize the retriever and Pinecone index

The retriever is flax-sentence-embeddings/all_datasets_v3_mpnet-base, which produces 768-dimensional embeddings. It receives the same device value as the NER model.

from sentence_transformers import SentenceTransformer

# load the model from huggingface
retriever = SentenceTransformer(
    'flax-sentence-embeddings/all_datasets_v3_mpnet-base',
    device=device
)
retriever

The Pinecone client reads the API key from the environment. The source falls back to a literal string 'PINECONE_API_KEY' if the variable is missing, so set the real key or replace that fallback before running.

import os
from pinecone import Pinecone

# initialize connection to pinecone (get API key at app.pinecone.io)
api_key = os.environ.get('PINECONE_API_KEY') or 'PINECONE_API_KEY'

# configure client
pc = Pinecone(api_key=api_key)

The next cell configures serverless deployment. The notebook defaults to AWS us-east-1, but you can set PINECONE_CLOUD and PINECONE_REGION for your Pinecone account.

from pinecone import ServerlessSpec

cloud = os.environ.get('PINECONE_CLOUD') or 'aws'
region = os.environ.get('PINECONE_REGION') or 'us-east-1'

spec = ServerlessSpec(cloud=cloud, region=region)

index_name = "ner-search"

The index must use dimension 768 and cosine metric to match the retriever. The wait loop polls until Pinecone reports the index is ready.

import time

# check if index already exists (it shouldn't if this is first time)
if index_name not in pc.list_indexes().names():
    # if does not exist, create index
    pc.create_index(
        index_name,
        dimension=768,
        metric='cosine',
        spec=spec
    )
    # wait for index to be initialized
    while not pc.describe_index(index_name).status['ready']:
        time.sleep(1)

# connect to index
index = pc.Index(index_name)
# view index stats
index.describe_index_stats()

The index_name value is arbitrary as long as it is unique in your Pinecone project. If the index already exists, the code skips creation and connects to the existing one.

Step 4: Embed documents and upsert with entity metadata

The helper function runs the NER pipeline on a batch and keeps only the word field from each prediction, discarding scores and token positions.

def extract_named_entities(text_batch):
    # extract named entities using the NER pipeline
    extracted_batch = nlp(text_batch)
    entities = []
    # loop through the results and only select the entity names
    for text in extracted_batch:
        ne = [entity["word"] for entity in text]
        entities.append(ne)
    return entities

The upsert loop processes batches of 64. For each batch it embeds title_text, extracts entities, de-duplicates them with set, drops the embedded text from the metadata, and writes records with string IDs. Only the original title and named_entities list remain as metadata.

from tqdm.auto import tqdm
import warnings
warnings.filterwarnings('ignore', category=UserWarning)

# we will use batches of 64
batch_size = 64

for i in tqdm(range(0, len(df), batch_size)):
    # find end of batch
    i_end = min(i+batch_size, len(df))
    # extract batch
    batch = df.iloc[i:i_end].copy()
    # generate embeddings for batch
    emb = retriever.encode(batch["title_text"].tolist()).tolist()
    # extract named entities from the batch
    entities = extract_named_entities(batch["title_text"].tolist())
    # remove duplicate entities from each record
    batch["named_entities"] = [list(set(entity)) for entity in entities]
    batch = batch.drop('title_text', axis=1)
    # get metadata
    meta = batch.to_dict(orient="records")
    # create unique IDs
    ids = [f"{idx}" for idx in range(i, i_end)]
    # add all to upsert list
    to_upsert = list(zip(ids, emb, meta))
    # upsert/insert these records to pinecone
    _ = index.upsert(vectors=to_upsert)

# check that we have all vectors in index
index.describe_index_stats()

Batch size 64 is the source configuration. If you run on a small GPU, reduce the batch size; the tradeoff is more API round trips. The warning filter is present in the source to keep the notebook output clean.

Step 5: Query with named entity filters

The search_pinecone function extracts entities from the query, embeds the query, and passes both the vector and a $in predicate on the named_entities metadata field. Only records whose named_entities list contains at least one extracted query entity are eligible for ranking.

from pprint import pprint

def search_pinecone(query):
    # extract named entities from the query
    ne = extract_named_entities([query])[0]
    # create embeddings for the query
    xq = retriever.encode(query).tolist()
    # query the pinecone index while applying named entity filter
    xc = index.query(xq, top_k=10, include_metadata=True, filter={"named_entities": {"$in": ne}})
    # extract article titles from the search result
    r = [x["metadata"]["title"] for x in xc["matches"]]
    return pprint({"Extracted Named Entities": ne, "Result": r})

The function prints the extracted entities and the returned titles. The source includes three example queries:

query = "What are the best places to visit in Greece?"
search_pinecone(query)

query = "What are the best places to visit in London?"
search_pinecone(query)

query = "Why does SpaceX want to build a city on Mars?"
search_pinecone(query)

These queries are designed to test whether the filter grounds the results in the query's specific entities rather than returning broad travel or space content.

Here is how a metadata-filtered entity-aware query compares with other retrieval strategies:

Step 6: Clean up

Delete the index when you are done experimenting. This is the source notebook's final step, and it prevents unused metadata and vectors from accumulating in your Pinecone project.

pc.delete_index(index_name)

What to watch out for

The biggest failure mode is an empty entity list. If the query contains no names, places, or organizations, the filter predicate {"named_entities": {"$in": ne}} runs with an empty set. That can return no results or produce an error. In production, fall back to unfiltered vector search or skip the metadata predicate until entities are available.

The general NER model recognizes person, location, organization, and miscellaneous entity types. It will not reliably capture domain-specific terms such as product names, library names, or medical conditions. For those cases, fine-tune a custom NER model on the target vocabulary or combine this metadata filter with a keyword extractor.

Set-based de-duplication changes entity order. That is harmless for filtering, but if you present entities to users in the order they appeared, sort them or preserve the original order before de-duplication.

Pinecone stores the named_entities list as per-vector metadata. If you index much longer documents, the entity list can grow and hit the metadata size limit. Truncation to 1,000 characters keeps this demo small; production pipelines should cap the number of stored entities or store only high-confidence predictions.

Running NER on every query adds inference latency. The source model is a fine-tuned BERT-base classifier, so it is heavier than a simple lexical parser. If you need a low query-time budget, cache frequent query entities, run the NER step asynchronously, or use a smaller NER model.

Indexing scale also matters. The notebook samples 20,000 rows from a larger CSV. Removing the sample means running NER and embedding over the full loaded dataset. If you use longer passages, token-level efficiency becomes an even larger cost driver; related production analysis such as Tokengate's 11.3x long-context prefill speedup is a useful reference for thinking about token budgets before scaling this pattern.

Where to go next

This notebook gives you a working entity-aware retrieval system. You can extend it by swapping in a domain-specific NER model, adding custom entity types from a knowledge graph, or combining the metadata filter with Pinecone hybrid search for lexical and vector scoring. For production, measure filter selectivity, cache frequent query NER results, and test metadata size as document length grows.