Update app.py

app.py

@@ -10,7 +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
+from langchain_community.embeddings import HuggingFaceEmbeddings
 
 # Set up API keys
 os.environ["OPENAI_API_KEY"] = st.secrets["OPENAI_API_KEY"]
@@ -21,16 +21,6 @@ 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 = []
@@ -51,22 +41,26 @@ def load_and_process_pdfs(pdf_links):
 def setup_vectorstore():
     LOCATION = ":memory:"
     COLLECTION_NAME = "AI_Ethics_Framework"
-    VECTOR_SIZE = 768  # Adjust this if your model's output size is different
-
+    
     qdrant_client = QdrantClient(location=LOCATION)
 
+    # Create the embeddings
+    embeddings = HuggingFaceEmbeddings(model_name="Technocoloredgeek/midterm-finetuned-embedding")
+    
+    # Get the vector size from the embeddings
+    VECTOR_SIZE = len(embeddings.embed_query("test"))
+
     # Create the collection
     qdrant_client.create_collection(
         collection_name=COLLECTION_NAME,
         vectors_config=VectorParams(size=VECTOR_SIZE, distance=Distance.COSINE),
     )
 
-    # Create the vector store with your custom embeddings
-    custom_embeddings = CustomHuggingFaceEmbeddings("Technocoloredgeek/midterm-finetuned-embedding")
+    # Create the vector store
     qdrant_vector_store = QdrantVectorStore(
         client=qdrant_client,
         collection_name=COLLECTION_NAME,
-        embedding=custom_embeddings
+        embeddings=embeddings
     )
 
     # Load and add documents