import torch
import torch.nn.functional as F

from lm_eval.api.registry import register_model
from lm_eval.models.huggingface import HFLM


@register_model("cloverlm")
class CloverLMHFLM(HFLM):
    def __init__(self, pad_multiple=128, **kwargs):
        super().__init__(**kwargs)
        self.pad_multiple = pad_multiple

    def _encode_pair(self, context, continuation):
        context_enc, continuation_enc = super()._encode_pair(context, continuation)

        if not continuation_enc and continuation:
            whole_enc = self.tok_encode(context + continuation)
            if len(whole_enc) > 1:
                continuation_enc = whole_enc[-1:]
                context_enc = whole_enc[:-1]
            elif whole_enc:
                continuation_enc = whole_enc
                context_enc = [self.prefix_token_id]
            else:
                continuation_enc = [self.prefix_token_id]
                context_enc = [self.prefix_token_id]

        return context_enc, continuation_enc

    def _model_call(self, inps: torch.Tensor, attn_mask: torch.Tensor = None, **kwargs):
        orig_len = inps.shape[1]
        remainder = orig_len % self.pad_multiple
        
        if remainder != 0:
            pad_len = self.pad_multiple - remainder
            inps = F.pad(inps, (0, pad_len), value=self.tokenizer.pad_token_id)
            if attn_mask is not None:
                attn_mask = F.pad(attn_mask, (0, pad_len), value=0)
                
        logits = super()._model_call(inps, attn_mask=attn_mask, **kwargs)
        if remainder != 0:
            logits = logits[:, :orig_len, :]
        return logits

    def _model_generate(self, context, max_length, **kwargs):
        orig_len = context.shape[1]
        remainder = orig_len % self.pad_multiple
        
        if remainder != 0:
            pad_len = self.pad_multiple - remainder
            context = F.pad(context, (pad_len, 0), value=self.tokenizer.pad_token_id)
            if "attention_mask" in kwargs and kwargs["attention_mask"] is not None:
                kwargs["attention_mask"] = F.pad(kwargs["attention_mask"], (pad_len, 0), value=0)
                
        out = super()._model_generate(context, max_length, **kwargs)
        if remainder != 0:
            out = out[:, pad_len:]
            
        return out


if __name__ == "__main__":
    from lm_eval.__main__ import cli_evaluate
    cli_evaluate()