Merge pull request #38 from one-some/ui2-probs

Fix probabilities
This commit is contained in:
ebolam
2022-08-17 18:27:37 -04:00
committed by GitHub

View File

@@ -1716,10 +1716,11 @@ def patch_transformers():
assert scores.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
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 = []
for token_id, score in sorted(enumerate(probs), key=lambda x: x[1], reverse=True)[:8]:
@@ -1729,7 +1730,12 @@ def patch_transformers():
"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
def new_get_logits_processor(*args, **kwargs) -> LogitsProcessorList: