File size: 4,031 Bytes
31c93b1 4aaae80 31c93b1 4aaae80 31c93b1 3f78ea8 31c93b1 3f78ea8 31c93b1 8ee8138 31c93b1 8ee8138 31c93b1 4aaae80 31c93b1 8ee8138 31c93b1 4aaae80 31c93b1 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """llama-cpp-python in-process backend."""
from __future__ import annotations
from hearthnet.services.llm.backends.base import BackendModel, ChatResult, Token
from hearthnet.services.llm.tokenizers import model_family
def _family(model_name: str) -> str:
return model_family(model_name)
class LlamaCppBackend:
name = "llama_cpp"
def __init__(self, model_path: str, n_ctx: int = 4096, n_gpu_layers: int = -1) -> None:
self._model_path = model_path
self._n_ctx = n_ctx
self._n_gpu_layers = n_gpu_layers
self._llm = None
model_name = model_path.split("/")[-1].split(".")[0]
self.models = [
BackendModel(
name=model_name,
family=_family(model_name),
context_length=n_ctx,
requires_internet=False,
)
]
def is_available(self) -> bool:
try:
from importlib.util import find_spec
from pathlib import Path
return Path(self._model_path).exists() and find_spec("llama_cpp") is not None
except ImportError:
return False
async def warm(self) -> None:
if not self.is_available():
return
import asyncio
loop = asyncio.get_running_loop()
await loop.run_in_executor(None, self._load_model)
def _load_model(self) -> None:
from llama_cpp import Llama
self._llm = Llama(
model_path=self._model_path,
n_ctx=self._n_ctx,
n_gpu_layers=self._n_gpu_layers,
verbose=False,
)
async def chat(
self,
messages: list[dict],
*,
model: str = "",
stream: bool = False,
temperature: float = 0.7,
max_tokens: int = 1024,
**kwargs,
):
import asyncio
import time
if self._llm is None:
await self.warm()
if self._llm is None:
raise RuntimeError("llama.cpp model not loaded")
t0 = time.monotonic()
loop = asyncio.get_running_loop()
if not stream:
result = await loop.run_in_executor(
None,
lambda: self._llm.create_chat_completion(
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
),
)
text = result["choices"][0]["message"]["content"]
ms = int((time.monotonic() - t0) * 1000)
return ChatResult(
text=text,
tokens_in=result["usage"]["prompt_tokens"],
tokens_out=result["usage"]["completion_tokens"],
model=self.models[0].name,
ms=ms,
)
return self._stream_chat(messages, temperature, max_tokens)
async def _stream_chat(self, messages, temperature, max_tokens):
import asyncio
loop = asyncio.get_running_loop()
result = await loop.run_in_executor(
None,
lambda: self._llm.create_chat_completion(
messages=messages,
temperature=temperature,
max_tokens=max_tokens,
stream=True,
),
)
for chunk in result:
delta = chunk["choices"][0].get("delta", {})
text = delta.get("content", "")
done = chunk["choices"][0]["finish_reason"] is not None
if text or done:
yield Token(text=text, stop=done)
async def complete(self, prompt: str, *, model: str = "", stream: bool = False, **kwargs):
messages = [{"role": "user", "content": prompt}]
return await self.chat(messages, model=model, stream=stream, **kwargs)
async def close(self) -> None:
self._llm = None
def health(self) -> dict:
return {
"backend": "llama_cpp",
"model_path": self._model_path,
"loaded": self._llm is not None,
}
|