Change model to Huggingface

textattack/commands/attack/attack_args_helpers.py

@@ -341,7 +341,8 @@ def parse_model_from_args(args):
             num_labels,
             model_path=args.model,
         )
-        model = textattack.models.wrappers.PyTorchModelWrapper(
+        # Logic to change this according to model being loaded
+        model = textattack.models.wrappers.HuggingFaceModelWrapper(
             model.model, model.tokenizer, batch_size=args.model_batch_size
         )
     else:

textattack/commands/eval_model/eval_model_command.py

@@ -40,7 +40,7 @@ class EvalModelCommand(TextAttackCommand):
 
     def get_preds(self, model, inputs):
         with torch.no_grad():
-            preds = model(inputs)
+            preds = textattack.shared.utils.batch_model_predict(model, inputs)
         return preds
 
     def test_model_on_dataset(self, args):

textattack/models/wrappers/pytorch_model_wrapper.py

@@ -33,7 +33,10 @@ class PyTorchModelWrapper(ModelWrapper):
 
     def __call__(self, text_input_list):
         model_device = next(self.model.parameters()).device
+<<<<<<< HEAD
         ids = self.encode(text_input_list)
+=======
+>>>>>>> Change model to Huggingface
         ids = torch.tensor(ids).to(model_device)
 
         with torch.no_grad():

textattack/shared/utils/tensor.py

@@ -11,7 +11,7 @@ def batch_model_predict(model_predict, inputs, batch_size=32):
     i = 0
     while i < len(inputs):
         batch = inputs[i : i + batch_size]
-        batch_preds = model_predict(batch)[0]
+        batch_preds = model_predict(batch)
 
         # Some seq-to-seq models will return a single string as a prediction
         # for a single-string list. Wrap these in a list.