Update README.md

README.md

@@ -48,19 +48,22 @@ if __name__ == '__main__':
 from transformers import AutoTokenizer, AutoModel, DataCollatorWithPadding
 import torch
 from torch.utils.data import DataLoader
+from tqdm import tqdm
+from datasets import Dataset
 
-device = torch.device('cuda')
+device = torch.device('cpu')
+batch_size=1
 
 # 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}')
-collator = DataCollatorWithPadding(tokenizer)
-model = AutoModel.from_pretrained('{MODEL_NAME}').to(device)
+tokenizer = AutoTokenizer.from_pretrained('sorryhyun/sentence-embedding-klue-large')
+collator = DataCollatorWithPadding(tokenizer=tokenizer)
+model = AutoModel.from_pretrained('sorryhyun/sentence-embedding-klue-large').to(device)
 
 tokenized_data = tokenizer(sentences, padding=True, truncation=True)
-tokenized_data = tokenized_data.remove_columns('text')
+tokenized_data = Dataset.from_dict(tokenized_data)
 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
@@ -69,7 +72,7 @@ start_idx = 0
 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)
+    representations, _ = model(**inputs, return_dict=False)
     attention_mask = inputs["attention_mask"]
     input_mask_expanded = (attention_mask.unsqueeze(-1).expand(representations.size()).to(representations.dtype))
     summed = torch.sum(representations * input_mask_expanded, 1)