chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,47 @@
|
||||
/*
|
||||
Copyright (c) Microsoft Corporation.
|
||||
Licensed under the MIT License.
|
||||
*/
|
||||
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
/*
|
||||
CPP Binding for CUDA OP
|
||||
*/
|
||||
|
||||
// CUDA forward declarations
|
||||
torch::Tensor ngram_repeat_block_cuda_forward(torch::Tensor tokens,
|
||||
torch::Tensor lprobs, int bsz,
|
||||
int step, int beam_size,
|
||||
int no_repeat_ngram_size);
|
||||
|
||||
#define CHECK_CUDA(x) \
|
||||
TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) \
|
||||
TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INPUT(x) \
|
||||
CHECK_CUDA(x); \
|
||||
CHECK_CONTIGUOUS(x)
|
||||
|
||||
// Input check and call to CUDA OP
|
||||
// Backward method not required
|
||||
torch::Tensor ngram_repeat_block_forward(torch::Tensor tokens,
|
||||
torch::Tensor lprobs, int bsz,
|
||||
int step, int beam_size,
|
||||
int no_repeat_ngram_size) {
|
||||
CHECK_INPUT(tokens);
|
||||
CHECK_INPUT(lprobs);
|
||||
assert(bsz > 0);
|
||||
assert(step >= 0);
|
||||
assert(beam_size > 0);
|
||||
assert(no_repeat_ngram_size > 0);
|
||||
|
||||
return ngram_repeat_block_cuda_forward(tokens, lprobs, bsz, step, beam_size,
|
||||
no_repeat_ngram_size);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("forward", &ngram_repeat_block_forward,
|
||||
"No Repeat Ngram Block forward (CUDA)");
|
||||
}
|
||||
@@ -0,0 +1,76 @@
|
||||
/*
|
||||
Copyright (c) Microsoft Corporation.
|
||||
Licensed under the MIT License.
|
||||
*/
|
||||
|
||||
/*
|
||||
Kernel implementation for blocking repeated n-grams.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <math.h>
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
|
||||
// Ban repeated ngrams of length = 'no_repeat_ngram_size'
|
||||
__global__ void banRepeatedTokens(long* __restrict__ tokens,
|
||||
float* __restrict__ lprobs,
|
||||
int max_predict_len, int vocab_size,
|
||||
int no_repeat_ngram_size) {
|
||||
auto row = blockIdx.x;
|
||||
auto col = threadIdx.x;
|
||||
auto start = row * (max_predict_len) + col;
|
||||
// Each thread compares ngram starting from
|
||||
// thread index with final ngram starting from
|
||||
// step - no_repeat_ngram_size +2
|
||||
auto check_start_pos = blockDim.x;
|
||||
auto lprob_start = row * vocab_size;
|
||||
bool is_banned = true;
|
||||
extern __shared__ long tokens_shm[];
|
||||
tokens_shm[col] = tokens[start];
|
||||
if (col == blockDim.x - 1) {
|
||||
for (int i=1; i<no_repeat_ngram_size; i++){
|
||||
if (col+i < max_predict_len){
|
||||
tokens_shm[col + i] = tokens[start + i];
|
||||
}
|
||||
}
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
for (int k = 0; k < no_repeat_ngram_size - 1; k++) {
|
||||
if (tokens_shm[col + k] != tokens_shm[check_start_pos + k]) {
|
||||
is_banned = false;
|
||||
}
|
||||
}
|
||||
if (is_banned == true) {
|
||||
auto token_to_be_banned = tokens_shm[col + no_repeat_ngram_size - 1];
|
||||
lprobs[lprob_start + token_to_be_banned] = -INFINITY;
|
||||
}
|
||||
}
|
||||
|
||||
// Allocate blocks and threads based on
|
||||
// batch size and sequence length and launch
|
||||
// kernel
|
||||
torch::Tensor ngram_repeat_block_cuda_forward(const torch::Tensor tokens,
|
||||
torch::Tensor lprobs, int bsz,
|
||||
int step, int beam_size,
|
||||
int no_repeat_ngram_size) {
|
||||
int threads = step - no_repeat_ngram_size + 2;
|
||||
if (threads <= 0) return lprobs;
|
||||
int max_predict_len = tokens.size(1);
|
||||
int vocab_size = lprobs.size(1);
|
||||
auto token_ptr = tokens.data_ptr<long>();
|
||||
auto lprob_ptr = lprobs.data_ptr<float>();
|
||||
int blocks = bsz * beam_size;
|
||||
int shared_mem_size = (step + 1) * sizeof(long);
|
||||
|
||||
// Launching N blocks where N is number of samples in a batch (beams*bsz)
|
||||
// Launching T threads where T is number of previous ngrams in a sample
|
||||
// Allocating shared mem per block for fastser access of input tokens since
|
||||
// each token will be accessed N times to compare with current Ngram where
|
||||
// N is Ngram size.
|
||||
banRepeatedTokens<<<blocks, threads, shared_mem_size>>>(
|
||||
token_ptr, lprob_ptr, max_predict_len, vocab_size, no_repeat_ngram_size);
|
||||
return lprobs;
|
||||
}
|
||||
@@ -0,0 +1,141 @@
|
||||
/**
|
||||
* Copyright 2017-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This source code is licensed under the license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <map>
|
||||
#include <array>
|
||||
#include <cstring>
|
||||
#include <cstdio>
|
||||
|
||||
typedef struct
|
||||
{
|
||||
size_t reflen;
|
||||
size_t predlen;
|
||||
size_t match1;
|
||||
size_t count1;
|
||||
size_t match2;
|
||||
size_t count2;
|
||||
size_t match3;
|
||||
size_t count3;
|
||||
size_t match4;
|
||||
size_t count4;
|
||||
} bleu_stat;
|
||||
|
||||
// left trim (remove pad)
|
||||
void bleu_ltrim(size_t* len, int** sent, int pad) {
|
||||
size_t start = 0;
|
||||
while(start < *len) {
|
||||
if (*(*sent + start) != pad) { break; }
|
||||
start++;
|
||||
}
|
||||
*sent += start;
|
||||
*len -= start;
|
||||
}
|
||||
|
||||
// right trim remove (eos)
|
||||
void bleu_rtrim(size_t* len, int** sent, int pad, int eos) {
|
||||
size_t end = *len - 1;
|
||||
while (end > 0) {
|
||||
if (*(*sent + end) != eos && *(*sent + end) != pad) { break; }
|
||||
end--;
|
||||
}
|
||||
*len = end + 1;
|
||||
}
|
||||
|
||||
// left and right trim
|
||||
void bleu_trim(size_t* len, int** sent, int pad, int eos) {
|
||||
bleu_ltrim(len, sent, pad);
|
||||
bleu_rtrim(len, sent, pad, eos);
|
||||
}
|
||||
|
||||
size_t bleu_hash(int len, int* data) {
|
||||
size_t h = 14695981039346656037ul;
|
||||
size_t prime = 0x100000001b3;
|
||||
char* b = (char*) data;
|
||||
size_t blen = sizeof(int) * len;
|
||||
|
||||
while (blen-- > 0) {
|
||||
h ^= *b++;
|
||||
h *= prime;
|
||||
}
|
||||
|
||||
return h;
|
||||
}
|
||||
|
||||
void bleu_addngram(
|
||||
size_t *ntotal, size_t *nmatch, size_t n,
|
||||
size_t reflen, int* ref, size_t predlen, int* pred) {
|
||||
|
||||
if (predlen < n) { return; }
|
||||
|
||||
predlen = predlen - n + 1;
|
||||
(*ntotal) += predlen;
|
||||
|
||||
if (reflen < n) { return; }
|
||||
|
||||
reflen = reflen - n + 1;
|
||||
|
||||
std::map<size_t, size_t> count;
|
||||
while (predlen > 0) {
|
||||
size_t w = bleu_hash(n, pred++);
|
||||
count[w]++;
|
||||
predlen--;
|
||||
}
|
||||
|
||||
while (reflen > 0) {
|
||||
size_t w = bleu_hash(n, ref++);
|
||||
if (count[w] > 0) {
|
||||
(*nmatch)++;
|
||||
count[w] -=1;
|
||||
}
|
||||
reflen--;
|
||||
}
|
||||
}
|
||||
|
||||
extern "C" {
|
||||
|
||||
#ifdef _WIN64
|
||||
__declspec(dllexport)
|
||||
#endif
|
||||
void bleu_zero_init(bleu_stat* stat) {
|
||||
std::memset(stat, 0, sizeof(bleu_stat));
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
__declspec(dllexport)
|
||||
#endif
|
||||
void bleu_one_init(bleu_stat* stat) {
|
||||
bleu_zero_init(stat);
|
||||
stat->count1 = 0;
|
||||
stat->count2 = 1;
|
||||
stat->count3 = 1;
|
||||
stat->count4 = 1;
|
||||
stat->match1 = 0;
|
||||
stat->match2 = 1;
|
||||
stat->match3 = 1;
|
||||
stat->match4 = 1;
|
||||
}
|
||||
|
||||
#ifdef _WIN64
|
||||
__declspec(dllexport)
|
||||
#endif
|
||||
void bleu_add(
|
||||
bleu_stat* stat,
|
||||
size_t reflen, int* ref, size_t predlen, int* pred, int pad, int eos) {
|
||||
|
||||
bleu_trim(&reflen, &ref, pad, eos);
|
||||
bleu_trim(&predlen, &pred, pad, eos);
|
||||
stat->reflen += reflen;
|
||||
stat->predlen += predlen;
|
||||
|
||||
bleu_addngram(&stat->count1, &stat->match1, 1, reflen, ref, predlen, pred);
|
||||
bleu_addngram(&stat->count2, &stat->match2, 2, reflen, ref, predlen, pred);
|
||||
bleu_addngram(&stat->count3, &stat->match3, 3, reflen, ref, predlen, pred);
|
||||
bleu_addngram(&stat->count4, &stat->match4, 4, reflen, ref, predlen, pred);
|
||||
}
|
||||
|
||||
}
|
||||
@@ -0,0 +1,37 @@
|
||||
/**
|
||||
* Copyright 2017-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This source code is licensed under the license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <Python.h>
|
||||
|
||||
|
||||
static PyMethodDef method_def[] = {
|
||||
{NULL, NULL, 0, NULL}
|
||||
};
|
||||
|
||||
static struct PyModuleDef module_def = {
|
||||
PyModuleDef_HEAD_INIT,
|
||||
"libbleu", /* name of module */
|
||||
NULL, /* module documentation, may be NULL */
|
||||
-1, /* size of per-interpreter state of the module,
|
||||
or -1 if the module keeps state in global variables. */
|
||||
method_def
|
||||
};
|
||||
|
||||
|
||||
#if PY_MAJOR_VERSION == 2
|
||||
PyMODINIT_FUNC init_libbleu()
|
||||
#else
|
||||
PyMODINIT_FUNC PyInit_libbleu()
|
||||
#endif
|
||||
{
|
||||
PyObject *m = PyModule_Create(&module_def);
|
||||
if (!m) {
|
||||
return NULL;
|
||||
}
|
||||
return m;
|
||||
}
|
||||
@@ -0,0 +1,231 @@
|
||||
/**
|
||||
* Copyright 2017-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This source code is licensed under the license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include <torch/torch.h> // @manual=//caffe2:torch_extension
|
||||
#include <pybind11/detail/common.h>
|
||||
#include <pybind11/pybind11.h>
|
||||
#include <vector>
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <iosfwd>
|
||||
#include <memory>
|
||||
#include <new>
|
||||
#include <string>
|
||||
#include <utility>
|
||||
|
||||
using namespace ::std;
|
||||
|
||||
vector<vector<uint32_t>> edit_distance2_with_dp(
|
||||
vector<uint32_t>& x,
|
||||
vector<uint32_t>& y) {
|
||||
uint32_t lx = x.size();
|
||||
uint32_t ly = y.size();
|
||||
vector<vector<uint32_t>> d(lx + 1, vector<uint32_t>(ly + 1));
|
||||
for (uint32_t i = 0; i < lx + 1; i++) {
|
||||
d[i][0] = i;
|
||||
}
|
||||
for (uint32_t j = 0; j < ly + 1; j++) {
|
||||
d[0][j] = j;
|
||||
}
|
||||
for (uint32_t i = 1; i < lx + 1; i++) {
|
||||
for (uint32_t j = 1; j < ly + 1; j++) {
|
||||
d[i][j] =
|
||||
min(min(d[i - 1][j], d[i][j - 1]) + 1,
|
||||
d[i - 1][j - 1] + 2 * (x.at(i - 1) == y.at(j - 1) ? 0 : 1));
|
||||
}
|
||||
}
|
||||
return d;
|
||||
}
|
||||
|
||||
vector<vector<uint32_t>> edit_distance2_backtracking(
|
||||
vector<vector<uint32_t>>& d,
|
||||
vector<uint32_t>& x,
|
||||
vector<uint32_t>& y,
|
||||
uint32_t terminal_symbol) {
|
||||
vector<uint32_t> seq;
|
||||
vector<vector<uint32_t>> edit_seqs(x.size() + 2, vector<uint32_t>());
|
||||
/*
|
||||
edit_seqs:
|
||||
0~x.size() cell is the insertion sequences
|
||||
last cell is the delete sequence
|
||||
*/
|
||||
|
||||
if (x.size() == 0) {
|
||||
edit_seqs.at(0) = y;
|
||||
return edit_seqs;
|
||||
}
|
||||
|
||||
uint32_t i = d.size() - 1;
|
||||
uint32_t j = d.at(0).size() - 1;
|
||||
|
||||
while ((i >= 0) && (j >= 0)) {
|
||||
if ((i == 0) && (j == 0)) {
|
||||
break;
|
||||
}
|
||||
|
||||
if ((j > 0) && (d.at(i).at(j - 1) < d.at(i).at(j))) {
|
||||
seq.push_back(1); // insert
|
||||
seq.push_back(y.at(j - 1));
|
||||
j--;
|
||||
} else if ((i > 0) && (d.at(i - 1).at(j) < d.at(i).at(j))) {
|
||||
seq.push_back(2); // delete
|
||||
seq.push_back(x.at(i - 1));
|
||||
i--;
|
||||
} else {
|
||||
seq.push_back(3); // keep
|
||||
seq.push_back(x.at(i - 1));
|
||||
i--;
|
||||
j--;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t prev_op, op, s, word;
|
||||
prev_op = 0, s = 0;
|
||||
for (uint32_t k = 0; k < seq.size() / 2; k++) {
|
||||
op = seq.at(seq.size() - 2 * k - 2);
|
||||
word = seq.at(seq.size() - 2 * k - 1);
|
||||
if (prev_op != 1) {
|
||||
s++;
|
||||
}
|
||||
if (op == 1) // insert
|
||||
{
|
||||
edit_seqs.at(s - 1).push_back(word);
|
||||
} else if (op == 2) // delete
|
||||
{
|
||||
edit_seqs.at(x.size() + 1).push_back(1);
|
||||
} else {
|
||||
edit_seqs.at(x.size() + 1).push_back(0);
|
||||
}
|
||||
|
||||
prev_op = op;
|
||||
}
|
||||
|
||||
for (uint32_t k = 0; k < edit_seqs.size(); k++) {
|
||||
if (edit_seqs[k].size() == 0) {
|
||||
edit_seqs[k].push_back(terminal_symbol);
|
||||
}
|
||||
}
|
||||
return edit_seqs;
|
||||
}
|
||||
|
||||
vector<vector<uint32_t>> edit_distance2_backtracking_with_delete(
|
||||
vector<vector<uint32_t>>& d,
|
||||
vector<uint32_t>& x,
|
||||
vector<uint32_t>& y,
|
||||
uint32_t terminal_symbol,
|
||||
uint32_t deletion_symbol) {
|
||||
vector<uint32_t> seq;
|
||||
vector<vector<uint32_t>> edit_seqs(x.size() + 1, vector<uint32_t>());
|
||||
/*
|
||||
edit_seqs:
|
||||
0~x.size() cell is the insertion sequences
|
||||
last cell is the delete sequence
|
||||
*/
|
||||
|
||||
if (x.size() == 0) {
|
||||
edit_seqs.at(0) = y;
|
||||
return edit_seqs;
|
||||
}
|
||||
|
||||
uint32_t i = d.size() - 1;
|
||||
uint32_t j = d.at(0).size() - 1;
|
||||
|
||||
while ((i >= 0) && (j >= 0)) {
|
||||
if ((i == 0) && (j == 0)) {
|
||||
break;
|
||||
}
|
||||
|
||||
if ((j > 0) && (d.at(i).at(j - 1) < d.at(i).at(j))) {
|
||||
seq.push_back(1); // insert
|
||||
seq.push_back(y.at(j - 1));
|
||||
j--;
|
||||
} else if ((i > 0) && (d.at(i - 1).at(j) < d.at(i).at(j))) {
|
||||
seq.push_back(2); // delete
|
||||
seq.push_back(x.at(i - 1));
|
||||
i--;
|
||||
} else {
|
||||
seq.push_back(3); // keep
|
||||
seq.push_back(x.at(i - 1));
|
||||
i--;
|
||||
j--;
|
||||
}
|
||||
}
|
||||
|
||||
uint32_t prev_op, op, s, word;
|
||||
prev_op = 0, s = 0;
|
||||
for (uint32_t k = 0; k < seq.size() / 2; k++) {
|
||||
op = seq.at(seq.size() - 2 * k - 2);
|
||||
word = seq.at(seq.size() - 2 * k - 1);
|
||||
if (prev_op != 1) {
|
||||
s++;
|
||||
}
|
||||
if (op == 1) // insert
|
||||
{
|
||||
edit_seqs.at(s - 1).push_back(word);
|
||||
} else if (op == 2) // delete
|
||||
{
|
||||
edit_seqs.at(s - 1).push_back(deletion_symbol);
|
||||
}
|
||||
|
||||
prev_op = op;
|
||||
}
|
||||
|
||||
for (uint32_t k = 0; k < edit_seqs.size(); k++) {
|
||||
if (edit_seqs.at(k).size() == 0) {
|
||||
edit_seqs.at(k).push_back(terminal_symbol);
|
||||
}
|
||||
}
|
||||
return edit_seqs;
|
||||
}
|
||||
|
||||
vector<uint32_t> compute_ed2(
|
||||
vector<vector<uint32_t>>& xs,
|
||||
vector<vector<uint32_t>>& ys) {
|
||||
vector<uint32_t> distances(xs.size());
|
||||
for (uint32_t i = 0; i < xs.size(); i++) {
|
||||
vector<vector<uint32_t>> d = edit_distance2_with_dp(xs.at(i), ys.at(i));
|
||||
distances.at(i) = d.at(xs.at(i).size()).at(ys.at(i).size());
|
||||
}
|
||||
return distances;
|
||||
}
|
||||
|
||||
vector<vector<vector<uint32_t>>> suggested_ed2_path(
|
||||
vector<vector<uint32_t>>& xs,
|
||||
vector<vector<uint32_t>>& ys,
|
||||
uint32_t terminal_symbol) {
|
||||
vector<vector<vector<uint32_t>>> seq(xs.size());
|
||||
for (uint32_t i = 0; i < xs.size(); i++) {
|
||||
vector<vector<uint32_t>> d = edit_distance2_with_dp(xs.at(i), ys.at(i));
|
||||
seq.at(i) =
|
||||
edit_distance2_backtracking(d, xs.at(i), ys.at(i), terminal_symbol);
|
||||
}
|
||||
return seq;
|
||||
}
|
||||
|
||||
vector<vector<vector<uint32_t>>> suggested_ed2_path_with_delete(
|
||||
vector<vector<uint32_t>>& xs,
|
||||
vector<vector<uint32_t>>& ys,
|
||||
uint32_t terminal_symbol,
|
||||
uint32_t deletion_symbol) {
|
||||
vector<vector<vector<uint32_t>>> seq(xs.size());
|
||||
for (uint32_t i = 0; i < xs.size(); i++) {
|
||||
vector<vector<uint32_t>> d = edit_distance2_with_dp(xs.at(i), ys.at(i));
|
||||
seq.at(i) = edit_distance2_backtracking_with_delete(
|
||||
d, xs.at(i), ys.at(i), terminal_symbol, deletion_symbol);
|
||||
}
|
||||
return seq;
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(libnat, m) {
|
||||
m.def("compute_ed2", &compute_ed2, "compute_ed2");
|
||||
m.def("suggested_ed2_path", &suggested_ed2_path, "suggested_ed2_path");
|
||||
m.def(
|
||||
"suggested_ed2_path_with_delete",
|
||||
&suggested_ed2_path_with_delete,
|
||||
"suggested_ed2_path_with_delete");
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
/**
|
||||
* Copyright 2017-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This source code is licensed under the license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
/*
|
||||
This code is partially adpoted from https://github.com/1ytic/pytorch-edit-distance
|
||||
*/
|
||||
|
||||
#include "edit_dist.h"
|
||||
#include <torch/types.h>
|
||||
|
||||
#ifndef TORCH_CHECK
|
||||
#define TORCH_CHECK AT_CHECK
|
||||
#endif
|
||||
|
||||
#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
|
||||
#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
|
||||
#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
|
||||
|
||||
|
||||
torch::Tensor LevenshteinDistance(
|
||||
torch::Tensor source,
|
||||
torch::Tensor target,
|
||||
torch::Tensor source_length,
|
||||
torch::Tensor target_length) {
|
||||
|
||||
CHECK_INPUT(source);
|
||||
CHECK_INPUT(target);
|
||||
CHECK_INPUT(source_length);
|
||||
CHECK_INPUT(target_length);
|
||||
return LevenshteinDistanceCuda(source, target, source_length, target_length);
|
||||
}
|
||||
|
||||
torch::Tensor GenerateDeletionLabel(
|
||||
torch::Tensor source,
|
||||
torch::Tensor operations) {
|
||||
|
||||
CHECK_INPUT(source);
|
||||
CHECK_INPUT(operations);
|
||||
return GenerateDeletionLabelCuda(source, operations);
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> GenerateInsertionLabel(
|
||||
torch::Tensor target,
|
||||
torch::Tensor operations) {
|
||||
|
||||
CHECK_INPUT(target);
|
||||
CHECK_INPUT(operations);
|
||||
return GenerateInsertionLabelCuda(target, operations);
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
m.def("levenshtein_distance", &LevenshteinDistance, "Levenshtein distance");
|
||||
m.def("generate_deletion_labels", &GenerateDeletionLabel, "Generate Deletion Label");
|
||||
m.def("generate_insertion_labels", &GenerateInsertionLabel, "Generate Insertion Label");
|
||||
}
|
||||
@@ -0,0 +1,332 @@
|
||||
/**
|
||||
* Copyright 2017-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This source code is licensed under the license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#include "edit_dist.h"
|
||||
#include <THC/THC.h>
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <device_launch_parameters.h>
|
||||
#include <utility> // std::pair
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void generate_deletion_label_kernel(
|
||||
const scalar_t* __restrict__ source,
|
||||
const size_t source_size,
|
||||
const size_t operation_size,
|
||||
int* __restrict__ operations,
|
||||
int* __restrict__ labels) {
|
||||
|
||||
const int index = blockIdx.x;
|
||||
const int offset = index * operation_size;
|
||||
const int offset_label = index * source_size;
|
||||
|
||||
for (int i = 0; i < source_size; i++) {
|
||||
labels[offset_label + i] = 0;
|
||||
}
|
||||
|
||||
int k = 0;
|
||||
for (int i = 0; i < operation_size; i++){
|
||||
if (operations[offset + i] == 0){
|
||||
break;
|
||||
} else if (operations[offset + i] == 1){
|
||||
continue;
|
||||
} else {
|
||||
labels[offset_label + k] = 3 - operations[offset + i];
|
||||
k++;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void generate_insertion_label_kernel(
|
||||
const scalar_t* __restrict__ target,
|
||||
const size_t target_size,
|
||||
const size_t operation_size,
|
||||
int* __restrict__ operations,
|
||||
int* __restrict__ labels,
|
||||
int* __restrict__ masks) {
|
||||
|
||||
const int index = blockIdx.x;
|
||||
const int offset = index * operation_size;
|
||||
const int offset_label = index * target_size;
|
||||
|
||||
int k = 0;
|
||||
int u = 0;
|
||||
int m = 0;
|
||||
|
||||
for (int i = 0; i < target_size; i++) {
|
||||
labels[offset_label + i] = 0;
|
||||
masks[offset_label + i] = 0;
|
||||
}
|
||||
|
||||
for (int i = 0; i < operation_size-1; i++){
|
||||
if (operations[offset + i] == 0){
|
||||
break;
|
||||
} else if (operations[offset + i] == 2){
|
||||
continue;
|
||||
} else if (operations[offset + i] == 1){
|
||||
masks[offset_label + m] = 1;
|
||||
u++; m++;
|
||||
} else {
|
||||
labels[offset_label + k] = u;
|
||||
masks[offset_label + m] = 0;
|
||||
k++; m++;
|
||||
u = 0;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void levenshtein_distance_kernel(
|
||||
const scalar_t* __restrict__ source,
|
||||
const scalar_t* __restrict__ target,
|
||||
const int* __restrict__ source_length,
|
||||
const int* __restrict__ target_length,
|
||||
const size_t source_size,
|
||||
const size_t target_size,
|
||||
int* __restrict__ operations,
|
||||
int* __restrict__ errors_curr) {
|
||||
|
||||
const int index = blockIdx.x;
|
||||
const int offset = index * (source_size + target_size);
|
||||
const int d = index * (source_size + 1) * (target_size + 1);
|
||||
const int t = target_size + 1;
|
||||
|
||||
auto err_idx = [d, t](int i, int j) { return d + i * t + j; };
|
||||
auto opt_idx = [offset](int k) { return offset + k; };
|
||||
|
||||
const int hyp_len = source_length[index];
|
||||
const int ref_len = target_length[index];
|
||||
const scalar_t* hyp_begin = source + index * source_size;
|
||||
const scalar_t* ref_begin = target + index * target_size;
|
||||
|
||||
// dynamic programming
|
||||
for (int i = 0; i <= hyp_len; i++){
|
||||
errors_curr[err_idx(i, 0)] = i;
|
||||
}
|
||||
for (int j = 0; j <= ref_len; j++){
|
||||
errors_curr[err_idx(0, j)] = j;
|
||||
}
|
||||
for (int i = 1; i <= hyp_len; i++){
|
||||
for (int j = 1; j <= ref_len; j++){
|
||||
errors_curr[err_idx(i, j)] = min(
|
||||
min(
|
||||
errors_curr[err_idx(i-1, j)],
|
||||
errors_curr[err_idx(i, j-1)]
|
||||
) + 1,
|
||||
errors_curr[err_idx(i-1, j-1)] + 2 * (
|
||||
*(hyp_begin+i-1) == *(ref_begin+j-1) ? 0 : 1
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// back-tracing
|
||||
int i = hyp_len;
|
||||
int j = ref_len;
|
||||
int o = hyp_len + ref_len;
|
||||
|
||||
for (int k = 0; k < source_size + target_size; k++) {
|
||||
operations[opt_idx(k)] = 0;
|
||||
}
|
||||
|
||||
while ((i >= 0) && (j >= 0)) {
|
||||
if ((i == 0) && (j == 0)) {
|
||||
break;
|
||||
}
|
||||
|
||||
if ((j > 0) && (errors_curr[err_idx(i, j-1)] < errors_curr[err_idx(i, j)])) {
|
||||
o--; operations[opt_idx(o)] = 1; j--; // insertion
|
||||
} else if ((i > 0) && (errors_curr[err_idx(i-1, j)] < errors_curr[err_idx(i, j)])) {
|
||||
o--; operations[opt_idx(o)] = 2; i--; // deletion
|
||||
} else {
|
||||
o--; operations[opt_idx(o)] = 3; i--; j--; // do nothing
|
||||
}
|
||||
}
|
||||
|
||||
// moving to the left
|
||||
for (int k = 0; k < hyp_len + ref_len; k++) {
|
||||
if (k + o < hyp_len + ref_len){
|
||||
operations[opt_idx(k)] = operations[opt_idx(k+o)];
|
||||
} else{
|
||||
operations[opt_idx(k)] = 0; // padding
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
template <typename scalar_t>
|
||||
__global__ void faster_levenshtein_distance_kernel(
|
||||
const scalar_t* __restrict__ source,
|
||||
const scalar_t* __restrict__ target,
|
||||
const int* __restrict__ source_length,
|
||||
const int* __restrict__ target_length,
|
||||
const size_t source_size,
|
||||
const size_t target_size,
|
||||
int* __restrict__ operations) {
|
||||
|
||||
extern __shared__ short errors[];
|
||||
auto errors_curr = errors;
|
||||
|
||||
const int index = blockIdx.x;
|
||||
const int offset = index * (source_size + target_size);
|
||||
const int t = target_size + 1;
|
||||
|
||||
auto err_idx = [t](int i, int j) { return i * t + j; };
|
||||
auto opt_idx = [offset](int k) { return offset + k; };
|
||||
|
||||
const int hyp_len = source_length[index];
|
||||
const int ref_len = target_length[index];
|
||||
const scalar_t* hyp_begin = source + index * source_size;
|
||||
const scalar_t* ref_begin = target + index * target_size;
|
||||
|
||||
// dynamic programming
|
||||
for (int i = 0; i <= hyp_len; i++){
|
||||
errors_curr[err_idx(i, 0)] = i;
|
||||
}
|
||||
for (int j = 0; j <= ref_len; j++){
|
||||
errors_curr[err_idx(0, j)] = j;
|
||||
}
|
||||
for (int i = 1; i <= hyp_len; i++){
|
||||
for (int j = 1; j <= ref_len; j++){
|
||||
errors_curr[err_idx(i, j)] = min(
|
||||
min(
|
||||
errors_curr[err_idx(i-1, j)],
|
||||
errors_curr[err_idx(i, j-1)]
|
||||
) + 1,
|
||||
errors_curr[err_idx(i-1, j-1)] + 2 * (
|
||||
*(hyp_begin+i-1) == *(ref_begin+j-1) ? 0 : 1
|
||||
)
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// back-tracing
|
||||
int i = hyp_len;
|
||||
int j = ref_len;
|
||||
int o = hyp_len + ref_len;
|
||||
|
||||
for (int k = 0; k < source_size + target_size; k++) {
|
||||
operations[opt_idx(k)] = 0;
|
||||
}
|
||||
|
||||
while ((i >= 0) && (j >= 0)) {
|
||||
if ((i == 0) && (j == 0)) {
|
||||
break;
|
||||
}
|
||||
|
||||
if ((j > 0) && (errors_curr[err_idx(i, j-1)] < errors_curr[err_idx(i, j)])) {
|
||||
o--; operations[opt_idx(o)] = 1; j--; // insertion
|
||||
} else if ((i > 0) && (errors_curr[err_idx(i-1, j)] < errors_curr[err_idx(i, j)])) {
|
||||
o--; operations[opt_idx(o)] = 2; i--; // deletion
|
||||
} else {
|
||||
o--; operations[opt_idx(o)] = 3; i--; j--; // do nothing
|
||||
}
|
||||
}
|
||||
|
||||
// moving to the left
|
||||
for (int k = 0; k < hyp_len + ref_len; k++) {
|
||||
if (k + o < hyp_len + ref_len){
|
||||
operations[opt_idx(k)] = operations[opt_idx(k+o)];
|
||||
} else{
|
||||
operations[opt_idx(k)] = 0; // padding
|
||||
}
|
||||
}
|
||||
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor GenerateDeletionLabelCuda(
|
||||
torch::Tensor source,
|
||||
torch::Tensor operations) {
|
||||
|
||||
const auto batch_size = source.size(0);
|
||||
at::TensorOptions options(source.device());
|
||||
options = options.dtype(at::ScalarType::Int);
|
||||
auto labels = torch::empty({batch_size, source.size(1)}, options);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(source.device().index());
|
||||
|
||||
AT_DISPATCH_ALL_TYPES(source.scalar_type(), "generate_deletion_labels", ([&] {
|
||||
generate_deletion_label_kernel<scalar_t><<<batch_size, 1, 0, stream>>>(
|
||||
source.data_ptr<scalar_t>(),
|
||||
source.size(1),
|
||||
operations.size(1),
|
||||
operations.data_ptr<int>(),
|
||||
labels.data_ptr<int>());
|
||||
}));
|
||||
|
||||
return labels;
|
||||
}
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> GenerateInsertionLabelCuda(
|
||||
torch::Tensor target,
|
||||
torch::Tensor operations) {
|
||||
|
||||
const auto batch_size = target.size(0);
|
||||
at::TensorOptions options(target.device());
|
||||
options = options.dtype(at::ScalarType::Int);
|
||||
auto labels = torch::empty({batch_size, target.size(1)}, options);
|
||||
auto masks = torch::empty({batch_size, target.size(1)}, options);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(target.device().index());
|
||||
|
||||
AT_DISPATCH_ALL_TYPES(target.scalar_type(), "generate_insertion_labels", ([&] {
|
||||
generate_insertion_label_kernel<scalar_t><<<batch_size, 1, 0, stream>>>(
|
||||
target.data_ptr<scalar_t>(),
|
||||
target.size(1),
|
||||
operations.size(1),
|
||||
operations.data_ptr<int>(),
|
||||
labels.data_ptr<int>(),
|
||||
masks.data_ptr<int>());
|
||||
}));
|
||||
|
||||
return std::make_pair(labels, masks);
|
||||
}
|
||||
|
||||
|
||||
torch::Tensor LevenshteinDistanceCuda(
|
||||
torch::Tensor source,
|
||||
torch::Tensor target,
|
||||
torch::Tensor source_length,
|
||||
torch::Tensor target_length) {
|
||||
|
||||
const auto batch_size = source.size(0);
|
||||
const auto shared_size = (source.size(1) + 1) * (target.size(1) + 1) * sizeof(short);
|
||||
|
||||
at::TensorOptions options(source.device());
|
||||
options = options.dtype(at::ScalarType::Int);
|
||||
auto operations = torch::empty({batch_size, source.size(1) + target.size(1)}, options);
|
||||
auto stream = at::cuda::getCurrentCUDAStream(source.device().index());
|
||||
|
||||
if (shared_size > 40000) {
|
||||
auto distances = torch::empty({batch_size, (source.size(1) + 1) * (target.size(1) + 1)}, options);
|
||||
AT_DISPATCH_ALL_TYPES(source.scalar_type(), "levenshtein_distance", ([&] {
|
||||
levenshtein_distance_kernel<scalar_t><<<batch_size, 1, 0, stream>>>(
|
||||
source.data_ptr<scalar_t>(),
|
||||
target.data_ptr<scalar_t>(),
|
||||
source_length.data_ptr<int>(),
|
||||
target_length.data_ptr<int>(),
|
||||
source.size(1),
|
||||
target.size(1),
|
||||
operations.data_ptr<int>(),
|
||||
distances.data_ptr<int>());
|
||||
}));
|
||||
} else {
|
||||
AT_DISPATCH_ALL_TYPES(source.scalar_type(), "faster_levenshtein_distance", ([&] {
|
||||
faster_levenshtein_distance_kernel<scalar_t><<<batch_size, 1, shared_size, stream>>>(
|
||||
source.data_ptr<scalar_t>(),
|
||||
target.data_ptr<scalar_t>(),
|
||||
source_length.data_ptr<int>(),
|
||||
target_length.data_ptr<int>(),
|
||||
source.size(1),
|
||||
target.size(1),
|
||||
operations.data_ptr<int>());
|
||||
}));
|
||||
}
|
||||
|
||||
return operations;
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
/**
|
||||
* Copyright 2017-present, Facebook, Inc.
|
||||
* All rights reserved.
|
||||
*
|
||||
* This source code is licensed under the license found in the
|
||||
* LICENSE file in the root directory of this source tree.
|
||||
*/
|
||||
|
||||
#pragma once
|
||||
|
||||
#include <torch/extension.h>
|
||||
|
||||
torch::Tensor LevenshteinDistanceCuda(
|
||||
torch::Tensor source,
|
||||
torch::Tensor target,
|
||||
torch::Tensor source_length,
|
||||
torch::Tensor target_length);
|
||||
|
||||
torch::Tensor GenerateDeletionLabelCuda(
|
||||
torch::Tensor source,
|
||||
torch::Tensor operations);
|
||||
|
||||
std::pair<torch::Tensor, torch::Tensor> GenerateInsertionLabelCuda(
|
||||
torch::Tensor source,
|
||||
torch::Tensor operations);
|
||||
Reference in New Issue
Block a user