|
|
| import re |
| import time |
| import torch |
|
|
| class SynapseController: |
| def __init__(self, model, tokenizer): |
| self.model = model |
| self.tokenizer = tokenizer |
| self.device = next(model.parameters()).device |
| self.eot_id = tokenizer.convert_tokens_to_ids("<|eot|>") |
|
|
| def _prompt(self, question): |
| return f"<|user|>\n{question}\n<|assistant|>\n" |
|
|
| def _clean(self, text): |
| text = str(text) |
| if "<|eot|>" in text: |
| text = text.split("<|eot|>")[0] |
| if "<|final|>" in text: |
| text = text.split("<|final|>")[-1] |
| return text.strip() |
|
|
| def _model(self, question, max_new_tokens=140): |
| inputs = self.tokenizer(self._prompt(question), return_tensors="pt").to(self.device) |
| input_len = inputs["input_ids"].shape[-1] |
| start = time.time() |
| with torch.no_grad(): |
| out = self.model.generate( |
| **inputs, |
| max_new_tokens=max_new_tokens, |
| min_new_tokens=4, |
| do_sample=False, |
| pad_token_id=self.tokenizer.pad_token_id, |
| eos_token_id=self.eot_id, |
| repetition_penalty=1.12, |
| no_repeat_ngram_size=4, |
| ) |
| raw = self.tokenizer.decode(out[0][input_len:], skip_special_tokens=False).strip() |
| return self._clean(raw), raw, round(time.time() - start, 3) |
|
|
| def _nsr(self, question): |
| q = str(question).strip() |
| m = re.search(r"Calculate:\s*(-?\d+)\s*([\+\-\*])\s*(-?\d+)", q) |
| if m: |
| a, op, b = int(m.group(1)), m.group(2), int(m.group(3)) |
| val = a + b if op == "+" else a - b if op == "-" else a * b |
| return True, f"Final answer: {val}", f"{a} {op} {b} = {val}" |
| m = re.search(r"A shop has\s+(\d+)\s+boxes\. Each box has\s+(\d+)\s+pencils\.\s+(\d+)\s+pencils are lost", q) |
| if m: |
| boxes, items, lost = int(m.group(1)), int(m.group(2)), int(m.group(3)) |
| val = boxes * items - lost |
| return True, f"Final answer: {val}", f"{boxes} * {items} - {lost} = {val}" |
| return False, "", "" |
|
|
| def _uqm(self, question): |
| q = str(question).lower() |
| risky = ["private password", "hidden bank pin", "silently think", "lost private letter", "unpublished diary", "unknown person"] |
| if any(x in q for x in risky): |
| return True, "I do not have sufficient information.", "Insufficient evidence or unknowable/private information." |
| return False, "", "" |
|
|
| def _tms(self, question): |
| q = str(question) |
| if "Context:" not in q or "Question:" not in q: |
| return False, "", "" |
| context = q.split("Context:", 1)[1].split("Question:", 1)[0].strip() |
| patterns = [ |
| r"called\s+([A-Z][A-Za-z0-9\-]+(?:\s+[A-Z][A-Za-z0-9\-]+)*)", |
| r"named\s+([A-Z][A-Za-z0-9\-]+(?:\s+[A-Z][A-Za-z0-9\-]+)*)", |
| r"in\s+([A-Z][A-Za-z0-9\-]+)\s+in\s+\d{4}", |
| r"code name\s+([A-Z][A-Za-z0-9\-]+(?:\-[A-Z][A-Za-z0-9\-]+)*)", |
| ] |
| for pat in patterns: |
| m = re.search(pat, context) |
| if m: |
| ans = m.group(1).strip() |
| return True, ans, f"Extracted from context: {ans}" |
| return False, "", "" |
|
|
| def generate(self, question, return_trace=True, **kwargs): |
| for route_name, fn in [("UQM_ABSTAIN", self._uqm), ("NSR_CALCULATOR", self._nsr), ("TMS_CONTEXT", self._tms)]: |
| ok, ans, trace = fn(question) |
| if ok: |
| result = {"answer": ans, "route": route_name, "trace": trace} |
| return result if return_trace else ans |
|
|
| ans, raw, latency = self._model(question, max_new_tokens=kwargs.get("max_new_tokens", 140)) |
| result = {"answer": ans, "route": "MODEL_FALLBACK", "trace": "Used pure model fallback.", "raw_answer": raw, "latency_sec": latency} |
| return result if return_trace else ans |
|
|
