File size: 4,939 Bytes
b0e3bba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6993f7f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
b0e3bba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
cf6adbe
b0e3bba
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
import os
from langchain_community.document_loaders import PyPDFLoader
from langchain_community.vectorstores import faiss
from langchain.memory import ConversationBufferMemory
from langchain_google_genai import ChatGoogleGenerativeAI, GoogleGenerativeAIEmbeddings
from tempfile import NamedTemporaryFile
from dotenv import load_dotenv
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain.chains import ConversationalRetrievalChain
import streamlit as st
import nest_asyncio

nest_asyncio.apply()
load_dotenv()

# Initialize app resources
st.set_page_config(page_title="StudyAssist", page_icon=":book:")
st.title("StudyAssist(pharmassist-v0)")
st.write(
    "An AI/RAG application to aid students in their studies, specially optimized for the pharm 028 students. In simpler terms, chat with your pdf"
)


@st.cache_resource
def initialize_resources():
    llm_gemini = ChatGoogleGenerativeAI(
        model="gemini-1.5-flash-latest", google_api_key=os.getenv("GOOGLE_API_KEY")
    )
    return llm_gemini


def get_retriever(pdf_file):
    with NamedTemporaryFile(suffix="pdf") as temp:
        temp.write(pdf_file.getvalue())
        pdf_loader = PyPDFLoader(temp.name, extract_images=True)
        pages = pdf_loader.load()

    # st.write(f"AI Chatbot for {course_material}")

    underlying_embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001")
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1000,
        chunk_overlap=20,
        length_function=len,
        is_separator_regex=False,
        separators="\n",
    )
    documents = text_splitter.split_documents(pages)
    vectorstore = faiss.FAISS.from_documents(documents, underlying_embeddings)
    doc_retiever = vectorstore.as_retriever(
        search_type="mmr", search_kwargs={"k": 5, "fetch_k": 10}
    )

    return doc_retiever


chat_model = initialize_resources()

# Streamlit UI
# Course list and pdf retrieval

courses = ["PMB", "PCL", "Kelechi_research"]  # "GSP", "CPM", "PCG",  "PCH"
course_pdfs = None
doc_retriever = None
conversational_chain = None

# course = st.sidebar.selectbox("Choose course", (courses))
# docs_path = f"pdfs/{course}"
# course_pdfs = os.listdir(docs_path)
# pdfs = [os.path.join(docs_path, pdf) for pdf in course_pdfs]

course_material = "{Not selected}"


# @st.cache_resource
def query_response(query, _retriever):
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    conversational_chain = ConversationalRetrievalChain.from_llm(
        llm=chat_model, retriever=_retriever, memory=memory, verbose=False
    )
    response = conversational_chain.run(query)

    return response


if "doc" not in st.session_state:
    st.session_state.doc = ""

course_material = st.file_uploader("or Upload your own pdf", type="pdf")

if st.session_state != "":
    try:
        with st.spinner("loading document.."):
            doc_retriever = get_retriever(course_material)
        st.success("File loading successful, vector db initialize")
    except Exception as e:
        st.error(e)

    # We store the conversation in the session state.
    # This will be use to render the chat conversation.
    # We initialize it with the first message we want to be greeted with.
    if "messages" not in st.session_state:
        st.session_state.messages = [
            {"role": "assistant", "content": "Yoo, How far boss?"}
        ]

    if "current_response" not in st.session_state:
        st.session_state.current_response = ""

    # We loop through each message in the session state and render it as
    # a chat message.
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # We take questions/instructions from the chat input to pass to the LLM
    if user_prompt := st.chat_input("Ask...", key="user_input"):
        # Add our input to the session state
        st.session_state.messages.append({"role": "user", "content": user_prompt})

        # Add our input to the chat window
        with st.chat_message("user"):
            st.markdown(user_prompt)

        # Pass our input to the llm chain and capture the final responses.
        # here once the llm has finished generating the complete response.
        response = query_response(user_prompt, doc_retriever)
        # Add the response to the session state
        st.session_state.messages.append({"role": "assistant", "content": response})

        # Add the response to the chat window
        with st.chat_message("assistant"):
            st.markdown(response)
#
st.write("")
st.write("")


st.markdown(
    """
    <div style="text-align: center; padding: 1rem;">
        Project by <a href="https://github.com/kelechi-c" target="_blank" style="color: white; font-weight: bold; text-decoration: none;">
         kelechi(tensor)</a>
    </div>
""",
    unsafe_allow_html=True,
)