hackeracademy's picture
Fix Gradio signature & set MPLCONFIGDIR
722f6b0
import os, gradio as gr, requests, tempfile, logging, time
from llama_cpp import Llama
MODEL_URL = (
"https://huggingface.co/fdtn-ai/Foundation-Sec-8B-Q4_K_M-GGUF/"
"resolve/main/foundation-sec-8b-q4_k_m.gguf"
)
CACHE_DIR = "/tmp"
MODEL_PATH = os.path.join(CACHE_DIR, "foundation-sec-8b-q4_k_m.gguf")
# silence matplotlib cache warning
os.environ["MPLCONFIGDIR"] = CACHE_DIR
# download once
if not os.path.exists(MODEL_PATH):
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)
logging.info("Download finished.")
llm = Llama(model_path=MODEL_PATH, n_ctx=4096, verbose=False)
# correct signature: message, history
def chat_fn(message, history):
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")
demo.launch(server_name="0.0.0.0", server_port=7860)