import torch
from transformers import AutoModel

from fewie.encoders.encoder import Encoder, EncoderOutput


class RobertaEncoder(Encoder):
    def __init__(self, model_name_or_path: str) -> None:
        super().__init__()
        self.model = AutoModel.from_pretrained(model_name_or_path)

    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        labels=None,
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ) -> EncoderOutput:

        output = self.model(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=False,
            output_hidden_states=False,
            return_dict=True,
        )

        return EncoderOutput(embeddings=output.last_hidden_state)
