File size: 5,288 Bytes
4c383ef | 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 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 | """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()
|