Spaces:
Runtime error
Runtime error
| """ | |
| Phase N - Interactive multi-turn chat REPL for The Abyss-1 (SFT Checkpoint) | |
| Builds The Abyss chat template on every turn, streams tokesn as they | |
| generate, and stops cleanly at <|endofturn|> or <|endoftext|> | |
| Usage (from the_final_product/): | |
| python -m ml.chat | |
| python -m ml.chat --checkpoint ml/checkpoints/sft_ckpt.pt | |
| python -m ml.chat --temperature 0.7 --max_tokens 512 | |
| Commands during chat: | |
| reset - clear conversation history | |
| quit - exit | |
| """ | |
| import argparse | |
| import sys | |
| import torch | |
| import torch.nn.functional as F | |
| from ml.model import TheAbyssModel | |
| from ml.tokenizer.tokenizer import AbyssTokenizer | |
| TOKENIZER_DIR = "ml/tokenizer/vocab" | |
| SYSTEM_TEXT = "You are The Abyss, and empty void of nothingness" | |
| def respond(model, tok, history, sys_ids, cfg, | |
| max_new_tokens, temperature, top_k, top_p, | |
| EOS, EOT, USR, AST, device): | |
| """ | |
| Build the full prompt from history, sample until EOT/EOS, | |
| stream each token to stdout. Returns list of response token ids. | |
| """ | |
| # ----- BUILD PROMPT ----- | |
| prompt = list(sys_ids) | |
| for turn in history: | |
| prompt += turn | |
| prompt += [AST] | |
| # If prompt is too long, drop oldest turns (keep system + recent turns) | |
| max_prompt = cfg.block_size - max_new_tokens - 1 | |
| if len(prompt) > max_prompt: | |
| prompt = list(sys_ids) | |
| tail = [] | |
| for turn in reversed(history): | |
| if len(sys_ids) + len(turn) + len(tail) + 1 > max_prompt: | |
| break | |
| tail = list(turn) + tail | |
| prompt = list(sys_ids) + tail + [AST] | |
| idx = torch.tensor([prompt], dtype=torch.long, device=device) | |
| # ----- SAMPLE ----- | |
| response_ids = [] | |
| for _ in range(max_new_tokens): | |
| idx_cond = idx[:, -cfg.block_size:] | |
| logits, _ = model(idx_cond) | |
| logits = logits[0, -1, :] / temperature # (V,) — 1D | |
| # top_k | |
| if top_k > 0: | |
| k = min(top_k, logits.size(-1)) | |
| kth = torch.topk(logits, k).values[-1] | |
| logits = logits.masked_fill(logits < kth, float("-inf")) | |
| # top_p (nucleus) | |
| if 0.0 < top_p < 1.0: | |
| s_logits, s_idx = torch.sort(logits, descending=True) | |
| cum = torch.cumsum(F.softmax(s_logits, dim=-1), dim=-1) | |
| mask = (cum - F.softmax(s_logits, dim=-1)) > top_p | |
| s_logits = s_logits.masked_fill(mask, float("-inf")) | |
| logits = torch.zeros_like(logits).scatter_(0, s_idx, s_logits) | |
| next_tok = torch.multinomial(F.softmax(logits, dim=-1), 1) # (1,) | |
| tok_id = next_tok.item() | |
| # STOP ON END-OF-TURN or EOS | |
| if tok_id == EOT or tok_id == EOS: | |
| break | |
| response_ids.append(tok_id) | |
| sys.stdout.write(tok.decode([tok_id])) | |
| sys.stdout.flush() | |
| idx = torch.cat([idx, next_tok.view(1, 1)], dim=1) | |
| print() | |
| return response_ids | |
| def main(): | |
| parser = argparse.ArgumentParser(description="Chat with The Abyss-1") | |
| parser.add_argument("--checkpoint", default="ml/checkpoints/sft_ckpt.pt") | |
| parser.add_argument("--max_tokens", type=int, default=256) | |
| parser.add_argument("--temperature", type=float, default=0.8) | |
| parser.add_argument("--top_k", type=int, default=50) | |
| parser.add_argument("--top_p", type=float, default=0.95) | |
| parser.add_argument("--device", | |
| default="cuda" if torch.cuda.is_available() else "cpu") | |
| args = parser.parse_args() | |
| # ----- LOAD MODEL ----- | |
| print(f"Loading {args.checkpoint} ...", flush=True) | |
| ck = torch.load(args.checkpoint, map_location=args.device, weights_only=False) | |
| cfg = ck["config"] | |
| model = TheAbyssModel(cfg).to(args.device) | |
| model.load_state_dict(ck["model"]) | |
| model.eval() | |
| print(f" iter={ck.get('iter', '?')} params={model.num_params():,}\n", flush=True) | |
| # ----- LOAD TOKENIZER ----- | |
| tok = AbyssTokenizer.load(TOKENIZER_DIR) | |
| EOS = tok.encode("<|endoftext|>")[0] | |
| USR = tok.encode("<|user|>")[0] | |
| AST = tok.encode("<|assistant|>")[0] | |
| EOT = tok.encode("<|endofturn|>")[0] | |
| SYS = tok.encode("<|system|>")[0] | |
| sys_ids = [SYS] + tok.encode(SYSTEM_TEXT) + [EOT] | |
| print("The Abyss Chat - type 'reset' to clear history, 'quit' to exit\n") | |
| print("-" *50) | |
| history = [] | |
| while True: | |
| try: | |
| user_input = input("\nYou: ").strip() | |
| except (EOFError, KeyboardInterrupt): | |
| print("\nBye.") | |
| break | |
| if not user_input: | |
| continue | |
| if user_input.lower() in ("quit", "exit"): | |
| print("Bye.") | |
| break | |
| if user_input.lower() == "reset": | |
| history.clear() | |
| print(" [history cleared]") | |
| continue | |
| # ENCODE AND STORE USER TURN | |
| user_turn = [USR] + tok.encode(user_input) + [EOT] | |
| history.append(user_turn) | |
| print("\nThe Abyss: ", end="", flush=True) | |
| response_ids = respond( | |
| model, tok, history, sys_ids, cfg, | |
| args.max_tokens, args.temperature, args.top_k, args.top_p, | |
| EOS, EOT, USR, AST, args.device | |
| ) | |
| # STORE ASSISTANT TURN | |
| history.append([AST] + response_ids + [EOT]) | |
| if __name__ == "__main__": | |
| main() |