upd: utils module

promptbench/utils/dataprocess.py

@@ -1,3 +1,6 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
 class DataProcessor:
  def __init__(self):
  self.INPUT_PROCESS_MAP = {

promptbench/utils/defense.py

@@ -5,13 +5,13 @@
 
 
 class Defense(object):
- def __init__(self, lang='en', defense_method='autocorrect'):
+ def __init__(self, defense_method='spellcorrect', lang='en'):
  self.defense_method = defense_method
- if self.defense_method == 'autocorrect':
+ if self.defense_method == 'spellcorrect':
  self.spell = Speller(lang=lang)
 
  def __call__(self, text):
- if self.defense_method == 'autocorrect':
+ if self.defense_method == 'spellcorrect':
  return self.spell(text)
  else:
  raise NotImplementedError

promptbench/utils/visualize.py

@@ -3,7 +3,7 @@
 
 class Visualizer:
 
- def __init__(self) -> None:
+ def __init__(self, model) -> None:
  self.model = model
 
  def vis_by_grad(self, tokenizer, input_sentence, label):

test.py

@@ -18,6 +18,12 @@
 output_t5 = model_t5("Translate English to German: 'Hello World'")
 print(output_t5)
 
+# test visualization, not tested yet
 from promptbench.utils import Visualizer
 x = Visualizer('google/flan-t5-large')
 print(x)
+
+# test defense of prompt input using spell correct
+from promptbench.utils import Defense
+x = Defense('spellcorrect')
+print(x('I am a student at the Univrsity of California, Berkeey.'))