Returns the tokenized representation of given input string(s) Parameters ---------- texts : Union[str, List[str]] An input string or a list of input strings to tokenize context_length : int The context length to use; all CLIP models use 77 as the context length
(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False)
| 195 | |
| 196 | |
| 197 | def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> Union[torch.IntTensor, torch.LongTensor]: |
| 198 | """ |
| 199 | Returns the tokenized representation of given input string(s) |
| 200 | |
| 201 | Parameters |
| 202 | ---------- |
| 203 | texts : Union[str, List[str]] |
| 204 | An input string or a list of input strings to tokenize |
| 205 | |
| 206 | context_length : int |
| 207 | The context length to use; all CLIP models use 77 as the context length |
| 208 | |
| 209 | truncate: bool |
| 210 | Whether to truncate the text in case its encoding is longer than the context length |
| 211 | |
| 212 | Returns |
| 213 | ------- |
| 214 | A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]. |
| 215 | We return LongTensor when torch version is <1.8.0, since older index_select requires indices to be long. |
| 216 | """ |
| 217 | if isinstance(texts, str): |
| 218 | texts = [texts] |
| 219 | |
| 220 | sot_token = _tokenizer.encoder["<|startoftext|>"] |
| 221 | eot_token = _tokenizer.encoder["<|endoftext|>"] |
| 222 | all_tokens = [[sot_token] + _tokenizer.encode(text) + [eot_token] for text in texts] |
| 223 | if packaging.version.parse(torch.__version__) < packaging.version.parse("1.8.0"): |
| 224 | result = torch.zeros(len(all_tokens), context_length, dtype=torch.long) |
| 225 | else: |
| 226 | result = torch.zeros(len(all_tokens), context_length, dtype=torch.int) |
| 227 | |
| 228 | for i, tokens in enumerate(all_tokens): |
| 229 | if len(tokens) > context_length: |
| 230 | if truncate: |
| 231 | tokens = tokens[:context_length] |
| 232 | tokens[-1] = eot_token |
| 233 | else: |
| 234 | raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}") |
| 235 | result[i, :len(tokens)] = torch.tensor(tokens) |
| 236 | |
| 237 | return result |