Update README.md

README.md

@@ -63,7 +63,7 @@ from transformers import AutoTokenizer, AutoModel
 tokenizer = AutoTokenizer.from_pretrained("nomic-ai/nomic-embed-code")
 model = AutoModel.from_pretrained("nomic-ai/nomic-embed-code")
 
-def last_token_pooling(model_output, attention_mask):
+def last_token_pooling(hidden_states, attention_mask):
     sequence_lengths = attention_mask.sum(-1) - 1
     return hidden_states[torch.arange(hidden_states.shape[0]), sequence_lengths]
 
@@ -74,7 +74,8 @@ code_snippets = queries + codes
 encoded_input = tokenizer(code_snippets, padding=True, truncation=True, return_tensors='pt')
 model.eval()
 with torch.no_grad():
-    model_output = model(**encoded_input)
+    model_output = model(**encoded_input)[0]
+
 embeddings = last_token_pooling(model_output, encoded_input['attention_mask'])
 embeddings = F.normalize(embeddings, p=2, dim=1)
 print(embeddings.shape)
@@ -95,7 +96,7 @@ model = SentenceTransformer("nomic-ai/nomic-embed-code")
 query_emb = model.encode(queries, prompt_name="query")
 code_emb = model.encode(code_snippets)
 
-similarity = model.similarity(query_emb, code_emb)
+similarity = model.similarity(query_emb[0], code_emb[0])
 print(similarity)
 ```