Update README.md

README.md

@@ -20,18 +20,23 @@ language:
 ```python
 from transformers import AutoTokenizer, AutoModel
 import torch
+device = torch.device('cuda')
 
 # Sentences we want sentence embeddings for
 sentences = ['This is an example sentence', 'Each sentence is converted']
 
 # Load model from HuggingFace Hub
 tokenizer = AutoTokenizer.from_pretrained('{MODEL_NAME}')
-model = AutoModel.from_pretrained('{MODEL_NAME}')
+model = AutoModel.from_pretrained('{MODEL_NAME}').to(device)
 
-encoded_input = tokenizer(sentences, padding=True, truncation=True, return_tensors='pt')
+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)
+start_idx = 0
 
 # I used mean-pool method for sentence representation
 with torch.no_grad():
+  for inputs in tqdm(dataloader):
     inputs = {k: v.to(device) for k, v in inputs.items()}
     representations, _ = self.model(**inputs, return_dict=False)
     attention_mask = inputs["attention_mask"]
@@ -39,7 +44,9 @@ with torch.no_grad():
     summed = torch.sum(representations * input_mask_expanded, 1)
     sum_mask = input_mask_expanded.sum(1)
     sum_mask = torch.clamp(sum_mask, min=1e-9)
-    res.extend(summed.cpu() / sum_mask.cpu()) # mean_pool
+    end_idx = start_idx + representations.shape[0]
+    all_outputs[start_idx:end_idx] = (summed / sum_mask)
+    start_idx = end_idx
 
 ```