Update README.md

README.md

@@ -45,8 +45,10 @@ if __name__ == '__main__':
 ## Usage (HuggingFace Transformers)
 
 ```python
-from transformers import AutoTokenizer, AutoModel
+from transformers import AutoTokenizer, AutoModel, DataCollatorWithPadding
 import torch
+from torch.utils.data import DataLoader
+
 device = torch.device('cuda')
 
 # Sentences we want sentence embeddings for
@@ -54,11 +56,13 @@ sentences = ['This is an example sentence', 'Each sentence is converted']
 
 # Load model from HuggingFace Hub
 tokenizer = AutoTokenizer.from_pretrained('{MODEL_NAME}')
+collator = DataCollatorWithPadding(tokenizer)
 model = AutoModel.from_pretrained('{MODEL_NAME}').to(device)
 
-tokenized_data = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
-dataloader = DataLoader(tokenized_data, batch_size=batch_size, pin_memory=True)
-all_outputs = torch.zeros((len(tokenized_data), self.hidden_size)).to(device)
+tokenized_data = tokenizer(sentences, padding=True, truncation=True)
+tokenized_data = tokenized_data.remove_columns('text')
+dataloader = DataLoader(tokenized_data, batch_size=batch_size, pin_memory=True, collate_fn=collator)
+all_outputs = torch.zeros((len(tokenized_data), 1024)).to(device)
 start_idx = 0
 
 # I used mean-pool method for sentence representation