| 1 | from transformers import AutoTokenizer |
| 2 | |
| 3 | |
| 4 | class LTXVGemmaTokenizer: |
| 5 | """ |
| 6 | Tokenizer wrapper for Gemma models compatible with LTXV processes. |
| 7 | This class wraps HuggingFace's `AutoTokenizer` for use with Gemma text encoders, |
| 8 | ensuring correct settings and output formatting for downstream consumption. |
| 9 | """ |
| 10 | |
| 11 | def __init__(self, tokenizer_path: str, max_length: int = 256): |
| 12 | """ |
| 13 | Initialize the tokenizer. |
| 14 | Args: |
| 15 | tokenizer_path (str): Path to the pretrained tokenizer files or model directory. |
| 16 | max_length (int, optional): Max sequence length for encoding. Defaults to 256. |
| 17 | """ |
| 18 | self.tokenizer = AutoTokenizer.from_pretrained( |
| 19 | tokenizer_path, local_files_only=True, model_max_length=max_length |
| 20 | ) |
| 21 | # Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much. |
| 22 | self.tokenizer.padding_side = "left" |
| 23 | if self.tokenizer.pad_token is None: |
| 24 | self.tokenizer.pad_token = self.tokenizer.eos_token |
| 25 | |
| 26 | self.max_length = max_length |
| 27 | |
| 28 | def tokenize_with_weights(self, text: str, return_word_ids: bool = False) -> dict[str, list[tuple[int, int]]]: |
| 29 | """ |
| 30 | Tokenize the given text and return token IDs and attention weights. |
| 31 | Args: |
| 32 | text (str): The input string to tokenize. |
| 33 | return_word_ids (bool, optional): If True, includes the token's position (index) in the output tuples. |
| 34 | If False (default), omits the indices. |
| 35 | Returns: |
| 36 | dict[str, list[tuple[int, int]]] OR dict[str, list[tuple[int, int, int]]]: |
| 37 | A dictionary with a "gemma" key mapping to: |
| 38 | - a list of (token_id, attention_mask) tuples if return_word_ids is False; |
| 39 | - a list of (token_id, attention_mask, index) tuples if return_word_ids is True. |
| 40 | Example: |
| 41 | >>> tokenizer = LTXVGemmaTokenizer("path/to/tokenizer", max_length=8) |
| 42 | >>> tokenizer.tokenize_with_weights("hello world") |
| 43 | {'gemma': [(1234, 1), (5678, 1), (2, 0), ...]} |
| 44 | """ |
| 45 | text = text.strip() |
| 46 | encoded = self.tokenizer( |
| 47 | text, |
| 48 | padding="max_length", |
| 49 | max_length=self.max_length, |
| 50 | truncation=True, |
| 51 | return_tensors="pt", |
| 52 | ) |
| 53 | input_ids = encoded.input_ids |
| 54 | attention_mask = encoded.attention_mask |
| 55 | tuples = [ |
| 56 | (token_id, attn, i) for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0], strict=True)) |
| 57 | ] |
| 58 | out = {"gemma": tuples} |
| 59 | |
| 60 | if not return_word_ids: |
| 61 | # Return only (token_id, attention_mask) pairs, omitting token position |
| 62 | out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()} |
| 63 | |
| 64 | return out |
| 65 |