Fix probabilities

This commit is contained in:
somebody
2022-08-17 14:13:04 -05:00
parent 601c20ff77
commit 7ed48652dd

View File

@@ -1717,10 +1717,11 @@ def patch_transformers():
assert scores.ndim == 2 assert scores.ndim == 2
assert input_ids.ndim == 2 assert input_ids.ndim == 2
if koboldai_vars.numseqs > 1 or not koboldai_vars.show_probs: if not koboldai_vars.show_probs:
return scores return scores
probs = F.softmax(scores, dim = -1).cpu().numpy()[0] for batch_index, batch in enumerate(scores):
probs = F.softmax(batch, dim = -1).cpu().numpy()
token_prob_info = [] token_prob_info = []
for token_id, score in sorted(enumerate(probs), key=lambda x: x[1], reverse=True)[:8]: for token_id, score in sorted(enumerate(probs), key=lambda x: x[1], reverse=True)[:8]:
@@ -1730,7 +1731,10 @@ def patch_transformers():
"score": float(score), "score": float(score),
}) })
koboldai_vars.token_stream_queue.probability_buffer = token_prob_info if len(scores) == 1:
koboldai_vars.actions.set_probabilities(token_prob_info)
else:
koboldai_vars.actions.set_option_probabilities(token_prob_info, batch_index)
return scores return scores
def new_get_logits_processor(*args, **kwargs) -> LogitsProcessorList: def new_get_logits_processor(*args, **kwargs) -> LogitsProcessorList: