I am trying to translate the Reformer machine_translation code
https://github.com/google/trax/blob/master/trax/models/reformer/machine_translation.ipynb
that was written for trax 1.2.9 to make it work with trax 1.4.1
However, my program crashes when I try to load model.pkl with this error:
"ModuleNotFoundError: No module named 'trax.history'"
My understanding is that trax.history in 1.2.9 has been moved to trax.supervised.history in 1.4.1.
I believe I need a new model.pkl for 1.4.1 to make it work. Where can I download the new model.pkl?
It would also be great if I can download the new config.gin as well. Thanks a lot in advance.
...
Environment information
OS: <your answer here>
Ubuntu 20.04
$ pip freeze | grep trax
# your output here
trax==1.4.1
$ pip freeze | grep tensor
# your output here
tensorboard @ file:///home/conda/feedstock_root/build_artifacts/tensorboard_1664238338171/work/tensorboard-2.10.1-py3-none-any.whl
tensorboard-data-server @ file:///home/conda/feedstock_root/build_artifacts/tensorboard-data-server_1649932776625/work/tensorboard_data_server-0.6.0-py3-none-manylinux2010_x86_64.whl
tensorboard-plugin-wit @ file:///home/conda/feedstock_root/build_artifacts/tensorboard-plugin-wit_1641458951060/work/tensorboard_plugin_wit-1.8.1-py3-none-any.whl
tensorflow==2.10.0
tensorflow-datasets==4.7.0
tensorflow-estimator @ file:///home/conda/feedstock_root/build_artifacts/tensorflow-split_1663957899180/work/tensorflow-estimator/wheel_dir/tensorflow_estimator-2.10.0-py2.py3-none-any.whl
tensorflow-hub==0.12.0
tensorflow-io-gcs-filesystem==0.27.0
tensorflow-metadata==1.10.0
tensorflow-text==2.10.0
$ pip freeze | grep jax
# your output here
jax @ file:///home/conda/feedstock_root/build_artifacts/jax_1665610009116/work
jaxlib==0.3.22
$ python -V
# your output here
Python 3.8.10
For bugs: reproduction and error logs
# Steps to reproduce:
...
import sys
import gin
import os
import pickle
import jax
import trax
import numpy as np
import jax.numpy as jnp
import sacrebleu
from trax.data.text_encoder import SubwordTextEncoder
from tensorflow.io.gfile import GFile
Load the source text and reference translations into Python
refs = []
for lineno, line in enumerate(sacrebleu.smart_open('wmt14-en-de.ref'), 1):
if line.endswith('\n'):
line = line[:-1]
refs.append(line)
srcs = []
for lineno, line in enumerate(sacrebleu.smart_open('wmt14-en-de.src'), 1):
if line.endswith('\n'):
line = line[:-1]
srcs.append(line)
input_ids = np.zeros((len(srcs), 128), dtype=jnp.int64)
for i, x in enumerate(srcs):
x = tokenizer.encode(x)
assert len(x) <= 127
input_ids[i, :len(x)] = x
input_ids[i, len(x)] = 1
We'll be using a pre-trained reversible transformer-base model.
First, load the config (which sets all needed hyperparameters).
with GFile('gs://trax-ml/reformer/mt/model.pkl', 'rb') as f:
model_weights = pickle.load(f)['weights']
# Error logs:
...
Traceback (most recent call last):
File "reformer.py", line 65, in
model_weights = pickle.load(f)['weights']
ModuleNotFoundError: No module named 'trax.history'
Description
I am trying to translate the Reformer machine_translation code https://github.com/google/trax/blob/master/trax/models/reformer/machine_translation.ipynb that was written for trax 1.2.9 to make it work with trax 1.4.1 However, my program crashes when I try to load model.pkl with this error: "ModuleNotFoundError: No module named 'trax.history'" My understanding is that trax.history in 1.2.9 has been moved to trax.supervised.history in 1.4.1. I believe I need a new model.pkl for 1.4.1 to make it work. Where can I download the new model.pkl? It would also be great if I can download the new config.gin as well. Thanks a lot in advance. ...
Environment information
For bugs: reproduction and error logs
import sys import gin import os import pickle import jax import trax import numpy as np import jax.numpy as jnp import sacrebleu from trax.data.text_encoder import SubwordTextEncoder from tensorflow.io.gfile import GFile
Load the source text and reference translations into Python
refs = [] for lineno, line in enumerate(sacrebleu.smart_open('wmt14-en-de.ref'), 1): if line.endswith('\n'): line = line[:-1] refs.append(line) srcs = [] for lineno, line in enumerate(sacrebleu.smart_open('wmt14-en-de.src'), 1): if line.endswith('\n'): line = line[:-1] srcs.append(line)
Set up our sub-word tokenizer
tokenizer = SubwordTextEncoder( 'gs://trax-ml/reformer/mt/vocab.translate_ende_wmt32k.32768.subwords')
Encode source sentences using the tokenizer
input_ids = np.zeros((len(srcs), 128), dtype=jnp.int64) for i, x in enumerate(srcs): x = tokenizer.encode(x) assert len(x) <= 127 input_ids[i, :len(x)] = x input_ids[i, len(x)] = 1
We'll be using a pre-trained reversible transformer-base model.
First, load the config (which sets all needed hyperparameters).
!gsutil cp gs://trax-ml/reformer/mt/config.gin ./config.gin
gin.parse_config_file('./config.gin')
Now we load the pre-trained model weights.
with GFile('gs://trax-ml/reformer/mt/model.pkl', 'rb') as f: model_weights = pickle.load(f)['weights']
Traceback (most recent call last): File "reformer.py", line 65, in
model_weights = pickle.load(f)['weights']
ModuleNotFoundError: No module named 'trax.history'