@@ -55,25 +55,45 @@ def __init__(
5555 Example:
5656 ::
5757
58- from sentence_transformers import SentenceTransformer, losses
59- from sentence_transformers.datasets import DenoisingAutoEncoderDataset
60- from torch.utils.data import DataLoader
58+ import random
59+ from datasets import Dataset
60+ from nltk import word_tokenize
61+ from nltk.tokenize.treebank import TreebankWordDetokenizer
62+ from sentence_transformers import SentenceTransformer
63+ from sentence_transformers.losses import DenoisingAutoEncoderLoss
64+ from sentence_transformers.trainer import SentenceTransformerTrainer
6165
6266 model_name = "bert-base-cased"
6367 model = SentenceTransformer(model_name)
68+
69+ def noise_transform(batch, del_ratio=0.6):
70+ texts = batch["text"]
71+ noisy_texts = []
72+ for text in texts:
73+ words = word_tokenize(text)
74+ n = len(words)
75+ if n == 0:
76+ noisy_texts.append(text)
77+ continue
78+
79+ kept_words = [word for word in words if random.random() < del_ratio]
80+ # Guarantee that at least one word remains
81+ if len(kept_words) == 0:
82+ noisy_texts.append(random.choice(words))
83+ continue
84+
85+ noisy_texts.append(TreebankWordDetokenizer().detokenize(kept_words))
86+ return {"noisy": noisy_texts, "text": texts}
87+
6488 train_sentences = [
6589 "First training sentence", "Second training sentence", "Third training sentence", "Fourth training sentence",
6690 ]
67- batch_size = 2
68- train_dataset = DenoisingAutoEncoderDataset(train_sentences)
69- train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, drop_last=True)
70- train_loss = losses.DenoisingAutoEncoderLoss(
71- model, decoder_name_or_path=model_name, tie_encoder_decoder=True
72- )
73- model.fit(
74- train_objectives=[(train_dataloader, train_loss)],
75- epochs=10,
76- )
91+
92+ train_dataset = Dataset.from_dict({"text": train_sentences})
93+ train_dataset.set_transform(transform=lambda batch: noise_transform(batch), columns=["text"], output_all_columns=True)
94+ train_loss = DenoisingAutoEncoderLoss(model, decoder_name_or_path=model_name, tie_encoder_decoder=True)
95+ trainer = SentenceTransformerTrainer(model=model, train_dataset=train_dataset, loss=train_loss)
96+ trainer.train()
7797 """
7898 super ().__init__ ()
7999
0 commit comments