Resolve #7 and #16

tensorops · Sep 3, 2022 · 2b5ba69 · 2b5ba69
1 parent 60e75d1
commit 2b5ba69
Show file tree

Hide file tree

Showing 16 changed files with 27 additions and 23 deletions.
diff --git a/examples/eng2fr_translation.py b/examples/eng2fr_translation.py
@@ -1,9 +1,8 @@
 # **NOTE**: This example will be heavily edited, hence, this is not an official part of the library at this time
 
 from data_loader import BaseDataset
-from layers.transformer_decoder import TransformerDecoder
-from layers.transformer_encoder import TransformerEncoder
 from training.base import Transformer, Trainer
+from transformerx.layers import TransformerEncoder, TransformerDecoder
 
 depth, n_blocks, dropout = 256, 2, 0.2
 ffn_num_hiddens, num_heads = 64, 4

diff --git a/layers/__init__.py b/layers/__init__.py
diff --git a/test_main.py b/test_main.py
@@ -1,17 +1,12 @@
-import os
-
 import numpy as np
 import pytest
 import tensorflow as tf
 
-from layers.addnorm import AddNorm
-from layers.multihead_attention import MultiHeadAttention
-from layers.positional_encoding import PositionalEncoding
-from layers.positionwise_ffn import PositionWiseFFN
-from layers.transformer_decoder_block import TransformerDecoderBlock
-from layers.transformer_encoder import TransformerEncoder
-from layers.transformer_encoder_block import TransformerEncoderBlock
-from txplot.plot_pe import Plot
+from transformerx.layers import (
+    MultiHeadAttention, PositionalEncoding, PositionWiseFFN, AddNorm,
+    TransformerEncoderBlock, TransformerEncoder, TransformerDecoderBlock,
+)
+from transformerx.txplot import Plot
 
 
 @pytest.fixture()

diff --git a/transformerx/layers/__init__.py b/transformerx/layers/__init__.py
@@ -0,0 +1,9 @@
+from .addnorm import AddNorm
+from .dot_product_attention import DotProductAttention
+from .multihead_attention import MultiHeadAttention
+from .positional_encoding import PositionalEncoding
+from .positionwise_ffn import PositionWiseFFN
+from .transformer_decoder import TransformerDecoder
+from .transformer_decoder_block import TransformerDecoderBlock
+from .transformer_encoder import TransformerEncoder
+from .transformer_encoder_block import TransformerEncoderBlock
diff --git a/layers/addnorm.py → transformerx/layers/addnorm.py b/layers/addnorm.py → transformerx/layers/addnorm.py
diff --git a/layers/dot_product_attention.py → transformerx/layers/dot_product_attention.py b/layers/dot_product_attention.py → transformerx/layers/dot_product_attention.py
diff --git a/layers/multihead_attention.py → transformerx/layers/multihead_attention.py b/layers/multihead_attention.py → transformerx/layers/multihead_attention.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
 from einops import rearrange
 
-from layers.dot_product_attention import DotProductAttention
+from transformerx.layers.dot_product_attention import DotProductAttention
 
 
 class MultiHeadAttention(tf.keras.layers.Layer):

diff --git a/layers/positional_encoding.py → transformerx/layers/positional_encoding.py b/layers/positional_encoding.py → transformerx/layers/positional_encoding.py
diff --git a/layers/positionwise_ffn.py → transformerx/layers/positionwise_ffn.py b/layers/positionwise_ffn.py → transformerx/layers/positionwise_ffn.py
diff --git a/layers/transformer_decoder.py → transformerx/layers/transformer_decoder.py b/layers/transformer_decoder.py → transformerx/layers/transformer_decoder.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
 
-from layers.positional_encoding import PositionalEncoding
-from layers.transformer_decoder_block import TransformerDecoderBlock
+from transformerx.layers.positional_encoding import PositionalEncoding
+from transformerx.layers.transformer_decoder_block import TransformerDecoderBlock
 
 
 class TransformerDecoder(tf.keras.layers.Layer):

diff --git a/layers/transformer_decoder_block.py → ...rmerx/layers/transformer_decoder_block.py b/layers/transformer_decoder_block.py → ...rmerx/layers/transformer_decoder_block.py
@@ -1,8 +1,8 @@
 import tensorflow as tf
 
-from layers.addnorm import AddNorm
-from layers.multihead_attention import MultiHeadAttention
-from layers.positionwise_ffn import PositionWiseFFN
+from transformerx.layers.addnorm import AddNorm
+from transformerx.layers.multihead_attention import MultiHeadAttention
+from transformerx.layers.positionwise_ffn import PositionWiseFFN
 
 
 class TransformerDecoderBlock(tf.keras.layers.Layer):

diff --git a/layers/transformer_encoder.py → transformerx/layers/transformer_encoder.py b/layers/transformer_encoder.py → transformerx/layers/transformer_encoder.py
@@ -1,7 +1,7 @@
 import tensorflow as tf
 
-from layers.positional_encoding import PositionalEncoding
-from layers.transformer_encoder_block import TransformerEncoderBlock
+from transformerx.layers.positional_encoding import PositionalEncoding
+from transformerx.layers.transformer_encoder_block import TransformerEncoderBlock
 
 
 class TransformerEncoder(tf.keras.layers.Layer):

diff --git a/layers/transformer_encoder_block.py → ...rmerx/layers/transformer_encoder_block.py b/layers/transformer_encoder_block.py → ...rmerx/layers/transformer_encoder_block.py
@@ -1,8 +1,8 @@
 import tensorflow as tf
 
-from layers.addnorm import AddNorm
-from layers.multihead_attention import MultiHeadAttention
-from layers.positionwise_ffn import PositionWiseFFN
+from transformerx.layers.addnorm import AddNorm
+from transformerx.layers.multihead_attention import MultiHeadAttention
+from transformerx.layers.positionwise_ffn import PositionWiseFFN
 
 
 class TransformerEncoderBlock(tf.keras.layers.Layer):

diff --git a/transformerx/txplot/__init__.py b/transformerx/txplot/__init__.py
@@ -0,0 +1 @@
+from .plot_pe import Plot
diff --git a/txplot/plot_pe.py → transformerx/txplot/plot_pe.py b/txplot/plot_pe.py → transformerx/txplot/plot_pe.py
diff --git a/txplot/__init__.py b/txplot/__init__.py