File size: 2,645 Bytes
4779c1f 8291b7d 4779c1f 8291b7d 4779c1f 8291b7d 4779c1f 8291b7d e947786 4779c1f 8291b7d 4779c1f 8291b7d 4779c1f bc89964 8291b7d 4779c1f 8291b7d 4779c1f 8291b7d 4779c1f 8291b7d 4779c1f 8291b7d 4779c1f 8291b7d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 |
import os
import streamlit as st
from llama_index import (
GPTVectorStoreIndex,
SimpleDirectoryReader,
ServiceContext,
StorageContext,
LLMPredictor,
load_index_from_storage,
)
from langchain.chat_models import ChatOpenAI
index_name = "./saved_index"
documents_folder = "./documents"
@st.cache_resource
def initialize_index(index_name, documents_folder):
llm_predictor = LLMPredictor(
llm=ChatOpenAI(model_name="gpt-3.5-turbo", temperature=0)
)
service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)
if os.path.exists(index_name):
index = load_index_from_storage(
StorageContext.from_defaults(persist_dir=index_name),
service_context=service_context,
)
else:
documents = SimpleDirectoryReader(documents_folder).load_data()
index = GPTVectorStoreIndex.from_documents(
documents, service_context=service_context
)
index.storage_context.persist(persist_dir=index_name)
return index
@st.cache_data(max_entries=200, persist=True)
def query_index(_index, query_text):
if _index is None:
return "Please initialize the index!"
response = _index.as_query_engine().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). Your query will be answered using the essay as context, using embeddings from text-ada-002 and LLM completions from gpt-3.5-turbo. You can read more about Llama Index and how this works in [our docs!](https://gpt-index.readthedocs.io/en/latest/index.html)"
)
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}"
)
|