code: def select_prompts_based_on_similarity(
      query, dialog_list, prompt_list, topic, tokenizer, encoder, topk
  ):
      """Select samples based on the similarity"""
  
      with torch.no_grad():
          # get the query embeddings
          query_ids = tokenizer.encode(query)
          query_ids = torch.LongTensor([query_ids]).cuda()
          query_emb = encoder(input_ids=query_ids).pooler_output
          query_emb = query_emb[0]
  
          # calculate embeddings for the samples in the database
          if topic in emb_dict:
              example_embeddings = emb_dict[topic]
              example_embeddings = example_embeddings.cuda()
          else:
              for idx, example in enumerate(dialog_list):
                  example_ids = tokenizer.encode(example)
                  example_ids = torch.LongTensor([example_ids]).cuda()
                  example_emb = encoder(
                      input_ids=example_ids
                  ).pooler_output
                  if idx == 0:
                      example_embeddings = example_emb
                  else:
                      example_embeddings = torch.cat(
                          (example_embeddings, example_emb), dim=0
                      )
              emb_dict[topic] = example_embeddings.cpu()
  
          # compare the similarity and select the topk samples
          similarity_list = example_embeddings.matmul(query_emb)
          _, indices = torch.topk(similarity_list, k=topk)
  
      indices = indices.tolist()
      indices = indices[::-1]  # reverse the order
      selected_prompts = []
      for index in indices:
          # index = index.item()
          selected_prompts.append(prompt_list[index])
  
      return selected_prompts
  
file_path: megatron-lm/tasks/msdp/preprocessing.py
imports: import torch
  import argparse
  from nltk import word_tokenize
  from tqdm import tqdm
  import numpy
  import json
label_desc: Use image node_id 5 for calling a global function with example usage: select_prompts_based_on_similarity(query, dialog_list, prompt_list, topic, tokenizer, encoder, topk) and returns: selected_prompts