mmarch/text-generation-webui
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity

Add /v1/internal/logits endpoint (#4650)

Browse source
This commit is contained in:
oobabooga 2023-11-18 23:19:31 -03:00 • committed by GitHub
parent 8f4f4daf8b
commit 0fa1af296c
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
4 changed files with 71 additions and 9 deletions
Download patch file Download diff file Expand all files Collapse all files

18
modules/logits.py
View file

@ -8,7 +8,7 @@ from modules.text_generation import generate_reply
global_scores = None
def get_next_logits(prompt, state, use_samplers, previous):
def get_next_logits(prompt, state, use_samplers, previous, return_dict=False):
if shared.model is None:
logger.error("No model is loaded! Select one in the Model tab.")
return 'Error: No model is loaded1 Select one in the Model tab.', previous
@ -56,8 +56,16 @@ def get_next_logits(prompt, state, use_samplers, previous):
topk_indices = [i.expand((1, 1)) for i in topk_indices]
tokens = [shared.tokenizer.decode(i) for i in topk_indices]
output = ''
for row in list(zip(topk_values, tokens)):
output += f"{row[0]} - {repr(row[1])}\n"
return output, previous
if return_dict:
output = {}
for row in list(zip(topk_values, tokens)):
output[row[1]] = row[0]
return output
else:
output = ''
for row in list(zip(topk_values, tokens)):
output += f"{row[0]} - {repr(row[1])}\n"
return output, previous