oracle / llm.py
vivek gangadharan
newversion
f0d3c87
Raw
History Blame Contribute Delete
6.79 kB
"""Local Llama via the llama.cpp runtime (earns the 'Llama Champion' badge).
Two ways to reach a model, tried in order; both are optional and the callers
always fall back to built-in phrasing if neither is available, so the game
never breaks:
1. In-process llama.cpp — `llama-cpp-python` loads the GGUF directly. This is
what runs on the Space: fully local, no network at play time (also satisfies
the 'Off the Grid' badge). The model is downloaded once from the Hub.
2. HTTP llama-server — an OpenAI-compatible endpoint (handy in local dev:
`llama-server -m model.gguf --port 8080`).
Nothing here imports llama_cpp at module load, so tests stay offline and fast.
"""
from __future__ import annotations
import json
import os
import threading
import urllib.request
# Built with Llama 🦙 — a 3B model fits the "build small" spirit (≤32B).
REPO = os.environ.get("ORACLE_LLAMA_REPO", "bartowski/Llama-3.2-3B-Instruct-GGUF")
FILE = os.environ.get("ORACLE_LLAMA_FILE", "Llama-3.2-3B-Instruct-Q4_K_M.gguf")
HTTP_URL = os.environ.get("ORACLE_LLAMA_URL", "http://localhost:8080/v1/chat/completions")
MODEL_NAME = os.environ.get("ORACLE_LLAMA_MODEL", "Llama-3.2-3B-Instruct")
def _default_threads() -> int:
"""Cap threads. On HF Spaces os.cpu_count() reports the whole host (16-32),
but the container only has ~2 cores; oversubscribing makes llama.cpp crawl."""
try:
n = len(os.sched_getaffinity(0)) # cores actually available to us
except (AttributeError, OSError):
n = os.cpu_count() or 2
return max(1, min(n, 4))
THREADS = int(os.environ.get("ORACLE_LLAMA_THREADS", str(_default_threads())))
_llm = None # cached llama_cpp.Llama instance
_inproc_failed = False
_load_lock = threading.Lock()
_active_mode = None # "in-process" | "http" — logged once so it's easy to debug
def status() -> str:
"""Human-readable description of how the model is (or will be) reached."""
if _llm is not None:
return f"IN-PROCESS llama.cpp · {FILE} · threads={THREADS}"
if _inproc_failed:
return f"in-process unavailable -> HTTP llama-server @ {HTTP_URL} (else built-in phrasing)"
return "not loaded yet"
def _announce(mode: str) -> None:
"""Print the active backend the first time it's actually used."""
global _active_mode
if _active_mode != mode:
_active_mode = mode
if mode == "in-process":
print(f"[llm] 🟢 MODE = IN-PROCESS llama.cpp ({FILE}, threads={THREADS})", flush=True)
else:
print(f"[llm] 🟡 MODE = HTTP llama-server @ {HTTP_URL}", flush=True)
def _model_cache_dir():
"""Where to download/keep the GGUF. Prefer a persistent location (the bucket)
so the ~2 GB model isn't re-downloaded on every cold start."""
explicit = os.environ.get("ORACLE_MODEL_DIR")
if explicit:
return explicit
data_dir = os.environ.get("ORACLE_DATA_DIR")
if data_dir: # e.g. a mounted bucket at /data
return os.path.join(data_dir, "models")
return None # fall back to the default HF cache
def _load_inproc():
"""Lazily download + load the GGUF through llama.cpp. None if unavailable."""
global _llm, _inproc_failed
if _llm is not None or _inproc_failed:
return _llm
with _load_lock: # only one thread loads the model
if _llm is not None or _inproc_failed:
return _llm
try:
from llama_cpp import Llama
from huggingface_hub import hf_hub_download
cache_dir = _model_cache_dir()
if cache_dir:
os.makedirs(cache_dir, exist_ok=True)
path = hf_hub_download(repo_id=REPO, filename=FILE, cache_dir=cache_dir)
_llm = Llama(model_path=path, n_ctx=2048,
n_threads=THREADS, verbose=False)
print(f"[llm] llama.cpp loaded {FILE} "
f"(threads={THREADS}, cache={cache_dir or 'default'})", flush=True)
except Exception as exc: # noqa: BLE001 — fall back to HTTP / built-in phrasing
print(f"[llm] in-process llama.cpp unavailable: {exc}")
_inproc_failed = True
return _llm
def warmup() -> None:
"""Load the model (and run one tiny generation) ahead of time, so the first
real question isn't blocked by the cold download+load. Call at startup."""
llm = _load_inproc()
if llm is not None:
try:
llm.create_chat_completion(
messages=[{"role": "user", "content": "hi"}], max_tokens=1)
print(f"[llm] warmup complete — {status()}", flush=True)
except Exception as exc: # noqa: BLE001
print(f"[llm] warmup skipped: {exc}")
else:
print(f"[llm] no in-process model — {status()}", flush=True)
def _chat_http(messages: list, temperature: float, max_tokens: int) -> str:
payload = {"model": MODEL_NAME, "messages": messages,
"temperature": temperature, "max_tokens": max_tokens}
req = urllib.request.Request(
HTTP_URL, data=json.dumps(payload).encode("utf-8"),
headers={"Content-Type": "application/json"}, method="POST")
with urllib.request.urlopen(req, timeout=90) as resp:
body = json.loads(resp.read().decode("utf-8"))
return body["choices"][0]["message"]["content"]
CHAT_TIMEOUT = float(os.environ.get("ORACLE_LLM_TIMEOUT", "25"))
def _chat_once(messages: list, temperature: float, max_tokens: int) -> str:
llm = _load_inproc()
if llm is not None:
_announce("in-process")
out = llm.create_chat_completion(
messages=messages, temperature=temperature, max_tokens=max_tokens)
return out["choices"][0]["message"]["content"]
_announce("http")
return _chat_http(messages, temperature, max_tokens)
def chat(messages: list, temperature: float = 0.4, max_tokens: int = 120,
timeout: float | None = None) -> str:
"""Run a chat completion through llama.cpp (in-process first, then HTTP).
Bounded by a timeout so a slow CPU generation can never hang the game — on
timeout it raises and the caller falls back to built-in phrasing. Raises if
neither backend is reachable, too.
"""
timeout = CHAT_TIMEOUT if timeout is None else timeout
box: dict = {}
def run():
try:
box["out"] = _chat_once(messages, temperature, max_tokens)
except Exception as exc: # noqa: BLE001
box["err"] = exc
t = threading.Thread(target=run, daemon=True)
t.start()
t.join(timeout)
if t.is_alive():
raise TimeoutError(f"llm.chat exceeded {timeout}s")
if "err" in box:
raise box["err"]
return box["out"]