use proper model output dim

finetune.py

@@ -12,8 +12,8 @@ import gzip
 
 model_name = 'cross-encoder/ms-marco-MiniLM-L-12-v2'
 
-train_batch_size = 32
-max_seq_length = 128
+train_batch_size = 8
+max_seq_length = 384
 num_epochs = 1
 warmup_steps = 1000
 model_save_path = '.'
@@ -27,13 +27,13 @@ class ESCIDataset(Dataset):
             for line in jsonfile.readlines():
                 query = json.loads(line)
                 for doc in query['e']:
-                    self.queries.append(InputExample(texts=[query['query'], doc['title']], label=1.0))
+                    self.queries.append(InputExample(texts=[query['query'], doc['title'] + ' ' + doc['desc']], label=1.0))
                 for doc in query['s']:
-                    self.queries.append(InputExample(texts=[query['query'], doc['title']], label=0.1))
+                    self.queries.append(InputExample(texts=[query['query'], doc['title'] + ' ' + doc['desc']], label=0.1))
                 for doc in query['c']:
-                    self.queries.append(InputExample(texts=[query['query'], doc['title']], label=0.01))
+                    self.queries.append(InputExample(texts=[query['query'], doc['title'] + ' ' + doc['desc']], label=0.01))
                 for doc in query['i']:
-                    self.queries.append(InputExample(texts=[query['query'], doc['title']], label=0.0))
+                    self.queries.append(InputExample(texts=[query['query'], doc['title'] + ' ' + doc['desc']], label=0.0))
 
     def __getitem__(self, item):
         return self.queries[item]
@@ -49,9 +49,9 @@ class ESCIEvalDataset(Dataset):
                 query = json.loads(line)
                 if len(query['e']) > 0 and len(query['i']) > 0:
                     for p in query['e']:
-                        positive = p['title']
+                        positive = p['title'] + ' ' + p['title']
                         for n in query['i']:
-                            negative = n['title']
+                            negative = n['title'] + ' ' + n['title']
                             self.queries.append(InputExample(texts=[query['query'], positive, negative]))
 
     def __getitem__(self, item):

onnx_convert.py

@@ -1,9 +1,9 @@
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, AutoModelForSequenceClassification
 import torch
 
 max_seq_length=128
 
-model = AutoModel.from_pretrained(".")
+model = AutoModelForSequenceClassification.from_pretrained(".")
 model.eval()
 
 inputs = {"input_ids": torch.ones(1, max_seq_length, dtype=torch.int64),

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:7dcfb2efa8e9be4d55c8353e38f61ccfd7223e0bfc2f24ab8af495b2cbbc8bc3
+oid sha256:f8eb5889a76cfd3d6beaaf62bf061723ebf7edd212329fc527ff36c5ed1b571a
 size 133514357

pytorch_model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:cb0312525f025d18e7013477ed8c389ad104591fdbfda838599762dad8608acb
-size 133694712
+oid sha256:5a0fe068eded0383c63e7e63e8d5fef4e6d30a5e4d3011b4e7d1602844fcd251
+size 133717601

test-small.json.gz

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

train-small.json.gz

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