train

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.

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)

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)