halyn commited on
Commit
01e574a
ยท
1 Parent(s): ee2bd18

modify 2 page to 1 main page

Browse files
Files changed (1) hide show
  1. app.py +68 -95
app.py CHANGED
@@ -49,93 +49,74 @@ def load_model():
49
  print(f"Error loading model: {e}")
50
  return None
51
 
52
- # QA ์ฒด์ธ ์„ค์ •
53
- def setup_qa_chain():
54
- global qa_chain
55
- try:
56
- pipe = load_model()
57
- except Exception as e:
58
- print(f"Error loading model: {e}")
59
- return
60
- llm = HuggingFacePipeline(pipeline=pipe)
61
- qa_chain = load_qa_chain(llm, chain_type="map_rerank")
62
-
63
-
64
- # st.session_state.paper_name = paper.name[:-4]
65
- st.session_state.page = "chat"
66
- st.success("PDF successfully processed! You can now ask questions.")
67
-
68
-
69
- # ๋ฉ”์ธ ํŽ˜์ด์ง€ UI
70
- def main_page():
71
  st.title("Welcome to GemmaPaperQA")
72
- st.subheader("Upload Your Paper")
73
-
74
- paper = st.file_uploader("Upload Here!", type="pdf", label_visibility="hidden")
75
- if paper:
76
- st.write(f"Upload complete! File name: {paper.name}")
77
- # ํŒŒ์ผ ํฌ๊ธฐ ํ™•์ธ
78
- file_size = paper.size # ํŒŒ์ผ ํฌ๊ธฐ๋ฅผ ํŒŒ์ผ ํฌ์ธํ„ฐ ์ด๋™ ์—†์ด ํ™•์ธ
79
- if file_size > 10 * 1024 * 1024: # 10MB ์ œํ•œ
80
- st.error("File is too large! Please upload a file smaller than 10MB.")
81
- return
82
-
83
- # ์ค‘๊ฐ„ ํ™•์ธ ์ ˆ์ฐจ - PDF ๋‚ด์šฉ ๋ฏธ๋ฆฌ๋ณด๊ธฐ
84
- with st.spinner('Processing PDF...'):
85
- try:
86
- paper.seek(0) # ํŒŒ์ผ ์ฝ๊ธฐ ํฌ์ธํ„ฐ๋ฅผ ์ฒ˜์Œ์œผ๋กœ ๋˜๋Œ๋ฆผ
87
- contents = paper.read()
88
- pdf_file = io.BytesIO(contents)
89
- text = load_pdf(pdf_file)
90
-
91
- # ํ…์ŠคํŠธ๊ฐ€ ์ถ”์ถœ๋˜์ง€ ์•Š์„ ๊ฒฝ์šฐ ์—๋Ÿฌ ์ฒ˜๋ฆฌ
92
- if len(text.strip()) == 0:
93
- st.error("The PDF appears to have no extractable text. Please check the file and try again.")
94
- return
95
-
96
- st.text_area("Preview of extracted text", text[:1000], height=200)
97
- st.write(f"Total characters extracted: {len(text)}")
98
- global knowledge_base
99
- if st.button("Proceed with this file"):
100
- chunks = split_text(text)
101
- knowledge_base = create_knowledge_base(chunks)
102
-
103
- if knowledge_base is None:
104
- st.error("Failed to create knowledge base.")
105
  return
106
-
107
- setup_qa_chain()
108
-
109
- except Exception as e:
110
- st.error(f"Failed to process the PDF: {str(e)}")
111
-
112
 
113
- # ์ฑ„ํŒ… ํŽ˜์ด์ง€ UI
114
- def chat_page():
115
- st.title(f"Ask anything about {st.session_state.paper_name}")
116
-
117
- if "messages" not in st.session_state:
118
- st.session_state.messages = []
119
-
120
- for message in st.session_state.messages:
121
- with st.chat_message(message["role"]):
122
- st.markdown(message["content"])
123
-
124
- if prompt := st.chat_input("Chat here!"):
125
- st.session_state.messages.append({"role": "user", "content": prompt})
126
-
127
- with st.chat_message("user"):
128
- st.markdown(prompt)
129
-
130
- response = get_response_from_model(prompt)
131
-
132
- with st.chat_message("assistant"):
133
- st.markdown(response)
134
-
135
- st.session_state.messages.append({"role": "assistant", "content": response})
136
-
137
- if st.button("Go back to main page"):
138
- st.session_state.page = "main"
 
 
 
 
 
 
 
 
 
 
 
 
139
 
140
  # ๋ชจ๋ธ ์‘๋‹ต ์ฒ˜๋ฆฌ
141
  def get_response_from_model(prompt):
@@ -156,15 +137,7 @@ def get_response_from_model(prompt):
156
  except Exception as e:
157
  return f"Error: {str(e)}"
158
 
159
- # ํŽ˜์ด์ง€ ์„ค์ •
160
- if "page" not in st.session_state:
161
- st.session_state.page = "main"
162
-
163
- if "paper_name" not in st.session_state:
164
- st.session_state.paper_name = ""
165
 
166
- # ํŽ˜์ด์ง€ ๋ Œ๋”๋ง
167
- if st.session_state.page == "main":
168
- main_page()
169
- elif st.session_state.page == "chat":
170
- chat_page()
 
49
  print(f"Error loading model: {e}")
50
  return None
51
 
52
+ # ํŽ˜์ด์ง€ UI
53
+ def main():
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
54
  st.title("Welcome to GemmaPaperQA")
55
+
56
+ # PDF ์—…๋กœ๋“œ ์„น์…˜
57
+ with st.expander("Upload Your Paper", expanded=True):
58
+ paper = st.file_uploader("Upload Here!", type="pdf", label_visibility="hidden")
59
+
60
+ if paper:
61
+ st.write(f"Upload complete! File name: {paper.name}")
62
+
63
+ # ํŒŒ์ผ ํฌ๊ธฐ ํ™•์ธ
64
+ file_size = paper.size # ํŒŒ์ผ ํฌ๊ธฐ๋ฅผ ํŒŒ์ผ ํฌ์ธํ„ฐ ์ด๋™ ์—†์ด ํ™•์ธ
65
+ if file_size > 10 * 1024 * 1024: # 10MB ์ œํ•œ
66
+ st.error("File is too large! Please upload a file smaller than 10MB.")
67
+ return
68
+
69
+ # PDF ํ…์ŠคํŠธ ๋ฏธ๋ฆฌ๋ณด๊ธฐ
70
+ with st.spinner('Processing PDF...'):
71
+ try:
72
+ paper.seek(0)
73
+ contents = paper.read()
74
+ pdf_file = io.BytesIO(contents)
75
+ text = load_pdf(pdf_file)
76
+ print("text:", text)
77
+
78
+ if len(text.strip()) == 0:
79
+ st.error("The PDF appears to have no extractable text. Please check the file and try again.")
 
 
 
 
 
 
 
 
80
  return
 
 
 
 
 
 
81
 
82
+ st.text_area("Preview of extracted text", text[:1000], height=200)
83
+ st.write(f"Total characters extracted: {len(text)}")
84
+
85
+ if st.button("Create Knowledge Base"):
86
+ global knowledge_base
87
+ chunks = split_text(text)
88
+ print("chunks:", chunks)
89
+ knowledge_base = create_knowledge_base(chunks)
90
+ print("knowledge_base:", knowledge_base)
91
+
92
+ if knowledge_base is None:
93
+ st.error("Failed to create knowledge base.")
94
+ return
95
+
96
+ # QA ์ฒด์ธ ์„ค์ •
97
+ global qa_chain
98
+ try:
99
+ pipe = load_model()
100
+ except Exception as e:
101
+ st.error(f"Error loading model: {e}")
102
+ return
103
+ llm = HuggingFacePipeline(pipeline=pipe)
104
+ qa_chain = load_qa_chain(llm, chain_type="map_rerank")
105
+
106
+ st.success("Knowledge base created! You can now ask questions.")
107
+
108
+ except Exception as e:
109
+ st.error(f"Failed to process the PDF: {str(e)}")
110
+
111
+ # ์งˆ๋ฌธ-์‘๋‹ต ์„น์…˜
112
+ if knowledge_base and qa_chain:
113
+ with st.expander("Ask Questions", expanded=True):
114
+ prompt = st.text_input("Chat here!")
115
+
116
+ if prompt:
117
+ response = get_response_from_model(prompt)
118
+ if response:
119
+ st.write(f"**Assistant**: {response}")
120
 
121
  # ๋ชจ๋ธ ์‘๋‹ต ์ฒ˜๋ฆฌ
122
  def get_response_from_model(prompt):
 
137
  except Exception as e:
138
  return f"Error: {str(e)}"
139
 
 
 
 
 
 
 
140
 
141
+ # ์•ฑ ์‹คํ–‰
142
+ if __name__ == "__main__":
143
+ main()