<commit_before>--[
-- NN package model helpers.
--
-- (c) 2014 <NAME>
--]
require "nn"
require "math"
require "torch"

require "RluMax"
require "SgdMomentum"

models = {}

---Append dataset information to a model.
function models.append_ds_info(ds, mlp)
    mlp.time_sig = ds.time_sig
    mlp.dims = ds.logical_dims
    return mlp
end

---Create a simple 2-layer NN.
--Each output channel is taken in isolation using this architecture. The
--filters in layer 2 are applied to each note channel and accumulated into the
--hidden features with num_hidden per note track. The output layer is a linear
--combination of these hidden features where each feature is still isolated to
--a single note.
function models.simple_2lnn_iso(ds, num_hidden)

    local mlp = models.append_ds_info(ds, nn.Sequential())
    mlp:add(nn.Linear(mlp.dims.input[2], num_hidden))
    mlp:add(nn.RluMax())
    mlp:add(nn.Linear(num_hidden, mlp.dims.output[2]))

    return mlp
end

---Create a simple 2-layer NN.
--Each output channel is a linear combination over all hidden features for all
--note tracks. Therefore, the activation of note 1 is influenced by the
--activation of all other notes. This requires more weights than the isolated
--architecture.
function models.simple_2lnn_cmb(ds, num_hidden)

    local mlp = models.append_ds_info(ds, nn.Sequential())
    local num_notes = mlp.dims.input[1]
    -- By transposing, we extract features over the notes for each time slice.
    mlp:add(nn.Transpose({1, 2}))
    mlp:add(nn.Linear(num_notes, num_hidden))
    mlp:add(nn.RluMax())
    -- Flatten the network.
    local flattened_len = mlp.dims.input[2] * num_hidden
    mlp:add(nn.Reshape(flattened_len))
    -- Allow output channels to take features across notes.
    mlp:add(nn.Linear(flattened_len, num_notes * mlp.dims.output[2]))
    mlp:add(nn.Reshape(num_notes, mlp.dims.output[2]))

    return mlp
end

---Create a simple 2-layer NN.
--Each output channel is a linear combination over all hidden features for all
--note tracks. Therefore, the activation of note 1 is influenced by the
--activation of all other notes. This requires more weights than the isolated
--architecture.
function models.simple_2lnn_iso_cmb(ds, num_hidden)

    local mlp = models.append_ds_info(ds, nn.Sequential())
    local num_notes = mlp.dims.input[1]
    mlp:add(nn.Linear(mlp.input.dims[2], num_hidden))
    mlp:add(nn.RluMax())
    -- Flatten the network.
    local flattened_len = num_notes * num_hidden
    mlp:add(nn.Reshape(flattened_len))
    -- Allow output channels to take features across notes.
    mlp:add(nn.Linear(flattened_len, num_notes * mlp.output.dims[2]))
    mlp:add(nn.Reshape(num_notes, mlp.output.dims[2]))

    return mlp
end

--- Train model using the given dataset.
function models.train_model(ds, model, criterion, train_args)

    local train = ds.data_train()
    local trainer = nn.SgdMomentum(model, criterion, train_args)
    trainer:train(train)

    local train_loss = 0
    for i = 1, train:size() do
        local X = train[i][1]
        local Y = train[i][2]
        local loss = criterion:forward(model:forward(X), Y)
        train_loss = train_loss + loss
    end
    local avg_train_loss = train_loss / train:size()

    local test_loss = 0
    local test = ds.data_test()
    for i = 1, test:size() do
        local X = test[i][1]
        local Y = test[i][2]
        local loss = criterion:forward(model:forward(X), Y)
        test_loss = test_loss + loss
    end
    local avg_test_loss = test_loss / test:size()<commit_msg>Fix dims usage in models.<commit_after>--[
-- NN package model helpers.
--
-- (c) 2014 <NAME>
--]
require "nn"
require "math"
require "torch"

require "RluMax"
require "SgdMomentum"

models = {}

---Append dataset information to a model.
function models.append_ds_info(ds, mlp)
    mlp.time_sig = ds.time_sig
    mlp.dims = ds.logical_dims
    return mlp
end

---Create a simple 2-layer NN.
--Each output channel is taken in isolation using this architecture. The
--filters in layer 2 are applied to each note channel and accumulated into the
--hidden features with num_hidden per note track. The output layer is a linear
--combination of these hidden features where each feature is still isolated to
--a single note.
function models.simple_2lnn_iso(ds, num_hidden)

    local mlp = models.append_ds_info(ds, nn.Sequential())
    mlp:add(nn.Linear(mlp.dims.input[2], num_hidden))
    mlp:add(nn.RluMax())
    mlp:add(nn.Linear(num_hidden, mlp.dims.output[2]))

    return mlp
end

---Create a simple 2-layer NN.
--Each output channel is a linear combination over all hidden features for all
--note tracks. Therefore, the activation of note 1 is influenced by the
--activation of all other notes. This requires more weights than the isolated
--architecture.
function models.simple_2lnn_cmb(ds, num_hidden)

    local mlp = models.append_ds_info(ds, nn.Sequential())
    local num_notes = mlp.dims.input[1]
    -- By transposing, we extract features over the notes for each time slice.
    mlp:add(nn.Transpose({1, 2}))
    mlp:add(nn.Linear(num_notes, num_hidden))
    mlp:add(nn.RluMax())
    -- Flatten the network.
    local flattened_len = mlp.dims.input[2] * num_hidden
    mlp:add(nn.Reshape(flattened_len))
    -- Allow output channels to take features across notes.
    mlp:add(nn.Linear(flattened_len, num_notes * mlp.dims.output[2]))
    mlp:add(nn.Reshape(num_notes, mlp.dims.output[2]))

    return mlp
end

---Create a simple 2-layer NN.
--Each output channel is a linear combination over all hidden features for all
--note tracks. Therefore, the activation of note 1 is influenced by the
--activation of all other notes. This requires more weights than the isolated
--architecture.
function models.simple_2lnn_iso_cmb(ds, num_hidden)

    local mlp = models.append_ds_info(ds, nn.Sequential())
    local num_notes = mlp.dims.input[1]
    mlp:add(nn.Linear(mlp.dims.input[2], num_hidden))
    mlp:add(nn.RluMax())
    -- Flatten the network.
    local flattened_len = num_notes * num_hidden
    mlp:add(nn.Reshape(flattened_len))
    -- Allow output channels to take features across notes.
    mlp:add(nn.Linear(flattened_len, num_notes * mlp.dims.output[2]))
    mlp:add(nn.Reshape(num_notes, mlp.dims.output[2]))

    return mlp
end

--- Train model using the given dataset.
function models.train_model(ds, model, criterion, train_args)

    local train = ds.data_train()
    local trainer = nn.SgdMomentum(model, criterion, train_args)
    trainer:train(train)

    local train_loss = 0
    for i = 1, train:size() do
        local X = train[i][1]
        local Y = train[i][2]
        local loss = criterion:forward(model:forward(X), Y)
        train_loss = train_loss + loss
    end
    local avg_train_loss = train_loss / train:size()

    local test_loss = 0
    local test = ds.data_test()
    for i = 1, test:size() do
        local X = test[i][1]
        local Y = test[i][2]
        local loss = criterion:forward(model:forward(X), Y)
        test_loss = test_loss + loss
    end
    local avg_test_loss = test_loss / test:size()
{'fraction_non_alphanumeric': 0.10815002820078962, 'fraction_numerical': 0.009306260575296108, 'mean_word_length': 3.4386733416770965, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, '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': '567352', 'n_tokens_mistral': 2406, 'n_tokens_neox': 2341, 'n_words': 717}