Update retrieval_qa_pipeline.py

retrieval_qa_pipeline.py

@@ -1,5 +1,3 @@
-# retrieval_qa_pipeline.py
-
 from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 from langchain.embeddings import HuggingFaceEmbeddings
 from langchain.vectorstores import FAISS
@@ -10,26 +8,26 @@ from datasets import load_dataset
 def load_model_and_tokenizer(model_name: str):
     """
     Load the pre-trained model and tokenizer from the Hugging Face Hub.
-    
+
     Args:
         model_name (str): The Hugging Face repository name of the model.
-        
+
     Returns:
         model: The loaded model.
         tokenizer: The loaded tokenizer.
     """
     print(f"Loading model and tokenizer from {model_name}...")
-    model = AutoModelForCausalLM.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     return model, tokenizer
 
 def load_dataset_from_hf(dataset_name: str):
     """
     Load the dataset from the Hugging Face Hub.
-    
+
     Args:
         dataset_name (str): The Hugging Face repository name of the dataset.
-        
+
     Returns:
         texts (list): The text descriptions from the dataset.
         metadatas (list): Metadata for each text (e.g., upf_code).
@@ -43,27 +41,27 @@ def load_dataset_from_hf(dataset_name: str):
 def load_faiss_index(faiss_index_path: str):
     """
     Load the FAISS index and associated embeddings.
-    
+
     Args:
         faiss_index_path (str): Path to the saved FAISS index.
-        
+
     Returns:
         vectorstore (FAISS): The FAISS vector store.
     """
     print(f"Loading FAISS index from {faiss_index_path}...")
     embeddings = HuggingFaceEmbeddings()  # Default embeddings
-    vectorstore = FAISS.load_local(faiss_index_path, embeddings)
+    vectorstore = FAISS.load_local(faiss_index_path, embeddings, allow_dangerous_deserialization=True)
     return vectorstore
 
 def build_retrieval_qa_pipeline(model, tokenizer, vectorstore):
     """
     Build the retrieval-based QA pipeline.
-    
+
     Args:
         model: The pre-trained model.
         tokenizer: The tokenizer associated with the model.
         vectorstore (FAISS): The FAISS vector store for retrieval.
-        
+
     Returns:
         qa_chain (RetrievalQA): The retrieval-based QA pipeline.
     """
@@ -77,11 +75,11 @@ def build_retrieval_qa_pipeline(model, tokenizer, vectorstore):
         top_p=0.95,
         repetition_penalty=1.15
     )
-    
+
     llm = HuggingFacePipeline(pipeline=hf_pipeline)
     retriever = vectorstore.as_retriever()
     qa_chain = RetrievalQA.from_chain_type(llm=llm, retriever=retriever)
-    
+
     return qa_chain
 
 def main():
@@ -89,21 +87,21 @@ def main():
     model_name = "anirudh248/upf_code_generator_final"
     dataset_name = "PranavKeshav/upf_dataset"
     faiss_index_path = "faiss_index"
-    
+
     print("Starting pipeline setup...")
-    
+
     # Load model and tokenizer
     model, tokenizer = load_model_and_tokenizer(model_name)
-    
+
     # Load dataset
     texts, metadatas = load_dataset_from_hf(dataset_name)
-    
+
     # Load FAISS index
     vectorstore = load_faiss_index(faiss_index_path)
-    
+
     # Build QA pipeline
     qa_chain = build_retrieval_qa_pipeline(model, tokenizer, vectorstore)
-    
+
     # Test the pipeline
     print("Pipeline is ready! You can now ask questions.")
     while True: