File size: 1,189 Bytes
8b41845 89dc002 8b41845 b39aef7 8b41845 60e3fe9 8b41845 | 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 | import torch
from transformers import PreTrainedModel
from transformers.modeling_outputs import BaseModelOutput
from spectre.model import SpectreImageFeatureExtractor
try:
from .configuration_spectre import SpectreConfig
except ImportError:
from configuration_spectre import SpectreConfig
class SpectreModel(PreTrainedModel):
config_class = SpectreConfig
base_model_prefix = "spectre"
main_input_name = "pixel_values"
def __init__(self, config):
super().__init__(config)
self.model = SpectreImageFeatureExtractor(
backbone_name=config.backbone_name,
backbone_kwargs=config.backbone_kwargs,
feature_combiner_name=config.feature_combiner_name,
feature_combiner_kwargs=config.feature_combiner_kwargs,
)
self.post_init()
def forward(
self,
pixel_values: torch.Tensor,
grid_size=None,
return_dict=False,
**kwargs,
):
outputs = self.model(pixel_values, grid_size=grid_size)
if not return_dict:
return outputs
return BaseModelOutput(last_hidden_state=outputs) |