mirror of
https://github.com/mcloh/decision-models-benchmark.git
synced 2026-11-05 22:04:12 +00:00
60 lines
2.5 KiB
Python
60 lines
2.5 KiB
Python
"""Adaptador do Rizzo Flow 4B (GGUF Q8_0 sobre llama.cpp) com prompts spark-decisions-v3 (tarefa 45)."""
|
|
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
from dmb.adapter import Adapter
|
|
from dmb.jev import ContextLimit, answer_from_distribution, option_keys
|
|
|
|
GGUF_FILE = "spark-x2.5-4b-rizzo-flow-lora-q8_0.gguf"
|
|
|
|
|
|
class RizzoFlowAdapter(Adapter):
|
|
name = "rizzo_flow"
|
|
probabilities_source = "native"
|
|
|
|
def __init__(self, *args, runtime_dir: str | Path | None = None, **kwargs):
|
|
super().__init__(*args, **kwargs)
|
|
self.runtime_dir = runtime_dir
|
|
|
|
def load(self) -> None:
|
|
from rizzo_flow.backend_llama import LlamaBackend
|
|
from rizzo_flow.compat import model_name
|
|
from rizzo_flow.engine import Engine
|
|
|
|
device = "cuda" if self.device == "cuda" else ("cpu" if self.device == "cpu" else "auto")
|
|
backend = LlamaBackend.load(self.model_dir / GGUF_FILE, device=device, ctx=self.context_limit,
|
|
runtime_dir=self.runtime_dir)
|
|
self.engine = Engine(backend, ctx=self.context_limit)
|
|
self.served = model_name(backend.metadata)
|
|
|
|
def native_context_limit(self) -> int:
|
|
return 2048 # comprimento máximo dos prompts de treino do LoRA
|
|
|
|
def count_tokens(self, text: str) -> int:
|
|
return len(self.engine.backend.tokenizer.encode(text))
|
|
|
|
def decide(self, state: str, questions: dict) -> tuple[dict, int]:
|
|
from rizzo_flow.compat import SystemOneRequest, from_native, to_native
|
|
|
|
request = SystemOneRequest.model_validate(
|
|
{"model": self.served, "state": state, "questions": questions})
|
|
native, options = to_native(request)
|
|
try:
|
|
raw = from_native(request, self.engine.decide(native), options, self.served)
|
|
except ValueError as error:
|
|
if "context" in str(error).lower() or "tokens" in str(error).lower():
|
|
raise ContextLimit(str(error)) from error
|
|
raise
|
|
raw = raw if isinstance(raw, dict) else raw.model_dump()
|
|
answers = {}
|
|
for name, question in questions.items():
|
|
a = raw["answers"][name]
|
|
if question["type"] == "noul":
|
|
p = float(a["noul"])
|
|
dist = {"true": p, "false": 1 - p}
|
|
else:
|
|
dist = {k: a["probabilities"][k] for k in option_keys(question)}
|
|
answers[name] = answer_from_distribution(question, dist, a.get("confidence"))
|
|
return answers, int(raw.get("usage", {}).get("input_tokens", 0))
|