Update code snippet to use sentence-level embeddings; update title

README.md

@@ -7,7 +7,7 @@ language:
 - code
 ---
 
-## CodeSage-Large
+## CodeSage-Base
 
 ### Model description
 CodeSage is a new family of open code embedding models with an encoder architecture that support a wide range of source code understanding tasks. It is introduced in the paper:
@@ -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 (356M 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-base"
@@ -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