File size: 3,949 Bytes
b399ea6 | 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 |
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
|