| """Inference engine (PRIVATE to Serving_Module). |
| |
| * Modes: non_thinking (direct answer), thinking (<think>...</think> then answer), max_thinking (long reasoning). |
| The mode is selected through the same `[MODE: ...]` system-prompt tag the model was fine-tuned with. |
| * Lag-free: KV-cache generation in a worker thread + TextIteratorStreamer, so tokens appear as they are produced. |
| * Stops cleanly on <|im_end|> (the stop token the SFT stage aligned) and never prints special tokens. |
| * Live `teach`: a question/answer pair is used IMMEDIATELY (direct hit or injected as a known fact) and is also |
| appended to chat.jsonl in the exact SFT schema, so the next training round learns it permanently. |
| |
| The chat format is intentionally duplicated here (it mirrors Training_Module's tokenizer_surgeon) because |
| modules may not import each other's internals. |
| """ |
| from __future__ import annotations |
|
|
| import difflib |
| import json |
| import re |
| import threading |
| import time |
| from collections import deque |
| from pathlib import Path |
| from typing import Callable, Deque, Dict, List, Optional, Tuple |
|
|
| IM_START, IM_END, ENDOFTEXT = "<|im_start|>", "<|im_end|>", "<|endoftext|>" |
| THINK_OPEN, THINK_CLOSE = "<think>", "</think>" |
| MODES = ("thinking", "non_thinking", "max_thinking") |
| DEFAULT_SYSTEM = "You are RSS AI, a helpful assistant created by RSS. You understand Hindi, Hinglish and English." |
| TOKEN_BUDGET = {"non_thinking": 220, "thinking": 420, "max_thinking": 800} |
| SAMPLING = {"non_thinking": 0.45, "thinking": 0.6, "max_thinking": 0.65} |
| _WS = re.compile(r"\s+") |
| _SPECIAL = re.compile(re.escape(IM_END) + "|" + re.escape(ENDOFTEXT) + "|" + re.escape(IM_START)) |
|
|
|
|
| def normalise_mode(value: str) -> Optional[str]: |
| value = (value or "").strip().lower().replace("-", "_").replace(" ", "_").lstrip("/") |
| return value if value in MODES else None |
|
|
|
|
| def _norm(text: str) -> str: |
| return _WS.sub(" ", text.strip()).casefold() |
|
|
|
|
| def split_thinking(text: str) -> Tuple[str, str, bool]: |
| """-> (thinking, answer, truncated). `truncated` = the token budget ended inside <think>.""" |
| text = _SPECIAL.sub("", text).strip() |
| if THINK_CLOSE in text: |
| head, _, answer = text.partition(THINK_CLOSE) |
| return head.replace(THINK_OPEN, "").strip(), answer.strip(), False |
| if THINK_OPEN in text: |
| return text.replace(THINK_OPEN, "").strip(), "", True |
| return "", text, False |
|
|
|
|
| class InferenceEngine: |
| def __init__(self, model, tokenizer, chat_file: Optional[str] = None, max_new_tokens: Optional[int] = None, |
| history_turns: int = 3, active_vocab_file: Optional[str] = None): |
| import torch |
| self.torch = torch |
| self.model, self.tokenizer = model, tokenizer |
| self.model.eval() |
| try: |
| self.model.config.use_cache = True |
| except Exception: |
| pass |
| try: |
| self.device = next(model.parameters()).device |
| except StopIteration: |
| self.device = torch.device("cpu") |
| vocab = tokenizer.get_vocab() |
| self.eos_ids = [vocab[t] for t in (IM_END, ENDOFTEXT) if t in vocab] or [tokenizer.eos_token_id] |
| self.pad_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else self.eos_ids[0] |
| self.max_new_tokens = max_new_tokens |
| self.active_token_ids: Optional[List[int]] = None |
| if active_vocab_file: |
| vocab_path = Path(active_vocab_file) |
| if vocab_path.is_file(): |
| try: |
| self.active_token_ids = sorted({int(x) for x in json.loads(vocab_path.read_text(encoding="utf-8"))}) |
| print(f"[Serving] active vocabulary mask loaded: {len(self.active_token_ids):,} token IDs", flush=True) |
| except Exception as exc: |
| raise RuntimeError(f"could not load active vocabulary manifest {vocab_path}: {exc}") from exc |
| self.chat_file = Path(chat_file) if chat_file else None |
| self.history: Dict[str, Deque[Tuple[str, str]]] = {} |
| self.history_turns = history_turns |
| self.taught: Dict[str, Tuple[str, str]] = {} |
| self.lock = threading.Lock() |
| self._load_taught() |
|
|
| |
| def _load_taught(self) -> None: |
| if not self.chat_file or not self.chat_file.exists(): |
| return |
| try: |
| with open(self.chat_file, "r", encoding="utf-8") as handle: |
| for line in handle: |
| if '"teach-' not in line: |
| continue |
| try: |
| rec = json.loads(line) |
| except ValueError: |
| continue |
| if str(rec.get("id", "")).startswith("teach-") and rec.get("user") and rec.get("assistant"): |
| self.taught[_norm(rec["user"])] = (rec["user"], rec["assistant"]) |
| except OSError: |
| pass |
|
|
| def teach(self, question: str, answer: str) -> str: |
| question, answer = question.strip(), answer.strip() |
| if not question or not answer: |
| return "teach: needs both a question and an answer -> teach: सवाल | जवाब" |
| self.taught[_norm(question)] = (question, answer) |
| saved = "" |
| if self.chat_file: |
| record = {"id": f"teach-{int(time.time() * 1000)}", "mode": "non_thinking", "system": DEFAULT_SYSTEM, |
| "user": question, "thinking": "", "assistant": answer} |
| try: |
| self.chat_file.parent.mkdir(parents=True, exist_ok=True) |
| with open(self.chat_file, "a", encoding="utf-8") as handle: |
| handle.write(json.dumps(record, ensure_ascii=False) + "\n") |
| saved = f" and saved to {self.chat_file.name} (next training round will learn it)" |
| except OSError as exc: |
| saved = f" (could not write {self.chat_file.name}: {exc})" |
| return f"learned: {question!r} -> {answer!r}{saved}" |
|
|
| def _taught_lookup(self, question: str) -> Tuple[Optional[str], List[Tuple[str, str]]]: |
| key = _norm(question) |
| if key in self.taught: |
| return self.taught[key][1], [] |
| scored = sorted(((difflib.SequenceMatcher(None, key, k).ratio(), k) for k in self.taught), reverse=True) |
| if scored and scored[0][0] >= 0.90: |
| return self.taught[scored[0][1]][1], [] |
| related = [self.taught[k] for ratio, k in scored[:3] if ratio >= 0.55] |
| return None, related |
|
|
| |
| def _render(self, user: str, mode: str, session: str, facts: List[Tuple[str, str]]) -> str: |
| system = DEFAULT_SYSTEM |
| if facts: |
| system += "\nKnown facts taught by the user: " + " | ".join(f"{q} -> {a}" for q, a in facts) |
| parts = [f"{IM_START}system\n{system}\n[MODE: {mode}]{IM_END}\n"] |
| for past_user, past_answer in self.history.get(session, []): |
| parts.append(f"{IM_START}user\n{past_user}{IM_END}\n{IM_START}assistant\n{past_answer}{IM_END}\n") |
| parts.append(f"{IM_START}user\n{user}{IM_END}\n{IM_START}assistant\n") |
| return "".join(parts) |
|
|
| def reset(self, session: str = "default") -> None: |
| self.history.pop(session, None) |
|
|
| |
| def generate(self, user: str, mode: str = "non_thinking", session: str = "default", |
| on_text: Optional[Callable[[str], None]] = None) -> Dict: |
| mode = normalise_mode(mode) or "non_thinking" |
| started = time.time() |
| direct, facts = self._taught_lookup(user) |
| if direct is not None: |
| if on_text: |
| on_text(direct) |
| self._remember(session, user, direct) |
| return {"mode": mode, "thinking": "", "answer": direct, "truncated": False, "source": "taught", |
| "tokens": 0, "seconds": round(time.time() - started, 2)} |
| prompt = self._render(user, mode, session, facts) |
| budget = self.max_new_tokens or TOKEN_BUDGET[mode] |
| with self.lock: |
| raw, tokens = self._run(prompt, mode, budget, on_text) |
| thinking, answer, truncated = split_thinking(raw) |
| if truncated or (not answer and mode != "non_thinking"): |
| answer = answer or "(token limit reached while thinking - ask again, or use mode non_thinking for a direct answer)" |
| self._remember(session, user, answer) |
| return {"mode": mode, "thinking": thinking, "answer": answer, "truncated": truncated, "source": "model", |
| "tokens": tokens, "seconds": round(time.time() - started, 2)} |
|
|
| def _remember(self, session: str, user: str, answer: str) -> None: |
| turns = self.history.setdefault(session, deque(maxlen=self.history_turns)) |
| turns.append((user, answer)) |
|
|
| def _run(self, prompt: str, mode: str, budget: int, on_text) -> Tuple[str, int]: |
| torch = self.torch |
| enc = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False).to(self.device) |
| kwargs = dict(**enc, max_new_tokens=budget, do_sample=True, temperature=SAMPLING[mode], top_p=0.9, |
| repetition_penalty=1.08, eos_token_id=self.eos_ids, pad_token_id=self.pad_id, use_cache=True) |
| if self.active_token_ids: |
| try: |
| from transformers import LogitsProcessor, LogitsProcessorList |
| except ImportError as exc: |
| raise RuntimeError("transformers is required to enforce the active-vocabulary generation mask") from exc |
| allowed_ids = self.active_token_ids |
|
|
| class _ActiveVocabLogitsProcessor(LogitsProcessor): |
| def __init__(self, ids): |
| self.ids = ids |
| self._masks = {} |
|
|
| def __call__(self, input_ids, scores): |
| key = (str(scores.device), int(scores.shape[-1])) |
| mask = self._masks.get(key) |
| if mask is None: |
| if self.ids[-1] >= scores.shape[-1]: |
| raise RuntimeError("active vocabulary manifest contains an ID outside model logits") |
| mask = torch.zeros(scores.shape[-1], dtype=torch.bool, device=scores.device) |
| mask[torch.as_tensor(self.ids, dtype=torch.long, device=scores.device)] = True |
| self._masks[key] = mask |
| scores.masked_fill_(~mask.unsqueeze(0), -float("inf")) |
| return scores |
|
|
| kwargs["logits_processor"] = LogitsProcessorList([_ActiveVocabLogitsProcessor(allowed_ids)]) |
| if on_text is None: |
| with torch.inference_mode(): |
| out = self.model.generate(**kwargs) |
| new_ids = out[0][enc["input_ids"].shape[1]:] |
| return self.tokenizer.decode(new_ids, skip_special_tokens=False), int(new_ids.shape[0]) |
|
|
| from transformers import TextIteratorStreamer |
| streamer = TextIteratorStreamer(self.tokenizer, skip_prompt=True, skip_special_tokens=False, timeout=900.0) |
| failure: List[BaseException] = [] |
|
|
| def worker() -> None: |
| try: |
| with torch.inference_mode(): |
| self.model.generate(**kwargs, streamer=streamer) |
| except BaseException as exc: |
| failure.append(exc) |
| streamer.text_queue.put(streamer.stop_signal) |
|
|
| thread = threading.Thread(target=worker, daemon=True) |
| thread.start() |
| collected: List[str] = [] |
| for chunk in streamer: |
| collected.append(chunk) |
| shown = _SPECIAL.sub("", chunk) |
| if shown: |
| on_text(shown) |
| thread.join() |
| if failure: |
| raise failure[0] |
| raw = "".join(collected) |
| return raw, len(self.tokenizer(raw, add_special_tokens=False)["input_ids"]) |
|
|
