transformer.hpp

#pragma once
 
#include "../layernorm.hpp"
#include "../linear.hpp"
#include "../model.hpp"
 
#include <cmath>
#include <vector>
 
namespace dl {
 
    constexpr double calcPosEncoding(size_t pos, size_t i, size_t dimModel) {
        return (i % 2 == 0) ? std::sin(pos / std::pow(10000, i / dimModel))
                            : std::sin(pos / std::pow(10000, (i - 1) / dimModel));
    }
 
    struct TransformerConf {
        struct {
            size_t model; 
            size_t key;   
            size_t value; 
            size_t inner; 
        } dimensions;
        size_t numEncoders;
        size_t numAttnHeads;
    };
 
    class TransformerEncoder final : public Model<TensorPtr(TensorPtr)> {
    public:
        TransformerConf conf;
        // Multi-Head Attention
        dl::Linear weightQuery;
        dl::Linear weightKey;
        dl::Linear weightValue;
        dl::Linear weightOut;
        dl::Linear weightIntermed;
        dl::LayerNorm mhaNorm;
        // FFN
        dl::Linear weightIntermedOut;
        dl::LayerNorm ffnNorm;
 
        TransformerEncoder(TransformerEncoder& other) = delete;
        TransformerEncoder(TransformerEncoder&& other) = delete;
 
    public:
        TransformerEncoder(TransformerConf conf) noexcept;
        virtual TensorPtr forward(TensorPtr input) override;
 
        const float dimKeysInvSqrt;
 
        TensorPtr scaledDotProductAttention(TensorPtr query, TensorPtr key, TensorPtr value) noexcept;
 
        TensorPtr multiHeadAttention(TensorPtr query, TensorPtr key, TensorPtr value) noexcept;
    };
 
    class Transformer final : public Model<TensorPtr(TensorPtr)> {
    public:
        const TransformerConf conf;
        std::vector<std::unique_ptr<TransformerEncoder>> encoders;
        dl::Linear weightOut;
 
    public:
        Transformer(TransformerConf conf) noexcept;
 
        virtual TensorPtr forward(TensorPtr input) override;
    };
}; // namespace dl