File size: 1,873 Bytes
3c3eabb |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 |
import gradio as gr
from transformers import GPT2Tokenizer, AutoModelForCausalLM
import numpy as np
tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
model = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer.pad_token_id = tokenizer.eos_token_id
# if prob > x, then label = y; sorted in descending probability order
probs_to_label = [
(0.1, "p >= 10%"),
(0.01, "p >= 1%"),
(1e-20, "p < 1%"),
]
label_to_color = {
"p >= 10%": "green",
"p >= 1%": "yellow",
"p < 1%": "red"
}
def get_tokens_and_scores(prompt):
inputs = tokenizer([prompt], return_tensors="pt")
outputs = model.generate(**inputs, max_new_tokens=50, return_dict_in_generate=True, output_scores=True, do_sample=True)
transition_scores = model.compute_transition_scores(
outputs.sequences, outputs.scores, normalize_logits=True
)
transition_proba = np.exp(transition_scores)
input_length = 1 if model.config.is_encoder_decoder else inputs.input_ids.shape[1]
generated_tokens = outputs.sequences[:, input_length:]
highlighted_out = [(tokenizer.decode(token), None) for token in inputs.input_ids]
for token, proba in zip(generated_tokens[0], transition_proba[0]):
this_label = None
assert 0. <= proba <= 1.0
for min_proba, label in probs_to_label:
if proba >= min_proba:
this_label = label
break
highlighted_out.append((tokenizer.decode(token), this_label))
return highlighted_out
demo = gr.Interface(
get_tokens_and_scores,
[
gr.Textbox(
label="Prompt",
lines=3,
value="Today is",
),
],
gr.HighlightedText(
label="Highlighted generation",
combine_adjacent=True,
show_legend=True,
).style(color_map=label_to_color),
)
if __name__ == "__main__":
demo.launch()
|