පෙළකැබලි වල බර්ට් කාවැද්දීම්

13from typing import List
14
15import torch
16from transformers import BertTokenizer, BertModel
17
18from labml import lab, monit

21class BERTChunkEmbeddings:

29    def __init__(self, device: torch.device):
30        self.device = device

33        with monit.section('Load BERT tokenizer'):
34            self.tokenizer = BertTokenizer.from_pretrained('bert-base-uncased',
35                                                           cache_dir=str(
36                                                               lab.get_data_path() / 'cache' / 'bert-tokenizer'))

39        with monit.section('Load BERT model'):
40            self.model = BertModel.from_pretrained("bert-base-uncased",
41                                                   cache_dir=str(lab.get_data_path() / 'cache' / 'bert-model'))

44            self.model.to(device)

46    @staticmethod
47    def _trim_chunk(chunk: str):

61        stripped = chunk.strip()

63        parts = stripped.split()

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

68        stripped = stripped.strip()

71        if not stripped:
72            return chunk

74        else:
75            return stripped

77    def __call__(self, chunks: List[str]):

83        with torch.no_grad():

85            trimmed_chunks = [self._trim_chunk(c) for c in chunks]

88            tokens = self.tokenizer(trimmed_chunks, return_tensors='pt', add_special_tokens=False, padding=True)

91            input_ids = tokens['input_ids'].to(self.device)
92            attention_mask = tokens['attention_mask'].to(self.device)
93            token_type_ids = tokens['token_type_ids'].to(self.device)

95            output = self.model(input_ids=input_ids,
96                                attention_mask=attention_mask,
97                                token_type_ids=token_type_ids)

100            state = output['last_hidden_state']

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

107            return emb

110def _test():

114    from labml.logger import inspect

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

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

125    encoded_input = bert.tokenizer(text, return_tensors='pt', add_special_tokens=False, padding=True)
126
127    inspect(encoded_input, _expand=True)

130    output = bert.model(input_ids=encoded_input['input_ids'].to(device),
131                        attention_mask=encoded_input['attention_mask'].to(device),
132                        token_type_ids=encoded_input['token_type_ids'].to(device))
133
134    inspect({'last_hidden_state': output['last_hidden_state'],
135             'pooler_output': output['pooler_output']},
136            _expand=True)

139    inspect(bert.tokenizer.convert_ids_to_tokens(encoded_input['input_ids'][0]), _n=-1)
140    inspect(bert.tokenizer.convert_ids_to_tokens(encoded_input['input_ids'][1]), _n=-1)

143    inspect(bert(text))

147if __name__ == '__main__':
148    _test()