Mirrors/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-11-21 23:57:58 +01:00
Code Issues Packages Projects Releases Wiki Activity

[OpenAI Extension] Add 'max_logits' parameter in logits endpoint (#4916)

Browse Source
This commit is contained in:
Kim Jaewon 2023-12-15 12:22:43 +09:00 committed by GitHub
parent eaa1fe67f3
commit e53f99faa0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 8 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
extensions/openai/logits.py
View File

@ -8,4 +8,4 @@ def _get_next_logits(body):
state = process_parameters(body) if use_samplers else {} state = process_parameters(body) if use_samplers else {}
state['stream'] = True 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)

5
extensions/openai/typing.py
View File

@ -1,6 +1,6 @@
import json import json
import time import time
from typing import List from typing import Dict, List
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@ -156,6 +156,7 @@ class TokenCountResponse(BaseModel):
class LogitsRequestParams(BaseModel): class LogitsRequestParams(BaseModel):
prompt: str prompt: str
use_samplers: bool = False use_samplers: bool = False
top_logits: int | None = 50
frequency_penalty: float | None = 0 frequency_penalty: float | None = 0
max_tokens: int | None = 16 max_tokens: int | None = 16
presence_penalty: float | None = 0 presence_penalty: float | None = 0
@ -168,7 +169,7 @@ class LogitsRequest(GenerationOptions, LogitsRequestParams):
class LogitsResponse(BaseModel): class LogitsResponse(BaseModel):
logits: dict logits: Dict[str, float]
class ModelInfoResponse(BaseModel): class ModelInfoResponse(BaseModel):

7
modules/logits.py
View File

@ -8,7 +8,7 @@ from modules.text_generation import generate_reply
global_scores = None 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: if shared.model is None:
logger.error("No model is loaded! Select one in the Model tab.") 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 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] scores = output['logits'][-1][-1]
probs = torch.softmax(scores, dim=-1, dtype=torch.float) probs = torch.softmax(scores, dim=-1, dtype=torch.float)
topk_values, topk_indices = torch.topk(probs, k=50, largest=True, sorted=True) topk_values, topk_indices = torch.topk(probs, k=top_logits, largest=True, sorted=True)
topk_values = [f"{float(i):.5f}" for i in topk_values]
if is_non_hf_exllamav1 or is_non_hf_llamacpp: if is_non_hf_exllamav1 or is_non_hf_llamacpp:
topk_indices = [i.expand((1, 1)) for i in topk_indices] 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] tokens = [shared.tokenizer.decode(i) for i in topk_indices]
if return_dict: if return_dict:
topk_values = [float(i) for i in topk_values]
output = {} output = {}
for row in list(zip(topk_values, tokens)): for row in list(zip(topk_values, tokens)):
output[row[1]] = row[0] output[row[1]] = row[0]
return output return output
else: else:
topk_values = [f"{float(i):.5f}" for i in topk_values]
output = '' output = ''
for row in list(zip(topk_values, tokens)): for row in list(zip(topk_values, tokens)):
output += f"{row[0]} - {repr(row[1])}\n" output += f"{row[0]} - {repr(row[1])}\n"