Train an Autoencoder for Learning Text in TensorFlow 2

Syllable-based (subword-based) encoder or autoencoder (word2vec) are useful for learning text in ML as they make words related to one another instead of difference indices in vocabulary with no relation at all. The following code train an autoencoder to encode a word index into vector.

Source code:

#!pip install tensorflow==2.0.0
%reset -f

#libs
import tensorflow as tf;

#data---------------------------------------------------------------------------
X = [[0],[1],[2]];
Y = [[0],[1],[2]];
Max_X = 2;
Max_Y = 2;

#normalise
for I in range(len(X)):
  #X[I][0] /= Max_X; #don't normalise X, the Embedding layer needs indices
  Y[I][0] /= Max_Y;
#end for

X = tf.convert_to_tensor(X,tf.float32);
Y = tf.convert_to_tensor(Y,tf.float32);

#model--------------------------------------------------------------------------
Layer1 = tf.keras.layers.Embedding(3,2);

W2 = tf.Variable(tf.random.uniform([2,1], -1,1));
B2 = tf.Variable(tf.random.uniform([  1], -1,1));

@tf.function
def feedforward(X):
  H1  = Layer1(X);
  H1  = tf.keras.backend.batch_flatten(H1);
  Out = tf.sigmoid(tf.matmul(H1,W2) + B2);
  return Out;
#end def

def encode(X): #encode number into vector
  return Layer1(X);

def decode(V): #decode vector back to number
  H1  = tf.keras.backend.batch_flatten(V);
  Out = tf.sigmoid(tf.matmul(H1,W2) + B2);
  return tf.round(Out*Max_Y);

#train--------------------------------------------------------------------------
def get_loss(Out):
  return tf.reduce_sum(tf.square(Y-Out));
#end def

Optim = tf.keras.optimizers.SGD(1e-1);
Steps = 1000;

for I in range(Steps):
  if I%(Steps/10)==0:
    Out  = feedforward(X);
    Loss = get_loss(Out);
    print("Loss:",Loss.numpy());
  #end if

  with tf.GradientTape() as T:
    Out  = feedforward(X);
    Loss = get_loss(Out);
  #end with

  Grads = T.gradient(Loss, Layer1.trainable_variables+[W2,B2]);
  Optim.apply_gradients(zip(Grads, Layer1.trainable_variables+[W2,B2]));
#end for

Out  = feedforward(X);
Loss = get_loss(Out);
print("Loss:",Loss.numpy(),"(Last)");
print("\nEval:");
print(tf.round(Out*Max_Y).numpy());

print("\nEncode value 1:");
Embeddings = encode(1);
print(Embeddings.numpy());

print("\nDecode embeddings of 1:");
print(decode([Embeddings]).numpy());

print("\nDone.");
#eof