use gpu and new qa_chain

app.py

@@ -1,9 +1,11 @@
 import io
 import os
+import torch
 import streamlit as st
 from PyPDF2 import PdfReader
 from langchain.text_splitter import CharacterTextSplitter
 from langchain.chains.question_answering import load_qa_chain
+from langchain.chains import load_qa_with_sources_chain
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import FAISS
 from langchain_community.llms import HuggingFacePipeline
@@ -41,6 +43,7 @@ def load_model():
     model_name = "google/gemma-2-2b"  # Hugging Face 모델 ID
     access_token = os.getenv("HF_TOKEN")
     try:
+        device = 0 if torch.cuda.is_available() else -1
         tokenizer = AutoTokenizer.from_pretrained(model_name, use_auth_token=access_token, clean_up_tokenization_spaces=False)
         model = AutoModelForCausalLM.from_pretrained(model_name, use_auth_token=access_token)
         return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.1)
@@ -57,7 +60,7 @@ def setup_qa_chain():
         print(f"Error loading model: {e}")
         return
     llm = HuggingFacePipeline(pipeline=pipe)
-    qa_chain = load_qa_chain(llm, chain_type="stuff")
+    qa_chain = load_qa_with_sources_chain(llm)
     
 # 메인 페이지 UI
 def main_page():