mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-10-01 01:26:03 -04:00
Add /v1/internal/logits endpoint (#4650)
This commit is contained in:
parent
8f4f4daf8b
commit
0fa1af296c
@ -97,6 +97,29 @@ curl http://127.0.0.1:5000/v1/chat/completions \
|
||||
}'
|
||||
```
|
||||
|
||||
#### Logits
|
||||
|
||||
```
|
||||
curl -k http://127.0.0.1:5000/v1/internal/logits \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "Who is best, Asuka or Rei? Answer:",
|
||||
"use_samplers": false
|
||||
}'
|
||||
```
|
||||
|
||||
#### Logits after sampling parameters
|
||||
|
||||
```
|
||||
curl -k http://127.0.0.1:5000/v1/internal/logits \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "Who is best, Asuka or Rei? Answer:",
|
||||
"use_samplers": true,
|
||||
"top_k": 3
|
||||
}'
|
||||
```
|
||||
|
||||
#### Python chat example
|
||||
|
||||
```python
|
||||
|
@ -16,6 +16,7 @@ from sse_starlette import EventSourceResponse
|
||||
import extensions.openai.completions as OAIcompletions
|
||||
import extensions.openai.embeddings as OAIembeddings
|
||||
import extensions.openai.images as OAIimages
|
||||
import extensions.openai.logits as OAIlogits
|
||||
import extensions.openai.models as OAImodels
|
||||
import extensions.openai.moderations as OAImoderations
|
||||
from extensions.openai.errors import ServiceUnavailableError
|
||||
@ -38,6 +39,8 @@ from .typing import (
|
||||
EncodeRequest,
|
||||
EncodeResponse,
|
||||
LoadModelRequest,
|
||||
LogitsRequest,
|
||||
LogitsResponse,
|
||||
ModelInfoResponse,
|
||||
TokenCountResponse,
|
||||
to_dict
|
||||
@ -242,6 +245,16 @@ async def handle_token_count(request_data: EncodeRequest):
|
||||
return JSONResponse(response)
|
||||
|
||||
|
||||
@app.post("/v1/internal/logits", response_model=LogitsResponse, dependencies=check_key)
|
||||
async def handle_logits(request_data: LogitsRequest):
|
||||
'''
|
||||
Given a prompt, returns the top 50 most likely logits as a dict.
|
||||
The keys are the tokens, and the values are the probabilities.
|
||||
'''
|
||||
response = OAIlogits._get_next_logits(to_dict(request_data))
|
||||
return JSONResponse(response)
|
||||
|
||||
|
||||
@app.post("/v1/internal/stop-generation", dependencies=check_key)
|
||||
async def handle_stop_generation(request: Request):
|
||||
stop_everything_event()
|
||||
|
@ -126,15 +126,15 @@ class EncodeRequest(BaseModel):
|
||||
text: str
|
||||
|
||||
|
||||
class DecodeRequest(BaseModel):
|
||||
tokens: List[int]
|
||||
|
||||
|
||||
class EncodeResponse(BaseModel):
|
||||
tokens: List[int]
|
||||
length: int
|
||||
|
||||
|
||||
class DecodeRequest(BaseModel):
|
||||
tokens: List[int]
|
||||
|
||||
|
||||
class DecodeResponse(BaseModel):
|
||||
text: str
|
||||
|
||||
@ -143,6 +143,24 @@ class TokenCountResponse(BaseModel):
|
||||
length: int
|
||||
|
||||
|
||||
class LogitsRequestParams(BaseModel):
|
||||
prompt: str
|
||||
use_samplers: bool = False
|
||||
frequency_penalty: float | None = 0
|
||||
max_tokens: int | None = 16
|
||||
presence_penalty: float | None = 0
|
||||
temperature: float | None = 1
|
||||
top_p: float | None = 1
|
||||
|
||||
|
||||
class LogitsRequest(GenerationOptions, LogitsRequestParams):
|
||||
pass
|
||||
|
||||
|
||||
class LogitsResponse(BaseModel):
|
||||
logits: dict
|
||||
|
||||
|
||||
class ModelInfoResponse(BaseModel):
|
||||
model_name: str
|
||||
lora_names: List[str]
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user