Views
No views yet
1import torch
2import torch.nn as nn
3from transformers import BertForTokenClassification, AutoModel, AutoTokenizer
4from razdel import sentenize
5
6class BertForTextSegmentationEmbeddings(nn.Module):
7 def __init__(self, config, embeddings_dim=768):
8 super(BertForTextSegmentationEmbeddings, self).__init__()
9
10 self.config = config
11 self.position_embeddings = torch.nn.Embedding(config.max_position_embeddings, config.hidden_size)
12
13 self.LayerNorm = torch.nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
14 self.dropout = torch.nn.Dropout(config.hidden_dropout_prob)
15
16 def forward(self, inputs_embeds, position_ids=None, input_ids=None, token_type_ids=None, past_key_values_length=None):
17 input_shape = inputs_embeds.size()[:-1]
18 seq_length = input_shape[1]
19 device = inputs_embeds.device
20
21 assert seq_length <= self.config.max_position_embeddings, \
22 f"Too long sequence is passed {seq_length}. Maximum allowed sequence length is {self.config.max_position_embeddings}"
23
24 if position_ids is None:
25 position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
26 position_ids = position_ids.unsqueeze(0).expand(input_shape)
27
28 position_embeddings = self.position_embeddings(position_ids)
29
30 embeddings = inputs_embeds + position_embeddings
31 embeddings = self.LayerNorm(embeddings)
32 embeddings = self.dropout(embeddings)
33 return embeddings
34
35class BertForTextSegmentation(BertForTokenClassification):
36 def __init__(self, config):
37 super(BertForTextSegmentation, self).__init__(config)
38
39 self.bert.base_model.embeddings = BertForTextSegmentationEmbeddings(config)
40
41 self.init_weights()
42
43def max_pooling(model_output, attention_mask):
44 token_embeddings = model_output[0] #First element of model_output contains all token embeddings
45 input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
46 token_embeddings[input_mask_expanded == 0] = -1e9 # Set padding tokens to large negative value
47 return torch.max(token_embeddings, 1)[0]
48
49def create_embeddings(sentences, tokenizer, model):
50 # Tokenize sentences
51 encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
52 # Compute token embeddings
53 with torch.no_grad():
54 model_output = model(**encoded_input.to(device))
55 # Perform pooling. In this case, max pooling.
56 sentence_embeddings = max_pooling(model_output, encoded_input['attention_mask'])
57
58 return sentence_embeddings
59
60device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
61
62emb_tokenizer = AutoTokenizer.from_pretrained("ai-forever/sbert_large_nlu_ru")
63emb_model = AutoModel.from_pretrained("ai-forever/sbert_large_nlu_ru")
64model = BertForTextSegmentation.from_pretrained("mlenjoyneer/RuTextSegModel")
65
66emb_model.to(device)
67model.to(device)
68
69text = """В Норильске за годы работы телефона доверия консультанты приняли в общей сложности порядка 75 тысяч обращений, сообщает «Заполярная Правда». Служба психологической помощи появилась в 2000 году. Руководитель службы профилактики наркомании Елена Слатвицкая рассказала журналистам, что в Заполярье настал период, когда ухудшается психо– эмоциональное состояние населения. Это происходит на входе в полярную ночь и на выходе из нее. Осень является кризисным моментом. Сейчас на телефоне доверия работают 15 специалистов. Каждый — под своим псевдонимом. Тему беседы определяет звонящий. Это могут быть наркомания и алкоголизм, ВИЧ–инфекция и прочие заболевания и зависимости, кризисы семейных отношений и многое другое. Сотрудники службы отметили, что больше стало звонков по поводу суицидальных намерений. Наибольшее количество обращений по суицидам пришлось на октябрь — ноябрь. Много звонков как от мужчин, так и от женщин с вопросами об одиночестве. Лидерами по количеству обращений пока остаются женщины. В сентябре в Норильске обнаружили тело девятиклассницы. По версии следствия, девочка сбросилась с крыши. В январе подросток нанес себе порезы стеклом от разбитой бутылки, пытаясь покончить с собой. Мальчик поссорился с матерью и в ходе ссоры нанес себе несколько порезов. Проводится расследование."""
70
71input_embeds = create_embeddings([s.text for s in sentenize(text)], emb_tokenizer, emb_model).unsqueeze(0)
72outputs = model(inputs_embeds=input_embeds)
73
74logits = outputs.logits.cpu()
75preds = logits.argmax(axis=2).tolist()[0] # [0, 0, 1, 0, 0, 0, 0, 1, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0]
76
77# true_labels = [0, 0, 1, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 1, 0, 0, 0, 0]| Train Dataset | Test Dataset | F1_total | F1_1 | Pk | Pk_5 | WinDiff | WinDiff_5 |
|---|---|---|---|---|---|---|---|
| News+Wiki | News | 0.88 | 0.80 | 0.16 | 0.11 | 0.20 | 0.35 |
| News+Wiki | Wiki | 0.89 | 0.80 | 0.18 | 0.16 | 0.09 | 0.19 |
In progress