Update app.py

app.py

@@ -5,18 +5,64 @@ from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 from transformers import pipeline
 from langchain_community.llms import HuggingFacePipeline
 
-# Initialize ChromaDB client
-chroma_client = chromadb.PersistentClient(path="data_db")
+# # Initialize ChromaDB client
+# chroma_client = chromadb.PersistentClient(path="data_db")
 
-# Define the embedding function
-sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-mpnet-base-v2")
+# # Define the embedding function
+# sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-mpnet-base-v2")
 
-# Get or create a collection
-collection = chroma_client.get_or_create_collection(name="my_collection", embedding_function=sentence_transformer_ef)
+# # Get or create a collection
+# collection = chroma_client.get_or_create_collection(name="my_collection", embedding_function=sentence_transformer_ef)
 
 # Streamlit UI elements
 st.title("ChromaDB and HuggingFace Pipeline Integration")
 query = st.text_input("Enter your query:", value="director")
+import csv
+import chromadb
+from chromadb.utils import embedding_functions
+
+
+
+with open('./output.csv' ,  encoding="utf-8") as file:
+    lines = csv.reader(file)
+
+    
+    documents = []
+
+   
+    metadatas = []
+
+    
+    ids = []
+    id = 1
+
+    
+    for i, line in enumerate(lines):
+        if i == 0:
+            
+            continue
+
+        documents.append(line[0])
+        metadatas.append({"item_id": line[1]})
+        ids.append(str(id))
+        id += 1
+
+
+chroma_client = chromadb.PersistentClient(path="db")
+
+
+sentence_transformer_ef = embedding_functions.SentenceTransformerEmbeddingFunction(model_name="all-mpnet-base-v2")
+
+
+collection = chroma_client.get_or_create_collection(name="my_collection", embedding_function=sentence_transformer_ef)
+
+
+collection.add(
+    documents=documents,
+    metadatas=metadatas,
+    ids=ids
+)
+
 
 if st.button("Search"):
     # Query the collection