Fix Gradio signature & set MPLCONFIGDIR

app.py

@@ -1,4 +1,4 @@
-import os, gradio as gr, requests, tempfile
+import os, gradio as gr, requests, tempfile, logging, time
 from llama_cpp import Llama
 
 MODEL_URL = (
@@ -6,29 +6,39 @@ MODEL_URL = (
     "resolve/main/foundation-sec-8b-q4_k_m.gguf"
 )
 
-# writable directory
 CACHE_DIR = "/tmp"
 MODEL_PATH = os.path.join(CACHE_DIR, "foundation-sec-8b-q4_k_m.gguf")
 
-# download only once
+# silence matplotlib cache warning
+os.environ["MPLCONFIGDIR"] = CACHE_DIR
+
+# download once
 if not os.path.exists(MODEL_PATH):
-    print("Downloading model … (~4.9 GB)")
+    logging.info("Downloading model …")
     with requests.get(MODEL_URL, stream=True) as r:
         r.raise_for_status()
         with open(MODEL_PATH, "wb") as f:
             for chunk in r.iter_content(chunk_size=8192):
                 f.write(chunk)
-    print("Download finished.")
+    logging.info("Download finished.")
 
-# load model
 llm = Llama(model_path=MODEL_PATH, n_ctx=4096, verbose=False)
 
+# correct signature: message, history
 def chat_fn(message, history):
-    messages = [{"role": "user", "content": message}]
-    out = llm.create_chat_completion(messages=messages, max_tokens=256, temperature=0.7)
+    messages = []
+    for human, ai in history:
+        messages.append({"role": "user", "content": human})
+        messages.append({"role": "assistant", "content": ai})
+    messages.append({"role": "user", "content": message})
+
+    out = llm.create_chat_completion(
+        messages=messages,
+        max_tokens=512,
+        temperature=0.7,
+        stream=False,
+    )
     return out["choices"][0]["message"]["content"]
 
 demo = gr.ChatInterface(chat_fn, title="Foundation-Sec-8B")
-
-# expose on 0.0.0.0:7860 (Gradio default)
 demo.launch(server_name="0.0.0.0", server_port=7860)