| 49 | |
| 50 | |
| 51 | class LTXVGemmaTokenizer: |
| 52 | def __init__(self, tokenizer_path: str, max_length: int = 1024): |
| 53 | self.tokenizer = AutoTokenizer.from_pretrained( |
| 54 | tokenizer_path, local_files_only=True, model_max_length=max_length |
| 55 | ) |
| 56 | # Gemma expects left padding for chat-style prompts; for plain text it doesn't matter much. |
| 57 | self.tokenizer.padding_side = "left" |
| 58 | if self.tokenizer.pad_token is None: |
| 59 | self.tokenizer.pad_token = self.tokenizer.eos_token |
| 60 | |
| 61 | self.max_length = max_length |
| 62 | |
| 63 | def tokenize_with_weights(self, text: str, return_word_ids: bool = False): |
| 64 | text = text.strip() |
| 65 | encoded = self.tokenizer( |
| 66 | text, |
| 67 | padding="max_length", |
| 68 | max_length=self.max_length, |
| 69 | truncation=True, |
| 70 | return_tensors="pt", |
| 71 | ) |
| 72 | input_ids = encoded.input_ids |
| 73 | attention_mask = encoded.attention_mask |
| 74 | tuples = [ |
| 75 | (token_id, attn, i) |
| 76 | for i, (token_id, attn) in enumerate(zip(input_ids[0], attention_mask[0])) |
| 77 | ] |
| 78 | out = {"gemma": tuples} |
| 79 | |
| 80 | if not return_word_ids: |
| 81 | out = {k: [(t, w) for t, w, _ in v] for k, v in out.items()} |
| 82 | |
| 83 | return out |
| 84 | |
| 85 | |
| 86 | class LTXVGemmaTextEncoderModel(torch.nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected