move success check code to setup_qa_chain

app.py

@@ -41,11 +41,10 @@ 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, token=access_token, clean_up_tokenization_spaces=False)
         model = AutoModelForCausalLM.from_pretrained(model_name, token=access_token)
         device = 0 if torch.cuda.is_available() else -1
-        return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.1)
+        return pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=150, temperature=0.1, device=device)
     except Exception as e:
         print(f"Error loading model: {e}")
         return None
@@ -60,6 +59,12 @@ def setup_qa_chain():
         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():
@@ -101,10 +106,6 @@ def main_page():
                     
                     setup_qa_chain()
 
-                    st.session_state.paper_name = paper.name[:-4]
-                    st.session_state.page = "chat"
-                    st.success("PDF successfully processed! You can now ask questions.")
-
             except Exception as e:
                 st.error(f"Failed to process the PDF: {str(e)}")