13fromtypingimportList1415importtorch16fromtransformersimportBertTokenizer,BertModel1718fromlabmlimportlab,monit

21classBERTChunkEmbeddings:

29def\_\_init\_\_(self,device:torch.device):30self.device=device

33withmonit.section('Load BERT tokenizer'):34self.tokenizer=BertTokenizer.from\_pretrained('bert-base-uncased',35cache\_dir=str(36lab.get\_data\_path()/'cache'/'bert-tokenizer'))

39withmonit.section('Load BERT model'):40self.model=BertModel.from\_pretrained("bert-base-uncased",41cache\_dir=str(lab.get\_data\_path()/'cache'/'bert-model'))

44self.model.to(device)

46@staticmethod47def\_trim\_chunk(chunk:str):

61stripped=chunk.strip()

63parts=stripped.split()

65stripped=stripped[len(parts[0]):-len(parts[-1])]

68stripped=stripped.strip()

71ifnotstripped:72returnchunk

74else:75returnstripped

77def\_\_call\_\_(self,chunks:List[str]):

83withtorch.no\_grad():

85trimmed\_chunks=[self.\_trim\_chunk(c)forcinchunks]

88tokens=self.tokenizer(trimmed\_chunks,return\_tensors='pt',add\_special\_tokens=False,padding=True)

91input\_ids=tokens['input\_ids'].to(self.device)92attention\_mask=tokens['attention\_mask'].to(self.device)93token\_type\_ids=tokens['token\_type\_ids'].to(self.device)

95output=self.model(input\_ids=input\_ids,96attention\_mask=attention\_mask,97token\_type\_ids=token\_type\_ids)

100state=output['last\_hidden\_state']

104emb=(state\*attention\_mask[:,:,None]).sum(dim=1)/attention\_mask[:,:,None].sum(dim=1)

107returnemb

110def\_test():

114fromlabml.loggerimportinspect

117device=torch.device('cuda:0')118bert=BERTChunkEmbeddings(device)

121text=["Replace me by any text you'd like.",122"Second sentence"]

125encoded\_input=bert.tokenizer(text,return\_tensors='pt',add\_special\_tokens=False,padding=True)126127inspect(encoded\_input,\_expand=True)

130output=bert.model(input\_ids=encoded\_input['input\_ids'].to(device),131attention\_mask=encoded\_input['attention\_mask'].to(device),132token\_type\_ids=encoded\_input['token\_type\_ids'].to(device))133134inspect({'last\_hidden\_state':output['last\_hidden\_state'],135'pooler\_output':output['pooler\_output']},136\_expand=True)

139inspect(bert.tokenizer.convert\_ids\_to\_tokens(encoded\_input['input\_ids'][0]),\_n=-1)140inspect(bert.tokenizer.convert\_ids\_to\_tokens(encoded\_input['input\_ids'][1]),\_n=-1)

143inspect(bert(text))

147if\_\_name\_\_=='\_\_main\_\_':148\_test()