import torch
from transformers import PreTrainedModel
from .configuration_x3d import X3DConfig
from .x3d import build_model


class X3DModel(PreTrainedModel):
    config_class = X3DConfig

    def __init__(self, config, **kwargs):
        super().__init__(config)
        self.model = build_model(config.cfg)

        checkpoint = kwargs.get("checkpoint", None)

        if checkpoint:
            checkpoint = torch.load(
                checkpoint, weights_only=True, map_location=torch.device("cpu"))
            self.model.load_state_dict(checkpoint["model_state"])

    def forward(self, input_video):
        outputs = self.model(input_video)
        return outputs