Atif20024's picture
Update app.py
bc58e4f verified
import streamlit as st
from doc_loading import read_pdf_text
from llm_functions import get_conversational_chain
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_openai import OpenAIEmbeddings
from langchain.vectorstores import FAISS
import warnings
warnings.filterwarnings('ignore')
from dotenv import load_dotenv
load_dotenv()
from dataclasses import dataclass
import os
api_key = os.getenv("OPENAI_API_KEY")
from openai import OpenAI
client = OpenAI(api_key=api_key)
st.set_page_config(page_title="Resume QA")
st.header("Feel free to ask any question")
st.write("Welcome!!")
# st.session_state
if "option" not in st.session_state:
option = st.selectbox("How are you going to input your document?",
("Upload LinkedIn PDF", "Upload xml"))
# Conditionally show components based on user's choice
file_name = ""
main_text = ""
if option == "Upload LinkedIn PDF":
uploaded_file = st.file_uploader("Choose a PDF file", type=["pdf"], key="uploading_pdf")
if uploaded_file is not None:
# Process the uploaded file
st.write("PDF uploaded successfully!")
main_text = read_pdf_text(uploaded_file)
else:
st.write("This functionality in under construction.")
if option == "Upload LinkedIn PDF":
if uploaded_file is not None and uploaded_file.type != "application/pdf":
st.error("Please choose a PDF file only.")
if len(main_text) > 0:
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=500,
chunk_overlap=50,
length_function=len,
is_separator_regex='\n'
)
texts = text_splitter.create_documents([main_text])
embeddings = OpenAIEmbeddings(api_key=api_key)
# store in vector db
db = FAISS.from_documents(texts, embeddings)
hash_name = f"{option.replace(' ', '-')}"
db.save_local(f'faiss_{hash_name}_index')
if "messages" not in st.session_state:
st.session_state.messages = []
st.write("What's your question? (try being a bit detailed about the question)")
@dataclass
class Message:
actor: str
payload: str
USER = "user"
ASSISTANT = "ai"
MESSAGES = "messages"
def initialize_session_state():
if MESSAGES not in st.session_state:
st.session_state[MESSAGES] = [Message(actor=ASSISTANT, payload="Hi!How can I help you?")]
if "llm_chain" not in st.session_state:
st.session_state["llm_chain"] = get_conversational_chain()
def get_llm_chain_from_session():
return st.session_state["llm_chain"]
initialize_session_state()
msg: Message
for msg in st.session_state[MESSAGES]:
st.chat_message(msg.actor).write(msg.payload)
prompt = st.chat_input("Enter a prompt here")
if prompt:
docs = db.similarity_search(prompt, k=5)
st.session_state[MESSAGES].append(Message(actor=USER, payload=prompt))
st.chat_message(USER).write(prompt)
with st.spinner("Please wait.."):
llm_chain = get_llm_chain_from_session()
response = llm_chain({"context": docs, "question": prompt})['text']
st.session_state[MESSAGES].append(Message(actor=ASSISTANT, payload=response))
st.chat_message(ASSISTANT).write(response)