agutfraind/llmscanner
1
1'''2LLM scanner streamlit app3 4streamlit run .\app.py5 6Functionality7- tokenize documents8- respond to queries9- generate new documents10 11Based on: 121. https://huggingface.co/spaces/llamaindex/llama_index_vector_demo132. https://github.com/logan-markewich/llama_index_starter_pack/blob/main/streamlit_term_definition/14 15TODO:16- customize to other [LLMs](https://gpt-index.readthedocs.io/en/latest/reference/llm_predictor.html#llama_index.llm_predictor.LLMPredictor) 17- guardrails on 18- prevent answers on facts outside the document (e.g. birthdate of Michael Jordan in the docs vs. the baseball player)19'''20 21import os22import streamlit as st23from llama_index import GPTVectorStoreIndex, SimpleDirectoryReader, ServiceContext, LLMPredictor, PromptHelper, readers24from llama_index import StorageContext, load_index_from_storage25 26from langchain import OpenAI, HuggingFaceHub27 28import app_constants29 30index_fpath = "./llamas_index"31documents_folder = "./documents" #initial documents - additional can be added via upload32 33if "dummy" not in st.session_state:34 st.session_state["dummy"] = "dummy"35 36#@st.cache_resource #st makes this globally available for all users and sessions 37def initialize_index(index_name, documents_folder, persisted_to_storage=True):38 """39 creates an index of the documents in the folder40 if the index exists, skipped41 """42 # set maximum input size43 max_input_size = 409644 # set number of output tokens45 num_outputs = 200046 # set maximum chunk overlap47 max_chunk_overlap = 2048 # set chunk size limit49 chunk_size_limit = 600 50 51 llm_predictor = LLMPredictor(llm=OpenAI(openai_api_key=api_key, #from env52 temperature=0.5, 53 model_name="text-davinci-003", 54 max_tokens=num_outputs)) 55 #wishlist: alternatives56 service_context = ServiceContext.from_defaults(llm_predictor=llm_predictor)57 if os.path.exists(index_name):58 storage_context = StorageContext.from_defaults(persist_dir=index_fpath)59 doc_index = load_index_from_storage(service_context=service_context, storage_context=storage_context)60 else:61 #st.info("Updating the document index")62 prompt_helper = PromptHelper(max_input_size, num_outputs, max_chunk_overlap, chunk_size_limit=chunk_size_limit)63 64 documents = SimpleDirectoryReader(documents_folder).load_data()65 doc_index = GPTVectorStoreIndex.from_documents(66 documents, llm_predictor=llm_predictor, prompt_helper=prompt_helper, 67 chunk_size_limit=512, service_context=service_context68 )69 if persisted_to_storage:70 doc_index.storage_context.persist(index_fpath)71 72 #avoid this side-effect: st.session_state["doc_index"] = "doc_index"73 return doc_index74 75#st returns data that's available for future caller76@st.cache_data(max_entries=200, persist=True) 77def query_index(_index, query_text):78 query_engine = _index.as_query_engine()79 response = query_engine.query(query_text)80 #response = _index.query(query_text)81 return str(response)82 83 84#page format is directly written her85st.title("LLM scanner")86st.markdown(87 (88 "This app allows you to query documents!\n\n"89 "Powered by [Llama Index](https://gpt-index.readthedocs.io/en/latest/index.html)"90 )91)92 93setup_tab, upload_tab, query_tab = st.tabs(94 ["Setup", "Index", "Query"]95)96 97with setup_tab:98 st.subheader("LLM Setup")99 api_key = st.text_input("Enter your OpenAI API key here", type="password")100 101 #wishlist llm_name = st.selectbox(102 # "Which LLM?", ["text-davinci-003", "gpt-3.5-turbo", "gpt-4"]103 #)104 #repo_id = "google/flan-t5-xl" # See https://huggingface.co/models?pipeline_tag=text-generation&sort=downloads for some other options105 #llm = HuggingFaceHub(repo_id=repo_id, model_kwargs={"temperature":0, "max_length":64})106 107 #model_temperature = st.slider(108 # "LLM Temperature", min_value=0.0, max_value=1.0, step=0.1109 #)110 111if api_key is not None and "doc_index" not in st.session_state: 112 st.session_state["doc_index"] = initialize_index(index_fpath, documents_folder, persisted_to_storage=False) 113 114 115with upload_tab:116 st.subheader("Upload documents")117 118 if st.button("Re-initialize index with pre-packaged documents"):119 st.session_state["doc_index"] = initialize_index(index_fpath, documents_folder, persisted_to_storage=False) 120 st.info('Documents in index: ' + str(st.session_state["doc_index"].docstore.docs.__len__()))121 122 if "doc_index" in st.session_state:123 doc_index = st.session_state["doc_index"]124 st.markdown(125 "Either upload a document, or enter the text manually."126 )127 uploaded_file = st.file_uploader(128 "Upload a document (pdf):", type=["pdf"]129 )130 document_text = st.text_area("Enter text")131 if st.button("Add document to index") and (uploaded_file or document_text):132 with st.spinner("Inserting (large files may be slow)..."):133 if document_text:134 doc_index.refresh([readers.Document(text=document_text)]) #tokenizes new documents135 st.info('Documents in index: ' + str(st.session_state["doc_index"].docstore.docs.__len__()))136 137 st.session_state["doc_index"] = doc_index138 if uploaded_file: 139 uploads_folder = "uploads/"140 if not os.path.exists(uploads_folder):141 os.mkdir(uploads_folder)142 #file_details = {"FileName":uploaded_file.name,"FileType":uploaded_file.type}143 with open(uploads_folder + "tmp.pdf", "wb") as f:144 f.write(uploaded_file.getbuffer())145 documents = SimpleDirectoryReader(uploads_folder).load_data()146 doc_index.refresh(documents) #tokenizes new documents147 st.session_state["doc_index"] = doc_index148 st.info('Documents in index: ' + str(st.session_state["doc_index"].docstore.docs.__len__()))149 150 st.session_state["doc_index"] = doc_index151 os.remove(uploads_folder + "tmp.pdf")152 153with query_tab:154 st.subheader("Query Tab")155 st.write("Enter a query about the included documents. Find [documentation here](https://huggingface.co/spaces/agutfraind/llmscanner)")156 157 doc_index = None158 #api_key = st.text_input("Enter your OpenAI API key here:", type="password")159 if api_key:160 os.environ['OPENAI_API_KEY'] = api_key161 #doc_index = initialize_index(index_fpath, documents_folder) 162 163 if doc_index is None:164 if "doc_index" in st.session_state:165 doc_index = st.session_state["doc_index"]166 st.info('Documents in index: ' + str(doc_index.docstore.docs.__len__())) 167 else:168 st.warning("Doc index is not available - initialize or upload")169 #st.warning("Please enter your api key first.")170 171 if doc_index and api_key:172 select_type_your_own = 'type your own...'173 options_for_queries = app_constants.canned_questions + [select_type_your_own]174 query_selection = st.selectbox("Select option", options=options_for_queries)175 query_text = None176 177 if query_selection == select_type_your_own: 178 query_text = st.text_input("Query text")179 else:180 query_text = query_selection 181 182 if st.button("Run Query") and (doc_index is not None) and (query_text is not None):183 response = query_index(doc_index, query_text)184 st.markdown(response)185 186 llm_col, embed_col = st.columns(2)187 with llm_col:188 st.markdown(f"LLM Tokens Used: {doc_index.service_context.llm_predictor._last_token_usage}")189 190 with embed_col:191 st.markdown(f"Embedding Tokens Used: {doc_index.service_context.embed_model._last_token_usage}")192 193