diff --git a/README.md b/README.md index af568c3..0b170d8 100644 --- a/README.md +++ b/README.md @@ -99,3 +99,36 @@ from towhee import ops op = ops.sentence_embedding.sentence_transformers().get_op() full_list = op.supported_model_names() ``` + +## Fine-tune + +### Get started + +In this example, we fine-tune operator in Semantic Textual Similarity (STS) task, which assigns a score on the similarity of two texts. +We use the STSbenchmark as training data to fine-tune. + +We only need to construct an op instance and pass in some configurations to train the specified task. +```python +import towhee +import os +from sentence_transformers import util + +op = towhee.ops.sentence_embedding.sentence_transformers(model_name='nli-distilroberta-base-v2').get_op() + +sts_dataset_path = 'datasets/stsbenchmark.tsv.gz' + +if not os.path.exists(sts_dataset_path): + util.http_get('https://sbert.net/datasets/stsbenchmark.tsv.gz', sts_dataset_path) + +training_config = { + 'sts_dataset_path': sts_dataset_path, + 'train_batch_size': 16, + 'num_epochs': 4, + 'model_save_path': './output' +} +op.train(training_config) +``` + +### Dive deep and customize your training +You can change the [training script](https://towhee.io/sentence-embedding/sentence_transformers/src/branch/main/train_sts_task.py) in your customer way. +Or your can refer to the original [sbert training guide](https://www.sbert.net/docs/training/overview.html) and [code example](https://github.com/UKPLab/sentence-transformers/tree/master/examples/training) for more information. \ No newline at end of file diff --git a/s_bert.py b/s_bert.py index f7db40d..c2bbb8d 100644 --- a/s_bert.py +++ b/s_bert.py @@ -249,3 +249,26 @@ class STransformers(NNOperator): else: log.error(f'Invalid or unsupported format "{format}".') return model_list + + def train(self, training_config=None, **kwargs): + from .train_sts_task import train_sts + train_sts(self._model, training_config) + + + +if __name__ == '__main__': + from sentence_transformers import util + op = STransformers(model_name='nli-distilroberta-base-v2') + # Check if dataset exsist. If not, download and extract it + sts_dataset_path = 'datasets/stsbenchmark.tsv.gz' + + if not os.path.exists(sts_dataset_path): + util.http_get('https://sbert.net/datasets/stsbenchmark.tsv.gz', sts_dataset_path) + + training_config = { + 'sts_dataset_path': sts_dataset_path, + 'train_batch_size': 16, + 'num_epochs': 4, + 'model_save_path': './output' + } + op.train(training_config) \ No newline at end of file diff --git a/train_sts_task.py b/train_sts_task.py new file mode 100644 index 0000000..869c96f --- /dev/null +++ b/train_sts_task.py @@ -0,0 +1,86 @@ +# This script is hacked and modified from https://github.com/UKPLab/sentence-transformers/blob/master/examples/training/sts/training_stsbenchmark_continue_training.py +# For more specified training tasks, please refer https://github.com/UKPLab/sentence-transformers/tree/master/examples/training + +from torch.utils.data import DataLoader +import math +from sentence_transformers import SentenceTransformer, LoggingHandler, losses, InputExample +from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator +import logging +from datetime import datetime +import os +import gzip +import csv + +#### Just some code to print debug information to stdout +logging.basicConfig(format='%(asctime)s - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + level=logging.INFO, + handlers=[LoggingHandler()]) +#### /print debug information to stdout + + + +def train_sts(model, training_config): + sts_dataset_path = training_config['sts_dataset_path'] + train_batch_size = training_config['train_batch_size'] + num_epochs = training_config['num_epochs'] + model_save_path = training_config['model_save_path'] + if not os.path.exists(model_save_path): + os.mkdir(model_save_path) + model_save_path = os.path.join('training_stsbenchmark_continue_training-' + datetime.now().strftime( + "%Y-%m-%d_%H-%M-%S")) + + + # Convert the dataset to a DataLoader ready for training + logging.info("Read STSbenchmark train dataset") + + train_samples = [] + dev_samples = [] + test_samples = [] + with gzip.open(sts_dataset_path, 'rt', encoding='utf8') as fIn: + reader = csv.DictReader(fIn, delimiter='\t', quoting=csv.QUOTE_NONE) + for row in reader: + score = float(row['score']) / 5.0 # Normalize score to range 0 ... 1 + inp_example = InputExample(texts=[row['sentence1'], row['sentence2']], label=score) + + if row['split'] == 'dev': + dev_samples.append(inp_example) + elif row['split'] == 'test': + test_samples.append(inp_example) + else: + train_samples.append(inp_example) + + + + train_dataloader = DataLoader(train_samples, shuffle=True, batch_size=train_batch_size) + train_loss = losses.CosineSimilarityLoss(model=model) + + + # Development set: Measure correlation between cosine score and gold labels + logging.info("Read STSbenchmark dev dataset") + evaluator = EmbeddingSimilarityEvaluator.from_input_examples(dev_samples, name='sts-dev') + + + # Configure the training. We skip evaluation in this example + warmup_steps = math.ceil(len(train_dataloader) * num_epochs * 0.1) #10% of train data for warm-up + logging.info("Warmup-steps: {}".format(warmup_steps)) + + + # Train the model + model.fit(train_objectives=[(train_dataloader, train_loss)], + evaluator=evaluator, + epochs=num_epochs, + evaluation_steps=1000, + warmup_steps=warmup_steps, + output_path=model_save_path) + + + ############################################################################## + # + # Load the stored model and evaluate its performance on STS benchmark dataset + # + ############################################################################## + + model = SentenceTransformer(model_save_path) + test_evaluator = EmbeddingSimilarityEvaluator.from_input_examples(test_samples, name='sts-test') + test_evaluator(model, output_path=model_save_path) \ No newline at end of file