AI/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2024-09-19 15:56:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add a simple logit viewer (#3636)

Browse Source
This commit is contained in:
oobabooga 2023-08-20 20:49:21 -03:00 committed by GitHub
parent 2c1fd0d72b
commit 120fb86c6a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
4 changed files with 45 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
css/main.css
View File

@ -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
View 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

7
modules/ui_default.py
View File

@ -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'))

7
modules/ui_notebook.py
View File

@ -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'))