modify code

app.py

@@ -10,9 +10,7 @@ from langchain_community.vectorstores import FAISS
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 
-# Global variables
-knowledge_base = None
-qa_chain = None
+# Global variables are no longer needed, we will use session state
 
 # PDF 파일 로드 및 텍스트 추출
 def load_pdf(pdf_file):
@@ -82,24 +80,22 @@ def main():
                     st.write(f"Total characters extracted: {len(text)}")
 
                     if st.button("Create Knowledge Base"):
-                        global knowledge_base
                         chunks = split_text(text)
-                        knowledge_base = create_knowledge_base(chunks)
-                        print("knowledge_base:", knowledge_base)
+                        st.session_state.knowledge_base = create_knowledge_base(chunks)
+                        print("knowledge_base:", st.session_state.knowledge_base)
 
-                        if knowledge_base is None:
+                        if st.session_state.knowledge_base is None:
                             st.error("Failed to create knowledge base.")
                             return
 
                         # QA 체인 설정
-                        global qa_chain
                         try:
                             pipe = load_model()
                         except Exception as e:
                             st.error(f"Error loading model: {e}")
                             return
                         llm = HuggingFacePipeline(pipeline=pipe)
-                        qa_chain = load_qa_chain(llm, chain_type="map_rerank")
+                        st.session_state.qa_chain = load_qa_chain(llm, chain_type="map_rerank")
 
                         st.success("Knowledge base created! You can now ask questions.")
 
@@ -107,7 +103,7 @@ def main():
                     st.error(f"Failed to process the PDF: {str(e)}")
 
     # 질문-응답 섹션
-    if knowledge_base and qa_chain:
+    if "knowledge_base" in st.session_state and "qa_chain" in st.session_state:
         with st.expander("Ask Questions", expanded=True):
             prompt = st.text_input("Chat here!")
 
@@ -121,14 +117,13 @@ def main():
 # 모델 응답 처리
 def get_response_from_model(prompt):
     try:
-        global knowledge_base, qa_chain
-        if not knowledge_base:
+        if "knowledge_base" not in st.session_state:
             return "No PDF has been uploaded yet."
-        if not qa_chain:
+        if "qa_chain" not in st.session_state:
             return "QA chain is not initialized."
 
-        docs = knowledge_base.similarity_search(prompt)
-        response = qa_chain.run(input_documents=docs, question=prompt)
+        docs = st.session_state.knowledge_base.similarity_search(prompt)
+        response = st.session_state.qa_chain.run(input_documents=docs, question=prompt)
 
         if "Helpful Answer:" in response:
             response = response.split("Helpful Answer:")[1].strip()