<commit_before>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import torch.nn as nn
import torch as T
from torch.autograd import Variable as var
import numpy as np

from torch.nn.utils.rnn import pad_packed_sequence as pad
from torch.nn.utils.rnn import pack_padded_sequence as pack
from torch.nn.utils.rnn import PackedSequence
from torch.nn.init import orthogonal, xavier_uniform

from .util import *
from .sparse_memory import SparseMemory


class SDNC(nn.Module):

  def __init__(
      self,
      input_size,
      hidden_size,
      rnn_type='lstm',
      num_layers=1,
      num_hidden_layers=2,
      bias=True,
      batch_first=True,
      dropout=0,
      bidirectional=False,
      nr_cells=5,
      sparse_reads=10,
      read_heads=4,
      cell_size=10,
      nonlinearity='tanh',
      gpu_id=-1,
      independent_linears=False,
      share_memory=True,
      debug=False,
      clip=20
  ):
    super(SDNC, self).__init__()
    # todo: separate weights and RNNs for the interface and output vectors

    self.input_size = input_size
    self.hidden_size = hidden_size
    self.rnn_type = rnn_type
    self.num_layers = num_layers
    self.num_hidden_layers = num_hidden_layers
    self.bias = bias
    self.batch_first = batch_first
    self.dropout = dropout
    self.bidirectional = bidirectional
    self.nr_cells = nr_cells
    self.sparse_reads = sparse_reads
    self.read_heads = read_heads
    self.cell_size = cell_size
    self.nonlinearity = nonlinearity
    self.gpu_id = gpu_id
    self.independent_linears = independent_linears
    self.share_memory = share_memory
    self.debug = debug
    self.clip = clip

    self.w = self.cell_size
    self.r = self.read_heads

    self.read_vectors_size = self.r * self.w
    self.output_size = self.hidden_size

    self.nn_input_size = self.input_size + self.read_vectors_size
    self.nn_output_size = self.output_size + self.read_vectors_size

    self.rnns = []
    self.memories = []

    for layer in range(self.num_layers):
      if self.rnn_type.lower() == 'rnn':
        self.rnns.append(nn.RNN((self.nn_input_size if layer == 0 else self.nn_output_size), self.output_size,
                                bias=self.bias, nonlinearity=self.nonlinearity, batch_first=True, dropout=self.dropout, num_layers=self.num_hidden_layers))
      elif self.rnn_type.lower() == 'gru':
        self.rnns.append(nn.GRU((self.nn_input_size if layer == 0 else self.nn_output_size),
                                self.output_size, bias=self.bias, batch_first=True, dropout=self.dropout, num_layers=self.num_hidden_layers))
      if self.rnn_type.lower() == 'lstm':
        self.rnns.append(nn.LSTM((self.nn_input_size if layer == 0 else self.nn_output_size),
                                 self.output_size, bias=self.bias, batch_first=True, dropout=self.dropout, num_layers=self.num_hidden_layers))
      setattr(self, self.rnn_type.lower() + '_layer_' + str(layer), self.rnns[layer])

      # memories for each layer
      if not self.share_memory:
        self.memories.append(
            SparseMemory(
                input_size=self.output_size,
                mem_size=self.nr_cells,
                cell_size=self.w,
                sparse_reads=self.sparse_reads,
                read_heads=self.read_heads,
                gpu_id=self.gpu_id,
                mem_gpu_id=self.gpu_id,
                independent_linears=self.independent_linears
            )
        )
        setattr(self, 'rnn_layer_memory_' + str(layer), self.memories[layer])

    # only one memory shared by all layers
    if self.share_memory:
      self.memories.append(
          SparseMemory(
              input_size=self.output_size,
              mem_size=self.nr_cells,
              cell_size=self.w,
              sparse_reads=self.sparse_reads,
              read_heads=self.read_heads,
              gpu_id=self.gpu_id,
              mem_gpu_id=self.gpu_id,
              independent_linears=self.independent_linears
          )
      )
      setattr(self, 'rnn_layer_memory_shared', self.memories[0])

    # final output layer
    self.output = nn.Linear(self.nn_output_size, self.input_size)
    orthogonal(self.output.weight)

    if self.gpu_id != -1:
      [x.cuda(self.gpu_id) for x in self.rnns]
      [x.cuda(self.gpu_id) for x in self.memories]

  def _init_hidden(self, hx, batch_size, reset_experience):
    # create empty hidden states if not provided
    if hx is None:
      hx = (None, None, None)
    (chx, mhx, last_read) = hx

    # initialize hidden state of the controller RNN
    if chx is None:
      h = cuda(T.zeros(self.num_hidden_layers, batch_size, self.output_size), gpu_id=self.gpu_id)
      xavier_uniform(h)

      chx = [ (h, h) if self.rnn_type.lower() == 'lstm' else h for x in range(self.num_layers)]

    # Last read vectors
    if last_read is None:
      last_read = cuda(T.zeros(batch_size, self.w * self.r), gpu_id=self.gpu_id)

    # memory states
    if mhx is None:
      if self.share_memory:
        mhx = self.memories[0].reset(batch_size, erase=reset_experience)
      else:
        mhx = [m.reset(batch_size, erase=reset_experience) for m in self.memories]
    else:
      if self.share_memory:
        mhx = self.memories[0].reset(batch_size, mhx, erase=reset_experience)
      else:
        mhx = [m.reset(batch_size, h, erase=reset_experience) for m, h in zip(self.memories, mhx)]

    return chx, mhx, last_read

  def _debug(self, mhx, debug_obj):
    if not debug_obj:
      debug_obj = {
          'memory': [],
          'visible_memory': [],
          'read_weights': [],
          'write_weights': [],
          'read_vectors': [],
          'last_used_mem': [],
          'usage': [],
          'read_positions': []
      }

    debug_obj['memory'].append(mhx['memory'][0].data.cpu().numpy())
    debug_obj['visible_memory'].append(mhx['visible_memory'][0].data.cpu().numpy())
    debug_obj['read_weights'].append(mhx['read_weights'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['write_weights'].append(mhx['write_weights'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['read_vectors'].append(mhx['read_vectors'][0].data.cpu().numpy())
    debug_obj['last_used_mem'].append(mhx['last_used_mem'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['usage'].append(mhx['usage'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['read_positions'].append(mhx['read_positions'][0].unsqueeze(0).data.cpu().numpy())

    return debug_obj

  def _layer_forward(self, input, layer, hx=(None, None), pass_through_memory=True):
    (chx, mhx) = hx

    # pass through the controller layer
    input, chx = self.rnns[layer](input.unsqueeze(1), chx)
    input = input.squeeze(1)

    # clip the controller output
    if self.clip != 0:
      output = T.clamp(input, -self.clip, self.clip)<commit_msg>revert to using FLANN indexes, in a faster way<commit_after>#!/usr/bin/env python3
# -*- coding: utf-8 -*-

import torch.nn as nn
import torch as T
from torch.autograd import Variable as var
import numpy as np

from torch.nn.utils.rnn import pad_packed_sequence as pad
from torch.nn.utils.rnn import pack_padded_sequence as pack
from torch.nn.utils.rnn import PackedSequence
from torch.nn.init import orthogonal, xavier_uniform

from .util import *
from .sparse_memory import SparseMemory


class SDNC(nn.Module):

  def __init__(
      self,
      input_size,
      hidden_size,
      rnn_type='lstm',
      num_layers=1,
      num_hidden_layers=2,
      bias=True,
      batch_first=True,
      dropout=0,
      bidirectional=False,
      nr_cells=5,
      sparse_reads=10,
      read_heads=4,
      cell_size=10,
      nonlinearity='tanh',
      gpu_id=-1,
      independent_linears=False,
      share_memory=True,
      debug=False,
      clip=20
  ):
    super(SDNC, self).__init__()
    # todo: separate weights and RNNs for the interface and output vectors

    self.input_size = input_size
    self.hidden_size = hidden_size
    self.rnn_type = rnn_type
    self.num_layers = num_layers
    self.num_hidden_layers = num_hidden_layers
    self.bias = bias
    self.batch_first = batch_first
    self.dropout = dropout
    self.bidirectional = bidirectional
    self.nr_cells = nr_cells
    self.sparse_reads = sparse_reads
    self.read_heads = read_heads
    self.cell_size = cell_size
    self.nonlinearity = nonlinearity
    self.gpu_id = gpu_id
    self.independent_linears = independent_linears
    self.share_memory = share_memory
    self.debug = debug
    self.clip = clip

    self.w = self.cell_size
    self.r = self.read_heads

    self.read_vectors_size = self.r * self.w
    self.output_size = self.hidden_size

    self.nn_input_size = self.input_size + self.read_vectors_size
    self.nn_output_size = self.output_size + self.read_vectors_size

    self.rnns = []
    self.memories = []

    for layer in range(self.num_layers):
      if self.rnn_type.lower() == 'rnn':
        self.rnns.append(nn.RNN((self.nn_input_size if layer == 0 else self.nn_output_size), self.output_size,
                                bias=self.bias, nonlinearity=self.nonlinearity, batch_first=True, dropout=self.dropout, num_layers=self.num_hidden_layers))
      elif self.rnn_type.lower() == 'gru':
        self.rnns.append(nn.GRU((self.nn_input_size if layer == 0 else self.nn_output_size),
                                self.output_size, bias=self.bias, batch_first=True, dropout=self.dropout, num_layers=self.num_hidden_layers))
      if self.rnn_type.lower() == 'lstm':
        self.rnns.append(nn.LSTM((self.nn_input_size if layer == 0 else self.nn_output_size),
                                 self.output_size, bias=self.bias, batch_first=True, dropout=self.dropout, num_layers=self.num_hidden_layers))
      setattr(self, self.rnn_type.lower() + '_layer_' + str(layer), self.rnns[layer])

      # memories for each layer
      if not self.share_memory:
        self.memories.append(
            SparseMemory(
                input_size=self.output_size,
                mem_size=self.nr_cells,
                cell_size=self.w,
                sparse_reads=self.sparse_reads,
                read_heads=self.read_heads,
                gpu_id=self.gpu_id,
                mem_gpu_id=self.gpu_id,
                independent_linears=self.independent_linears
            )
        )
        setattr(self, 'rnn_layer_memory_' + str(layer), self.memories[layer])

    # only one memory shared by all layers
    if self.share_memory:
      self.memories.append(
          SparseMemory(
              input_size=self.output_size,
              mem_size=self.nr_cells,
              cell_size=self.w,
              sparse_reads=self.sparse_reads,
              read_heads=self.read_heads,
              gpu_id=self.gpu_id,
              mem_gpu_id=self.gpu_id,
              independent_linears=self.independent_linears
          )
      )
      setattr(self, 'rnn_layer_memory_shared', self.memories[0])

    # final output layer
    self.output = nn.Linear(self.nn_output_size, self.input_size)
    orthogonal(self.output.weight)

    if self.gpu_id != -1:
      [x.cuda(self.gpu_id) for x in self.rnns]
      [x.cuda(self.gpu_id) for x in self.memories]

  def _init_hidden(self, hx, batch_size, reset_experience):
    # create empty hidden states if not provided
    if hx is None:
      hx = (None, None, None)
    (chx, mhx, last_read) = hx

    # initialize hidden state of the controller RNN
    if chx is None:
      h = cuda(T.zeros(self.num_hidden_layers, batch_size, self.output_size), gpu_id=self.gpu_id)
      xavier_uniform(h)

      chx = [ (h, h) if self.rnn_type.lower() == 'lstm' else h for x in range(self.num_layers)]

    # Last read vectors
    if last_read is None:
      last_read = cuda(T.zeros(batch_size, self.w * self.r), gpu_id=self.gpu_id)

    # memory states
    if mhx is None:
      if self.share_memory:
        mhx = self.memories[0].reset(batch_size, erase=reset_experience)
      else:
        mhx = [m.reset(batch_size, erase=reset_experience) for m in self.memories]
    else:
      if self.share_memory:
        mhx = self.memories[0].reset(batch_size, mhx, erase=reset_experience)
      else:
        mhx = [m.reset(batch_size, h, erase=reset_experience) for m, h in zip(self.memories, mhx)]

    return chx, mhx, last_read

  def _debug(self, mhx, debug_obj):
    if not debug_obj:
      debug_obj = {
          'memory': [],
          'visible_memory': [],
          'read_weights': [],
          'write_weights': [],
          'read_vectors': [],
          'least_used_mem': [],
          'usage': [],
          'read_positions': []
      }

    debug_obj['memory'].append(mhx['memory'][0].data.cpu().numpy())
    debug_obj['visible_memory'].append(mhx['visible_memory'][0].data.cpu().numpy())
    debug_obj['read_weights'].append(mhx['read_weights'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['write_weights'].append(mhx['write_weights'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['read_vectors'].append(mhx['read_vectors'][0].data.cpu().numpy())
    debug_obj['least_used_mem'].append(mhx['least_used_mem'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['usage'].append(mhx['usage'][0].unsqueeze(0).data.cpu().numpy())
    debug_obj['read_positions'].append(mhx['read_positions'][0].unsqueeze(0).data.cpu().numpy())

    return debug_obj

  def _layer_forward(self, input, layer, hx=(None, None), pass_through_memory=True):
    (chx, mhx) = hx

    # pass through the controller layer
    input, chx = self.rnns[layer](input.unsqueeze(1), chx)
    input = input.squeeze(1)

    # clip the controller output
    if self.clip != 0:
      output = T.clamp(input, -self.clip, self.clip)
{'fraction_non_alphanumeric': 0.11752230510457803, 'fraction_numerical': 0.005411730291063332, 'mean_word_length': 2.7828492392807744, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1395344', 'n_tokens_mistral': 4791, 'n_tokens_neox': 4602, 'n_words': 849}