Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Add truncate_text option to tokenize #126

Merged
merged 2 commits into from
Jul 19, 2021
Merged

Add truncate_text option to tokenize #126

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a2f6647
Add truncate_text option to tokenize
rom1504 Jul 8, 2021
624f90e
add doc, rename to just "truncate", use eot_token
jongwook Jul 19, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 9 additions & 2 deletions clip/clip.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -161,7 +161,7 @@ def patch_float(module):
return model, _transform(model.input_resolution.item())


def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.LongTensor:
def tokenize(texts: Union[str, List[str]], context_length: int = 77, truncate: bool = False) -> torch.LongTensor:
"""
Returns the tokenized representation of given input string(s)

Expand All @@ -173,6 +173,9 @@ def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.Lo
context_length : int
The context length to use; all CLIP models use 77 as the context length

truncate: bool
Whether to truncate the text in case its encoding is longer than the context length

Returns
-------
A two-dimensional tensor containing the resulting tokens, shape = [number of input strings, context_length]
Expand All @@ -187,7 +190,11 @@ def tokenize(texts: Union[str, List[str]], context_length: int = 77) -> torch.Lo

for i, tokens in enumerate(all_tokens):
if len(tokens) > context_length:
raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
if truncate:
tokens = tokens[:context_length]
tokens[-1] = eot_token
else:
raise RuntimeError(f"Input {texts[i]} is too long for context length {context_length}")
result[i, :len(tokens)] = torch.tensor(tokens)

return result