From cc312c1cc614bc371cd85dcbcbc1d3ba1590f364 Mon Sep 17 00:00:00 2001 From: jingyaogong Date: Fri, 18 Sep 2026 16:33:09 +0800 Subject: [PATCH] [fix] unify Ctrl+C handling --- eval_llm.py | 14 +++++--------- model/model_minimind.py | 2 +- 2 files changed, 6 insertions(+), 10 deletions(-) diff --git a/eval_llm.py b/eval_llm.py index 8588898..13b076c 100755 --- a/eval_llm.py +++ b/eval_llm.py @@ -64,14 +64,7 @@ def main(): input_mode = int(input('[0] 自动测试\n[1] 手动输入\n')) streamer = TextStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) - def _next_prompt(): - try: - return input('💬: ') - except (EOFError, KeyboardInterrupt): - print() - return '' - - prompt_iter = prompts if input_mode == 0 else iter(_next_prompt, '') + prompt_iter = prompts if input_mode == 0 else iter(lambda: input('💬: '), '') for prompt in prompt_iter: setup_seed(random.randint(0, 31415926)) if input_mode == 0: print(f'💬: {prompt}') @@ -98,4 +91,7 @@ def main(): print(f'\n[Speed]: {gen_tokens / (time.time() - st):.2f} tokens/s\n\n') if args.show_speed else print('\n\n') if __name__ == "__main__": - main() \ No newline at end of file + try: + main() + except (EOFError, KeyboardInterrupt): + print() \ No newline at end of file diff --git a/model/model_minimind.py b/model/model_minimind.py index b2e314e..0f0c3ef 100755 --- a/model/model_minimind.py +++ b/model/model_minimind.py @@ -163,7 +163,7 @@ class MOEFeedForward(nn.Module): topk_weight = topk_weight / (topk_weight.sum(dim=-1, keepdim=True) + 1e-20) else: top1 = torch.topk(F.softmax(self.gate(x_flat.detach()), dim=-1), k=1, dim=-1, sorted=False)[0] - topk_weight = top1 - top1.detach() + 1.0 + topk_weight = top1 - top1.detach() + 1.0 # k=1: 1.0 forward on purpose, gradient via straight-through y = torch.zeros_like(x_flat) for i, expert in enumerate(self.experts): mask = (topk_idx == i)