<commit_before>#pragma once

#include "layers/attention.h"
#include "models/encdec.h"

namespace marian {

class EncoderStateS2S : public EncoderState {
private:
  Expr context_;
  Expr mask_;
  Ptr<data::CorpusBatch> batch_;

public:
  EncoderStateS2S(Expr context, Expr mask, Ptr<data::CorpusBatch> batch)
      : context_(context), mask_(mask), batch_(batch) {}

  Expr getContext() { return context_; }
  Expr getMask() { return mask_; }

  virtual const std::vector<size_t>& getSourceWords() {
    return batch_->front()->indeces();
  }
};

class DecoderStateS2S : public DecoderState {
private:
  std::vector<Expr> states_;
  Expr probs_;
  Ptr<EncoderState> encState_;

public:
  DecoderStateS2S(const std::vector<Expr> states,
                  Expr probs,
                  Ptr<EncoderState> encState)
      : states_(states), probs_(probs), encState_(encState) {}

  Ptr<EncoderState> getEncoderState() { return encState_; }
  Expr getProbs() { return probs_; }
  void setProbs(Expr probs) { probs_ = probs; }

  Ptr<DecoderState> select(const std::vector<size_t>& selIdx) {
    int numSelected = selIdx.size();
    int dimState = states_[0]->shape()[1];

    std::vector<Expr> selectedStates;
    for(auto state : states_) {
      selectedStates.push_back(
          reshape(rows(state, selIdx), {1, dimState, 1, numSelected}));
    }

    return New<DecoderStateS2S>(selectedStates, probs_, encState_);
  }

  const std::vector<Expr>& getStates() { return states_; }
};

class EncoderS2S : public EncoderBase {
public:
  template <class... Args>
  EncoderS2S(Ptr<Config> options, Args... args)
      : EncoderBase(options, args...) {}

  Ptr<EncoderState> build(Ptr<ExpressionGraph> graph,
                          Ptr<data::CorpusBatch> batch,
                          size_t batchIdx) {
    using namespace keywords;

    int dimSrcVoc = options_->get<std::vector<int>>("dim-vocabs")[batchIdx];
    int dimSrcEmb = options_->get<int>("dim-emb");

    int dimEncState = options_->get<int>("dim-rnn");
    bool layerNorm = options_->get<bool>("layer-normalization");
    bool skipDepth = options_->get<bool>("skip");
    size_t encoderLayers = options_->get<size_t>("layers-enc");

    bool amun = options_->get<std::string>("type") == "amun";
    UTIL_THROW_IF2(amun && options_->get<int>("layers-enc") > 1,
                   "--type amun does not currently support multiple encoder "
                   "layers, use --type s2s");
    UTIL_THROW_IF2(amun && options_->get<bool>("skip"),
                   "--type amun does not currently support skip connections, "
                   "use --type s2s");

    float dropoutRnn = inference_ ? 0 : options_->get<float>("dropout-rnn");
    float dropoutSrc = inference_ ? 0 : options_->get<float>("dropout-src");

    auto xEmb = Embedding(prefix_ + "_Wemb", dimSrcVoc, dimSrcEmb)(graph);

    Expr x, xMask;
    std::tie(x, xMask) = prepareSource(xEmb, batch, batchIdx);

    if(dropoutSrc) {
      int dimBatch = x->shape()[0];
      int srcWords = x->shape()[2];
      auto srcWordDrop = graph->dropout(dropoutSrc, {1, 1, srcWords});
      x = dropout(x, mask = srcWordDrop);
    }

    auto xFw = RNN<GRU>(graph,
                        prefix_ + "_bi",
                        dimSrcEmb,
                        dimEncState,
                        normalize = layerNorm,
                        dropout_prob = dropoutRnn)(x);

    auto xBw = RNN<GRU>(graph,
                        prefix_ + "_bi_r",
                        dimSrcEmb,
                        dimEncState,
                        normalize = layerNorm,
                        direction = dir::backward,
                        dropout_prob = dropoutRnn)(x, mask = xMask);

    if(encoderLayers > 1) {
      auto xBi = concatenate({xFw, xBw}, axis = 1);

      Expr xContext;
      std::vector<Expr> states;
      std::tie(xContext, states) = MLRNN<GRU>(graph,
                                              prefix_,
                                              encoderLayers - 1,
                                              2 * dimEncState,
                                              dimEncState,
                                              normalize = layerNorm,
                                              skip = skipDepth,
                                              dropout_prob = dropoutRnn)(xBi);
      return New<EncoderStateS2S>(xContext, xMask, batch);
    } else {
      auto xContext = concatenate({xFw, xBw}, axis = 1);
      return New<EncoderStateS2S>(xContext, xMask, batch);
    }
  }
};

class DecoderS2S : public DecoderBase {
private:
  Ptr<GlobalAttention> attention_;
  Ptr<RNN<CGRU>> rnnL1;
  Ptr<MLRNN<GRU>> rnnLn;
  Expr tiedOutputWeights_;

public:
  template <class... Args>
  DecoderS2S(Ptr<Config> options, Args... args)
      : DecoderBase(options, args...) {}

  virtual Ptr<DecoderState> startState(Ptr<EncoderState> encState) {
    using namespace keywords;

    auto meanContext = weighted_average(
        encState->getContext(), encState->getMask(), axis = 2);

    bool layerNorm = options_->get<bool>("layer-normalization");
    auto start = Dense(prefix_ + "_ff_state",
                       options_->get<int>("dim-rnn"),
                       activation = act::tanh,
                       normalize = layerNorm)(meanContext);

    std::vector<Expr> startStates(options_->get<size_t>("layers-dec"), start);
    return New<DecoderStateS2S>(startStates, nullptr, encState);
  }

  virtual Ptr<DecoderState> step(Ptr<ExpressionGraph> graph,
                                 Ptr<DecoderState> state) {
    using namespace keywords;

    int dimTrgVoc = options_->get<std::vector<int>>("dim-vocabs").back();

    int dimTrgEmb
        = options_->get<int>("dim-emb") + options_->get<int>("dim-pos");

    int dimDecState = options_->get<int>("dim-rnn");
    bool layerNorm = options_->get<bool>("layer-normalization");
    bool skipDepth = options_->get<bool>("skip");
    size_t decoderLayers = options_->get<size_t>("layers-dec");

    float dropoutRnn = inference_ ? 0 : options_->get<float>("dropout-rnn");
    float dropoutTrg = inference_ ? 0 : options_->get<float>("dropout-trg");

    bool amun = options_->get<std::string>("type") == "amun";
    UTIL_THROW_IF2(options_->get<bool>("skip"),
                   "--type amun does not currently support skip connections, "
                   "use --type s2s");
    UTIL_THROW_IF2(options_->get<int>("layers-dec") > 1,
                   "--type amun does not currently support multiple decoder "
                   "layers, use --type s2s");
    UTIL_THROW_IF2(options_->get<bool>("tied-embeddings"),
                   "--type amun does not currently support tied embeddings, "
                   "use --type s2s");

    auto stateS2S = std::dynamic_pointer_cast<DecoderStateS2S>(state);

    auto embeddings = stateS2S->getTargetEmbeddings();

    if(dropoutTrg) {
      int dimBatch = embeddings->shape()[0];
      int trgWords = embeddings->shape()[2];
      auto trgWordDrop = graph->dropout(dropoutTrg, {1, 1, trgWords});
      embeddings = dropout(embeddings, mask = trgWordDrop);
    }

    if(!attention_)
      attention_ = New<GlobalAttention>(prefix_,
                                        state->getEncoderState(),
                                        dimDecState,
                                        dropout_prob = dropoutRnn,
                                        normalize = layerNorm);

    if(!rnnL1)
      rnnL1 = New<RNN<CGRU>>(graph,
                             prefix_,
                             dimTrgEmb,
                             dimDecState,
                             attention_,<commit_msg>missing condition on amun type<commit_after>#pragma once

#include "layers/attention.h"
#include "models/encdec.h"

namespace marian {

class EncoderStateS2S : public EncoderState {
private:
  Expr context_;
  Expr mask_;
  Ptr<data::CorpusBatch> batch_;

public:
  EncoderStateS2S(Expr context, Expr mask, Ptr<data::CorpusBatch> batch)
      : context_(context), mask_(mask), batch_(batch) {}

  Expr getContext() { return context_; }
  Expr getMask() { return mask_; }

  virtual const std::vector<size_t>& getSourceWords() {
    return batch_->front()->indeces();
  }
};

class DecoderStateS2S : public DecoderState {
private:
  std::vector<Expr> states_;
  Expr probs_;
  Ptr<EncoderState> encState_;

public:
  DecoderStateS2S(const std::vector<Expr> states,
                  Expr probs,
                  Ptr<EncoderState> encState)
      : states_(states), probs_(probs), encState_(encState) {}

  Ptr<EncoderState> getEncoderState() { return encState_; }
  Expr getProbs() { return probs_; }
  void setProbs(Expr probs) { probs_ = probs; }

  Ptr<DecoderState> select(const std::vector<size_t>& selIdx) {
    int numSelected = selIdx.size();
    int dimState = states_[0]->shape()[1];

    std::vector<Expr> selectedStates;
    for(auto state : states_) {
      selectedStates.push_back(
          reshape(rows(state, selIdx), {1, dimState, 1, numSelected}));
    }

    return New<DecoderStateS2S>(selectedStates, probs_, encState_);
  }

  const std::vector<Expr>& getStates() { return states_; }
};

class EncoderS2S : public EncoderBase {
public:
  template <class... Args>
  EncoderS2S(Ptr<Config> options, Args... args)
      : EncoderBase(options, args...) {}

  Ptr<EncoderState> build(Ptr<ExpressionGraph> graph,
                          Ptr<data::CorpusBatch> batch,
                          size_t batchIdx) {
    using namespace keywords;

    int dimSrcVoc = options_->get<std::vector<int>>("dim-vocabs")[batchIdx];
    int dimSrcEmb = options_->get<int>("dim-emb");

    int dimEncState = options_->get<int>("dim-rnn");
    bool layerNorm = options_->get<bool>("layer-normalization");
    bool skipDepth = options_->get<bool>("skip");
    size_t encoderLayers = options_->get<size_t>("layers-enc");

    bool amun = options_->get<std::string>("type") == "amun";
    UTIL_THROW_IF2(amun && options_->get<int>("layers-enc") > 1,
                   "--type amun does not currently support multiple encoder "
                   "layers, use --type s2s");
    UTIL_THROW_IF2(amun && options_->get<bool>("skip"),
                   "--type amun does not currently support skip connections, "
                   "use --type s2s");

    float dropoutRnn = inference_ ? 0 : options_->get<float>("dropout-rnn");
    float dropoutSrc = inference_ ? 0 : options_->get<float>("dropout-src");

    auto xEmb = Embedding(prefix_ + "_Wemb", dimSrcVoc, dimSrcEmb)(graph);

    Expr x, xMask;
    std::tie(x, xMask) = prepareSource(xEmb, batch, batchIdx);

    if(dropoutSrc) {
      int dimBatch = x->shape()[0];
      int srcWords = x->shape()[2];
      auto srcWordDrop = graph->dropout(dropoutSrc, {1, 1, srcWords});
      x = dropout(x, mask = srcWordDrop);
    }

    auto xFw = RNN<GRU>(graph,
                        prefix_ + "_bi",
                        dimSrcEmb,
                        dimEncState,
                        normalize = layerNorm,
                        dropout_prob = dropoutRnn)(x);

    auto xBw = RNN<GRU>(graph,
                        prefix_ + "_bi_r",
                        dimSrcEmb,
                        dimEncState,
                        normalize = layerNorm,
                        direction = dir::backward,
                        dropout_prob = dropoutRnn)(x, mask = xMask);

    if(encoderLayers > 1) {
      auto xBi = concatenate({xFw, xBw}, axis = 1);

      Expr xContext;
      std::vector<Expr> states;
      std::tie(xContext, states) = MLRNN<GRU>(graph,
                                              prefix_,
                                              encoderLayers - 1,
                                              2 * dimEncState,
                                              dimEncState,
                                              normalize = layerNorm,
                                              skip = skipDepth,
                                              dropout_prob = dropoutRnn)(xBi);
      return New<EncoderStateS2S>(xContext, xMask, batch);
    } else {
      auto xContext = concatenate({xFw, xBw}, axis = 1);
      return New<EncoderStateS2S>(xContext, xMask, batch);
    }
  }
};

class DecoderS2S : public DecoderBase {
private:
  Ptr<GlobalAttention> attention_;
  Ptr<RNN<CGRU>> rnnL1;
  Ptr<MLRNN<GRU>> rnnLn;
  Expr tiedOutputWeights_;

public:
  template <class... Args>
  DecoderS2S(Ptr<Config> options, Args... args)
      : DecoderBase(options, args...) {}

  virtual Ptr<DecoderState> startState(Ptr<EncoderState> encState) {
    using namespace keywords;

    auto meanContext = weighted_average(
        encState->getContext(), encState->getMask(), axis = 2);

    bool layerNorm = options_->get<bool>("layer-normalization");
    auto start = Dense(prefix_ + "_ff_state",
                       options_->get<int>("dim-rnn"),
                       activation = act::tanh,
                       normalize = layerNorm)(meanContext);

    std::vector<Expr> startStates(options_->get<size_t>("layers-dec"), start);
    return New<DecoderStateS2S>(startStates, nullptr, encState);
  }

  virtual Ptr<DecoderState> step(Ptr<ExpressionGraph> graph,
                                 Ptr<DecoderState> state) {
    using namespace keywords;

    int dimTrgVoc = options_->get<std::vector<int>>("dim-vocabs").back();

    int dimTrgEmb
        = options_->get<int>("dim-emb") + options_->get<int>("dim-pos");

    int dimDecState = options_->get<int>("dim-rnn");
    bool layerNorm = options_->get<bool>("layer-normalization");
    bool skipDepth = options_->get<bool>("skip");
    size_t decoderLayers = options_->get<size_t>("layers-dec");

    float dropoutRnn = inference_ ? 0 : options_->get<float>("dropout-rnn");
    float dropoutTrg = inference_ ? 0 : options_->get<float>("dropout-trg");

    bool amun = options_->get<std::string>("type") == "amun";
    UTIL_THROW_IF2(amun && options_->get<bool>("skip"),
                   "--type amun does not currently support skip connections, "
                   "use --type s2s");
    UTIL_THROW_IF2(amun && options_->get<int>("layers-dec") > 1,
                   "--type amun does not currently support multiple decoder "
                   "layers, use --type s2s");
    UTIL_THROW_IF2(amun && options_->get<bool>("tied-embeddings"),
                   "--type amun does not currently support tied embeddings, "
                   "use --type s2s");

    auto stateS2S = std::dynamic_pointer_cast<DecoderStateS2S>(state);

    auto embeddings = stateS2S->getTargetEmbeddings();

    if(dropoutTrg) {
      int dimBatch = embeddings->shape()[0];
      int trgWords = embeddings->shape()[2];
      auto trgWordDrop = graph->dropout(dropoutTrg, {1, 1, trgWords});
      embeddings = dropout(embeddings, mask = trgWordDrop);
    }

    if(!attention_)
      attention_ = New<GlobalAttention>(prefix_,
                                        state->getEncoderState(),
                                        dimDecState,
                                        dropout_prob = dropoutRnn,
                                        normalize = layerNorm);

    if(!rnnL1)
      rnnL1 = New<RNN<CGRU>>(graph,
                             prefix_,
                             dimTrgEmb,
                             dimDecState,
                             attention_,
{'fraction_non_alphanumeric': 0.12796362455342644, 'fraction_numerical': 0.006755440077947385, 'mean_word_length': 2.1555646648903464, 'pattern_counts': {'":': 0, '<': 145, '<?xml version=': 0, '>': 231, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 38, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '183048', 'n_tokens_mistral': 4655, 'n_tokens_neox': 4436, 'n_words': 940}