Back to Llama Index

Simple Vector Store

docs/examples/vector_stores/SimpleIndexDemo.ipynb

0.14.214.4 KB
Original Source

<a href="https://colab.research.google.com/github/run-llama/llama_index/blob/main/docs/examples/vector_stores/SimpleIndexDemo.ipynb" target="_parent"></a>

Simple Vector Store

If you're opening this Notebook on colab, you will probably need to install LlamaIndex 🦙.

python
!pip install llama-index
python
import os
import openai

os.environ["OPENAI_API_KEY"] = "sk-..."
openai.api_key = os.environ["OPENAI_API_KEY"]

Load documents, build the VectorStoreIndex

python
import nltk

nltk.download("stopwords")
python
import llama_index.core
python
import logging
import sys

logging.basicConfig(stream=sys.stdout, level=logging.INFO)
logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    load_index_from_storage,
    StorageContext,
)
from IPython.display import Markdown, display

Download Data

python
!mkdir -p 'data/paul_graham/'
!wget 'https://raw.githubusercontent.com/run-llama/llama_index/main/docs/examples/data/paul_graham/paul_graham_essay.txt' -O 'data/paul_graham/paul_graham_essay.txt'
python
# load documents
documents = SimpleDirectoryReader("./data/paul_graham/").load_data()
python
index = VectorStoreIndex.from_documents(documents)
python
# save index to disk
index.set_index_id("vector_index")
index.storage_context.persist("./storage")
python
# rebuild storage context
storage_context = StorageContext.from_defaults(persist_dir="storage")
# load index
index = load_index_from_storage(storage_context, index_id="vector_index")

Query Index

python
# set Logging to DEBUG for more detailed outputs
query_engine = index.as_query_engine(response_mode="tree_summarize")
response = query_engine.query("What did the author do growing up?")
python
display(Markdown(f"<b>{response}</b>"))

Query Index with SVM/Linear Regression

Use Karpathy's SVM-based approach. Set query as positive example, all other datapoints as negative examples, and then fit a hyperplane.

python
query_modes = [
    "svm",
    "linear_regression",
    "logistic_regression",
]
for query_mode in query_modes:
    # set Logging to DEBUG for more detailed outputs
    query_engine = index.as_query_engine(vector_store_query_mode=query_mode)
    response = query_engine.query("What did the author do growing up?")
    print(f"Query mode: {query_mode}")
    display(Markdown(f"<b>{response}</b>"))
python
display(Markdown(f"<b>{response}</b>"))
python
print(response.source_nodes[0].text)

Query Index with custom embedding string

python
from llama_index.core import QueryBundle
python
query_bundle = QueryBundle(
    query_str="What did the author do growing up?",
    custom_embedding_strs=["The author grew up painting."],
)
query_engine = index.as_query_engine()
response = query_engine.query(query_bundle)
python
display(Markdown(f"<b>{response}</b>"))

Use maximum marginal relevance

Instead of ranking vectors purely by similarity, adds diversity to the documents by penalizing documents similar to ones that have already been found based on <a href="https://www.cs.cmu.edu/~jgc/publication/The_Use_MMR_Diversity_Based_LTMIR_1998.pdf">MMR</a> . A lower mmr_treshold increases diversity.

python
query_engine = index.as_query_engine(
    vector_store_query_mode="mmr", vector_store_kwargs={"mmr_threshold": 0.2}
)
response = query_engine.query("What did the author do growing up?")

Get Sources

python
print(response.get_formatted_sources())

Query Index with Filters

We can also filter our queries using metadata

python
from llama_index.core import Document

doc = Document(text="target", metadata={"tag": "target"})

index.insert(doc)
python
from llama_index.core.vector_stores import ExactMatchFilter, MetadataFilters

filters = MetadataFilters(
    filters=[ExactMatchFilter(key="tag", value="target")]
)

retriever = index.as_retriever(
    similarity_top_k=20,
    filters=filters,
)

source_nodes = retriever.retrieve("What did the author do growing up?")
python
# retrieves only our target node, even though we set the top k to 20
print(len(source_nodes))
python
print(source_nodes[0].text)
print(source_nodes[0].metadata)