Update app.py

app.py

@@ -1,14 +1,14 @@
 import os
 import streamlit as st
-from huggingface_hub import login
+from huggingface_hub import login, hf_hub_download
 from datasets import load_dataset
 from llama_cpp import Llama
-from huggingface_hub import hf_hub_download
 import chromadb
+from chromadb.config import Settings  # Added import for Settings
 from sentence_transformers import SentenceTransformer
 
 # Load Hugging Face token from environment variable
-hf_token = os.getenv("HF_TOKEN")  
+hf_token = os.getenv("HF_TOKEN")
 if hf_token:
     login(token=hf_token)
 else:
@@ -22,6 +22,7 @@ llm = Llama(
     model_path=hf_hub_download(
         repo_id="TheBloke/CapybaraHermes-2.5-Mistral-7B-GGUF",
         filename="capybarahermes-2.5-mistral-7b.Q2_K.gguf",
+        token=hf_token  # Ensure the token is passed for authentication
     ),
     n_ctx=2048,
 )
@@ -30,9 +31,13 @@ llm = Llama(
 class VectorStore:
     def __init__(self, collection_name):
         self.embedding_model = SentenceTransformer('sentence-transformers/multi-qa-MiniLM-L6-cos-v1')
-        self.chroma_client = chromadb.PersistentClient(path="./chroma_db")  # Ensure persistence
+        # Use Settings to configure persistence
+        self.chroma_client = chromadb.Client(Settings(
+            chroma_db_impl="duckdb+parquet",
+            persist_directory="./chroma_db"  # Ensure persistence
+        ))
         if collection_name in [c.name for c in self.chroma_client.list_collections()]:
-            self.chroma_client.delete_collection(collection_name)
+            self.chroma_client.delete_collection(name=collection_name)
         self.collection = self.chroma_client.create_collection(name=collection_name)
 
     def populate_vectors(self, dataset):
@@ -60,7 +65,7 @@ class VectorStore:
     def search_context(self, query, n_results=1):
         query_embedding = self.embedding_model.encode([query]).tolist()
         results = self.collection.query(query_embeddings=query_embedding, n_results=n_results)
-        return results['documents']
+        return results['documents'][0]  # Adjusted to access the correct document
 
 # Initialize and populate vector store
 vector_store = VectorStore("embedding_vector")