Docs/add nnlm docs

CHANGELOG.md

@@ -2,3 +2,7 @@
 
 
 ## [Unreleased]
+
+
+### Added
+- Add NNLM model

docs/models/nnlm.md

@@ -0,0 +1,3 @@
+# NNLM
+
+::: toynlp.nnlm.model

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

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