cache model with respecting zerogpu cuda usage

app.py

@@ -59,9 +59,20 @@ def load_model():
     generator = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, attn_implementation="eager", use_flash_attention_2=False) # True for flash-attn2 else False
     generator_mini = pipeline(task="text-generation", model=ankerbot_model, tokenizer=ankerbot_tokenizer, torch_dtype=torch.float16, attn_implementation="eager", use_flash_attention_2=False) # True for flash-attn2 else False
     return (generator, generator_mini)
-generator = load_model()
+    
+model_cache = None
+
+@spaces.GPU
+def get_model():
+    global _model_cache
+    if _model_cache is None:
+        # Load model only if it's not already loaded
+        print("Loading model for the first time...")
+        _model_cache = load_model()
+    return _model_cache
+
 @spaces.GPU
-def generate_response(query, context, prompts, max_tokens, temperature, top_p):
+def generate_response(query, context, prompts, max_tokens, temperature, top_p, generator):
     system_message_support = f"""<|im_start|>system
         Rolle: Du bist der KI-Assistent für Kundenservice, der im Namen des Unternehmens und Gewürzmanufaktur Ankerkraut handelt und Antworten aus der Ich-Perspektive, basierend auf den bereitgestellten Informationen gibt.
         Oberstes Ziel: Beantworte die folgende Frage präzise, indem du den Kontext zusammenfasst.
@@ -119,7 +130,7 @@ def generate_response(query, context, prompts, max_tokens, temperature, top_p):
 
     print("Prompt: ", system_message)
 
-    response = generator[0](system_message, do_sample=True, top_p=top_p, max_new_tokens=max_tokens, temperature=temperature)[0]["generated_text"]
+    response = generator(system_message, do_sample=True, top_p=top_p, max_new_tokens=max_tokens, temperature=temperature)[0]["generated_text"]
 
     # Extract only the assistant's response
     if "assistant" in response:
@@ -153,7 +164,7 @@ def respond(
     temperature,
     top_p,
 ):  
-
+    generator = get_model()
     if system_message == None: 
         system_message = f"""<|im_start|>system Rolle: Du bist ein KI-Assistent der die vom Kunden formuliert Frage in Stichworte verwandelt die für eine Vektorsuche verwendet werden.
             Oberstes Ziel: Suche Schlüsselbegriffe aus der Frage heraus und gebe diese als Hauptbegriff aus. Suche zusätzlich ähnliche Begriffe aus.
@@ -173,7 +184,7 @@ def respond(
         colleciton_name = "products"
 
     context = search_qdrant_with_context(query + " " + refined_context[0]["generated_text"].split("assistant\n").pop(), collection_name)
-    answer = generate_response(query, context, max_tokens, temperature, top_p)
+    answer = generate_response(query, context, max_tokens, temperature, top_p, generator[0])
     full_conv = f"Nutzer:{query};Assistent:{answer}"
     if len(last_messages) > 5:
         last_messages.pop(0)