from gettext import npgettext
from prompt_injection.evaluators.base import PromptEvaluator

from transformers import GPT2Tokenizer, GPT2LMHeadModel,GPT2Model



class GPT2SequenceLengthPromptEvaluator(PromptEvaluator):
    def __init__(self) -> None:
       super().__init__()

    def __calculate_sequence_length(self,sentence, model_name='gpt2'):
        # Load pre-trained model and tokenizer
        tokenizer = GPT2Tokenizer.from_pretrained(model_name)
        inputs = tokenizer(sentence, return_tensors='pt')
        return inputs['input_ids'].shape[1]

    def eval_sample(self,sample):
        try:
          return self.__calculate_sequence_length(sample)
        except Exception as err:
          print(err)

          return npgettext.nan

    def get_name(self):
        return 'Length'