Function run_passkey

src/main.cpp:433–512

}

void run_passkey(
  Session& session,
  const int n_junk,
  const int passkey_pos
) {
  Model& model = session.model;
  InferenceState& state = session.state;
  Sampler& sampler = session.sampler;
  Tokenizer& tokenizer = session.tokenizer;

  std::cout << "Model active bytes with full context window: " << model.active_bytes(model.config->max_seq_len) << std::endl;

  const std::string PROMPT_PREFIX = 
    "There is an important info hidden inside a lot of irrelevant text. "
    "Find it and memorize them. I will quiz you about the important information there.";
  const std::string PROMPT_SUFFIX = " What is the pass key? The pass key is";

  const int passkey = std::rand() % 50000 + 1;
  const int pos = passkey_pos == -1 ? std::rand() % n_junk : passkey_pos;

  std::string prompt = PROMPT_PREFIX;
  for (int i = 0; i < n_junk; i++) {
    if (i % n_junk == pos) {
      prompt += " The pass key is " + std::to_string(passkey) + ". Remember it. " + std::to_string(passkey) + " is the pass key.";
    }
    prompt += " The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again.";
  }
  prompt += PROMPT_SUFFIX;

  std::vector<int> encoding;
  {
    uint64_t encode_start_ms = get_timestamp_ms();
    encoding = tokenizer.encode(prompt, true);
    uint64_t encode_end_ms = get_timestamp_ms();

    uint64_t encoding_ms = encode_end_ms - encode_start_ms;
    std::cout << fmt::format(
      "Encoding stats: ({} tokens, throughput: {:.5}tok/s, latency: {:.5}s/tok, total: {:.5}s)\n",
      encoding.size(),
      encoding.size() / (encoding_ms / 1000.0),
      (encoding_ms / 1000.0) / encoding.size(),
      encoding_ms / 1000.0
    ) << std::endl;
  }

  // Allow max 16 steps to generate passkey
  const size_t MAX_GENERATION_STEPS = 16;

  std::cout << fmt::format(
    "Passkey test:\n"
    "  prompt: {} tokens\n"
    "  passkey: {}\n"
    "  passkey token index: ~{}\n",
    encoding.size(),
    passkey,
    (int)(((float)pos) / n_junk * encoding.size())
  ) << std::endl;