Upload model

config.json

@@ -42,7 +42,7 @@
   "summary_type": "cls_index",
   "summary_use_proj": true,
   "torch_dtype": "float32",
-  "transformers_version": "4.31.0",
+  "transformers_version": "4.35.2",
   "use_cache": true,
   "vocab_size": 50264
 }

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a540f442266937e77756ce8730b0f0492dc80e51c21522a7e86984e3c6d21bfd
+size 682724236

modeling_backpack_gpt2_nli.py

@@ -52,4 +52,12 @@ class BackpackGPT2NLIModel(GPT2PreTrainedModel):
             loss = self.loss_func(flat_logits, flat_labels)
             return {'logits': logits, 'loss': loss}
         else:
-            return {'logits': logits}
+            return {'logits': logits}
+        
+
+    def predict(self, input_ids=None, attention_mask=None):
+        logits = self.forward(input_ids, attention_mask, labels=None)
+        p = torch.argmax(p, axis=1)
+        labels = [self.config.id2label[index] for index in p]
+        return labels
+