RAG / app.py
mahynski's picture
testing llamaparse
1ab5fce
raw
history blame
3.49 kB
import tempfile
import os
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.core import (
VectorStoreIndex,
Settings,
)
from llama_parse import LlamaParse
from streamlit_pdf_viewer import pdf_viewer
# Global configurations
from llama_index.core import set_global_handler
set_global_handler("langfuse")
st.set_page_config(layout="wide")
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=0
)
# 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']
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 = 4096
# Create LLM
if provider == 'openai':
llm = OpenAI(
model=llm_name,
temperature=temperature,
max_tokens=max_tokens
)
# Global tokenization needs to be consistent with LLM
# https://docs.llamaindex.ai/en/stable/module_guides/models/llms/
Settings.tokenizer = tiktoken.encoding_for_model(llm_name).encode
Settings.num_output = max_tokens
Settings.context_window = 4096 # max possible
# Enter LLM Token
llm_token = st.text_input(
"Enter your LLM token",
value=None
)
# Enter parsing Token
parse_token = st.text_input(
"Enter your LlamaParse token",
value=None
)
uploaded_file = st.file_uploader(
"Choose a PDF file to upload",
type=['pdf'],
accept_multiple_files=False
)
if uploaded_file is not None:
# Parse the file
temp_dir = tempfile.TemporaryDirectory()
parser = LlamaParse(
api_key=parse_token, # can also be set in your env as LLAMA_CLOUD_API_KEY
result_type="text" # "markdown" and "text" are available
)
filename = os.path.join('./', uploaded_file.name)
with open(filename, "wb") as f:
f.write(uploaded_file.getvalue())
parsed_document = parser.load_data(filename)
temp_dir.cleanup()
col1, col2 = st.columns(2)
with col1:
st.write(uploaded_file)
st.write(parsed_document)
with col2:
if uploaded_file is not None:
# Display the pdf
bytes_data = uploaded_file.getvalue()
pdf_viewer(input=bytes_data, width=700)