StateNameModule.cpp 1.19 KB
Newer Older
Franck Dary's avatar
Franck Dary committed
1
2
#include "StateNameModule.hpp"

Franck Dary's avatar
Franck Dary committed
3
StateNameModuleImpl::StateNameModuleImpl(std::string name, const std::string & definition)
Franck Dary's avatar
Franck Dary committed
4
{
Franck Dary's avatar
Franck Dary committed
5
6
  setName(name);
  std::regex regex("(?:(?:\\s|\\t)*)Out\\{(.*)\\}(?:(?:\\s|\\t)*)");
Franck Dary's avatar
Franck Dary committed
7
8
9
10
  if (!util::doIfNameMatch(regex, definition, [this,&definition](auto sm)
        {
          try
          {
Franck Dary's avatar
Franck Dary committed
11
            outSize = std::stoi(sm.str(1));
Franck Dary's avatar
Franck Dary committed
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
          } catch (std::exception & e) {util::myThrow(fmt::format("{} in '{}'",e.what(),definition));}
        }))
    util::myThrow(fmt::format("invalid definition '{}'", definition));
}

torch::Tensor StateNameModuleImpl::forward(torch::Tensor input)
{
  return embeddings(input.narrow(1,firstInputIndex,1).squeeze(1));
}

std::size_t StateNameModuleImpl::getOutputSize()
{
  return outSize;
}

std::size_t StateNameModuleImpl::getInputSize()
{
  return 1;
}

Franck Dary's avatar
Franck Dary committed
32
void StateNameModuleImpl::addToContext(std::vector<std::vector<long>> & context, const Config & config)
Franck Dary's avatar
Franck Dary committed
33
{
Franck Dary's avatar
Franck Dary committed
34
  auto & dict = getDict();
Franck Dary's avatar
Franck Dary committed
35
  for (auto & contextElement : context)
36
    contextElement.emplace_back(dict.getIndexOrInsert(config.getState(), ""));
Franck Dary's avatar
Franck Dary committed
37
38
}

39
void StateNameModuleImpl::registerEmbeddings()
Franck Dary's avatar
Franck Dary committed
40
{
41
  embeddings = register_module("embeddings", WordEmbeddings(getDict().size(), outSize));
Franck Dary's avatar
Franck Dary committed
42
43
}