Repository: openspeech-team/openspeech (language: python)

Issue:
Change `log_steps()` => `info()` in `OpenspeechModel`
# 🚀 Change `log_steps()` => `info()` in `OpenspeechModel`

- Current implementation of `OpenspeechModel` logging

```python
def log_steps(...):
    self.log(f"{stage}_wer", wer)
    self.log(f"{stage}_cer", cer)
    ...
```

- Modified version

```python
def info(self, dictionary: dict) -> None:
    r"""
    Logging information from dictionary.

    Args:
        dictionary (dict): dictionary contains information.
    """
    for key, value in dictionary.items():
        self.log(key, value, prog_bar=True)
```

## Motivation

- For scalability

## Your contribution

- implementation


Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 7fd9edb..e56235c 100644
--- a/README.md
+++ b/README.md
@@ -85,8 +85,8 @@ Because of these advantages, many end-to-end speech recognition related open sou
   
 ## Why shouldn't I use OpenSpeech?
   
-- This library provides code for learning ASR models, but does not provide APIs by pre-trained models.  
-- We do not provide pre-training mechanisms such as Wav2vec 2.0 since pre-training costs a lot of computation. Though computation optimization is very important, and this library does not provide that optimization.
+- This library provides code for training ASR models, but does not provide APIs by pre-trained models.  
+- This library does not provides pre-trained models.
   
 ## Model architectures  
   
@@ -95,6 +95,7 @@ We support all the models below. Note that, the important concepts of the model
 1. [**DeepSpeech2**](https://sooftware.github.io/OpenSpeech/DeepSpeech2.html) (from Baidu Research) released with paper [Deep Speech 2: End-to-End Speech Recognition in
 English and Mandarin](https://arxiv.org/abs/1512.02595.pdf), by Dario Amodei, Rishita Anubhai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Jingdong Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Erich Elsen, Jesse Engel, Linxi Fan, Christopher Fougner, Tony Han, Awni Hannun, Billy Jun, Patrick LeGresley, Libby Lin, Sharan Narang, Andrew Ng, Sherjil Ozair, Ryan Prenger, Jonathan Raiman, Sanjeev Satheesh, David Seetapun, Shubho Sengupta, Yi Wang, Zhiqian Wang, Chong Wang, Bo Xiao, Dani Yogatama, Jun Zhan, Zhenyao Zhu. 
 2. [**RNN-Transducer**](https://sooftware.github.io/OpenSpeech/RNN%20Transducer.html) (from University of Toronto) released with paper [Sequence Transduction with Recurrent Neural Networks](https://arxiv.org/abs/1211.3711.pdf), by Alex Graves.
+3. [**LSTM Language Model**]() (from RWTH Aachen University) released with paper [LSTM Neural Networks for Language Modeling](http://www-i6.informatik.rwth-aachen.de/publications/download/820/Sundermeyer-2012.pdf), by  Martin Sundermeyer, Ralf Schluter, and Hermann Ney.  
 3. [**Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Listen%20Attend%20Spell.html) (from Carnegie Mellon University and Google Brain) released with paper [Listen, Attend and Spell](https://arxiv.org/abs/1508.01211), by William Chan, Navdeep Jaitly, Quoc V. Le, Oriol Vinyals.  
 4. [**Location-aware attention based Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Listen%20Attend%20Spell%20(location-aware).html) (from University of Wrocław and Jacobs University and Universite de Montreal) released with paper [Attention-Based Models for Speech Recognition](https://arxiv.org/abs/1506.07503), by Jan Chorowski, Dzmitry Bahdanau, Dmitriy Serdyuk, Kyunghyun Cho, Yoshua Bengio.  
 5. [**Joint CTC-Attention based Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Joint%20CTC%20LAS.html) (from Mitsubishi Electric Research Laboratories and Carnegie Mellon University) released with paper [Joint CTC-Attention based End-to-End Speech Recognition using Multi-task Learning](https://arxiv.org/abs/1609.06773), by Suyoun Kim, Takaaki Hori, Shinji Watanabe.  
@@ -104,6 +105,7 @@ English and Mandarin](https://arxiv.org/abs/1512.02595.pdf), by Dario Amodei, Ri
 9. [**VGG-Transformer**](https://sooftware.github.io/OpenSpeech/VGG%20Transformer.html) (from Facebook AI Research) released with paper [Transformers with convolutional context for ASR](https://arxiv.org/abs/1904.11660), by Abdelrahman Mohamed, Dmytro Okhonko, Luke Zettlemoyer.  
 10. [**Transformer with CTC**](https://sooftware.github.io/OpenSpeech/Transformer%20with%20CTC.html) (from NTT Communication Science Laboratories, Waseda University, Center for Language and Speech Processing, Johns Hopkins University) released with paper [Improving Transformer-based End-to-End Speech Recognition with Connectionist Temporal Classification and Language Model Integration](https://www.isca-speech.org/archive/Interspeech_2019/pdfs/1938.pdf), by Shigeki Karita, Nelson Enrique Yalta Soplin, Shinji Watanabe, Marc Delcroix, Atsunori Ogawa, Tomohiro Nakatani.
 11. [**Joint CTC-Attention based Transformer**](https://sooftware.github.io/OpenSpeech/Joint%20CTC%20Transformer.html)(from NTT Corporation) released with paper [Self-Distillation for Improving CTC-Transformer-based ASR Systems](https://www.isca-speech.org/archive/Interspeech_2020/pdfs/1223.pdf), by Takafumi Moriya, Tsubasa Ochiai, Shigeki Karita, Hiroshi Sato, Tomohiro Tanaka, Takanori Ashihara, Ryo Masumura, Yusuke Shinohara, Marc Delcroix.
+12. [**Transformer Language Model**]() (from Amazon Web Services) released with paper [Language Models with Transformers](https://arxiv.org/abs/1904.09408), by Chenguang Wang, Mu Li, Alexander J. Smola.
 12. [**Jasper**](https://sooftware.github.io/OpenSpeech/Jasper10x5.html) (from NVIDIA and New York University) released with paper [Jasper: An End-to-End Convolutional Neural Acoustic Model](https://arxiv.org/pdf/1904.03288.pdf), by Jason Li, Vitaly Lavrukhin, Boris Ginsburg, Ryan Leary, Oleksii Kuchaiev, Jonathan M. Cohen, Huyen Nguyen, Ravi Teja Gadde.   
 13. [**QuartzNet**](https://sooftware.github.io/OpenSpeech/QuartzNet15x5.html) (from NVIDIA and Univ. of Illinois and Univ. of Saint Petersburg) released with paper [QuartzNet: Deep Automatic Speech Recognition with 1D Time-Channel Separable Convolutions](https://arxiv.org/abs/1910.10261.pdf), by Samuel Kriman, Stanislav Beliaev, Boris Ginsburg, Jocelyn Huang, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Yang Zhang.  
 14. [**Transformer Transducer**](https://sooftware.github.io/OpenSpeech/Transformer%20Transducer.html) (from Facebook AI) released with paper [Transformer-Transducer:
@@ -154,7 +156,7 @@ Please send e-mail including the approved screenshot to sh951011@gmail.com.
   
 ### Manifest File
   
-- Manifest file format:
+- Acoustic model manifest file format:
   
 ```
 LibriSpeech/test-other/8188/269288/8188-269288-0052.flac        ▁ANNIE ' S ▁MANNER ▁WAS ▁VERY ▁MYSTERIOUS       4039 20 5 531 17 84 2352
@@ -249,6 +251,33 @@ $ python ./openspeech_cli/hydra_eval.py \
     eval.ensemble_method=weighted \
     eval.manifest_file_path=$MANIFEST_FILE_PATH  
 ```
+
+### Language model training example
+  
+Language model training requires only data to be prepared in the following format:  
+  
+```
+openspeech is a framework for making end-to-end speech recognizers.
+end to end automatic speech recognition is an emerging paradigm in the field of neural network-based speech recognition that offers multiple benefits.
+because of these advantages, many end-to-end speech recognition related open sources have emerged.
+...
+...
+```
+  
+Note that you need to use the same vocabulary as the acoustic model.
+  
+- Example: Train the `lstm_lm` model:
+```
+$ python ./openspeech_cli/hydra_lm_train.py \
+    dataset=lm \
+    dataset.dataset_path=../../../lm.txt \
+    vocab=kspon_character \
+    vocab.vocab_path=../../../labels.csv \
+    model=lstm_lm \
+    lr_scheduler=tri_stage \
+    trainer=gpu \
+    criterion=perplexity
+```
   
 ## Installation
   
diff --git a/openspeech/configs/lm_train.yaml b/openspeech/configs/lm_train.yaml
new file mode 100644
index 0000000..e84f32b
--- /dev/null
+++ b/openspeech/configs/lm_train.yaml
@@ -0,0 +1,9 @@
+# @package _group_
+
+defaults:
+  - dataset: null
+  - criterion: null
+  - lr_scheduler: null
+  - model: null
+  - trainer: gpu
+  - vocab: null
diff --git a/openspeech/criterion/__init__.py b/openspeech/criterion/__init__.py
index 39e9e89..0e9e9c7 100644
--- a/openspeech/criterion/__init__.py
+++ b/openspeech/criterion/__init__.py
@@ -81,11 +81,13 @@ def register_criterion_cls(cls):
 from .joint_ctc_cross_entropy.configuration import JointCTCCrossEntropyLossConfigs
 from .label_smoothed_cross_entropy.configuration import LabelSmoothedCrossEntropyLossConfigs
 from .transducer.configuration import TransducerLossConfigs
+from .perplexity.perplexity import PerplexityLossConfigs
 from .cross_entropy.cross_entropy import CrossEntropyLoss
 from .ctc.ctc import CTCLoss
 from .joint_ctc_cross_entropy.joint_ctc_cross_entropy import JointCTCCrossEntropyLoss
 from .label_smoothed_cross_entropy.label_smoothed_cross_entropy import LabelSmoothedCrossEntropyLoss
 from .transducer.transducer import TransducerLoss
+from .perplexity.perplexity import Perplexity
 
 __all__ = [
     "CrossEntropyLossConfigs",
diff --git a/openspeech/criterion/perplexity/__init__.py b/openspeech/criterion/perplexity/__init__.py
new file mode 100644
index 0000000..9a083c6
--- /dev/null
+++ b/openspeech/criterion/perplexity/__init__.py
@@ -0,0 +1,21 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
\ No newline at end of file
diff --git a/openspeech/criterion/perplexity/configuration.py b/openspeech/criterion/perplexity/configuration.py
new file mode 100644
index 0000000..0f8428c
--- /dev/null
+++ b/openspeech/criterion/perplexity/configuration.py
@@ -0,0 +1,48 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from dataclasses import dataclass, field
+
+from ...dataclass.configurations import OpenspeechDataclass
+
+
+@dataclass
+class PerplexityLossConfigs(OpenspeechDataclass):
+    r"""
+    This is the configuration class to store the configuration of a
+    :class: `~openspeech.criterion.Perplexity`.
+
+    It is used to initiated an `CrossEntropyLoss` criterion.
+
+    Configuration objects inherit from :class: `~openspeech.dataclass.configs.OpenspeechDataclass`.
+
+    Configurations:
+        criterion_name (str): name of criterion (default: perplexity)
+        reduction (str): reduction method of criterion (default: mean)
+    """
+    criterion_name: str = field(
+        default="perplexity", metadata={"help": "Criterion name for training"}
+    )
+    reduction: str = field(
+        default="mean", metadata={"help": "Reduction method of criterion"}
+    )
+
diff --git a/openspeech/criterion/perplexity/perplexity.py b/openspeech/criterion/perplexity/perplexity.py
new file mode 100644
index 0000000..438fb7e
--- /dev/null
+++ b/openspeech/criterion/perplexity/perplexity.py
@@ -0,0 +1,83 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+from omegaconf import DictConfig
+
+from .. import register_criterion
+from ..perplexity.configuration import PerplexityLossConfigs
+from ...vocabs.vocab import Vocabulary
+
+
+@register_criterion("perplexity", dataclass=PerplexityLossConfigs)
+class Perplexity(nn.Module):
+    r"""
+    Language model perplexity loss.
+    Perplexity is the token averaged likelihood.  When the averaging options are the
+    same, it is the exponential of negative log-likelihood.
+
+    Configurations:
+        criterion_name (str): name of criterion
+        reduction (str): reduction method of criterion
+
+    Args:
+        configs (DictConfig): hydra configuration set
+        vocab (Vocabulary): the set of unique words used in the text corpus
+
+    Inputs: logits, targets
+        - **logits** (torch.FloatTensor): probability distribution value from model and it has a logarithm shape.
+            The `FloatTensor` of size ``(batch, seq_length, num_classes)``
+        - **targets** (torch.LongTensor): ground-truth encoded to integers which directly point a word in label
+            The `LongTensor` of size ``(batch, target_length)``
+
+    Returns: loss
+        * loss (float): loss for training
+    """
+    def __init__(
+            self,
+            configs: DictConfig,
+            vocab: Vocabulary,
+    ) -> None:
+        super(Perplexity, self).__init__()
+        self.cross_entropy_loss = nn.CrossEntropyLoss(
+            reduction=configs.criterion.reduction,
+            ignore_index=vocab.pad_id,
+        )
+
+    def forward(self, logits: Tensor, targets: Tensor) -> Tensor:
+        max_target_length = targets.size(1)
+        max_logits_length = logits.size(1)
+
+        if max_logits_length > max_target_length:
+            logits = logits[:, :max_target_length, :]
+        elif max_target_length > max_logits_length:
+            targets = targets[:, :max_logits_length]
+
+        logits = logits.contiguous().view(-1, logits.size(-1))
+
+        cross_entropy_loss = self.cross_entropy_loss(
+            logits.contiguous().view(-1, logits.size(-1)),
+            targets.contiguous().view(-1),
+        )
+        return torch.exp(cross_entropy_loss)
diff --git a/openspeech/data/data_loader.py b/openspeech/data/audio/data_loader.py
similarity index 82%
rename from openspeech/data/data_loader.py
rename to openspeech/data/audio/data_loader.py
index e73886b..45a26fa 100644
--- a/openspeech/data/data_loader.py
+++ b/openspeech/data/audio/data_loader.py
@@ -19,10 +19,10 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
-from typing import Tuple
 
 import torch
 import numpy as np
+from typing import Tuple
 from torch.utils.data import DataLoader, Sampler
 
 
@@ -106,35 +106,6 @@ def __init__(
         self.collate_fn = _collate_fn
 
 
-class BucketingSampler(Sampler):
-    r"""
-    Samples batches assuming they are in order of size to batch similarly sized samples together.
-
-    Args:
-        data_source (torch.utils.data.Dataset): dataset to sample from
-        batch_size (int): size of batch
-        drop_last (bool): flat indication whether to drop last batch or not
-    """
-    def __init__(self, data_source, batch_size: int = 32, drop_last: bool = False) -> None:
-        super(BucketingSampler, self).__init__(data_source)
-        self.batch_size = batch_size
-        self.data_source = data_source
-        ids = list(range(0, len(data_source)))
-        self.bins = [ids[i:i + batch_size] for i in range(0, len(ids), batch_size)]
-        self.drop_last = drop_last
-
-    def __iter__(self):
-        for ids in self.bins:
-            np.random.shuffle(ids)
-            yield ids
-
-    def __len__(self):
-        return len(self.bins)
-
-    def shuffle(self, epoch):
-        np.random.shuffle(self.bins)
-
-
 def load_dataset(manifest_file_path: str) -> Tuple[list, list]:
     """
     Provides dictionary of filename and labels.
diff --git a/openspeech/data/dataset.py b/openspeech/data/audio/dataset.py
similarity index 97%
rename from openspeech/data/dataset.py
rename to openspeech/data/audio/dataset.py
index 545a129..dec6f3d 100644
--- a/openspeech/data/dataset.py
+++ b/openspeech/data/audio/dataset.py
@@ -29,9 +29,9 @@
 from torch import Tensor
 from torch.utils.data import Dataset
 
-from . import AUDIO_FEATURE_TRANSFORM_REGISTRY
-from .audio.augment import JoiningAugment, NoiseInjector, SpecAugment, TimeStretchAugment
-from .audio.load import load_audio
+from openspeech.data import AUDIO_FEATURE_TRANSFORM_REGISTRY
+from openspeech.data.audio.augment import JoiningAugment, NoiseInjector, SpecAugment, TimeStretchAugment
+from openspeech.data.audio.load import load_audio
 
 logger = logging.getLogger(__name__)
 
diff --git a/openspeech/data/sampler.py b/openspeech/data/sampler.py
new file mode 100644
index 0000000..e1d1e80
--- /dev/null
+++ b/openspeech/data/sampler.py
@@ -0,0 +1,54 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import numpy as np
+
+from torch.utils.data import Sampler
+
+
+class BucketingSampler(Sampler):
+    r"""
+    Samples batches assuming they are in order of size to batch similarly sized samples together.
+
+    Args:
+        data_source (torch.utils.data.Dataset): dataset to sample from
+        batch_size (int): size of batch
+        drop_last (bool): flat indication whether to drop last batch or not
+    """
+    def __init__(self, data_source, batch_size: int = 32, drop_last: bool = False) -> None:
+        super(BucketingSampler, self).__init__(data_source)
+        self.batch_size = batch_size
+        self.data_source = data_source
+        ids = list(range(0, len(data_source)))
+        self.bins = [ids[i:i + batch_size] for i in range(0, len(ids), batch_size)]
+        self.drop_last = drop_last
+
+    def __iter__(self):
+        for ids in self.bins:
+            np.random.shuffle(ids)
+            yield ids
+
+    def __len__(self):
+        return len(self.bins)
+
+    def shuffle(self, epoch):
+        np.random.shuffle(self.bins)
diff --git a/openspeech/data/text/data_loader.py b/openspeech/data/text/data_loader.py
new file mode 100644
index 0000000..151af59
--- /dev/null
+++ b/openspeech/data/text/data_loader.py
@@ -0,0 +1,87 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+from torch.utils.data import DataLoader
+
+
+def _collate_fn(batch, pad_id: int = 0):
+    r"""
+    Functions that pad to the maximum sequence length
+
+    Args:
+        batch (tuple): tuple contains input and target tensors
+        pad_id (int): identification of pad token
+
+    Returns:
+        inputs (torch.FloatTensor): tensor contains input sequences.
+    """
+    def seq_length_(p):
+        return len(p[0])
+
+    batch = sorted(batch, key=lambda sample: sample[0].size(0), reverse=True)
+
+    seq_lengths = [len(s[0]) for s in batch]
+
+    max_seq_sample = max(batch, key=seq_length_)[0]
+    max_seq_size = max_seq_sample.size(0)
+
+    batch_size = len(batch)
+
+    inputs = torch.zeros(batch_size, max_seq_size).fill_(pad_id).long()
+    targets = torch.zeros(batch_size, max_seq_size).fill_(pad_id).long()
+
+    for x in range(batch_size):
+        sample = batch[x]
+        input_var = sample[0]
+        target = sample[1]
+        inputs[x].narrow(0, 0, len(input_var)).copy_(torch.LongTensor(input_var))
+        targets[x].narrow(0, 0, len(target)).copy_(torch.LongTensor(target))
+
+    seq_lengths = torch.IntTensor(seq_lengths)
+
+    return inputs, seq_lengths, targets
+
+
+class TextDataLoader(DataLoader):
+    r"""
+    Text Data Loader
+
+    Args:
+        dataset (torch.utils.data.Dataset): dataset from which to load the data.
+        num_workers (int): how many subprocesses to use for data loading.
+        batch_sampler (torch.utils.data.sampler.Sampler): defines the strategy to draw samples from the dataset.
+    """
+    def __init__(
+            self,
+            dataset: torch.utils.data.Dataset,
+            num_workers: int,
+            batch_sampler: torch.utils.data.sampler.Sampler,
+            **kwargs,
+    ) -> None:
+        super(TextDataLoader, self).__init__(
+            dataset=dataset,
+            num_workers=num_workers,
+            batch_sampler=batch_sampler,
+            **kwargs,
+        )
+        self.collate_fn = _collate_fn
diff --git a/openspeech/data/text/dataset.py b/openspeech/data/text/dataset.py
new file mode 100644
index 0000000..548f424
--- /dev/null
+++ b/openspeech/data/text/dataset.py
@@ -0,0 +1,70 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+import logging
+
+from torch.utils.data import Dataset
+
+logger = logging.getLogger(__name__)
+
+
+class TextDataset(Dataset):
+    def __init__(self, transcripts: list, vocab):
+        super(TextDataset, self).__init__()
+        self.transcripts = transcripts
+        self.vocab = vocab
+        self.sos_id = vocab.sos_id
+        self.eos_id = vocab.eos_id
+
+    def _get_inputs(self, transcript):
+        tokens = transcript.split(' ')
+        transcript = [int(self.sos_id)]
+
+        for token in tokens:
+            transcript.append(int(token))
+
+        return transcript
+
+    def _get_targets(self, transcript):
+        tokens = transcript.split(' ')
+        transcript = list()
+
+        for token in tokens:
+            transcript.append(int(token))
+
+        transcript.append(int(self.eos_id))
+
+        return transcript
+
+    def __getitem__(self, idx):
+        transcript = self.vocab.string_to_label(self.transcripts[idx])
+        inputs = torch.LongTensor(self._get_inputs(transcript))
+        targets = torch.LongTensor(self._get_targets(transcript))
+        return inputs, targets
+
+    def __len__(self):
+        return len(self.transcripts)
+
+    def count(self):
+        return len(self.transcripts)
+
diff --git a/openspeech/dataclass/__init__.py b/openspeech/dataclass/__init__.py
index 6ed175f..21ffe8f 100644
--- a/openspeech/dataclass/__init__.py
+++ b/openspeech/dataclass/__init__.py
@@ -25,6 +25,7 @@
     LibriSpeechConfigs,
     KsponSpeechConfigs,
     AIShellConfigs,
+    LMConfigs,
     CPUTrainerConfigs,
     GPUTrainerConfigs,
     TPUTrainerConfigs,
@@ -35,7 +36,7 @@
     EnsembleEvaluationConfigs,
 )
 
-OPENSPEECH_CONFIGS = [
+OPENSPEECH_TRAIN_CONFIGS = [
     "audio",
     "augment",
     "dataset",
@@ -46,12 +47,22 @@
     "vocab",
 ]
 
+OPENSPEECH_LM_TRAIN_CONFIGS = [
+    "dataset",
+    "model",
+    "criterion",
+    "lr_scheduler",
+    "trainer",
+    "vocab",
+]
+
 DATASET_DATACLASS_REGISTRY = {
     "kspon": KsponSpeechConfigs,
     "libri": LibriSpeechConfigs,
     "aishell": AIShellConfigs,
     "ksponspeech": KsponSpeechConfigs,
     "librispeech": LibriSpeechConfigs,
+    "lm": LMConfigs,
 }
 TRAINER_DATACLASS_REGISTRY = {
     "cpu": CPUTrainerConfigs,
diff --git a/openspeech/dataclass/configurations.py b/openspeech/dataclass/configurations.py
index 578c8ef..8268a7e 100644
--- a/openspeech/dataclass/configurations.py
+++ b/openspeech/dataclass/configurations.py
@@ -63,7 +63,7 @@ def _get_help(self, attribute_name: str) -> Any:
 class LibriSpeechConfigs(OpenspeechDataclass):
     """ Configuration dataclass that common used """
     dataset: str = field(
-        default="librispeech", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell)"}
+        default="librispeech", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell, lm)"}
     )
     dataset_path: str = field(
         default=MISSING, metadata={"help": "Path of dataset"}
@@ -80,7 +80,7 @@ class LibriSpeechConfigs(OpenspeechDataclass):
 class KsponSpeechConfigs(OpenspeechDataclass):
     """ Configuration dataclass that common used """
     dataset: str = field(
-        default="ksponspeech", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell)"}
+        default="ksponspeech", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell, lm)"}
     )
     dataset_path: str = field(
         default=MISSING, metadata={"help": "Path of dataset"}
@@ -103,7 +103,7 @@ class KsponSpeechConfigs(OpenspeechDataclass):
 class AIShellConfigs(OpenspeechDataclass):
     """ Configuration dataclass that common used """
     dataset: str = field(
-        default="aishell", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell)"}
+        default="aishell", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell, lm)"}
     )
     dataset_path: str = field(
         default=MISSING, metadata={"help": "Path of dataset"}
@@ -116,6 +116,22 @@ class AIShellConfigs(OpenspeechDataclass):
     )
 
 
+@dataclass
+class LMConfigs(OpenspeechDataclass):
+    dataset: str = field(
+        default="lm", metadata={"help": "Select dataset for training (librispeech, ksponspeech, aishell, lm)"}
+    )
+    dataset_path: str = field(
+        default=MISSING, metadata={"help": "Path of dataset"}
+    )
+    valid_ratio: float = field(
+        default=0.05, metadata={"help": "Ratio of validation data"}
+    )
+    test_ratio: float = field(
+        default=0.05, metadata={"help": "Ratio of test data"}
+    )
+
+
 @dataclass
 class AugmentConfigs(OpenspeechDataclass):
     apply_spec_augment: bool = field(
@@ -368,7 +384,7 @@ class EnsembleEvaluationConfigs(OpenspeechDataclass):
 
 
 def generate_openspeech_configs_with_help():
-    from openspeech.dataclass import OPENSPEECH_CONFIGS, TRAINER_DATACLASS_REGISTRY
+    from openspeech.dataclass import OPENSPEECH_TRAIN_CONFIGS, TRAINER_DATACLASS_REGISTRY
     from openspeech.models import MODEL_DATACLASS_REGISTRY
     from openspeech.criterion import CRITERION_DATACLASS_REGISTRY
     from openspeech.data import AUDIO_FEATURE_TRANSFORM_DATACLASS_REGISTRY
@@ -387,7 +403,7 @@ def generate_openspeech_configs_with_help():
     }
 
     with open("../../../configuration.md", "w") as f:
-        for group in OPENSPEECH_CONFIGS:
+        for group in OPENSPEECH_TRAIN_CONFIGS:
             dataclass_registry = registries[group]
 
             f.write(f"## `{group}`\n")
diff --git a/openspeech/dataclass/initialize.py b/openspeech/dataclass/initialize.py
index 2a0c7f6..238547d 100644
--- a/openspeech/dataclass/initialize.py
+++ b/openspeech/dataclass/initialize.py
@@ -23,9 +23,9 @@
 from hydra.core.config_store import ConfigStore
 
 
-def hydra_init() -> None:
+def hydra_train_init() -> None:
     r""" initialize ConfigStore for hydra-train """
-    from openspeech.dataclass import OPENSPEECH_CONFIGS, DATASET_DATACLASS_REGISTRY, TRAINER_DATACLASS_REGISTRY
+    from openspeech.dataclass import OPENSPEECH_TRAIN_CONFIGS, DATASET_DATACLASS_REGISTRY, TRAINER_DATACLASS_REGISTRY
     from openspeech.models import MODEL_DATACLASS_REGISTRY
     from openspeech.criterion import CRITERION_DATACLASS_REGISTRY
     from openspeech.data import AUDIO_FEATURE_TRANSFORM_DATACLASS_REGISTRY
@@ -46,7 +46,32 @@ def hydra_init() -> None:
 
     cs = ConfigStore.instance()
 
-    for group in OPENSPEECH_CONFIGS:
+    for group in OPENSPEECH_TRAIN_CONFIGS:
+        dataclass_registry = registries[group]
+
+        for k, v in dataclass_registry.items():
+            cs.store(group=group, name=k, node=v, provider="openspeech")
+
+
+def hydra_lm_train_init() -> None:
+    from openspeech.dataclass import OPENSPEECH_LM_TRAIN_CONFIGS, DATASET_DATACLASS_REGISTRY, TRAINER_DATACLASS_REGISTRY
+    from openspeech.models import MODEL_DATACLASS_REGISTRY
+    from openspeech.criterion import CRITERION_DATACLASS_REGISTRY
+    from openspeech.optim.scheduler import SCHEDULER_DATACLASS_REGISTRY
+    from openspeech.vocabs import VOCAB_DATACLASS_REGISTRY
+
+    registries = {
+        "dataset": DATASET_DATACLASS_REGISTRY,
+        "trainer": TRAINER_DATACLASS_REGISTRY,
+        "model": MODEL_DATACLASS_REGISTRY,
+        "criterion": CRITERION_DATACLASS_REGISTRY,
+        "lr_scheduler": SCHEDULER_DATACLASS_REGISTRY,
+        "vocab": VOCAB_DATACLASS_REGISTRY,
+    }
+
+    cs = ConfigStore.instance()
+
+    for group in OPENSPEECH_LM_TRAIN_CONFIGS:
         dataclass_registry = registries[group]
 
         for k, v in dataclass_registry.items():
diff --git a/openspeech/datasets/aishell/__init__.py b/openspeech/datasets/aishell/__init__.py
index e69de29..9a083c6 100644
--- a/openspeech/datasets/aishell/__init__.py
+++ b/openspeech/datasets/aishell/__init__.py
@@ -0,0 +1,21 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
\ No newline at end of file
diff --git a/openspeech/datasets/aishell/lit_data_module.py b/openspeech/datasets/aishell/lit_data_module.py
index e659af6..bc35f98 100644
--- a/openspeech/datasets/aishell/lit_data_module.py
+++ b/openspeech/datasets/aishell/lit_data_module.py
@@ -29,9 +29,10 @@
 from typing import Optional, Tuple
 from torch.utils.data import DataLoader
 
-from openspeech.data.dataset import SpeechToTextDataset
+from openspeech.data.audio.dataset import SpeechToTextDataset
 from openspeech.datasets import register_data_module
-from openspeech.data.data_loader import BucketingSampler, AudioDataLoader
+from openspeech.data.sampler import BucketingSampler
+from openspeech.data.audio.data_loader import AudioDataLoader
 from openspeech.vocabs.vocab import Vocabulary
 from openspeech.vocabs import VOCAB_REGISTRY
 from openspeech.datasets.aishell.preprocess import (
diff --git a/openspeech/datasets/ksponspeech/lit_data_module.py b/openspeech/datasets/ksponspeech/lit_data_module.py
index 000f491..840688e 100644
--- a/openspeech/datasets/ksponspeech/lit_data_module.py
+++ b/openspeech/datasets/ksponspeech/lit_data_module.py
@@ -25,10 +25,11 @@
 import pytorch_lightning as pl
 from typing import Optional
 from omegaconf import DictConfig
-from openspeech.data.dataset import SpeechToTextDataset
+from openspeech.data.audio.dataset import SpeechToTextDataset
 
 from openspeech.datasets import register_data_module
-from openspeech.data.data_loader import AudioDataLoader, BucketingSampler
+from openspeech.data.sampler import BucketingSampler
+from openspeech.data.audio.data_loader import AudioDataLoader
 from openspeech.datasets.ksponspeech.preprocess.preprocess import preprocess, preprocess_test_data
 from openspeech.datasets.ksponspeech.preprocess.character import generate_character_script, generate_character_labels
 from openspeech.datasets.ksponspeech.preprocess.grapheme import sentence_to_grapheme
@@ -56,7 +57,7 @@ def __init__(self, configs: DictConfig) -> None:
         super(LightningKsponSpeechDataModule, self).__init__()
         self.configs = configs
         self.dataset = dict()
-        self.logger = logging.getLogger("openspeech.data.ksponspeech.lit_data_module")
+        self.logger = logging.getLogger(__name__)
         self.encoding = 'utf-8' if self.configs.vocab.unit == 'kspon_subword' else 'cp949'
 
     def _generate_manifest_files(self, manifest_file_path: str) -> None:
@@ -122,7 +123,7 @@ def prepare_data(self):
             self.logger.info("Manifest file is not exists !!\n"
                              "Generate manifest files..")
             if not os.path.exists(self.configs.dataset.dataset_path):
-                raise ValueError("Dataset path is not valid.")
+                raise FileNotFoundError
             self._generate_manifest_files(self.configs.dataset.manifest_file_path)
         return VOCAB_REGISTRY[self.configs.vocab.unit](self.configs)
 
diff --git a/openspeech/datasets/ksponspeech/preprocess/grapheme.py b/openspeech/datasets/ksponspeech/preprocess/grapheme.py
index 4811d77..1e36c63 100644
--- a/openspeech/datasets/ksponspeech/preprocess/grapheme.py
+++ b/openspeech/datasets/ksponspeech/preprocess/grapheme.py
@@ -20,7 +20,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import os
 import pandas as pd
 import unicodedata
 
diff --git a/openspeech/datasets/language_model/__init__.py b/openspeech/datasets/language_model/__init__.py
new file mode 100644
index 0000000..9a083c6
--- /dev/null
+++ b/openspeech/datasets/language_model/__init__.py
@@ -0,0 +1,21 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
\ No newline at end of file
diff --git a/openspeech/datasets/language_model/lit_data_module.py b/openspeech/datasets/language_model/lit_data_module.py
new file mode 100644
index 0000000..62b0727
--- /dev/null
+++ b/openspeech/datasets/language_model/lit_data_module.py
@@ -0,0 +1,105 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import pytorch_lightning as pl
+import logging
+import os
+import random
+from omegaconf import DictConfig
+from typing import Optional
+
+from openspeech.data.sampler import BucketingSampler
+from openspeech.data.text.data_loader import TextDataLoader
+from openspeech.data.text.dataset import TextDataset
+from openspeech.datasets import register_data_module
+from openspeech.vocabs import VOCAB_REGISTRY
+from openspeech.vocabs.vocab import Vocabulary
+
+
+@register_data_module('lm')
+class LightningLanguageModelDataModule(pl.LightningDataModule):
+    def __init__(self, configs: DictConfig) -> None:
+        super(LightningLanguageModelDataModule, self).__init__()
+        self.configs = configs
+        self.dataset = dict()
+        self.logger = logging.getLogger(__name__)
+
+    def prepare_data(self):
+        if not os.path.exists(self.configs.dataset.dataset_path):
+            raise FileNotFoundError
+        return VOCAB_REGISTRY[self.configs.vocab.unit](self.configs)
+
+    def setup(self, stage: Optional[str] = None, vocab: Vocabulary = None):
+        num_total_transcripts = 0
+        transcripts = list()
+
+        with open(self.configs.dataset.dataset_path, encoding=self.configs.vocab.encoding) as f:
+            for line in f.readlines():
+                transcripts.append(line)
+                num_total_transcripts += 1
+
+        random.shuffle(transcripts)
+
+        train_ratio = (1 - self.configs.dataset.valid_ratio - self.configs.dataset.test_ratio)
+
+        num_train_transcripts = int(num_total_transcripts * train_ratio)
+        num_valid_transcripts = int(num_total_transcripts * self.configs.dataset.valid_ratio)
+
+        valid_end_idx = num_train_transcripts + num_valid_transcripts
+
+        transcripts = {
+            "train": transcripts[:num_train_transcripts],
+            "valid": transcripts[num_train_transcripts:valid_end_idx],
+            "test": transcripts[valid_end_idx:]
+        }
+
+        for stage in transcripts.keys():
+            self.dataset[stage] = TextDataset(
+                transcripts=transcripts[stage],
+                vocab=vocab,
+            )
+
+    def train_dataloader(self) -> TextDataLoader:
+        train_sampler = BucketingSampler(self.dataset['train'], batch_size=self.configs.trainer.batch_size)
+        return TextDataLoader(
+            dataset=self.dataset['train'],
+            num_workers=self.configs.trainer.num_workers,
+            batch_sampler=train_sampler,
+        )
+
+    def val_dataloader(self) -> TextDataLoader:
+        r""" Return data loader for validation. """
+        valid_sampler = BucketingSampler(self.dataset['valid'], batch_size=self.configs.trainer.batch_size)
+        return TextDataLoader(
+            dataset=self.dataset['valid'],
+            num_workers=self.configs.trainer.num_workers,
+            batch_sampler=valid_sampler,
+        )
+
+    def test_dataloader(self) -> TextDataLoader:
+        r""" Return data loader for training. """
+        train_sampler = BucketingSampler(self.dataset['test'], batch_size=self.configs.trainer.batch_size)
+        return TextDataLoader(
+            dataset=self.dataset['test'],
+            num_workers=self.configs.trainer.num_workers,
+            batch_sampler=train_sampler,
+        )
diff --git a/openspeech/datasets/librispeech/__init__.py b/openspeech/datasets/librispeech/__init__.py
index e69de29..9a083c6 100644
--- a/openspeech/datasets/librispeech/__init__.py
+++ b/openspeech/datasets/librispeech/__init__.py
@@ -0,0 +1,21 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
\ No newline at end of file
diff --git a/openspeech/datasets/librispeech/lit_data_module.py b/openspeech/datasets/librispeech/lit_data_module.py
index bec24c3..8c49725 100644
--- a/openspeech/datasets/librispeech/lit_data_module.py
+++ b/openspeech/datasets/librispeech/lit_data_module.py
@@ -28,13 +28,14 @@
 import pytorch_lightning as pl
 from typing import Tuple, Optional
 from omegaconf import DictConfig
-from openspeech.data.dataset import SpeechToTextDataset
+from openspeech.data.audio.dataset import SpeechToTextDataset
 from torch.utils.data import DataLoader
 
 from openspeech.datasets import register_data_module
 from openspeech.vocabs import VOCAB_REGISTRY
 from openspeech.vocabs.vocab import Vocabulary
-from openspeech.data.data_loader import BucketingSampler, AudioDataLoader
+from openspeech.data.sampler import BucketingSampler
+from openspeech.data.audio.data_loader import AudioDataLoader
 
 
 @register_data_module('librispeech')
@@ -127,11 +128,8 @@ def prepare_data(self) -> Vocabulary:
         Prepare librispeech data
 
         Returns:
-            vocab (Vocabulary): vocab class of KsponSpeech.
+            vocab (Vocabulary): vocab class of LibriSpeech.
         """
-        if not os.path.exists(self.configs.dataset.dataset_path):
-            raise ValueError("Dataset path is not valid.")
-
         if self.configs.vocab.unit == 'libri_subword':
             from openspeech.datasets.librispeech.preprocess.subword import generate_manifest_files
         elif self.configs.vocab.unit == 'libri_character':
diff --git a/openspeech/decoders/__init__.py b/openspeech/decoders/__init__.py
index 6bd92d9..41bf8e3 100644
--- a/openspeech/decoders/__init__.py
+++ b/openspeech/decoders/__init__.py
@@ -21,7 +21,7 @@
 # SOFTWARE.
 
 from .openspeech_decoder import OpenspeechDecoder
-from .lstm_decoder import LSTMDecoder
+from .lstm_attention_decoder import LSTMAttentionDecoder
 from .rnn_transducer_decoder import RNNTransducerDecoder
 from .transformer_decoder import TransformerDecoder
 from .transformer_transducer_decoder import TransformerTransducerDecoder
diff --git a/openspeech/decoders/lstm_decoder.py b/openspeech/decoders/lstm_attention_decoder.py
similarity index 99%
rename from openspeech/decoders/lstm_decoder.py
rename to openspeech/decoders/lstm_attention_decoder.py
index 6cc9fa7..40c7a78 100644
--- a/openspeech/decoders/lstm_decoder.py
+++ b/openspeech/decoders/lstm_attention_decoder.py
@@ -36,7 +36,7 @@
 )
 
 
-class LSTMDecoder(OpenspeechDecoder):
+class LSTMAttentionDecoder(OpenspeechDecoder):
     r"""
     Converts higher level features (from encoders) into output utterances
     by specifying a probability distribution over sequences of characters.
@@ -85,7 +85,7 @@ def __init__(
             rnn_type: str = 'lstm',
             dropout_p: float = 0.3,
     ) -> None:
-        super(LSTMDecoder, self).__init__()
+        super(LSTMAttentionDecoder, self).__init__()
         self.hidden_state_dim = hidden_state_dim
         self.num_classes = num_classes
         self.num_heads = num_heads
diff --git a/openspeech/lm/__init__.py b/openspeech/lm/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/openspeech/lm/lstm_for_causal_lm.py b/openspeech/lm/lstm_for_causal_lm.py
new file mode 100644
index 0000000..d527991
--- /dev/null
+++ b/openspeech/lm/lstm_for_causal_lm.py
@@ -0,0 +1,137 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+import torch.nn as nn
+import random
+
+from openspeech.lm.openspeech_for_causal_lm import OpenspeechCausalLMBase
+from openspeech.modules import Linear, View
+from typing import Optional, Tuple
+
+
+class LSTMForCausalLM(OpenspeechCausalLMBase):
+    supported_rnns = {
+        'lstm': nn.LSTM,
+        'gru': nn.GRU,
+        'rnn': nn.RNN,
+    }
+
+    def __init__(
+            self,
+            num_classes: int,
+            max_length: int = 128,
+            hidden_state_dim: int = 1024,
+            pad_id: int = 0,
+            sos_id: int = 1,
+            eos_id: int = 2,
+            num_layers: int = 2,
+            rnn_type: str = 'lstm',
+            dropout_p: float = 0.3,
+    ) -> None:
+        super(LSTMForCausalLM, self).__init__()
+        self.hidden_state_dim = hidden_state_dim
+        self.num_classes = num_classes
+        self.num_layers = num_layers
+        self.max_length = max_length
+        self.eos_id = eos_id
+        self.sos_id = sos_id
+        self.pad_id = pad_id
+        self.embedding = nn.Embedding(num_classes, hidden_state_dim)
+        self.input_dropout = nn.Dropout(dropout_p)
+        rnn_cell = self.supported_rnns[rnn_type.lower()]
+        self.rnn = rnn_cell(
+            input_size=hidden_state_dim,
+            hidden_size=hidden_state_dim,
+            num_layers=num_layers,
+            bias=True,
+            batch_first=True,
+            dropout=dropout_p,
+            bidirectional=False,
+        )
+
+        self.fc = nn.Sequential(
+            Linear(hidden_state_dim, hidden_state_dim),
+            nn.Tanh(),
+            View(shape=(-1, self.hidden_state_dim), contiguous=True),
+            Linear(hidden_state_dim, num_classes),
+        )
+
+    def forward_step(
+            self,
+            input_var: torch.Tensor,
+            hidden_states: Optional[torch.Tensor],
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        batch_size, output_lengths = input_var.size(0), input_var.size(1)
+
+        embedded = self.embedding(input_var)
+        embedded = self.input_dropout(embedded)
+
+        if self.training:
+            self.rnn.flatten_parameters()
+
+        outputs, hidden_states = self.rnn(embedded, hidden_states)
+
+        step_outputs = self.fc(outputs.reshape(-1, self.hidden_state_dim)).log_softmax(dim=-1)
+        step_outputs = step_outputs.view(batch_size, output_lengths, -1).squeeze(1)
+
+        return step_outputs, hidden_states
+
+    def forward(
+            self,
+            inputs: torch.Tensor,
+            teacher_forcing_ratio: float = 1.0,
+    ) -> torch.Tensor:
+        """
+        Forward propagate a `encoder_outputs` for training.
+
+        Args:
+            inputs (torch.LongTensr): A input sequence passed to decoders. `IntTensor` of size ``(batch, seq_length)``
+            teacher_forcing_ratio (float): ratio of teacher forcing
+
+        Returns:
+            * logits (torch.FloatTensor): Log probability of model predictions.
+        """
+        batch_size = inputs.size(0)
+        logits, hidden_states = list(), None
+        use_teacher_forcing = True if random.random() < teacher_forcing_ratio else False
+
+        if use_teacher_forcing:
+            inputs = inputs[inputs != self.eos_id].view(batch_size, -1)
+            step_outputs, hidden_states = self.forward_step(input_var=inputs, hidden_states=hidden_states)
+
+            for di in range(step_outputs.size(1)):
+                step_output = step_outputs[:, di, :]
+                logits.append(step_output)
+
+        else:
+            input_var = inputs[:, 0].unsqueeze(1)
+            for di in range(self.max_length):
+                step_output, hidden = self.forward_step(input_var=input_var, hidden_states=hidden_states)
+
+                step_output = step_output.squeeze(1)
+                logits.append(step_output)
+                input_var = logits[-1].topk(1)[1]
+
+        logits = torch.stack(logits, dim=1)
+
+        return logits
diff --git a/openspeech/lm/openspeech_for_causal_lm.py b/openspeech/lm/openspeech_for_causal_lm.py
new file mode 100644
index 0000000..93761dd
--- /dev/null
+++ b/openspeech/lm/openspeech_for_causal_lm.py
@@ -0,0 +1,45 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch.nn as nn
+
+
+class OpenspeechCausalLMBase(nn.Module):
+    r""" Interface of OpenSpeech decoder. """
+    def __init__(self):
+        super(OpenspeechCausalLMBase, self).__init__()
+
+    def count_parameters(self) -> int:
+        r""" Count parameters of decoders """
+        return sum([p.numel for p in self.parameters()])
+
+    def update_dropout(self, dropout_p: float) -> None:
+        r""" Update dropout probability of decoders """
+        for name, child in self.named_children():
+            if isinstance(child, nn.Dropout):
+                child.p = dropout_p
+
+    def forward_step(self, *args, **kwargs):
+        raise NotImplementedError
+
+    def forward(self, *args, **kwargs):
+        raise NotImplementedError
diff --git a/openspeech/lm/transformer_for_causal_lm.py b/openspeech/lm/transformer_for_causal_lm.py
new file mode 100644
index 0000000..e556bfa
--- /dev/null
+++ b/openspeech/lm/transformer_for_causal_lm.py
@@ -0,0 +1,137 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple
+
+from openspeech.lm.openspeech_for_causal_lm import OpenspeechCausalLMBase
+from openspeech.modules import (
+    TransformerEmbedding,
+    PositionalEncoding,
+    Linear,
+    PositionwiseFeedForward,
+    MultiHeadAttention,
+    get_attn_pad_mask,
+    get_attn_subsequent_mask,
+)
+
+
+class TransformerForCausalLMLayer(nn.Module):
+    def __init__(
+            self,
+            d_model: int = 768,
+            num_attention_heads: int = 8,
+            d_ff: int = 2048,
+            dropout_p: float = 0.3,
+    ) -> None:
+        super(TransformerForCausalLMLayer, self).__init__()
+        self.attention_prenorm = nn.LayerNorm(d_model)
+        self.attention = MultiHeadAttention(d_model, num_attention_heads)
+        self.feed_forward_prenorm = nn.LayerNorm(d_model)
+        self.feed_forward = PositionwiseFeedForward(d_model=d_model, d_ff=d_ff, dropout_p=dropout_p)
+
+    def forward(
+            self,
+            inputs: torch.Tensor,
+            mask: Optional[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        residual = inputs
+        inputs = self.attention_prenorm(inputs)
+        outputs, _ = self.attention(inputs, inputs, inputs, mask)
+        outputs += residual
+
+        residual = outputs
+        outputs = self.feed_forward_prenorm(outputs)
+        outputs = self.feed_forward(outputs)
+        outputs += residual
+
+        return outputs
+
+
+class TransformerForCausalLM(OpenspeechCausalLMBase):
+    def __init__(
+            self,
+            num_classes: int,
+            max_length: int = 128,
+            d_model: int = 768,
+            num_attention_heads: int = 8,
+            d_ff: int = 1536,
+            pad_id: int = 0,
+            sos_id: int = 1,
+            eos_id: int = 2,
+            num_layers: int = 2,
+            dropout_p: float = 0.3,
+    ):
+        super(TransformerForCausalLM, self).__init__()
+        self.d_model = d_model
+        self.num_classes = num_classes
+        self.num_layers = num_layers
+        self.max_length = max_length
+        self.eos_id = eos_id
+        self.sos_id = sos_id
+        self.pad_id = pad_id
+        self.embedding = TransformerEmbedding(num_classes, pad_id, d_model)
+        self.positional_encoding = PositionalEncoding(d_model)
+        self.input_dropout = nn.Dropout(p=dropout_p)
+        self.layers = nn.ModuleList([
+            TransformerForCausalLMLayer(
+                d_model=d_model,
+                num_attention_heads=num_attention_heads,
+                d_ff=d_ff,
+                dropout_p=dropout_p,
+            ) for _ in range(num_layers)
+        ])
+        self.fc = nn.Sequential(
+            nn.LayerNorm(d_model),
+            Linear(d_model, d_model, bias=False),
+            nn.Tanh(),
+            Linear(d_model, num_classes, bias=False),
+        )
+
+    def forward_step(self, inputs, input_lengths):
+        pad_mask = get_attn_pad_mask(
+            inputs, input_lengths, inputs.size(1)
+        )
+        subsequent_mask = get_attn_subsequent_mask(inputs)
+        mask = torch.gt((pad_mask + subsequent_mask), 0)
+
+        outputs = self.embedding(inputs) + self.positional_encoding(inputs.size(1))
+        outputs = self.input_dropout(outputs)
+
+        for layer in self.layers:
+            outputs = layer(inputs=outputs, mask=mask)
+
+        step_outputs = self.fc(outputs).log_softmax(dim=-1)
+
+        return step_outputs
+
+    def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor) -> torch.Tensor:
+        logits = list()
+
+        step_outputs = self.forward_step(inputs, input_lengths)
+
+        for di in range(step_outputs.size(1)):
+            step_output = step_outputs[:, di, :]
+            logits.append(step_output)
+
+        return torch.stack(logits, dim=1)
diff --git a/openspeech/models/README.md b/openspeech/models/README.md
index 9b21646..b31007d 100644
--- a/openspeech/models/README.md
+++ b/openspeech/models/README.md
@@ -2,23 +2,25 @@
   
 We support all the models below. Note that, the important concepts of the model have been implemented to match, but the details of the implementation may vary.  
   
-1. [**DeepSpeech2**]() (from Baidu Research) released with paper [Deep Speech 2: End-to-End Speech Recognition in
+1. [**DeepSpeech2**](https://sooftware.github.io/OpenSpeech/DeepSpeech2.html) (from Baidu Research) released with paper [Deep Speech 2: End-to-End Speech Recognition in
 English and Mandarin](https://arxiv.org/abs/1512.02595.pdf), by Dario Amodei, Rishita Anubhai, Eric Battenberg, Carl Case, Jared Casper, Bryan Catanzaro, Jingdong Chen, Mike Chrzanowski, Adam Coates, Greg Diamos, Erich Elsen, Jesse Engel, Linxi Fan, Christopher Fougner, Tony Han, Awni Hannun, Billy Jun, Patrick LeGresley, Libby Lin, Sharan Narang, Andrew Ng, Sherjil Ozair, Ryan Prenger, Jonathan Raiman, Sanjeev Satheesh, David Seetapun, Shubho Sengupta, Yi Wang, Zhiqian Wang, Chong Wang, Bo Xiao, Dani Yogatama, Jun Zhan, Zhenyao Zhu. 
-2. [**RNN-Transducer**]() (from University of Toronto) released with paper [Sequence Transduction with Recurrent Neural Networks](https://arxiv.org/abs/1211.3711.pdf), by Alex Graves.
-3. [**Listen Attend Spell**](https://sooftware.github.io/openspeech/Listen%20Attend%20Spell.html) (from Carnegie Mellon University and Google Brain) released with paper [Listen, Attend and Spell](https://arxiv.org/abs/1508.01211), by William Chan, Navdeep Jaitly, Quoc V. Le, Oriol Vinyals.  
-4. [**Location-aware attention based Listen Attend Spell**]() (from University of Wrocław and Jacobs University and Universite de Montreal ) released with paper [Attention-Based Models for Speech Recognition](https://arxiv.org/abs/1506.07503), by Jan Chorowski, Dzmitry Bahdanau, Dmitriy Serdyuk, Kyunghyun Cho, Yoshua Bengio.  
-5. [**Joint CTC-Attention based Listen Attend Spell**]() (from Mitsubishi Electric Research Laboratories and Carnegie Mellon University) released with paper [Joint CTC-Attention based End-to-End Speech Recognition using Multi-task Learning](https://arxiv.org/abs/1609.06773), by Suyoun Kim, Takaaki Hori, Shinji Watanabe.  
-6. [**Deep CNN Encoder with Joint CTC-Attention Listen Attend Spell**]() (from Mitsubishi Electric Research Laboratories and Massachusetts Institute of Technology and Carnegie Mellon University) released with paper [Advances in Joint CTC-Attention based End-to-End Speech Recognition with a Deep CNN Encoder and RNN-LM](https://arxiv.org/abs/1706.02737), by Takaaki Hori, Shinji Watanabe, Yu Zhang, William Chan.
-7. [**Multi-head attention based Listen Attend Spell**]() (from Google) released with paper [State-of-the-art Speech Recognition With Sequence-to-Sequence Models](https://arxiv.org/abs/1712.01769), by Chung-Cheng Chiu, Tara N. Sainath, Yonghui Wu, Rohit Prabhavalkar, Patrick Nguyen, Zhifeng Chen, Anjuli Kannan, Ron J. Weiss, Kanishka Rao, Ekaterina Gonina, Navdeep Jaitly, Bo Li, Jan Chorowski, Michiel Bacchiani.  
-8. [**Speech-Transformer**]() (from University of Chinese Academy of Sciences and Institute of Automation and Chinese Academy of Sciences) released with paper [Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for Speech Recognition](https://ieeexplore.ieee.org/document/8462506), by Linhao Dong; Shuang Xu; Bo Xu.
-9. [**VGG-Transformer**]() (from Facebook AI Research) released with paper [Transformers with convolutional context for ASR](https://arxiv.org/abs/1904.11660), by Abdelrahman Mohamed, Dmytro Okhonko, Luke Zettlemoyer.  
-10. [**Transformer with CTC**](https://github.com/sooftware/openspeech-dev/issues/151) (from NTT Communication Science Laboratories, Waseda University, Center for Language and Speech Processing, Johns Hopkins University) released with paper [Improving Transformer-based End-to-End Speech Recognition with Connectionist Temporal Classification and Language Model Integration](https://www.isca-speech.org/archive/Interspeech_2019/pdfs/1938.pdf), by Shigeki Karita, Nelson Enrique Yalta Soplin, Shinji Watanabe, Marc Delcroix, Atsunori Ogawa, Tomohiro Nakatani.
-11. [**Joint CTC-Attention based Transformer**](https://github.com/sooftware/openspeech-dev/issues/151)(from NTT Corporation) released with paper [Self-Distillation for Improving CTC-Transformer-based ASR Systems](https://www.isca-speech.org/archive/Interspeech_2020/pdfs/1223.pdf), by Takafumi Moriya, Tsubasa Ochiai, Shigeki Karita, Hiroshi Sato, Tomohiro Tanaka, Takanori Ashihara, Ryo Masumura, Yusuke Shinohara, Marc Delcroix.
-12. [**Jasper**]() (from NVIDIA and New York University) released with paper [Jasper: An End-to-End Convolutional Neural Acoustic Model](https://arxiv.org/pdf/1904.03288.pdf), by Jason Li, Vitaly Lavrukhin, Boris Ginsburg, Ryan Leary, Oleksii Kuchaiev, Jonathan M. Cohen, Huyen Nguyen, Ravi Teja Gadde.   
-13. [**QuartzNet**]() (from NVIDIA and Univ. of Illinois and Univ. of Saint Petersburg) released with paper [QuartzNet: Deep Automatic Speech Recognition with 1D Time-Channel Separable Convolutions](https://arxiv.org/abs/1910.10261.pdf), by Samuel Kriman, Stanislav Beliaev, Boris Ginsburg, Jocelyn Huang, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Yang Zhang.  
-14. [**Transformer Transducer**]() (from Facebook AI) released with paper [Transformer-Transducer:
+2. [**RNN-Transducer**](https://sooftware.github.io/OpenSpeech/RNN%20Transducer.html) (from University of Toronto) released with paper [Sequence Transduction with Recurrent Neural Networks](https://arxiv.org/abs/1211.3711.pdf), by Alex Graves.
+3. [**LSTM Language Model**]() (from RWTH Aachen University) released with paper [LSTM Neural Networks for Language Modeling](http://www-i6.informatik.rwth-aachen.de/publications/download/820/Sundermeyer-2012.pdf), by  Martin Sundermeyer, Ralf Schluter, and Hermann Ney.  
+3. [**Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Listen%20Attend%20Spell.html) (from Carnegie Mellon University and Google Brain) released with paper [Listen, Attend and Spell](https://arxiv.org/abs/1508.01211), by William Chan, Navdeep Jaitly, Quoc V. Le, Oriol Vinyals.  
+4. [**Location-aware attention based Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Listen%20Attend%20Spell%20(location-aware).html) (from University of Wrocław and Jacobs University and Universite de Montreal) released with paper [Attention-Based Models for Speech Recognition](https://arxiv.org/abs/1506.07503), by Jan Chorowski, Dzmitry Bahdanau, Dmitriy Serdyuk, Kyunghyun Cho, Yoshua Bengio.  
+5. [**Joint CTC-Attention based Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Joint%20CTC%20LAS.html) (from Mitsubishi Electric Research Laboratories and Carnegie Mellon University) released with paper [Joint CTC-Attention based End-to-End Speech Recognition using Multi-task Learning](https://arxiv.org/abs/1609.06773), by Suyoun Kim, Takaaki Hori, Shinji Watanabe.  
+6. [**Deep CNN Encoder with Joint CTC-Attention Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Deep%20CNN%20with%20Joint%20CTC%20LAS.html) (from Mitsubishi Electric Research Laboratories and Massachusetts Institute of Technology and Carnegie Mellon University) released with paper [Advances in Joint CTC-Attention based End-to-End Speech Recognition with a Deep CNN Encoder and RNN-LM](https://arxiv.org/abs/1706.02737), by Takaaki Hori, Shinji Watanabe, Yu Zhang, William Chan.
+7. [**Multi-head attention based Listen Attend Spell**](https://sooftware.github.io/OpenSpeech/Listen%20Attend%20Spell%20(multi-head).html) (from Google) released with paper [State-of-the-art Speech Recognition With Sequence-to-Sequence Models](https://arxiv.org/abs/1712.01769), by Chung-Cheng Chiu, Tara N. Sainath, Yonghui Wu, Rohit Prabhavalkar, Patrick Nguyen, Zhifeng Chen, Anjuli Kannan, Ron J. Weiss, Kanishka Rao, Ekaterina Gonina, Navdeep Jaitly, Bo Li, Jan Chorowski, Michiel Bacchiani.  
+8. [**Speech-Transformer**](https://sooftware.github.io/OpenSpeech/Transformer.html) (from University of Chinese Academy of Sciences and Institute of Automation and Chinese Academy of Sciences) released with paper [Speech-Transformer: A No-Recurrence Sequence-to-Sequence Model for Speech Recognition](https://ieeexplore.ieee.org/document/8462506), by Linhao Dong; Shuang Xu; Bo Xu.
+9. [**VGG-Transformer**](https://sooftware.github.io/OpenSpeech/VGG%20Transformer.html) (from Facebook AI Research) released with paper [Transformers with convolutional context for ASR](https://arxiv.org/abs/1904.11660), by Abdelrahman Mohamed, Dmytro Okhonko, Luke Zettlemoyer.  
+10. [**Transformer with CTC**](https://sooftware.github.io/OpenSpeech/Transformer%20with%20CTC.html) (from NTT Communication Science Laboratories, Waseda University, Center for Language and Speech Processing, Johns Hopkins University) released with paper [Improving Transformer-based End-to-End Speech Recognition with Connectionist Temporal Classification and Language Model Integration](https://www.isca-speech.org/archive/Interspeech_2019/pdfs/1938.pdf), by Shigeki Karita, Nelson Enrique Yalta Soplin, Shinji Watanabe, Marc Delcroix, Atsunori Ogawa, Tomohiro Nakatani.
+11. [**Joint CTC-Attention based Transformer**](https://sooftware.github.io/OpenSpeech/Joint%20CTC%20Transformer.html)(from NTT Corporation) released with paper [Self-Distillation for Improving CTC-Transformer-based ASR Systems](https://www.isca-speech.org/archive/Interspeech_2020/pdfs/1223.pdf), by Takafumi Moriya, Tsubasa Ochiai, Shigeki Karita, Hiroshi Sato, Tomohiro Tanaka, Takanori Ashihara, Ryo Masumura, Yusuke Shinohara, Marc Delcroix.
+12. [**Transformer Language Model**]() (from Amazon Web Services) released with paper [Language Models with Transformers](https://arxiv.org/abs/1904.09408), by Chenguang Wang, Mu Li, Alexander J. Smola.
+12. [**Jasper**](https://sooftware.github.io/OpenSpeech/Jasper10x5.html) (from NVIDIA and New York University) released with paper [Jasper: An End-to-End Convolutional Neural Acoustic Model](https://arxiv.org/pdf/1904.03288.pdf), by Jason Li, Vitaly Lavrukhin, Boris Ginsburg, Ryan Leary, Oleksii Kuchaiev, Jonathan M. Cohen, Huyen Nguyen, Ravi Teja Gadde.   
+13. [**QuartzNet**](https://sooftware.github.io/OpenSpeech/QuartzNet15x5.html) (from NVIDIA and Univ. of Illinois and Univ. of Saint Petersburg) released with paper [QuartzNet: Deep Automatic Speech Recognition with 1D Time-Channel Separable Convolutions](https://arxiv.org/abs/1910.10261.pdf), by Samuel Kriman, Stanislav Beliaev, Boris Ginsburg, Jocelyn Huang, Oleksii Kuchaiev, Vitaly Lavrukhin, Ryan Leary, Jason Li, Yang Zhang.  
+14. [**Transformer Transducer**](https://sooftware.github.io/OpenSpeech/Transformer%20Transducer.html) (from Facebook AI) released with paper [Transformer-Transducer:
 End-to-End Speech Recognition with Self-Attention](https://arxiv.org/abs/1910.12977.pdf), by Ching-Feng Yeh, Jay Mahadeokar, Kaustubh Kalgaonkar, Yongqiang Wang, Duc Le, Mahaveer Jain, Kjell Schubert, Christian Fuegen, Michael L. Seltzer.  
-15. [**Conformer**]() (from Google) released with paper [Conformer: Convolution-augmented Transformer for Speech Recognition](https://arxiv.org/abs/2005.08100), by Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, Ruoming Pang.  
-16. [**Conformer with CTC**]() (from Northwestern Polytechnical University and University of Bordeaux and Johns Hopkins University and Human Dataware Lab and Kyoto University and NTT Corporation and Shanghai Jiao Tong University and  Chinese Academy of Sciences) released with paper [RECENT DEVELOPMENTS ON ESPNET TOOLKIT BOOSTED BY CONFORMER](https://arxiv.org/abs/2010.13956.pdf), by Pengcheng Guo, Florian Boyer, Xuankai Chang, Tomoki Hayashi, Yosuke Higuchi, Hirofumi Inaguma, Naoyuki Kamo, Chenda Li, Daniel Garcia-Romero, Jiatong Shi, Jing Shi, Shinji Watanabe, Kun Wei, Wangyou Zhang, Yuekai Zhang.
-17. [**Conformer with LSTM Decoder**]() (from IBM Research AI) released with paper [On the limit of English conversational speech recognition](https://arxiv.org/abs/2105.00982.pdf), by Zoltán Tüske, George Saon, Brian Kingsbury.
-  
+15. [**Conformer**](https://sooftware.github.io/OpenSpeech/Conformer%20Transducer.html) (from Google) released with paper [Conformer: Convolution-augmented Transformer for Speech Recognition](https://arxiv.org/abs/2005.08100), by Anmol Gulati, James Qin, Chung-Cheng Chiu, Niki Parmar, Yu Zhang, Jiahui Yu, Wei Han, Shibo Wang, Zhengdong Zhang, Yonghui Wu, Ruoming Pang.  
+16. [**Conformer with CTC**](https://sooftware.github.io/OpenSpeech/Conformer.html) (from Northwestern Polytechnical University and University of Bordeaux and Johns Hopkins University and Human Dataware Lab and Kyoto University and NTT Corporation and Shanghai Jiao Tong University and  Chinese Academy of Sciences) released with paper [Recent Developments on ESPNET Toolkit Boosted by Conformer](https://arxiv.org/abs/2010.13956.pdf), by Pengcheng Guo, Florian Boyer, Xuankai Chang, Tomoki Hayashi, Yosuke Higuchi, Hirofumi Inaguma, Naoyuki Kamo, Chenda Li, Daniel Garcia-Romero, Jiatong Shi, Jing Shi, Shinji Watanabe, Kun Wei, Wangyou Zhang, Yuekai Zhang.
+17. [**Conformer with LSTM Decoder**](https://sooftware.github.io/OpenSpeech/Conformer%20LSTM.html) (from IBM Research AI) released with paper [On the limit of English conversational speech recognition](https://arxiv.org/abs/2105.00982.pdf), by Zoltán Tüske, George Saon, Brian Kingsbury.
+  
\ No newline at end of file
diff --git a/openspeech/models/__init__.py b/openspeech/models/__init__.py
index 849b583..9dabd8c 100644
--- a/openspeech/models/__init__.py
+++ b/openspeech/models/__init__.py
@@ -125,6 +125,8 @@ def register_model_cls(cls):
 from .transformer_with_ctc.configurations import TransformerWithCTCConfigs
 from .vgg_transformer.model import VGGTransformerModel
 from .vgg_transformer.configurations import VGGTransformerConfigs
+from .lstm_lm.model import LSTMLanguageModel
+from .lstm_lm.configurations import LSTMLanguageModelConfigs
 
 __all__ = [
     "OpenspeechModel",
@@ -172,4 +174,6 @@ def register_model_cls(cls):
     "TransformerTransducerConfigs",
     "TransformerTransducerModel",
     "TransformerWithCTCModel",
+    "LSTMLanguageModel",
+    "LSTMLanguageModelConfigs",
 ]
diff --git a/openspeech/models/conformer_lstm/model.py b/openspeech/models/conformer_lstm/model.py
index a66bdd1..753a6d6 100644
--- a/openspeech/models/conformer_lstm/model.py
+++ b/openspeech/models/conformer_lstm/model.py
@@ -29,7 +29,7 @@
 from openspeech.models.conformer_lstm.configurations import ConformerLSTMConfigs
 from openspeech.encoders import ConformerEncoder
 from openspeech.models import OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.vocabs.vocab import Vocabulary
 
 
@@ -71,7 +71,7 @@ def build_model(self):
             half_step_residual=self.configs.model.half_step_residual,
             joint_ctc_attention=False,
         )
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=self.configs.model.encoder_dim,
diff --git a/openspeech/models/deep_cnn_with_joint_ctc_listen_attend_spell/model.py b/openspeech/models/deep_cnn_with_joint_ctc_listen_attend_spell/model.py
index eb5776b..deaff9b 100644
--- a/openspeech/models/deep_cnn_with_joint_ctc_listen_attend_spell/model.py
+++ b/openspeech/models/deep_cnn_with_joint_ctc_listen_attend_spell/model.py
@@ -26,7 +26,7 @@
 from collections import OrderedDict
 
 from openspeech.models import register_model, OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.encoders import ConvolutionalLSTMEncoder
 from openspeech.models.deep_cnn_with_joint_ctc_listen_attend_spell.configurations import \
     DeepCNNWithJointCTCListenAttendSpellConfigs
@@ -69,7 +69,7 @@ def build_model(self):
         decoder_hidden_state_dim = self.configs.model.hidden_state_dim << 1 \
             if self.configs.model.encoder_bidirectional \
             else self.configs.model.hidden_state_dim
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=decoder_hidden_state_dim,
diff --git a/openspeech/models/joint_ctc_conformer_lstm/model.py b/openspeech/models/joint_ctc_conformer_lstm/model.py
index 4acb760..65ec3a4 100644
--- a/openspeech/models/joint_ctc_conformer_lstm/model.py
+++ b/openspeech/models/joint_ctc_conformer_lstm/model.py
@@ -28,7 +28,7 @@
 from openspeech.models import register_model
 from openspeech.encoders import ConformerEncoder
 from openspeech.models import OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.models.joint_ctc_conformer_lstm.configurations import JointCTCConformerLSTMConfigs
 from openspeech.vocabs.vocab import Vocabulary
 
@@ -71,7 +71,7 @@ def build_model(self):
             half_step_residual=self.configs.model.half_step_residual,
             joint_ctc_attention=True,
         )
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=self.configs.model.encoder_dim,
diff --git a/openspeech/models/joint_ctc_listen_attend_spell/model.py b/openspeech/models/joint_ctc_listen_attend_spell/model.py
index 53764b3..2ad3f87 100644
--- a/openspeech/models/joint_ctc_listen_attend_spell/model.py
+++ b/openspeech/models/joint_ctc_listen_attend_spell/model.py
@@ -26,7 +26,7 @@
 from collections import OrderedDict
 
 from openspeech.models import register_model, OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.encoders import LSTMEncoder
 from openspeech.models.joint_ctc_listen_attend_spell.configurations import JointCTCListenAttendSpellConfigs
 from openspeech.vocabs.vocab import Vocabulary
@@ -68,7 +68,7 @@ def build_model(self):
         decoder_hidden_state_dim = self.configs.model.hidden_state_dim << 1 \
             if self.configs.model.encoder_bidirectional \
             else self.configs.model.hidden_state_dim
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=decoder_hidden_state_dim,
diff --git a/openspeech/models/listen_attend_spell/model.py b/openspeech/models/listen_attend_spell/model.py
index f835d3f..0bb05fe 100644
--- a/openspeech/models/listen_attend_spell/model.py
+++ b/openspeech/models/listen_attend_spell/model.py
@@ -26,7 +26,7 @@
 from collections import OrderedDict
 
 from openspeech.models import register_model, OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.encoders import LSTMEncoder
 from openspeech.models.listen_attend_spell.configurations import ListenAttendSpellConfigs
 from openspeech.vocabs.vocab import Vocabulary
@@ -68,7 +68,7 @@ def build_model(self):
         decoder_hidden_state_dim = self.configs.model.hidden_state_dim << 1 \
             if self.configs.model.encoder_bidirectional \
             else self.configs.model.hidden_state_dim
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=decoder_hidden_state_dim,
diff --git a/openspeech/models/listen_attend_spell_with_location_aware/model.py b/openspeech/models/listen_attend_spell_with_location_aware/model.py
index cfed86c..d079b30 100644
--- a/openspeech/models/listen_attend_spell_with_location_aware/model.py
+++ b/openspeech/models/listen_attend_spell_with_location_aware/model.py
@@ -26,7 +26,7 @@
 from collections import OrderedDict
 
 from openspeech.models import register_model, OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.encoders import LSTMEncoder
 from openspeech.models.listen_attend_spell_with_location_aware.configurations import ListenAttendSpellWithLocationAwareConfigs
 from openspeech.vocabs.vocab import Vocabulary
@@ -68,7 +68,7 @@ def build_model(self):
         decoder_hidden_state_dim = self.configs.model.hidden_state_dim << 1 \
             if self.configs.model.encoder_bidirectional \
             else self.configs.model.hidden_state_dim
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=decoder_hidden_state_dim,
diff --git a/openspeech/models/listen_attend_spell_with_multi_head/model.py b/openspeech/models/listen_attend_spell_with_multi_head/model.py
index de1cca1..73c71e9 100644
--- a/openspeech/models/listen_attend_spell_with_multi_head/model.py
+++ b/openspeech/models/listen_attend_spell_with_multi_head/model.py
@@ -26,7 +26,7 @@
 from collections import OrderedDict
 
 from openspeech.models import register_model, OpenspeechEncoderDecoderModel
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 from openspeech.encoders import LSTMEncoder
 from openspeech.models.listen_attend_spell_with_multi_head.configurations import ListenAttendSpellWithMultiHeadConfigs
 from openspeech.vocabs.vocab import Vocabulary
@@ -68,7 +68,7 @@ def build_model(self):
         decoder_hidden_state_dim = self.configs.model.hidden_state_dim << 1 \
             if self.configs.model.encoder_bidirectional \
             else self.configs.model.hidden_state_dim
-        self.decoder = LSTMDecoder(
+        self.decoder = LSTMAttentionDecoder(
             num_classes=self.num_classes,
             max_length=self.configs.model.max_length,
             hidden_state_dim=decoder_hidden_state_dim,
diff --git a/openspeech/models/lstm_lm/configurations.py b/openspeech/models/lstm_lm/configurations.py
new file mode 100644
index 0000000..56cc385
--- /dev/null
+++ b/openspeech/models/lstm_lm/configurations.py
@@ -0,0 +1,71 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from dataclasses import dataclass, field
+
+from openspeech.dataclass.configurations import OpenspeechDataclass
+
+
+@dataclass
+class LSTMLanguageModelConfigs(OpenspeechDataclass):
+    r"""
+    This is the configuration class to store the configuration of
+    a :class:`~openspeech.models.ListenAttendSpellWithMultiHead`.
+
+    It is used to initiated an `ListenAttendSpellWithMultiHead` model.
+
+    Configuration objects inherit from :class: `~openspeech.dataclass.configs.OpenspeechDataclass`.
+
+    Args:
+        model_name (str): Model name (default: lstm_lm)
+        num_layers (int): The number of lstm layers. (default: 3)
+        hidden_state_dim (int): The hidden state dimension of model. (default: 512)
+        dropout_p (float): The dropout probability of encoder. (default: 0.3)
+        rnn_type (str): Type of rnn cell (rnn, lstm, gru) (default: lstm)
+        max_length (int): Max decoding length. (default: 128)
+        teacher_forcing_ratio (float): The ratio of teacher forcing. (default: 1.0)
+        optimizer (str): Optimizer for training. (default: adam)
+    """
+    model_name: str = field(
+        default="lstm_lm", metadata={"help": "Model name"}
+    )
+    num_layers: int = field(
+        default=3, metadata={"help": "The number of encoder layers."}
+    )
+    hidden_state_dim: int = field(
+        default=512, metadata={"help": "The hidden state dimension of encoder."}
+    )
+    dropout_p: float = field(
+        default=0.3, metadata={"help": "The dropout probability of encoder."}
+    )
+    rnn_type: str = field(
+        default="lstm", metadata={"help": "Type of rnn cell (rnn, lstm, gru)"}
+    )
+    max_length: int = field(
+        default=128, metadata={"help": "Max decoding length."}
+    )
+    teacher_forcing_ratio: float = field(
+        default=1.0, metadata={"help": "The ratio of teacher forcing. "}
+    )
+    optimizer: str = field(
+        default="adam", metadata={"help": "Optimizer for training."}
+    )
diff --git a/openspeech/models/lstm_lm/model.py b/openspeech/models/lstm_lm/model.py
new file mode 100644
index 0000000..2fd6e17
--- /dev/null
+++ b/openspeech/models/lstm_lm/model.py
@@ -0,0 +1,136 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+from omegaconf import DictConfig
+from collections import OrderedDict
+from typing import Dict
+
+from openspeech.lm.lstm_for_causal_lm import LSTMForCausalLM
+from openspeech.models import register_model, OpenspeechModel
+from openspeech.models.lstm_lm.configurations import LSTMLanguageModelConfigs
+from openspeech.vocabs.vocab import Vocabulary
+
+
+@register_model('lstm_lm', dataclass=LSTMLanguageModelConfigs)
+class LSTMLanguageModel(OpenspeechModel):
+    def __init__(self, configs: DictConfig, vocab: Vocabulary, ) -> None:
+        super(LSTMLanguageModel, self).__init__(configs, vocab)
+
+    def build_model(self):
+        self.lm = LSTMForCausalLM(
+            num_classes=self.num_classes,
+            max_length=self.configs.model.max_length,
+            hidden_state_dim=self.configs.model.hidden_state_dim,
+            pad_id=self.vocab.pad_id,
+            sos_id=self.vocab.sos_id,
+            eos_id=self.vocab.eos_id,
+            dropout_p=self.configs.model.dropout_p,
+            num_layers=self.configs.model.num_layers,
+            rnn_type=self.configs.model.rnn_type,
+        )
+
+    def collect_outputs(
+            self,
+            stage: str,
+            logits: torch.Tensor,
+            targets: torch.Tensor,
+    ) -> OrderedDict:
+        perplexity = self.criterion(logits, targets[:, 1:])
+        predictions = logits.max(-1)[1]
+
+        self.info({
+            f"{stage}_perplexity": perplexity,
+            "learning_rate": self.get_lr(),
+        })
+
+        return OrderedDict({
+            "perplexity": perplexity,
+            "logits": logits,
+            "targets": targets,
+            "predictions": predictions,
+        })
+
+    def forward(self, inputs: torch.Tensor) -> Dict[str, torch.Tensor]:
+        logits = self.lm(inputs, teacher_forcing_ratio=0.0)
+        predictions = logits.max(-1)[1]
+        return {
+            "predictions": predictions,
+            "logits": logits,
+        }
+
+    def training_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
+        r"""
+        Forward propagate a `inputs` and `targets` pair for training.
+
+        Inputs:
+            train_batch (tuple): A train batch contains `inputs`
+            batch_idx (int): The index of batch
+
+        Returns:
+            loss (torch.Tensor): loss for training
+        """
+        inputs, input_lengths, targets = batch
+        logits = self.lm(inputs, teacher_forcing_ratio=self.teacher_forcing_ratio)
+        return self.collect_outputs(
+            stage='train',
+            logits=logits,
+            targets=targets,
+        )
+
+    def validation_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
+        r"""
+        Forward propagate a `inputs` and `targets` pair for validation.
+
+        Inputs:
+            train_batch (tuple): A train batch contains `inputs`
+            batch_idx (int): The index of batch
+
+        Returns:
+            loss (torch.Tensor): loss for training
+        """
+        inputs, input_lengths, targets = batch
+        logits = self.lm(inputs, teacher_forcing_ratio=0.0)
+        return self.collect_outputs(
+            stage='val',
+            logits=logits,
+            targets=targets,
+        )
+
+    def test_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
+        r"""
+        Forward propagate a `inputs` and `targets` pair for test.
+
+        Inputs:
+            train_batch (tuple): A train batch contains `inputs`
+            batch_idx (int): The index of batch
+
+        Returns:
+            loss (torch.Tensor): loss for training
+        """
+        inputs, input_lengths, targets = batch
+        logits = self.lm(inputs, teacher_forcing_ratio=0.0)
+        return self.collect_outputs(
+            stage='test',
+            logits=logits,
+            targets=targets,
+        )
diff --git a/openspeech/models/openspeech_ctc_model.py b/openspeech/models/openspeech_ctc_model.py
index 955c329..165d323 100644
--- a/openspeech/models/openspeech_ctc_model.py
+++ b/openspeech/models/openspeech_ctc_model.py
@@ -78,18 +78,17 @@ def collect_outputs(
         wer = self.wer_metric(targets[:, 1:], predictions)
         cer = self.cer_metric(targets[:, 1:], predictions)
 
-        self.log_steps(stage, wer, cer, loss)
-
-        progress_bar_dict = {
+        self.info({
+            f"{stage}_wer": wer,
+            f"{stage}_cer": cer,
             f"{stage}_loss": loss,
-            "wer": wer,
-            "cer": cer,
-        }
+            "learning_rate": self.get_lr(),
+        })
 
         return OrderedDict({
             "loss": loss,
-            "progress_bar": progress_bar_dict,
-            "log": progress_bar_dict
+            "wer": wer,
+            "cer": cer,
         })
 
     def forward(self, inputs: torch.FloatTensor, input_lengths: torch.IntTensor) -> Dict[str, torch.Tensor]:
@@ -156,7 +155,7 @@ def validation_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
         inputs, targets, input_lengths, target_lengths = batch
         logits, output_lengths = self.encoder(inputs, input_lengths)
         return self.collect_outputs(
-            stage='valid',
+            stage='val',
             logits=logits,
             output_lengths=output_lengths,
             targets=targets,
diff --git a/openspeech/models/openspeech_encoder_decoder_model.py b/openspeech/models/openspeech_encoder_decoder_model.py
index 9592c65..de45acc 100644
--- a/openspeech/models/openspeech_encoder_decoder_model.py
+++ b/openspeech/models/openspeech_encoder_decoder_model.py
@@ -82,23 +82,27 @@ def collect_outputs(
         else:
             raise ValueError(f"Unsupported criterion: {self.criterion}")
 
-        y_hats = logits.max(-1)[1]
+        predictions = logits.max(-1)[1]
 
-        wer = self.wer_metric(targets[:, 1:], y_hats)
-        cer = self.cer_metric(targets[:, 1:], y_hats)
+        wer = self.wer_metric(targets[:, 1:], predictions)
+        cer = self.cer_metric(targets[:, 1:], predictions)
 
-        self.log_steps(stage, wer, cer, loss, cross_entropy_loss, ctc_loss)
-
-        progress_bar_dict = {
+        self.info({
             f"{stage}_loss": loss,
-            "wer": wer,
-            "cer": cer,
-        }
+            f"{stage}_cross_entropy_loss": cross_entropy_loss,
+            f"{stage}_ctc_loss": ctc_loss,
+            f"{stage}_wer": wer,
+            f"{stage}_cer": cer,
+        })
 
         return OrderedDict({
             "loss": loss,
-            "progress_bar": progress_bar_dict,
-            "log": progress_bar_dict,
+            "cross_entropy_loss": cross_entropy_loss,
+            "ctc_loss": ctc_loss,
+            "predictions": predictions,
+            "targets": targets,
+            "logits": logits,
+            "learning_rate": self.get_lr(),
         })
 
     def forward(self, inputs: Tensor, input_lengths: Tensor) -> Dict[str, Tensor]:
@@ -193,7 +197,7 @@ def validation_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
             teacher_forcing_ratio=0.0,
         )
         return self.collect_outputs(
-            stage='valid',
+            stage='val',
             logits=logits,
             encoder_logits=encoder_logits,
             encoder_output_lengths=encoder_output_lengths,
diff --git a/openspeech/models/openspeech_model.py b/openspeech/models/openspeech_model.py
index a2bc2df..b63dd53 100644
--- a/openspeech/models/openspeech_model.py
+++ b/openspeech/models/openspeech_model.py
@@ -23,14 +23,15 @@
 import torch
 import torch.nn as nn
 import pytorch_lightning as pl
-from typing import Dict, Optional, List, Tuple
+from typing import Dict
 from omegaconf import DictConfig
 from torch import Tensor
+from torch.optim import Adam, Adagrad, Adadelta, Adamax, AdamW, SGD, ASGD
 
+from openspeech.optim import AdamP, RAdam, Novograd
 from openspeech.criterion import CRITERION_REGISTRY
 from openspeech.metrics import WordErrorRate, CharacterErrorRate
 from openspeech.optim.scheduler import SCHEDULER_REGISTRY
-from openspeech.utils import get_class_name
 from openspeech.vocabs.vocab import Vocabulary
 
 
@@ -70,36 +71,15 @@ def build_model(self):
     def set_beam_decoder(self, beam_size: int = 3):
         raise NotImplementedError
 
-    def log_steps(
-            self,
-            stage: str,
-            wer: float,
-            cer: float,
-            loss: Optional[float] = None,
-            cross_entropy_loss: Optional[float] = None,
-            ctc_loss: Optional[float] = None,
-    ) -> None:
+    def info(self, dictionary: dict) -> None:
         r"""
-        Provides log dictionary.
+        Logging information from dictionary.
 
         Args:
-            stage (str): current stage (train, valid, test)
-            wer (float): word error rate
-            cer (float): character error rate
-            loss (float): loss of model's prediction
-            cross_entropy_loss (Optional, float): cross entropy loss of model's prediction
-            ctc_loss (Optional, float): ctc loss of model's prediction
+            dictionary (dict): dictionary contains information.
         """
-        self.log(f"{stage}_wer", wer)
-        self.log(f"{stage}_cer", cer)
-        if loss is not None:
-            self.log(f"{stage}_loss", loss)
-        if cross_entropy_loss is not None:
-            self.log(f"{stage}_cross_entropy_loss", cross_entropy_loss)
-        if ctc_loss is not None:
-            self.log(f"{stage}_ctc_loss", ctc_loss)
-        if hasattr(self, "optimizer"):
-            self.log("current_lr", self.get_lr())
+        for key, value in dictionary.items():
+            self.log(key, value, prog_bar=True)
 
     def forward(self, inputs: torch.FloatTensor, input_lengths: torch.LongTensor) -> Dict[str, Tensor]:
         r"""
@@ -141,18 +121,6 @@ def validation_step(self, batch: tuple, batch_idx: int):
         """
         raise NotImplementedError
 
-    def validation_epoch_end(self, outputs: dict) -> dict:
-        self.current_val_loss = torch.stack([output['loss'] for output in outputs]).mean()
-
-        if get_class_name(self.scheduler) == "WarmupReduceLROnPlateauScheduler" \
-            or get_class_name(self.scheduler) == "ReduceLROnPlateauScheduler":
-            self.scheduler.step(self.current_val_loss)
-
-        return {
-            'loss': self.current_val_loss,
-            'log': {'val_loss': self.current_val_loss},
-        }
-
     def test_step(self, batch: tuple, batch_idx: int):
         r"""
         Forward propagate a `inputs` and `targets` pair for test.
@@ -172,12 +140,9 @@ def configure_optimizers(self):
 
 
         Returns:
-            - **Two lists** - The first list has multiple optimizers, and the second has multiple LR schedulers
+            - **Dictionary** - The first item has multiple optimizers, and the second has multiple LR schedulers
                 (or multiple ``lr_dict``).
         """
-        from torch.optim import Adam, Adagrad, Adadelta, Adamax, AdamW, SGD, ASGD
-        from openspeech.optim import AdamP, RAdam, Novograd
-
         SUPPORTED_OPTIMIZERS = {
             "adam": Adam,
             "adamp": AdamP,
@@ -199,16 +164,21 @@ def configure_optimizers(self):
             self.parameters(),
             lr=self.configs.lr_scheduler.lr,
         )
-        self.scheduler = SCHEDULER_REGISTRY[self.configs.lr_scheduler.scheduler_name](
-            optimizer=self.optimizer,
-            configs=self.configs,
-        )
+        scheduler = SCHEDULER_REGISTRY[self.configs.lr_scheduler.scheduler_name](self.optimizer, self.configs)
+
+        if self.configs.lr_scheduler.scheduler_name in ("reduce_lr_on_plateau", "warmup_reduce_lr_on_plateau"):
+            lr_scheduler = {
+                'scheduler': scheduler,
+                'monitor': 'val_loss',
+            }
+        else:
+            lr_scheduler = {
+                'scheduler': scheduler
+            }
+
         return {
             'optimizer': self.optimizer,
-            'lr_scheduler': {
-                'scheduler': self.scheduler,
-                'monitor': 'metric_to_track',
-            }
+            'lr_scheduler': lr_scheduler
         }
 
     def configure_criterion(self, criterion_name: str) -> nn.Module:
diff --git a/openspeech/models/openspeech_transducer_model.py b/openspeech/models/openspeech_transducer_model.py
index 656a7f3..3abec9b 100644
--- a/openspeech/models/openspeech_transducer_model.py
+++ b/openspeech/models/openspeech_transducer_model.py
@@ -93,35 +93,38 @@ def collect_outputs(
             wer = self.wer_metric(targets[:, 1:], predictions)
             cer = self.cer_metric(targets[:, 1:], predictions)
 
-            self.log_steps(stage, wer, cer, loss)
-
-            tqdm_dict = {
+            self.info({
                 f"{stage}_loss": loss,
-                "wer": wer,
-                "cer": cer,
-            }
+                f"{stage}_wer": wer,
+                f"{stage}_cer": cer,
+                "learning_rate": self.get_lr(),
+            })
 
             return OrderedDict({
                 "loss": loss,
-                "progress_bar": tqdm_dict,
-                "log": tqdm_dict,
+                "wer": wer,
+                "cer": cer,
+                "predictions": predictions,
+                "targets": targets,
+                "logits": logits,
             })
 
         else:
             wer = self.wer_metric(targets[:, 1:], predictions)
             cer = self.cer_metric(targets[:, 1:], predictions)
-            self.log_steps(stage, wer, cer)
 
-            progress_bar_dict = {
-                f"{stage}_loss": None,
-                "wer": wer,
-                "cer": cer,
-            }
+            self.info({
+                f"{stage}_wer": wer,
+                f"{stage}_cer": cer,
+            })
 
             return OrderedDict({
                 "loss": None,
-                "progress_bar": progress_bar_dict,
-                "log": progress_bar_dict,
+                "wer": wer,
+                "cer": cer,
+                "predictions": predictions,
+                "targets": targets,
+                "logits": logits,
             })
 
     def _expand_for_joint(self, encoder_outputs: Tensor, decoder_outputs: Tensor) -> Tuple[Tensor, Tensor]:
@@ -282,7 +285,7 @@ def validation_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
         predictions = torch.stack(predictions)
 
         return self.collect_outputs(
-            'valid',
+            'val',
             logits=None,
             input_lengths=input_lengths,
             targets=targets,
diff --git a/openspeech/models/transformer_lm/configurations.py b/openspeech/models/transformer_lm/configurations.py
new file mode 100644
index 0000000..5e63de8
--- /dev/null
+++ b/openspeech/models/transformer_lm/configurations.py
@@ -0,0 +1,71 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+from dataclasses import dataclass, field
+
+from openspeech.dataclass.configurations import OpenspeechDataclass
+
+
+@dataclass
+class TransformerLanguageModelConfigs(OpenspeechDataclass):
+    r"""
+    This is the configuration class to store the configuration of
+    a :class:`~openspeech.models.ListenAttendSpellWithMultiHead`.
+
+    It is used to initiated an `ListenAttendSpellWithMultiHead` model.
+
+    Configuration objects inherit from :class: `~openspeech.dataclass.configs.OpenspeechDataclass`.
+
+    Args:
+        model_name (str): Model name (default: transformer_lm)
+        num_layers (int): The number of lstm layers. (default: 6)
+        d_model (int): The dimension of model. (default: 768)
+        dropout_p (float): The dropout probability of encoder. (default: 0.3)
+        d_ff (int): Dimenstion of feed forward network. (default: 2048)
+        num_attention_heads (int): The number of attention heads. (default: 8)
+        max_length (int): Max decoding length. (default: 128)
+        optimizer (str): Optimizer for training. (default: adam)
+    """
+    model_name: str = field(
+        default="transformer_lm", metadata={"help": "Model name"}
+    )
+    num_layers: int = field(
+        default=6, metadata={"help": "The number of encoder layers."}
+    )
+    d_model: int = field(
+        default=768, metadata={"help": "The dimension of model."}
+    )
+    d_ff: int = field(
+        default=1536, metadata={"help": "The dimenstion of feed forward network."}
+    )
+    num_attention_heads: int = field(
+        default=8, metadata={"help": "The number of attention heads."}
+    )
+    dropout_p: float = field(
+        default=0.3, metadata={"help": "The dropout probability of encoder."}
+    )
+    max_length: int = field(
+        default=128, metadata={"help": "Max decoding length."}
+    )
+    optimizer: str = field(
+        default="adam", metadata={"help": "Optimizer for training."}
+    )
diff --git a/openspeech/models/transformer_lm/model.py b/openspeech/models/transformer_lm/model.py
new file mode 100644
index 0000000..0ed5ab6
--- /dev/null
+++ b/openspeech/models/transformer_lm/model.py
@@ -0,0 +1,137 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import torch
+from omegaconf import DictConfig
+from collections import OrderedDict
+from typing import Dict
+
+from openspeech.lm.transformer_for_causal_lm import TransformerForCausalLM
+from openspeech.models import register_model, OpenspeechModel
+from openspeech.models.transformer_lm.configurations import TransformerLanguageModelConfigs
+from openspeech.vocabs.vocab import Vocabulary
+
+
+@register_model('transformer_lm', dataclass=TransformerLanguageModelConfigs)
+class TransformerLanguageModel(OpenspeechModel):
+    def __init__(self, configs: DictConfig, vocab: Vocabulary, ) -> None:
+        super(TransformerLanguageModel, self).__init__(configs, vocab)
+
+    def build_model(self):
+        self.lm = TransformerForCausalLM(
+            num_classes=self.num_classes,
+            max_length=self.configs.model.max_length,
+            d_model=self.configs.model.d_model,
+            d_ff=self.configs.model.d_ff,
+            num_attention_heads=self.configs.model.num_attention_heads,
+            pad_id=self.vocab.pad_id,
+            sos_id=self.vocab.sos_id,
+            eos_id=self.vocab.eos_id,
+            dropout_p=self.configs.model.dropout_p,
+            num_layers=self.configs.model.num_layers,
+        )
+
+    def collect_outputs(
+            self,
+            stage: str,
+            logits: torch.Tensor,
+            targets: torch.Tensor,
+    ) -> OrderedDict:
+        perplexity = self.criterion(logits, targets[:, 1:])
+        predictions = logits.max(-1)[1]
+
+        self.info({
+            f"{stage}_perplexity": perplexity,
+            "learning_rate": self.get_lr(),
+        })
+
+        return OrderedDict({
+            "perplexity": perplexity,
+            "logits": logits,
+            "targets": targets,
+            "predictions": predictions,
+        })
+
+    def forward(self, inputs: torch.Tensor, input_lengths: torch.Tensor) -> Dict[str, torch.Tensor]:
+        logits = self.lm(inputs, input_lengths)
+        predictions = logits.max(-1)[1]
+        return {
+            "predictions": predictions,
+            "logits": logits,
+        }
+
+    def training_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
+        r"""
+        Forward propagate a `inputs` and `targets` pair for training.
+
+        Inputs:
+            train_batch (tuple): A train batch contains `inputs`
+            batch_idx (int): The index of batch
+
+        Returns:
+            loss (torch.Tensor): loss for training
+        """
+        inputs, input_lengths, targets = batch
+        logits = self.lm(inputs, input_lengths)
+        return self.collect_outputs(
+            stage='train',
+            logits=logits,
+            targets=targets,
+        )
+
+    def validation_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
+        r"""
+        Forward propagate a `inputs` and `targets` pair for validation.
+
+        Inputs:
+            train_batch (tuple): A train batch contains `inputs`
+            batch_idx (int): The index of batch
+
+        Returns:
+            loss (torch.Tensor): loss for training
+        """
+        inputs, input_lengths, targets = batch
+        logits = self.lm(inputs, input_lengths)
+        return self.collect_outputs(
+            stage='val',
+            logits=logits,
+            targets=targets,
+        )
+
+    def test_step(self, batch: tuple, batch_idx: int) -> OrderedDict:
+        r"""
+        Forward propagate a `inputs` and `targets` pair for test.
+
+        Inputs:
+            train_batch (tuple): A train batch contains `inputs`
+            batch_idx (int): The index of batch
+
+        Returns:
+            loss (torch.Tensor): loss for training
+        """
+        inputs, input_lengths, targets = batch
+        logits = self.lm(inputs, input_lengths)
+        return self.collect_outputs(
+            stage='test',
+            logits=logits,
+            targets=targets,
+        )
diff --git a/openspeech/optim/scheduler/lr_scheduler.py b/openspeech/optim/scheduler/lr_scheduler.py
index d670640..2e87552 100644
--- a/openspeech/optim/scheduler/lr_scheduler.py
+++ b/openspeech/optim/scheduler/lr_scheduler.py
@@ -20,8 +20,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+from torch.optim.lr_scheduler import _LRScheduler
 
-class LearningRateScheduler(object):
+
+class LearningRateScheduler(_LRScheduler):
     r"""
     Provides inteface of learning rate scheduler.
 
diff --git a/openspeech/optim/scheduler/reduce_lr_on_plateau_scheduler.py b/openspeech/optim/scheduler/reduce_lr_on_plateau_scheduler.py
index 27494de..99aad48 100644
--- a/openspeech/optim/scheduler/reduce_lr_on_plateau_scheduler.py
+++ b/openspeech/optim/scheduler/reduce_lr_on_plateau_scheduler.py
@@ -20,10 +20,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-import torch
 from dataclasses import dataclass, field
 from omegaconf import DictConfig
 from torch.optim import Optimizer
+from typing import Optional
 
 from openspeech.dataclass.configurations import LearningRateSchedulerConfigs
 from openspeech.optim.scheduler import register_scheduler
@@ -66,7 +66,7 @@ def __init__(
         self.val_loss = 100.0
         self.count = 0
 
-    def step(self, val_loss: float):
+    def step(self, val_loss: Optional[float] = None):
         if val_loss is not None:
             if self.val_loss < val_loss:
                 self.count += 1
diff --git a/openspeech/optim/scheduler/tri_stage_lr_scheduler.py b/openspeech/optim/scheduler/tri_stage_lr_scheduler.py
index 0dcd075..aaf5337 100644
--- a/openspeech/optim/scheduler/tri_stage_lr_scheduler.py
+++ b/openspeech/optim/scheduler/tri_stage_lr_scheduler.py
@@ -37,6 +37,9 @@ class TriStageLRSchedulerConfigs(LearningRateSchedulerConfigs):
     scheduler_name: str = field(
         default="tri_stage", metadata={"help": "Name of learning rate scheduler."}
     )
+    init_lr: float = field(
+        default=1e-7, metadata={"help": "Initial learning rate."}
+    )
     init_lr_scale: float = field(
         default=0.01, metadata={"help": "Initial learning rate scale."}
     )
diff --git a/openspeech/optim/scheduler/warmup_scheduler.py b/openspeech/optim/scheduler/warmup_scheduler.py
index 8a190dc..004fd79 100644
--- a/openspeech/optim/scheduler/warmup_scheduler.py
+++ b/openspeech/optim/scheduler/warmup_scheduler.py
@@ -40,7 +40,7 @@ class WarmupLRSchedulerConfigs(LearningRateSchedulerConfigs):
         default=1e-04, metadata={"help": "Maximum learning rate."}
     )
     init_lr: float = field(
-        default=1e-10, metadata={"help": "Initial learning rate."}
+        default=1e-7, metadata={"help": "Initial learning rate."}
     )
     warmup_steps: int = field(
         default=4000, metadata={"help": "Warmup the learning rate linearly for the first N updates"}
diff --git a/openspeech/search/beam_search_lstm.py b/openspeech/search/beam_search_lstm.py
index 9459ec9..a54bc2d 100644
--- a/openspeech/search/beam_search_lstm.py
+++ b/openspeech/search/beam_search_lstm.py
@@ -23,7 +23,7 @@
 import torch
 
 from openspeech.search.base import OpenspeechBeamSearchBase
-from openspeech.decoders import LSTMDecoder
+from openspeech.decoders import LSTMAttentionDecoder
 
 
 class BeamSearchLSTM(OpenspeechBeamSearchBase):
@@ -47,7 +47,7 @@ class BeamSearchLSTM(OpenspeechBeamSearchBase):
     Returns:
         * logits (torch.FloatTensor): Log probability of model predictions.
     """
-    def __init__(self, decoder: LSTMDecoder, beam_size: int):
+    def __init__(self, decoder: LSTMAttentionDecoder, beam_size: int):
         super(BeamSearchLSTM, self).__init__(decoder, beam_size)
         self.hidden_state_dim = decoder.hidden_state_dim
         self.num_layers = decoder.num_layers
diff --git a/openspeech/utils.py b/openspeech/utils.py
index 5e9db01..fbbae40 100644
--- a/openspeech/utils.py
+++ b/openspeech/utils.py
@@ -27,7 +27,7 @@
 from collections import OrderedDict
 from typing import Tuple, Union, Iterable
 from omegaconf import DictConfig, OmegaConf
-
+from pytorch_lightning.callbacks import LearningRateMonitor
 
 PYTORCH_IMPORT_ERROR = """
 Openspeech requires the PyTorch library but it was not found in your environment. Checkout the instructions on the
@@ -93,6 +93,17 @@
 DUMMY_TARGET_LENGTHS = torch.IntTensor([9, 8, 7])
 DUMMY_TRANSCRIPTS = "OPENSPEECH IS AWESOME"
 
+DUMMY_LM_INPUTS = torch.LongTensor([
+    [2, 3, 3, 3, 3, 3, 2, 2, 0],
+    [2, 3, 3, 3, 3, 3, 2, 3, 2],
+    [2, 3, 3, 3, 3, 3, 2, 2, 0],
+])
+DYMMY_LM_INPUT_LENGTHS = torch.IntTensor([9, 8, 7])
+DUMMY_LM_TARGETS = torch.LongTensor([
+    [3, 3, 3, 3, 3, 2, 2, 1, 0],
+    [3, 3, 3, 3, 3, 2, 1, 2, 0],
+    [3, 3, 3, 3, 3, 2, 2, 0, 1],
+])
 
 def is_pytorch_available():
     return importlib.util.find_spec("torch") is not None
@@ -192,7 +203,7 @@ def parse_configs(configs: DictConfig) -> Tuple[Union[TensorBoardLogger, bool],
     if configs.trainer.logger == "tensorboard":
         logger = TensorBoardLogger("logs/")
     elif configs.trainer.logger == "wandb":
-        logger = WandbLogger("logs/")
+        logger = WandbLogger(project=f"{configs.model.model_name}-{configs.dataset.dataset}", job_type='train')
     else:
         logger = True
 
@@ -216,7 +227,8 @@ def get_pl_trainer(
                              gradient_clip_val=configs.trainer.gradient_clip_val,
                              logger=logger,
                              auto_scale_batch_size=configs.trainer.auto_scale_batch_size,
-                             max_epochs=configs.trainer.max_epochs)
+                             max_epochs=configs.trainer.max_epochs,
+                             callbacks=[LearningRateMonitor(logging_interval='step')])
     elif configs.trainer.name == "gpu":
         trainer = pl.Trainer(accelerator=configs.trainer.accelerator,
                              gpus=num_devices,
@@ -226,7 +238,8 @@ def get_pl_trainer(
                              gradient_clip_val=configs.trainer.gradient_clip_val,
                              logger=logger,
                              auto_scale_batch_size=configs.trainer.auto_scale_batch_size,
-                             max_epochs=configs.trainer.max_epochs)
+                             max_epochs=configs.trainer.max_epochs,
+                             callbacks=[LearningRateMonitor(logging_interval='step')])
     elif configs.trainer.name == "tpu":
         trainer = pl.Trainer(accelerator=configs.trainer.accelerator,
                              tpu_cores=configs.trainer.tpu_cores,
@@ -236,7 +249,8 @@ def get_pl_trainer(
                              gradient_clip_val=configs.trainer.gradient_clip_val,
                              logger=logger,
                              auto_scale_batch_size=configs.trainer.auto_scale_batch_size,
-                             max_epochs=configs.trainer.max_epochs)
+                             max_epochs=configs.trainer.max_epochs,
+                             callbacks=[LearningRateMonitor(logging_interval='step')])
     elif configs.trainer.name == "gpu-fp16":
         trainer = pl.Trainer(precision=configs.trainer.precision,
                              accelerator=configs.trainer.accelerator,
@@ -248,7 +262,8 @@ def get_pl_trainer(
                              gradient_clip_val=configs.trainer.gradient_clip_val,
                              logger=logger,
                              auto_scale_batch_size=configs.trainer.auto_scale_batch_size,
-                             max_epochs=configs.trainer.max_epochs)
+                             max_epochs=configs.trainer.max_epochs,
+                             callbacks=[LearningRateMonitor(logging_interval='step')])
     elif configs.trainer.name == "tpu-fp16":
         trainer = pl.Trainer(precision=configs.trainer.precision,
                              accelerator=configs.trainer.accelerator,
@@ -259,7 +274,8 @@ def get_pl_trainer(
                              gradient_clip_val=configs.trainer.gradient_clip_val,
                              logger=logger,
                              auto_scale_batch_size=configs.trainer.auto_scale_batch_size,
-                             max_epochs=configs.trainer.max_epochs)
+                             max_epochs=configs.trainer.max_epochs,
+                             callbacks=[LearningRateMonitor(logging_interval='step')])
     elif configs.trainer.name == "cpu-fp64":
         trainer = pl.Trainer(precision=configs.trainer.precision,
                              accelerator=configs.trainer.accelerator,
@@ -269,7 +285,8 @@ def get_pl_trainer(
                              gradient_clip_val=configs.trainer.gradient_clip_val,
                              logger=logger,
                              auto_scale_batch_size=configs.trainer.auto_scale_batch_size,
-                             max_epochs=configs.trainer.max_epochs)
+                             max_epochs=configs.trainer.max_epochs,
+                             callbacks=[LearningRateMonitor(logging_interval='step')])
     else:
         raise ValueError(f"Unsupported trainer: {configs.trainer.name}")
 
diff --git a/openspeech/vocabs/aishell/character.py b/openspeech/vocabs/aishell/character.py
index 073f8b6..6cce1de 100644
--- a/openspeech/vocabs/aishell/character.py
+++ b/openspeech/vocabs/aishell/character.py
@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import torch
 import csv
 from dataclasses import dataclass, field
 from omegaconf import DictConfig
@@ -55,6 +56,9 @@ def __init__(self, configs: DictConfig):
             encoding=configs.vocab.encoding,
         )
         self.labels = self.vocab_dict.keys()
+        self.sos_token = configs.vocab.sos_token
+        self.eos_token = configs.vocab.eos_token
+        self.pad_token = configs.vocab.pad_token
         self.sos_id = int(self.vocab_dict[configs.vocab.sos_token])
         self.eos_id = int(self.vocab_dict[configs.vocab.eos_token])
         self.pad_id = int(self.vocab_dict[configs.vocab.pad_token])
@@ -96,6 +100,17 @@ def label_to_string(self, labels):
             sentences.append(sentence)
         return sentences
 
+    def string_to_label(self, sentence):
+        label = str()
+
+        for ch in sentence:
+            try:
+                label += (str(self.vocab_dict[ch]) + ' ')
+            except KeyError:
+                continue
+
+        return label[:-1]
+
     def load_vocab(self, vocab_path, encoding='utf-8'):
         r"""
         Provides char2id, id2char
diff --git a/openspeech/vocabs/ksponspeech/character.py b/openspeech/vocabs/ksponspeech/character.py
index 2d7436f..e879c8b 100644
--- a/openspeech/vocabs/ksponspeech/character.py
+++ b/openspeech/vocabs/ksponspeech/character.py
@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import torch
 import csv
 from omegaconf import DictConfig
 from dataclasses import dataclass, MISSING, field
@@ -95,6 +96,17 @@ def label_to_string(self, labels):
             sentences.append(sentence)
         return sentences
 
+    def string_to_label(self, sentence):
+        label = str()
+
+        for ch in sentence:
+            try:
+                label += (str(self.vocab_dict[ch]) + ' ')
+            except KeyError:
+                continue
+
+        return label[:-1]
+
     def load_vocab(self, vocab_path, encoding='utf-8'):
         r"""
         Provides char2id, id2char
diff --git a/openspeech/vocabs/ksponspeech/grapheme.py b/openspeech/vocabs/ksponspeech/grapheme.py
index dee4e55..075a25f 100644
--- a/openspeech/vocabs/ksponspeech/grapheme.py
+++ b/openspeech/vocabs/ksponspeech/grapheme.py
@@ -20,8 +20,10 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import torch
 import csv
-from dataclasses import dataclass, MISSING, field
+import unicodedata
+from dataclasses import dataclass, field
 from omegaconf import DictConfig
 
 from openspeech.dataclass.configurations import VocabularyConfigs
@@ -95,6 +97,17 @@ def label_to_string(self, labels):
             sentences.append(sentence)
         return sentences
 
+    def string_to_label(self, sentence):
+        label = str()
+
+        for ch in sentence:
+            try:
+                label += (str(self.vocab_dict[ch]) + ' ')
+            except KeyError:
+                continue
+
+        return label[:-1]
+
     def load_vocab(self, vocab_path, encoding='utf-8'):
         """
         Provides char2id, id2char
diff --git a/openspeech/vocabs/ksponspeech/subword.py b/openspeech/vocabs/ksponspeech/subword.py
index a6df731..e324aea 100644
--- a/openspeech/vocabs/ksponspeech/subword.py
+++ b/openspeech/vocabs/ksponspeech/subword.py
@@ -20,8 +20,9 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import torch
 import sentencepiece as spm
-from dataclasses import dataclass, MISSING, field
+from dataclasses import dataclass, field
 from omegaconf import DictConfig
 
 from openspeech.dataclass.configurations import VocabularyConfigs
@@ -93,3 +94,8 @@ def label_to_string(self, labels):
                 sentence = self.sp.DecodeIds([int(l) for l in label])
             sentences.append(sentence)
         return sentences
+
+    def string_to_label(self, sentence):
+        text = " ".join(self.sp.EncodeAsPieces(sentence))
+        label = " ".join([str(self.sp.PieceToId(token)) for token in text])
+        return label
diff --git a/openspeech/vocabs/librispeech/character.py b/openspeech/vocabs/librispeech/character.py
index 6609395..e172e43 100644
--- a/openspeech/vocabs/librispeech/character.py
+++ b/openspeech/vocabs/librispeech/character.py
@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import torch
 import csv
 from omegaconf import DictConfig
 from dataclasses import dataclass, MISSING, field
@@ -95,6 +96,17 @@ def label_to_string(self, labels):
             sentences.append(sentence)
         return sentences
 
+    def string_to_label(self, sentence):
+        label = str()
+
+        for ch in sentence:
+            try:
+                label += (str(self.vocab_dict[ch]) + ' ')
+            except KeyError:
+                continue
+
+        return label[:-1]
+
     def load_vocab(self, vocab_path, encoding='utf-8'):
         r"""
         Provides char2id, id2char
diff --git a/openspeech/vocabs/librispeech/subword.py b/openspeech/vocabs/librispeech/subword.py
index 065934f..c2ddeab 100644
--- a/openspeech/vocabs/librispeech/subword.py
+++ b/openspeech/vocabs/librispeech/subword.py
@@ -89,3 +89,8 @@ def label_to_string(self, labels):
             return sentences
         else:
             raise ValueError("Unsupported label's shape")
+
+    def string_to_label(self, sentence):
+        text = " ".join(self.sp.EncodeAsPieces(sentence))
+        label = " ".join([str(self.sp.PieceToId(token)) for token in text])
+        return label
diff --git a/openspeech/vocabs/vocab.py b/openspeech/vocabs/vocab.py
index 331cce7..e9e4b1b 100644
--- a/openspeech/vocabs/vocab.py
+++ b/openspeech/vocabs/vocab.py
@@ -34,3 +34,6 @@ def __init__(self, *args, **kwargs):
 
     def label_to_string(self, labels):
         raise NotImplementedError
+
+    def string_to_label(self, labels):
+        raise NotImplementedError
diff --git a/openspeech_cli/hydra_ensemble_eval.py b/openspeech_cli/hydra_ensemble_eval.py
index 58ac2a0..19257a3 100644
--- a/openspeech_cli/hydra_ensemble_eval.py
+++ b/openspeech_cli/hydra_ensemble_eval.py
@@ -28,8 +28,9 @@
 from pytorch_lightning.utilities import rank_zero_info
 
 from openspeech.metrics import WordErrorRate, CharacterErrorRate
-from openspeech.data.dataset import SpeechToTextDataset
-from openspeech.data.data_loader import load_dataset, AudioDataLoader, BucketingSampler
+from openspeech.data.audio.dataset import SpeechToTextDataset
+from openspeech.data.sampler import BucketingSampler
+from openspeech.data.audio.data_loader import load_dataset, AudioDataLoader
 from openspeech.dataclass.initialize import hydra_eval_init
 from openspeech.models import MODEL_REGISTRY
 from openspeech.search.ensemble_search import EnsembleSearch, WeightedEnsembleSearch
diff --git a/openspeech_cli/hydra_eval.py b/openspeech_cli/hydra_eval.py
index 404b0ee..f548357 100644
--- a/openspeech_cli/hydra_eval.py
+++ b/openspeech_cli/hydra_eval.py
@@ -28,8 +28,9 @@
 from openspeech.metrics import WordErrorRate, CharacterErrorRate
 from pytorch_lightning.utilities import rank_zero_info
 
-from openspeech.data.dataset import SpeechToTextDataset
-from openspeech.data.data_loader import load_dataset, AudioDataLoader, BucketingSampler
+from openspeech.data.audio.dataset import SpeechToTextDataset
+from openspeech.data.sampler import BucketingSampler
+from openspeech.data.audio.data_loader import load_dataset, AudioDataLoader
 from openspeech.dataclass.initialize import hydra_eval_init
 from openspeech.models import MODEL_REGISTRY
 
diff --git a/openspeech_cli/hydra_lm_train.py b/openspeech_cli/hydra_lm_train.py
new file mode 100644
index 0000000..2bef2ca
--- /dev/null
+++ b/openspeech_cli/hydra_lm_train.py
@@ -0,0 +1,56 @@
+# MIT License
+#
+# Copyright (c) 2021 Soohwan Kim and Sangchun Ha and Soyoung Cho
+#
+# Permission is hereby granted, free of charge, to any person obtaining a copy
+# of this software and associated documentation files (the "Software"), to deal
+# in the Software without restriction, including without limitation the rights
+# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+# copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+#
+# The above copyright notice and this permission notice shall be included in all
+# copies or substantial portions of the Software.
+#
+# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE.
+
+import os
+import hydra
+import pytorch_lightning as pl
+from omegaconf import DictConfig, OmegaConf
+from pytorch_lightning.utilities import rank_zero_info
+
+from openspeech.dataclass.initialize import hydra_lm_train_init
+from openspeech.datasets import DATA_MODULE_REGISTRY
+from openspeech.models import MODEL_REGISTRY
+from openspeech.utils import parse_configs, get_pl_trainer
+
+
+@hydra.main(config_path=os.path.join("..", "openspeech", "configs"), config_name="lm_train")
+def hydra_main(configs: DictConfig) -> None:
+    rank_zero_info(OmegaConf.to_yaml(configs))
+    pl.seed_everything(configs.trainer.seed)
+
+    logger, num_devices = parse_configs(configs)
+
+    data_module = DATA_MODULE_REGISTRY[configs.dataset.dataset](configs)
+    vocab = data_module.prepare_data()
+    data_module.setup(vocab=vocab)
+
+    model = MODEL_REGISTRY[configs.model.model_name](configs=configs, vocab=vocab)
+    model.build_model()
+
+    trainer = get_pl_trainer(configs, num_devices, logger)
+    trainer.fit(model, data_module)
+    trainer.test()
+
+
+if __name__ == '__main__':
+    hydra_lm_train_init()
+    hydra_main()
diff --git a/openspeech_cli/hydra_train.py b/openspeech_cli/hydra_train.py
index 187bde2..636599f 100644
--- a/openspeech_cli/hydra_train.py
+++ b/openspeech_cli/hydra_train.py
@@ -22,13 +22,13 @@
 
 import os
 import hydra
+import wandb
 import pytorch_lightning as pl
-import warnings
 from omegaconf import DictConfig, OmegaConf
 from pytorch_lightning.utilities import rank_zero_info
 
 from openspeech.datasets import DATA_MODULE_REGISTRY
-from openspeech.dataclass.initialize import hydra_init
+from openspeech.dataclass.initialize import hydra_train_init
 from openspeech.models import MODEL_REGISTRY
 from openspeech.utils import parse_configs, get_pl_trainer
 
@@ -53,6 +53,5 @@ def hydra_main(configs: DictConfig) -> None:
 
 
 if __name__ == '__main__':
-    warnings.filterwarnings("ignore")
-    hydra_init()
+    hydra_train_init()
     hydra_main()


Test patch:
diff --git a/tests/test_lstm_for_causal_lm.py b/tests/test_lstm_for_causal_lm.py
new file mode 100644
index 0000000..1685fe7
--- /dev/null
+++ b/tests/test_lstm_for_causal_lm.py
@@ -0,0 +1,94 @@
+import unittest
+import logging
+import torch
+
+from openspeech.criterion import Perplexity, PerplexityLossConfigs
+from openspeech.lm.lstm_for_causal_lm import LSTMForCausalLM
+from openspeech.utils import DUMMY_TARGETS, build_dummy_configs
+from openspeech.vocabs.ksponspeech.character import KsponSpeechCharacterVocabulary
+
+logger = logging.getLogger(__name__)
+
+
+class TestLSTMForCausalLM(unittest.TestCase):
+    def test_lstm_forward(self):
+        configs = build_dummy_configs(
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+
+        model = LSTMForCausalLM(
+            num_classes=4,
+            max_length=32,
+            hidden_state_dim=64,
+            pad_id=0,
+            sos_id=1,
+            eos_id=2,
+            rnn_type='lstm',
+        )
+
+        criterion = Perplexity(configs, vocab)
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+
+        for i in range(3):
+            logits = model(DUMMY_TARGETS, teacher_forcing_ratio=1.0)
+
+            loss = criterion(logits, DUMMY_TARGETS[:, 1:])
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            logits = model(DUMMY_TARGETS, teacher_forcing_ratio=1.0)
+
+            loss = criterion(logits, DUMMY_TARGETS[:, 1:])
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            logits = model(DUMMY_TARGETS, teacher_forcing_ratio=0.0)
+
+            loss = criterion(logits, DUMMY_TARGETS[:, 1:])
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+    def test_gru_forward(self):
+        configs = build_dummy_configs(
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+
+        model = LSTMForCausalLM(
+            num_classes=4,
+            max_length=32,
+            hidden_state_dim=64,
+            pad_id=0,
+            sos_id=1,
+            eos_id=2,
+            rnn_type='gru',
+        )
+
+        criterion = Perplexity(configs, vocab)
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+
+        for i in range(3):
+            logits = model(DUMMY_TARGETS, teacher_forcing_ratio=1.0)
+
+            loss = criterion(logits, DUMMY_TARGETS[:, 1:])
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            logits = model(DUMMY_TARGETS, teacher_forcing_ratio=0.0)
+
+            loss = criterion(logits, DUMMY_TARGETS[:, 1:])
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_lstm_lm.py b/tests/test_lstm_lm.py
new file mode 100644
index 0000000..155c7b5
--- /dev/null
+++ b/tests/test_lstm_lm.py
@@ -0,0 +1,77 @@
+import unittest
+import logging
+import torch
+
+from openspeech.criterion import Perplexity, PerplexityLossConfigs
+from openspeech.models.lstm_lm.configurations import LSTMLanguageModelConfigs
+from openspeech.models.lstm_lm.model import LSTMLanguageModel
+from openspeech.utils import DUMMY_LM_INPUTS, DUMMY_LM_TARGETS, build_dummy_configs
+from openspeech.vocabs.ksponspeech.character import KsponSpeechCharacterVocabulary
+
+logger = logging.getLogger(__name__)
+
+
+class TestLSTMLanguageModel(unittest.TestCase):
+    def test_forward(self):
+        configs = build_dummy_configs(
+            model_configs=LSTMLanguageModelConfigs,
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+        model = LSTMLanguageModel(configs, vocab)
+        model.build_model()
+
+        criterion = Perplexity(configs, vocab)
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+        ptimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+
+        for i in range(3):
+            outputs = model(DUMMY_LM_INPUTS)
+            loss = criterion(outputs['logits'], DUMMY_LM_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            outputs = model(DUMMY_LM_INPUTS)
+            loss = criterion(outputs['logits'], DUMMY_LM_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            outputs = model(DUMMY_LM_INPUTS)
+            loss = criterion(outputs['logits'], DUMMY_LM_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+    def test_training_step(self):
+        configs = build_dummy_configs(
+            model_configs=LSTMLanguageModelConfigs,
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+        model = LSTMLanguageModel(configs, vocab)
+        model.build_model()
+
+        for i in range(5):
+            outputs = model.training_step(batch=(DUMMY_LM_INPUTS, DUMMY_LM_TARGETS), batch_idx=i)
+            assert type(outputs["perplexity"].item()) == float
+
+    def test_validation_step(self):
+        configs = build_dummy_configs(
+            model_configs=LSTMLanguageModelConfigs,
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+        model = LSTMLanguageModel(configs, vocab)
+        model.build_model()
+
+        for i in range(5):
+            outputs = model.validation_step(batch=(DUMMY_LM_INPUTS, DUMMY_LM_TARGETS), batch_idx=i)
+            assert type(outputs["perplexity"].item()) == float
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_transformer_for_causal_lm.py b/tests/test_transformer_for_causal_lm.py
new file mode 100644
index 0000000..603825c
--- /dev/null
+++ b/tests/test_transformer_for_causal_lm.py
@@ -0,0 +1,74 @@
+import unittest
+import logging
+import torch
+
+from openspeech.criterion import Perplexity, PerplexityLossConfigs
+from openspeech.lm.transformer_for_causal_lm import TransformerForCausalLM
+from openspeech.utils import build_dummy_configs, DUMMY_LM_INPUTS, DUMMY_INPUT_LENGTHS, DUMMY_TARGETS
+from openspeech.vocabs.ksponspeech.character import KsponSpeechCharacterVocabulary
+
+logger = logging.getLogger(__name__)
+
+
+class TestTransformerForCausalLM(unittest.TestCase):
+    def test_forward(self):
+        configs = build_dummy_configs(
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+
+        model = TransformerForCausalLM(
+            num_classes=4,
+            max_length=32,
+            d_model=64,
+            num_heads=4,
+            d_ff=128,
+            pad_id=0,
+            sos_id=1,
+            eos_id=2,
+            num_layers=3,
+        )
+
+        criterion = Perplexity(configs, vocab)
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+
+        for i in range(3):
+            logits = model(DUMMY_LM_INPUTS, DUMMY_INPUT_LENGTHS)
+
+            loss = criterion(logits, DUMMY_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+    def test_forward_step(self):
+        configs = build_dummy_configs(
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+
+        model = TransformerForCausalLM(
+            num_classes=4,
+            max_length=32,
+            d_model=64,
+            num_heads=4,
+            d_ff=128,
+            pad_id=0,
+            sos_id=1,
+            eos_id=2,
+            num_layers=3,
+        )
+
+        criterion = Perplexity(configs, vocab)
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+
+        for i in range(1, 4):
+            input_lengths = torch.IntTensor([i] * 3)
+            logits = model.forward_step(DUMMY_LM_INPUTS[:i], input_lengths)
+            loss = criterion(logits, DUMMY_TARGETS[:i])
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/tests/test_transformer_lm.py b/tests/test_transformer_lm.py
new file mode 100644
index 0000000..d87c1b5
--- /dev/null
+++ b/tests/test_transformer_lm.py
@@ -0,0 +1,83 @@
+import unittest
+import logging
+import torch
+
+from openspeech.criterion import Perplexity, PerplexityLossConfigs
+from openspeech.models.transformer_lm.configurations import TransformerLanguageModelConfigs
+from openspeech.models.transformer_lm.model import TransformerLanguageModel
+from openspeech.utils import DUMMY_LM_INPUTS, DUMMY_LM_TARGETS, DYMMY_LM_INPUT_LENGTHS, build_dummy_configs
+from openspeech.vocabs.ksponspeech.character import KsponSpeechCharacterVocabulary
+
+logger = logging.getLogger(__name__)
+
+
+class TestTransformerLanguageModel(unittest.TestCase):
+    def test_forward(self):
+        configs = build_dummy_configs(
+            model_configs=TransformerLanguageModelConfigs,
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+        model = TransformerLanguageModel(configs, vocab)
+        model.build_model()
+        model.configure_optimizers()
+
+        criterion = Perplexity(configs, vocab)
+        optimizer = torch.optim.Adam(model.parameters(), lr=1e-04)
+
+        for i in range(3):
+            outputs = model(DUMMY_LM_INPUTS, DYMMY_LM_INPUT_LENGTHS)
+            loss = criterion(outputs['logits'], DUMMY_LM_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            outputs = model(DUMMY_LM_INPUTS, DYMMY_LM_INPUT_LENGTHS)
+            loss = criterion(outputs['logits'], DUMMY_LM_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+        for i in range(3):
+            outputs = model(DUMMY_LM_INPUTS, DYMMY_LM_INPUT_LENGTHS)
+            loss = criterion(outputs['logits'], DUMMY_LM_TARGETS)
+            loss.backward()
+            optimizer.step()
+            assert type(loss.item()) == float
+
+    def test_training_step(self):
+        configs = build_dummy_configs(
+            model_configs=TransformerLanguageModelConfigs,
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+        model = TransformerLanguageModel(configs, vocab)
+        model.build_model()
+        model.configure_optimizers()
+
+        for i in range(5):
+            outputs = model.training_step(
+                batch=(DUMMY_LM_INPUTS, DYMMY_LM_INPUT_LENGTHS, DUMMY_LM_TARGETS), batch_idx=i
+            )
+            assert type(outputs["perplexity"].item()) == float
+
+    def test_validation_step(self):
+        configs = build_dummy_configs(
+            model_configs=TransformerLanguageModelConfigs,
+            criterion_configs=PerplexityLossConfigs(),
+        )
+        vocab = KsponSpeechCharacterVocabulary(configs)
+        model = TransformerLanguageModel(configs, vocab)
+        model.build_model()
+        model.configure_optimizers()
+
+        for i in range(5):
+            outputs = model.training_step(
+                batch=(DUMMY_LM_INPUTS, DYMMY_LM_INPUT_LENGTHS, DUMMY_LM_TARGETS), batch_idx=i
+            )
+            assert type(outputs["perplexity"].item()) == float
+
+
+if __name__ == '__main__':
+    unittest.main()