Adapters
khulnasoft commited on
Commit
3287a71
·
verified ·
1 Parent(s): fa1c6ee

Create gpt2_sequence_length.py

Browse files
prompt_injection/evaluators/gpt2_sequence_length.py ADDED
@@ -0,0 +1,28 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+
2
+ from gettext import npgettext
3
+ from prompt_injection.evaluators.base import PromptEvaluator
4
+
5
+ from transformers import GPT2Tokenizer, GPT2LMHeadModel,GPT2Model
6
+
7
+
8
+
9
+ class GPT2SequenceLengthPromptEvaluator(PromptEvaluator):
10
+ def __init__(self) -> None:
11
+ super().__init__()
12
+
13
+ def __calculate_sequence_length(self,sentence, model_name='gpt2'):
14
+ # Load pre-trained model and tokenizer
15
+ tokenizer = GPT2Tokenizer.from_pretrained(model_name)
16
+ inputs = tokenizer(sentence, return_tensors='pt')
17
+ return inputs['input_ids'].shape[1]
18
+
19
+ def eval_sample(self,sample):
20
+ try:
21
+ return self.__calculate_sequence_length(sample)
22
+ except Exception as err:
23
+ print(err)
24
+
25
+ return npgettext.nan
26
+
27
+ def get_name(self):
28
+ return 'Length'