File size: 12,096 Bytes
9d614f1 | 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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 | """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() # one generation at a time (CPU / shared weights)
self._load_taught()
# ------------------------------------------------------------------ teach
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
# ------------------------------------------------------------------ prompt
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)
# ------------------------------------------------------------------ generation
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: # unblock the consumer, then re-raise below
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"])
|