Views
No views yet
1tokenizer = AutoTokenizer.from_pretrained('Gigon/ruBert_for_construction_resourse_classification')
2model = BertForSequenceClassification.from_pretrained('Gigon/ruBert_for_construction_resourse_classification', num_labels=2893)
3
4def evaluate(text, model, tokenizer, device):
5 with torch.no_grad():
6 model.eval()
7 model.to(device)
8 text_tokenized = tokenizer(text, return_tensors='pt')
9 output = model(text_tokenized['input_ids'], attention_mask=text_tokenized['attention_mask'], token_type_ids=text_tokenized['token_type_ids'])
10 return output
11
12def classify(text, n, KSR, model, tokenizer, device):
13 output = evaluate(text, model, tokenizer, device)
14 output = torch.nn.functional.normalize(output.logits)
15 result = sorted(list(zip(KSR['code_KSR'], KSR['name_KSR'], output.numpy()[0][1:])), key=lambda item: -item[2])[:n]
16 return result
17
18def cousine_sim(sentence, candidate, model, tokenizer, device):
19 output_sentence = evaluate(sentence, model.bert, tokenizer, device)
20 sentence_emb = torch.nn.functional.normalize(output_sentence.last_hidden_state[:,0,:])
21 output_candidate = evaluate(candidate, model.bert, tokenizer, device)
22 candidate_emb = torch.nn.functional.normalize(output_candidate.last_hidden_state[:,0,:])
23 result = torch.matmul(sentence_emb[0], candidate_emb[0])
24 return result
25
26print(classify('мастика из битума', 3, KSR, model, tokenizer, 'cpu'))
27
28[('01.2.03.03', 'Группа 01.2.03.03: Мастики битумосодержащие', 0.19304055),
29 ('14.5.04.08', 'Группа 14.5.04.08: Мастики, не включенные в группы', 0.14259844),
30 ('01.5.01.01', 'Группа 01.5.01.01: Краски разметочные дорожные', 0.046304163)]
31
32print(classify('армокаркас из арматурной стали класса А500', 3, KSR, model, tokenizer, 'cpu'))
33
34[('08.4.02.03', 'Группа 08.4.02.03: Каркасы арматурные', 0.102737375),
35 ('08.4.02.04', 'Группа 08.4.02.04: Каркасы металлические', 0.055409335),
36 ('07.2.07.04', 'Группа 07.2.07.04: Конструкции индивидуальные сварные', 0.048963238)]
37
38print(cousine_sim('бетон класса B40', 'смесь бетонная тяжелая', model, tokenizer, 'cpu').item())
39
400.6447
41