1 /* 2 * Copyright (c) Meta Platforms, Inc. and affiliates. 3 * All rights reserved. 4 * 5 * This source code is licensed under the BSD-style license found in the 6 * LICENSE file in the root directory of this source tree. 7 */ 8 9 #pragma once 10 11 #include <algorithm> 12 #include <vector> 13 14 class BasicSampler { 15 public: sample(const std::vector<float> & logits)16 int64_t sample(const std::vector<float>& logits) { 17 // Find the token with the highest log probability. 18 int64_t max_index = 19 std::max_element(logits.begin(), logits.end()) - logits.begin(); 20 return max_index; 21 } 22 }; 23