Llama_Index Starter Example using DBRX

llama-index_starter_dbrx.py

#
# LlamaIndex Starter Example with DBRX
#

# Based on the LlamaIndex Starter Example
# https://docs.llamaindex.ai/en/stable/getting_started/starter_example/
#
# Ensure you have installed both llama_index and Databricks integration
# e.g., pip install llama_index llama-index-llms-databricks
# 
# Ensure you have configured the following environment variables
# export DATABRICKS_TOKEN=... 
# export DATABRICKS_SERVING_ENDPOINT=... 
#

# Load data 
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
documents = SimpleDirectoryReader("data").load_data()

# Configure to use DBRX
from llama_index.llms.databricks import Databricks

# define LLM
Settings.llm = Databricks(model="databricks-dbrx-instruct")

# Define index
index = VectorStoreIndex.from_documents(documents)

# Query your data
my_question = "What did the author do growing up?"
query_engine = index.as_query_engine()
response = query_engine.query(my_question)
print("===========================================")
print("Question: %s \n" % my_question)
print("DBRX Response: %s" % response)
print("===========================================")