fix: use correct variable for logprobs in batch generation (ml-explore#800)

luqmannurhakimbazman · web-flow · commit d9846d37cb73 · 2026-01-23T07:24:39.000-08:00
diff --git a/mlx_lm/server.py b/mlx_lm/server.py
@@ -702,10 +702,10 @@ def progress_callback(info):
                         top_tokens = None
                         if args.logprobs > 0:
                             sorted_indices = mx.argpartition(
-                                -gen.logprobs, kth=args.logprobs - 1
+                                -r.logprobs, kth=args.logprobs - 1
                             )
                             top_indices = sorted_indices[: args.logprobs]
-                            top_logprobs = gen.logprobs[top_indices]
+                            top_logprobs = r.logprobs[top_indices]
                             top_token_info = zip(
                                 top_indices.tolist(), top_logprobs.tolist()
                             )

Original file line number	Diff line number	Diff line change
`@@ -702,10 +702,10 @@ def progress_callback(info):`
`702`	`702`	`top_tokens = None`
`703`	`703`	`if args.logprobs > 0:`
`704`	`704`	`sorted_indices = mx.argpartition(`
`705`		`- -gen.logprobs, kth=args.logprobs - 1`
	`705`	`+ -r.logprobs, kth=args.logprobs - 1`
`706`	`706`	`)`
`707`	`707`	`top_indices = sorted_indices[: args.logprobs]`
`708`		`- top_logprobs = gen.logprobs[top_indices]`
	`708`	`+ top_logprobs = r.logprobs[top_indices]`
`709`	`709`	`top_token_info = zip(`
`710`	`710`	`top_indices.tolist(), top_logprobs.tolist()`
`711`	`711`	`)`