rag.py: use optionally HuggingFaceHub zephyr model

src/rag.py

@@ -6,9 +6,9 @@ from langchain.text_splitter import RecursiveCharacterTextSplitter
 
 # from langchain.embeddings import OpenAIEmbeddings
 from langchain_community.embeddings import HuggingFaceEmbeddings
-
 from langchain.vectorstores import Chroma
 from langchain.chat_models import ChatOpenAI
+from langchain.llms import HuggingFaceHub
 from langchain.chains import RetrievalQA
 from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
 from tqdm import tqdm
@@ -24,7 +24,13 @@ class RAG():
         self.pdfs = pdfs  # Source PDFs to encode in vectorestore
         self.k = 3  # Number of relevant chunks to retrieve
 
-        # Load environment variables that should contain a 'OPENAI_API_KEY'
+        # Constants
+        # self.use_model = 'gpt-4o-mini'
+        self.use_model = 'zephyr-7b-alpha'
+
+        # Load environment variables that should contain:
+        #  - 'OPENAI_API_KEY' for OpenAI models
+        #  - 'HUGGINGFACEHUB_API_TOKEN' for HuggingFace models
         dotenv.load_dotenv(dotenv.find_dotenv())
 
         # Placeholders:
@@ -59,20 +65,30 @@ class RAG():
 
     def create_embeddings(self):
         # embeddings = OpenAIEmbeddings()
+        print ('Using Embeddings from HuggingFace: all-MiniLM-L6-v2')
         embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2")
         return embeddings
 
     def create_retriever(self, texts, embeddings):
         # Create embeddings and vector store
+        print ('Creating vectore store with Chroma')
         vectorstore = Chroma.from_documents(texts, embeddings)
         retriever = vectorstore.as_retriever(search_kwargs={"k": self.k})
         return retriever
 
     def create_llm(self):
         # Create the language model
-        llm = ChatOpenAI(
-            model_name="gpt-4o-mini",
-            temperature=0)
+        if self.use_model == 'gpt-4o-mini':
+            print(f'As llm, using OpenAI model: {self.use_model}')
+            llm = ChatOpenAI(
+                model_name="gpt-4o-mini",
+                temperature=0)
+        elif self.use_model ==  'zephyr-7b-alpha':
+            print(f'As llm, using HF model: {self.use_model}')
+            llm = HuggingFaceHub(
+                repo_id="huggingfaceh4/zephyr-7b-alpha", 
+                model_kwargs={"temperature": 0.5, "max_length": 64,"max_new_tokens":512}
+                )
         return llm
 
     def create_QAbot(self, retriever, llm):