Load SortformerEncLabelModel from .nemo file or HuggingFace.
(nemo_path: Optional[str] = None, hf_model: Optional[str] = None)
| 129 | |
| 130 | |
| 131 | def load_model(nemo_path: Optional[str] = None, hf_model: Optional[str] = None): |
| 132 | """Load SortformerEncLabelModel from .nemo file or HuggingFace.""" |
| 133 | from nemo.collections.asr.models import SortformerEncLabelModel |
| 134 | |
| 135 | if nemo_path: |
| 136 | model = SortformerEncLabelModel.restore_from( |
| 137 | nemo_path, map_location="cpu", strict=False |
| 138 | ) |
| 139 | elif hf_model: |
| 140 | model = SortformerEncLabelModel.from_pretrained(hf_model, map_location="cpu") |
| 141 | else: |
| 142 | model = SortformerEncLabelModel.from_pretrained( |
| 143 | "nvidia/diar_streaming_sortformer_4spk-v2", map_location="cpu" |
| 144 | ) |
| 145 | |
| 146 | model.eval() |
| 147 | model.cpu() |
| 148 | return model |
| 149 | |
| 150 | |
| 151 | def _rel_shift_export(self, x): |