nexora/inference.py

"""Interchangeable local HF and explicit HTTP chat backends."""
from pathlib import Path
import json
import os
import time
import urllib.request
import urllib.parse
import threading


class HTTPBackend:
    def __init__(self, url, model, *, allow_network=False, timeout=60, api_key_env="NEXORA_API_KEY"):
        parsed = urllib.parse.urlparse(url)
        if parsed.scheme not in {"http", "https"} or parsed.username or parsed.password:
            raise ValueError("Expected HTTP(S) URL without embedded credentials")
        if parsed.hostname not in {"localhost", "127.0.0.1", "::1"} and not allow_network:
            raise PermissionError("Remote inference requires explicit allow_network")
        self.url, self.model, self.timeout, self.api_key_env = url.rstrip("/"), model, timeout, api_key_env

    def complete(self, messages, schema=None):
        body = {"model": self.model, "messages": messages, "temperature": 0, "max_tokens": 512}
        if schema:
            body["response_format"] = {"type": "json_schema", "json_schema": {"name": "nexora_action", "schema": schema}}
        headers = {"Content-Type": "application/json"}
        key = os.environ.get(self.api_key_env)
        if key:
            headers["Authorization"] = "Bearer " + key
        request = urllib.request.Request(self.url + "/chat/completions", data=json.dumps(body).encode(), headers=headers)
        # Redirects are disabled so an approved local URL cannot redirect prompts elsewhere.
        class NoRedirect(urllib.request.HTTPRedirectHandler):
            def redirect_request(self, *args, **kwargs):
                return None
        opener = urllib.request.build_opener(urllib.request.ProxyHandler({}), NoRedirect())
        with opener.open(request, timeout=self.timeout) as response:
            raw = response.read(2_000_001)
            if len(raw) > 2_000_000:
                raise ValueError("Inference response too large")
        return json.loads(raw)["choices"][0]["message"]["content"]


class HFBackend:
    def __init__(self, model_path, *, device="cpu", threads=4, max_new_tokens=256, context_limit=4096):
        import torch
        from transformers import AutoConfig, AutoTokenizer, AutoModelForCausalLM, AutoModelForImageTextToText
        torch.set_num_threads(threads)
        self.torch = torch
        self.tokenizer = AutoTokenizer.from_pretrained(model_path, local_files_only=True, trust_remote_code=False)
        dtype = torch.float32 if device == "cpu" else torch.float16
        config = AutoConfig.from_pretrained(model_path, local_files_only=True, trust_remote_code=False)
        factory = AutoModelForImageTextToText if config.model_type == "qwen3_5" else AutoModelForCausalLM
        self.model = factory.from_pretrained(model_path, local_files_only=True, trust_remote_code=False, dtype=dtype).to(device).eval()
        self.max_new_tokens, self.context_limit = max_new_tokens, context_limit
        self.last_metrics = {}

    def _inputs(self, messages, schema=None):
        if schema:
            messages = [*messages, {"role": "user", "content": "Return JSON only matching: " + json.dumps(schema)}]
        text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True, enable_thinking=False)
        inputs = self.tokenizer(text, return_tensors="pt").to(self.model.device)
        count = inputs["input_ids"].shape[1]
        if count + self.max_new_tokens > self.context_limit:
            raise ValueError("Request exceeds configured context budget; retrieve less context")
        return inputs, count

    def complete(self, messages, schema=None):
        inputs, count = self._inputs(messages, schema)
        start = time.perf_counter()
        with self.torch.inference_mode():
            result = self.model.generate(**inputs, max_new_tokens=self.max_new_tokens, do_sample=False,
                                         pad_token_id=self.tokenizer.eos_token_id)
        if self.model.device.type == "cuda":
            self.torch.cuda.synchronize()
        elapsed = time.perf_counter() - start
        generated = result[0, count:]
        self.last_metrics = {"input_tokens": count, "output_tokens": len(generated), "seconds": elapsed,
                             "tokens_per_second": len(generated) / max(elapsed, 1e-9),
                             "ttft": None, "note": "Non-streaming adapter; TTFT not measured"}
        return self.tokenizer.decode(generated, skip_special_tokens=True)

    def stream(self, messages):
        """One request at a time; caller must serialize access to this backend."""
        from transformers import TextIteratorStreamer, StoppingCriteria, StoppingCriteriaList
        inputs, count = self._inputs(messages)
        cancelled = threading.Event()
        errors = []
        class Stop(StoppingCriteria):
            def __call__(self, input_ids, scores, **kwargs):
                return cancelled.is_set()
        streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=True, timeout=120)
        def generate():
            try:
                with self.torch.inference_mode():
                    self.model.generate(**inputs, max_new_tokens=self.max_new_tokens, do_sample=False,
                                        pad_token_id=self.tokenizer.eos_token_id, streamer=streamer,
                                        stopping_criteria=StoppingCriteriaList([Stop()]))
            except Exception as exc:
                errors.append(exc)
                streamer.on_finalized_text("", stream_end=True)
        start, first = time.perf_counter(), None
        worker = threading.Thread(target=generate, daemon=True)
        worker.start()
        chunks = []
        try:
            for chunk in streamer:
                if chunk:
                    first = time.perf_counter() if first is None else first
                    chunks.append(chunk)
                    yield chunk
            if errors:
                raise errors[0]
        finally:
            cancelled.set()
            worker.join(timeout=120)
            if worker.is_alive():
                raise RuntimeError("Generation did not stop; restart inference worker")
            elapsed = time.perf_counter()-start
            output_tokens = len(self.tokenizer.encode("".join(chunks), add_special_tokens=False))
            self.last_metrics = {"input_tokens": count, "output_tokens_retokenized": output_tokens,
                                 "seconds": elapsed, "time_to_first_text_seconds": None if first is None else first-start,
                                 "note": "Text chunks can buffer multiple tokens; first text is not exact first-token latency"}


def tiny_generate(checkpoint, prompt, max_new_tokens=64):
    import torch
    from safetensors.torch import load_file
    from .model import NexoraLM, ModelConfig
    from .tokenizer import ByteTokenizer
    torch.set_num_threads(4)
    path = Path(checkpoint)
    model = NexoraLM(ModelConfig(**json.loads((path / "config.json").read_text())))
    model.load_state_dict(load_file(str(path / "model.safetensors")))
    tok = ByteTokenizer()
    ids = torch.tensor([[tok.bos_id, *tok.encode(prompt)]])
    out = model.generate(ids, max_new_tokens=max_new_tokens)
    return tok.decode(out[0].tolist())