switch to cos loss with exp mapping

.gitignore

@@ -1,2 +1,4 @@
 .venv
-venv
+venv
+cos-exp
+cos-lin

finetune_cos.py

@@ -0,0 +1,87 @@
+from sentence_transformers import SentenceTransformer, LoggingHandler, util, models, evaluation, losses, InputExample, CrossEncoder
+from torch import nn
+import csv
+from torch.utils.data import DataLoader, Dataset
+import torch
+from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator, SentenceEvaluator, SimilarityFunction, RerankingEvaluator
+from sentence_transformers.cross_encoder.evaluation import CERerankingEvaluator
+import logging
+import json
+import random
+import gzip
+
+model_name = 'sentence-transformers/all-MiniLM-L6-v2'
+
+train_batch_size = 100
+max_seq_length = 128
+num_epochs = 1
+warmup_steps = 1000
+model_save_path = 'cos-exp'
+lr = 2e-5
+
+class ESCIDataset(Dataset):
+    def __init__(self, input):
+        self.queries = []
+        with gzip.open(input) as jsonfile:
+            for line in jsonfile.readlines():
+                query = json.loads(line)
+                for p in query['e']:
+                    positive = p['title']
+                    self.queries.append(InputExample(texts=[query['query'], positive], label=1.0))
+                for p in query['s']:
+                    positive = p['title']
+                    self.queries.append(InputExample(texts=[query['query'], positive], label=0.1))
+                for p in query['c']:
+                    positive = p['title']
+                    self.queries.append(InputExample(texts=[query['query'], positive], label=0.01))
+                for p in query['i']:
+                    positive = p['title']
+                    self.queries.append(InputExample(texts=[query['query'], positive], label=0.0))
+
+    def __getitem__(self, item):
+        return self.queries[item]
+
+    def __len__(self):
+        return len(self.queries)
+
+
+model = SentenceTransformer(model_name, device='cpu')
+model.max_seq_length = max_seq_length
+
+
+train_dataset = ESCIDataset(input='train-small.json.gz')
+eval_dataset = ESCIDataset(input='test-small.json.gz')
+train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=train_batch_size)
+train_loss = losses.CosineSimilarityLoss(model=model)
+
+# samples = {}
+# for query in eval_dataset.queries:
+#     qstr = query.texts[0]
+#     sample = samples.get(qstr, {'query': qstr})
+#     positive = sample.get('positive', [])
+#     positive.append(query.texts[1])
+#     sample['positive'] = positive
+#     negative = sample.get('negative', [])
+#     negative.append(query.texts[2])
+#     sample['negative'] = negative
+#     samples[qstr] = sample
+
+# evaluator = RerankingEvaluator(samples=samples,name='esci')
+
+# Train the model
+
+model.fit(train_objectives=[(train_dataloader, train_loss)],
+          epochs=num_epochs,
+          warmup_steps=warmup_steps,
+          use_amp=True,
+#          checkpoint_path=model_save_path,
+#          checkpoint_save_steps=len(train_dataloader),
+          optimizer_params = {'lr': lr},
+#          evaluator=evaluator,
+#          evaluation_steps=1000,
+          output_path=model_save_path
+          )
+
+# Save the model
+
+model.save(model_save_path)

pytorch_model.bin

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:32be7d4ded580fffe92bfff7e8dc865b17dc0cfbba3dd598865adc54dd89d0c3
+oid sha256:6c0cc55f150c9ab5404bc6b64cd1b9399b81bfe3cf083a1d30c89e9cd7d4235a
 size 90891565

pytorch_model.onnx

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b1abb2047d142c643ed10c33b6a1517171dab1ffb3d3a57bc2043437d3f5bf77
+oid sha256:1264e3429772fc7d27c6c2f6e9bbf04e4358821fb380dfcff18e5f96d14f8f32
 size 90984263