Update app.py

app.py

@@ -1,8 +1,8 @@
 import gradio as gr
-
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-@gr.cache(allow_output_mutation=True)
+model_cache = {}
+
 def load_model():
     model_id = "Tech-Meld/Hajax_Chat_1.0"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
@@ -15,17 +15,20 @@ def get_response(input_text, model, tokenizer):
     response = tokenizer.decode(outputs[:, inputs.shape[-1]:][0], skip_special_tokens=True)
     return response
 
-model, tokenizer = load_model()
-
 def chat(input_text):
-  response = get_response(input_text, model, tokenizer)
-  return response
+    global model_cache
+    if "model" not in model_cache:
+        model_cache["model"], model_cache["tokenizer"] = load_model()
+    model = model_cache["model"]
+    tokenizer = model_cache["tokenizer"]
+    response = get_response(input_text, model, tokenizer)
+    return response
 
 iface = gr.Interface(
     chat,
     "text",
     "text",
-    title="Chat with Hajax",
+    title="Chat with AI",
     description="Type your message and press Enter to chat with the AI.",
 )
 iface.launch()