Views
No views yet
1import torch
2import numpy as np
3from roformer import RoFormerForCausalLM, RoFormerConfig
4from transformers import BertTokenizer
5
6device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
7pretrained_model = "junnyu/roformer_chinese_sim_char_base"
8tokenizer = BertTokenizer.from_pretrained(pretrained_model)
9config = RoFormerConfig.from_pretrained(pretrained_model)
10config.is_decoder = True
11config.eos_token_id = tokenizer.sep_token_id
12config.pooler_activation = "linear"
13model = RoFormerForCausalLM.from_pretrained(pretrained_model, config=config)
14model.to(device)
15model.eval()
16
17def gen_synonyms(text, n=100, k=20):
18 ''''含义： 产生sent的n个相似句，然后返回最相似的k个。
19 做法：用seq2seq生成，并用encoder算相似度并排序。
20 '''
21 # 寻找所有相似的句子
22 r = []
23 inputs1 = tokenizer(text, return_tensors="pt")
24 for _ in range(n):
25 inputs1.to(device)
26 output = tokenizer.batch_decode(model.generate(**inputs1, top_p=0.95, do_sample=True, max_length=128), skip_special_tokens=True)[0].replace(" ","").replace(text, "") # 去除空格，去除原始text文本。
27 r.append(output)
28
29 # 对相似的句子进行排序
30 r = [i for i in set(r) if i != text and len(i) > 0]
31 r = [text] + r
32 inputs2 = tokenizer(r, padding=True, return_tensors="pt")
33 with torch.no_grad():
34 inputs2.to(device)
35 outputs = model(**inputs2)
36 Z = outputs.pooler_output.cpu().numpy()
37 Z /= (Z**2).sum(axis=1, keepdims=True)**0.5
38 argsort = np.dot(Z[1:], -Z[0]).argsort()
39
40 return [r[i + 1] for i in argsort[:k]]
41
42out = gen_synonyms("广州和深圳哪个好？")
43print(out)
44# ['深圳和广州哪个好？',
45# '广州和深圳哪个好',
46# '深圳和广州哪个好',
47# '深圳和广州哪个比较好。',
48# '深圳和广州哪个最好？',
49# '深圳和广州哪个比较好',
50# '广州和深圳那个比较好',
51# '深圳和广州哪个更好？',
52# '深圳与广州哪个好',
53# '深圳和广州，哪个比较好',
54# '广州与深圳比较哪个好',
55# '深圳和广州哪里比较好',
56# '深圳还是广州比较好？',
57# '广州和深圳哪个地方好一些？',
58# '广州好还是深圳好？',
59# '广州好还是深圳好呢？',
60# '广州与深圳哪个地方好点？',
61# '深圳好还是广州好',
62# '广州好还是深圳好',
63# '广州和深圳哪个城市好？']