| 13 | |
| 14 | |
| 15 | class T5TextProcessingEngine: |
| 16 | def __init__(self, text_encoder, tokenizer, min_length=256, end_with_pad=False, add_special_tokens=False): |
| 17 | super().__init__() |
| 18 | |
| 19 | self.text_encoder = text_encoder.transformer |
| 20 | self.tokenizer = tokenizer |
| 21 | |
| 22 | self.min_length = min_length |
| 23 | self.end_with_pad = end_with_pad |
| 24 | self.add_special_tokens = add_special_tokens |
| 25 | self.id_end = 1 |
| 26 | self.id_pad = 0 |
| 27 | |
| 28 | # vocab = self.tokenizer.get_vocab() |
| 29 | |
| 30 | # self.comma_token = vocab.get(',</w>', None) |
| 31 | |
| 32 | def tokenize(self, texts): |
| 33 | tokenized = self.tokenizer(texts, truncation=False, add_special_tokens=self.add_special_tokens)["input_ids"] |
| 34 | return tokenized |
| 35 | |
| 36 | def tokenize_for_UI(self, prompt): |
| 37 | parsed = parsing.parse_prompt_attention(prompt, "Ignore") |
| 38 | text = "".join([text for text, _ in parsed if text != "BREAK"]) |
| 39 | length = len(self.tokenizer(text, truncation=False, add_special_tokens=self.add_special_tokens)["input_ids"]) |
| 40 | if self.end_with_pad: |
| 41 | length += 1 |
| 42 | return 1 + length |
| 43 | |
| 44 | def _process_tokens(self, tokens): |
| 45 | attention_masks = [] |
| 46 | |
| 47 | for x in tokens: |
| 48 | attention_mask = [] |
| 49 | eos = False |
| 50 | |
| 51 | for y in x: |
| 52 | if isinstance(y, int): |
| 53 | attention_mask.append(0 if eos else 1) |
| 54 | if not eos and int(y) == self.id_end: |
| 55 | eos = True |
| 56 | |
| 57 | attention_masks.append(attention_mask) |
| 58 | |
| 59 | return torch.tensor(attention_masks, dtype=torch.long) |
| 60 | |
| 61 | def encode_with_transformers(self, tokens, attention_mask=None): |
| 62 | device = memory_management.get_torch_device() |
| 63 | offload_device = memory_management.text_encoder_offload_device() |
| 64 | tokens = tokens.to(device) |
| 65 | self.text_encoder.shared = self.text_encoder.shared.to(device=device, dtype=torch.float32) |
| 66 | |
| 67 | if attention_mask is not None: |
| 68 | attention_mask = attention_mask.to(device) |
| 69 | z = self.text_encoder(input_ids=tokens, attention_mask=attention_mask) |
| 70 | else: |
| 71 | z = self.text_encoder(input_ids=tokens,) |
| 72 | |