Views
No views yet
1from transformers import AutoModelForMaskedLM,AutoTokenizer
2import torch
3import numpy as np
4
5model = AutoModelForMaskedLM.from_pretrained("aken12/splade-japanese")
6tokenizer = AutoTokenizer.from_pretrained("aken12/splade-japanese")
7
8query = "筑波大学では何の研究が行われているか？"
9
10def encode_query(query, tokenizer, model):
11 encoded_input = tokenizer(query, return_tensors="pt")
12 with torch.no_grad():
13 output = model(**encoded_input, return_dict=True).logits
14 aggregated_output, _ = torch.max(torch.log(1 + torch.relu(output)) * encoded_input['attention_mask'].unsqueeze(-1), dim=1)
15 return aggregated_output
16
17def get_topk_tokens(reps, vocab_dict, topk):
18 topk_values, topk_indices = torch.topk(reps, topk, dim=1)
19 values = np.rint(topk_values.numpy() * 100).astype(int)
20 dict_splade = {vocab_dict[id_token.item()]: int(value_token) for id_token, value_token in zip(topk_indices[0], values[0]) if value_token > 0}
21 return dict_splade
22
23
24vocab_dict = {v: k for k, v in tokenizer.get_vocab().items()}
25topk = len(vocab_dict) // 1000
26
27
28model_output = encode_query(query, tokenizer, model)
29
30
31dict_splade = get_topk_tokens(model_output, vocab_dict, topk)
32
33
34for token, value in dict_splade.items():
35 print(token, value)に 250
が 248
は 247
の 247
、 244
と 240
を 239
。 239
も 238
で 237
から 221
や 219
な 206
筑波 204
( 204
・ 202
て 197
へ 191
にて 189
など 188
) 186
まで 184
た 182
この 171
- 170
「 170
より 166
その 165
: 163
」 161