Use autosummary for torchaudio.prototyoe.models documentation (#3084)

Summary: Pull Request resolved: https://github.com/pytorch/audio/pull/3084

Reviewed By: mthrok

Differential Revision: D43550150

Pulled By: nateanl

fbshipit-source-id: 5c5e3d9461e375be202493e3399ff38ce5cd7690

docs/source/_templates/autosummary/model_class.rst

@@ -48,6 +48,19 @@
         ],
     }
 -%}
+{%- set prototype_factory = {
+    "torchaudio.models.Wav2Vec2Model": [
+        "emformer_hubert_model",
+        "emformer_hubert_base",
+        "conformer_wav2vec2_model",
+        "conformer_wav2vec2_base",
+    ],
+    "torchaudio.models.RNNT": [
+        "conformer_rnnt_model",
+        "conformer_rnnt_base",
+    ],
+}
+-%}
 {%- set utils={
         "torchaudio.models.Wav2Vec2Model": [
             "~torchaudio.models.wav2vec2.utils.import_fairseq_model",
@@ -122,6 +135,25 @@ Factory Functions
 
 {%- endif %}
 
+.. PROTOTYPE FACTORY FUNCTIONS
+
+{%- if prototype_factory[fullname] %}
+
+Prototype Factory Functions
+===========================
+
+.. currentmodule:: torchaudio.prototype.models
+
+.. autosummary::
+   :toctree: ../generated
+   :nosignatures:
+
+{% for item in prototype_factory[fullname] %}
+   {{["~torchaudio.prototype.models", item] | join('.')}}
+{%- endfor %}
+
+{%- endif %}
+
 .. UTILITY FUNCTIONS
 
 {%- if utils[fullname] %}
@@ -129,6 +161,8 @@ Factory Functions
 Utility Functions
 =================
 
+.. currentmodule:: torchaudio.models
+
 .. autosummary::
    :toctree: ../generated
    :nosignatures:

docs/source/_templates/autosummary/prototype_model_class.rst

+..
+  autogenerated from source/_templates/autosummary/prototype_model_class.rst
+
+.. currentmodule:: torchaudio.prototype.models
+
+..
+
+{%- set methods=["forward"] %}
+{%- set helpers={
+        "torchaudio.models.RNNTBeamSearch": [
+            "Hypothesis",
+        ],
+    }
+-%}
+{%- set factory={
+        "torchaudio.prototype.models.SquimObjective": [
+            "squim_objective_model",
+            "squim_objective_base",
+        ],
+        "torchaudio.prototype.models.ConformerWav2Vec2PretrainModel": [
+            "conformer_wav2vec2_pretrain_model",
+            "conformer_wav2vec2_pretrain_base",
+            "conformer_wav2vec2_pretrain_large",
+        ],
+        "torchaudio.prototype.models.HiFiGANVocoder": [
+            "hifigan_vocoder",
+            "hifigan_vocoder_v1",
+            "hifigan_vocoder_v2",
+            "hifigan_vocoder_v3",
+        ],
+    }
+-%}
+
+{%- if name in ["ConvEmformer"] %}
+  {{ methods.extend(["infer"]) }}
+{%- endif %}
+
+.. TITLE
+
+{{ name | underline }}
+
+.. CLASS DEFINITIONS
+
+.. autoclass:: {{ fullname }}
+
+Methods
+=======
+
+{% for item in methods %}
+
+{{item | underline("-") }}
+
+.. container:: py attribute
+
+   .. automethod:: {{[fullname, item] | join('.')}}
+
+{%- endfor %}
+
+.. HELPER STRUCTURES
+
+{%- if helpers[fullname] %}
+
+Support Structures
+==================
+
+{%- for item in helpers[fullname] %}
+
+{{item | underline("-") }}
+
+.. container:: py attribute
+
+   .. autodata:: {{["torchaudio.prototype.models", item] | join('.')}}
+      :no-value:
+
+{%- endfor %}
+
+{%- endif %}
+
+.. FACTORY FUNCTIONS
+
+{%- if factory[fullname] %}
+
+Factory Functions
+=================
+
+.. autosummary::
+   :toctree: ../generated
+   :nosignatures:
+
+{% for item in factory[fullname] %}
+   {{["~torchaudio.prototype.models", item] | join('.')}}
+{%- endfor %}
+
+{%- endif %}

docs/source/prototype.models.rst

@@ -4,108 +4,38 @@ torchaudio.prototype.models
 .. py:module:: torchaudio.prototype.models
 .. currentmodule:: torchaudio.prototype.models
 
-conformer_rnnt_model
-~~~~~~~~~~~~~~~~~~~~
 
-.. autofunction:: conformer_rnnt_model
+The ``torchaudio.prototype.models`` subpackage contains definitions of models for addressing common audio tasks.
 
-conformer_rnnt_base
-~~~~~~~~~~~~~~~~~~~
+.. note::
+   For models with pre-trained parameters, please refer to :mod:`torchaudio.prototype.pipelines` module.
 
-.. autofunction:: conformer_rnnt_base
+Model defintions are responsible for constructing computation graphs and executing them.
 
-emformer_hubert_model
-~~~~~~~~~~~~~~~~~~~~~
+Some models have complex structure and variations.
+For such models, factory functions are provided.
 
-.. autofunction:: emformer_hubert_model
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/prototype_model_class.rst
 
-emformer_hubert_base
-~~~~~~~~~~~~~~~~~~~~
+  ConformerWav2Vec2PretrainModel
+  ConvEmformer
+  HiFiGANVocoder
+  SquimObjective
 
-.. autofunction:: emformer_hubert_base
+Prototype Factory Functions of Beta Models
+==========================================
 
-ConvEmformer
-~~~~~~~~~~~~
+.. currentmodule:: torchaudio.models
 
-.. autoclass:: ConvEmformer
+Some model definitions are in beta, but there are new factory functions that are still in prototype. Please check "Prototype Factory Functions" section in each model.
 
-  .. automethod:: forward
+.. autosummary::
+  :toctree: generated
+  :nosignatures:
+  :template: autosummary/model_class.rst
 
-  .. automethod:: infer
-
-ConformerWav2Vec2PretrainModel
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: ConformerWav2Vec2PretrainModel
-
-  .. automethod:: forward
-
-conformer_wav2vec2_model
-~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: conformer_wav2vec2_model
-
-conformer_wav2vec2_base
-~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: conformer_wav2vec2_base
-
-conformer_wav2vec2_pretrain_model
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: conformer_wav2vec2_pretrain_model
-
-conformer_wav2vec2_pretrain_base
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: conformer_wav2vec2_pretrain_base
-
-conformer_wav2vec2_pretrain_large
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: conformer_wav2vec2_pretrain_large
-
-HiFiGANVocoder
-~~~~~~~~~~~~~~
-
-.. autoclass:: HiFiGANVocoder
-
-  .. automethod:: forward
-
-hifigan_vocoder
-~~~~~~~~~~~~~~~
-
-.. autofunction:: hifigan_vocoder
-
-hifigan_vocoder_v1
-~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: hifigan_vocoder_v1
-
-hifigan_vocoder_v2
-~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: hifigan_vocoder_v2
-
-hifigan_vocoder_v3
-~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: hifigan_vocoder_v3
-
-
-SQUIM_OBJECTIVE
-~~~~~~~~~~~~~~~
-
-.. autoclass:: SQUIM_OBJECTIVE
-
-  .. automethod:: forward
-
-squim_objective_model
-~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: squim_objective_model
-
-squim_objective_base
-~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: squim_objective_base
+  Wav2Vec2Model
+  RNNT