| |
|
|
| #include <pybind11/numpy.h> |
| #include <pybind11/pybind11.h> |
|
|
| #include <algorithm> |
| #include <cstdint> |
| #include <functional> |
| #include <queue> |
| #include <tuple> |
| #include <vector> |
|
|
| #ifdef _OPENMP |
| #include <omp.h> |
| #endif |
|
|
| namespace py = pybind11; |
|
|
| namespace { |
|
|
| struct Replay { |
| int32_t dimension; |
| int32_t vocab_size; |
| std::vector<int32_t> offsets; |
| std::vector<int32_t> rule; |
| std::vector<int32_t> anchor; |
|
|
| int32_t child(int32_t left, int32_t right) const { |
| if (left > right) std::swap(left, right); |
| return rule[static_cast<size_t>(left) * vocab_size + right]; |
| } |
|
|
| std::vector<int32_t> encode(const int32_t* code) const { |
| std::vector<int32_t> active; |
| active.reserve(dimension); |
| for (int32_t axis = 0; axis < dimension; ++axis) { |
| if (code[axis] < 0) continue; |
| active.push_back(offsets[axis] + code[axis]); |
| } |
| using Candidate = std::tuple<int32_t, int32_t, int32_t>; |
| std::priority_queue<Candidate, std::vector<Candidate>, std::greater<Candidate>> heap; |
| auto offer = [&](int32_t left, int32_t right) { |
| const int32_t value = child(left, right); |
| if (value >= 0) heap.emplace(value, std::min(left, right), std::max(left, right)); |
| }; |
| for (int32_t a = 0; a < static_cast<int32_t>(active.size()); ++a) { |
| for (int32_t b = a + 1; b < static_cast<int32_t>(active.size()); ++b) { |
| offer(active[a], active[b]); |
| } |
| } |
| while (!heap.empty()) { |
| const auto [child_token, left, right] = heap.top(); |
| heap.pop(); |
| auto left_it = std::find(active.begin(), active.end(), left); |
| auto right_it = std::find(active.begin(), active.end(), right); |
| if (left_it == active.end() || right_it == active.end()) continue; |
| *left_it = child_token; |
| active.erase(right_it); |
| for (const int32_t other : active) { |
| if (other != child_token) offer(child_token, other); |
| } |
| } |
| std::sort(active.begin(), active.end(), [&](int32_t left, int32_t right) { |
| if (anchor[left] != anchor[right]) return anchor[left] < anchor[right]; |
| return left < right; |
| }); |
| return active; |
| } |
| }; |
|
|
| } |
|
|
| py::tuple spatial_encode_instances( |
| const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& codes, |
| const py::array_t<int64_t, py::array::c_style | py::array::forcecast>& row_anchors, |
| const std::vector<int32_t>& bins, |
| const py::array_t<int32_t, py::array::c_style | py::array::forcecast>& merges, |
| int threads) { |
| if (codes.ndim() != 2 || row_anchors.ndim() != 1 || |
| row_anchors.shape(0) != codes.shape(0) || |
| codes.shape(1) != static_cast<int64_t>(bins.size()) || merges.ndim() != 2 || |
| merges.shape(1) != 2) { |
| throw std::invalid_argument("invalid spatial replay inputs"); |
| } |
| Replay replay; |
| replay.dimension = static_cast<int32_t>(bins.size()); |
| replay.offsets.assign(replay.dimension + 1, 0); |
| for (int32_t axis = 0; axis < replay.dimension; ++axis) { |
| replay.offsets[axis + 1] = replay.offsets[axis] + bins[axis]; |
| } |
| const int32_t atoms = replay.offsets.back(); |
| replay.vocab_size = atoms + static_cast<int32_t>(merges.shape(0)); |
| replay.rule.assign(static_cast<size_t>(replay.vocab_size) * replay.vocab_size, -1); |
| replay.anchor.assign(replay.vocab_size, -1); |
| for (int32_t axis = 0; axis < replay.dimension; ++axis) { |
| for (int32_t token = replay.offsets[axis]; token < replay.offsets[axis + 1]; ++token) { |
| replay.anchor[token] = axis; |
| } |
| } |
| auto rules = merges.unchecked<2>(); |
| for (int32_t index = 0; index < merges.shape(0); ++index) { |
| const int32_t left = rules(index, 0), right = rules(index, 1), child = atoms + index; |
| if (left < 0 || right < 0 || left >= child || right >= child) { |
| throw std::invalid_argument("invalid spatial merge rule"); |
| } |
| replay.rule[static_cast<size_t>(std::min(left, right)) * replay.vocab_size + |
| std::max(left, right)] = child; |
| replay.anchor[child] = std::min(replay.anchor[left], replay.anchor[right]); |
| } |
|
|
| const int64_t rows = codes.shape(0); |
| std::vector<int32_t> counts(rows); |
| const int32_t* source = codes.data(); |
| for (int64_t row = 0; row < rows; ++row) { |
| int32_t active = 0; |
| for (int32_t axis = 0; axis < replay.dimension; ++axis) { |
| const int32_t value = source[row * replay.dimension + axis]; |
| if (value < -1 || value >= bins[axis]) { |
| throw std::invalid_argument("a cell index is outside its coordinate's bin count"); |
| } |
| active += value >= 0; |
| } |
| if (active == 0) throw std::invalid_argument("every row needs an active coordinate"); |
| } |
| #ifndef _OPENMP |
| (void)threads; |
| #endif |
| { |
| py::gil_scoped_release release; |
| #ifdef _OPENMP |
| #pragma omp parallel for schedule(static) num_threads(threads > 0 ? threads : omp_get_max_threads()) |
| #endif |
| for (int64_t row = 0; row < rows; ++row) { |
| counts[row] = static_cast<int32_t>( |
| replay.encode(source + row * replay.dimension).size()); |
| } |
| } |
| std::vector<int64_t> offsets(rows + 1, 0); |
| for (int64_t row = 0; row < rows; ++row) offsets[row + 1] = offsets[row] + counts[row]; |
| py::array_t<int32_t> tokens(offsets.back()); |
| py::array_t<int64_t> anchors(offsets.back()); |
| int32_t* output_tokens = tokens.mutable_data(); |
| int64_t* output_anchors = anchors.mutable_data(); |
| const int64_t* input_anchors = row_anchors.data(); |
| { |
| py::gil_scoped_release release; |
| #ifdef _OPENMP |
| #pragma omp parallel for schedule(static) num_threads(threads > 0 ? threads : omp_get_max_threads()) |
| #endif |
| for (int64_t row = 0; row < rows; ++row) { |
| const auto encoded = replay.encode(source + row * replay.dimension); |
| const int64_t begin = offsets[row]; |
| for (int32_t index = 0; index < static_cast<int32_t>(encoded.size()); ++index) { |
| output_tokens[begin + index] = encoded[index]; |
| output_anchors[begin + index] = input_anchors[row]; |
| } |
| } |
| } |
| return py::make_tuple(tokens, anchors); |
| } |
|
|
