Update code snippet to use sentence-level embeddings

README.md

@@ -24,7 +24,7 @@ This checkpoint is first trained on code data via masked language modeling (MLM)
 ### How to use
 This checkpoint consists of an encoder (130M model), which can be used to extract code embeddings of 1024 dimension. It can be easily loaded using the AutoModel functionality and employs the Starcoder tokenizer (https://arxiv.org/pdf/2305.06161.pdf).
 
-```
+```python
 from transformers import AutoModel, AutoTokenizer
 
 checkpoint = "codesage/codesage-small"
@@ -33,10 +33,10 @@ device = "cuda"  # for GPU usage or "cpu" for CPU usage
 tokenizer = AutoTokenizer.from_pretrained(checkpoint, trust_remote_code=True)
 model = AutoModel.from_pretrained(checkpoint, trust_remote_code=True).to(device)
 
-inputs = tokenizer.encode("def print_hello_world():\tprint('Hello World!')", return_tensors="pt").to(device)
-embedding = model(inputs)[0]
-print(f'Dimension of the embedding: {embedding[0].size()}')
-# Dimension of the embedding: torch.Size([13, 1024])
+inputs = tokenizer("def print_hello_world():\tprint('Hello World!')", return_tensors="pt").to(device)
+embedding = model(**inputs).pooler_output
+print(f'Dimension of the embedding: {embedding.size()}')
+# Dimension of the embedding: torch.Size([1, 1024])
 ```
 
 ### BibTeX entry and citation info