From df54044de582d6f5fd122c97c5198a4206bb87a0 Mon Sep 17 00:00:00 2001 From: Jael Gu Date: Mon, 13 Mar 2023 14:40:35 +0800 Subject: [PATCH] Remove dc2 for towhee1.0 Signed-off-by: Jael Gu --- README.md | 4 ++-- s_bert.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/README.md b/README.md index e30dde3..2fe3c33 100644 --- a/README.md +++ b/README.md @@ -20,7 +20,7 @@ to generate a text embedding for the sentence "This is a sentence.". *Write a pipeline with explicit inputs/outputs name specifications:* ```python -from towhee.dc2 import pipe, ops, DataCollection +from towhee import pipe, ops, DataCollection p = ( pipe.input('sentence') @@ -140,4 +140,4 @@ new_op = towhee.ops.sentence_embedding.sbert(model_name=model_path).get_op() ### Dive deep and customize your training You can change the [training script](https://towhee.io/sentence-embedding/sbert/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 +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. diff --git a/s_bert.py b/s_bert.py index c2bbb8d..1a73b58 100644 --- a/s_bert.py +++ b/s_bert.py @@ -85,6 +85,8 @@ class STransformers(NNOperator): else: sentences = txt inputs = self.tokenize(sentences) + for k, v in inputs.items(): + inputs[k] = v.to(self.device) embs = self.model(**inputs).cpu().detach().numpy() if isinstance(txt, str): embs = embs.squeeze(0) @@ -271,4 +273,4 @@ if __name__ == '__main__': 'num_epochs': 4, 'model_save_path': './output' } - op.train(training_config) \ No newline at end of file + op.train(training_config)