fix indexing token_logprobs after sorting

wu-qing-157 · wu-qing-157 · commit 9e61661518d7 · 2023-07-07T10:18:49.000Z
diff --git a/llama_cpp/llama.py b/llama_cpp/llama.py
@@ -958,7 +958,7 @@ def _create_completion(
                                 )
                             ],
                             "text_offset": [text_offset],
-                            "token_logprobs": [sorted_logprobs[int(token)][0]],
+                            "token_logprobs": [current_logprobs[int(token)]],
                             "top_logprobs": [top_logprob],
                         }
                     returned_tokens += 1
@@ -1033,7 +1033,7 @@ def _create_completion(
                             self.detokenize([token]).decode("utf-8", errors="ignore")
                         ],
                         "text_offset": [text_offset],
-                        "token_logprobs": [sorted_logprobs[int(token)][0]],
+                        "token_logprobs": [current_logprobs[int(token)]],
                         "top_logprobs": [top_logprob],
                     }
 
@@ -1131,7 +1131,7 @@ def _create_completion(
                         zip(logprobs_token, range(len(logprobs_token))), reverse=True
                     )
                 )
-                token_logprobs.append(sorted_logprobs[int(token)][0])
+                token_logprobs.append(logprobs_token[int(token)])
                 top_logprob: Optional[Dict[str, float]] = {
                     self.detokenize([i]).decode("utf-8", errors="ignore"): logprob
                     for logprob, i in sorted_logprobs[:logprobs]

Original file line number	Diff line number	Diff line change
`@@ -958,7 +958,7 @@ def _create_completion(`
`958`	`958`	`)`
`959`	`959`	`],`
`960`	`960`	`"text_offset": [text_offset],`
`961`		`- "token_logprobs": [sorted_logprobs[int(token)][0]],`
	`961`	`+ "token_logprobs": [current_logprobs[int(token)]],`
`962`	`962`	`"top_logprobs": [top_logprob],`
`963`	`963`	`}`
`964`	`964`	`returned_tokens += 1`
`@@ -1033,7 +1033,7 @@ def _create_completion(`
`1033`	`1033`	`self.detokenize([token]).decode("utf-8", errors="ignore")`
`1034`	`1034`	`],`
`1035`	`1035`	`"text_offset": [text_offset],`
`1036`		`- "token_logprobs": [sorted_logprobs[int(token)][0]],`
	`1036`	`+ "token_logprobs": [current_logprobs[int(token)]],`
`1037`	`1037`	`"top_logprobs": [top_logprob],`
`1038`	`1038`	`}`
`1039`	`1039`
`@@ -1131,7 +1131,7 @@ def _create_completion(`
`1131`	`1131`	`zip(logprobs_token, range(len(logprobs_token))), reverse=True`
`1132`	`1132`	`)`
`1133`	`1133`	`)`
`1134`		`- token_logprobs.append(sorted_logprobs[int(token)][0])`
	`1134`	`+ token_logprobs.append(logprobs_token[int(token)])`
`1135`	`1135`	`top_logprob: Optional[Dict[str, float]] = {`
`1136`	`1136`	`self.detokenize([i]).decode("utf-8", errors="ignore"): logprob`
`1137`	`1137`	`for logprob, i in sorted_logprobs[:logprobs]`