Ask For the problem of Transformer-XL to get text embedding

[graph-0]

Ask For the problem of Transformer-XL to get text embedding

[GRAPH-0]

An example code

tokenizer = XLNetTokenizer.from_pretrained('xlnet-base-cased')
model = XLNetModel.from_pretrained('xlnet-base-cased',
                                   output_hidden_states=True,
                                   output_attentions=True).to(device)

lang_features = dict()
with open(args.input_file) as fin:
    # read your text file as fin; node_id + text
    lines = fin.readlines()
    for line in tqdm(lines[1:]):
        l = line.split('\t')
        node_id = l[0]
        input_ids = torch.tensor([tokenizer.encode(l[1])]).to(device)
        input_ids = input_ids[:, :args.max_word]
        all_hidden_states, all_attentions = model(input_ids)[-2:]
        rep = (all_hidden_states[-2][0] * all_attentions[-2][0].mean(dim=0).mean(dim=0).view(-1, 1)).sum(dim=0)
        lang_features[node_id] = rep.detach().cpu().numpy()

[eddy60506]

thx! I will try this code.