modify 2 page to 1 main page

app.py

@@ -49,93 +49,74 @@ def load_model():
         print(f"Error loading model: {e}")
         return None
 
-# QA 체인 설정
-def setup_qa_chain():
-    global qa_chain
-    try:
-        pipe = load_model()  
-    except Exception as e:
-        print(f"Error loading model: {e}")
-        return
-    llm = HuggingFacePipeline(pipeline=pipe)
-    qa_chain = load_qa_chain(llm, chain_type="map_rerank")
-
-    
-    # st.session_state.paper_name = paper.name[:-4]
-    st.session_state.page = "chat"
-    st.success("PDF successfully processed! You can now ask questions.")
-
-
-# 메인 페이지 UI
-def main_page():
+# 페이지 UI
+def main():
     st.title("Welcome to GemmaPaperQA")
-    st.subheader("Upload Your Paper")
-
-    paper = st.file_uploader("Upload Here!", type="pdf", label_visibility="hidden")
-    if paper:
-        st.write(f"Upload complete! File name: {paper.name}")
-        # 파일 크기 확인
-        file_size = paper.size  # 파일 크기를 파일 포인터 이동 없이 확인
-        if file_size > 10 * 1024 * 1024:  # 10MB 제한
-            st.error("File is too large! Please upload a file smaller than 10MB.")
-            return
-
-        # 중간 확인 절차 - PDF 내용 미리보기
-        with st.spinner('Processing PDF...'):
-            try:
-                paper.seek(0)  # 파일 읽기 포인터를 처음으로 되돌림
-                contents = paper.read()
-                pdf_file = io.BytesIO(contents)
-                text = load_pdf(pdf_file)
-
-                # 텍스트가 추출되지 않을 경우 에러 처리
-                if len(text.strip()) == 0:
-                    st.error("The PDF appears to have no extractable text. Please check the file and try again.")
-                    return
-
-                st.text_area("Preview of extracted text", text[:1000], height=200)
-                st.write(f"Total characters extracted: {len(text)}")
-                global knowledge_base
-                if st.button("Proceed with this file"):
-                    chunks = split_text(text)
-                    knowledge_base = create_knowledge_base(chunks)
-                    
-                    if knowledge_base is None:
-                        st.error("Failed to create knowledge base.")
+    
+    # PDF 업로드 섹션
+    with st.expander("Upload Your Paper", expanded=True):
+        paper = st.file_uploader("Upload Here!", type="pdf", label_visibility="hidden")
+        
+        if paper:
+            st.write(f"Upload complete! File name: {paper.name}")
+            
+            # 파일 크기 확인
+            file_size = paper.size  # 파일 크기를 파일 포인터 이동 없이 확인
+            if file_size > 10 * 1024 * 1024:  # 10MB 제한
+                st.error("File is too large! Please upload a file smaller than 10MB.")
+                return
+
+            # PDF 텍스트 미리보기
+            with st.spinner('Processing PDF...'):
+                try:
+                    paper.seek(0)
+                    contents = paper.read()
+                    pdf_file = io.BytesIO(contents)
+                    text = load_pdf(pdf_file)
+                    print("text:", text)
+
+                    if len(text.strip()) == 0:
+                        st.error("The PDF appears to have no extractable text. Please check the file and try again.")
                         return
-                    
-                    setup_qa_chain()
-
-            except Exception as e:
-                st.error(f"Failed to process the PDF: {str(e)}")
-
 
-# 채팅 페이지 UI
-def chat_page():
-    st.title(f"Ask anything about {st.session_state.paper_name}")
-
-    if "messages" not in st.session_state:
-        st.session_state.messages = []
-
-    for message in st.session_state.messages:
-        with st.chat_message(message["role"]):
-            st.markdown(message["content"])
-
-    if prompt := st.chat_input("Chat here!"):
-        st.session_state.messages.append({"role": "user", "content": prompt})
-
-        with st.chat_message("user"):
-            st.markdown(prompt)
-
-        response = get_response_from_model(prompt)
-
-        with st.chat_message("assistant"):
-            st.markdown(response)
-
-        st.session_state.messages.append({"role": "assistant", "content": response})
-
-    if st.button("Go back to main page"):
-        st.session_state.page = "main"
+                    st.text_area("Preview of extracted text", text[:1000], height=200)
+                    st.write(f"Total characters extracted: {len(text)}")
+
+                    if st.button("Create Knowledge Base"):
+                        global knowledge_base
+                        chunks = split_text(text)
+                        print("chunks:", chunks)
+                        knowledge_base = create_knowledge_base(chunks)
+                        print("knowledge_base:", knowledge_base)
+
+                        if 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.success("Knowledge base created! You can now ask questions.")
+
+                except Exception as e:
+                    st.error(f"Failed to process the PDF: {str(e)}")
+
+    # 질문-응답 섹션
+    if knowledge_base and qa_chain:
+        with st.expander("Ask Questions", expanded=True):
+            prompt = st.text_input("Chat here!")
+
+            if prompt:
+                response = get_response_from_model(prompt)
+                if response:
+                    st.write(f"**Assistant**: {response}")
 
 # 모델 응답 처리
 def get_response_from_model(prompt):
@@ -156,15 +137,7 @@ def get_response_from_model(prompt):
     except Exception as e:
         return f"Error: {str(e)}"
 
-# 페이지 설정
-if "page" not in st.session_state:
-    st.session_state.page = "main"
-
-if "paper_name" not in st.session_state:
-    st.session_state.paper_name = ""
 
-# 페이지 렌더링
-if st.session_state.page == "main":
-    main_page()
-elif st.session_state.page == "chat":
-    chat_page()
+# 앱 실행
+if __name__ == "__main__":
+    main()