| """ |
| The main idea for this code is to provide a way for users to not need to bother with the hassle of multiple tokens for a concept by typing |
| a photo of <concept>_0 <concept>_1 ... and so on |
| and instead just do |
| a photo of <concept> |
| which gets translated to the above. This needs to work for both inference and training. |
| For inference, |
| the tokenizer encodes the text. So, we would want logic for our tokenizer to replace the placeholder token with |
| it's underlying vectors |
| For training, |
| we would want to abstract away some logic like |
| 1. Adding tokens |
| 2. Updating gradient mask |
| 3. Saving embeddings |
| to our Util class here. |
| so |
| TODO: |
| 1. have tokenizer keep track of concept, multiconcept pairs and replace during encode call x |
| 2. have mechanism for adding tokens x |
| 3. have mech for saving emebeddings x |
| 4. get mask to update x |
| 5. Loading tokens from embedding x |
| 6. Integrate to training x |
| 7. Test |
| """ |
|
|
| import copy |
| import random |
|
|
| from transformers import CLIPTokenizer |
|
|
|
|
| class MultiTokenCLIPTokenizer(CLIPTokenizer): |
| def __init__(self, *args, **kwargs): |
| super().__init__(*args, **kwargs) |
| self.token_map = {} |
|
|
| def try_adding_tokens(self, placeholder_token, *args, **kwargs): |
| num_added_tokens = super().add_tokens(placeholder_token, *args, **kwargs) |
| if num_added_tokens == 0: |
| raise ValueError( |
| f"The tokenizer already contains the token {placeholder_token}. Please pass a different" |
| " `placeholder_token` that is not already in the tokenizer." |
| ) |
|
|
| def add_placeholder_tokens(self, placeholder_token, *args, num_vec_per_token=1, **kwargs): |
| output = [] |
| if num_vec_per_token == 1: |
| self.try_adding_tokens(placeholder_token, *args, **kwargs) |
| output.append(placeholder_token) |
| else: |
| output = [] |
| for i in range(num_vec_per_token): |
| ith_token = placeholder_token + f"_{i}" |
| self.try_adding_tokens(ith_token, *args, **kwargs) |
| output.append(ith_token) |
| |
| for token in self.token_map: |
| if token in placeholder_token: |
| raise ValueError( |
| f"The tokenizer already has placeholder token {token} that can get confused with" |
| f" {placeholder_token}keep placeholder tokens independent" |
| ) |
| self.token_map[placeholder_token] = output |
|
|
| def replace_placeholder_tokens_in_text(self, text, vector_shuffle=False, prop_tokens_to_load=1.0): |
| """ |
| Here, we replace the placeholder tokens in text recorded in token_map so that the text_encoder |
| can encode them |
| vector_shuffle was inspired by https://github.com/rinongal/textual_inversion/pull/119 |
| where shuffling tokens were found to force the model to learn the concepts more descriptively. |
| """ |
| if isinstance(text, list): |
| output = [] |
| for i in range(len(text)): |
| output.append(self.replace_placeholder_tokens_in_text(text[i], vector_shuffle=vector_shuffle)) |
| return output |
| for placeholder_token in self.token_map: |
| if placeholder_token in text: |
| tokens = self.token_map[placeholder_token] |
| tokens = tokens[: 1 + int(len(tokens) * prop_tokens_to_load)] |
| if vector_shuffle: |
| tokens = copy.copy(tokens) |
| random.shuffle(tokens) |
| text = text.replace(placeholder_token, " ".join(tokens)) |
| return text |
|
|
| def __call__(self, text, *args, vector_shuffle=False, prop_tokens_to_load=1.0, **kwargs): |
| return super().__call__( |
| self.replace_placeholder_tokens_in_text( |
| text, vector_shuffle=vector_shuffle, prop_tokens_to_load=prop_tokens_to_load |
| ), |
| *args, |
| **kwargs, |
| ) |
|
|
| def encode(self, text, *args, vector_shuffle=False, prop_tokens_to_load=1.0, **kwargs): |
| return super().encode( |
| self.replace_placeholder_tokens_in_text( |
| text, vector_shuffle=vector_shuffle, prop_tokens_to_load=prop_tokens_to_load |
| ), |
| *args, |
| **kwargs, |
| ) |
|
|
