mirror of
https://github.com/oobabooga/text-generation-webui.git
synced 2024-12-25 13:58:56 +01:00
Add a simple logit viewer (#3636)
This commit is contained in:
parent
2c1fd0d72b
commit
120fb86c6a
15
css/main.css
15
css/main.css
@ -116,7 +116,20 @@ div.svelte-15lo0d8 > *, div.svelte-15lo0d8 > .form > * {
|
||||
height: calc(100dvh - 241px);
|
||||
}
|
||||
|
||||
.textbox_default textarea, .textbox_default_output textarea, .textbox textarea {
|
||||
.textbox_logits textarea {
|
||||
height: calc(100dvh - 241px);
|
||||
}
|
||||
|
||||
.textbox_logits_notebook textarea {
|
||||
height: calc(100dvh - 292px);
|
||||
}
|
||||
|
||||
.textbox_default textarea,
|
||||
.textbox_default_output textarea,
|
||||
.textbox_logits textarea,
|
||||
.textbox_logits_notebook textarea,
|
||||
.textbox textarea
|
||||
{
|
||||
font-size: 16px !important;
|
||||
color: #46464A !important;
|
||||
}
|
||||
|
19
modules/logits.py
Normal file
19
modules/logits.py
Normal file
@ -0,0 +1,19 @@
|
||||
import torch
|
||||
|
||||
from modules import shared
|
||||
|
||||
|
||||
def get_next_logits(prompt):
|
||||
tokens = shared.tokenizer.encode(prompt, return_tensors='pt').cuda()
|
||||
output = shared.model(input_ids=tokens)
|
||||
|
||||
scores = output['logits'][-1][-1]
|
||||
probs = torch.softmax(scores, dim=-1, dtype=torch.float)
|
||||
|
||||
topk_values, topk_indices = torch.topk(probs, k=20, largest=True, sorted=True)
|
||||
topk_values = [f"{float(i):.5f}" % i for i in topk_values]
|
||||
output = ''
|
||||
for row in list(zip(topk_values, shared.tokenizer.convert_ids_to_tokens(topk_indices))):
|
||||
output += f"{row[0]} {row[1]}\n"
|
||||
|
||||
return output
|
@ -1,6 +1,6 @@
|
||||
import gradio as gr
|
||||
|
||||
from modules import shared, ui, utils
|
||||
from modules import logits, shared, ui, utils
|
||||
from modules.prompts import count_tokens, load_prompt
|
||||
from modules.text_generation import (
|
||||
generate_reply_wrapper,
|
||||
@ -43,6 +43,10 @@ def create_ui():
|
||||
with gr.Tab('HTML'):
|
||||
shared.gradio['html-default'] = gr.HTML()
|
||||
|
||||
with gr.Tab('Logits'):
|
||||
shared.gradio['get_logits-default'] = gr.Button('Get next token probabilities')
|
||||
shared.gradio['logits-default'] = gr.Textbox(lines=23, label='Output', elem_classes=['textbox_logits', 'add_scrollbar'])
|
||||
|
||||
|
||||
def create_event_handlers():
|
||||
shared.gradio['Generate-default'].click(
|
||||
@ -80,3 +84,4 @@ def create_event_handlers():
|
||||
lambda: gr.update(visible=True), None, gradio('file_deleter'))
|
||||
|
||||
shared.gradio['count_tokens-default'].click(count_tokens, gradio('textbox-default'), gradio('status-default'), show_progress=False)
|
||||
shared.gradio['get_logits-default'].click(logits.get_next_logits, gradio('textbox-default'), gradio('logits-default'))
|
||||
|
@ -1,6 +1,6 @@
|
||||
import gradio as gr
|
||||
|
||||
from modules import shared, ui, utils
|
||||
from modules import logits, shared, ui, utils
|
||||
from modules.prompts import count_tokens, load_prompt
|
||||
from modules.text_generation import (
|
||||
generate_reply_wrapper,
|
||||
@ -27,6 +27,10 @@ def create_ui():
|
||||
with gr.Tab('HTML'):
|
||||
shared.gradio['html-notebook'] = gr.HTML()
|
||||
|
||||
with gr.Tab('Logits'):
|
||||
shared.gradio['get_logits-notebook'] = gr.Button('Get next token probabilities')
|
||||
shared.gradio['logits-notebook'] = gr.Textbox(lines=23, label='Output', elem_classes=['textbox_logits_notebook', 'add_scrollbar'])
|
||||
|
||||
with gr.Row():
|
||||
shared.gradio['Generate-notebook'] = gr.Button('Generate', variant='primary', elem_classes='small-button')
|
||||
shared.gradio['Stop-notebook'] = gr.Button('Stop', elem_classes='small-button', elem_id='stop')
|
||||
@ -83,3 +87,4 @@ def create_event_handlers():
|
||||
lambda: gr.update(visible=True), None, gradio('file_deleter'))
|
||||
|
||||
shared.gradio['count_tokens-notebook'].click(count_tokens, gradio('textbox-notebook'), gradio('status-notebook'), show_progress=False)
|
||||
shared.gradio['get_logits-notebook'].click(logits.get_next_logits, gradio('textbox-notebook'), gradio('logits-notebook'))
|
||||
|
Loading…
Reference in New Issue
Block a user