Views
No views yet
1
2
3
4import toolz
5import torch
6batch_size = 6
7
8def roberta_similarity_batches(to_predict):
9 batches = toolz.partition(batch_size, to_predict)
10 similarity_scores = []
11 for batch in batches:
12 sentences = [(sentence_similarity["sent1"], sentence_similarity["sent2"]) for sentence_similarity in batch]
13 batch_scores = similarity_roberta(model, tokenizer,sentences)
14 similarity_scores = similarity_scores + batch_scores[0].cpu().squeeze(axis=1).tolist()
15 return similarity_scores
16
17def similarity_roberta(model, tokenizer, sent_pairs):
18 batch_token = tokenizer(sent_pairs, padding='max_length', truncation=True, max_length=500)
19 res = model(torch.tensor(batch_token['input_ids']).cuda(), attention_mask=torch.tensor(batch_token["attention_mask"]).cuda())
20 return res
21
22similarity_roberta(model, tokenizer, [('NEW YORK--(BUSINESS WIRE)--Rosen Law Firm, a global investor rights law firm, announces it is investigating potential securities claims on behalf of shareholders of Vale S.A. ( VALE ) resulting from allegations that Vale may have issued materially misleading business information to the investing public',
23 'EQUITY ALERT: Rosen Law Firm Announces Investigation of Securities Claims Against Vale S.A. – VALE')])
24