Spaces:
Running
Running
| 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 ["<EOS>", "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>" | |
| 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("<UNK>", 1 if "<UNK>" 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() | |