Changing the embedding model

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import os
 from langchain_community.document_loaders import PyMuPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_openai import ChatOpenAI
 from langchain_qdrant import QdrantVectorStore
 from langchain.prompts import ChatPromptTemplate
 from langchain_core.output_parsers import StrOutputParser
@@ -10,6 +10,7 @@ from langchain_core.runnables import RunnablePassthrough
 from qdrant_client import QdrantClient
 from qdrant_client.http.models import Distance, VectorParams
 from operator import itemgetter
+from sentence_transformers import SentenceTransformer
 
 # Set up API keys
 os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
@@ -20,6 +21,16 @@ pdf_links = [
     "https://www.whitehouse.gov/wp-content/uploads/2022/10/Blueprint-for-an-AI-Bill-of-Rights.pdf"
 ]
 
+class CustomHuggingFaceEmbeddings:
+    def __init__(self, model_name):
+        self.model = SentenceTransformer(model_name)
+
+    def embed_documents(self, texts):
+        return self.model.encode(texts).tolist()
+
+    def embed_query(self, text):
+        return self.model.encode(text).tolist()
+
 @st.cache_resource
 def load_and_process_pdfs(pdf_links):
     documents = []
@@ -40,7 +51,7 @@ def load_and_process_pdfs(pdf_links):
 def setup_vectorstore():
     LOCATION = ":memory:"
     COLLECTION_NAME = "AI_Ethics_Framework"
-    VECTOR_SIZE = 1536
+    VECTOR_SIZE = 768  # Adjust this if your model's output size is different
 
     qdrant_client = QdrantClient(location=LOCATION)
 
@@ -50,11 +61,12 @@ def setup_vectorstore():
         vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
     )
 
-    # Create the vector store
+    # Create the vector store with your custom embeddings
+    custom_embeddings = CustomHuggingFaceEmbeddings("Technocoloredgeek/midterm-finetuned-embedding")
     qdrant_vector_store = QdrantVectorStore(
         client=qdrant_client,
         collection_name=COLLECTION_NAME,
-        embedding=OpenAIEmbeddings()
+        embedding=custom_embeddings
     )
 
     # Load and add documents
@@ -69,7 +81,6 @@ def create_rag_pipeline(_vector_store):
     
     template = """
     You are an expert AI assistant with deep knowledge of business, technology, and entrepreneurship. Your task is to provide accurate, insightful answers based solely on the given context. Follow these guidelines:
-
     1. Analyze the question carefully to understand the core information being sought.
     2. Thoroughly examine the provided context, identifying key relevant information.
     3. Formulate a clear, concise answer that directly addresses the question.
@@ -79,13 +90,10 @@ def create_rag_pipeline(_vector_store):
     7. If asked for an opinion or recommendation, base it strictly on insights from the context.
     8. Use a confident, authoritative tone while maintaining accuracy.
     9. If you cannot provide a clear answer to the question, reply with "I don't know".
-
     Question:
     {question}
-
     Context:
     {context}
-
     Answer:
     """