code update

app.py

@@ -1,6 +1,5 @@
 import os
 import io
-import requests
 import streamlit as st
 from PyPDF2 import PdfReader
 from langchain.text_splitter import CharacterTextSplitter
@@ -14,74 +13,90 @@ from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 knowledge_base = None
 qa_chain = None
 
+# PDF 파일 로드 및 텍스트 추출
 def load_pdf(pdf_file):
-    """
-    Load and extract text from a PDF.
-    """
     pdf_reader = PdfReader(pdf_file)
     text = "".join(page.extract_text() for page in pdf_reader.pages)
     return text
 
+# 텍스트를 청크로 분할
 def split_text(text):
-    """
-    Split the extracted text into chunks.
-    """
     text_splitter = CharacterTextSplitter(
         separator="\n", chunk_size=1000, chunk_overlap=200, length_function=len
     )
     return text_splitter.split_text(text)
 
+# FAISS 벡터 저장소 생성
 def create_knowledge_base(chunks):
-    """
-    Create a FAISS knowledge base from text chunks.
-    """
     embeddings = HuggingFaceEmbeddings()
     return FAISS.from_texts(chunks, embeddings)
 
-def load_model(model_path):
-    """
-    Load the HuggingFace model and tokenizer, and create a text-generation pipeline.
-    """
-    tokenizer = AutoTokenizer.from_pretrained(model_path)
-    model = AutoModelForCausalLM.from_pretrained(model_path)
+# Hugging Face 모델 로드
+def load_model():
+    model_name = "halyn/gemma2-2b-it-finetuned-paperqa"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
     return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.1)
 
+# QA 체인 설정
 def setup_qa_chain():
-    """
-    Set up the question-answering chain.
-    """
     global qa_chain
-    pipe = load_model(MODEL_PATH)
+    pipe = load_model()
     llm = HuggingFacePipeline(pipeline=pipe)
     qa_chain = load_qa_chain(llm, chain_type="stuff")
+    
 
-# Streamlit UI
+
+# 메인 페이지 UI
 def main_page():
     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 is {paper.name}")
-        st.write("Please click the button below.")
-
-        if st.button("Click Here :)"):
+        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:
-                # PDF 파일 처리
+                paper.seek(0)  # 파일 읽기 포인터를 처음으로 되돌림
                 contents = paper.read()
                 pdf_file = io.BytesIO(contents)
                 text = load_pdf(pdf_file)
-                chunks = split_text(text)
-                global knowledge_base
-                knowledge_base = create_knowledge_base(chunks)
-
-                st.success("PDF successfully processed! You can now ask questions.")
-                st.session_state.paper_name = paper.name[:-4]
-                st.session_state.page = "chat"
-                setup_qa_chain()
+
+                # 텍스트가 추출되지 않을 경우 에러 처리
+                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)}")
+
+                if st.button("Proceed with this file"):
+                    chunks = split_text(text)
+                    global knowledge_base
+                    knowledge_base = create_knowledge_base(chunks)
+                    
+                    if knowledge_base is None:
+                        st.error("Failed to create knowledge base.")
+                        return
+                    
+                    st.session_state.paper_name = paper.name[:-4]
+                    st.session_state.page = "chat"
+                    setup_qa_chain()
+                    st.success("PDF successfully processed! You can now ask questions.")
+                    
             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}")
 
@@ -108,6 +123,7 @@ def chat_page():
     if st.button("Go back to main page"):
         st.session_state.page = "main"
 
+# 모델 응답 처리
 def get_response_from_model(prompt):
     try:
         global knowledge_base, qa_chain
@@ -126,11 +142,10 @@ def get_response_from_model(prompt):
     except Exception as e:
         return f"Error: {str(e)}"
 
-# Streamlit - 초기 페이지 설정
+# 페이지 설정
 if "page" not in st.session_state:
     st.session_state.page = "main"
 
-# paper_name 초기화
 if "paper_name" not in st.session_state:
     st.session_state.paper_name = ""
 

requirements.txt

@@ -5,4 +5,4 @@ transformers==4.31.0
 torch==2.0.1
 faiss-cpu==1.7.4
 requests==2.31.0
-huggingface-hub==0.16.4
+huggingface-hub==0.16.4