| """Keep the decision model loaded for interactive questions.""" | |
| import argparse | |
| import sys | |
| import time | |
| from pathlib import Path | |
| ROOT = Path(__file__).resolve().parents[1] | |
| sys.path.insert(0, str(ROOT)) | |
| import torch | |
| from mugi_decision.encoding import DecisionEncoder | |
| from mugi_decision.runtime import ( | |
| configure_inference, | |
| dtype_for, | |
| gpu_device, | |
| load_model, | |
| move_batch, | |
| ) | |
| def main(): | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--model", default=str(ROOT)) | |
| parser.add_argument("--max-length", type=int, default=1024) | |
| parser.add_argument("--branch-batch-size", type=int, default=8) | |
| parser.add_argument("--no-cache", action="store_true") | |
| parser.add_argument( | |
| "--optimization", choices=["auto", "none", "hip", "fast"], default="auto" | |
| ) | |
| args = parser.parse_args() | |
| if args.branch_batch_size < 1: | |
| parser.error("--branch-batch-size must be positive") | |
| device = gpu_device() | |
| print("正在加载第 2 轮最佳模型；后续提问无需重新加载。", flush=True) | |
| model, tokenizer = load_model(args.model, device, dtype_for("bf16")) | |
| model.eval() | |
| model, mode = configure_inference(model, args.optimization) | |
| base = model.get_base_model() if hasattr(model, "get_base_model") else model | |
| print(f"推理模式：{mode}；使用 --optimization none 可切回参考实现。") | |
| encoder = DecisionEncoder(tokenizer, base.config.include_enum, args.max_length) | |
| print("问题可输入多行，以空行结束；选项每行一个，以空行结束。输入 /quit 退出。") | |
| print("含空行的长题请先输入 /paste，以 /end 结束，避免段落被提前截断。") | |
| print("百分比是候选间的 softmax 相对分数，不是校准后的正确率。") | |
| try: | |
| while True: | |
| print("\n问题描述：", flush=True) | |
| lines = [] | |
| while True: | |
| line = input() | |
| if line.strip() == "/quit": | |
| return | |
| if line.strip() == "/paste" and not lines: | |
| print("粘贴完整问题，允许空行；独占一行输入 /end 结束。") | |
| while True: | |
| line = input() | |
| if line.strip() == "/quit": | |
| return | |
| if line.strip() == "/end": | |
| break | |
| lines.append(line) | |
| break | |
| if not line.strip(): | |
| break | |
| lines.append(line) | |
| if not lines: | |
| continue | |
| print("选项（至少两个，每行一个）：", flush=True) | |
| options = [] | |
| while True: | |
| line = input(f"{len(options) + 1}> ") | |
| if line.strip() == "/quit": | |
| return | |
| if not line.strip(): | |
| break | |
| options.append({"id": str(len(options) + 1), "text": line.strip()}) | |
| record = {"description": "\n".join(lines), "options": options} | |
| try: | |
| prefix, suffixes, options, _ = encoder.encode(record) | |
| except ValueError as exc: | |
| print(f"输入未处理：{exc}") | |
| continue | |
| torch.cuda.synchronize() | |
| started = time.monotonic() | |
| with torch.inference_mode(): | |
| if args.no_cache: | |
| chunks = [] | |
| for start in range(0, len(options), args.branch_batch_size): | |
| batch = tokenizer.pad( | |
| { | |
| "input_ids": [ | |
| prefix + s | |
| for s in suffixes[ | |
| start : start + args.branch_batch_size | |
| ] | |
| ] | |
| }, | |
| padding=True, | |
| return_tensors="pt", | |
| ) | |
| chunks.append(model(**move_batch(batch, device)).flat_scores) | |
| scores = torch.cat(chunks) | |
| else: | |
| scores = base.score_cached(prefix, suffixes, args.branch_batch_size) | |
| if not torch.isfinite(scores).all(): | |
| raise FloatingPointError("模型评分包含非有限值") | |
| probs = torch.softmax(scores.float(), dim=-1).cpu().tolist() | |
| values = scores.float().cpu().tolist() | |
| winner = max(range(len(options)), key=lambda i: values[i]) | |
| print(f"\n选择 {options[winner]['id']}：{options[winner]['text']}") | |
| for i in sorted(range(len(options)), key=lambda i: values[i], reverse=True): | |
| print( | |
| f" {options[i]['id']:>2}. {probs[i]:7.2%} 分数 {values[i]:+.4f} {options[i]['text']}" | |
| ) | |
| print(f"耗时 {time.monotonic() - started:.3f} 秒") | |
| except (EOFError, KeyboardInterrupt): | |
| print("\n已退出。") | |
| if __name__ == "__main__": | |
| main() | |