File size: 5,518 Bytes
0c7bd3c 1ab5fce 1e49352 549dce5 1e6d996 64b8f12 feeb9a7 0c7bd3c 7a9e2a5 a487212 e4585e1 a487212 8178bd6 a487212 23d4dee a487212 e9802f8 cbbfa1c a487212 a94912c e9802f8 a487212 e9802f8 a487212 a94912c a487212 db64926 a94912c e4585e1 a487212 e9802f8 a487212 e9802f8 cbbfa1c a487212 a94912c a487212 e9802f8 a487212 e9802f8 a487212 a94912c a487212 0c7bd3c 1ab5fce c1ff6b3 a487212 |
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 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 |
import tempfile
import os
import tiktoken
import streamlit as st
from llama_index.llms.gemini import Gemini
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.llms.mistralai import MistralAI
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import (
VectorStoreIndex,
Settings,
)
from llama_parse import LlamaParse
from streamlit_pdf_viewer import pdf_viewer
def main():
with st.sidebar:
st.title('Document Summarization and QA System')
# st.markdown('''
# ## About this application
# Upload a pdf to ask questions about it. This retrieval-augmented generation (RAG) workflow uses:
# - [Streamlit](https://streamlit.io/)
# - [LlamaIndex](https://docs.llamaindex.ai/en/stable/)
# - [OpenAI](https://platform.openai.com/docs/models)
# ''')
# st.write('Made by ***Nate Mahynski***')
# st.write('nathan.mahynski@nist.gov')
# Select Provider
provider = st.selectbox(
label="Select LLM Provider",
options=['google', 'huggingface', 'mistralai', 'openai'],
index=3
)
# Select LLM
if provider == 'google':
llm_list = ['gemini']
elif provider == 'huggingface':
llm_list = []
elif provider == 'mistralai':
llm_list =[]
elif provider == 'openai':
llm_list = ['gpt-3.5-turbo', 'gpt-4', 'gpt-4-turbo', 'gpt-4o', 'gpt-4o-mini']
else:
llm_list = []
llm_name = st.selectbox(
label="Select LLM Model",
options=llm_list,
index=0
)
# Temperature
temperature = st.slider(
"Temperature",
min_value=0.0,
max_value=1.0,
value=0.0,
step=0.05,
)
max_output_tokens = 2048
# Enter LLM API Key
llm_key = st.text_input(
"Enter your LLM API Key",
value=None,
)
# Create LLM
# Global tokenization needs to be consistent with LLM for token counting
# https://docs.llamaindex.ai/en/stable/module_guides/models/llms/
if llm_key is not None:
if provider == 'openai':
os.environ["OPENAI_API_KEY"] = str(llm_key)
Settings.llm = OpenAI(
model=llm_name,
temperature=temperature,
max_tokens=max_output_tokens
)
Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
Settings.num_output = max_output_tokens
Settings.context_window = 4096 # max possible
Settings.embed_model = OpenAIEmbedding()
elif provider == 'huggingface':
os.environ['HFTOKEN'] = str(llm_key)
# Enter Parsing API Key
parse_key = st.text_input(
"Enter your LlamaParse API Key",
value=None,
)
uploaded_file = st.file_uploader(
"Choose a PDF file to upload",
type=['pdf'],
accept_multiple_files=False
)
parsed_document = None
if uploaded_file is not None:
# Parse the file
parser = LlamaParse(
api_key=parse_key, # Can also be set in your env as LLAMA_CLOUD_API_KEY
result_type="text" # "markdown" and "text" are available
)
# Create a temporary directory to save the file then load and parse it
temp_dir = tempfile.TemporaryDirectory()
temp_filename = os.path.join(temp_dir.name, uploaded_file.name)
with open(temp_filename, "wb") as f:
f.write(uploaded_file.getvalue())
parsed_document = parser.load_data(temp_filename)
temp_dir.cleanup()
col1, col2 = st.columns(2)
with col1:
st.markdown(
"""
# Instructions
1. Obtain an [API Key](https://cloud.llamaindex.ai/api-key) from LlamaParse to parse your document.
2. Obtain a similar API Key from your preferred LLM provider.
3. Make selections at the left and upload a document to use a context.
4. Begin asking questions below!
"""
)
st.divider()
prompt_txt = 'Summarize this document in a 3-5 sentences.'
prompt = st.text_area(
label="Enter your query.",
key="prompt_widget",
value=prompt_txt
)
if parsed_document is not None:
index = VectorStoreIndex.from_documents(parsed_document)
query_engine = index.as_query_engine()
response = query_engine.query(prompt)
st.write(response.response)
with col2:
tab1, tab2 = st.tabs(["Uploaded File", "Parsed File",])
with tab1:
if uploaded_file is not None: # Display the pdf
bytes_data = uploaded_file.getvalue()
pdf_viewer(input=bytes_data, width=700)
with tab2:
if parsed_document is not None: # Showed the raw parsing result
st.write(parsed_document)
if __name__ == '__main__':
# Global configurations
from llama_index.core import set_global_handler
set_global_handler("langfuse")
st.set_page_config(layout="wide")
main() |