Update pdfchatbot.py

pdfchatbot.py

@@ -6,16 +6,15 @@ import weaviate
 import os
 from PIL import Image
 from config import MODEL_CONFIG
-from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_openai import OpenAIEmbeddings
 from langchain_weaviate.vectorstores import WeaviateVectorStore
 from langchain.text_splitter import CharacterTextSplitter
-from langchain_community.llms.huggingface_pipeline import HuggingFacePipeline
 from langchain.chains import ConversationalRetrievalChain
 from langchain_community.document_loaders import PyPDFLoader
 from langchain.prompts import PromptTemplate
-from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
 
 os.environ["HUGGINGFACE_API_TOKEN"] = os.getenv("HUGGINGFACE_API_TOKEN")
+os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
 
 class PDFChatBot:
     def __init__(self):
@@ -66,7 +65,7 @@ class PDFChatBot:
         """
         Load embeddings from Hugging Face and set in the config file.
         """
-        self.embeddings = HuggingFaceEmbeddings(model_name=MODEL_CONFIG.MODEL_EMBEDDINGS)
+        self.embeddings = OpenAIEmbeddings(model=MODEL_CONFIG.MODEL_EMBEDDINGS)
 
     def load_vectordb(self):
         """
@@ -82,42 +81,15 @@ class PDFChatBot:
         
         self.vectordb = WeaviateVectorStore.from_documents(docs, self.embeddings, client=weaviate_client)
 
-    def load_tokenizer(self):
-        """
-        Load the tokenizer from Hugging Face and set in the config file.
-        """
-        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_CONFIG.AUTO_TOKENIZER, token=os.getenv("HUGGINGFACE_API_TOKEN"))
-
-    def load_model(self):
-        """
-        Load the causal language model from Hugging Face and set in the config file.
-        """
-        self.model = AutoModelForCausalLM.from_pretrained(
-            MODEL_CONFIG.MODEL_LLM,
-            device_map='auto',
-            torch_dtype=torch.float32,
-            token=os.getenv("HUGGINGFACE_API_TOKEN"),
-            load_in_8bit=False
-        )
-
-    def create_pipeline(self):
-        """
-        Create a pipeline for text generation using the loaded model and tokenizer.
-        """
-        pipe = pipeline(
-            model=self.model,
-            task='text-generation',
-            tokenizer=self.tokenizer,
-            max_new_tokens=200
-        )
-        self.pipeline = HuggingFacePipeline(pipeline=pipe)
 
     def create_chain(self):
         """
         Create a Conversational Retrieval Chain
         """
+        llm = OpenAI(openai_api_key=os.getenv("OPENAI_API_KEY"))
+        
         self.chain = ConversationalRetrievalChain.from_llm(
-            self.pipeline,
+            llm,
             chain_type="stuff",
             retriever=self.vectordb.as_retriever(search_kwargs={"k": 1}),
             condense_question_prompt=self.prompt,
@@ -135,9 +107,6 @@ class PDFChatBot:
         self.documents = PyPDFLoader(file.name).load()
         self.load_embeddings()
         self.load_vectordb()
-        self.load_tokenizer()
-        self.load_model()
-        self.create_pipeline()
         self.create_chain()
 
     def generate_response(self, history, query, file):