75 lines
2.4 KiB
Python
75 lines
2.4 KiB
Python
from __future__ import annotations
|
|
|
|
import os
|
|
import time
|
|
from pathlib import Path
|
|
|
|
try:
|
|
from llama_cpp import Llama
|
|
except ImportError as exc:
|
|
raise RuntimeError(
|
|
"llama-cpp-python with CUDA support is required for the CUDA engine runtime."
|
|
) from exc
|
|
|
|
from engine_loader import BaseRuntime
|
|
|
|
|
|
class EngineRuntime(BaseRuntime):
|
|
name = "llama_cpp_cuda"
|
|
|
|
def __init__(self, *, context_window: int = 4096, gpu_layers: int = -1) -> None:
|
|
super().__init__()
|
|
self.context_window = context_window
|
|
self.gpu_layers = gpu_layers
|
|
self.threads = max(os.cpu_count() or 1, 1)
|
|
self._llm: Llama | None = None
|
|
self._model_path: Path | None = None
|
|
|
|
def load_model(self, model_path: Path) -> None:
|
|
if self._model_path == model_path and self._llm is not None:
|
|
return
|
|
|
|
self.unload()
|
|
self._llm = Llama(
|
|
model_path=str(model_path),
|
|
n_ctx=self.context_window,
|
|
n_threads=self.threads,
|
|
n_gpu_layers=self.gpu_layers,
|
|
use_gpu=True,
|
|
verbose=False,
|
|
)
|
|
self._model_path = model_path
|
|
|
|
def generate(self, prompt: str, *, temperature: float) -> dict:
|
|
if self._llm is None:
|
|
raise RuntimeError("Model must be loaded before calling generate().")
|
|
|
|
start_time = time.time()
|
|
response = self._llm.create_chat_completion(
|
|
messages=[{"role": "user", "content": prompt}],
|
|
temperature=temperature,
|
|
)
|
|
elapsed = max(time.time() - start_time, 1e-5)
|
|
|
|
content = response["choices"][0]["message"]["content"]
|
|
usage = response.get("usage", {})
|
|
completion_tokens = usage.get("completion_tokens", 0)
|
|
total_tokens = usage.get("total_tokens", completion_tokens)
|
|
|
|
return {
|
|
"response": content,
|
|
"metrics": {
|
|
"tokens_per_second": round(completion_tokens / elapsed, 2) if completion_tokens else 0,
|
|
"total_tokens": total_tokens,
|
|
"time_to_first_token": round(usage.get("prompt_eval_duration", 0) / 1_000_000_000, 2)
|
|
if "prompt_eval_duration" in usage
|
|
else 0,
|
|
"stop_reason": response["choices"][0].get("finish_reason", "unknown"),
|
|
},
|
|
}
|
|
|
|
def unload(self) -> None:
|
|
if self._llm is not None:
|
|
self._llm = None
|
|
self._model_path = None
|