Skip to content

Commit 8c4be2c

Browse files
omkar-334tomaarsen
andauthored
update docstring for DenoisingAutoEncoderLoss.py (#3652)
* update docstring for `DenoisingAutoEncoderLoss.py` * Shrink example; it's meant to be small --------- Co-authored-by: Tom Aarsen <Cubiegamedev@gmail.com>
1 parent 0389228 commit 8c4be2c

1 file changed

Lines changed: 33 additions & 13 deletions

File tree

sentence_transformers/losses/DenoisingAutoEncoderLoss.py

Lines changed: 33 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)