cudf/cpp/tests/large_strings/bpe_tests.cpp at 66d2171b36ea9183ad3bf76fea3c41419c2cdcdf · PointKernel/cudf · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
/*
 * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 */

#include "large_strings_fixture.hpp"

#include <tests/text/bpe_data_generator.hpp>

#include <cudf_test/byte_literals.hpp>
#include <cudf_test/column_utilities.hpp>
#include <cudf_test/column_wrapper.hpp>
#include <cudf_test/large_memory.hpp>

#include <cudf/column/column_factories.hpp>
#include <cudf/concatenate.hpp>
#include <cudf/copying.hpp>
#include <cudf/scalar/scalar.hpp>
#include <cudf/strings/strings_column_view.hpp>

#include <nvtext/byte_pair_encoding.hpp>

#include <algorithm>
#include <limits>
#include <string>
#include <vector>

struct BPETest : public cudf::test::StringsLargeTest {};

namespace {
std::vector<cudf::size_type> make_splits(cudf::size_type rows, int multiplier)
{
  std::vector<cudf::size_type> splits;
  for (int n = 1; n < multiplier; ++n) {
    splits.push_back(rows * n);
  }
  return splits;
}
}  // namespace

// Requires about 8GB of device memory.
TEST_F(BPETest, LargeOutput)
{
  // synthetic GPT-2 style table and byte-mapped text
  namespace bpe = cudf::test::bpe;
  bpe::text_generator gen(3);
  auto const text   = bpe::gpt2_byte_map(gen.generate(4'000'000));
  auto const merges = bpe::train_merges(gen.generate(1'000'000), 4000);
  std::vector<std::string> rows;
  std::size_t pos = 0;
  while (pos < text.size()) {
    auto end = std::min(pos + 4096, text.size());
    while (end < text.size() && (static_cast<unsigned char>(text[end]) & 0xC0) == 0x80) {
      ++end;
    }
    rows.push_back(text.substr(pos, end - pos));
    pos = end;
  }
  auto const input    = cudf::test::strings_column_wrapper(rows.begin(), rows.end());
  auto const mpt      = cudf::test::strings_column_wrapper(merges.begin(), merges.end());
  auto const mp       = nvtext::load_merge_pairs(cudf::strings_column_view(mpt));
  auto const expected = nvtext::byte_pair_encoding(cudf::strings_column_view(input), *mp);

  // replicate the input so the input and output chars exceed max size_type
  auto const view             = cudf::column_view(input);
  int constexpr max_size_type = std::numeric_limits<cudf::size_type>::max();
  auto const multiplier =
    static_cast<int>(max_size_type /
                     cudf::strings_column_view(view).chars_size(cudf::get_default_stream())) +
    1;
  auto const large_input = cudf::concatenate(std::vector<cudf::column_view>(multiplier, view));
  auto const sv          = cudf::strings_column_view(large_input->view());
  EXPECT_EQ(sv.offsets().type(), cudf::data_type{cudf::type_id::INT64});

  auto const result = nvtext::byte_pair_encoding(sv, *mp);
  EXPECT_EQ(cudf::strings_column_view(result->view()).offsets().type(),
            cudf::data_type{cudf::type_id::INT64});
  for (auto const& c : cudf::split(result->view(), make_splits(view.size(), multiplier))) {
    CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(c, expected->view());
  }
}

TEST_F(BPETest, ManySegments)
{
  using namespace cudf::test::literals;
  CUDF_TEST_REQUIRE_DEVICE_MEMORY(40_GiB);
  // every character is its own (unpairable) segment so the number
  // of segments exceeds max size_type
  auto const mpt = cudf::test::strings_column_wrapper({"a b", "ab c"});
  auto const mp  = nvtext::load_merge_pairs(cudf::strings_column_view(mpt));

  auto const rows     = std::vector<std::string>(1000, std::string(1000, 'x'));
  auto const input    = cudf::test::strings_column_wrapper(rows.begin(), rows.end());
  std::string encoded = "x";
  for (int i = 1; i < 1000; ++i) {
    encoded += " x";
  }
  auto const encoded_rows = std::vector<std::string>(1000, encoded);
  auto const expected =
    cudf::test::strings_column_wrapper(encoded_rows.begin(), encoded_rows.end());

  auto const view             = cudf::column_view(input);
  int constexpr max_size_type = std::numeric_limits<cudf::size_type>::max();
  auto const multiplier =
    static_cast<int>(max_size_type /
                     cudf::strings_column_view(view).chars_size(cudf::get_default_stream())) +
    1;
  auto const large_input = cudf::concatenate(std::vector<cudf::column_view>(multiplier, view));

  auto const result =
    nvtext::byte_pair_encoding(cudf::strings_column_view(large_input->view()), *mp);
  for (auto const& c : cudf::split(result->view(), make_splits(view.size(), multiplier))) {
    CUDF_TEST_EXPECT_COLUMNS_EQUIVALENT(c, expected);
  }
}