import os, re, glob, json, requests, tempfile, torch, unicodedata from bot_model import GPT from unit_parser import ( canonicalize_nl, try_solve_unit_like, is_addition_query, canonicalize_addition, sum_addends, ) # ----------------- Token helpers ----------------- def decode(ids, itos): return "".join(itos[i] for i in ids) MAX_NEW_TOKENS = 512 def greedy_until_eos(model, idx, eos_id=None, max_new_tokens=MAX_NEW_TOKENS): model.eval() with torch.no_grad(): gen = idx for _ in range(max_new_tokens): idx_cond = gen[:, -model.pos_emb.num_embeddings:] logits = model(idx_cond)[:, -1, :] next_id = torch.argmax(logits, dim=-1, keepdim=True) gen = torch.cat([gen, next_id], dim=1) if eos_id is not None and gen[0, -1].item() == eos_id: break return gen # ----------------- Input normalization ----------------- SMART_MAP = { "\u2013": "-", # en dash "\u2014": "-", # em dash "\u2212": "-", # minus sign "\u00D7": "x", # multiplication sign "\u00F7": "/", # division sign "\u2018": "'", "\u2019": "'", "\u201C": '"', "\u201D": '"', "\u00A0": " ", # non-breaking space } def normalize_input(s: str) -> str: s = unicodedata.normalize("NFKC", s) for k, v in SMART_MAP.items(): s = s.replace(k, v) return " ".join(s.split()) # ----------------- Stop helper ----------------- def _stop_at(resp: str) -> str: for marker in ["", "final answer:"]: if marker in resp: return resp.split(marker)[0] + (marker if marker == "final answer:" else "") return resp # ----------------- Checkpoint loader ----------------- def _find_ckpt_local(): cand = sorted(glob.glob("checkpoints/*.pt")) or sorted(glob.glob("*.pt")) return cand[0] if cand else None def _load_ckpt(): local = _find_ckpt_local() if local and os.path.exists(local): return torch.load(local, map_location="cpu") raise FileNotFoundError("No checkpoint found locally in checkpoints/ or repo root.") _checkpoint = _load_ckpt() # ----------------- Vocab loading ----------------- def _normalize_to_itos(x): if x is None: return None if isinstance(x, dict): if x and all(isinstance(k, int) for k in x.keys()): return x if x and all(isinstance(k, str) for k in x.keys()): return {i: ch for ch, i in x.items()} return None if isinstance(x, list) and x: return {i: ch for i, ch in enumerate(x)} return None def _normalize_to_stoi(x): if x is None: return None if isinstance(x, dict): if x and all(isinstance(k, str) for k in x.keys()): return x if x and all(isinstance(k, int) for k in x.keys()): return {ch: i for i, ch in x.items()} return None if isinstance(x, list) and x: return {ch: i for i, ch in enumerate(x)} return None def _load_vocab_json(): cand = [os.environ.get("VOCAB_JSON_PATH",""), "vocab.json", "checkpoints/vocab.json"] for p in cand: p = p.strip() if not p: continue if os.path.exists(p): with open(p, "r", encoding="utf-8") as f: data = json.load(f) itos_raw = data.get("itos") stoi_raw = data.get("stoi") itos = _normalize_to_itos(itos_raw) stoi = _normalize_to_stoi(stoi_raw) if itos is None and isinstance(stoi, dict) and stoi: itos = {i: ch for ch, i in stoi.items()} if stoi is None and isinstance(itos, dict) and itos: stoi = {ch: i for i, ch in itos.items()} if isinstance(stoi, dict) and stoi and isinstance(itos, dict) and itos: return stoi, itos return None, None def _build_vocab(ckpt): v = ckpt.get("vocab", None) stoi = itos = None if isinstance(v, (tuple, list)) and len(v) >= 2: a, b = v[0], v[1] itos = _normalize_to_itos(a) or _normalize_to_itos(b) stoi = _normalize_to_stoi(a) or _normalize_to_stoi(b) elif isinstance(v, dict) and v is not None: itos = _normalize_to_itos(v.get("itos", None)) or _normalize_to_itos(v) stoi = _normalize_to_stoi(v.get("stoi", None)) or _normalize_to_stoi(v) else: itos = _normalize_to_itos(ckpt.get("itos", None)) stoi = _normalize_to_stoi(ckpt.get("stoi", None)) if itos is None and isinstance(stoi, dict) and stoi: itos = {i: ch for ch, i in stoi.items()} if stoi is None and isinstance(itos, dict) and itos: stoi = {ch: i for i, ch in itos.items()} if not (isinstance(stoi, dict) and stoi and isinstance(itos, dict) and itos): stoi, itos = _load_vocab_json() if not (isinstance(stoi, dict) and stoi and isinstance(itos, dict) and itos): raise KeyError("No usable vocab found: checkpoint has empty/None vocab and vocab.json missing/invalid.") return stoi, itos stoi, itos = _build_vocab(_checkpoint) expected = _checkpoint["model"]["token_emb.weight"].shape[0] actual = len(stoi) if expected != actual: raise RuntimeError(f"Vocab size mismatch: checkpoint expects {expected}, but vocab.json has {actual}.") EOS = "" EOS_ID = stoi.get(EOS, None) # ----------------- Build model ----------------- block_size = _checkpoint.get("block_size", None) if block_size is None: pe = _checkpoint["model"].get("pos_emb.weight", None) block_size = pe.shape[0] if hasattr(pe, "shape") else 320 if "model_cfg" in _checkpoint: mc = _checkpoint["model_cfg"] model = GPT(len(stoi), d_model=mc.get("d_model",192), num_heads=mc.get("num_heads",6), num_layers=mc.get("num_layers",4), max_len=block_size, dropout=mc.get("dropout",0.1)) else: model = GPT(len(stoi), d_model=192, num_heads=6, num_layers=4, max_len=block_size) model.load_state_dict(_checkpoint["model"]) model.eval() # ----------------- Public entry point ----------------- def _encode(s: str): unk = stoi.get("", 1 if "" in stoi else 0) return [stoi.get(c, unk) for c in s] def chatbot_reply(user_text: str, history=None) -> str: q_raw = normalize_input(user_text.strip()) q = canonicalize_nl(q_raw) or q_raw unit_det = try_solve_unit_like(q_raw) if unit_det: return f"q: {q_raw} a: {unit_det}" if is_addition_query(q_raw): can = canonicalize_addition(q_raw) total = sum_addends(q_raw) return f"q: {can} a: {can} = {total}. final answer: {total}" prompt = f"q: {q} a:" if not q.lower().startswith("q:") else q idx = torch.tensor([_encode(prompt)], dtype=torch.long) out = greedy_until_eos(model, idx, EOS_ID) resp = decode(out[0].tolist(), itos) resp = _stop_at(resp) if EOS in resp: resp = resp.split(EOS)[0] return resp.split("a:",1)[-1].strip()