<commit_before>-- Criterion based on the warp-ctc library
-- by Baidu: https://github.com/baidu-research/warp-ctc
-- Formatting of the input can be seen here:
-- https://github.com/baidu-research/warp-ctc/blob/master/torch_binding/TUTORIAL.md
require 'warp_ctc'

local CTCCriterion, parent = torch.class('nn.CTCCriterion', 'nn.Criterion')

function CTCCriterion:__init()
    parent.__init(self)
    -- GPU inputs (preallocate)
    self.acts = torch.Tensor()
    self.convertedGradients = torch.Tensor()
end

function CTCCriterion:updateOutput(output, labels)
    local tensorSizes = output:size()
    local acts = self:createCTCBatch(output, tensorSizes)
    local sizes = {}
    for x = 1, tensorSizes[1] do
        table.insert(sizes, tensorSizes[2])
    end
    if (output:type() == 'torch.CudaTensor') then
        local grads = torch.CudaTensor()
        self.output = sumCosts(gpu_ctc(acts, grads, labels, sizes))
    else
        local grads = torch.DoubleTensor()
        self.output = sumCosts(cpu_ctc(acts, grads, labels, sizes))
    end

    return self.output
end


function CTCCriterion:updateGradInput(output, labels)
    local tensorSizes = output:size()
    local acts = self:createCTCBatch(output, tensorSizes)
    local sizes = {}
    for x = 1, tensorSizes[1] do
        table.insert(sizes, tensorSizes[2])
    end
    local grads = acts:clone():zero()
    if (output:type() == 'torch.CudaTensor') then
        gpu_ctc(acts, grads, labels, sizes)
    else
        cpu_ctc(acts, grads, labels, sizes)
    end
    self.gradInput = self:revertBatching(grads, tensorSizes)
    return self.gradInput
end

--If batching occurs multiple costs are returned. We sum the costs and return.
function sumCosts(list)
    local acc
    for k, v in ipairs(list) do
        if 1 == k then
            acc = v
        else
            acc = acc + v
        end
    end
    return acc
end

function CTCCriterion:createCTCBatch(output, sizes)
    self.acts:resize(sizes[1] * sizes[2], sizes[3]):zero()
    local counter = 1
    for i = 1, sizes[2] do
        for j = 1, sizes[1] do
            self.acts[counter] = output[j][i]
            counter = counter + 1
        end
    end
    return self.acts
end

function CTCCriterion:revertBatching(gradients, sizes)
    self.convertedGradients:resize(sizes[1], sizes[2], sizes[3]):zero()
    local counter = 1
    for i = 1, sizes[2] do
        for j = 1, sizes[1] do
            self.convertedGradients[j][i] = gradients[counter]
            counter = counter + 1
        end
    end
    return self.convertedGradients
end<commit_msg>fixed type issues<commit_after>-- Criterion based on the warp-ctc library
-- by Baidu: https://github.com/baidu-research/warp-ctc
-- Formatting of the input can be seen here:
-- https://github.com/baidu-research/warp-ctc/blob/master/torch_binding/TUTORIAL.md
require 'warp_ctc'

local CTCCriterion, parent = torch.class('nn.CTCCriterion', 'nn.Criterion')

function CTCCriterion:__init()
    parent.__init(self)
    -- GPU inputs (preallocate)
    self.acts = torch.Tensor()
    self.convertedGradients = torch.Tensor()
end

function CTCCriterion:updateOutput(output, labels)
    local tensorSizes = output:size()
    local acts = self:createCTCBatch(output, tensorSizes)
    local sizes = {}
    for x = 1, tensorSizes[1] do
        table.insert(sizes, tensorSizes[2])
    end
    if (output:type() == 'torch.CudaTensor') then
        local grads = torch.CudaTensor()
        self.output = sumCosts(gpu_ctc(acts, grads, labels, sizes))
    else
        local grads = torch.Tensor()
        self.output = sumCosts(cpu_ctc(acts:float(), grads:float(), labels, sizes))
    end

    return self.output
end


function CTCCriterion:updateGradInput(output, labels)
    local tensorSizes = output:size()
    local acts = self:createCTCBatch(output, tensorSizes)
    local sizes = {}
    for x = 1, tensorSizes[1] do
        table.insert(sizes, tensorSizes[2])
    end
    local grads = acts:clone():zero()
    if (output:type() == 'torch.CudaTensor') then
        gpu_ctc(acts, grads, labels, sizes)
    else
        cpu_ctc(acts:float(), grads:float(), labels, sizes)
    end
    self.gradInput = self:revertBatching(grads, tensorSizes):typeAs(output)
    return self.gradInput
end

--If batching occurs multiple costs are returned. We sum the costs and return.
function sumCosts(list)
    local acc
    for k, v in ipairs(list) do
        if 1 == k then
            acc = v
        else
            acc = acc + v
        end
    end
    return acc
end

function CTCCriterion:createCTCBatch(output, sizes)
    self.acts:resize(sizes[1] * sizes[2], sizes[3]):zero()
    local counter = 1
    for i = 1, sizes[2] do
        for j = 1, sizes[1] do
            self.acts[counter] = output[j][i]
            counter = counter + 1
        end
    end
    return self.acts
end

function CTCCriterion:revertBatching(gradients, sizes)
    self.convertedGradients:resize(sizes[1], sizes[2], sizes[3]):zero()
    local counter = 1
    for i = 1, sizes[2] do
        for j = 1, sizes[1] do
            self.convertedGradients[j][i] = gradients[counter]
            counter = counter + 1
        end
    end
    return self.convertedGradients
end
{'fraction_non_alphanumeric': 0.11196466295371615, 'fraction_numerical': 0.009602458229306703, 'mean_word_length': 2.854922279792746, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 4, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2797907', 'n_tokens_mistral': 1709, 'n_tokens_neox': 1692, 'n_words': 443}