Update README.md

README.md

@@ -107,10 +107,7 @@ contexts = [
 
 ## convert query into a format as follows:
 ## user: {user}\nagent: {agent}\nuser: {user}
-formatted_query = ""
-for turn in query:
-    formatted_query += turn['role'] + ": " + turn['content'] + "\n"
-formatted_query = formatted_query.strip()
+formatted_query = '\n'.join([turn['role'] + ": " + turn['content'] for turn in messages]).strip()
 
 ## get query and context embeddings
 query_input = tokenizer(formatted_query, return_tensors='pt')
@@ -118,9 +115,11 @@ ctx_input = tokenizer(contexts, padding=True, return_tensors='pt')
 query_emb = query_encoder(**query_input).last_hidden_state[:, 0, :]
 ctx_emb = context_encoder(**ctx_input).last_hidden_state[:, 0, :]
 
-# Compute similarity scores using dot product
-score1 = query_emb @ ctx_emb[0]
-score2 = query_emb @ ctx_emb[1]
+## Compute similarity scores using dot product
+similarities = query_emb.matmul(ctx_emb.transpose(0, 1)) # (1, num_ctx)
+
+## rank the similarity (from highest to lowest)
+ranked_results = torch.argsort(similarities, dim=-1, descending=True) # (1, num_ctx)
 ```
 
 ## License