Views
No views yet
title-paragraph separation task.wikitext1from transformers import XLMRobertaPreTrainedModel, XLMRobertaModel, AutoTokenizer
2from nltk.tokenize import line_tokenize
3import torch
4import torch.nn as nn
5import torch.nn.functional as F
6from torch.utils.data import DataLoader
7from datasets import Dataset
8
9
10#Utility Functions
11
12
13def get_default_device():
14 if torch.cuda.is_available():
15 return torch.device('cuda')
16 elif torch.backends.mps.is_available():
17 return torch.device('mps')
18 else:
19 return torch.device('cpu')
20
21def to_device(data, device):
22 if isinstance(data, (list,tuple)):
23 return [to_device(x, device) for x in data]
24 elif isinstance(data, dict):
25 return {'input_ids':to_device(data['input_ids'],device),'attention_mask':to_device(data['attention_mask'],device)}
26 return data.to(device)
27
28class DeviceDataLoader():
29 def __init__(self, dl, device):
30 self.dl = dl
31 self.device = device
32
33 def __iter__(self):
34 for b in self.dl:
35 yield to_device(b, self.device)
36
37 def __len__(self):
38 return len(self.dl)
39
40class IsoBN(nn.Module):
41 def __init__(self, hidden_size):
42 """Init method"""
43 super().__init__()
44 self.register_parameter(name='cov', param=torch.nn.Parameter(torch.zeros(hidden_size, hidden_size)))
45 self.register_parameter(name='std', param=torch.nn.Parameter(torch.zeros(hidden_size)))
46
47 self.cov.requires_grad = False
48 self.std.requires_grad = False
49
50 def forward(self, input, momentum: float = 0.05, eps: float = 1e-3, beta: float = 0.5):
51 """Forward method"""
52 if self.training:
53 x = input.detach()
54 n = x.size(0)
55 mean = x.mean(dim=0)
56 y = x - mean.unsqueeze(0)
57 std = (y ** 2).mean(0) ** 0.5
58 cov = (y.t() @ y) / n
59 self.cov.data += momentum * (cov.data - self.cov.data)
60 self.std.data += momentum * (std.data - self.std.data)
61 corr = torch.clamp(self.cov / torch.ger(self.std, self.std), -1, 1)
62 gamma = (corr ** 2).mean(1)
63 denorm = (gamma * self.std)
64 scale = 1 / (denorm + eps) ** beta
65 E = torch.diag(self.cov).sum()
66 new_E = (torch.diag(self.cov) * (scale ** 2)).sum()
67 m = (E / (new_E + eps)) ** 0.5
68 scale *= m
69 return input * scale.unsqueeze(0).detach()
70
71class e5_base_CTSEG(XLMRobertaPreTrainedModel):
72
73 def __init__(self, config):
74
75 super().__init__(config)
76
77 self.e5 = XLMRobertaModel(config).from_pretrained('intfloat/multilingual-e5-base')
78 self.dropout = nn.Dropout(0.5)
79 self.linear_1 = nn.Linear(768,256)
80 self.linear_2 = nn.Linear(256,128)
81 self.linear_3 = nn.Linear(128,2)
82 self.relu = nn.ReLU()
83 self.isobn = IsoBN(768)
84
85 def forward(self, sent):
86
87 sent['input_ids'] = sent['input_ids'].reshape(sent['input_ids'].shape[0],-1)
88 sent['attention_mask'] = sent['attention_mask'].reshape(sent['attention_mask'].shape[0],-1)
89
90 hs= self.e5(input_ids=sent['input_ids'], attention_mask=sent['attention_mask'])
91 cls_hs = hs.last_hidden_state[:, 0]
92 cls_hs = self.isobn(cls_hs)
93
94
95 out = self.linear_1(cls_hs)
96 out = self.relu(out)
97 out = self.dropout(out)
98 out = self.linear_2(out)
99 out = self.relu(out)
100 out = self.dropout(out)
101 out = self.linear_3(out)
102
103
104 return out
105
106
107 def training_step(self, sent, labels):
108 out = self.forward(sent)
109 loss = F.cross_entropy(out, labels)
110 return loss
111
112 def validation_step(self, sent, labels):
113 out = self.forward(sent)
114 loss = F.cross_entropy(out, labels)
115 acc = accuracy(out, labels)
116 return {'val_acc':acc,'val_loss':loss.detach()}
117
118 def validation_epoch_end(self, metrics):
119 batch_losses = [x['val_loss'] for x in metrics]
120 batch_accs = [x['val_acc'] for x in metrics]
121
122 epoch_loss = torch.stack(batch_losses).mean().item()
123 epoch_acc = torch.stack(batch_accs).mean().item()
124
125 return {'val_loss':epoch_loss, 'val_acc':epoch_acc}
126
127 def epoch_end(self, epoch, result):
128
129 print("Epoch [{}], train_loss: {:.4f}, val_loss: {:.4f}, val_acc: {:.4f}".format(
130 epoch, result['train_loss'], result['val_loss'], result['val_acc']))
131
132 def evaluate(self, val_loader):
133 self.eval()
134 metrics = [self.validation_step(sent,labels.type(torch.LongTensor).to(device, non_blocking=True)) for sent,labels in val_loader]
135 return self.validation_epoch_end(metrics)
136def accuracy(out, labels):
137 return (out.argmax(dim=1) == labels).sum()/labels.numel()
1381tokenizer = AutoTokenizer.from_pretrained('intfloat/multilingual-e5-base')
2model = e5_base_CTSEG.from_pretrained('ProfessorBob/title-par-segmentation')
3device = get_default_device()
4to_device(model,device)1def infer_block(
2 chunks,
3 batch_size: int = 8,
4 return_probability: bool = False,
5 tokenizer = tokenizer
6):
7 """ Bulk Infer function"""
8
9 tok_text_bulk = tokenizer(
10 ['query: ' + sent[0] +'[SEP]'+ sent[1] for sent in chunks],
11 padding='max_length',
12 truncation=True,
13 return_tensors='pt'
14 )
15 sentences = Dataset.from_dict({
16 'input_ids': tok_text_bulk['input_ids'],
17 'attention_mask': tok_text_bulk['attention_mask']
18 })
19 sentences.set_format(
20 'torch',
21 columns=['input_ids','attention_mask']
22 )
23 sentences = DataLoader(
24 sentences,
25 batch_size=batch_size,
26 pin_memory=True
27 )
28 sentences = DeviceDataLoader(sentences, device)
29 preds = list()
30 model.eval()
31 with torch.no_grad():
32 for i, batch in enumerate(sentences):
33 out = model(batch)
34 if return_probability:
35 preds.extend((out.softmax(dim=1).cpu()[:, 1]).tolist())
36 else:
37 preds.extend(out.argmax(dim=1).cpu().tolist())
38
39 if device == torch.device('cuda'):
40 torch.cuda.empty_cache()
41 assert len(preds) == len(chunks)
42
43 return preds, out
44
45def segmentation_pipeline(text):
46
47 block = line_tokenize(text)
48 chunks = [
49 (u, v) for u, v in zip(block[:-1], block[1:])
50 ]
51 preds, out = infer_block(chunks,return_probability=False)
52 cut_idx = [i+1 for i, value in enumerate(preds) if value == 1]
53 cut_idx = [0]+cut_idx+[len(block)]
54 seg = [block[cut_idx[i]:cut_idx[i+1]] for i in range(len(cut_idx)-1)]
55
56 return seg1
2mixed_string = """
3
4Ancient Foundations (3000 BCE - 600 CE)
5
6In the dawn of human civilization, mathematics emerged as an essential tool for commerce, construction, and astronomy. Explore the mathematical innovations of ancient cultures such as the Babylonians, Egyptians, and Greeks, laying the groundwork for numerical systems, geometry, and the Pythagorean theorem.
7
8The Golden Age of Islamic Mathematics (700 CE - 1300 CE)
9
10Delve into the intellectual flourishing during the Islamic Golden Age, where scholars like Al-Khwarizmi and Omar Khayyam made groundbreaking contributions to algebra, trigonometry, and the development of algorithms. Discover how these advancements paved the way for the Renaissance in Europe.
11
12"""
131Block 1
2-----
3Ancient Foundations (3000 BCE - 600 CE)
4
5Block 2
6-----
7In the dawn of human civilization, mathematics emerged as an essential tool for commerce, construction, and astronomy. Explore the mathematical innovations of ancient cultures such as the Babylonians, Egyptians, and Greeks, laying the groundwork for numerical systems, geometry, and the Pythagorean theorem.
8
9Block 3
10-----
11The Golden Age of Islamic Mathematics (700 CE - 1300 CE)
12
13Block 4
14-----
15Delve into the intellectual flourishing during the Islamic Golden Age, where scholars like Al-Khwarizmi and Omar Khayyam made groundbreaking contributions to algebra, trigonometry, and the development of algorithms. Discover how these advancements paved the way for the Renaissance in Europe.
16