<commit_before>#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn

from functools import lru_cache


def opt_to_kwargs(opt):
    """Get kwargs for seq2seq from opt."""
    kwargs = {}
    for k in ['mem_size', 'time_features', 'position_encoding', 'hops']:
        if k in opt:
            kwargs[k] = opt[k]
    return kwargs


class MemNN(nn.Module):
    """Memory Network module."""

    def __init__(
        self, num_features, embedding_size, hops=1,
        mem_size=32, time_features=False, position_encoding=False,
        dropout=0, padding_idx=0,
    ):
        """Initialize memnn model.

        See cmdline args in MemnnAgent for description of arguments.
        """
        super().__init__()

        # prepare features
        self.hops = hops

        # time features: we learn an embedding for each memory slot
        self.extra_features = 0
        if time_features:
            self.extra_features += mem_size
            self.time_features = torch.LongTensor(
                range(num_features, num_features + mem_size))

        def embedding(use_extra_feats=True):
            if use_extra_feats:
                return Embed(num_features + self.extra_features,
                             embedding_size,
                             position_encoding=position_encoding,
                             padding_idx=padding_idx)
            else:
                return Embed(num_features, embedding_size,<commit_msg>Fixing argument name mismatch for MemNN (#1588)<commit_after>#!/usr/bin/env python3

# Copyright (c) Facebook, Inc. and its affiliates.
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

import torch
import torch.nn as nn

from functools import lru_cache


def opt_to_kwargs(opt):
    """Get kwargs for seq2seq from opt."""
    kwargs = {}
    for k in ['memsize', 'time_features', 'position_encoding', 'hops']:
        if k in opt:
            kwargs[k] = opt[k]
    return kwargs


class MemNN(nn.Module):
    """Memory Network module."""

    def __init__(
        self, num_features, embedding_size, hops=1,
        memsize=32, time_features=False, position_encoding=False,
        dropout=0, padding_idx=0,
    ):
        """Initialize memnn model.

        See cmdline args in MemnnAgent for description of arguments.
        """
        super().__init__()

        # prepare features
        self.hops = hops

        # time features: we learn an embedding for each memory slot
        self.extra_features = 0
        if time_features:
            self.extra_features += memsize
            self.time_features = torch.LongTensor(
                range(num_features, num_features + memsize))

        def embedding(use_extra_feats=True):
            if use_extra_feats:
                return Embed(num_features + self.extra_features,
                             embedding_size,
                             position_encoding=position_encoding,
                             padding_idx=padding_idx)
            else:
                return Embed(num_features, embedding_size,
{'fraction_non_alphanumeric': 0.08074344911639245, 'fraction_numerical': 0.006093845216331505, 'mean_word_length': 2.1719806763285026, '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': '2621812', 'n_tokens_mistral': 912, 'n_tokens_neox': 882, 'n_words': 271}