|
import os |
|
import streamlit as st |
|
from llama_index import GPTSimpleVectorIndex, SimpleDirectoryReader |
|
|
|
index_name = "./index.json" |
|
documents_folder = "./documents" |
|
|
|
|
|
@st.cache_resource |
|
def initialize_index(index_name, documents_folder): |
|
if os.path.exists(index_name): |
|
index = GPTSimpleVectorIndex.load_from_disk(index_name) |
|
else: |
|
documents = SimpleDirectoryReader(documents_folder).load_data() |
|
index = GPTSimpleVectorIndex.from_documents(documents) |
|
index.save_to_disk(index_name) |
|
|
|
return index |
|
|
|
|
|
@st.cache_data(max_entries=200, persist=True) |
|
def query_index(_index, query_text): |
|
response = _index.query(query_text) |
|
return str(response) |
|
|
|
|
|
st.title("π¦ Llama Index Demo π¦") |
|
st.header("Welcome to the Llama Index Streamlit Demo") |
|
st.write("Enter a query about Paul Graham's essays. You can check out the original essay [here](https://raw.githubusercontent.com/jerryjliu/llama_index/main/examples/paul_graham_essay/data/paul_graham_essay.txt)") |
|
|
|
index = None |
|
api_key = st.text_input("Enter your OpenAI API key here:", type="password") |
|
if api_key: |
|
os.environ['OPENAI_API_KEY'] = api_key |
|
index = initialize_index(index_name, documents_folder) |
|
|
|
|
|
if index is None: |
|
st.warning("Please enter your api key first.") |
|
|
|
text = st.text_input("Query text:", value="What did the author do growing up?") |
|
|
|
if st.button("Run Query") and text is not None: |
|
response = query_index(index, text) |
|
st.markdown(response) |
|
|
|
llm_col, embed_col = st.columns(2) |
|
with llm_col: |
|
st.markdown(f"LLM Tokens Used: {index.service_context.llm_predictor._last_token_usage}") |
|
|
|
with embed_col: |
|
st.markdown(f"Embedding Tokens Used: {index.service_context.embed_model._last_token_usage}") |
|
|
|
|