| from typing import Any, Dict, List, Optional, Tuple, Union |
|
|
| import torch |
| from torch import nn |
| from transformers import AutoConfig, AutoModel, AutoProcessor |
|
|
|
|
| class VisionTransformer(nn.Module): |
| """Huggingface AutoModel to generate token embeddings. |
| Loads the correct class, e.g. BERT / RoBERTa etc. |
| |
| Args: |
| model_name_or_path: Huggingface models name |
| (https://huggingface.co/models) |
| model_args: Keyword arguments passed to the Huggingface |
| Transformers model |
| tokenizer_args: Keyword arguments passed to the Huggingface |
| Transformers tokenizer |
| config_args: Keyword arguments passed to the Huggingface |
| Transformers config |
| cache_dir: Cache dir for Huggingface Transformers to store/load |
| models |
| """ |
|
|
| def __init__( |
| self, |
| model_name_or_path: str, |
| model_args: Optional[Dict[str, Any]] = None, |
| tokenizer_args: Optional[Dict[str, Any]] = None, |
| config_args: Optional[Dict[str, Any]] = None, |
| cache_dir: Optional[str] = None, |
| ) -> None: |
| super(VisionTransformer, self).__init__() |
| if model_args is None: |
| model_args = {} |
| if tokenizer_args is None: |
| tokenizer_args = {} |
| if config_args is None: |
| config_args = {} |
|
|
| self.config = AutoConfig.from_pretrained(model_name_or_path, **config_args, cache_dir=cache_dir) |
| self.model = AutoModel.from_pretrained(model_name_or_path, config=self.config, **model_args, cache_dir=cache_dir) |
| self.processor = AutoProcessor.from_pretrained(model_name_or_path, config=self.config, **tokenizer_args, cache_dir=cache_dir) |
|
|
| def forward(self, features: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: |
| """Returns token_embeddings, cls_token""" |
| output_states = self.model(pixel_values=features["pixel_values"], return_dict=False)[0] |
| features.update({"token_embeddings": output_states}) |
| return features |
|
|
| def get_word_embedding_dimension(self) -> int: |
| return self.config.hidden_size |
|
|
| def tokenize( |
| self, texts: Union[List[str], List[Dict], List[Tuple[str, str]]], padding: Union[str, bool] = True |
| ) -> Dict[str, torch.Tensor]: |
| return self.processor(texts, return_tensors="pt") |
|
|
| def get_config_dict(self) -> Dict[str, Any]: |
| return {key: self.__dict__[key] for key in self.config_keys} |
|
|
| def save(self, output_path: str, safe_serialization: bool = True) -> None: |
| self.model.save_pretrained(output_path, safe_serialization=safe_serialization) |
| self.processor.save_pretrained(output_path) |
|
|
