Add a simple logit viewer (#3636)

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;
 }

modules/logits.py

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

modules/ui_default.py

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

modules/ui_notebook.py

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