mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-10-01 01:26:03 -04:00
[OpenAI Extension] Add 'max_logits' parameter in logits endpoint (#4916)
This commit is contained in:
parent
eaa1fe67f3
commit
e53f99faa0
@ -8,4 +8,4 @@ def _get_next_logits(body):
|
||||
state = process_parameters(body) if use_samplers else {}
|
||||
state['stream'] = True
|
||||
|
||||
return get_next_logits(body['prompt'], state, use_samplers, "", return_dict=True)
|
||||
return get_next_logits(body['prompt'], state, use_samplers, "", top_logits=body['top_logits'], return_dict=True)
|
||||
|
@ -1,6 +1,6 @@
|
||||
import json
|
||||
import time
|
||||
from typing import List
|
||||
from typing import Dict, List
|
||||
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
@ -156,6 +156,7 @@ class TokenCountResponse(BaseModel):
|
||||
class LogitsRequestParams(BaseModel):
|
||||
prompt: str
|
||||
use_samplers: bool = False
|
||||
top_logits: int | None = 50
|
||||
frequency_penalty: float | None = 0
|
||||
max_tokens: int | None = 16
|
||||
presence_penalty: float | None = 0
|
||||
@ -168,7 +169,7 @@ class LogitsRequest(GenerationOptions, LogitsRequestParams):
|
||||
|
||||
|
||||
class LogitsResponse(BaseModel):
|
||||
logits: dict
|
||||
logits: Dict[str, float]
|
||||
|
||||
|
||||
class ModelInfoResponse(BaseModel):
|
||||
|
@ -8,7 +8,7 @@ from modules.text_generation import generate_reply
|
||||
global_scores = None
|
||||
|
||||
|
||||
def get_next_logits(prompt, state, use_samplers, previous, return_dict=False):
|
||||
def get_next_logits(prompt, state, use_samplers, previous, top_logits=50, 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
|
||||
@ -50,8 +50,7 @@ def get_next_logits(prompt, state, use_samplers, previous, return_dict=False):
|
||||
scores = output['logits'][-1][-1]
|
||||
|
||||
probs = torch.softmax(scores, dim=-1, dtype=torch.float)
|
||||
topk_values, topk_indices = torch.topk(probs, k=50, largest=True, sorted=True)
|
||||
topk_values = [f"{float(i):.5f}" for i in topk_values]
|
||||
topk_values, topk_indices = torch.topk(probs, k=top_logits, largest=True, sorted=True)
|
||||
if is_non_hf_exllamav1 or is_non_hf_llamacpp:
|
||||
topk_indices = [i.expand((1, 1)) for i in topk_indices]
|
||||
|
||||
@ -61,12 +60,14 @@ def get_next_logits(prompt, state, use_samplers, previous, return_dict=False):
|
||||
tokens = [shared.tokenizer.decode(i) for i in topk_indices]
|
||||
|
||||
if return_dict:
|
||||
topk_values = [float(i) for i in topk_values]
|
||||
output = {}
|
||||
for row in list(zip(topk_values, tokens)):
|
||||
output[row[1]] = row[0]
|
||||
|
||||
return output
|
||||
else:
|
||||
topk_values = [f"{float(i):.5f}" for i in topk_values]
|
||||
output = ''
|
||||
for row in list(zip(topk_values, tokens)):
|
||||
output += f"{row[0]} - {repr(row[1])}\n"
|
||||
|
Loading…
Reference in New Issue
Block a user