metarank
/

esci-MiniLM-L6-v2

Sentence Similarity

sentence-transformers

feature-extraction

Inference Endpoints

Model card Files Files and versions Community

esci-MiniLM-L6-v2 / finetune_cos.py

shuttie's picture

switch to cos loss with exp mapping

a16d1d7 12 months ago

raw history blame contribute delete

3 kB

	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)