Update README.md

README.md

@@ -91,29 +91,22 @@ How to Use Bloomz-3b-reranking
 The following example utilizes the API Pipeline of the Transformers library.
 
 ```python
-import numpy as np
 from transformers import pipeline
-from scipy.spatial.distance import cdist
 
-retriever = pipeline('feature-extraction', 'cmarkea/bloomz-3b-retriever')
-
-# Inportant: take only last token!
-infer = lambda x: [ii[0][-1] for ii in retriever(x)]
-
-list_of_contexts = [...]
-emb_contexts = np.concatenate(infer(list_of_contexts), axis=0)
-list_of_queries = [...]
-emb_queries = np.concatenate(infer(list_of_queries), axis=0)
-
-# Important: take l2 distance!
-dist = cdist(emb_queries, emb_contexts, 'euclidean') 
-top_k = lambda x: [
-    [list_of_contexts[qq] for qq in ii]
-    for ii in dist.argsort(axis=-1)[:,:x]
-]
-
-# top 5 nearest contexts for each queries
-top_contexts = top_k(5)
+reranker = pipeline('feature-extraction', 'cmarkea/bloomz-3b-retriever', top_k=None)
+
+similarity = reranker(
+    [dict(text=ii, text_pair=query) for ii in context_list]
+)
+context_reranked = sorted(
+    filter(lambda x: x[0]['label'] == "LABEL_1", zip(similarity, context_list)),
+    key=lambda x: x[0]
+)
+score, context_cleaned = zip(
+    *filter(
+        lambda x: x[0] >= 0.8
+    )
+)
 ```
 
 Citation