|
|
| """
|
| AENEA BHR-0.1 Inference Engine
|
| ================================
|
| Production-ready inference wrapper for the BHR decision model.
|
|
|
| Usage:
|
| from bhr_engine import BHREngine
|
|
|
| engine = BHREngine("bhr_best.pt", "tokenizers/QT-Cittern-12k")
|
|
|
| result = engine.decide(
|
| state='text_A: The cat sat on the mat.\\ntext_B: The mat was empty.',
|
| question='Does text_A contradict text_B? Is "contradiction" the correct answer?',
|
| options=["yes", "no"]
|
| )
|
|
|
| print(result.decision) # "yes"
|
| print(result.confidence) # 0.93
|
| print(result.probabilities) # {"yes": 0.93, "no": 0.07}
|
| """
|
|
|
| import time
|
| import torch
|
| import torch.nn.functional as F
|
| from dataclasses import dataclass, field
|
| from typing import List, Optional, Dict, Union
|
| from transformers import AutoTokenizer
|
|
|
| from cartan_olm import CartanConfig, CartanLM
|
|
|
|
|
|
|
|
|
|
|
|
|
| @dataclass
|
| class DecisionResult:
|
| """Result of a single decision."""
|
| decision: str
|
| confidence: float
|
| probabilities: Dict[str, float]
|
| latency_ms: float
|
| kind: str
|
| options: List[str]
|
|
|
| def __repr__(self):
|
| return (f"DecisionResult(decision='{self.decision}', "
|
| f"confidence={self.confidence:.4f}, "
|
| f"latency={self.latency_ms:.1f}ms)")
|
|
|
|
|
| @dataclass
|
| class BatchResult:
|
| """Result of a batch decision."""
|
| results: List[DecisionResult]
|
| total_latency_ms: float
|
|
|
| def __repr__(self):
|
| return (f"BatchResult(n={len(self.results)}, "
|
| f"total_latency={self.total_latency_ms:.1f}ms)")
|
|
|
|
|
|
|
|
|
|
|
|
|
| class BHREngine:
|
| """
|
| Production inference engine for BHR-0.1 decision model.
|
|
|
| Args:
|
| checkpoint_path: Path to the BHR checkpoint (.pt file)
|
| tokenizer_path: Path to the tokenizer directory
|
| device: "cpu" or "cuda"
|
| max_seq_len: Maximum sequence length for state+question
|
| max_option_len: Maximum token length per option
|
| """
|
|
|
| def __init__(
|
| self,
|
| checkpoint_path: str = "bhr_best.pt",
|
| tokenizer_path: str = "tokenizers/QT-Cittern-12k",
|
| device: str = "cpu",
|
| max_seq_len: int = 1536,
|
| max_option_len: int = 64,
|
| ):
|
| self.device = device
|
| self.max_seq_len = max_seq_len
|
| self.max_option_len = max_option_len
|
|
|
|
|
| checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
|
|
|
|
|
| self.config = CartanConfig.from_dict(checkpoint["config"])
|
| self.config.decision_mode = True
|
|
|
|
|
| self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
|
| if self.tokenizer.pad_token is None:
|
| self.tokenizer.pad_token = self.tokenizer.eos_token
|
| self.config.vocab_size = len(self.tokenizer)
|
|
|
|
|
| self.model = CartanLM(self.config).to(device)
|
| state_dict = {
|
| (k[len("_orig_mod."):] if k.startswith("_orig_mod.") else k): v
|
| for k, v in checkpoint["model_state_dict"].items()
|
| }
|
| self.model.load_state_dict(state_dict, strict=False)
|
| self.model.eval()
|
|
|
|
|
| self.checkpoint_metadata = {
|
| "step": checkpoint.get("step"),
|
| "epoch": checkpoint.get("epoch"),
|
| "test_accuracy": checkpoint.get("test_accuracy"),
|
| "test_brier": checkpoint.get("test_brier"),
|
| "test_ece": checkpoint.get("test_ece"),
|
| }
|
|
|
| self.pad_id = self.tokenizer.pad_token_id or 0
|
| self._param_count = sum(p.numel() for p in self.model.parameters())
|
|
|
| print(f"BHR Engine initialized: {self._param_count / 1e6:.1f}M params on {device}")
|
|
|
| @property
|
| def param_count_m(self) -> float:
|
| """Parameter count in millions."""
|
| return self._param_count / 1e6
|
|
|
| def _encode_state_question(self, state: str, question: str) -> torch.Tensor:
|
| """Encode state + question into input_ids tensor."""
|
| state_text = f"State: {state}\nQuestion: {question}"
|
| ids = self.tokenizer.encode(
|
| state_text,
|
| add_special_tokens=False,
|
| truncation=True,
|
| max_length=self.max_seq_len
|
| )
|
| return torch.tensor([ids], dtype=torch.long, device=self.device)
|
|
|
| def _encode_options(self, options: List[str]) -> torch.Tensor:
|
| """Encode and pad options into option_ids tensor."""
|
| token_lists = []
|
| for opt in options:
|
| ids = self.tokenizer.encode(
|
| str(opt),
|
| add_special_tokens=False,
|
| truncation=True,
|
| max_length=self.max_option_len
|
| )
|
| if not ids:
|
| ids = [self.pad_id]
|
| token_lists.append(ids)
|
|
|
| max_len = max(len(t) for t in token_lists)
|
| padded = []
|
| for tokens in token_lists:
|
| padded.append(tokens + [self.pad_id] * (max_len - len(tokens)))
|
|
|
| return torch.tensor([padded], dtype=torch.long, device=self.device)
|
|
|
| def _detect_kind(self, options: List[str]) -> str:
|
| """Detect the decision kind from options."""
|
| if len(options) == 2:
|
| lowered = [o.lower() for o in options]
|
| if set(lowered) == {"yes", "no"}:
|
| return "yesno"
|
| if set(lowered) == {"true", "false"}:
|
| return "yesno"
|
|
|
| try:
|
| [float(o) for o in options]
|
| return "score"
|
| except ValueError:
|
| pass
|
| if any("level" in o.lower() for o in options):
|
| return "score"
|
| return "choice"
|
|
|
| def decide(
|
| self,
|
| state: str,
|
| question: str,
|
| options: List[str],
|
| return_logits: bool = False,
|
| ) -> DecisionResult:
|
| """
|
| Make a single decision.
|
|
|
| Args:
|
| state: The context/state to evaluate
|
| question: The question to answer
|
| options: List of possible answers
|
| return_logits: If True, include raw logits in result
|
|
|
| Returns:
|
| DecisionResult with decision, confidence, and probabilities
|
| """
|
| if not options:
|
| raise ValueError("options must not be empty")
|
| if len(options) == 1:
|
| return DecisionResult(
|
| decision=options[0],
|
| confidence=1.0,
|
| probabilities={options[0]: 1.0},
|
| latency_ms=0.0,
|
| kind=self._detect_kind(options),
|
| options=options,
|
| )
|
|
|
| input_ids = self._encode_state_question(state, question)
|
| option_ids = self._encode_options(options)
|
|
|
| start = time.perf_counter()
|
| with torch.no_grad():
|
| output = self.model(
|
| input_ids,
|
| option_ids=option_ids,
|
| pad_token_id=self.pad_id
|
| )
|
| latency = (time.perf_counter() - start) * 1000
|
|
|
| logits = output["logits"][0].float()
|
| probs = F.softmax(logits, dim=-1).cpu().tolist()
|
| best_idx = int(torch.argmax(logits).item())
|
|
|
| prob_dict = {opt: p for opt, p in zip(options, probs)}
|
|
|
| result = DecisionResult(
|
| decision=options[best_idx],
|
| confidence=probs[best_idx],
|
| probabilities=prob_dict,
|
| latency_ms=latency,
|
| kind=self._detect_kind(options),
|
| options=options,
|
| )
|
|
|
| return result
|
|
|
| def batch_decide(
|
| self,
|
| items: List[Dict[str, Union[str, List[str]]]],
|
| ) -> BatchResult:
|
| """
|
| Make decisions on a batch of items.
|
|
|
| Each item should have:
|
| - "state": str
|
| - "question": str
|
| - "options": List[str]
|
|
|
| Returns:
|
| BatchResult with list of DecisionResults
|
| """
|
| results = []
|
| start = time.perf_counter()
|
|
|
| for item in items:
|
| result = self.decide(
|
| state=item["state"],
|
| question=item["question"],
|
| options=item["options"],
|
| )
|
| results.append(result)
|
|
|
| total_latency = (time.perf_counter() - start) * 1000
|
| return BatchResult(results=results, total_latency_ms=total_latency)
|
|
|
| def yesno(self, state: str, question: str) -> DecisionResult:
|
| """Convenience method for yes/no decisions."""
|
| return self.decide(state, question, ["yes", "no"])
|
|
|
| def choose(self, state: str, question: str, options: List[str]) -> DecisionResult:
|
| """Convenience method for choice decisions."""
|
| return self.decide(state, question, options)
|
|
|
| def score(self, state: str, question: str, options: List[str]) -> DecisionResult:
|
| """Convenience method for score decisions."""
|
| return self.decide(state, question, options)
|
|
|
| def info(self) -> Dict:
|
| """Return engine metadata."""
|
| return {
|
| "model": "AENEA BHR-0.1",
|
| "parameters_m": round(self.param_count_m, 1),
|
| "architecture": "CartanLM (Orthogonal)",
|
| "device": self.device,
|
| "max_seq_len": self.max_seq_len,
|
| "max_option_len": self.max_option_len,
|
| **self.checkpoint_metadata,
|
| }
|
|
|
|
|
|
|
|
|
|
|
|
|
| if __name__ == "__main__":
|
| import argparse
|
|
|
| parser = argparse.ArgumentParser(description="BHR-0.1 Inference Engine")
|
| parser.add_argument("--checkpoint", default="bhr_best.pt")
|
| parser.add_argument("--tokenizer", default="tokenizers/QT-Cittern-12k")
|
| parser.add_argument("--device", default="cpu")
|
| parser.add_argument("--interactive", action="store_true")
|
| args = parser.parse_args()
|
|
|
| engine = BHREngine(args.checkpoint, args.tokenizer, args.device)
|
| print(f"\nEngine info: {engine.info()}\n")
|
|
|
| if args.interactive:
|
| print("BHR Interactive Mode (type 'quit' to exit)")
|
| print("=" * 60)
|
| while True:
|
| state = input("\nState: ").strip()
|
| if state.lower() == "quit":
|
| break
|
| question = input("Question: ").strip()
|
| if question.lower() == "quit":
|
| break
|
| options_str = input("Options (comma-separated): ").strip()
|
| if options_str.lower() == "quit":
|
| break
|
| options = [o.strip() for o in options_str.split(",") if o.strip()]
|
| if not options:
|
| continue
|
|
|
| result = engine.decide(state, question, options)
|
| print(f"\n Decision: {result.decision}")
|
| print(f" Confidence: {result.confidence:.4f}")
|
| print(f" Kind: {result.kind}")
|
| print(f" Latency: {result.latency_ms:.1f}ms")
|
| print(f" Distribution:")
|
| for opt, prob in sorted(result.probabilities.items(), key=lambda x: -x[1]):
|
| bar = "█" * int(prob * 30)
|
| print(f" {prob:6.4f} {bar:<30s} {opt}")
|
| else:
|
|
|
| result = engine.yesno(
|
| state='{"user": {"age": 25}, "rule": "Allow entry if age >= 18."}',
|
| question="Following the rule, is entry allowed?"
|
| )
|
| print(f"Demo result: {result}") |
