diff --git a/CHANGELOG.md b/CHANGELOG.md index 591fd50..8ba302d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,3 +2,7 @@ ## [Unreleased] + + +### Added +- Add NNLM model diff --git a/docs/models/nnlm.md b/docs/models/nnlm.md new file mode 100644 index 0000000..96ea5b2 --- /dev/null +++ b/docs/models/nnlm.md @@ -0,0 +1,3 @@ +# NNLM + +::: toynlp.nnlm.model diff --git a/mkdocs.yml b/mkdocs.yml index 6488387..743906b 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -8,6 +8,8 @@ copyright: Copyright © 2025 Xiangzhuang Shen # Page tree nav: - Home: index.md + - Models: + - NNLM: models/nnlm.md - Changelog: changelog.md - Author's website: https://shenxiangzhuang.github.io/ @@ -100,3 +102,31 @@ plugins: bib_file: "docs/references.bib" csl_file: "docs/ieee.csl" cite_inline: false + - mkdocstrings: + handlers: + python: + import: + - https://docs.python-requests.org/en/master/objects.inv + - https://docs.python.org/3/objects.inv + - https://mkdocstrings.github.io/objects.inv + - https://mkdocstrings.github.io/griffe/objects.inv + - https://python-markdown.github.io/objects.inv + options: + docstring_style: google + docstring_options: + ignore_init_summary: true + docstring_section_style: spacy + show_docstring_classes: true + show_docstring_attributes: true + heading_level: 2 + members_order: source + merge_init_into_class: true + separate_signature: true + show_root_heading: false + show_root_full_path: true + show_signature_annotations: true + signature_crossrefs: true + show_root_members_full_path: true + show_root_toc_entry: true + show_symbol_type_heading: true + show_symbol_type_toc: true diff --git a/toynlp/nnlm/README.md b/playground/nnlm/README.md similarity index 100% rename from toynlp/nnlm/README.md rename to playground/nnlm/README.md diff --git a/toynlp/nnlm/__init__.py b/toynlp/nnlm/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/toynlp/nnlm/model.py b/toynlp/nnlm/model.py index 86e6d00..da55e1b 100644 --- a/toynlp/nnlm/model.py +++ b/toynlp/nnlm/model.py @@ -8,6 +8,12 @@ def __init__( self, config: ModelConfig, ): + """ + The Neural Network Language Model (NNLM) model. + + Args: + config: ModelConfig, the model configuration. + """ super(NNLM, self).__init__() self.with_direct_connection = config.with_direct_connection self.with_dropout = config.with_dropout @@ -32,6 +38,16 @@ def __init__( self.dropout = torch.nn.Dropout(config.dropout_rate) def forward(self, tokens: torch.Tensor) -> torch.Tensor: + """ + Forward pass of the model. + + Args: + tokens: torch.Tensor, (batch_size, seq_len-1), the input tokens. + + Returns: + torch.Tensor, (batch_size, vocab_size), the logits. + + """ # tokens: (batch_size, seq_len-1) -> x: (batch_size, seq_len-1, embedding_dim) x = self.C(tokens) b, _, _ = x.shape