transformer-lm-japanese-1.0b

transformer-lm-japanese-1.0b

This is a JAX/Flax-based transformer language model trained on a Japanese dataset. It is based on the official Flax example code (lm1b).

Source Code

We've modified Flax's 'lm1b' example to train on Japanese dataset. You can find the code on Github.

• transformer-lm-japanese

Our Blog Post

• 【0.1Bから作るLLM】 JAX/Flaxで作るTransformer言語モデル

Model Details

Usage: FlaxAutoModel

Requirements:

pip install transformers>=4.39.0
pip install jax==0.4.31
pip install flax==0.8.3
pip install sentencepiece==0.1.99

# For CPU
pip install -U "jax[cpu]==0.4.31"

# For GPU
pip install -U "jax[cuda12]==0.4.31"

Note: Set trust_remote_code=True to load our custom model.

from transformers import AutoTokenizer, FlaxAutoModelForCausalLM

tokenizer = AutoTokenizer.from_pretrained("fukugawa/transformer-lm-japanese-1.0b", trust_remote_code=True)
model = FlaxAutoModelForCausalLM.from_pretrained("fukugawa/transformer-lm-japanese-1.0b", trust_remote_code=True)

text = "日本の首都は、"
token_ids = tokenizer.encode(text, return_tensors="jax", add_special_tokens=False)

output_ids = model.generate(
  token_ids,
  do_sample=True,
  temperature=0.6,
  top_k=20,
  max_new_tokens=100
)

output = tokenizer.decode(output_ids[0][0], skip_special_tokens=True)
print(output)

We tested text generation in a Python 3.10 environment on GCP as follows

• GPU Type: NVIDIA L4 (x 1)
• Machine Type: g2-standard-16 (16 CPUs, 64GB Memory)
• Disk: 256GB
• OS: Ubuntu 22.04 LTS x86/64

Dataset

• wiki40b/ja (2.19GB)

Tokenization

• sentencepiece

Author

Ryoichi Fukugawa