code: def prompt_selection_for_knowledge_generation(
      test_datapath,
      train_datapath,
      model_path,
      output_prompt_path,
      data_type,
  ):
      """Selecting prompts for the knowledge generation"""
  
      print("> Selecting prompts for the knowledge generation")
  
      (
          train_data_by_topic,
          dialog_data_by_topic,
          dialog_examples,
      ) = get_database(test_datapath, train_datapath, data_type)
  
      from transformers import DPRQuestionEncoderTokenizer
  
      print("> loading tokenizer and encoder")
      tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
          "facebook/dpr-question_encoder-single-nq-base"
      )
      encoder = torch.load(model_path).cuda()
  
      print("> getting dialog embeddings")
      with torch.no_grad():
          for idx, example in tqdm(enumerate(dialog_examples)):
              dialog = example[1]
              dialog_ids = tokenizer.encode(dialog)
              dialog_ids = torch.LongTensor([dialog_ids]).cuda()
              dialog_emb = encoder(input_ids=dialog_ids).pooler_output
  
              if idx == 0:
                  dialog_embeddings = dialog_emb
              else:
                  dialog_embeddings = torch.cat(
                      (dialog_embeddings, dialog_emb), dim=0
                  )
  
      print("> reading test data from %s" % test_datapath)
      prompt_list_for_each_sample = []
      with open(test_datapath, "r") as f:
          for i, line in tqdm(enumerate(f)):
              line = line.strip()
  
              splits = line.split("\t")
              topic = splits[0]
              turns = splits[1].split(" [SEP] ")[-3:]
  
              # get the query sentence
              query_sent = ""
              if data_type != "seen":
                  query_sent += "( " + topic + " ) "
              for i, turn in enumerate(turns):
                  if i != 0:
                      query_sent += " "
                  query_sent += turn
  
              if topic not in train_data_by_topic:
                  # get the query embedding
                  query_ids = tokenizer.encode(query_sent)
                  query_ids = torch.LongTensor([query_ids]).cuda()
                  query_emb = encoder(input_ids=query_ids).pooler_output
                  query_emb = query_emb[0]
  
                  # calculate the similarity
                  similarity_list = dialog_embeddings.matmul(query_emb)
                  _, indices = torch.sort(similarity_list)
                  indices = indices.tolist()
                  selected_topics = {}
                  selected_prompts = []
                  num_prompt = 0
                  for index in indices:
                      example = dialog_examples[index]
                      topic_temp = example[0]
                      if topic_temp not in selected_topics:
                          selected_topics[topic_temp] = True
                          selected_prompts.append(example[2])
                          num_prompt += 1
                          if num_prompt == 10:
                              break
  
                  # get the selected samples
                  example_list = selected_prompts[::-1]
                  key = topic + " " + turns[-1]
                  prompt_list_for_each_sample.append(
                      {key: example_list}
                  )
  
              else:
                  num_data_sample = min(
                      len(train_data_by_topic[topic]), 10
                  )
                  total_example_list = train_data_by_topic[topic]
  
                  dialog_list = dialog_data_by_topic[topic]
                  assert len(dialog_list) == len(
                      train_data_by_topic[topic]
                  )
  
                  # calculate the similarity
                  example_list = select_prompts_based_on_similarity(
                      query_sent,
                      dialog_list,
                      total_example_list,
                      topic,
                      tokenizer,
                      encoder,
                      topk=num_data_sample,
                  )
  
                  key = topic + " " + turns[-1]
                  prompt_list_for_each_sample.append(
                      {key: example_list}
                  )
  
      print("writing to %s" % output_prompt_path)
      with open(output_prompt_path, "w") as f:
          for instance in tqdm(prompt_list_for_each_sample):
              json.dump(instance, f)
              f.write("\n")
  
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 6 for calling a global function with example usage: prompt_selection_for_knowledge_generation(test_datapath, train_datapath, model_path, output_prompt_path, data_type) without return types