Update README.md

README.md

@@ -73,7 +73,7 @@ model = AutoModel.from_pretrained('cmarkea/bloomz-3b-retriever-v2')
 
 def infer(txt: Union[str, List[str]]):
     tok = tokenizer(txt, padding=True, return_tensors='pt')
-    with torch.no_grad():
+    with torch.inference_mode():
         embedding = model(**tok)
     # Inportant: take only last token!
     return embedding.get('last_hidden_state')[:,-1,:].numpy()
@@ -104,7 +104,7 @@ from scipy.spatial.distance import cdist
 retriever = pipeline('feature-extraction', 'cmarkea/bloomz-3b-retriever-v2')
 
 # Inportant: take only last token!
-infer = lambda x: [ii[0][-1] for ii in retriever(x)]
+infer = lambda x: [np.array(ii[0][-1]).reshape(1,-1) for ii in retriever(x)]
 
 list_of_contexts: List[str] = [...]
 emb_contexts = np.concatenate(infer(list_of_contexts), axis=0)