File size: 2,264 Bytes
3e0b0bf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | """Lightweight text processor for QaDiT (T5 tokenizer wrapper)."""
from __future__ import annotations
from typing import List, Optional, Union
from transformers import AutoTokenizer
from transformers.feature_extraction_utils import BatchFeature
try:
from .configuration_qadit import QaDiTConfig
except ImportError:
from configuration_qadit import QaDiTConfig
class QaDiTProcessor:
"""Tokenize captions the same way training / ``generate`` expect.
Usage::
processor = QaDiTProcessor.from_pretrained("USER/qadit")
inputs = processor("A dog barks", return_tensors="pt")
"""
def __init__(self, tokenizer, text_max_length: int = 64):
self.tokenizer = tokenizer
self.text_max_length = text_max_length
@classmethod
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
trust = kwargs.pop("trust_remote_code", True)
try:
config = QaDiTConfig.from_pretrained(
pretrained_model_name_or_path,
trust_remote_code=trust,
)
text_model = config.text_model
text_max_length = config.text_max_length
except Exception:
text_model = "google/flan-t5-large"
text_max_length = 64
tokenizer = AutoTokenizer.from_pretrained(text_model)
return cls(tokenizer=tokenizer, text_max_length=text_max_length)
def __call__(
self,
text: Union[str, List[str]],
padding: str = "max_length",
truncation: bool = True,
max_length: Optional[int] = None,
return_tensors: Optional[str] = "pt",
**kwargs,
) -> BatchFeature:
if isinstance(text, str):
text = [text]
encoded = self.tokenizer(
text,
padding=padding,
truncation=truncation,
max_length=max_length or self.text_max_length,
return_tensors=return_tensors,
**kwargs,
)
return BatchFeature(data=dict(encoded))
def batch_decode(self, *args, **kwargs):
return self.tokenizer.batch_decode(*args, **kwargs)
def decode(self, *args, **kwargs):
return self.tokenizer.decode(*args, **kwargs)
__all__ = ["QaDiTProcessor"]
|