1pip install git+https://github.com/JunnYu/GAU-alpha-pytorch.git
2or
3pip install gau_alpha
1import torch
2
3from gau_alpha import GAUAlphaForMaskedLM, GAUAlphaTokenizer
4
5text = "今天[MASK]很好,我[MASK]去公园玩。"
6tokenizer = GAUAlphaTokenizer.from_pretrained(
7 "junnyu/chinese_GAU-alpha-char_L-24_H-768"
8)
9pt_model = GAUAlphaForMaskedLM.from_pretrained(
10 "junnyu/chinese_GAU-alpha-char_L-24_H-768"
11)
12
13pt_inputs = tokenizer(text, return_tensors="pt")
14
15with torch.no_grad():
16 pt_outputs = pt_model(**pt_inputs).logits[0]
17pt_outputs_sentence = "pytorch: "
18for i, id in enumerate(tokenizer.encode(text)):
19 if id == tokenizer.mask_token_id:
20 val, idx = pt_outputs[i].softmax(-1).topk(k=5)
21 tokens = tokenizer.convert_ids_to_tokens(idx)
22 new_tokens = []
23 for v, t in zip(val.cpu(), tokens):
24 new_tokens.append(f"{t}+{round(v.item(),4)}")
25 pt_outputs_sentence += "[" + "||".join(new_tokens) + "]"
26 else:
27 pt_outputs_sentence += "".join(
28 tokenizer.convert_ids_to_tokens([id], skip_special_tokens=True)
29 )
30print(pt_outputs_sentence)
31# pytorch: 今天[天+0.8657||气+0.0535||阳+0.0165||,+0.0126||晴+0.0111]很好,我[要+0.4619||想+0.4352||又+0.0252||就+0.0157||跑+0.0064]去公园玩。
1@techreport{gau-alpha,
2 title={GAU-α: GAU-based Transformers for NLP - ZhuiyiAI},
3 author={Jianlin Su, Shengfeng Pan, Bo Wen, Yunfeng Liu},
4 year={2022},
5 url="https://github.com/ZhuiyiTechnology/GAU-alpha",
6}