token-importance-nb

import%20marimo%0A%0A__generated_with%20%3D%20%220.22.5%22%0Aapp%20%3D%20marimo.App(width%3D%22medium%22)%0A%0A%0A%40app.cell%0Aasync%20def%20_()%3A%0A%20%20%20%20import%20micropip%0A%20%20%20%20await%20micropip.install(%22transformers_js_py%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20import%20marimo%20as%20mo%0A%20%20%20%20import%20numpy%20as%20np%0A%20%20%20%20from%20transformers_js_py%20import%20import_transformers_js%0A%0A%20%20%20%20return%20import_transformers_js%2C%20mo%2C%20np%0A%0A%0A%40app.cell%0Aasync%20def%20_(import_transformers_js)%3A%0A%20%20%20%20transformers%20%3D%20await%20import_transformers_js()%0A%20%20%20%20return%20(transformers%2C)%0A%0A%0A%40app.cell%0Adef%20_(transformers)%3A%0A%20%20%20%20Tensor%20%3D%20transformers.Tensor%0A%20%20%20%20return%20(Tensor%2C)%0A%0A%0A%40app.cell%0Adef%20_(Tensor%2C%20np)%3A%0A%20%20%20%20def%20tensor_from_list(lst%2C%20dtype%3D%22int64%22)%3A%0A%20%20%20%20%20%20%20%20arr%20%3D%20np.array(lst%2C%20dtype%3Dnp.int64)%0A%20%20%20%20%20%20%20%20dims%20%3D%20list(arr.shape)%0A%20%20%20%20%20%20%20%20flat%20%3D%20arr.flatten()%0A%20%20%20%20%20%20%20%20return%20Tensor(dtype%2C%20flat%2C%20dims)%0A%0A%20%20%20%20return%20(tensor_from_list%2C)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20model_id%20%3D%20%22onnx-community%2FQwen2.5-0.5B-Instruct-ONNX%22%0A%20%20%20%20return%20(model_id%2C)%0A%0A%0A%40app.cell%0Aasync%20def%20_(model_id%2C%20transformers)%3A%0A%20%20%20%20AutoTokenizer%20%3D%20transformers.AutoTokenizer%0A%20%20%20%20tokenizer%20%3D%20await%20AutoTokenizer.from_pretrained(model_id)%0A%0A%20%20%20%20AutoModelForCausalLM%20%3D%20transformers.AutoModelForCausalLM%0A%20%20%20%20model%20%3D%20await%20AutoModelForCausalLM.from_pretrained(model_id)%0A%20%20%20%20return%20model%2C%20tokenizer%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%20Token%20importance%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20text%20%3D%20mo.ui.text_area(label%3D%22Input%20some%20text%3A%22%2C%20value%3D%22Question%3A%20What%20is%20the%20capital%20of%20France%5CnAnswer%3A%20The%20capital%20of%20France%20is%22%2C%20debounce%3DTrue%2C%20placeholder%3D%22Type%20some%20text...%22)%0A%20%20%20%20text%0A%20%20%20%20return%20(text%2C)%0A%0A%0A%40app.cell%0Adef%20_(text%2C%20tokenizer)%3A%0A%20%20%20%20inputs%20%3D%20tokenizer(text.value)%0A%20%20%20%20return%20(inputs%2C)%0A%0A%0A%40app.cell%0Adef%20_(inputs)%3A%0A%20%20%20%20input_ids%20%3D%20inputs%5B%22input_ids%22%5D.tolist()%0A%20%20%20%20return%20(input_ids%2C)%0A%0A%0A%40app.cell%0Adef%20_(input_ids)%3A%0A%20%20%20%20ids%20%3D%20input_ids%5B0%5D%0A%20%20%20%20return%20(ids%2C)%0A%0A%0A%40app.cell%0Adef%20_(ids%2C%20tokenizer)%3A%0A%20%20%20%20lst%20%3D%20%5Bids%5D%20%2B%20%5B%5Btokenizer.pad_token_id%5D%20%2B%20ids%5B%3Ai%5D%20%2B%20ids%5Bi%2B1%3A%5D%20for%20i%20in%20range(len(ids))%5D%0A%20%20%20%20attention_mask%20%3D%20%5B%5B1%5D%20*%20len(ids)%5D%20%2B%20%5B%5B0%5D%20%2B%20%5B1%5D%20*%20(len(ids)%20-%201)%20for%20_%20in%20range(len(ids))%5D%0A%20%20%20%20return%20attention_mask%2C%20lst%0A%0A%0A%40app.cell%0Aasync%20def%20_(attention_mask%2C%20lst%2C%20model%2C%20tensor_from_list)%3A%0A%20%20%20%20output%20%3D%20await%20model(input_ids%20%3D%20tensor_from_list(lst)%2C%20attention_mask%20%3D%20tensor_from_list(attention_mask))%0A%20%20%20%20return%20(output%2C)%0A%0A%0A%40app.cell%0Adef%20_(output)%3A%0A%20%20%20%20output_logits_lst%20%3D%20output%5B%22logits%22%5D.tolist()%0A%20%20%20%20return%20(output_logits_lst%2C)%0A%0A%0A%40app.cell%0Adef%20_(output_logits_lst)%3A%0A%20%20%20%20last_logits%20%3D%20output_logits_lst%5B0%5D%5B-1%5D%0A%20%20%20%20return%20(last_logits%2C)%0A%0A%0A%40app.cell%0Adef%20_()%3A%0A%20%20%20%20from%20scipy.special%20import%20softmax%0A%0A%20%20%20%20return%20(softmax%2C)%0A%0A%0A%40app.cell%0Adef%20_(last_logits%2C%20softmax)%3A%0A%20%20%20%20probs%20%3D%20softmax(last_logits)%0A%20%20%20%20return%20(probs%2C)%0A%0A%0A%40app.cell%0Adef%20_(np%2C%20probs)%3A%0A%20%20%20%20selected_token%20%3D%20np.argmax(probs).item()%0A%20%20%20%20return%20(selected_token%2C)%0A%0A%0A%40app.cell%0Adef%20_(probs%2C%20selected_token)%3A%0A%20%20%20%20base_prob%20%3D%20probs%5Bselected_token%5D.item()%0A%20%20%20%20return%20(base_prob%2C)%0A%0A%0A%40app.cell%0Adef%20_(base_prob%2C%20mo%2C%20selected_token%2C%20tokenizer)%3A%0A%20%20%20%20mo.Html(f'next%20token%20prediction%3A%20%3Cspan%20style%3D%22color%3A%20blue%3B%22%3E%7Btokenizer.decode(%5Bselected_token%5D)%7D%3C%2Fspan%3E%3Cbr%3Enext%20token%20probability%3A%20%7Bbase_prob%3A.3f%7D')%0A%20%20%20%20return%0A%0A%0A%40app.function%0Adef%20score_to_color(score%3A%20float)%20-%3E%20tuple%5Bint%2C%20int%2C%20int%5D%3A%0A%20%20%20%20%22%22%22Map%20score%20in%20%5B-1%2C%201%5D%20to%20(R%2C%20G%2C%20B)%20with%20red%3D-1%2C%20white%3D0%2C%20blue%3D1.%22%22%22%0A%20%20%20%20score%20%3D%20max(-1.0%2C%20min(1.0%2C%20score))%20%20%23%20clamp%0A%0A%20%20%20%20if%20score%20%3C%200%3A%0A%20%20%20%20%20%20%20%20%23%20white%20(255%2C255%2C255)%20%E2%86%92%20blue%20(0%2C0%2C255)%0A%20%20%20%20%20%20%20%20t%20%3D%20score%20%2B%201%20%20%20%20%20%20%20%20%20%20%23%200..1%0A%20%20%20%20%20%20%20%20r%2C%20g%2C%20b%20%3D%20int(255%20*%20t)%2C%20int(255%20*%20t)%2C%20255%0A%20%20%20%20else%3A%0A%20%20%20%20%20%20%20%20%23%20red%20(255%2C0%2C0)%20%E2%86%92%20white%20(255%2C255%2C255)%0A%20%20%20%20%20%20%20%20t%20%3D%20score%20%20%20%20%20%20%20%20%20%20%20%20%20%20%23%200..1%0A%20%20%20%20%20%20%20%20r%2C%20g%2C%20b%20%3D%20255%2C%20int(255%20*%20(1%20-%20t))%2C%20int(255%20*%20(1%20-%20t))%0A%0A%20%20%20%20return%20r%2C%20g%2C%20b%0A%0A%0A%40app.function%0Adef%20score_to_hex(score%3A%20float)%20-%3E%20str%3A%0A%20%20%20%20r%2C%20g%2C%20b%20%3D%20score_to_color(score)%0A%20%20%20%20return%20f%22%23%7Br%3A02x%7D%7Bg%3A02x%7D%7Bb%3A02x%7D%22%0A%0A%0A%40app.cell%0Adef%20_(base_prob%2C%20ids%2C%20output_logits_lst%2C%20selected_token%2C%20softmax%2C%20tokenizer)%3A%0A%20%20%20%20alo%20%3D%20%5B%5D%0A%20%20%20%20tokens_html%20%3D%20%22%22%0A%20%20%20%20for%20k%20in%20range(len(output_logits_lst)-1)%3A%0A%20%20%20%20%20%20%20%20token%20%3D%20tokenizer.decode(%5Bids%5Bk%5D%5D)%0A%20%20%20%20%20%20%20%20last_logits_loop%20%3D%20output_logits_lst%5Bk%2B1%5D%5B-1%5D%0A%20%20%20%20%20%20%20%20probs_loop%20%3D%20softmax(last_logits_loop)%0A%20%20%20%20%20%20%20%20alo.append(%7B%22token%22%3A%20token%2C%20%22score%22%3A%20f%22%7B-(probs_loop%5Bselected_token%5D.item()-base_prob)%3A.3f%7D%22%7D)%0A%20%20%20%20%20%20%20%20color%20%3D%20score_to_hex(probs_loop%5Bselected_token%5D.item()%20-%20base_prob)%0A%20%20%20%20%20%20%20%20if%20'%5Cn'%20in%20token%3A%0A%20%20%20%20%20%20%20%20%20%20%20%20token%20%3D%20token.replace(%22%5Cn%22%2C%20%22%3Cbr%3E%22)%0A%20%20%20%20%20%20%20%20token_html%20%3D%20f'%3Cspan%20style%3D%22background-color%3A%20%7Bcolor%7D%3B%22%3E%7Btoken%7D%3C%2Fspan%3E'%0A%20%20%20%20%20%20%20%20tokens_html%20%2B%3D%20token_html%0A%20%20%20%20return%20alo%2C%20tokens_html%0A%0A%0A%40app.cell%0Adef%20_(mo)%3A%0A%20%20%20%20mo.md(%22%22%22%0A%20%20%20%20%23%23%20Token%20importance%20(leave-one-out)%0A%20%20%20%20%22%22%22)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(mo%2C%20tokens_html)%3A%0A%20%20%20%20mo.Html(tokens_html)%0A%20%20%20%20return%0A%0A%0A%40app.cell%0Adef%20_(alo%2C%20mo)%3A%0A%20%20%20%20mo.ui.table(data%3Dalo%2C%20pagination%3DFalse%2C%20selection%3DNone)%0A%20%20%20%20return%0A%0A%0Aif%20__name__%20%3D%3D%20%22__main__%22%3A%0A%20%20%20%20app.run()%0A