code update

app.py

@@ -1,4 +1,3 @@
-import os
 import io
 import streamlit as st
 from PyPDF2 import PdfReader
@@ -8,6 +7,7 @@ from langchain.vectorstores import FAISS
 from langchain.chains.question_answering import load_qa_chain
 from langchain.llms import HuggingFacePipeline
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
+from peft import PeftModel, PeftConfig
 
 # Global variables
 knowledge_base = None
@@ -22,7 +22,10 @@ def load_pdf(pdf_file):
 # 텍스트를 청크로 분할
 def split_text(text):
     text_splitter = CharacterTextSplitter(
-        separator="\n", chunk_size=1000, chunk_overlap=200, length_function=len
+        separator="\n", 
+        chunk_size=1000, 
+        chunk_overlap=200, 
+        length_function=len
     )
     return text_splitter.split_text(text)
 
@@ -35,9 +38,11 @@ def create_knowledge_base(chunks):
 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)
+    config = PeftConfig.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path)
+    model = PeftModel.from_pretrained(model, model_name)
 
+    return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.1)
 # QA 체인 설정
 def setup_qa_chain():
     global qa_chain
@@ -46,7 +51,6 @@ def setup_qa_chain():
     qa_chain = load_qa_chain(llm, chain_type="stuff")
     
 
-
 # 메인 페이지 UI
 def main_page():
     st.title("Welcome to GemmaPaperQA")
@@ -90,7 +94,7 @@ def main_page():
                     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)}")
 

requirements.txt

@@ -9,4 +9,5 @@ torch==2.4.1
 faiss-cpu==1.8.0.post1
 requests==2.32.3
 huggingface-hub==0.25.1
-sentence-transformers==3.1.1
+sentence-transformers==3.1.1
+peft==0.2.0