Update README.md

README.md

@@ -107,30 +107,27 @@ 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')
 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, :]
+query_emb = query_encoder(**query_input).last_hidden_state[:, 0, :] # (1, emb_dim)
+ctx_emb = context_encoder(**ctx_input).last_hidden_state[:, 0, :] # (num_ctx, emb_dim)
 
-# 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
 Dragon-multiturn is built on top of [Dragon](https://arxiv.org/abs/2302.07452). We refer users to the original license of the Dragon model.
 
-
 ## Correspondence to
 Zihan Liu (zihanl@nvidia.com), Wei Ping (wping@nvidia.com)
 
-
 ## Citation
 <pre>
 @article{liu2024chatqa,