Update Colab Script: TensorFlow 2.x Compatibility and Minor Model Fixes

[Yuta-chan]

While running the Colab script in the repository, I encountered several issues due to outdated code that is no longer compatible with the current environment and libraries. Specifically, the script relies on TensorFlow 1.x and does not function properly with TensorFlow 2.x, which is now the default version in Google Colab.

Additionally, there are issues in the model initialization where the input shape and the final output layer are incompatible with TensorFlow 2.x.

Changes Made

1. Initial Cell Setup:

Replaced the TensorFlow version check (originally tailored for TensorFlow 1.x) with a simplified check for whether the code is running in Google Colab:

try: import google.colab IN_COLAB = True except: IN_COLAB = False import os, sys import tensorflow as tf from tqdm import tqdm_notebook as tqdm import numpy as np import matplotlib.pyplot as plt

This ensures compatibility with TensorFlow 2.x in Colab without relying on deprecated version checks.

2.Model Initialization:

Fixed the input layer definition to specify the shape as a tuple (x,) rather than just x, as TensorFlow 2.x requires the input shape to be a tuple:
inputs = tf.keras.Input(shape=(3 + 3*2*L_embed,))

3. Output Layer Fix:
   Added a Lambda layer to handle the output concatenation properly, as TensorFlow 2.x requires the tf.concat operation to be wrapped within a Lambda layer:

outputs = tf.keras.layers.Lambda(lambda x: tf.concat(x, axis=-1))([outputs, inputs])

The original code directly applied tf.concat outside a layer, which caused an error in TensorFlow 2.x.

---

Therefore, I suggest updating the Colab script in the repository with these modifications to ensure compatibility with current environments.

---

pd: it's my first time writing an issue. I hope it has been done properly. Please notify me otherwise.