Update app.py

app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import time
 import re
 import io
+import pickle
 #from PIL import Image, ImageDraw, ImageOps, ImageFont
 #import base64
 import tempfile
@@ -32,6 +33,7 @@ from transformers import pipeline
 from huggingface_hub import InferenceApi
 from utils import *
 from beschreibungen import *
+
  
 
 #Konstanten
@@ -94,11 +96,7 @@ DOCS_DIR = "chroma/kkg"
 ###############################################
 #globale Variablen
 ##############################################
-#nur bei ersten Anfrage splitten der Dokumente - um die Vektordatenbank entsprechend zu füllen
-#splittet = False
-#DB für Vektorstore
-vektordatenbank = None
-retriever = None
+
 
 #############################################
 # Allgemeine Konstanten
@@ -213,7 +211,7 @@ def reset_textbox():
 
 ####################################################
 #aus einem Text-Prompt die Antwort von KI bekommen
-def generate_text (prompt, chatbot, history, vektordatenbank, retriever, top_p=0.6, temperature=0.2, max_new_tokens=4048, max_context_length_tokens=2048, repetition_penalty=1.3, top_k=35):
+def generate_text (prompt, chatbot, history, retriever, top_p=0.6, temperature=0.2, max_new_tokens=4048, max_context_length_tokens=2048, repetition_penalty=1.3, top_k=35):
     if (prompt == ""):
         raise gr.Error("Prompt ist erforderlich.")
 
@@ -271,13 +269,21 @@ def generate_text (prompt, chatbot, history, vektordatenbank, retriever, top_p=0
         raise gr.Error(e)
     return result, False
 
+########################################################
+#Vektorstore speichern - bzw. laden
+def save_vectorstore(vectorstore, filename="vectorstore.pkl"):
+    with open(filename, "wb") as f:
+        pickle.dump(vectorstore, f)
 
+def load_vectorstore(filename="vectorstore.pkl"):
+    if os.path.exists(filename):
+        with open(filename, "rb") as f:
+            return pickle.load(f)
+    return None
 
 ##############################################################	
 #Eingaben der GUI verarbeiten	
 def generate_auswahl(prompt_in, file, file_history, chatbot, history, anzahl_docs=4, top_p=0.6, temperature=0.5, max_new_tokens=4048, max_context_length_tokens=2048, repetition_penalty=1.3,top_k=5,  validate=False):
-    global vektordatenbank, retriever
-
     #nur wenn man sich validiert hat, kann die Anwendung los legen
     if (validate and not prompt_in == "" and not prompt_in == None):
         # Vektorstore initialisieren
@@ -287,17 +293,25 @@ def generate_auswahl(prompt_in, file, file_history, chatbot, history, anzahl_doc
         #prompt normalisieren bevor er an die KIs geht
         prompt = preprocess_text(prompt_in)
         
-        #muss nur einmal ausgeführt werden... 
-		#?????????????????????????????????????????????? Nicht passend zum Promt???????????????????????????
-        if not vektordatenbank or not SPLIT_TO_ORIGINAL_MAPPING:
+        # Versuchen, den Vektorstore zu laden - nur wenn Dateien neu hochgeladen werden, werden die Splits neu erstellt
+        vektordatenbank = load_vectorstore()
+
+        if vektordatenbank is None:
             print("db neu aufbauen!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!1")
+            #Splits zu allen Dokumenten in den Verzeichnissen erstellen
             PREPROCESSED_SPLITS, SPLIT_TO_ORIGINAL_MAPPING = document_loading_splitting()
             if PREPROCESSED_SPLITS:
-                vektordatenbank, retriever = document_storage_chroma(PREPROCESSED_SPLITS)
+                #Vektordatenbank zu den Splits erstellen
+                vektordatenbank = document_storage_chroma(PREPROCESSED_SPLITS)
+                # Speichern des Vektorstores
+                save_vectorstore(vektordatenbank)
+
+        #Retriever erstellen, um die relevanten Slpits zu einem Prompt zu suchen.... (retrieven)
+        retriever = vektordatenbank.as_retriever(search_kwargs = {"k": ANZAHL_DOCS})        
 
         #kein Bild hochgeladen -> auf Text antworten...
         status = "Antwort der Vektordatenbank"
-        results, status = generate_text(prompt, chatbot, history,vektordatenbank, retriever, top_p=0.6, temperature=0.5, max_new_tokens=4048, max_context_length_tokens=2048, repetition_penalty=1.3, top_k=3)
+        results, status = generate_text(prompt, chatbot, history, retriever, top_p=0.6, temperature=0.5, max_new_tokens=4048, max_context_length_tokens=2048, repetition_penalty=1.3, top_k=3)
 
         #in results sind die preprocessed Splits enthalten, dargestellt werden sollen die orginalen:
         relevant_docs_org=[]