Create gpt2_perplexity.py

prompt_injection/evaluators/gpt2_perplexity.py

@@ -0,0 +1,32 @@
+from prompt_injection.evaluators.base import PromptEvaluator
+from transformers import GPT2Tokenizer, GPT2LMHeadModel,GPT2Model
+import torch
+import numpy as np
+
+class GPT2PerplexityEvaluator(PromptEvaluator):
+  def __init__(self,model_name='gpt2') -> None:
+      super().__init__()
+      self.model_name=model_name
+      self.tokenizer_gpt2 = GPT2Tokenizer.from_pretrained('gpt2')
+      self.model_gpt2 = GPT2LMHeadModel.from_pretrained('gpt2')
+
+  def calculate_perplexity(self,sentence):
+    inputs = self.tokenizer_gpt2(sentence, return_tensors='pt')
+    input_ids = inputs['input_ids']
+    with torch.no_grad():
+        outputs = self.model_gpt2(input_ids, labels=input_ids)
+    # Calculate the loss
+    loss = outputs.loss
+    perplexity = torch.exp(loss)
+    return  perplexity
+
+
+  def eval_sample(self,sample):
+      try:
+        return self.calculate_perplexity(sample)
+      except Exception as err:
+        print(err)
+        return np.nan
+
+  def get_name(self):
+      return 'Perplexity'