Update utils.py

utils.py

@@ -397,6 +397,7 @@ def document_storage_chroma(splits):
 
 ########################################################
 #Vektorstore speichern - bzw. laden
+"""
 def save_splits_and_metadata(splits, directory="chroma/kkg", filename="splits_and_metadata.pkl"):
     # Erstellen des Verzeichnisses, falls es nicht existiert
     if not os.path.exists(directory):
@@ -427,6 +428,66 @@ def load_vectorstore():
         PREPROCESSED_SPLITS, SPLIT_TO_ORIGINAL_MAPPING = splits_and_metadata
         return document_storage_chroma(PREPROCESSED_SPLITS)
     return None
+"""
+
+def save_split_to_original_mapping(mapping, directory="chroma/kkg", mapping_filename="mapping.pkl"):
+    # Erstellen des Verzeichnisses, falls es nicht existiert
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+    
+    # Speichern des Mappings
+    mapping_filepath = os.path.join(directory, mapping_filename)
+    with open(mapping_filepath, "wb") as f:
+        pickle.dump(mapping, f)
+    
+    # Hochladen der Mapping-Datei zum Hugging Face Space
+    upload_file_to_huggingface(mapping_filepath, f"{directory}/{mapping_filename}")
+
+def load_split_to_original_mapping(directory="chroma/kkg", mapping_filename="mapping.pkl"):
+    # Laden des Mappings
+    mapping_filepath = os.path.join(directory, mapping_filename)
+    
+    if os.path.exists(mapping_filepath):
+        with open(mapping_filepath, "rb") as f:
+            return pickle.load(f)
+    return None
+
+def save_vectorstore(vectorstore, directory="chroma/kkg", splits_filename="splits_and_metadata.pkl", vectorstore_filename="vectorstore.pkl"):
+    # Erstellen des Verzeichnisses, falls es nicht existiert
+    if not os.path.exists(directory):
+        os.makedirs(directory)
+    
+    # Speichern der Splits und Metadaten
+    splits_filepath = os.path.join(directory, splits_filename)
+    with open(splits_filepath, "wb") as f:
+        pickle.dump((vectorstore.documents, SPLIT_TO_ORIGINAL_MAPPING), f)
+    
+    # Speichern des Vektorstores (ohne die SQLite-Verbindung)
+    vectorstore_filepath = os.path.join(directory, vectorstore_filename)
+    with open(vectorstore_filepath, "wb") as f:
+        pickle.dump(vectorstore, f)
+
+    # Hochladen der gespeicherten Dateien zum Hugging Face Space
+    upload_file_to_huggingface(splits_filepath, f"{directory}/{splits_filename}")
+    upload_file_to_huggingface(vectorstore_filepath, f"{directory}/{vectorstore_filename}")
+
+def load_vectorstore(directory="chroma/kkg", splits_filename="splits_and_metadata.pkl", vectorstore_filename="vectorstore.pkl"):
+    # Laden der Splits und Metadaten
+    splits_filepath = os.path.join(directory, splits_filename)
+    vectorstore_filepath = os.path.join(directory, vectorstore_filename)
+    
+    if os.path.exists(splits_filepath) and os.path.exists(vectorstore_filepath):
+        with open(splits_filepath, "rb") as f:
+            documents, SPLIT_TO_ORIGINAL_MAPPING = pickle.load(f)
+        
+        with open(vectorstore_filepath, "rb") as f:
+            vectorstore = pickle.load(f)
+            vectorstore.documents = documents  # Setze die Dokumente im Vektorstore
+            return vectorstore, SPLIT_TO_ORIGINAL_MAPPING
+    return None, None
+
+
+
 
 # Beispiel-Upload-Funktion
 def upload_file_to_huggingface(file_path, upload_path):