diff --git a/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/INSTALLER b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/INSTALLER new file mode 100644 index 0000000000000000000000000000000000000000..a1b589e38a32041e49332e5e81c2d363dc418d68 --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/INSTALLER @@ -0,0 +1 @@ +pip diff --git a/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/METADATA b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/METADATA new file mode 100644 index 0000000000000000000000000000000000000000..0700cabfe978ba3e103a9c43790d698a0072fff8 --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/METADATA @@ -0,0 +1,35 @@ +Metadata-Version: 2.1 +Name: nvidia-cuda-runtime-cu12 +Version: 12.4.127 +Summary: CUDA Runtime native Libraries +Home-page: https://developer.nvidia.com/cuda-zone +Author: Nvidia CUDA Installer Team +Author-email: cuda_installer@nvidia.com +License: NVIDIA Proprietary Software +Keywords: cuda,nvidia,runtime,machine learning,deep learning +Classifier: Development Status :: 4 - Beta +Classifier: Intended Audience :: Developers +Classifier: Intended Audience :: Education +Classifier: Intended Audience :: Science/Research +Classifier: License :: Other/Proprietary License +Classifier: Natural Language :: English +Classifier: Programming Language :: Python :: 3 +Classifier: Programming Language :: Python :: 3.5 +Classifier: Programming Language :: Python :: 3.6 +Classifier: Programming Language :: Python :: 3.7 +Classifier: Programming Language :: Python :: 3.8 +Classifier: Programming Language :: Python :: 3.9 +Classifier: Programming Language :: Python :: 3.10 +Classifier: Programming Language :: Python :: 3.11 +Classifier: Programming Language :: Python :: 3 :: Only +Classifier: Topic :: Scientific/Engineering +Classifier: Topic :: Scientific/Engineering :: Mathematics +Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence +Classifier: Topic :: Software Development +Classifier: Topic :: Software Development :: Libraries +Classifier: Operating System :: Microsoft :: Windows +Classifier: Operating System :: POSIX :: Linux +Requires-Python: >=3 +License-File: License.txt + +CUDA Runtime native Libraries diff --git a/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/RECORD b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/RECORD new file mode 100644 index 0000000000000000000000000000000000000000..31e7830df4181d01a134de21d16a0ffea44c431d --- /dev/null +++ b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/RECORD @@ -0,0 +1,107 @@ +nvidia/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/__pycache__/__init__.cpython-312.pyc,, +nvidia/cuda_runtime/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/cuda_runtime/__pycache__/__init__.cpython-312.pyc,, +nvidia/cuda_runtime/include/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/cuda_runtime/include/__pycache__/__init__.cpython-312.pyc,, +nvidia/cuda_runtime/include/builtin_types.h,sha256=JxT9Vf2q2snxTBOL9ACzNmYzTWACO2VOVUu1KdFt7_g,3150 +nvidia/cuda_runtime/include/channel_descriptor.h,sha256=no_vNky02LeMLI0CF8GDVGHaPm_uRUGcVUMYdt_Xn4U,21482 +nvidia/cuda_runtime/include/common_functions.h,sha256=22LTZRVcPZzEH6MJda7nNMCvMgIjSTe0OKR7sEQj6kc,3410 +nvidia/cuda_runtime/include/cooperative_groups.h,sha256=JUBW-C1x_7WWuNOaoorTKQab0qzrykkG8oAw1mEHZ2s,60332 +nvidia/cuda_runtime/include/cooperative_groups/details/async.h,sha256=xsEHCZP3nuEY3l2p8SU2d1226XiXumUvDP_Gyh8PdVY,19122 +nvidia/cuda_runtime/include/cooperative_groups/details/coalesced_reduce.h,sha256=pBQgFY7i64V87XNATg1UEIQHVNYOItQtHjS5B4yn8pc,4257 +nvidia/cuda_runtime/include/cooperative_groups/details/coalesced_scan.h,sha256=DfZv5d5W0XJv-tZVhgrIdjLjs6aCx_u0oy1lDIpjo1Q,7314 +nvidia/cuda_runtime/include/cooperative_groups/details/driver_abi.h,sha256=v-ZUb4UgGKJk6NR2WCWHD3x_42y-togI1urFn70Gi-g,3964 +nvidia/cuda_runtime/include/cooperative_groups/details/functional.h,sha256=2BV8i8Bidz0kgxuYkJCAbwFxOIZRyzHgG-c_rVKhRzc,8905 +nvidia/cuda_runtime/include/cooperative_groups/details/helpers.h,sha256=K9jvxnXc5-6Fum1KG4EQKJJrVZ4BhHOSAJbZR4uDL0c,26476 +nvidia/cuda_runtime/include/cooperative_groups/details/info.h,sha256=Ij_cqIrcXCcwlaQqCL7AHzMD4H89y0tJeQXCbjTGsFo,12578 +nvidia/cuda_runtime/include/cooperative_groups/details/invoke.h,sha256=Osq3K-tZuXHVCMQJ708PjPo-BwMhjhjApO4b0TYLFJg,8616 +nvidia/cuda_runtime/include/cooperative_groups/details/memory.h,sha256=WU28eUcYLA1z131VYGulR4eVCSN9xK9KSxbV656YPs0,5484 +nvidia/cuda_runtime/include/cooperative_groups/details/partitioning.h,sha256=4UXuvUmZvGANy0hd4erdBNllpgnn4K4qFWWlfzAsHO8,7125 +nvidia/cuda_runtime/include/cooperative_groups/details/reduce.h,sha256=UfMezM5pqRIotJjmuFgOmiMvbu49sYgjraHutmVVr0w,22111 +nvidia/cuda_runtime/include/cooperative_groups/details/scan.h,sha256=-Ttwb2AfEEY_tsmqJjR2dojkPpoRx387SoqxgvfdBtQ,17166 +nvidia/cuda_runtime/include/cooperative_groups/details/sync.h,sha256=zoiBicvB7rlXa_r_VSNuvHVwrLIM7EjF_KdmhvPj1LM,10638 +nvidia/cuda_runtime/include/cooperative_groups/memcpy_async.h,sha256=erOIHuObdfxRhBWfrXE3wsZF4B2GUuqwzQrsPwKPpbg,2960 +nvidia/cuda_runtime/include/cooperative_groups/reduce.h,sha256=B0hgDkqM-6ueqTTgb3b34A0RH4vGz8mBf5e2jT1dJ1o,2949 +nvidia/cuda_runtime/include/cooperative_groups/scan.h,sha256=2EU6T5cWNwftm2B7FicV31PojoI61yo5fHXGRYkGk40,2940 +nvidia/cuda_runtime/include/cuComplex.h,sha256=WpcgpaiPhU_o9sTPMcNTEZuyXDIc8x3sz4dUWSztL2g,12186 +nvidia/cuda_runtime/include/cuda.h,sha256=29OuNnfs8Hb2sqCXHUKy3VudXxzN8050d0oW_C33ysE,1048458 +nvidia/cuda_runtime/include/cudaEGL.h,sha256=_CwaQ4cEP1vfNyBSSd5qFxznPCYOovF6Cpj-QWSIBq4,39544 +nvidia/cuda_runtime/include/cudaEGLTypedefs.h,sha256=xF_FAN1Kar9oyHJ3cCU7jztTpxX8WylpiuYyYpGGHek,5645 +nvidia/cuda_runtime/include/cudaGL.h,sha256=gMT1HPGa-siuji0gAsKYr4X45Lc29HKglC_ttNSGyUM,22501 +nvidia/cuda_runtime/include/cudaGLTypedefs.h,sha256=dClpQI-LuXgF9rPSBsj7OkIg8g_fXDjT0hLZS8TGpOg,6576 +nvidia/cuda_runtime/include/cudaProfilerTypedefs.h,sha256=F2aWLIKv_AhNbxNOaZVcRsxIh0kuscnV8UMWWxkBAlY,3297 +nvidia/cuda_runtime/include/cudaTypedefs.h,sha256=0hWYyV-KM7R5Qjagz9UP1ldhAZDHGIcJmYtYvB_nwNc,110387 +nvidia/cuda_runtime/include/cudaVDPAU.h,sha256=Np7Nc2Wjaz--hkpbhW6f9aapr-NbcPDAgkot0sJerco,12694 +nvidia/cuda_runtime/include/cudaVDPAUTypedefs.h,sha256=wz8nyOUdwM9mH9JO3QZW-A9dyxt-IufSX7nggSXpCNs,4144 +nvidia/cuda_runtime/include/cuda_awbarrier.h,sha256=3ZH-ZlXODhSiwSY9rqSni_EQwi25QMHP6Tm-zOdxBwE,9340 +nvidia/cuda_runtime/include/cuda_awbarrier_helpers.h,sha256=OCskCts5bCKl_RKBe9M74zKSIsVpePn44S_aJp1tFXE,12489 +nvidia/cuda_runtime/include/cuda_awbarrier_primitives.h,sha256=n5__E1jYYDhlgH-f3u8MQjtz57UZ7v5VshhMye1eicM,4699 +nvidia/cuda_runtime/include/cuda_bf16.h,sha256=CUmV313-EoESxCNDiAC1azacAWxzzVWj7vE7bZNoebI,188194 +nvidia/cuda_runtime/include/cuda_bf16.hpp,sha256=Z2hjeJHZMwa-VQCgd5BNsJjo9rRfK2wqJH942HTVrQE,134181 +nvidia/cuda_runtime/include/cuda_device_runtime_api.h,sha256=bIhfusirXe5-osOTPAILDh6pY8MW1hefyZvTD_IzgqM,46249 +nvidia/cuda_runtime/include/cuda_egl_interop.h,sha256=PNWYns30MIytJQHSOh7UbZYlaTX5e0bavzK14tde_C8,37109 +nvidia/cuda_runtime/include/cuda_fp16.h,sha256=4yaQaQw0__T0TOOr37gtz-l1zxOc8nH-qhJ6RSk30OU,178258 +nvidia/cuda_runtime/include/cuda_fp16.hpp,sha256=wgSLe8EsVPul3XEpm4kD0k9jmmxlyHNiiZPpGpDU8dQ,119751 +nvidia/cuda_runtime/include/cuda_fp8.h,sha256=Q3OP5o_3rSYbKtVIlcXVr_CncU3SPM-09j605e2Zegw,13833 +nvidia/cuda_runtime/include/cuda_fp8.hpp,sha256=b-PcyZgei5MmIp6op0QQ40BgNupO_ei648hG_dUS-FQ,64246 +nvidia/cuda_runtime/include/cuda_gl_interop.h,sha256=VQEswFeOBF6JN6Q0pdlkvc5WT7bD1FnTfKewvANulCc,19150 +nvidia/cuda_runtime/include/cuda_occupancy.h,sha256=Kr9HyOe-hlRjBAzbINwUYkNgbbIgIjuvKs09UZhMYQo,67179 +nvidia/cuda_runtime/include/cuda_pipeline.h,sha256=0enXG49wN4JajlQi3ahbp2ei_ufTY_Mznic7zfWmKHM,8130 +nvidia/cuda_runtime/include/cuda_pipeline_helpers.h,sha256=bo1L7e6vCuM-K3Il8K1z4wJUja5DyXQKdo_hSWUME-E,13852 +nvidia/cuda_runtime/include/cuda_pipeline_primitives.h,sha256=FnJJtuV6rHr6LgL56XDwilcSbFr6W1Hj6mf1AJaMI20,8675 +nvidia/cuda_runtime/include/cuda_runtime.h,sha256=a-OXWPsmKSPst7mRCCxHNZV7m-uRLCAY8oGRi-dJzPA,90683 +nvidia/cuda_runtime/include/cuda_runtime_api.h,sha256=7Ys9yv_2trFEVybtbh-UJKnDKG8fHWvUjSX4cgZGCck,608580 +nvidia/cuda_runtime/include/cuda_surface_types.h,sha256=Mw5Lo4b8Q-f9mogOvATGyHhu9d2t2K6XOxuqtZrSh3A,3688 +nvidia/cuda_runtime/include/cuda_texture_types.h,sha256=ITbX-JNnP7Rm-JSgNVdJ9pq6k8FVor8RbnruDsKq6sk,3688 +nvidia/cuda_runtime/include/cuda_vdpau_interop.h,sha256=bXQanWc2IFXZAKWNGl2xAz9nLvFmQpWyGrsDvfeS9FA,7727 +nvidia/cuda_runtime/include/cudart_platform.h,sha256=YN6sKhB0b9w5tGX1IYL7ulJVPrWAiX9A44qLv4EtW5Q,2717 +nvidia/cuda_runtime/include/device_atomic_functions.h,sha256=OR2jNSfSKzaFri74zh4Vtz5M0z9UDBU3rKeC1rYaVQs,9500 +nvidia/cuda_runtime/include/device_atomic_functions.hpp,sha256=0e7MOiNNUnnloXpB_r9WT5YOws5cxgzQQAzRCYvgaFA,10486 +nvidia/cuda_runtime/include/device_double_functions.h,sha256=KUxId5Z1fx8SWfLRTxPD7RB-zN7zslzb4n7JaJLfL3I,3452 +nvidia/cuda_runtime/include/device_functions.h,sha256=bWSrhTYE9NQlss7xMSMEVusvto9j2fgUDXWVH2W_cOA,3410 +nvidia/cuda_runtime/include/device_launch_parameters.h,sha256=H1_CC-vvAaS26ys4XsTFkMgTxUTciAjdjswjizkisvQ,3846 +nvidia/cuda_runtime/include/device_types.h,sha256=2LFxoZBJPoA5V0H1EbKTEaXDi3GDJPtzOPdRHDaucIQ,3588 +nvidia/cuda_runtime/include/driver_functions.h,sha256=cN3IjRAz2Mj2Pj35SyxJIkZNDDusnJqaqzBdMzpQKbA,4625 +nvidia/cuda_runtime/include/driver_types.h,sha256=4eBQ10Nzgfs2BlxGaGHVMWLvnJfKrEnMml9zfFi0DyA,177782 +nvidia/cuda_runtime/include/host_config.h,sha256=BscH_GazAZbbotddVzL5RmafbQ-QjRx8f-I1O01IBW8,3380 +nvidia/cuda_runtime/include/host_defines.h,sha256=bBQwQF5C1N1c2qpLV56g1c-weu9Ysgz-gIf2Kn3uz_A,3386 +nvidia/cuda_runtime/include/library_types.h,sha256=p6746aCd_A_1VlgKRhLJChzeZ4tN7e4HBH2Hm7hDjbU,4836 +nvidia/cuda_runtime/include/math_constants.h,sha256=cV6hAyQe8X7f7MBtaKjjIJq3BycOUDp6I5cizJX5HLw,7608 +nvidia/cuda_runtime/include/math_functions.h,sha256=5XcC6j-fJKttvhwc4hZNoLHNw808a2ZYIOtZ7ry7yd0,3398 +nvidia/cuda_runtime/include/mma.h,sha256=IY_VenxuEncwGq92MhrWUb-Xswh0ekAXLy9Rbxhxa2Y,2932 +nvidia/cuda_runtime/include/sm_20_atomic_functions.h,sha256=x4ycINVq__l9B4SQPD-I48jQbKxxdBmgp8Vf2GO0Qfg,4478 +nvidia/cuda_runtime/include/sm_20_atomic_functions.hpp,sha256=1l5NLM8DhDbqYZ_E51LoqElQJXObkbwo57d3r-4uEbE,4107 +nvidia/cuda_runtime/include/sm_20_intrinsics.h,sha256=a4jDSp_DUW0d09g5wgEm_I7bGTAe73HKRinkhBKQBis,51048 +nvidia/cuda_runtime/include/sm_20_intrinsics.hpp,sha256=BhEBuXSKBsNGJDBJDtYL0cGRI3wX_w_OIgA5D-YxIWk,7694 +nvidia/cuda_runtime/include/sm_30_intrinsics.h,sha256=b6W8Vxp9vD9OCJI6lZuGyZYXEdQ3Ei8PTAloHNkwCcQ,16978 +nvidia/cuda_runtime/include/sm_30_intrinsics.hpp,sha256=yX0ebd265tJ-BDhvluP2BhadPuWXpRZPI2eeQFFt5ys,24567 +nvidia/cuda_runtime/include/sm_32_atomic_functions.h,sha256=HGnZgQHACE2AAb6zabGUURc53IsVZelc2BSJqvs9OgY,5703 +nvidia/cuda_runtime/include/sm_32_atomic_functions.hpp,sha256=CQTTvOEYp-s5hqAgLvAon11vLYDrDp8cTHdel-XRzBQ,6592 +nvidia/cuda_runtime/include/sm_32_intrinsics.h,sha256=Xdkogdsjy1vh8u3eGu0i5xTmHxBGAjj6_vVGR-spdOE,33539 +nvidia/cuda_runtime/include/sm_32_intrinsics.hpp,sha256=Gl8aSLDLcit4W3pKQS19GsDG8RYcwD65HwYB_CeZe8M,70616 +nvidia/cuda_runtime/include/sm_35_atomic_functions.h,sha256=a3XoEsKRCEOf0Q_5Y__rMfmC4pScv4VkUggVgVJVn44,2909 +nvidia/cuda_runtime/include/sm_35_intrinsics.h,sha256=0mS5-LCgvZiTvL7-MG_4YwI-zWGvM-s4xyRuMkunMC8,2664 +nvidia/cuda_runtime/include/sm_60_atomic_functions.h,sha256=_anfNaJsvQpDEorYeUKIkbizYkwrinBcG_ZCiECtLqI,13178 +nvidia/cuda_runtime/include/sm_60_atomic_functions.hpp,sha256=cgIKddDn2B3QzYlzeBILAP1IRys74QCCxsH0QqaVGls,22903 +nvidia/cuda_runtime/include/sm_61_intrinsics.h,sha256=h_MBL1UUDxQX_qOddSImzqyFjcrhhm_63G97pGDyreU,10902 +nvidia/cuda_runtime/include/sm_61_intrinsics.hpp,sha256=N-nQvcBsPMT2Umy5zR69c9K1q366W-Jqe7NpoLTqTmg,6787 +nvidia/cuda_runtime/include/surface_functions.h,sha256=b1O82SAvEgWWxA9uZTWQcGimzZUoem2QbAET3wh3fZc,6782 +nvidia/cuda_runtime/include/surface_indirect_functions.h,sha256=vy9QuFVV-ezZP-x2RT9RLp2qIUgdngACOCmalSfVFPA,10877 +nvidia/cuda_runtime/include/surface_types.h,sha256=XkFXD1nHbeSMgajR-UJE9uQ7TByzJnjdnUL4-yGiufk,4530 +nvidia/cuda_runtime/include/texture_fetch_functions.h,sha256=KLCmUxf5aY5_UalX8tSFB6e4TrjA8hyUPxLOkMFltAo,12468 +nvidia/cuda_runtime/include/texture_indirect_functions.h,sha256=lH_y3Ni-hq4RZ0_PMFbBM0th5-OmTn3TtqtpkHHhA8w,21163 +nvidia/cuda_runtime/include/texture_types.h,sha256=73ntVyg8r8fzKy5VIk6yuvC45GDeWepaLIqIk-M3Ri8,6360 +nvidia/cuda_runtime/include/vector_functions.h,sha256=WypGkL-IDbGOlay7g_G0p3HO7OLGRE0Do__JtiFoWxY,8003 +nvidia/cuda_runtime/include/vector_functions.hpp,sha256=afXhNSd3LFTZo96EPtesTLfvxd4nTmLVzgkj967rTRg,10060 +nvidia/cuda_runtime/include/vector_types.h,sha256=6CJ4yt3KD7zQVfm1NhrgqNYYEDEIZWwaivlFx12nhNg,13396 +nvidia/cuda_runtime/lib/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia/cuda_runtime/lib/__pycache__/__init__.cpython-312.pyc,, +nvidia/cuda_runtime/lib/libcudart.so.12,sha256=h3QiT1sRpzsV0HSj_M5zJzIsXEz9_ZJNaoJnee7JaP4,707904 +nvidia_cuda_runtime_cu12-12.4.127.dist-info/INSTALLER,sha256=zuuue4knoyJ-UwPPXg8fezS7VCrXJQrAP7zeNuwvFQg,4 +nvidia_cuda_runtime_cu12-12.4.127.dist-info/License.txt,sha256=rW9YU_ugyg0VnQ9Y1JrkmDDC-Mk_epJki5zpCttMbM0,59262 +nvidia_cuda_runtime_cu12-12.4.127.dist-info/METADATA,sha256=AcAo5k7tWjg-wgkViCrnv6J1v-QkwuViwNkUy8Pfk_0,1507 +nvidia_cuda_runtime_cu12-12.4.127.dist-info/RECORD,, +nvidia_cuda_runtime_cu12-12.4.127.dist-info/REQUESTED,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0 +nvidia_cuda_runtime_cu12-12.4.127.dist-info/WHEEL,sha256=XDTs3wIbcE-BcRO08VJlZpA6z9OaC1mOKPCGGGwuM2g,109 +nvidia_cuda_runtime_cu12-12.4.127.dist-info/top_level.txt,sha256=fTkAtiFuL16nUrB9ytDDtpytz2t0B4NvYTnRzwAhO14,7 diff --git a/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/REQUESTED b/lib/python3.12/site-packages/nvidia_cuda_runtime_cu12-12.4.127.dist-info/REQUESTED new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/sniffio/__init__.py b/lib/python3.12/site-packages/sniffio/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..63f2f19e409c7f4b3c6c064022e8f104227873aa --- /dev/null +++ b/lib/python3.12/site-packages/sniffio/__init__.py @@ -0,0 +1,17 @@ +"""Top-level package for sniffio.""" + +__all__ = [ + "current_async_library", + "AsyncLibraryNotFoundError", + "current_async_library_cvar", + "thread_local", +] + +from ._version import __version__ + +from ._impl import ( + current_async_library, + AsyncLibraryNotFoundError, + current_async_library_cvar, + thread_local, +) diff --git a/lib/python3.12/site-packages/sniffio/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sniffio/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7ab5a63418d9698486108bfe3b5f9c5a95d755d6 Binary files /dev/null and b/lib/python3.12/site-packages/sniffio/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sniffio/__pycache__/_impl.cpython-312.pyc b/lib/python3.12/site-packages/sniffio/__pycache__/_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e6762aec4df9678972add2b64e5a4a29aecf0405 Binary files /dev/null and b/lib/python3.12/site-packages/sniffio/__pycache__/_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sniffio/__pycache__/_version.cpython-312.pyc b/lib/python3.12/site-packages/sniffio/__pycache__/_version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7dfc37ff69af435e090ba4b72474ae91a337666b Binary files /dev/null and b/lib/python3.12/site-packages/sniffio/__pycache__/_version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sniffio/_impl.py b/lib/python3.12/site-packages/sniffio/_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..c1a7bbf218ba985b87cd1d9b23da69222894c1dd --- /dev/null +++ b/lib/python3.12/site-packages/sniffio/_impl.py @@ -0,0 +1,95 @@ +from contextvars import ContextVar +from typing import Optional +import sys +import threading + +current_async_library_cvar = ContextVar( + "current_async_library_cvar", default=None +) # type: ContextVar[Optional[str]] + + +class _ThreadLocal(threading.local): + # Since threading.local provides no explicit mechanism is for setting + # a default for a value, a custom class with a class attribute is used + # instead. + name = None # type: Optional[str] + + +thread_local = _ThreadLocal() + + +class AsyncLibraryNotFoundError(RuntimeError): + pass + + +def current_async_library() -> str: + """Detect which async library is currently running. + + The following libraries are currently supported: + + ================ =========== ============================ + Library Requires Magic string + ================ =========== ============================ + **Trio** Trio v0.6+ ``"trio"`` + **Curio** - ``"curio"`` + **asyncio** ``"asyncio"`` + **Trio-asyncio** v0.8.2+ ``"trio"`` or ``"asyncio"``, + depending on current mode + ================ =========== ============================ + + Returns: + A string like ``"trio"``. + + Raises: + AsyncLibraryNotFoundError: if called from synchronous context, + or if the current async library was not recognized. + + Examples: + + .. code-block:: python3 + + from sniffio import current_async_library + + async def generic_sleep(seconds): + library = current_async_library() + if library == "trio": + import trio + await trio.sleep(seconds) + elif library == "asyncio": + import asyncio + await asyncio.sleep(seconds) + # ... and so on ... + else: + raise RuntimeError(f"Unsupported library {library!r}") + + """ + value = thread_local.name + if value is not None: + return value + + value = current_async_library_cvar.get() + if value is not None: + return value + + # Need to sniff for asyncio + if "asyncio" in sys.modules: + import asyncio + try: + current_task = asyncio.current_task # type: ignore[attr-defined] + except AttributeError: + current_task = asyncio.Task.current_task # type: ignore[attr-defined] + try: + if current_task() is not None: + return "asyncio" + except RuntimeError: + pass + + # Sniff for curio (for now) + if 'curio' in sys.modules: + from curio.meta import curio_running + if curio_running(): + return 'curio' + + raise AsyncLibraryNotFoundError( + "unknown async library, or not in async context" + ) diff --git a/lib/python3.12/site-packages/sniffio/_tests/__init__.py b/lib/python3.12/site-packages/sniffio/_tests/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/sniffio/_tests/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/sniffio/_tests/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ea6bd9a778937bce40d98c5d60adc9ab87acc44a Binary files /dev/null and b/lib/python3.12/site-packages/sniffio/_tests/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sniffio/_tests/__pycache__/test_sniffio.cpython-312.pyc b/lib/python3.12/site-packages/sniffio/_tests/__pycache__/test_sniffio.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cedef0d0e284ddcc5b0782b3f05a6f77e23737b1 Binary files /dev/null and b/lib/python3.12/site-packages/sniffio/_tests/__pycache__/test_sniffio.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/sniffio/_tests/test_sniffio.py b/lib/python3.12/site-packages/sniffio/_tests/test_sniffio.py new file mode 100644 index 0000000000000000000000000000000000000000..02945a947527a9bfb396d0cd35ffee22eda665bb --- /dev/null +++ b/lib/python3.12/site-packages/sniffio/_tests/test_sniffio.py @@ -0,0 +1,84 @@ +import os +import sys + +import pytest + +from .. import ( + current_async_library, AsyncLibraryNotFoundError, + current_async_library_cvar, thread_local +) + + +def test_basics_cvar(): + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + token = current_async_library_cvar.set("generic-lib") + try: + assert current_async_library() == "generic-lib" + finally: + current_async_library_cvar.reset(token) + + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + +def test_basics_tlocal(): + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + old_name, thread_local.name = thread_local.name, "generic-lib" + try: + assert current_async_library() == "generic-lib" + finally: + thread_local.name = old_name + + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + +def test_asyncio(): + import asyncio + + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + ran = [] + + async def this_is_asyncio(): + assert current_async_library() == "asyncio" + # Call it a second time to exercise the caching logic + assert current_async_library() == "asyncio" + ran.append(True) + + asyncio.run(this_is_asyncio()) + assert ran == [True] + + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + +@pytest.mark.skipif( + sys.version_info >= (3, 12), + reason= + "curio broken on 3.12 (https://github.com/python-trio/sniffio/pull/42)", +) +def test_curio(): + import curio + + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() + + ran = [] + + async def this_is_curio(): + assert current_async_library() == "curio" + # Call it a second time to exercise the caching logic + assert current_async_library() == "curio" + ran.append(True) + + curio.run(this_is_curio) + assert ran == [True] + + with pytest.raises(AsyncLibraryNotFoundError): + current_async_library() diff --git a/lib/python3.12/site-packages/sniffio/_version.py b/lib/python3.12/site-packages/sniffio/_version.py new file mode 100644 index 0000000000000000000000000000000000000000..0495d10545c9fd515ed51e890309d2b66e2c30bb --- /dev/null +++ b/lib/python3.12/site-packages/sniffio/_version.py @@ -0,0 +1,3 @@ +# This file is imported from __init__.py and exec'd from setup.py + +__version__ = "1.3.1" diff --git a/lib/python3.12/site-packages/sniffio/py.typed b/lib/python3.12/site-packages/sniffio/py.typed new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchaudio/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1f8c1c595439d5863adfe0139e8105cf8c794934 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/__pycache__/kaldi_io.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/__pycache__/kaldi_io.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1a570550a89aeceec5264c2831ecf9bae08e98f Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/__pycache__/kaldi_io.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/__pycache__/version.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/__pycache__/version.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..131227f0dad365146e76ac12ea64e4b1c454e105 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/__pycache__/version.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a7886f162d444b2ab93ec070025d3a0db28b208d Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/_no_backend.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/_no_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..b437928be20e04ca1be6920ff45b99c98d605646 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/_no_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/_sox_io_backend.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/_sox_io_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..69ee9549ae5bad2f5836cc02773dc7add73e484b Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/_sox_io_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/common.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/common.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..18c90579946bbd6412433b2fc46955301384ae56 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/common.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/no_backend.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/no_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8adedbebab88a3102acb5b80be04ebefa3269361 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/no_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/soundfile_backend.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/soundfile_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9d0f818b6562a50cd252f08dfe469bef765464a Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/soundfile_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/__pycache__/sox_io_backend.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/sox_io_backend.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..24c6544a483e297569947ddb26fdea50829be4e9 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/backend/__pycache__/sox_io_backend.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/backend/no_backend.py b/lib/python3.12/site-packages/torchaudio/backend/no_backend.py new file mode 100644 index 0000000000000000000000000000000000000000..2645a86bc80538fa0522f5eb80e304881f30acc7 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/backend/no_backend.py @@ -0,0 +1,14 @@ +def __getattr__(name: str): + import warnings + + warnings.warn( + "Torchaudio's I/O functions now support par-call bakcend dispatch. " + "Importing backend implementation directly is no longer guaranteed to work. " + "Please use `backend` keyword with load/save/info function, instead of " + "calling the udnerlying implementation directly.", + stacklevel=2, + ) + + from . import _no_backend + + return getattr(_no_backend, name) diff --git a/lib/python3.12/site-packages/torchaudio/io/__init__.py b/lib/python3.12/site-packages/torchaudio/io/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..5486424dd21a0f1aeb8ea2cead737309fb7f9a7a --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/io/__init__.py @@ -0,0 +1,13 @@ +from torio.io import CodecConfig, StreamingMediaDecoder as StreamReader, StreamingMediaEncoder as StreamWriter + +from ._effector import AudioEffector +from ._playback import play_audio + + +__all__ = [ + "AudioEffector", + "StreamReader", + "StreamWriter", + "CodecConfig", + "play_audio", +] diff --git a/lib/python3.12/site-packages/torchaudio/io/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/io/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..dac9095e3d95b346a1125690dcf23d0eda3f078e Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/io/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/io/__pycache__/_effector.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/io/__pycache__/_effector.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cc287181e56c7d85da153f38e6c9c2539832038b Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/io/__pycache__/_effector.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/io/__pycache__/_playback.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/io/__pycache__/_playback.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a3b36ba18198d3fa9d44658d2705583ee026e2ad Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/io/__pycache__/_playback.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/io/_effector.py b/lib/python3.12/site-packages/torchaudio/io/_effector.py new file mode 100644 index 0000000000000000000000000000000000000000..74255684c8fa75789e88fc224bcdac12aa1b29cf --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/io/_effector.py @@ -0,0 +1,347 @@ +import io +from typing import Iterator, List, Optional + +import torch +from torch import Tensor + +from torio.io._streaming_media_decoder import _get_afilter_desc, StreamingMediaDecoder as StreamReader +from torio.io._streaming_media_encoder import CodecConfig, StreamingMediaEncoder as StreamWriter + + +class _StreamingIOBuffer: + """Streaming Bytes IO buffer. Data are dropped when read.""" + + def __init__(self): + self._buffer: List(bytes) = [] + + def write(self, b: bytes): + if b: + self._buffer.append(b) + return len(b) + + def pop(self, n): + """Pop the oldest byte string. It does not necessary return the requested amount""" + if not self._buffer: + return b"" + if len(self._buffer[0]) <= n: + return self._buffer.pop(0) + ret = self._buffer[0][:n] + self._buffer[0] = self._buffer[0][n:] + return ret + + +def _get_sample_fmt(dtype: torch.dtype): + types = { + torch.uint8: "u8", + torch.int16: "s16", + torch.int32: "s32", + torch.float32: "flt", + torch.float64: "dbl", + } + if dtype not in types: + raise ValueError(f"Unsupported dtype is provided {dtype}. Supported dtypes are: {types.keys()}") + return types[dtype] + + +class _AudioStreamingEncoder: + """Given a waveform, encode on-demand and return bytes""" + + def __init__( + self, + src: Tensor, + sample_rate: int, + effect: str, + muxer: str, + encoder: Optional[str], + codec_config: Optional[CodecConfig], + frames_per_chunk: int, + ): + self.src = src + self.buffer = _StreamingIOBuffer() + self.writer = StreamWriter(self.buffer, format=muxer) + self.writer.add_audio_stream( + num_channels=src.size(1), + sample_rate=sample_rate, + format=_get_sample_fmt(src.dtype), + encoder=encoder, + filter_desc=effect, + codec_config=codec_config, + ) + self.writer.open() + self.fpc = frames_per_chunk + + # index on the input tensor (along time-axis) + # we use -1 to indicate that we finished iterating the tensor and + # the writer is closed. + self.i_iter = 0 + + def read(self, n): + while not self.buffer._buffer and self.i_iter >= 0: + self.writer.write_audio_chunk(0, self.src[self.i_iter : self.i_iter + self.fpc]) + self.i_iter += self.fpc + if self.i_iter >= self.src.size(0): + self.writer.flush() + self.writer.close() + self.i_iter = -1 + return self.buffer.pop(n) + + +def _encode( + src: Tensor, + sample_rate: int, + effect: str, + muxer: str, + encoder: Optional[str], + codec_config: Optional[CodecConfig], +): + buffer = io.BytesIO() + writer = StreamWriter(buffer, format=muxer) + writer.add_audio_stream( + num_channels=src.size(1), + sample_rate=sample_rate, + format=_get_sample_fmt(src.dtype), + encoder=encoder, + filter_desc=effect, + codec_config=codec_config, + ) + with writer.open(): + writer.write_audio_chunk(0, src) + buffer.seek(0) + return buffer + + +def _get_muxer(dtype: torch.dtype): + # TODO: check if this works in Windows. + types = { + torch.uint8: "u8", + torch.int16: "s16le", + torch.int32: "s32le", + torch.float32: "f32le", + torch.float64: "f64le", + } + if dtype not in types: + raise ValueError(f"Unsupported dtype is provided {dtype}. Supported dtypes are: {types.keys()}") + return types[dtype] + + +class AudioEffector: + """Apply various filters and/or codecs to waveforms. + + .. versionadded:: 2.1 + + Args: + effect (str or None, optional): Filter expressions or ``None`` to apply no filter. + See https://ffmpeg.org/ffmpeg-filters.html#Audio-Filters for the + details of filter syntax. + + format (str or None, optional): When provided, encode the audio into the + corresponding format. Default: ``None``. + + encoder (str or None, optional): When provided, override the encoder used + by the ``format``. Default: ``None``. + + codec_config (CodecConfig or None, optional): When provided, configure the encoding codec. + Should be provided in conjunction with ``format`` option. + + pad_end (bool, optional): When enabled, and if the waveform becomes shorter after applying + effects/codec, then pad the end with silence. + + Example - Basic usage + To use ``AudioEffector``, first instantiate it with a set of + ``effect`` and ``format``. + + >>> # instantiate the effector + >>> effector = AudioEffector(effect=..., format=...) + + Then, use :py:meth:`~AudioEffector.apply` or :py:meth:`~AudioEffector.stream` + method to apply them. + + >>> # Apply the effect to the whole waveform + >>> applied = effector.apply(waveform, sample_rate) + + >>> # Apply the effect chunk-by-chunk + >>> for chunk in effector.stream(waveform, sample_rate): + >>> ... + + Example - Applying effects + Please refer to + https://ffmpeg.org/ffmpeg-filters.html#Filtergraph-description + for the overview of filter description, and + https://ffmpeg.org/ffmpeg-filters.html#toc-Audio-Filters + for the list of available filters. + + Tempo - https://ffmpeg.org/ffmpeg-filters.html#atempo + + >>> AudioEffector(effect="atempo=1.5") + + Echo - https://ffmpeg.org/ffmpeg-filters.html#aecho + + >>> AudioEffector(effect="aecho=0.8:0.88:60:0.4") + + Flanger - https://ffmpeg.org/ffmpeg-filters.html#flanger + + >>> AudioEffector(effect="aflanger") + + Vibrato - https://ffmpeg.org/ffmpeg-filters.html#vibrato + + >>> AudioEffector(effect="vibrato") + + Tremolo - https://ffmpeg.org/ffmpeg-filters.html#tremolo + + >>> AudioEffector(effect="vibrato") + + You can also apply multiple effects at once. + + >>> AudioEffector(effect="") + + Example - Applying codec + One can apply codec using ``format`` argument. ``format`` can be + audio format or container format. If the container format supports + multiple encoders, you can specify it with ``encoder`` argument. + + Wav format + (no compression is applied but samples are converted to + 16-bit signed integer) + + >>> AudioEffector(format="wav") + + Ogg format with default encoder + + >>> AudioEffector(format="ogg") + + Ogg format with vorbis + + >>> AudioEffector(format="ogg", encoder="vorbis") + + Ogg format with opus + + >>> AudioEffector(format="ogg", encoder="opus") + + Webm format with opus + + >>> AudioEffector(format="webm", encoder="opus") + + Example - Applying codec with configuration + Reference: https://trac.ffmpeg.org/wiki/Encode/MP3 + + MP3 with default config + + >>> AudioEffector(format="mp3") + + MP3 with variable bitrate + + >>> AudioEffector(format="mp3", codec_config=CodecConfig(qscale=5)) + + MP3 with constant bitrate + + >>> AudioEffector(format="mp3", codec_config=CodecConfig(bit_rate=32_000)) + """ + + def __init__( + self, + effect: Optional[str] = None, + format: Optional[str] = None, + *, + encoder: Optional[str] = None, + codec_config: Optional[CodecConfig] = None, + pad_end: bool = True, + ): + if format is None: + if encoder is not None or codec_config is not None: + raise ValueError("`encoder` and/or `condec_config` opions are provided without `format` option.") + self.effect = effect + self.format = format + self.encoder = encoder + self.codec_config = codec_config + self.pad_end = pad_end + + def _get_reader(self, waveform, sample_rate, output_sample_rate, frames_per_chunk=None): + num_frames, num_channels = waveform.shape + + if self.format is not None: + muxer = self.format + encoder = self.encoder + option = {} + # Some formats are headerless, so need to provide these infomation. + if self.format == "mulaw": + option = {"sample_rate": f"{sample_rate}", "channels": f"{num_channels}"} + + else: # PCM + muxer = _get_muxer(waveform.dtype) + encoder = None + option = {"sample_rate": f"{sample_rate}", "channels": f"{num_channels}"} + + if frames_per_chunk is None: + src = _encode(waveform, sample_rate, self.effect, muxer, encoder, self.codec_config) + else: + src = _AudioStreamingEncoder( + waveform, sample_rate, self.effect, muxer, encoder, self.codec_config, frames_per_chunk + ) + + output_sr = sample_rate if output_sample_rate is None else output_sample_rate + filter_desc = _get_afilter_desc(output_sr, _get_sample_fmt(waveform.dtype), num_channels) + if self.pad_end: + filter_desc = f"{filter_desc},apad=whole_len={num_frames}" + + reader = StreamReader(src, format=muxer, option=option) + reader.add_audio_stream(frames_per_chunk or -1, -1, filter_desc=filter_desc) + return reader + + def apply(self, waveform: Tensor, sample_rate: int, output_sample_rate: Optional[int] = None) -> Tensor: + """Apply the effect and/or codecs to the whole tensor. + + Args: + waveform (Tensor): The input waveform. Shape: ``(time, channel)`` + sample_rate (int): Sample rate of the input waveform. + output_sample_rate (int or None, optional): Output sample rate. + If provided, override the output sample rate. + Otherwise, the resulting tensor is resampled to have + the same sample rate as the input. + Default: ``None``. + + Returns: + Tensor: + Resulting Tensor. Shape: ``(time, channel)``. The number of frames + could be different from that of the input. + """ + if waveform.ndim != 2: + raise ValueError(f"Expected the input waveform to be 2D. Found: {waveform.ndim}") + + if waveform.numel() == 0: + return waveform + + reader = self._get_reader(waveform, sample_rate, output_sample_rate) + reader.process_all_packets() + (applied,) = reader.pop_chunks() + return Tensor(applied) + + def stream( + self, waveform: Tensor, sample_rate: int, frames_per_chunk: int, output_sample_rate: Optional[int] = None + ) -> Iterator[Tensor]: + """Apply the effect and/or codecs to the given tensor chunk by chunk. + + Args: + waveform (Tensor): The input waveform. Shape: ``(time, channel)`` + sample_rate (int): Sample rate of the waveform. + frames_per_chunk (int): The number of frames to return at a time. + output_sample_rate (int or None, optional): Output sample rate. + If provided, override the output sample rate. + Otherwise, the resulting tensor is resampled to have + the same sample rate as the input. + Default: ``None``. + + Returns: + Iterator[Tensor]: + Series of processed chunks. Shape: ``(time, channel)``, where the + the number of frames matches ``frames_per_chunk`` except the + last chunk, which could be shorter. + """ + if waveform.ndim != 2: + raise ValueError(f"Expected the input waveform to be 2D. Found: {waveform.ndim}") + + if waveform.numel() == 0: + return waveform + + reader = self._get_reader(waveform, sample_rate, output_sample_rate, frames_per_chunk) + for (applied,) in reader.stream(): + yield Tensor(applied) diff --git a/lib/python3.12/site-packages/torchaudio/io/_playback.py b/lib/python3.12/site-packages/torchaudio/io/_playback.py new file mode 100644 index 0000000000000000000000000000000000000000..7183ee3ba8cad8e842d066f6aaa9067687b9476b --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/io/_playback.py @@ -0,0 +1,72 @@ +import warnings +from sys import platform +from typing import Optional + +import torch +import torchaudio + +dict_format = { + torch.uint8: "u8", + torch.int16: "s16", + torch.int32: "s32", + torch.int64: "s64", + torch.float32: "flt", + torch.float64: "dbl", +} + + +def play_audio( + waveform: torch.Tensor, + sample_rate: Optional[float], + device: Optional[str] = None, +) -> None: + """Plays audio through specified or available output device. + + .. warning:: + This function is currently only supported on MacOS, and requires + libavdevice (FFmpeg) with ``audiotoolbox`` output device. + + .. note:: + This function can play up to two audio channels. + + Args: + waveform: Tensor containing the audio to play. + Expected shape: `(time, num_channels)`. + sample_rate: Sample rate of the audio to play. + device: Output device to use. If None, the default device is used. + """ + + if platform == "darwin": + device = device or "audiotoolbox" + path = "-" + else: + raise ValueError(f"This function only supports MacOS, but current OS is {platform}") + + available_devices = list(torchaudio.utils.ffmpeg_utils.get_output_devices().keys()) + if device not in available_devices: + raise ValueError(f"Device {device} is not available. Available devices are: {available_devices}") + + if waveform.dtype not in dict_format: + raise ValueError(f"Unsupported type {waveform.dtype}. The list of supported types is: {dict_format.keys()}") + format = dict_format[waveform.dtype] + + if waveform.ndim != 2: + raise ValueError(f"Expected 2D tensor with shape `(time, num_channels)`, got {waveform.ndim}D tensor instead") + + time, num_channels = waveform.size() + if num_channels > 2: + warnings.warn( + f"Expected up to 2 channels, got {num_channels} channels instead. " + "Only the first 2 channels will be played.", + stacklevel=2, + ) + + # Write to speaker device + s = torchaudio.io.StreamWriter(dst=path, format=device) + s.add_audio_stream(sample_rate, num_channels, format=format) + + # write audio to the device + block_size = 256 + with s.open(): + for i in range(0, time, block_size): + s.write_audio_chunk(0, waveform[i : i + block_size, :]) diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/__init__.py b/lib/python3.12/site-packages/torchaudio/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..efec1f3521e760803e095efb71f164ed268896f1 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/__init__.py @@ -0,0 +1,102 @@ +from ._source_separation_pipeline import ( + CONVTASNET_BASE_LIBRI2MIX, + HDEMUCS_HIGH_MUSDB, + HDEMUCS_HIGH_MUSDB_PLUS, + SourceSeparationBundle, +) +from ._squim_pipeline import SQUIM_OBJECTIVE, SQUIM_SUBJECTIVE, SquimObjectiveBundle, SquimSubjectiveBundle +from ._tts import ( + TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH, + TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH, + TACOTRON2_WAVERNN_CHAR_LJSPEECH, + TACOTRON2_WAVERNN_PHONE_LJSPEECH, + Tacotron2TTSBundle, +) +from ._wav2vec2.impl import ( + HUBERT_ASR_LARGE, + HUBERT_ASR_XLARGE, + HUBERT_BASE, + HUBERT_LARGE, + HUBERT_XLARGE, + MMS_FA, + VOXPOPULI_ASR_BASE_10K_DE, + VOXPOPULI_ASR_BASE_10K_EN, + VOXPOPULI_ASR_BASE_10K_ES, + VOXPOPULI_ASR_BASE_10K_FR, + VOXPOPULI_ASR_BASE_10K_IT, + WAV2VEC2_ASR_BASE_100H, + WAV2VEC2_ASR_BASE_10M, + WAV2VEC2_ASR_BASE_960H, + WAV2VEC2_ASR_LARGE_100H, + WAV2VEC2_ASR_LARGE_10M, + WAV2VEC2_ASR_LARGE_960H, + WAV2VEC2_ASR_LARGE_LV60K_100H, + WAV2VEC2_ASR_LARGE_LV60K_10M, + WAV2VEC2_ASR_LARGE_LV60K_960H, + WAV2VEC2_BASE, + WAV2VEC2_LARGE, + WAV2VEC2_LARGE_LV60K, + WAV2VEC2_XLSR53, + WAV2VEC2_XLSR_1B, + WAV2VEC2_XLSR_2B, + WAV2VEC2_XLSR_300M, + Wav2Vec2ASRBundle, + Wav2Vec2Bundle, + Wav2Vec2FABundle, + WAVLM_BASE, + WAVLM_BASE_PLUS, + WAVLM_LARGE, +) +from .rnnt_pipeline import EMFORMER_RNNT_BASE_LIBRISPEECH, RNNTBundle + + +__all__ = [ + "Wav2Vec2Bundle", + "Wav2Vec2ASRBundle", + "Wav2Vec2FABundle", + "WAV2VEC2_BASE", + "WAV2VEC2_LARGE", + "WAV2VEC2_LARGE_LV60K", + "WAV2VEC2_ASR_BASE_10M", + "WAV2VEC2_ASR_BASE_100H", + "WAV2VEC2_ASR_BASE_960H", + "WAV2VEC2_ASR_LARGE_10M", + "WAV2VEC2_ASR_LARGE_100H", + "WAV2VEC2_ASR_LARGE_960H", + "WAV2VEC2_ASR_LARGE_LV60K_10M", + "WAV2VEC2_ASR_LARGE_LV60K_100H", + "WAV2VEC2_ASR_LARGE_LV60K_960H", + "WAV2VEC2_XLSR53", + "WAV2VEC2_XLSR_300M", + "WAV2VEC2_XLSR_1B", + "WAV2VEC2_XLSR_2B", + "VOXPOPULI_ASR_BASE_10K_EN", + "VOXPOPULI_ASR_BASE_10K_ES", + "VOXPOPULI_ASR_BASE_10K_DE", + "VOXPOPULI_ASR_BASE_10K_FR", + "VOXPOPULI_ASR_BASE_10K_IT", + "HUBERT_BASE", + "HUBERT_LARGE", + "HUBERT_XLARGE", + "HUBERT_ASR_LARGE", + "HUBERT_ASR_XLARGE", + "MMS_FA", + "WAVLM_BASE", + "WAVLM_BASE_PLUS", + "WAVLM_LARGE", + "Tacotron2TTSBundle", + "TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH", + "TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH", + "TACOTRON2_WAVERNN_CHAR_LJSPEECH", + "TACOTRON2_WAVERNN_PHONE_LJSPEECH", + "RNNTBundle", + "EMFORMER_RNNT_BASE_LIBRISPEECH", + "SourceSeparationBundle", + "CONVTASNET_BASE_LIBRI2MIX", + "HDEMUCS_HIGH_MUSDB_PLUS", + "HDEMUCS_HIGH_MUSDB", + "SQUIM_OBJECTIVE", + "SQUIM_SUBJECTIVE", + "SquimObjectiveBundle", + "SquimSubjectiveBundle", +] diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..92473e7f058487bfbfa8908b2d0032a8b67cdddd Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/_source_separation_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/_source_separation_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2e3993183b84b48ba6553584e2830f429bc91f30 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/_source_separation_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/_squim_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/_squim_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..5efc4cb89a98737097345fb708dcd77f9d35f126 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/_squim_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..57c3c2f6aee50e940fb408949c00b72a8cd9dd67 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_source_separation_pipeline.py b/lib/python3.12/site-packages/torchaudio/pipelines/_source_separation_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..ae92e21831307f91450b32a73563c0011e455753 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_source_separation_pipeline.py @@ -0,0 +1,109 @@ +from dataclasses import dataclass +from functools import partial +from typing import Callable + +import torch +import torchaudio + +from torchaudio.models import conv_tasnet_base, hdemucs_high + + +@dataclass +class SourceSeparationBundle: + """Dataclass that bundles components for performing source separation. + + Example + >>> import torchaudio + >>> from torchaudio.pipelines import CONVTASNET_BASE_LIBRI2MIX + >>> import torch + >>> + >>> # Build the separation model. + >>> model = CONVTASNET_BASE_LIBRI2MIX.get_model() + >>> 100%|███████████████████████████████|19.1M/19.1M [00:04<00:00, 4.93MB/s] + >>> + >>> # Instantiate the test set of Libri2Mix dataset. + >>> dataset = torchaudio.datasets.LibriMix("/home/datasets/", subset="test") + >>> + >>> # Apply source separation on mixture audio. + >>> for i, data in enumerate(dataset): + >>> sample_rate, mixture, clean_sources = data + >>> # Make sure the shape of input suits the model requirement. + >>> mixture = mixture.reshape(1, 1, -1) + >>> estimated_sources = model(mixture) + >>> score = si_snr_pit(estimated_sources, clean_sources) # for demonstration + >>> print(f"Si-SNR score is : {score}.) + >>> break + >>> Si-SNR score is : 16.24. + >>> + """ + + _model_path: str + _model_factory_func: Callable[[], torch.nn.Module] + _sample_rate: int + + @property + def sample_rate(self) -> int: + """Sample rate of the audio that the model is trained on. + + :type: int + """ + return self._sample_rate + + def get_model(self) -> torch.nn.Module: + """Construct the model and load the pretrained weight.""" + model = self._model_factory_func() + path = torchaudio.utils.download_asset(self._model_path) + state_dict = torch.load(path) + model.load_state_dict(state_dict) + model.eval() + return model + + +CONVTASNET_BASE_LIBRI2MIX = SourceSeparationBundle( + _model_path="models/conv_tasnet_base_libri2mix.pt", + _model_factory_func=partial(conv_tasnet_base, num_sources=2), + _sample_rate=8000, +) +CONVTASNET_BASE_LIBRI2MIX.__doc__ = """Pre-trained Source Separation pipeline with *ConvTasNet* +:cite:`Luo_2019` trained on *Libri2Mix dataset* :cite:`cosentino2020librimix`. + +The source separation model is constructed by :func:`~torchaudio.models.conv_tasnet_base` +and is trained using the training script ``lightning_train.py`` +`here `__ +with default arguments. + +Please refer to :class:`SourceSeparationBundle` for usage instructions. +""" + + +HDEMUCS_HIGH_MUSDB_PLUS = SourceSeparationBundle( + _model_path="models/hdemucs_high_trained.pt", + _model_factory_func=partial(hdemucs_high, sources=["drums", "bass", "other", "vocals"]), + _sample_rate=44100, +) +HDEMUCS_HIGH_MUSDB_PLUS.__doc__ = """Pre-trained music source separation pipeline with +*Hybrid Demucs* :cite:`defossez2021hybrid` trained on both training and test sets of +MUSDB-HQ :cite:`MUSDB18HQ` and an additional 150 extra songs from an internal database +that was specifically produced for Meta. + +The model is constructed by :func:`~torchaudio.models.hdemucs_high`. + +Training was performed in the original HDemucs repository `here `__. + +Please refer to :class:`SourceSeparationBundle` for usage instructions. +""" + + +HDEMUCS_HIGH_MUSDB = SourceSeparationBundle( + _model_path="models/hdemucs_high_musdbhq_only.pt", + _model_factory_func=partial(hdemucs_high, sources=["drums", "bass", "other", "vocals"]), + _sample_rate=44100, +) +HDEMUCS_HIGH_MUSDB.__doc__ = """Pre-trained music source separation pipeline with +*Hybrid Demucs* :cite:`defossez2021hybrid` trained on the training set of MUSDB-HQ :cite:`MUSDB18HQ`. + +The model is constructed by :func:`~torchaudio.models.hdemucs_high`. +Training was performed in the original HDemucs repository `here `__. + +Please refer to :class:`SourceSeparationBundle` for usage instructions. +""" diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_squim_pipeline.py b/lib/python3.12/site-packages/torchaudio/pipelines/_squim_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..0c70db4aef70397d33dcb3d3b28131221cef52c8 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_squim_pipeline.py @@ -0,0 +1,156 @@ +from dataclasses import dataclass + +import torch +import torchaudio + +from torchaudio.models import squim_objective_base, squim_subjective_base, SquimObjective, SquimSubjective + + +@dataclass +class SquimObjectiveBundle: + """Data class that bundles associated information to use pretrained + :py:class:`~torchaudio.models.SquimObjective` model. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + This bundle can estimate objective metric scores for speech enhancement, such as STOI, PESQ, Si-SDR. + A typical use case would be a flow like `waveform -> list of scores`. Please see below for the code example. + + Example: Estimate the objective metric scores for the input waveform. + >>> import torch + >>> import torchaudio + >>> from torchaudio.pipelines import SQUIM_OBJECTIVE as bundle + >>> + >>> # Load the SquimObjective bundle + >>> model = bundle.get_model() + Downloading: "https://download.pytorch.org/torchaudio/models/squim_objective_dns2020.pth" + 100%|████████████| 28.2M/28.2M [00:03<00:00, 9.24MB/s] + >>> + >>> # Resample audio to the expected sampling rate + >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate) + >>> + >>> # Estimate objective metric scores + >>> scores = model(waveform) + >>> print(f"STOI: {scores[0].item()}, PESQ: {scores[1].item()}, SI-SDR: {scores[2].item()}.") + """ # noqa: E501 + + _path: str + _sample_rate: float + + def get_model(self) -> SquimObjective: + """Construct the SquimObjective model, and load the pretrained weight. + + Returns: + Variation of :py:class:`~torchaudio.models.SquimObjective`. + """ + model = squim_objective_base() + path = torchaudio.utils.download_asset(f"models/{self._path}") + state_dict = torch.load(path, weights_only=True) + model.load_state_dict(state_dict) + model.eval() + return model + + @property + def sample_rate(self): + """Sample rate of the audio that the model is trained on. + + :type: float + """ + return self._sample_rate + + +SQUIM_OBJECTIVE = SquimObjectiveBundle( + "squim_objective_dns2020.pth", + _sample_rate=16000, +) +SQUIM_OBJECTIVE.__doc__ = """SquimObjective pipeline trained using approach described in + :cite:`kumar2023torchaudio` on the *DNS 2020 Dataset* :cite:`reddy2020interspeech`. + + The underlying model is constructed by :py:func:`torchaudio.models.squim_objective_base`. + The weights are under `Creative Commons Attribution 4.0 International License + `__. + + Please refer to :py:class:`SquimObjectiveBundle` for usage instructions. + """ + + +@dataclass +class SquimSubjectiveBundle: + """Data class that bundles associated information to use pretrained + :py:class:`~torchaudio.models.SquimSubjective` model. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + This bundle can estimate subjective metric scores for speech enhancement, such as MOS. + A typical use case would be a flow like `waveform -> score`. Please see below for the code example. + + Example: Estimate the subjective metric scores for the input waveform. + >>> import torch + >>> import torchaudio + >>> from torchaudio.pipelines import SQUIM_SUBJECTIVE as bundle + >>> + >>> # Load the SquimSubjective bundle + >>> model = bundle.get_model() + Downloading: "https://download.pytorch.org/torchaudio/models/squim_subjective_bvcc_daps.pth" + 100%|████████████| 360M/360M [00:09<00:00, 41.1MB/s] + >>> + >>> # Resample audio to the expected sampling rate + >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate) + >>> # Use a clean reference (doesn't need to be the reference for the waveform) as the second input + >>> reference = torchaudio.functional.resample(reference, sample_rate, bundle.sample_rate) + >>> + >>> # Estimate subjective metric scores + >>> score = model(waveform, reference) + >>> print(f"MOS: {score}.") + """ # noqa: E501 + + _path: str + _sample_rate: float + + def get_model(self) -> SquimSubjective: + """Construct the SquimSubjective model, and load the pretrained weight. + Returns: + Variation of :py:class:`~torchaudio.models.SquimObjective`. + """ + model = squim_subjective_base() + path = torchaudio.utils.download_asset(f"models/{self._path}") + state_dict = torch.load(path, weights_only=True) + model.load_state_dict(state_dict) + model.eval() + return model + + @property + def sample_rate(self): + """Sample rate of the audio that the model is trained on. + + :type: float + """ + return self._sample_rate + + +SQUIM_SUBJECTIVE = SquimSubjectiveBundle( + "squim_subjective_bvcc_daps.pth", + _sample_rate=16000, +) +SQUIM_SUBJECTIVE.__doc__ = """SquimSubjective pipeline trained + as described in :cite:`manocha2022speech` and :cite:`kumar2023torchaudio` + on the *BVCC* :cite:`cooper2021voices` and *DAPS* :cite:`mysore2014can` datasets. + + The underlying model is constructed by :py:func:`torchaudio.models.squim_subjective_base`. + The weights are under `Creative Commons Attribution Non Commercial 4.0 International + `__. + + Please refer to :py:class:`SquimSubjectiveBundle` for usage instructions. + """ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__init__.py b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..02851f596ceb281acc75c4d6a1aaf17eeee4a809 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__init__.py @@ -0,0 +1,16 @@ +from .impl import ( + TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH, + TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH, + TACOTRON2_WAVERNN_CHAR_LJSPEECH, + TACOTRON2_WAVERNN_PHONE_LJSPEECH, +) +from .interface import Tacotron2TTSBundle + + +__all__ = [ + "Tacotron2TTSBundle", + "TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH", + "TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH", + "TACOTRON2_WAVERNN_CHAR_LJSPEECH", + "TACOTRON2_WAVERNN_PHONE_LJSPEECH", +] diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..1de6a0fa4e4c9e4fd8ad3773b82752bc719f7634 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/impl.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..cb3a7f77ec6e46d68a963696c4d5548ecc74cc41 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/interface.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/interface.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ec42e94c01f1aafcce10ee59472aaf7ffaf7c517 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/interface.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..da402149651458c380dbe40647ee64a9b9f52f0f Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/impl.py b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/impl.py new file mode 100644 index 0000000000000000000000000000000000000000..b8542286242dcbb2036fff49c1d0e11fbbf9258b --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/impl.py @@ -0,0 +1,385 @@ +import re +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple, Union + +import torch +from torch import Tensor +from torchaudio._internal import load_state_dict_from_url +from torchaudio.functional import mu_law_decoding +from torchaudio.models import Tacotron2, WaveRNN +from torchaudio.transforms import GriffinLim, InverseMelScale + +from . import utils +from .interface import Tacotron2TTSBundle + +__all__ = [] + +_BASE_URL = "https://download.pytorch.org/torchaudio/models" + + +################################################################################ +# Pipeline implementation - Text Processor +################################################################################ + + +class _EnglishCharProcessor(Tacotron2TTSBundle.TextProcessor): + def __init__(self): + super().__init__() + self._tokens = utils._get_chars() + self._mapping = {s: i for i, s in enumerate(self._tokens)} + + @property + def tokens(self): + return self._tokens + + def __call__(self, texts: Union[str, List[str]]) -> Tuple[Tensor, Tensor]: + if isinstance(texts, str): + texts = [texts] + indices = [[self._mapping[c] for c in t.lower() if c in self._mapping] for t in texts] + return utils._to_tensor(indices) + + +class _EnglishPhoneProcessor(Tacotron2TTSBundle.TextProcessor): + def __init__(self, *, dl_kwargs=None): + super().__init__() + self._tokens = utils._get_phones() + self._mapping = {p: i for i, p in enumerate(self._tokens)} + self._phonemizer = utils._load_phonemizer("en_us_cmudict_forward.pt", dl_kwargs=dl_kwargs) + self._pattern = r"(\[[A-Z]+?\]|[_!'(),.:;? -])" + + @property + def tokens(self): + return self._tokens + + def __call__(self, texts: Union[str, List[str]]) -> Tuple[Tensor, Tensor]: + if isinstance(texts, str): + texts = [texts] + + indices = [] + for phones in self._phonemizer(texts, lang="en_us"): + # '[F][UW][B][AA][R]!' -> ['F', 'UW', 'B', 'AA', 'R', '!'] + ret = [re.sub(r"[\[\]]", "", r) for r in re.findall(self._pattern, phones)] + indices.append([self._mapping[p] for p in ret]) + return utils._to_tensor(indices) + + +################################################################################ +# Pipeline implementation - Vocoder +################################################################################ + + +class _WaveRNNVocoder(torch.nn.Module, Tacotron2TTSBundle.Vocoder): + def __init__(self, model: WaveRNN, min_level_db: Optional[float] = -100): + super().__init__() + self._sample_rate = 22050 + self._model = model + self._min_level_db = min_level_db + + @property + def sample_rate(self): + return self._sample_rate + + def forward(self, mel_spec, lengths=None): + mel_spec = torch.exp(mel_spec) + mel_spec = 20 * torch.log10(torch.clamp(mel_spec, min=1e-5)) + if self._min_level_db is not None: + mel_spec = (self._min_level_db - mel_spec) / self._min_level_db + mel_spec = torch.clamp(mel_spec, min=0, max=1) + waveform, lengths = self._model.infer(mel_spec, lengths) + waveform = utils._unnormalize_waveform(waveform, self._model.n_bits) + waveform = mu_law_decoding(waveform, self._model.n_classes) + waveform = waveform.squeeze(1) + return waveform, lengths + + +class _GriffinLimVocoder(torch.nn.Module, Tacotron2TTSBundle.Vocoder): + def __init__(self): + super().__init__() + self._sample_rate = 22050 + self._inv_mel = InverseMelScale( + n_stft=(1024 // 2 + 1), + n_mels=80, + sample_rate=self.sample_rate, + f_min=0.0, + f_max=8000.0, + mel_scale="slaney", + norm="slaney", + ) + self._griffin_lim = GriffinLim( + n_fft=1024, + power=1, + hop_length=256, + win_length=1024, + ) + + @property + def sample_rate(self): + return self._sample_rate + + def forward(self, mel_spec, lengths=None): + mel_spec = torch.exp(mel_spec) + mel_spec = mel_spec.clone().detach().requires_grad_(True) + spec = self._inv_mel(mel_spec) + spec = spec.detach().requires_grad_(False) + waveforms = self._griffin_lim(spec) + return waveforms, lengths + + +################################################################################ +# Bundle classes mixins +################################################################################ + + +class _CharMixin: + def get_text_processor(self) -> Tacotron2TTSBundle.TextProcessor: + return _EnglishCharProcessor() + + +class _PhoneMixin: + def get_text_processor(self, *, dl_kwargs=None) -> Tacotron2TTSBundle.TextProcessor: + return _EnglishPhoneProcessor(dl_kwargs=dl_kwargs) + + +@dataclass +class _Tacotron2Mixin: + _tacotron2_path: str + _tacotron2_params: Dict[str, Any] + + def get_tacotron2(self, *, dl_kwargs=None) -> Tacotron2: + model = Tacotron2(**self._tacotron2_params) + url = f"{_BASE_URL}/{self._tacotron2_path}" + dl_kwargs = {} if dl_kwargs is None else dl_kwargs + state_dict = load_state_dict_from_url(url, **dl_kwargs) + model.load_state_dict(state_dict) + model.eval() + return model + + +@dataclass +class _WaveRNNMixin: + _wavernn_path: Optional[str] + _wavernn_params: Optional[Dict[str, Any]] + + def get_vocoder(self, *, dl_kwargs=None): + wavernn = self._get_wavernn(dl_kwargs=dl_kwargs) + return _WaveRNNVocoder(wavernn) + + def _get_wavernn(self, *, dl_kwargs=None): + model = WaveRNN(**self._wavernn_params) + url = f"{_BASE_URL}/{self._wavernn_path}" + dl_kwargs = {} if dl_kwargs is None else dl_kwargs + state_dict = load_state_dict_from_url(url, **dl_kwargs) + model.load_state_dict(state_dict) + model.eval() + return model + + +class _GriffinLimMixin: + def get_vocoder(self, **_): + return _GriffinLimVocoder() + + +################################################################################ +# Bundle classes +################################################################################ + + +@dataclass +class _Tacotron2WaveRNNCharBundle(_WaveRNNMixin, _Tacotron2Mixin, _CharMixin, Tacotron2TTSBundle): + pass + + +@dataclass +class _Tacotron2WaveRNNPhoneBundle(_WaveRNNMixin, _Tacotron2Mixin, _PhoneMixin, Tacotron2TTSBundle): + pass + + +@dataclass +class _Tacotron2GriffinLimCharBundle(_GriffinLimMixin, _Tacotron2Mixin, _CharMixin, Tacotron2TTSBundle): + pass + + +@dataclass +class _Tacotron2GriffinLimPhoneBundle(_GriffinLimMixin, _Tacotron2Mixin, _PhoneMixin, Tacotron2TTSBundle): + pass + + +################################################################################ +# Instantiate bundle objects +################################################################################ + + +TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH = _Tacotron2GriffinLimCharBundle( + _tacotron2_path="tacotron2_english_characters_1500_epochs_ljspeech.pth", + _tacotron2_params=utils._get_taco_params(n_symbols=38), +) +TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH.__doc__ = """Character-based TTS pipeline with :py:class:`~torchaudio.models.Tacotron2` trained on *LJSpeech* :cite:`ljspeech17` for 1,500 epochs, and +:py:class:`~torchaudio.transforms.GriffinLim` as vocoder. + +The text processor encodes the input texts character-by-character. + +You can find the training script `here `__. +The default parameters were used. + +Please refer to :func:`torchaudio.pipelines.Tacotron2TTSBundle` for the usage. + +Example - "Hello world! T T S stands for Text to Speech!" + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + + +Example - "The examination and testimony of the experts enabled the Commission to conclude that five shots may have been fired," + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_GRIFFINLIM_CHAR_LJSPEECH_v2.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + +""" # noqa: E501 + +TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH = _Tacotron2GriffinLimPhoneBundle( + _tacotron2_path="tacotron2_english_phonemes_1500_epochs_ljspeech.pth", + _tacotron2_params=utils._get_taco_params(n_symbols=96), +) +TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH.__doc__ = """Phoneme-based TTS pipeline with :py:class:`~torchaudio.models.Tacotron2` trained on *LJSpeech* :cite:`ljspeech17` for 1,500 epochs and +:py:class:`~torchaudio.transforms.GriffinLim` as vocoder. + +The text processor encodes the input texts based on phoneme. +It uses `DeepPhonemizer `__ to convert +graphemes to phonemes. +The model (*en_us_cmudict_forward*) was trained on +`CMUDict `__. + +You can find the training script `here `__. +The text processor is set to the *"english_phonemes"*. + +Please refer to :func:`torchaudio.pipelines.Tacotron2TTSBundle` for the usage. + +Example - "Hello world! T T S stands for Text to Speech!" + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + + +Example - "The examination and testimony of the experts enabled the Commission to conclude that five shots may have been fired," + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_GRIFFINLIM_PHONE_LJSPEECH_v2.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + + +""" # noqa: E501 + +TACOTRON2_WAVERNN_CHAR_LJSPEECH = _Tacotron2WaveRNNCharBundle( + _tacotron2_path="tacotron2_english_characters_1500_epochs_wavernn_ljspeech.pth", + _tacotron2_params=utils._get_taco_params(n_symbols=38), + _wavernn_path="wavernn_10k_epochs_8bits_ljspeech.pth", + _wavernn_params=utils._get_wrnn_params(), +) +TACOTRON2_WAVERNN_CHAR_LJSPEECH.__doc__ = """Character-based TTS pipeline with :py:class:`~torchaudio.models.Tacotron2` trained on *LJSpeech* :cite:`ljspeech17` for 1,500 epochs and :py:class:`~torchaudio.models.WaveRNN` vocoder trained on 8 bits depth waveform of *LJSpeech* :cite:`ljspeech17` for 10,000 epochs. + +The text processor encodes the input texts character-by-character. + +You can find the training script `here `__. +The following parameters were used; ``win_length=1100``, ``hop_length=275``, ``n_fft=2048``, +``mel_fmin=40``, and ``mel_fmax=11025``. + +You can find the training script `here `__. + +Please refer to :func:`torchaudio.pipelines.Tacotron2TTSBundle` for the usage. + +Example - "Hello world! T T S stands for Text to Speech!" + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_WAVERNN_CHAR_LJSPEECH.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + + +Example - "The examination and testimony of the experts enabled the Commission to conclude that five shots may have been fired," + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_WAVERNN_CHAR_LJSPEECH_v2.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + +""" # noqa: E501 + +TACOTRON2_WAVERNN_PHONE_LJSPEECH = _Tacotron2WaveRNNPhoneBundle( + _tacotron2_path="tacotron2_english_phonemes_1500_epochs_wavernn_ljspeech.pth", + _tacotron2_params=utils._get_taco_params(n_symbols=96), + _wavernn_path="wavernn_10k_epochs_8bits_ljspeech.pth", + _wavernn_params=utils._get_wrnn_params(), +) +TACOTRON2_WAVERNN_PHONE_LJSPEECH.__doc__ = """Phoneme-based TTS pipeline with :py:class:`~torchaudio.models.Tacotron2` trained on *LJSpeech* :cite:`ljspeech17` for 1,500 epochs, and +:py:class:`~torchaudio.models.WaveRNN` vocoder trained on 8 bits depth waveform of *LJSpeech* :cite:`ljspeech17` for 10,000 epochs. + +The text processor encodes the input texts based on phoneme. +It uses `DeepPhonemizer `__ to convert +graphemes to phonemes. +The model (*en_us_cmudict_forward*) was trained on +`CMUDict `__. + +You can find the training script for Tacotron2 `here `__. +The following parameters were used; ``win_length=1100``, ``hop_length=275``, ``n_fft=2048``, +``mel_fmin=40``, and ``mel_fmax=11025``. + +You can find the training script for WaveRNN `here `__. + +Please refer to :func:`torchaudio.pipelines.Tacotron2TTSBundle` for the usage. + +Example - "Hello world! T T S stands for Text to Speech!" + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_WAVERNN_PHONE_LJSPEECH.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + + + +Example - "The examination and testimony of the experts enabled the Commission to conclude that five shots may have been fired," + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/TACOTRON2_WAVERNN_PHONE_LJSPEECH_v2.png + :alt: Spectrogram generated by Tacotron2 + + .. raw:: html + + +""" # noqa: E501 diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/interface.py b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/interface.py new file mode 100644 index 0000000000000000000000000000000000000000..564f236bc7c239d17dc82db04c350a9ccc618841 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/interface.py @@ -0,0 +1,255 @@ +from abc import ABC, abstractmethod +from typing import List, Optional, Tuple, Union + +from torch import Tensor +from torchaudio.models import Tacotron2 + + +class _TextProcessor(ABC): + @property + @abstractmethod + def tokens(self): + """The tokens that the each value in the processed tensor represent. + + :type: List[str] + """ + + @abstractmethod + def __call__(self, texts: Union[str, List[str]]) -> Tuple[Tensor, Tensor]: + """Encode the given (batch of) texts into numerical tensors + + Args: + text (str or list of str): The input texts. + + Returns: + (Tensor, Tensor): + Tensor: + The encoded texts. Shape: `(batch, max length)` + Tensor: + The valid length of each sample in the batch. Shape: `(batch, )`. + """ + + +class _Vocoder(ABC): + @property + @abstractmethod + def sample_rate(self): + """The sample rate of the resulting waveform + + :type: float + """ + + @abstractmethod + def __call__(self, specgrams: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]: + """Generate waveform from the given input, such as spectrogram + + Args: + specgrams (Tensor): + The input spectrogram. Shape: `(batch, frequency bins, time)`. + The expected shape depends on the implementation. + lengths (Tensor, or None, optional): + The valid length of each sample in the batch. Shape: `(batch, )`. + (Default: `None`) + + Returns: + (Tensor, Optional[Tensor]): + Tensor: + The generated waveform. Shape: `(batch, max length)` + Tensor or None: + The valid length of each sample in the batch. Shape: `(batch, )`. + """ + + +class Tacotron2TTSBundle(ABC): + """Data class that bundles associated information to use pretrained Tacotron2 and vocoder. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + Please see below for the usage and the available values. + + Example - Character-based TTS pipeline with Tacotron2 and WaveRNN + >>> import torchaudio + >>> + >>> text = "Hello, T T S !" + >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH + >>> + >>> # Build processor, Tacotron2 and WaveRNN model + >>> processor = bundle.get_text_processor() + >>> tacotron2 = bundle.get_tacotron2() + Downloading: + 100%|███████████████████████████████| 107M/107M [00:01<00:00, 87.9MB/s] + >>> vocoder = bundle.get_vocoder() + Downloading: + 100%|███████████████████████████████| 16.7M/16.7M [00:00<00:00, 78.1MB/s] + >>> + >>> # Encode text + >>> input, lengths = processor(text) + >>> + >>> # Generate (mel-scale) spectrogram + >>> specgram, lengths, _ = tacotron2.infer(input, lengths) + >>> + >>> # Convert spectrogram to waveform + >>> waveforms, lengths = vocoder(specgram, lengths) + >>> + >>> torchaudio.save('hello-tts.wav', waveforms, vocoder.sample_rate) + + Example - Phoneme-based TTS pipeline with Tacotron2 and WaveRNN + >>> + >>> # Note: + >>> # This bundle uses pre-trained DeepPhonemizer as + >>> # the text pre-processor. + >>> # Please install deep-phonemizer. + >>> # See https://github.com/as-ideas/DeepPhonemizer + >>> # The pretrained weight is automatically downloaded. + >>> + >>> import torchaudio + >>> + >>> text = "Hello, TTS!" + >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_PHONE_LJSPEECH + >>> + >>> # Build processor, Tacotron2 and WaveRNN model + >>> processor = bundle.get_text_processor() + Downloading: + 100%|███████████████████████████████| 63.6M/63.6M [00:04<00:00, 15.3MB/s] + >>> tacotron2 = bundle.get_tacotron2() + Downloading: + 100%|███████████████████████████████| 107M/107M [00:01<00:00, 87.9MB/s] + >>> vocoder = bundle.get_vocoder() + Downloading: + 100%|███████████████████████████████| 16.7M/16.7M [00:00<00:00, 78.1MB/s] + >>> + >>> # Encode text + >>> input, lengths = processor(text) + >>> + >>> # Generate (mel-scale) spectrogram + >>> specgram, lengths, _ = tacotron2.infer(input, lengths) + >>> + >>> # Convert spectrogram to waveform + >>> waveforms, lengths = vocoder(specgram, lengths) + >>> + >>> torchaudio.save('hello-tts.wav', waveforms, vocoder.sample_rate) + """ + + # Using the inner class so that these interfaces are not directly exposed on + # `torchaudio.pipelines`, but still listed in documentation. + # The thing is, text processing and vocoder are generic and we do not know what kind of + # new text processing and vocoder will be added in the future, so we want to make these + # interfaces specific to this Tacotron2TTS pipeline. + + class TextProcessor(_TextProcessor): + """Interface of the text processing part of Tacotron2TTS pipeline + + See :func:`torchaudio.pipelines.Tacotron2TTSBundle.get_text_processor` for the usage. + """ + + class Vocoder(_Vocoder): + """Interface of the vocoder part of Tacotron2TTS pipeline + + See :func:`torchaudio.pipelines.Tacotron2TTSBundle.get_vocoder` for the usage. + """ + + @abstractmethod + def get_text_processor(self, *, dl_kwargs=None) -> TextProcessor: + """Create a text processor + + For character-based pipeline, this processor splits the input text by character. + For phoneme-based pipeline, this processor converts the input text (grapheme) to + phonemes. + + If a pre-trained weight file is necessary, + :func:`torch.hub.download_url_to_file` is used to downloaded it. + + Args: + dl_kwargs (dictionary of keyword arguments,): + Passed to :func:`torch.hub.download_url_to_file`. + + Returns: + TextProcessor: + A callable which takes a string or a list of strings as input and + returns Tensor of encoded texts and Tensor of valid lengths. + The object also has ``tokens`` property, which allows to recover the + tokenized form. + + Example - Character-based + >>> text = [ + >>> "Hello World!", + >>> "Text-to-speech!", + >>> ] + >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH + >>> processor = bundle.get_text_processor() + >>> input, lengths = processor(text) + >>> + >>> print(input) + tensor([[19, 16, 23, 23, 26, 11, 34, 26, 29, 23, 15, 2, 0, 0, 0], + [31, 16, 35, 31, 1, 31, 26, 1, 30, 27, 16, 16, 14, 19, 2]], + dtype=torch.int32) + >>> + >>> print(lengths) + tensor([12, 15], dtype=torch.int32) + >>> + >>> print([processor.tokens[i] for i in input[0, :lengths[0]]]) + ['h', 'e', 'l', 'l', 'o', ' ', 'w', 'o', 'r', 'l', 'd', '!'] + >>> print([processor.tokens[i] for i in input[1, :lengths[1]]]) + ['t', 'e', 'x', 't', '-', 't', 'o', '-', 's', 'p', 'e', 'e', 'c', 'h', '!'] + + Example - Phoneme-based + >>> text = [ + >>> "Hello, T T S !", + >>> "Text-to-speech!", + >>> ] + >>> bundle = torchaudio.pipelines.TACOTRON2_WAVERNN_PHONE_LJSPEECH + >>> processor = bundle.get_text_processor() + Downloading: + 100%|███████████████████████████████| 63.6M/63.6M [00:04<00:00, 15.3MB/s] + >>> input, lengths = processor(text) + >>> + >>> print(input) + tensor([[54, 20, 65, 69, 11, 92, 44, 65, 38, 2, 0, 0, 0, 0], + [81, 40, 64, 79, 81, 1, 81, 20, 1, 79, 77, 59, 37, 2]], + dtype=torch.int32) + >>> + >>> print(lengths) + tensor([10, 14], dtype=torch.int32) + >>> + >>> print([processor.tokens[i] for i in input[0]]) + ['HH', 'AH', 'L', 'OW', ' ', 'W', 'ER', 'L', 'D', '!', '_', '_', '_', '_'] + >>> print([processor.tokens[i] for i in input[1]]) + ['T', 'EH', 'K', 'S', 'T', '-', 'T', 'AH', '-', 'S', 'P', 'IY', 'CH', '!'] + """ + + @abstractmethod + def get_vocoder(self, *, dl_kwargs=None) -> Vocoder: + """Create a vocoder module, based off of either WaveRNN or GriffinLim. + + If a pre-trained weight file is necessary, + :func:`torch.hub.load_state_dict_from_url` is used to downloaded it. + + Args: + dl_kwargs (dictionary of keyword arguments): + Passed to :func:`torch.hub.load_state_dict_from_url`. + + Returns: + Vocoder: + A vocoder module, which takes spectrogram Tensor and an optional + length Tensor, then returns resulting waveform Tensor and an optional + length Tensor. + """ + + @abstractmethod + def get_tacotron2(self, *, dl_kwargs=None) -> Tacotron2: + """Create a Tacotron2 model with pre-trained weight. + + Args: + dl_kwargs (dictionary of keyword arguments): + Passed to :func:`torch.hub.load_state_dict_from_url`. + + Returns: + Tacotron2: + The resulting model. + """ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_tts/utils.py b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..ef3ecb31335ae0cf9950d040fa11c21fb3403b25 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_tts/utils.py @@ -0,0 +1,228 @@ +import logging +import os + +import torch +from torchaudio._internal import download_url_to_file, module_utils as _mod_utils + + +def _get_chars(): + return ( + "_", + "-", + "!", + "'", + "(", + ")", + ",", + ".", + ":", + ";", + "?", + " ", + "a", + "b", + "c", + "d", + "e", + "f", + "g", + "h", + "i", + "j", + "k", + "l", + "m", + "n", + "o", + "p", + "q", + "r", + "s", + "t", + "u", + "v", + "w", + "x", + "y", + "z", + ) + + +def _get_phones(): + return ( + "_", + "-", + "!", + "'", + "(", + ")", + ",", + ".", + ":", + ";", + "?", + " ", + "AA", + "AA0", + "AA1", + "AA2", + "AE", + "AE0", + "AE1", + "AE2", + "AH", + "AH0", + "AH1", + "AH2", + "AO", + "AO0", + "AO1", + "AO2", + "AW", + "AW0", + "AW1", + "AW2", + "AY", + "AY0", + "AY1", + "AY2", + "B", + "CH", + "D", + "DH", + "EH", + "EH0", + "EH1", + "EH2", + "ER", + "ER0", + "ER1", + "ER2", + "EY", + "EY0", + "EY1", + "EY2", + "F", + "G", + "HH", + "IH", + "IH0", + "IH1", + "IH2", + "IY", + "IY0", + "IY1", + "IY2", + "JH", + "K", + "L", + "M", + "N", + "NG", + "OW", + "OW0", + "OW1", + "OW2", + "OY", + "OY0", + "OY1", + "OY2", + "P", + "R", + "S", + "SH", + "T", + "TH", + "UH", + "UH0", + "UH1", + "UH2", + "UW", + "UW0", + "UW1", + "UW2", + "V", + "W", + "Y", + "Z", + "ZH", + ) + + +def _to_tensor(indices): + lengths = torch.tensor([len(i) for i in indices], dtype=torch.int32) + values = [torch.tensor(i) for i in indices] + values = torch.nn.utils.rnn.pad_sequence(values, batch_first=True) + return values, lengths + + +def _load_phonemizer(file, dl_kwargs): + if not _mod_utils.is_module_available("dp"): + raise RuntimeError("DeepPhonemizer is not installed. Please install it.") + + from dp.phonemizer import Phonemizer + + # By default, dp issues DEBUG level log. + logger = logging.getLogger("dp") + orig_level = logger.level + logger.setLevel(logging.INFO) + try: + url = f"https://public-asai-dl-models.s3.eu-central-1.amazonaws.com/DeepPhonemizer/{file}" + directory = os.path.join(torch.hub.get_dir(), "checkpoints") + os.makedirs(directory, exist_ok=True) + path = os.path.join(directory, file) + if not os.path.exists(path): + dl_kwargs = {} if dl_kwargs is None else dl_kwargs + download_url_to_file(url, path, **dl_kwargs) + return Phonemizer.from_checkpoint(path) + finally: + logger.setLevel(orig_level) + + +def _unnormalize_waveform(waveform: torch.Tensor, bits: int) -> torch.Tensor: + r"""Transform waveform [-1, 1] to label [0, 2 ** bits - 1]""" + waveform = torch.clamp(waveform, -1, 1) + waveform = (waveform + 1.0) * (2**bits - 1) / 2 + return torch.clamp(waveform, 0, 2**bits - 1).int() + + +def _get_taco_params(n_symbols): + return { + "mask_padding": False, + "n_mels": 80, + "n_frames_per_step": 1, + "symbol_embedding_dim": 512, + "encoder_embedding_dim": 512, + "encoder_n_convolution": 3, + "encoder_kernel_size": 5, + "decoder_rnn_dim": 1024, + "decoder_max_step": 2000, + "decoder_dropout": 0.1, + "decoder_early_stopping": True, + "attention_rnn_dim": 1024, + "attention_hidden_dim": 128, + "attention_location_n_filter": 32, + "attention_location_kernel_size": 31, + "attention_dropout": 0.1, + "prenet_dim": 256, + "postnet_n_convolution": 5, + "postnet_kernel_size": 5, + "postnet_embedding_dim": 512, + "gate_threshold": 0.5, + "n_symbol": n_symbols, + } + + +def _get_wrnn_params(): + return { + "upsample_scales": [5, 5, 11], + "n_classes": 2**8, # n_bits = 8 + "hop_length": 275, + "n_res_block": 10, + "n_rnn": 512, + "n_fc": 512, + "kernel_size": 5, + "n_freq": 80, + "n_hidden": 128, + "n_output": 128, + } diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__init__.py b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..d9165935c536ca226463798d10582548558a9bf7 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/aligner.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/aligner.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..a54fd10eef21a4aaffb60b2410c886e320585382 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/aligner.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/impl.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..84ff50e72d20e7184019a910ad90a124f71c367d Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/utils.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/utils.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..2cce43b0de37a3ed8594517d209a522f6c6dce7b Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/__pycache__/utils.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/aligner.py b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/aligner.py new file mode 100644 index 0000000000000000000000000000000000000000..3655d5bae88181796d6d889013b4438d0ea014b3 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/aligner.py @@ -0,0 +1,87 @@ +from abc import ABC, abstractmethod +from typing import Dict, List + +import torch +import torchaudio.functional as F +from torch import Tensor +from torchaudio.functional import TokenSpan + + +class ITokenizer(ABC): + @abstractmethod + def __call__(self, transcript: List[str]) -> List[List[str]]: + """Tokenize the given transcript (list of word) + + .. note:: + + The toranscript must be normalized. + + Args: + transcript (list of str): Transcript (list of word). + + Returns: + (list of int): List of token sequences + """ + + +class Tokenizer(ITokenizer): + def __init__(self, dictionary: Dict[str, int]): + self.dictionary = dictionary + + def __call__(self, transcript: List[str]) -> List[List[int]]: + return [[self.dictionary[c] for c in word] for word in transcript] + + +def _align_emission_and_tokens(emission: Tensor, tokens: List[int], blank: int = 0): + device = emission.device + emission = emission.unsqueeze(0) + targets = torch.tensor([tokens], dtype=torch.int32, device=device) + + aligned_tokens, scores = F.forced_align(emission, targets, blank=blank) + + scores = scores.exp() # convert back to probability + aligned_tokens, scores = aligned_tokens[0], scores[0] # remove batch dimension + return aligned_tokens, scores + + +class IAligner(ABC): + @abstractmethod + def __call__(self, emission: Tensor, tokens: List[List[int]]) -> List[List[TokenSpan]]: + """Generate list of time-stamped token sequences + + Args: + emission (Tensor): Sequence of token probability distributions in log-domain. + Shape: `(time, tokens)`. + tokens (list of integer sequence): Tokenized transcript. + Output from :py:class:`torchaudio.pipelines.Wav2Vec2FABundle.Tokenizer`. + + Returns: + (list of TokenSpan sequence): Tokens with time stamps and scores. + """ + + +def _unflatten(list_, lengths): + assert len(list_) == sum(lengths) + i = 0 + ret = [] + for l in lengths: + ret.append(list_[i : i + l]) + i += l + return ret + + +def _flatten(nested_list): + return [item for list_ in nested_list for item in list_] + + +class Aligner(IAligner): + def __init__(self, blank): + self.blank = blank + + def __call__(self, emission: Tensor, tokens: List[List[int]]) -> List[List[TokenSpan]]: + if emission.ndim != 2: + raise ValueError(f"The input emission must be 2D. Found: {emission.shape}") + + aligned_tokens, scores = _align_emission_and_tokens(emission, _flatten(tokens), self.blank) + spans = F.merge_tokens(aligned_tokens, scores) + return _unflatten(spans, [len(ts) for ts in tokens]) diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/impl.py b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/impl.py new file mode 100644 index 0000000000000000000000000000000000000000..d60fa8adb94e92e1a479fe94e09f521d0fe50056 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/impl.py @@ -0,0 +1,1699 @@ +from dataclasses import dataclass +from typing import Any, Dict, Optional, Tuple + +from torch.nn import Module + +from . import aligner, utils + + +__all__ = [] # type: ignore + + +@dataclass +class Wav2Vec2Bundle: + """Data class that bundles associated information to use pretrained :py:class:`~torchaudio.models.Wav2Vec2Model`. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + Please see below for the usage and the available values. + + Example - Feature Extraction + >>> import torchaudio + >>> + >>> bundle = torchaudio.pipelines.HUBERT_BASE + >>> + >>> # Build the model and load pretrained weight. + >>> model = bundle.get_model() + Downloading: + 100%|███████████████████████████████| 360M/360M [00:06<00:00, 60.6MB/s] + >>> + >>> # Resample audio to the expected sampling rate + >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate) + >>> + >>> # Extract acoustic features + >>> features, _ = model.extract_features(waveform) + """ # noqa: E501 + + _path: str + _params: Dict[str, Any] + _sample_rate: float + _normalize_waveform: bool + _model_type: str + + @property + def sample_rate(self) -> float: + """Sample rate of the audio that the model is trained on. + + :type: float + """ + return self._sample_rate + + def _get_state_dict(self, dl_kwargs): + # Note: This method is overridden in ASR bundle + return utils._get_state_dict(self._path, dl_kwargs) + + def get_model(self, *, dl_kwargs=None) -> Module: + """Construct the model and load the pretrained weight. + + The weight file is downloaded from the internet and cached with + :func:`torch.hub.load_state_dict_from_url` + + Args: + dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`. + + Returns: + Variation of :py:class:`~torchaudio.models.Wav2Vec2Model`. + + For the models listed below, an additional layer normalization is performed on the input. + + For all other models, a :py:class:`~torchaudio.models.Wav2Vec2Model` instance is returned. + + - WAV2VEC2_LARGE_LV60K + - WAV2VEC2_ASR_LARGE_LV60K_10M + - WAV2VEC2_ASR_LARGE_LV60K_100H + - WAV2VEC2_ASR_LARGE_LV60K_960H + - WAV2VEC2_XLSR53 + - WAV2VEC2_XLSR_300M + - WAV2VEC2_XLSR_1B + - WAV2VEC2_XLSR_2B + - HUBERT_LARGE + - HUBERT_XLARGE + - HUBERT_ASR_LARGE + - HUBERT_ASR_XLARGE + - WAVLM_LARGE + """ + model = utils._get_model(self._model_type, self._params) + state_dict = self._get_state_dict(dl_kwargs) + model.load_state_dict(state_dict) + if self._normalize_waveform: + model = utils._extend_model(model, normalize_waveform=True) + model.eval() + return model + + +@dataclass +class Wav2Vec2ASRBundle(Wav2Vec2Bundle): + """Data class that bundles associated information to use pretrained + :py:class:`~torchaudio.models.Wav2Vec2Model`. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + Please see below for the usage and the available values. + + Example - ASR + >>> import torchaudio + >>> + >>> bundle = torchaudio.pipelines.HUBERT_ASR_LARGE + >>> + >>> # Build the model and load pretrained weight. + >>> model = bundle.get_model() + Downloading: + 100%|███████████████████████████████| 1.18G/1.18G [00:17<00:00, 73.8MB/s] + >>> + >>> # Check the corresponding labels of the output. + >>> labels = bundle.get_labels() + >>> print(labels) + ('-', '|', 'E', 'T', 'A', 'O', 'N', 'I', 'H', 'S', 'R', 'D', 'L', 'U', 'M', 'W', 'C', 'F', 'G', 'Y', 'P', 'B', 'V', 'K', "'", 'X', 'J', 'Q', 'Z') + >>> + >>> # Resample audio to the expected sampling rate + >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate) + >>> + >>> # Infer the label probability distribution + >>> emissions, _ = model(waveform) + >>> + >>> # Pass emission to decoder + >>> # `ctc_decode` is for illustration purpose only + >>> transcripts = ctc_decode(emissions, labels) + """ # noqa: E501 + + _labels: Tuple[str, ...] + _remove_aux_axis: Tuple[int, ...] = (1, 2, 3) + + def get_labels( + self, + *, + blank: str = "-", + ) -> Tuple[str, ...]: + """The output class labels. + + The first is blank token, and it is customizable. + + Args: + blank (str, optional): Blank token. (default: ``'-'``) + + Returns: + Tuple[str, ...]: + For models fine-tuned on ASR, returns the tuple of strings representing + the output class labels. + + Example + >>> from torchaudio.pipelines import HUBERT_ASR_LARGE as bundle + >>> bundle.get_labels() + ('-', '|', 'E', 'T', 'A', 'O', 'N', 'I', 'H', 'S', 'R', 'D', 'L', 'U', 'M', 'W', 'C', 'F', 'G', 'Y', 'P', 'B', 'V', 'K', "'", 'X', 'J', 'Q', 'Z') + """ # noqa: E501 + return (blank, *self._labels) + + def _get_state_dict(self, dl_kwargs): + return utils._get_state_dict(self._path, dl_kwargs, self._remove_aux_axis) + + +WAV2VEC2_BASE = Wav2Vec2Bundle( + _path="wav2vec2_fairseq_base_ls960.pth", + _params={ + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_BASE.__doc__ = """Wav2vec 2.0 model ("base" architecture), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), not fine-tuned. + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_BASE_10M = Wav2Vec2ASRBundle( + _path="wav2vec2_fairseq_base_ls960_asr_ll10m.pth", + _params={ + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_BASE_10M.__doc__ = """Wav2vec 2.0 model ("base" architecture with an extra linear module), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), and +fine-tuned for ASR on 10 minutes of transcribed audio from *Libri-Light* dataset +:cite:`librilight` ("train-10min" subset). + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_BASE_100H = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_base_ls960_asr_ls100.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) + +WAV2VEC2_ASR_BASE_100H.__doc__ = """Wav2vec 2.0 model ("base" architecture with an extra linear module), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), and +fine-tuned for ASR on 100 hours of transcribed audio from "train-clean-100" subset. + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_BASE_960H = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_base_ls960_asr_ls960.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_BASE_960H.__doc__ = """Wav2vec 2.0 model ("base" architecture with an extra linear module), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), and +fine-tuned for ASR on the same audio with the corresponding transcripts. + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_LARGE = Wav2Vec2Bundle( + "wav2vec2_fairseq_large_ls960.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.2, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_LARGE.__doc__ = """Wav2vec 2.0 model ("large" architecture), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), not fine-tuned. + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_LARGE_10M = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_large_ls960_asr_ll10m.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.2, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_LARGE_10M.__doc__ = """Wav2vec 2.0 model ("large" architecture with an extra linear module), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), and +fine-tuned for ASR on 10 minutes of transcribed audio from *Libri-Light* dataset +:cite:`librilight` ("train-10min" subset). + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_LARGE_100H = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_large_ls960_asr_ls100.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.2, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_LARGE_100H.__doc__ = """Wav2vec 2.0 model ("large" architecture with an extra linear module), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), and +fine-tuned for ASR on 100 hours of transcribed audio from +the same dataset ("train-clean-100" subset). + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_LARGE_960H = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_large_ls960_asr_ls960.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.2, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_LARGE_960H.__doc__ = """Wav2vec 2.0 model ("large" architecture with an extra linear module), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), and +fine-tuned for ASR on the same audio with the corresponding transcripts. + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_LARGE_LV60K = Wav2Vec2Bundle( + "wav2vec2_fairseq_large_lv60k.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +WAV2VEC2_LARGE_LV60K.__doc__ = """Wav2vec 2.0 model ("large-lv60k" architecture), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* dataset :cite:`librilight`, +not fine-tuned. + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_LARGE_LV60K_10M = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_large_lv60k_asr_ll10m.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_LARGE_LV60K_10M.__doc__ = """Wav2vec 2.0 model ("large-lv60k" architecture with an extra linear module), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* dataset :cite:`librilight`, and +fine-tuned for ASR on 10 minutes of transcribed audio from the same dataset ("train-10min" subset). + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_LARGE_LV60K_100H = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_large_lv60k_asr_ls100.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_LARGE_LV60K_100H.__doc__ = """Wav2vec 2.0 model ("large-lv60k" architecture with an extra linear module), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* dataset :cite:`librilight`, and +fine-tuned for ASR on 100 hours of transcribed audio from +*LibriSpeech* dataset :cite:`7178964` ("train-clean-100" subset). + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_ASR_LARGE_LV60K_960H = Wav2Vec2ASRBundle( + "wav2vec2_fairseq_large_lv60k_asr_ls960.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +WAV2VEC2_ASR_LARGE_LV60K_960H.__doc__ = """Wav2vec 2.0 model ("large-lv60k" architecture with an extra linear module), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* :cite:`librilight` dataset, and +fine-tuned for ASR on 960 hours of transcribed audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"). + +Originally published by the authors of *wav2vec 2.0* :cite:`baevski2020wav2vec` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +WAV2VEC2_XLSR53 = Wav2Vec2Bundle( + "wav2vec2_fairseq_large_xlsr53.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +WAV2VEC2_XLSR53.__doc__ = """Wav2vec 2.0 model ("base" architecture), +pre-trained on 56,000 hours of unlabeled audio from multiple datasets ( +*Multilingual LibriSpeech* :cite:`Pratap_2020`, +*CommonVoice* :cite:`ardila2020common` and +*BABEL* :cite:`Gales2014SpeechRA`), +not fine-tuned. + +Originally published by the authors of +*Unsupervised Cross-lingual Representation Learning for Speech Recognition* +:cite:`conneau2020unsupervised` under MIT License and redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +HUBERT_BASE = Wav2Vec2Bundle( + "hubert_fairseq_base_ls960.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +HUBERT_BASE.__doc__ = """HuBERT model ("base" architecture), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"), not fine-tuned. + +Originally published by the authors of *HuBERT* :cite:`hsu2021hubert` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +HUBERT_LARGE = Wav2Vec2Bundle( + "hubert_fairseq_large_ll60k.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +HUBERT_LARGE.__doc__ = """HuBERT model ("large" architecture), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* dataset :cite:`librilight`, +not fine-tuned. + +Originally published by the authors of *HuBERT* :cite:`hsu2021hubert` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +HUBERT_XLARGE = Wav2Vec2Bundle( + "hubert_fairseq_xlarge_ll60k.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1280, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 48, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 5120, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +HUBERT_XLARGE.__doc__ = """HuBERT model ("extra large" architecture), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* dataset :cite:`librilight`, +not fine-tuned. + +Originally published by the authors of *HuBERT* :cite:`hsu2021hubert` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + +HUBERT_ASR_LARGE = Wav2Vec2ASRBundle( + "hubert_fairseq_large_ll60k_asr_ls960.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.1, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +HUBERT_ASR_LARGE.__doc__ = """HuBERT model ("large" architecture), +pre-trained on 60,000 hours of unlabeled audio from *Libri-Light* dataset :cite:`librilight`, and +fine-tuned for ASR on 960 hours of transcribed audio from *LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"). + +Originally published by the authors of *HuBERT* :cite:`hsu2021hubert` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +HUBERT_ASR_XLARGE = Wav2Vec2ASRBundle( + "hubert_fairseq_xlarge_ll60k_asr_ls960.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1280, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 48, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 5120, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.1, + "aux_num_out": 29, + }, + _labels=utils._get_en_labels(), + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +HUBERT_ASR_XLARGE.__doc__ = """HuBERT model ("extra large" architecture), +pre-trained on 60,000 hours of unlabeled audio from +*Libri-Light* dataset :cite:`librilight`, and +fine-tuned for ASR on 960 hours of transcribed audio from +*LibriSpeech* dataset :cite:`7178964` +(the combination of "train-clean-100", "train-clean-360", and "train-other-500"). + +Originally published by the authors of *HuBERT* :cite:`hsu2021hubert` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + + +VOXPOPULI_ASR_BASE_10K_DE = Wav2Vec2ASRBundle( + "wav2vec2_voxpopuli_base_10k_asr_de.pt", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.1, + "aux_num_out": 32, + }, + _labels=utils._get_de_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _remove_aux_axis=(1, 2, 3, 35), + _model_type="Wav2Vec2", +) +VOXPOPULI_ASR_BASE_10K_DE.__doc__ = """wav2vec 2.0 model ("base" architecture), +pre-trained on 10k hours of unlabeled audio from *VoxPopuli* dataset :cite:`voxpopuli` +("10k" subset, consisting of 23 languages), and +fine-tuned for ASR on 282 hours of transcribed audio from "de" subset. + +Originally published by the authors of *VoxPopuli* :cite:`voxpopuli` under CC BY-NC 4.0 and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + + +VOXPOPULI_ASR_BASE_10K_EN = Wav2Vec2ASRBundle( + "wav2vec2_voxpopuli_base_10k_asr_en.pt", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.1, + "aux_num_out": 28, + }, + _labels=utils._get_vp_en_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _remove_aux_axis=(1, 2, 3, 31), + _model_type="Wav2Vec2", +) +VOXPOPULI_ASR_BASE_10K_EN.__doc__ = """wav2vec 2.0 model ("base" architecture), +pre-trained on 10k hours of unlabeled audio from *VoxPopuli* dataset :cite:`voxpopuli` +("10k" subset, consisting of 23 languages), and +fine-tuned for ASR on 543 hours of transcribed audio from "en" subset. + +Originally published by the authors of *VoxPopuli* :cite:`voxpopuli` under CC BY-NC 4.0 and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + + +VOXPOPULI_ASR_BASE_10K_ES = Wav2Vec2ASRBundle( + "wav2vec2_voxpopuli_base_10k_asr_es.pt", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.1, + "aux_num_out": 35, + }, + _labels=utils._get_es_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _remove_aux_axis=(1, 2, 3, 35), + _model_type="Wav2Vec2", +) +VOXPOPULI_ASR_BASE_10K_ES.__doc__ = """wav2vec 2.0 model ("base" architecture), +pre-trained on 10k hours of unlabeled audio from *VoxPopuli* dataset :cite:`voxpopuli` +("10k" subset, consisting of 23 languages), and +fine-tuned for ASR on 166 hours of transcribed audio from "es" subset. + +Originally published by the authors of *VoxPopuli* :cite:`voxpopuli` under CC BY-NC 4.0 and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + +VOXPOPULI_ASR_BASE_10K_FR = Wav2Vec2ASRBundle( + "wav2vec2_voxpopuli_base_10k_asr_fr.pt", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.1, + "aux_num_out": 43, + }, + _labels=utils._get_fr_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _model_type="Wav2Vec2", +) +VOXPOPULI_ASR_BASE_10K_FR.__doc__ = """wav2vec 2.0 model ("base" architecture), +pre-trained on 10k hours of unlabeled audio from *VoxPopuli* dataset :cite:`voxpopuli` +("10k" subset, consisting of 23 languages), and +fine-tuned for ASR on 211 hours of transcribed audio from "fr" subset. + +Originally published by the authors of *VoxPopuli* :cite:`voxpopuli` under CC BY-NC 4.0 and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + + +VOXPOPULI_ASR_BASE_10K_IT = Wav2Vec2ASRBundle( + "wav2vec2_voxpopuli_base_10k_asr_it.pt", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.1, + "aux_num_out": 37, + }, + _labels=utils._get_it_labels(), + _sample_rate=16000, + _normalize_waveform=False, + _remove_aux_axis=(1, 2, 3), + _model_type="Wav2Vec2", +) +VOXPOPULI_ASR_BASE_10K_IT.__doc__ = """wav2vec 2.0 model ("base" architecture), +pre-trained on 10k hours of unlabeled audio from *VoxPopuli* dataset :cite:`voxpopuli` +("10k" subset, consisting of 23 languages), and +fine-tuned for ASR on 91 hours of transcribed audio from "it" subset. + +Originally published by the authors of *VoxPopuli* :cite:`voxpopuli` under CC BY-NC 4.0 and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2ASRBundle` for the usage. +""" # noqa: E501 + + +WAVLM_BASE = Wav2Vec2Bundle( + "wavlm_base.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_max_distance": 800, + "encoder_num_buckets": 320, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": None, + }, + _model_type="WavLM", + _sample_rate=16000, + _normalize_waveform=False, +) +WAVLM_BASE.__doc__ = """WavLM Base model ("base" architecture), +pre-trained on 960 hours of unlabeled audio from *LibriSpeech* dataset :cite:`7178964`, not fine-tuned. + +Originally published by the authors of *WavLM* :cite:`chen2022wavlm` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + + +WAVLM_BASE_PLUS = Wav2Vec2Bundle( + "wavlm_base_plus.pth", + { + "extractor_mode": "group_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 768, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 12, + "encoder_num_heads": 12, + "encoder_max_distance": 800, + "encoder_num_buckets": 320, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 3072, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": False, + "encoder_layer_drop": 0.05, + "aux_num_out": None, + }, + _model_type="WavLM", + _sample_rate=16000, + _normalize_waveform=False, +) +WAVLM_BASE_PLUS.__doc__ = """WavLM Base+ model ("base" architecture), +pre-trained on 60,000 hours of Libri-Light dataset :cite:`librilight`, 10,000 hours of GigaSpeech :cite:`GigaSpeech2021`, +and 24,000 hours of *VoxPopuli* :cite:`voxpopuli`, not fine-tuned. + +Originally published by the authors of *WavLM* :cite:`chen2022wavlm` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + + +WAVLM_LARGE = Wav2Vec2Bundle( + "wavlm_large.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": False, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_max_distance": 800, + "encoder_num_buckets": 320, + "encoder_attention_dropout": 0.1, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.1, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.05, + "aux_num_out": None, + }, + _model_type="WavLM", + _sample_rate=16000, + _normalize_waveform=True, +) +WAVLM_LARGE.__doc__ = """WavLM Large model ("large" architecture), +pre-trained on 60,000 hours of Libri-Light dataset :cite:`librilight`, 10,000 hours of GigaSpeech :cite:`GigaSpeech2021`, +and 24,000 hours of *VoxPopuli* :cite:`voxpopuli`, not fine-tuned. + +Originally published by the authors of *WavLM* :cite:`chen2022wavlm` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for the usage. +""" # noqa: E501 + + +WAV2VEC2_XLSR_300M = Wav2Vec2Bundle( + "wav2vec2_xlsr_300m.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _model_type="Wav2Vec2", + _sample_rate=16000, + _normalize_waveform=True, +) +WAV2VEC2_XLSR_300M.__doc__ = """XLS-R model with 300 million parameters, +pre-trained on 436,000 hours of unlabeled audio from multiple datasets ( +*Multilingual LibriSpeech* :cite:`Pratap_2020`, +*CommonVoice* :cite:`ardila2020common`, +*VoxLingua107* :cite:`valk2021voxlingua107`, +*BABEL* :cite:`Gales2014SpeechRA`, and +*VoxPopuli* :cite:`voxpopuli`) in 128 languages, +not fine-tuned. + +Originally published by the authors of *XLS-R* :cite:`babu2021xls` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for usage details. +""" # noqa: E501 + + +WAV2VEC2_XLSR_1B = Wav2Vec2Bundle( + "wav2vec2_xlsr_1b.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1280, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 48, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 5120, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _model_type="Wav2Vec2", + _sample_rate=16000, + _normalize_waveform=True, +) +WAV2VEC2_XLSR_1B.__doc__ = """XLS-R model with 1 billion parameters, +pre-trained on 436,000 hours of unlabeled audio from multiple datasets ( +*Multilingual LibriSpeech* :cite:`Pratap_2020`, +*CommonVoice* :cite:`ardila2020common`, +*VoxLingua107* :cite:`valk2021voxlingua107`, +*BABEL* :cite:`Gales2014SpeechRA`, and +*VoxPopuli* :cite:`voxpopuli`) in 128 languages, +not fine-tuned. + +Originally published by the authors of *XLS-R* :cite:`babu2021xls` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for usage details. +""" # noqa: E501 + +WAV2VEC2_XLSR_2B = Wav2Vec2Bundle( + "wav2vec2_xlsr_2b.pth", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1920, + "encoder_projection_dropout": 0.1, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 48, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 7680, + "encoder_ff_interm_dropout": 0.0, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.0, + "aux_num_out": None, + }, + _model_type="Wav2Vec2", + _sample_rate=16000, + _normalize_waveform=True, +) +WAV2VEC2_XLSR_2B.__doc__ = """XLS-R model with 2 billion parameters, +pre-trained on 436,000 hours of unlabeled audio from multiple datasets ( +*Multilingual LibriSpeech* :cite:`Pratap_2020`, +*CommonVoice* :cite:`ardila2020common`, +*VoxLingua107* :cite:`valk2021voxlingua107`, +*BABEL* :cite:`Gales2014SpeechRA`, and +*VoxPopuli* :cite:`voxpopuli`) in 128 languages, +not fine-tuned. + +Originally published by the authors of *XLS-R* :cite:`babu2021xls` under MIT License and +redistributed with the same license. +[`License `__, +`Source `__] + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2Bundle` for usage details. +""" # noqa: E501 + + +@dataclass +class Wav2Vec2FABundle(Wav2Vec2ASRBundle): + """Data class that bundles associated information to use pretrained :py:class:`~torchaudio.models.Wav2Vec2Model` for forced alignment. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + Please see below for the usage and the available values. + + Example - Feature Extraction + >>> import torchaudio + >>> + >>> bundle = torchaudio.pipelines.MMS_FA + >>> + >>> # Build the model and load pretrained weight. + >>> model = bundle.get_model() + Downloading: + 100%|███████████████████████████████| 1.18G/1.18G [00:05<00:00, 216MB/s] + >>> + >>> # Resample audio to the expected sampling rate + >>> waveform = torchaudio.functional.resample(waveform, sample_rate, bundle.sample_rate) + >>> + >>> # Estimate the probability of token distribution + >>> emission, _ = model(waveform) + >>> + >>> # Generate frame-wise alignment + >>> alignment, scores = torchaudio.functional.forced_align( + >>> emission, targets, input_lengths, target_lengths, blank=0) + >>> + """ # noqa: E501 + + class Tokenizer(aligner.ITokenizer): + """Interface of the tokenizer""" + + class Aligner(aligner.IAligner): + """Interface of the aligner""" + + def get_labels(self, star: Optional[str] = "*", blank: str = "-") -> Tuple[str, ...]: + """Get the labels corresponding to the feature dimension of emission. + + The first is blank token, and it is customizable. + + Args: + star (str or None, optional): Change or disable star token. (default: ``"*"``) + blank (str, optional): Change the blank token. (default: ``'-'``) + + Returns: + Tuple[str, ...]: + For models fine-tuned on ASR, returns the tuple of strings representing + the output class labels. + + Example + >>> from torchaudio.pipelines import MMS_FA as bundle + >>> bundle.get_labels() + ('-', 'a', 'i', 'e', 'n', 'o', 'u', 't', 's', 'r', 'm', 'k', 'l', 'd', 'g', 'h', 'y', 'b', 'p', 'w', 'c', 'v', 'j', 'z', 'f', "'", 'q', 'x', '*') + >>> bundle.get_labels(star=None) + ('-', 'a', 'i', 'e', 'n', 'o', 'u', 't', 's', 'r', 'm', 'k', 'l', 'd', 'g', 'h', 'y', 'b', 'p', 'w', 'c', 'v', 'j', 'z', 'f', "'", 'q', 'x') + """ # noqa: E501 + labels = super().get_labels(blank=blank) + return labels if star is None else (*labels, star) + + def get_model(self, with_star: bool = True, *, dl_kwargs=None) -> Module: + """Construct the model and load the pretrained weight. + + The weight file is downloaded from the internet and cached with + :func:`torch.hub.load_state_dict_from_url` + + Args: + with_star (bool, optional): If enabled, the last dimension of output layer is + extended by one, which corresponds to `star` token. + dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`. + + Returns: + Variation of :py:class:`~torchaudio.models.Wav2Vec2Model`. + + .. note:: + + The model created with this method returns probability in log-domain, + (i.e. :py:func:`torch.nn.functional.log_softmax` is applied), whereas + the other Wav2Vec2 models returns logit. + """ + model = utils._get_model(self._model_type, self._params) + state_dict = utils._get_state_dict(self._path, dl_kwargs, self._remove_aux_axis) + model.load_state_dict(state_dict) + model = utils._extend_model( + model, normalize_waveform=self._normalize_waveform, apply_log_softmax=True, append_star=with_star + ) + model.eval() + return model + + def get_dict(self, star: Optional[str] = "*", blank: str = "-") -> Dict[str, int]: + """Get the mapping from token to index (in emission feature dim) + + Args: + star (str or None, optional): Change or disable star token. (default: ``"*"``) + blank (str, optional): Change the blank token. (default: ``'-'``) + + Returns: + Tuple[str, ...]: + For models fine-tuned on ASR, returns the tuple of strings representing + the output class labels. + + Example + >>> from torchaudio.pipelines import MMS_FA as bundle + >>> bundle.get_dict() + {'-': 0, 'a': 1, 'i': 2, 'e': 3, 'n': 4, 'o': 5, 'u': 6, 't': 7, 's': 8, 'r': 9, 'm': 10, 'k': 11, 'l': 12, 'd': 13, 'g': 14, 'h': 15, 'y': 16, 'b': 17, 'p': 18, 'w': 19, 'c': 20, 'v': 21, 'j': 22, 'z': 23, 'f': 24, "'": 25, 'q': 26, 'x': 27, '*': 28} + >>> bundle.get_dict(star=None) + {'-': 0, 'a': 1, 'i': 2, 'e': 3, 'n': 4, 'o': 5, 'u': 6, 't': 7, 's': 8, 'r': 9, 'm': 10, 'k': 11, 'l': 12, 'd': 13, 'g': 14, 'h': 15, 'y': 16, 'b': 17, 'p': 18, 'w': 19, 'c': 20, 'v': 21, 'j': 22, 'z': 23, 'f': 24, "'": 25, 'q': 26, 'x': 27} + """ # noqa: E501 + return {k: i for i, k in enumerate(self.get_labels(star=star, blank=blank))} + + def get_tokenizer(self) -> Tokenizer: + """Instantiate a Tokenizer. + + Returns: + Tokenizer + """ + return aligner.Tokenizer(self.get_dict()) + + def get_aligner(self) -> Aligner: + """Instantiate an Aligner. + + Returns: + Aligner + """ + return aligner.Aligner(blank=0) + + +MMS_FA = Wav2Vec2FABundle( + "https://dl.fbaipublicfiles.com/mms/torchaudio/ctc_alignment_mling_uroman/model.pt", + { + "extractor_mode": "layer_norm", + "extractor_conv_layer_config": [ + (512, 10, 5), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 3, 2), + (512, 2, 2), + (512, 2, 2), + ], + "extractor_conv_bias": True, + "encoder_embed_dim": 1024, + "encoder_projection_dropout": 0.0, + "encoder_pos_conv_kernel": 128, + "encoder_pos_conv_groups": 16, + "encoder_num_layers": 24, + "encoder_num_heads": 16, + "encoder_attention_dropout": 0.0, + "encoder_ff_interm_features": 4096, + "encoder_ff_interm_dropout": 0.1, + "encoder_dropout": 0.0, + "encoder_layer_norm_first": True, + "encoder_layer_drop": 0.1, + "aux_num_out": 28, + }, + _labels=utils._get_mms_labels(), + _sample_rate=16000, + _normalize_waveform=True, + _model_type="Wav2Vec2", +) +MMS_FA.__doc__ = """ +Trained on 31K hours of data in 1,130 languages from *Scaling Speech Technology to 1,000+ Languages* :cite:`pratap2023scaling`. + +Published by the authors of *Scaling Speech Technology to 1,000+ Languages* :cite:`pratap2023scaling` under [`CC-BY-NC 4.0 License `__]. + +Please refer to :py:class:`torchaudio.pipelines.Wav2Vec2FABundle` for usage details. + +.. note:: + + Unlike other Wav2Vec2 bundles, this model does not have a token for word boundary (like `|`). This makes the post-processing of alignments slightly different. +""" # noqa: E501 diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/utils.py b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..e690e8103c7a47a01d719e746e6c98a9c7f6c8db --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/_wav2vec2/utils.py @@ -0,0 +1,346 @@ +from typing import List, Optional, Tuple + +import torch +from torch import nn, Tensor + +from torchaudio._internal import load_state_dict_from_url +from torchaudio.models import wav2vec2_model, Wav2Vec2Model, wavlm_model + + +def _get_model(type_, params): + factories = { + "Wav2Vec2": wav2vec2_model, + "WavLM": wavlm_model, + } + if type_ not in factories: + raise ValueError(f"Supported model types are {tuple(factories.keys())}. Found: {type_}") + factory = factories[type_] + return factory(**params) + + +class _Wav2Vec2Model(nn.Module): + """Wrapper class for :py:class:`~torchaudio.models.Wav2Vec2Model`. + + This is used for layer normalization at the input + """ + + def __init__(self, model: Wav2Vec2Model, normalize_waveform: bool, apply_log_softmax: bool, append_star: bool): + super().__init__() + self.model = model + self.normalize_waveform = normalize_waveform + self.apply_log_softmax = apply_log_softmax + self.append_star = append_star + + def forward(self, waveforms: Tensor, lengths: Optional[Tensor] = None) -> Tuple[Tensor, Optional[Tensor]]: + if self.normalize_waveform: + waveforms = nn.functional.layer_norm(waveforms, waveforms.shape) + output, output_lengths = self.model(waveforms, lengths) + if self.apply_log_softmax: + output = torch.nn.functional.log_softmax(output, dim=-1) + if self.append_star: + star_dim = torch.zeros((1, output.size(1), 1), dtype=output.dtype, device=output.device) + output = torch.cat((output, star_dim), dim=-1) + return output, output_lengths + + @torch.jit.export + def extract_features( + self, + waveforms: Tensor, + lengths: Optional[Tensor] = None, + num_layers: Optional[int] = None, + ) -> Tuple[List[Tensor], Optional[Tensor]]: + if self.normalize_waveform: + waveforms = nn.functional.layer_norm(waveforms, waveforms.shape) + return self.model.extract_features(waveforms, lengths, num_layers) + + +def _extend_model(module, normalize_waveform, apply_log_softmax=False, append_star=False): + """Add extra transformations to the model""" + return _Wav2Vec2Model(module, normalize_waveform, apply_log_softmax, append_star) + + +def _remove_aux_axes(state_dict, axes): + # Remove the seemingly unnecessary axis + # For ASR task, the pretrained weights originated from fairseq has unrelated dimensions at index 1, 2, 3 + # It's originated from the Dictionary implementation of fairseq, which was intended for NLP tasks, + # but not used during the ASR training. + # https://github.com/pytorch/fairseq/blob/c5ff181125c7e6126b49a85e5ebdd5f5b6a07914/fairseq/data/dictionary.py#L21-L37 + # https://github.com/pytorch/fairseq/blob/c5ff181125c7e6126b49a85e5ebdd5f5b6a07914/fairseq/criterions/ctc.py#L126-L129 + # + # Also, some pretrained weights originated from voxpopuli has an extra dimensions that almost never used and + # that resembles mistake. + # The label `1` shows up in the training dataset of German (1 out of 16M), + # English (1 / 28M), Spanish (1 / 9.4M), Romanian (1 / 4.7M) and Polish (6 / 5.8M) + for key in ["aux.weight", "aux.bias"]: + mat = state_dict[key] + state_dict[key] = torch.stack([mat[i] for i in range(mat.size(0)) if i not in axes]) + + +def _get_state_dict(url, dl_kwargs, remove_axes=None): + if not url.startswith("https"): + url = f"https://download.pytorch.org/torchaudio/models/{url}" + dl_kwargs = {} if dl_kwargs is None else dl_kwargs + state_dict = load_state_dict_from_url(url, **dl_kwargs) + if remove_axes: + _remove_aux_axes(state_dict, remove_axes) + return state_dict + + +def _get_en_labels(): + return ( + "|", + "E", + "T", + "A", + "O", + "N", + "I", + "H", + "S", + "R", + "D", + "L", + "U", + "M", + "W", + "C", + "F", + "G", + "Y", + "P", + "B", + "V", + "K", + "'", + "X", + "J", + "Q", + "Z", + ) + + +def _get_de_labels(): + return ( + "|", + "e", + "n", + "i", + "r", + "s", + "t", + "a", + "d", + "h", + "u", + "l", + "g", + "c", + "m", + "o", + "b", + "w", + "f", + "k", + "z", + "p", + "v", + "ü", + "ä", + "ö", + "j", + "ß", + "y", + "x", + "q", + ) + + +def _get_vp_en_labels(): + return ( + "|", + "e", + "t", + "o", + "i", + "a", + "n", + "s", + "r", + "h", + "l", + "d", + "c", + "u", + "m", + "p", + "f", + "g", + "w", + "y", + "b", + "v", + "k", + "x", + "j", + "q", + "z", + ) + + +def _get_es_labels(): + return ( + "|", + "e", + "a", + "o", + "s", + "n", + "r", + "i", + "l", + "d", + "c", + "t", + "u", + "p", + "m", + "b", + "q", + "y", + "g", + "v", + "h", + "ó", + "f", + "í", + "á", + "j", + "z", + "ñ", + "é", + "x", + "ú", + "k", + "w", + "ü", + ) + + +def _get_fr_labels(): + return ( + "|", + "e", + "s", + "n", + "i", + "t", + "r", + "a", + "o", + "u", + "l", + "d", + "c", + "p", + "m", + "é", + "v", + "q", + "f", + "g", + "b", + "h", + "x", + "à", + "j", + "è", + "y", + "ê", + "z", + "ô", + "k", + "ç", + "œ", + "û", + "ù", + "î", + "â", + "w", + "ï", + "ë", + "ü", + "æ", + ) + + +def _get_it_labels(): + return ( + "|", + "e", + "i", + "a", + "o", + "n", + "t", + "r", + "l", + "s", + "c", + "d", + "u", + "p", + "m", + "g", + "v", + "h", + "z", + "f", + "b", + "q", + "à", + "è", + "ù", + "é", + "ò", + "ì", + "k", + "y", + "x", + "w", + "j", + "ó", + "í", + "ï", + ) + + +def _get_mms_labels(): + return ( + "a", + "i", + "e", + "n", + "o", + "u", + "t", + "s", + "r", + "m", + "k", + "l", + "d", + "g", + "h", + "y", + "b", + "p", + "w", + "c", + "v", + "j", + "z", + "f", + "'", + "q", + "x", + ) diff --git a/lib/python3.12/site-packages/torchaudio/pipelines/rnnt_pipeline.py b/lib/python3.12/site-packages/torchaudio/pipelines/rnnt_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..11b5a479f3785241a00313a85ead1405b2f673cd --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/pipelines/rnnt_pipeline.py @@ -0,0 +1,380 @@ +import json +import math +from abc import ABC, abstractmethod +from dataclasses import dataclass +from functools import partial +from typing import Callable, List, Tuple + +import torch +import torchaudio +from torchaudio._internal import module_utils +from torchaudio.models import emformer_rnnt_base, RNNT, RNNTBeamSearch + + +__all__ = [] + +_decibel = 2 * 20 * math.log10(torch.iinfo(torch.int16).max) +_gain = pow(10, 0.05 * _decibel) + + +def _piecewise_linear_log(x): + x[x > math.e] = torch.log(x[x > math.e]) + x[x <= math.e] = x[x <= math.e] / math.e + return x + + +class _FunctionalModule(torch.nn.Module): + def __init__(self, functional): + super().__init__() + self.functional = functional + + def forward(self, input): + return self.functional(input) + + +class _GlobalStatsNormalization(torch.nn.Module): + def __init__(self, global_stats_path): + super().__init__() + + with open(global_stats_path) as f: + blob = json.loads(f.read()) + + self.register_buffer("mean", torch.tensor(blob["mean"])) + self.register_buffer("invstddev", torch.tensor(blob["invstddev"])) + + def forward(self, input): + return (input - self.mean) * self.invstddev + + +class _FeatureExtractor(ABC): + @abstractmethod + def __call__(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Generates features and length output from the given input tensor. + + Args: + input (torch.Tensor): input tensor. + + Returns: + (torch.Tensor, torch.Tensor): + torch.Tensor: + Features, with shape `(length, *)`. + torch.Tensor: + Length, with shape `(1,)`. + """ + + +class _TokenProcessor(ABC): + @abstractmethod + def __call__(self, tokens: List[int], **kwargs) -> str: + """Decodes given list of tokens to text sequence. + + Args: + tokens (List[int]): list of tokens to decode. + + Returns: + str: + Decoded text sequence. + """ + + +class _ModuleFeatureExtractor(torch.nn.Module, _FeatureExtractor): + """``torch.nn.Module``-based feature extraction pipeline. + + Args: + pipeline (torch.nn.Module): module that implements feature extraction logic. + """ + + def __init__(self, pipeline: torch.nn.Module) -> None: + super().__init__() + self.pipeline = pipeline + + def forward(self, input: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Generates features and length output from the given input tensor. + + Args: + input (torch.Tensor): input tensor. + + Returns: + (torch.Tensor, torch.Tensor): + torch.Tensor: + Features, with shape `(length, *)`. + torch.Tensor: + Length, with shape `(1,)`. + """ + features = self.pipeline(input) + length = torch.tensor([features.shape[0]]) + return features, length + + +class _SentencePieceTokenProcessor(_TokenProcessor): + """SentencePiece-model-based token processor. + + Args: + sp_model_path (str): path to SentencePiece model. + """ + + def __init__(self, sp_model_path: str) -> None: + if not module_utils.is_module_available("sentencepiece"): + raise RuntimeError("SentencePiece is not available. Please install it.") + + import sentencepiece as spm + + self.sp_model = spm.SentencePieceProcessor(model_file=sp_model_path) + self.post_process_remove_list = { + self.sp_model.unk_id(), + self.sp_model.eos_id(), + self.sp_model.pad_id(), + } + + def __call__(self, tokens: List[int], lstrip: bool = True) -> str: + """Decodes given list of tokens to text sequence. + + Args: + tokens (List[int]): list of tokens to decode. + lstrip (bool, optional): if ``True``, returns text sequence with leading whitespace + removed. (Default: ``True``). + + Returns: + str: + Decoded text sequence. + """ + filtered_hypo_tokens = [ + token_index for token_index in tokens[1:] if token_index not in self.post_process_remove_list + ] + output_string = "".join(self.sp_model.id_to_piece(filtered_hypo_tokens)).replace("\u2581", " ") + + if lstrip: + return output_string.lstrip() + else: + return output_string + + +@dataclass +class RNNTBundle: + """Dataclass that bundles components for performing automatic speech recognition (ASR, speech-to-text) + inference with an RNN-T model. + + More specifically, the class provides methods that produce the featurization pipeline, + decoder wrapping the specified RNN-T model, and output token post-processor that together + constitute a complete end-to-end ASR inference pipeline that produces a text sequence + given a raw waveform. + + It can support non-streaming (full-context) inference as well as streaming inference. + + Users should not directly instantiate objects of this class; rather, users should use the + instances (representing pre-trained models) that exist within the module, + e.g. :data:`torchaudio.pipelines.EMFORMER_RNNT_BASE_LIBRISPEECH`. + + Example + >>> import torchaudio + >>> from torchaudio.pipelines import EMFORMER_RNNT_BASE_LIBRISPEECH + >>> import torch + >>> + >>> # Non-streaming inference. + >>> # Build feature extractor, decoder with RNN-T model, and token processor. + >>> feature_extractor = EMFORMER_RNNT_BASE_LIBRISPEECH.get_feature_extractor() + 100%|███████████████████████████████| 3.81k/3.81k [00:00<00:00, 4.22MB/s] + >>> decoder = EMFORMER_RNNT_BASE_LIBRISPEECH.get_decoder() + Downloading: "https://download.pytorch.org/torchaudio/models/emformer_rnnt_base_librispeech.pt" + 100%|███████████████████████████████| 293M/293M [00:07<00:00, 42.1MB/s] + >>> token_processor = EMFORMER_RNNT_BASE_LIBRISPEECH.get_token_processor() + 100%|███████████████████████████████| 295k/295k [00:00<00:00, 25.4MB/s] + >>> + >>> # Instantiate LibriSpeech dataset; retrieve waveform for first sample. + >>> dataset = torchaudio.datasets.LIBRISPEECH("/home/librispeech", url="test-clean") + >>> waveform = next(iter(dataset))[0].squeeze() + >>> + >>> with torch.no_grad(): + >>> # Produce mel-scale spectrogram features. + >>> features, length = feature_extractor(waveform) + >>> + >>> # Generate top-10 hypotheses. + >>> hypotheses = decoder(features, length, 10) + >>> + >>> # For top hypothesis, convert predicted tokens to text. + >>> text = token_processor(hypotheses[0][0]) + >>> print(text) + he hoped there would be stew for dinner turnips and carrots and bruised potatoes and fat mutton pieces to [...] + >>> + >>> + >>> # Streaming inference. + >>> hop_length = EMFORMER_RNNT_BASE_LIBRISPEECH.hop_length + >>> num_samples_segment = EMFORMER_RNNT_BASE_LIBRISPEECH.segment_length * hop_length + >>> num_samples_segment_right_context = ( + >>> num_samples_segment + EMFORMER_RNNT_BASE_LIBRISPEECH.right_context_length * hop_length + >>> ) + >>> + >>> # Build streaming inference feature extractor. + >>> streaming_feature_extractor = EMFORMER_RNNT_BASE_LIBRISPEECH.get_streaming_feature_extractor() + >>> + >>> # Process same waveform as before, this time sequentially across overlapping segments + >>> # to simulate streaming inference. Note the usage of ``streaming_feature_extractor`` and ``decoder.infer``. + >>> state, hypothesis = None, None + >>> for idx in range(0, len(waveform), num_samples_segment): + >>> segment = waveform[idx: idx + num_samples_segment_right_context] + >>> segment = torch.nn.functional.pad(segment, (0, num_samples_segment_right_context - len(segment))) + >>> with torch.no_grad(): + >>> features, length = streaming_feature_extractor(segment) + >>> hypotheses, state = decoder.infer(features, length, 10, state=state, hypothesis=hypothesis) + >>> hypothesis = hypotheses[0] + >>> transcript = token_processor(hypothesis[0]) + >>> if transcript: + >>> print(transcript, end=" ", flush=True) + he hoped there would be stew for dinner turn ips and car rots and bru 'd oes and fat mut ton pieces to [...] + """ + + class FeatureExtractor(_FeatureExtractor): + """Interface of the feature extraction part of RNN-T pipeline""" + + class TokenProcessor(_TokenProcessor): + """Interface of the token processor part of RNN-T pipeline""" + + _rnnt_path: str + _rnnt_factory_func: Callable[[], RNNT] + _global_stats_path: str + _sp_model_path: str + _right_padding: int + _blank: int + _sample_rate: int + _n_fft: int + _n_mels: int + _hop_length: int + _segment_length: int + _right_context_length: int + + def _get_model(self) -> RNNT: + model = self._rnnt_factory_func() + path = torchaudio.utils.download_asset(self._rnnt_path) + state_dict = torch.load(path) + model.load_state_dict(state_dict) + model.eval() + return model + + @property + def sample_rate(self) -> int: + """Sample rate (in cycles per second) of input waveforms. + + :type: int + """ + return self._sample_rate + + @property + def n_fft(self) -> int: + """Size of FFT window to use. + + :type: int + """ + return self._n_fft + + @property + def n_mels(self) -> int: + """Number of mel spectrogram features to extract from input waveforms. + + :type: int + """ + return self._n_mels + + @property + def hop_length(self) -> int: + """Number of samples between successive frames in input expected by model. + + :type: int + """ + return self._hop_length + + @property + def segment_length(self) -> int: + """Number of frames in segment in input expected by model. + + :type: int + """ + return self._segment_length + + @property + def right_context_length(self) -> int: + """Number of frames in right contextual block in input expected by model. + + :type: int + """ + return self._right_context_length + + def get_decoder(self) -> RNNTBeamSearch: + """Constructs RNN-T decoder. + + Returns: + RNNTBeamSearch + """ + model = self._get_model() + return RNNTBeamSearch(model, self._blank) + + def get_feature_extractor(self) -> FeatureExtractor: + """Constructs feature extractor for non-streaming (full-context) ASR. + + Returns: + FeatureExtractor + """ + local_path = torchaudio.utils.download_asset(self._global_stats_path) + return _ModuleFeatureExtractor( + torch.nn.Sequential( + torchaudio.transforms.MelSpectrogram( + sample_rate=self.sample_rate, n_fft=self.n_fft, n_mels=self.n_mels, hop_length=self.hop_length + ), + _FunctionalModule(lambda x: x.transpose(1, 0)), + _FunctionalModule(lambda x: _piecewise_linear_log(x * _gain)), + _GlobalStatsNormalization(local_path), + _FunctionalModule(lambda x: torch.nn.functional.pad(x, (0, 0, 0, self._right_padding))), + ) + ) + + def get_streaming_feature_extractor(self) -> FeatureExtractor: + """Constructs feature extractor for streaming (simultaneous) ASR. + + Returns: + FeatureExtractor + """ + local_path = torchaudio.utils.download_asset(self._global_stats_path) + return _ModuleFeatureExtractor( + torch.nn.Sequential( + torchaudio.transforms.MelSpectrogram( + sample_rate=self.sample_rate, n_fft=self.n_fft, n_mels=self.n_mels, hop_length=self.hop_length + ), + _FunctionalModule(lambda x: x.transpose(1, 0)), + _FunctionalModule(lambda x: _piecewise_linear_log(x * _gain)), + _GlobalStatsNormalization(local_path), + ) + ) + + def get_token_processor(self) -> TokenProcessor: + """Constructs token processor. + + Returns: + TokenProcessor + """ + local_path = torchaudio.utils.download_asset(self._sp_model_path) + return _SentencePieceTokenProcessor(local_path) + + +EMFORMER_RNNT_BASE_LIBRISPEECH = RNNTBundle( + _rnnt_path="models/emformer_rnnt_base_librispeech.pt", + _rnnt_factory_func=partial(emformer_rnnt_base, num_symbols=4097), + _global_stats_path="pipeline-assets/global_stats_rnnt_librispeech.json", + _sp_model_path="pipeline-assets/spm_bpe_4096_librispeech.model", + _right_padding=4, + _blank=4096, + _sample_rate=16000, + _n_fft=400, + _n_mels=80, + _hop_length=160, + _segment_length=16, + _right_context_length=4, +) +EMFORMER_RNNT_BASE_LIBRISPEECH.__doc__ = """ASR pipeline based on Emformer-RNNT, +pretrained on *LibriSpeech* dataset :cite:`7178964`, +capable of performing both streaming and non-streaming inference. + +The underlying model is constructed by :py:func:`torchaudio.models.emformer_rnnt_base` +and utilizes weights trained on LibriSpeech using training script ``train.py`` +`here `__ with default arguments. + +Please refer to :py:class:`RNNTBundle` for usage instructions. +""" diff --git a/lib/python3.12/site-packages/torchaudio/prototype/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/lib/python3.12/site-packages/torchaudio/prototype/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..7433182d1ca313ee931fc7386452f76bc638b442 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/datasets/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/datasets/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0e4a6194f48027caf10f3dcbbada53719a14d4a2 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/datasets/__init__.py @@ -0,0 +1,4 @@ +from .musan import Musan + + +__all__ = ["Musan"] diff --git a/lib/python3.12/site-packages/torchaudio/prototype/datasets/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/datasets/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..784b03406538bfef5f4fd83e95adfcfce1713727 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/datasets/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/datasets/__pycache__/musan.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/datasets/__pycache__/musan.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e1190ac5c81072a84d96817b44d9e7762b15a0e7 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/datasets/__pycache__/musan.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/datasets/musan.py b/lib/python3.12/site-packages/torchaudio/prototype/datasets/musan.py new file mode 100644 index 0000000000000000000000000000000000000000..c4592bb3e4097f51064bfac01467873ba7263ec8 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/datasets/musan.py @@ -0,0 +1,67 @@ +from pathlib import Path +from typing import Tuple, Union + +import torch +from torch.utils.data import Dataset +from torchaudio.datasets.utils import _load_waveform + + +_SUBSETS = ["music", "noise", "speech"] +_SAMPLE_RATE = 16_000 + + +class Musan(Dataset): + r"""*MUSAN* :cite:`musan2015` dataset. + + Args: + root (str or Path): Root directory where the dataset's top-level directory exists. + subset (str): Subset of the dataset to use. Options: [``"music"``, ``"noise"``, ``"speech"``]. + """ + + def __init__(self, root: Union[str, Path], subset: str): + if subset not in _SUBSETS: + raise ValueError(f"Invalid subset '{subset}' given. Please provide one of {_SUBSETS}") + + subset_path = Path(root) / subset + self._walker = [str(p) for p in subset_path.glob("*/*.*")] + + def get_metadata(self, n: int) -> Tuple[str, int, str]: + r"""Get metadata for the n-th sample in the dataset. Returns filepath instead of waveform, + but otherwise returns the same fields as :py:func:`__getitem__`. + + Args: + n (int): Index of sample to be loaded. + + Returns: + (str, int, str): + str + Path to audio. + int + Sample rate. + str + File name. + """ + audio_path = self._walker[n] + return audio_path, _SAMPLE_RATE, Path(audio_path).name + + def __getitem__(self, n: int) -> Tuple[torch.Tensor, int, str]: + r"""Return the n-th sample in the dataset. + + Args: + n (int): Index of sample to be loaded. + + Returns: + (torch.Tensor, int, str): + torch.Tensor + Waveform. + int + Sample rate. + str + File name. + """ + audio_path, sample_rate, filename = self.get_metadata(n) + path = Path(audio_path) + return _load_waveform(path.parent, path.name, sample_rate), sample_rate, filename + + def __len__(self) -> int: + return len(self._walker) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/functional/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..20bc181731eba87faeb77a36e7e1cdce4101f496 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/functional/__init__.py @@ -0,0 +1,26 @@ +from ._dsp import ( + adsr_envelope, + exp_sigmoid, + extend_pitch, + filter_waveform, + frequency_impulse_response, + oscillator_bank, + sinc_impulse_response, +) +from ._rir import ray_tracing, simulate_rir_ism +from .functional import barkscale_fbanks, chroma_filterbank + + +__all__ = [ + "adsr_envelope", + "exp_sigmoid", + "barkscale_fbanks", + "chroma_filterbank", + "extend_pitch", + "filter_waveform", + "frequency_impulse_response", + "oscillator_bank", + "ray_tracing", + "sinc_impulse_response", + "simulate_rir_ism", +] diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..9156d9342a71e91976f6e3bfef93329b11fa6d8a Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/_dsp.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/_dsp.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..52fe81250d91e31fd88175826dcb0ee6c6ef99d0 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/_dsp.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/_rir.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/_rir.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..3a4a0e84de8ab085b8089f5c474a3d44cf5524db Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/_rir.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/functional.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/functional.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..ce0e9504175a42ce4daa74c200c2d87e3fa1e1ce Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/functional/__pycache__/functional.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/_dsp.py b/lib/python3.12/site-packages/torchaudio/prototype/functional/_dsp.py new file mode 100644 index 0000000000000000000000000000000000000000..72b1a153f57eaec1b464ad42199cf6f6e331ae26 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/functional/_dsp.py @@ -0,0 +1,433 @@ +import warnings +from typing import List, Optional, Union + +import torch + +from torchaudio.functional import fftconvolve + + +def oscillator_bank( + frequencies: torch.Tensor, + amplitudes: torch.Tensor, + sample_rate: float, + reduction: str = "sum", + dtype: Optional[torch.dtype] = torch.float64, +) -> torch.Tensor: + """Synthesize waveform from the given instantaneous frequencies and amplitudes. + + .. devices:: CPU CUDA + + .. properties:: Autograd TorchScript + + Note: + The phase information of the output waveform is found by taking the cumulative sum + of the given instantaneous frequencies (``frequencies``). + This incurs roundoff error when the data type does not have enough precision. + Using ``torch.float64`` can work around this. + + The following figure shows the difference between ``torch.float32`` and + ``torch.float64`` when generating a sin wave of constant frequency and amplitude + with sample rate 8000 [Hz]. + Notice that ``torch.float32`` version shows artifacts that are not seen in + ``torch.float64`` version. + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/oscillator_precision.png + + Args: + frequencies (Tensor): Sample-wise oscillator frequencies (Hz). Shape `(..., time, N)`. + amplitudes (Tensor): Sample-wise oscillator amplitude. Shape: `(..., time, N)`. + sample_rate (float): Sample rate + reduction (str): Reduction to perform. + Valid values are ``"sum"``, ``"mean"`` or ``"none"``. Default: ``"sum"`` + dtype (torch.dtype or None, optional): The data type on which cumulative sum operation is performed. + Default: ``torch.float64``. Pass ``None`` to disable the casting. + + Returns: + Tensor: + The resulting waveform. + + If ``reduction`` is ``"none"``, then the shape is + `(..., time, N)`, otherwise the shape is `(..., time)`. + """ + if frequencies.shape != amplitudes.shape: + raise ValueError( + "The shapes of `frequencies` and `amplitudes` must match. " + f"Found: {frequencies.shape} and {amplitudes.shape} respectively." + ) + reductions = ["sum", "mean", "none"] + if reduction not in reductions: + raise ValueError(f"The value of reduction must be either {reductions}. Found: {reduction}") + + invalid = torch.abs(frequencies) >= sample_rate / 2 + if torch.any(invalid): + warnings.warn( + "Some frequencies are above nyquist frequency. " + "Setting the corresponding amplitude to zero. " + "This might cause numerically unstable gradient." + ) + amplitudes = torch.where(invalid, 0.0, amplitudes) + + pi2 = 2.0 * torch.pi + freqs = frequencies * pi2 / sample_rate % pi2 + phases = torch.cumsum(freqs, dim=-2, dtype=dtype) + if dtype is not None and freqs.dtype != dtype: + phases = phases.to(freqs.dtype) + + waveform = amplitudes * torch.sin(phases) + if reduction == "sum": + return waveform.sum(-1) + if reduction == "mean": + return waveform.mean(-1) + return waveform + + +def adsr_envelope( + num_frames: int, + *, + attack: float = 0.0, + hold: float = 0.0, + decay: float = 0.0, + sustain: float = 1.0, + release: float = 0.0, + n_decay: int = 2, + dtype: Optional[torch.dtype] = None, + device: Optional[torch.device] = None, +): + """Generate ADSR Envelope + + .. devices:: CPU CUDA + + Args: + num_frames (int): The number of output frames. + attack (float, optional): + The relative *time* it takes to reach the maximum level from + the start. (Default: ``0.0``) + hold (float, optional): + The relative *time* the maximum level is held before + it starts to decay. (Default: ``0.0``) + decay (float, optional): + The relative *time* it takes to sustain from + the maximum level. (Default: ``0.0``) + sustain (float, optional): The relative *level* at which + the sound should sustain. (Default: ``1.0``) + + .. Note:: + The duration of sustain is derived as `1.0 - (The sum of attack, hold, decay and release)`. + + release (float, optional): The relative *time* it takes for the sound level to + reach zero after the sustain. (Default: ``0.0``) + n_decay (int, optional): The degree of polynomial decay. Default: ``2``. + dtype (torch.dtype, optional): the desired data type of returned tensor. + Default: if ``None``, uses a global default + (see :py:func:`torch.set_default_tensor_type`). + device (torch.device, optional): the desired device of returned tensor. + Default: if ``None``, uses the current device for the default tensor type + (see :py:func:`torch.set_default_tensor_type`). + device will be the CPU for CPU tensor types and the current CUDA + device for CUDA tensor types. + + Returns: + Tensor: ADSR Envelope. Shape: `(num_frames, )` + + Example + .. image:: https://download.pytorch.org/torchaudio/doc-assets/adsr_examples.png + + """ + if not 0 <= attack <= 1: + raise ValueError(f"The value of `attack` must be within [0, 1]. Found: {attack}") + if not 0 <= decay <= 1: + raise ValueError(f"The value of `decay` must be within [0, 1]. Found: {decay}") + if not 0 <= sustain <= 1: + raise ValueError(f"The value of `sustain` must be within [0, 1]. Found: {sustain}") + if not 0 <= hold <= 1: + raise ValueError(f"The value of `hold` must be within [0, 1]. Found: {hold}") + if not 0 <= release <= 1: + raise ValueError(f"The value of `release` must be within [0, 1]. Found: {release}") + if attack + decay + release + hold > 1: + raise ValueError("The sum of `attack`, `hold`, `decay` and `release` must not exceed 1.") + + nframes = num_frames - 1 + num_a = int(nframes * attack) + num_h = int(nframes * hold) + num_d = int(nframes * decay) + num_r = int(nframes * release) + + # Initialize with sustain + out = torch.full((num_frames,), float(sustain), device=device, dtype=dtype) + + # attack + if num_a > 0: + torch.linspace(0.0, 1.0, num_a + 1, out=out[: num_a + 1]) + + # hold + if num_h > 0: + out[num_a : num_a + num_h + 1] = 1.0 + + # decay + if num_d > 0: + # Compute: sustain + (1.0 - sustain) * (linspace[1, 0] ** n_decay) + i = num_a + num_h + decay = out[i : i + num_d + 1] + torch.linspace(1.0, 0.0, num_d + 1, out=decay) + decay **= n_decay + decay *= 1.0 - sustain + decay += sustain + + # sustain is handled by initialization + + # release + if num_r > 0: + torch.linspace(sustain, 0, num_r + 1, out=out[-num_r - 1 :]) + + return out + + +def extend_pitch( + base: torch.Tensor, + pattern: Union[int, List[float], torch.Tensor], +): + """Extend the given time series values with multipliers of them. + + .. devices:: CPU CUDA + + .. properties:: Autograd TorchScript + + Given a series of fundamental frequencies (pitch), this function appends + its harmonic overtones or inharmonic partials. + + Args: + base (torch.Tensor): + Base time series, like fundamental frequencies (Hz). Shape: `(..., time, 1)`. + pattern (int, list of floats or torch.Tensor): + If ``int``, the number of pitch series after the operation. + `pattern - 1` tones are added, so that the resulting Tensor contains + up to `pattern`-th overtones of the given series. + + If list of float or ``torch.Tensor``, it must be one dimensional, + representing the custom multiplier of the fundamental frequency. + + Returns: + Tensor: Oscillator frequencies (Hz). Shape: `(..., time, num_tones)`. + + Example + >>> # fundamental frequency + >>> f0 = torch.linspace(1, 5, 5).unsqueeze(-1) + >>> f0 + tensor([[1.], + [2.], + [3.], + [4.], + [5.]]) + >>> # Add harmonic overtones, up to 3rd. + >>> f = extend_pitch(f0, 3) + >>> f.shape + torch.Size([5, 3]) + >>> f + tensor([[ 1., 2., 3.], + [ 2., 4., 6.], + [ 3., 6., 9.], + [ 4., 8., 12.], + [ 5., 10., 15.]]) + >>> # Add custom (inharmonic) partials. + >>> f = extend_pitch(f0, torch.tensor([1, 2.1, 3.3, 4.5])) + >>> f.shape + torch.Size([5, 4]) + >>> f + tensor([[ 1.0000, 2.1000, 3.3000, 4.5000], + [ 2.0000, 4.2000, 6.6000, 9.0000], + [ 3.0000, 6.3000, 9.9000, 13.5000], + [ 4.0000, 8.4000, 13.2000, 18.0000], + [ 5.0000, 10.5000, 16.5000, 22.5000]]) + """ + if isinstance(pattern, torch.Tensor): + mult = pattern + elif isinstance(pattern, int): + mult = torch.linspace(1.0, float(pattern), pattern, device=base.device, dtype=base.dtype) + else: + mult = torch.tensor(pattern, dtype=base.dtype, device=base.device) + h_freq = base @ mult.unsqueeze(0) + return h_freq + + +def sinc_impulse_response(cutoff: torch.Tensor, window_size: int = 513, high_pass: bool = False): + """Create windowed-sinc impulse response for given cutoff frequencies. + + .. devices:: CPU CUDA + + .. properties:: Autograd TorchScript + + Args: + cutoff (Tensor): Cutoff frequencies for low-pass sinc filter. + + window_size (int, optional): Size of the Hamming window to apply. Must be odd. + (Default: 513) + + high_pass (bool, optional): + If ``True``, convert the resulting filter to high-pass. + Otherwise low-pass filter is returned. Default: ``False``. + + Returns: + Tensor: A series of impulse responses. Shape: `(..., window_size)`. + """ + if window_size % 2 == 0: + raise ValueError(f"`window_size` must be odd. Given: {window_size}") + + half = window_size // 2 + device, dtype = cutoff.device, cutoff.dtype + idx = torch.linspace(-half, half, window_size, device=device, dtype=dtype) + + filt = torch.special.sinc(cutoff.unsqueeze(-1) * idx.unsqueeze(0)) + filt = filt * torch.hamming_window(window_size, device=device, dtype=dtype, periodic=False).unsqueeze(0) + filt = filt / filt.sum(dim=-1, keepdim=True).abs() + + # High pass IR is obtained by subtracting low_pass IR from delta function. + # https://courses.engr.illinois.edu/ece401/fa2020/slides/lec10.pdf + if high_pass: + filt = -filt + filt[..., half] = 1.0 + filt[..., half] + return filt + + +def frequency_impulse_response(magnitudes): + """Create filter from desired frequency response + + Args: + magnitudes: The desired frequency responses. Shape: `(..., num_fft_bins)` + + Returns: + Tensor: Impulse response. Shape `(..., 2 * (num_fft_bins - 1))` + """ + if magnitudes.min() < 0.0: + # Negative magnitude does not make sense but allowing so that autograd works + # around 0. + # Should we raise error? + warnings.warn("The input frequency response should not contain negative values.") + ir = torch.fft.fftshift(torch.fft.irfft(magnitudes), dim=-1) + device, dtype = magnitudes.device, magnitudes.dtype + window = torch.hann_window(ir.size(-1), periodic=False, device=device, dtype=dtype).expand_as(ir) + return ir * window + + +def _overlap_and_add(waveform, stride): + num_frames, frame_size = waveform.shape[-2:] + numel = (num_frames - 1) * stride + frame_size + buffer = torch.zeros(waveform.shape[:-2] + (numel,), device=waveform.device, dtype=waveform.dtype) + for i in range(num_frames): + start = i * stride + end = start + frame_size + buffer[..., start:end] += waveform[..., i, :] + return buffer + + +def filter_waveform(waveform: torch.Tensor, kernels: torch.Tensor, delay_compensation: int = -1): + """Applies filters along time axis of the given waveform. + + This function applies the given filters along time axis in the following manner: + + 1. Split the given waveform into chunks. The number of chunks is equal to the number of given filters. + 2. Filter each chunk with corresponding filter. + 3. Place the filtered chunks at the original indices while adding up the overlapping parts. + 4. Crop the resulting waveform so that delay introduced by the filter is removed and its length + matches that of the input waveform. + + The following figure illustrates this. + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/filter_waveform.png + + .. note:: + + If the number of filters is one, then the operation becomes stationary. + i.e. the same filtering is applied across the time axis. + + Args: + waveform (Tensor): Shape `(..., time)`. + kernels (Tensor): Impulse responses. + Valid inputs are 2D tensor with shape `(num_filters, filter_length)` or + `(N+1)`-D tensor with shape `(..., num_filters, filter_length)`, where `N` is + the dimension of waveform. + + In case of 2D input, the same set of filters is used across channels and batches. + Otherwise, different sets of filters are applied. In this case, the shape of + the first `N-1` dimensions of filters must match (or be broadcastable to) that of waveform. + + delay_compensation (int): Control how the waveform is cropped after full convolution. + If the value is zero or positive, it is interpreted as the length of crop at the + beginning of the waveform. The value cannot be larger than the size of filter kernel. + Otherwise the initial crop is ``filter_size // 2``. + When cropping happens, the waveform is also cropped from the end so that the + length of the resulting waveform matches the input waveform. + + Returns: + Tensor: `(..., time)`. + """ + if kernels.ndim not in [2, waveform.ndim + 1]: + raise ValueError( + "`kernels` must be 2 or N+1 dimension where " + f"N is the dimension of waveform. Found: {kernels.ndim} (N={waveform.ndim})" + ) + + num_filters, filter_size = kernels.shape[-2:] + num_frames = waveform.size(-1) + + if delay_compensation > filter_size: + raise ValueError( + "When `delay_compenstation` is provided, it cannot be larger than the size of filters." + f"Found: delay_compensation={delay_compensation}, filter_size={filter_size}" + ) + + # Transform waveform's time axis into (num_filters x chunk_length) with optional padding + chunk_length = num_frames // num_filters + if num_frames % num_filters > 0: + chunk_length += 1 + num_pad = chunk_length * num_filters - num_frames + waveform = torch.nn.functional.pad(waveform, [0, num_pad], "constant", 0) + chunked = waveform.unfold(-1, chunk_length, chunk_length) + assert chunked.numel() >= waveform.numel() + + # Broadcast kernels + if waveform.ndim + 1 > kernels.ndim: + expand_shape = waveform.shape[:-1] + kernels.shape + kernels = kernels.expand(expand_shape) + + convolved = fftconvolve(chunked, kernels) + restored = _overlap_and_add(convolved, chunk_length) + + # Trim in a way that the number of samples are same as input, + # and the filter delay is compensated + if delay_compensation >= 0: + start = delay_compensation + else: + start = filter_size // 2 + num_crops = restored.size(-1) - num_frames + end = num_crops - start + result = restored[..., start:-end] + return result + + +def exp_sigmoid( + input: torch.Tensor, exponent: float = 10.0, max_value: float = 2.0, threshold: float = 1e-7 +) -> torch.Tensor: + """Exponential Sigmoid pointwise nonlinearity. + Implements the equation: + ``max_value`` * sigmoid(``input``) ** (log(``exponent``)) + ``threshold`` + + The output has a range of [``threshold``, ``max_value``]. + ``exponent`` controls the slope of the output. + + .. devices:: CPU CUDA + + Args: + input (Tensor): Input Tensor + exponent (float, optional): Exponent. Controls the slope of the output + max_value (float, optional): Maximum value of the output + threshold (float, optional): Minimum value of the output + + Returns: + Tensor: Exponential Sigmoid output. Shape: same as input + + """ + + return max_value * torch.pow( + torch.nn.functional.sigmoid(input), + torch.log(torch.tensor(exponent, device=input.device, dtype=input.dtype)), + ) + torch.tensor(threshold, device=input.device, dtype=input.dtype) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/_rir.py b/lib/python3.12/site-packages/torchaudio/prototype/functional/_rir.py new file mode 100644 index 0000000000000000000000000000000000000000..0e67a5494d204182d83cc09166064ea9d4355176 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/functional/_rir.py @@ -0,0 +1,379 @@ +import math +from typing import Optional, Tuple, Union + +import torch +import torchaudio +from torch import Tensor + + +def _compute_image_sources( + room: torch.Tensor, + source: torch.Tensor, + max_order: int, + absorption: torch.Tensor, + scatter: Optional[torch.Tensor] = None, +) -> Tuple[Tensor, Tensor]: + """Compute image sources in a shoebox-like room. + + Args: + room (torch.Tensor): The 1D Tensor to determine the room size. The shape is + `(D,)`, where ``D`` is 2 if room is a 2D room, or 3 if room is a 3D room. + source (torch.Tensor): The coordinate of the sound source. Tensor with dimensions + `(D)`. + max_order (int): The maximum number of reflections of the source. + absorption (torch.Tensor): The absorption coefficients of wall materials. + ``absorption`` is a Tensor with dimensions `(num_band, num_wall)`. + The shape options are ``[(1, 4), (1, 6), (7, 4), (7, 6)]``. + ``num_band`` is `1` if the coefficients is the same for all frequencies, or is `7` + if the coefficients are different to different frequencies. `7` refers to the default number + of octave bands. (See note in `simulate_rir_ism` method). + ``num_wall`` is `4` if the room is a 2D room, representing absorption coefficients + of ``"west"``, ``"east"``, ``"south"``, and ``"north"`` walls, respectively. + Or it is `6` if the room is a 3D room, representing absorption coefficients + of ``"west"``, ``"east"``, ``"south"``, ``"north"``, ``"floor"``, and ``"ceiling"``, respectively. + scatter (torch.Tensor): The scattering coefficients of wall materials. + The shape of ``scatter`` must match that of ``absorption``. If ``None``, it is not + used in image source computation. (Default: ``None``) + + Returns: + (torch.Tensor): The coordinates of all image sources within ``max_order`` number of reflections. + Tensor with dimensions `(num_image_source, D)`. + (torch.Tensor): The attenuation of corresponding image sources. Tensor with dimensions + `(num_band, num_image_source)`. + """ + if scatter is None: + tr = torch.sqrt(1 - absorption) + else: + tr = torch.sqrt(1 - absorption) * torch.sqrt(1 - scatter) + + ind = torch.arange(-max_order, max_order + 1, device=source.device) + if room.shape[0] == 2: + XYZ = torch.meshgrid(ind, ind, indexing="ij") + else: + XYZ = torch.meshgrid(ind, ind, ind, indexing="ij") + XYZ = torch.stack([c.reshape((-1,)) for c in XYZ], dim=-1) + XYZ = XYZ[XYZ.abs().sum(dim=-1) <= max_order] + + # compute locations of image sources + d = room[None, :] + s = source[None, :] + img_loc = torch.where(XYZ % 2 == 1, d * (XYZ + 1) - s, d * XYZ + s) + + # attenuation + exp_lo = abs(torch.floor((XYZ / 2))) + exp_hi = abs(torch.floor((XYZ + 1) / 2)) + t_lo = tr[:, ::2].unsqueeze(1).repeat(1, XYZ.shape[0], 1) # (num_band, left walls) + t_hi = tr[:, 1::2].unsqueeze(1).repeat(1, XYZ.shape[0], 1) # (num_band, right walls) + att = torch.prod((t_lo**exp_lo) * (t_hi**exp_hi), dim=-1) # (num_band, num_image_source) + return img_loc, att + + +def _hann(x: torch.Tensor, T: int): + """Compute the Hann window where the values are truncated based on window length. + torch.hann_window can only sample window function at integer points, the method is to sample + continuous window function at non-integer points. + + Args: + x (torch.Tensor): The fractional component of time delay Tensor. + T (torch.Tensor): The window length of sinc function. + + Returns: + (torch.Tensor): The hann window Tensor where values outside + the sinc window (`T`) is set to zero. + """ + y = torch.where( + torch.abs(x) <= T / 2, + 0.5 * (1 + torch.cos(2 * math.pi * x / T)), + x.new_zeros(1), + ) + return y + + +def _frac_delay(delay: torch.Tensor, delay_i: torch.Tensor, delay_filter_length: int): + """Compute fractional delay of impulse response signal. + + Args: + delay (torch.Tensor): The time delay Tensor in samples. + delay_i (torch.Tensor): The integer part of delay. + delay_filter_length (int): The window length for sinc function. + + Returns: + (torch.Tensor): The impulse response Tensor for all image sources. + """ + if delay_filter_length % 2 != 1: + raise ValueError("The filter length must be odd") + + pad = delay_filter_length // 2 + n = torch.arange(-pad, pad + 1, device=delay.device) + delay_i[..., None] + delay = delay[..., None] + + return torch.special.sinc(n - delay) * _hann(n - delay, 2 * pad) + + +def _adjust_coeff(coeffs: Union[float, torch.Tensor], name: str) -> torch.Tensor: + """Validates and converts absorption or scattering parameters to a tensor with appropriate shape + + Args: + coeff (float or torch.Tensor): The absorption coefficients of wall materials. + + If the dtype is ``float``, the absorption coefficient is identical for all walls and + all frequencies. + + If ``absorption`` is a 1D Tensor, the shape must be `(2*dim,)`, + where the values represent absorption coefficients of ``"west"``, ``"east"``, + ``"south"``, ``"north"``, ``"floor"``, and ``"ceiling"``, respectively. + + If ``absorption`` is a 2D Tensor, the shape must be `(7, 2*dim)`, + where 7 represents the number of octave bands. + + Returns: + (torch.Tensor): The expanded coefficient. + The shape is `(1, 6)` for single octave band case, and + `(7, 6)` for multi octave band case. + """ + num_walls = 6 + if isinstance(coeffs, float): + if coeffs < 0: + raise ValueError(f"`{name}` must be non-negative. Found: {coeffs}") + return torch.full((1, num_walls), coeffs) + if isinstance(coeffs, Tensor): + if torch.any(coeffs < 0): + raise ValueError(f"`{name}` must be non-negative. Found: {coeffs}") + if coeffs.ndim == 1: + if coeffs.numel() != num_walls: + raise ValueError( + f"The shape of `{name}` must be ({num_walls},) when it is a 1D Tensor. " + f"Found the shape {coeffs.shape}." + ) + return coeffs.unsqueeze(0) + if coeffs.ndim == 2: + if coeffs.shape[1] != num_walls: + raise ValueError( + f"The shape of `{name}` must be (NUM_BANDS, {num_walls}) when it " + f"is a 2D Tensor. Found: {coeffs.shape}." + ) + return coeffs + raise TypeError(f"`{name}` must be float or Tensor.") + + +def _validate_inputs( + room: torch.Tensor, + source: torch.Tensor, + mic_array: torch.Tensor, +): + """Validate dimensions of input arguments, and normalize different kinds of absorption into the same dimension. + + Args: + room (torch.Tensor): The size of the room. width, length (and height) + source (torch.Tensor): Sound source coordinates. Tensor with dimensions `(dim,)`. + mic_array (torch.Tensor): Microphone coordinates. Tensor with dimensions `(channel, dim)`. + """ + if not (room.ndim == 1 and room.numel() == 3): + raise ValueError(f"`room` must be a 1D Tensor with 3 elements. Found {room.shape}.") + if not (source.ndim == 1 and source.numel() == 3): + raise ValueError(f"`source` must be 1D Tensor with 3 elements. Found {source.shape}.") + if not (mic_array.ndim == 2 and mic_array.shape[1] == 3): + raise ValueError(f"`mic_array` must be a 2D Tensor with shape (num_channels, 3). Found {mic_array.shape}.") + + +def simulate_rir_ism( + room: torch.Tensor, + source: torch.Tensor, + mic_array: torch.Tensor, + max_order: int, + absorption: Union[float, torch.Tensor], + output_length: Optional[int] = None, + delay_filter_length: int = 81, + center_frequency: Optional[torch.Tensor] = None, + sound_speed: float = 343.0, + sample_rate: float = 16000.0, +) -> Tensor: + r"""Compute Room Impulse Response (RIR) based on the *image source method* :cite:`allen1979image`. + The implementation is based on *pyroomacoustics* :cite:`scheibler2018pyroomacoustics`. + + .. devices:: CPU + + .. properties:: TorchScript + + Args: + room (torch.Tensor): Room coordinates. The shape of `room` must be `(3,)` which represents + three dimensions of the room. + source (torch.Tensor): Sound source coordinates. Tensor with dimensions `(3,)`. + mic_array (torch.Tensor): Microphone coordinates. Tensor with dimensions `(channel, 3)`. + max_order (int): The maximum number of reflections of the source. + absorption (float or torch.Tensor): The *absorption* :cite:`wiki:Absorption_(acoustics)` + coefficients of wall materials for sound energy. + If the dtype is ``float``, the absorption coefficient is identical for all walls and + all frequencies. + If ``absorption`` is a 1D Tensor, the shape must be `(6,)`, where the values represent + absorption coefficients of ``"west"``, ``"east"``, ``"south"``, ``"north"``, ``"floor"``, + and ``"ceiling"``, respectively. + If ``absorption`` is a 2D Tensor, the shape must be `(7, 6)`, where 7 represents the number of octave bands. + output_length (int or None, optional): The output length of simulated RIR signal. If ``None``, + the length is defined as + + .. math:: + \frac{\text{max\_d} \cdot \text{sample\_rate}}{\text{sound\_speed}} + \text{delay\_filter\_length} + + where ``max_d`` is the maximum distance between image sources and microphones. + delay_filter_length (int, optional): The filter length for computing sinc function. (Default: ``81``) + center_frequency (torch.Tensor, optional): The center frequencies of octave bands for multi-band walls. + Only used when ``absorption`` is a 2D Tensor. + sound_speed (float, optional): The speed of sound. (Default: ``343.0``) + sample_rate (float, optional): The sample rate of the generated room impulse response signal. + (Default: ``16000.0``) + + Returns: + (torch.Tensor): The simulated room impulse response waveform. Tensor with dimensions + `(channel, rir_length)`. + + Note: + If ``absorption`` is a 2D Tensor and ``center_frequency`` is set to ``None``, the center frequencies + of octave bands are fixed to ``[125.0, 250.0, 500.0, 1000.0, 2000.0, 4000.0, 8000.0]``. + Users need to tune the values of ``absorption`` to the corresponding frequencies. + """ + _validate_inputs(room, source, mic_array) + absorption = _adjust_coeff(absorption, "absorption") + img_location, att = _compute_image_sources(room, source, max_order, absorption) + + # compute distances between image sources and microphones + vec = img_location[:, None, :] - mic_array[None, :, :] + dist = torch.linalg.norm(vec, dim=-1) # (image_source, channel) + + img_src_att = att[..., None] / dist[None, ...] # (band, image_source, channel) + + # separate delays in integer / frac part + delay = dist * sample_rate / sound_speed # distance to delay in samples + delay_i = torch.ceil(delay) # integer part + + # compute the shorts IRs corresponding to each image source + irs = img_src_att[..., None] * _frac_delay(delay, delay_i, delay_filter_length)[None, ...] + + rir_length = int(delay_i.max() + irs.shape[-1]) + rir = torch.ops.torchaudio._simulate_rir(irs, delay_i.type(torch.int32), rir_length) + + # multi-band processing + if absorption.shape[0] > 1: + if center_frequency is None: + center = torch.tensor( + [125.0, 250.0, 500.0, 1000.0, 2000.0, 4000.0, 8000.0], dtype=room.dtype, device=room.device + ) + else: + center = center_frequency + # n_fft is set to 512 by default. + filters = torch.ops.torchaudio._make_rir_filter(center, sample_rate, n_fft=512) + rir = torchaudio.functional.fftconvolve(rir, filters.unsqueeze(1).repeat(1, rir.shape[1], 1), mode="same") + + # sum up rir signals of all image sources into one waveform. + rir = rir.sum(0) + + if output_length is not None: + if output_length > rir.shape[-1]: + rir = torch.nn.functional.pad(rir, (0, output_length - rir.shape[-1]), "constant", 0.0) + else: + rir = rir[..., :output_length] + + return rir + + +def ray_tracing( + room: torch.Tensor, + source: torch.Tensor, + mic_array: torch.Tensor, + num_rays: int, + absorption: Union[float, torch.Tensor] = 0.0, + scattering: Union[float, torch.Tensor] = 0.0, + mic_radius: float = 0.5, + sound_speed: float = 343.0, + energy_thres: float = 1e-7, + time_thres: float = 10.0, + hist_bin_size: float = 0.004, +) -> torch.Tensor: + r"""Compute energy histogram via ray tracing. + + The implementation is based on *pyroomacoustics* :cite:`scheibler2018pyroomacoustics`. + + ``num_rays`` rays are casted uniformly in all directions from the source; + when a ray intersects a wall, it is reflected and part of its energy is absorbed. + It is also scattered (sent directly to the microphone(s)) according to the ``scattering`` + coefficient. + When a ray is close to the microphone, its current energy is recorded in the output + histogram for that given time slot. + + .. devices:: CPU + + .. properties:: TorchScript + + Args: + room (torch.Tensor): Room coordinates. The shape of `room` must be `(3,)` which represents + three dimensions of the room. + source (torch.Tensor): Sound source coordinates. Tensor with dimensions `(3,)`. + mic_array (torch.Tensor): Microphone coordinates. Tensor with dimensions `(channel, 3)`. + absorption (float or torch.Tensor, optional): The absorption coefficients of wall materials. + (Default: ``0.0``). + If the type is ``float``, the absorption coefficient is identical to all walls and + all frequencies. + If ``absorption`` is a 1D Tensor, the shape must be `(6,)`, representing absorption + coefficients of ``"west"``, ``"east"``, ``"south"``, ``"north"``, ``"floor"``, and + ``"ceiling"``, respectively. + If ``absorption`` is a 2D Tensor, the shape must be `(num_bands, 6)`. + ``num_bands`` is the number of frequency bands (usually 7). + scattering(float or torch.Tensor, optional): The scattering coefficients of wall materials. (Default: ``0.0``) + The shape and type of this parameter is the same as for ``absorption``. + mic_radius(float, optional): The radius of the microphone in meters. (Default: 0.5) + sound_speed (float, optional): The speed of sound in meters per second. (Default: ``343.0``) + energy_thres (float, optional): The energy level below which we stop tracing a ray. (Default: ``1e-7``) + The initial energy of each ray is ``2 / num_rays``. + time_thres (float, optional): The maximal duration for which rays are traced. (Unit: seconds) (Default: 10.0) + hist_bin_size (float, optional): The size of each bin in the output histogram. (Unit: seconds) (Default: 0.004) + + Returns: + (torch.Tensor): The 3D histogram(s) where the energy of the traced ray is recorded. + Each bin corresponds to a given time slot. + The shape is `(channel, num_bands, num_bins)`, where + ``num_bins = ceil(time_thres / hist_bin_size)``. + If both ``absorption`` and ``scattering`` are floats, then ``num_bands == 1``. + """ + if time_thres < hist_bin_size: + raise ValueError( + "`time_thres` must be greater than `hist_bin_size`. " + f"Found: hist_bin_size={hist_bin_size}, time_thres={time_thres}." + ) + + if room.dtype != source.dtype or source.dtype != mic_array.dtype: + raise ValueError( + "dtype of `room`, `source` and `mic_array` must match. " + f"Found: `room` ({room.dtype}), `source` ({source.dtype}) and " + f"`mic_array` ({mic_array.dtype})" + ) + + _validate_inputs(room, source, mic_array) + absorption = _adjust_coeff(absorption, "absorption").to(room.dtype) + scattering = _adjust_coeff(scattering, "scattering").to(room.dtype) + + # Bring absorption and scattering to the same shape + if absorption.shape[0] == 1 and scattering.shape[0] > 1: + absorption = absorption.expand(scattering.shape) + if scattering.shape[0] == 1 and absorption.shape[0] > 1: + scattering = scattering.expand(absorption.shape) + if absorption.shape != scattering.shape: + raise ValueError( + "`absorption` and `scattering` must be broadcastable to the same number of bands and walls. " + f"Inferred shapes absorption={absorption.shape} and scattering={scattering.shape}" + ) + + histograms = torch.ops.torchaudio.ray_tracing( + room, + source, + mic_array, + num_rays, + absorption, + scattering, + mic_radius, + sound_speed, + energy_thres, + time_thres, + hist_bin_size, + ) + + return histograms diff --git a/lib/python3.12/site-packages/torchaudio/prototype/functional/functional.py b/lib/python3.12/site-packages/torchaudio/prototype/functional/functional.py new file mode 100644 index 0000000000000000000000000000000000000000..0805a252af4ef6946606a32a532188cd937321b8 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/functional/functional.py @@ -0,0 +1,190 @@ +import math +import warnings +from typing import Optional + +import torch +from torchaudio.functional.functional import _create_triangular_filterbank + + +def _hz_to_bark(freqs: float, bark_scale: str = "traunmuller") -> float: + r"""Convert Hz to Barks. + + Args: + freqs (float): Frequencies in Hz + bark_scale (str, optional): Scale to use: ``traunmuller``, ``schroeder`` or ``wang``. (Default: ``traunmuller``) + + Returns: + barks (float): Frequency in Barks + """ + + if bark_scale not in ["schroeder", "traunmuller", "wang"]: + raise ValueError('bark_scale should be one of "schroeder", "traunmuller" or "wang".') + + if bark_scale == "wang": + return 6.0 * math.asinh(freqs / 600.0) + elif bark_scale == "schroeder": + return 7.0 * math.asinh(freqs / 650.0) + # Traunmuller Bark scale + barks = ((26.81 * freqs) / (1960.0 + freqs)) - 0.53 + # Bark value correction + if barks < 2: + barks += 0.15 * (2 - barks) + elif barks > 20.1: + barks += 0.22 * (barks - 20.1) + + return barks + + +def _bark_to_hz(barks: torch.Tensor, bark_scale: str = "traunmuller") -> torch.Tensor: + """Convert bark bin numbers to frequencies. + + Args: + barks (torch.Tensor): Bark frequencies + bark_scale (str, optional): Scale to use: ``traunmuller``,``schroeder`` or ``wang``. (Default: ``traunmuller``) + + Returns: + freqs (torch.Tensor): Barks converted in Hz + """ + + if bark_scale not in ["schroeder", "traunmuller", "wang"]: + raise ValueError('bark_scale should be one of "traunmuller", "schroeder" or "wang".') + + if bark_scale == "wang": + return 600.0 * torch.sinh(barks / 6.0) + elif bark_scale == "schroeder": + return 650.0 * torch.sinh(barks / 7.0) + # Bark value correction + if any(barks < 2): + idx = barks < 2 + barks[idx] = (barks[idx] - 0.3) / 0.85 + elif any(barks > 20.1): + idx = barks > 20.1 + barks[idx] = (barks[idx] + 4.422) / 1.22 + + # Traunmuller Bark scale + freqs = 1960 * ((barks + 0.53) / (26.28 - barks)) + + return freqs + + +def _hz_to_octs(freqs, tuning=0.0, bins_per_octave=12): + a440 = 440.0 * 2.0 ** (tuning / bins_per_octave) + return torch.log2(freqs / (a440 / 16)) + + +def barkscale_fbanks( + n_freqs: int, + f_min: float, + f_max: float, + n_barks: int, + sample_rate: int, + bark_scale: str = "traunmuller", +) -> torch.Tensor: + r"""Create a frequency bin conversion matrix. + + .. devices:: CPU + + .. properties:: TorchScript + + .. image:: https://download.pytorch.org/torchaudio/doc-assets/bark_fbanks.png + :alt: Visualization of generated filter bank + + Args: + n_freqs (int): Number of frequencies to highlight/apply + f_min (float): Minimum frequency (Hz) + f_max (float): Maximum frequency (Hz) + n_barks (int): Number of mel filterbanks + sample_rate (int): Sample rate of the audio waveform + bark_scale (str, optional): Scale to use: ``traunmuller``,``schroeder`` or ``wang``. (Default: ``traunmuller``) + + Returns: + torch.Tensor: Triangular filter banks (fb matrix) of size (``n_freqs``, ``n_barks``) + meaning number of frequencies to highlight/apply to x the number of filterbanks. + Each column is a filterbank so that assuming there is a matrix A of + size (..., ``n_freqs``), the applied result would be + ``A * barkscale_fbanks(A.size(-1), ...)``. + + """ + + # freq bins + all_freqs = torch.linspace(0, sample_rate // 2, n_freqs) + + # calculate bark freq bins + m_min = _hz_to_bark(f_min, bark_scale=bark_scale) + m_max = _hz_to_bark(f_max, bark_scale=bark_scale) + + m_pts = torch.linspace(m_min, m_max, n_barks + 2) + f_pts = _bark_to_hz(m_pts, bark_scale=bark_scale) + + # create filterbank + fb = _create_triangular_filterbank(all_freqs, f_pts) + + if (fb.max(dim=0).values == 0.0).any(): + warnings.warn( + "At least one bark filterbank has all zero values. " + f"The value for `n_barks` ({n_barks}) may be set too high. " + f"Or, the value for `n_freqs` ({n_freqs}) may be set too low." + ) + + return fb + + +def chroma_filterbank( + sample_rate: int, + n_freqs: int, + n_chroma: int, + *, + tuning: float = 0.0, + ctroct: float = 5.0, + octwidth: Optional[float] = 2.0, + norm: int = 2, + base_c: bool = True, +): + """Create a frequency-to-chroma conversion matrix. Implementation adapted from librosa. + + Args: + sample_rate (int): Sample rate. + n_freqs (int): Number of input frequencies. + n_chroma (int): Number of output chroma. + tuning (float, optional): Tuning deviation from A440 in fractions of a chroma bin. (Default: 0.0) + ctroct (float, optional): Center of Gaussian dominance window to weight filters by, in octaves. (Default: 5.0) + octwidth (float or None, optional): Width of Gaussian dominance window to weight filters by, in octaves. + If ``None``, then disable weighting altogether. (Default: 2.0) + norm (int, optional): order of norm to normalize filter bank by. (Default: 2) + base_c (bool, optional): If True, then start filter bank at C. Otherwise, start at A. (Default: True) + + Returns: + torch.Tensor: Chroma filter bank, with shape `(n_freqs, n_chroma)`. + """ + # Skip redundant upper half of frequency range. + freqs = torch.linspace(0, sample_rate // 2, n_freqs)[1:] + freq_bins = n_chroma * _hz_to_octs(freqs, bins_per_octave=n_chroma, tuning=tuning) + freq_bins = torch.cat((torch.tensor([freq_bins[0] - 1.5 * n_chroma]), freq_bins)) + freq_bin_widths = torch.cat( + ( + torch.maximum(freq_bins[1:] - freq_bins[:-1], torch.tensor(1.0)), + torch.tensor([1]), + ) + ) + + # (n_freqs, n_chroma) + D = freq_bins.unsqueeze(1) - torch.arange(0, n_chroma) + + n_chroma2 = round(n_chroma / 2) + + # Project to range [-n_chroma/2, n_chroma/2 - 1] + D = torch.remainder(D + n_chroma2, n_chroma) - n_chroma2 + + fb = torch.exp(-0.5 * (2 * D / torch.tile(freq_bin_widths.unsqueeze(1), (1, n_chroma))) ** 2) + fb = torch.nn.functional.normalize(fb, p=norm, dim=1) + + if octwidth is not None: + fb *= torch.tile( + torch.exp(-0.5 * (((freq_bins.unsqueeze(1) / n_chroma - ctroct) / octwidth) ** 2)), + (1, n_chroma), + ) + + if base_c: + fb = torch.roll(fb, -3 * (n_chroma // 12), dims=1) + + return fb diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cc1b62974644a672bce2916c0f6d04e80e55a2e8 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/__init__.py @@ -0,0 +1,36 @@ +from ._conformer_wav2vec2 import ( + conformer_wav2vec2_base, + conformer_wav2vec2_model, + conformer_wav2vec2_pretrain_base, + conformer_wav2vec2_pretrain_large, + conformer_wav2vec2_pretrain_model, + ConformerWav2Vec2PretrainModel, +) +from ._emformer_hubert import emformer_hubert_base, emformer_hubert_model +from .conv_emformer import ConvEmformer +from .hifi_gan import hifigan_vocoder, hifigan_vocoder_v1, hifigan_vocoder_v2, hifigan_vocoder_v3, HiFiGANVocoder +from .rnnt import conformer_rnnt_base, conformer_rnnt_biasing, conformer_rnnt_biasing_base, conformer_rnnt_model +from .rnnt_decoder import Hypothesis, RNNTBeamSearchBiasing + +__all__ = [ + "conformer_rnnt_base", + "conformer_rnnt_model", + "conformer_rnnt_biasing", + "conformer_rnnt_biasing_base", + "ConvEmformer", + "conformer_wav2vec2_model", + "conformer_wav2vec2_base", + "conformer_wav2vec2_pretrain_model", + "conformer_wav2vec2_pretrain_base", + "conformer_wav2vec2_pretrain_large", + "ConformerWav2Vec2PretrainModel", + "emformer_hubert_base", + "emformer_hubert_model", + "Hypothesis", + "RNNTBeamSearchBiasing", + "HiFiGANVocoder", + "hifigan_vocoder_v1", + "hifigan_vocoder_v2", + "hifigan_vocoder_v3", + "hifigan_vocoder", +] diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..060d762dd2871555e0773fd1e40933eb95ac1b98 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/_conformer_wav2vec2.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/_conformer_wav2vec2.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..12c040becbf0a43ffabbef5d99020091dcf63ac1 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/_conformer_wav2vec2.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/_emformer_hubert.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/_emformer_hubert.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8d0c03bfe70cc11b3a09e35e5689a56be52a9834 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/_emformer_hubert.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/conv_emformer.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/conv_emformer.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..36f5036271d13534a30d2cf3585f6f8e194d91ab Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/conv_emformer.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/hifi_gan.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/hifi_gan.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..56611f9ef189225521f2270a5cc1d2f344a75324 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/hifi_gan.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/rnnt.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/rnnt.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..6ff08b497ea0a93b3ecc3e54d94911562b122674 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/rnnt.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/rnnt_decoder.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/rnnt_decoder.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..37a99be062405277ad644e1291012cfc32d17a9a Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/models/__pycache__/rnnt_decoder.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/_conformer_wav2vec2.py b/lib/python3.12/site-packages/torchaudio/prototype/models/_conformer_wav2vec2.py new file mode 100644 index 0000000000000000000000000000000000000000..b1ea86a81c831a8f346dd1290e221ece67be4734 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/_conformer_wav2vec2.py @@ -0,0 +1,794 @@ +from typing import List, Optional, Tuple, Union + +import torch +from torch import nn, Tensor +from torch.nn import Module, ModuleList +from torchaudio.models import Wav2Vec2Model +from torchaudio.models.conformer import ConformerLayer +from torchaudio.models.rnnt import _TimeReduction +from torchaudio.models.wav2vec2 import components + + +def _buffered_arange(max) -> Tensor: + """Compute arange using a buffered tensor across function calls. + Produces same result as torch.arange(end=max). + + Args: + max (int): Ending value for arange. + """ + if not hasattr(_buffered_arange, "buf"): + _buffered_arange.buf = torch.LongTensor() + if max > _buffered_arange.buf.numel(): + _buffered_arange.buf.resize_(max) + torch.arange(max, out=_buffered_arange.buf) + return _buffered_arange.buf[:max] + + +def _sample_negatives(input: Tensor, num_negatives: int, cross_sample_negatives: int) -> Tuple[Tensor, Tensor]: + """Sample negative examples from masked input. + + Args: + input (Tensor): Tensor of dimension `(batch, frame, dim)`. + num_negatives (int): Number of negative examples to sample. + cross_sample_negatives (int): Number of negative examples to cross sample. + + Returns: + (Tensor, Tensor): + Tensor + The negative samples. + Tensor + The indices of the negative samples. + """ + if num_negatives == 0 and cross_sample_negatives == 0: + return ( + torch.zeros(0).to(input.device, input.dtype), + torch.zeros(0).to(input.device, input.dtype), + ) + + B, T, D = input.shape + input = input.view(-1, D) + + cross_high = T * B + high = T + + assert high > 1 + + if num_negatives > 0: + tszs = _buffered_arange(T).unsqueeze(-1).expand(-1, num_negatives).flatten() + + neg_idxs = torch.randint(low=0, high=high - 1, size=(B, num_negatives * T)) + neg_idxs[neg_idxs >= tszs] += 1 + + if cross_sample_negatives > 0: + tszs = _buffered_arange(T).unsqueeze(-1).expand(-1, cross_sample_negatives).flatten() + + cross_neg_idxs = torch.randint(low=0, high=cross_high - 1, size=(B, cross_sample_negatives * T)) + cross_neg_idxs[cross_neg_idxs >= tszs] += 1 + + if num_negatives > 0: + neg_idxs = neg_idxs + (torch.arange(B).unsqueeze(1) * high) + else: + neg_idxs = cross_neg_idxs + + if cross_sample_negatives > 0 and num_negatives > 0: + neg_idxs = torch.cat([neg_idxs, cross_neg_idxs], dim=1) + + negs = input[neg_idxs.view(-1)] + negs = negs.view(B, T, num_negatives + cross_sample_negatives, D).permute(2, 0, 1, 3) # NxBxCxT + + return negs, neg_idxs + + +class NegativeSampler(Module): + r"""Applies preprocessing to input and then computes negative sampling. + + Args: + preprocessor (nn.Module): Transforms input tensor prior to negative sampling. + num_negatives (int): Number of negative examples to sample. + cross_sample_negatives (int): Number of negative examples to cross sample. + """ + + def __init__( + self, + preprocessor: Module, + num_negatives: int, + cross_sample_negatives: int, + ): + super().__init__() + self.preprocessor = preprocessor + self.num_negatives = num_negatives + self.cross_sample_negatives = cross_sample_negatives + + def forward(self, input: Tensor) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """ + Args: + input (Tensor): Tensor of dimension `(B, T, D)`. + + Returns: + (Tensor, Tensor, Optional[Tensor]): + Tensor + The input tensor after preprocessing, prior to being sampled. + Tensor + The negative samples. + Tensor + The indices of the negative samples. + """ + preprocessed = self.preprocessor(input) + negs, neg_idxs = _sample_negatives(preprocessed, self.num_negatives, self.cross_sample_negatives) + return preprocessed, negs, neg_idxs + + +class FeatureEncoder(Module): + """Feature Encoder class, consisting of time reduction and linear layer. + + Args: + stride (int): Number of frames to merge for the output frame. + input_dim (int): Input dimension of the tensor. + output_dim (int): Output dimension of the tensor. + """ + + def __init__(self, input_dim: int, output_dim: int, stride: int): + super().__init__() + self.time_reduction_layer = _TimeReduction(stride=stride) + self.linear_layer = nn.Linear(input_dim * stride, output_dim) + + def forward( + self, + x: Tensor, + lengths: Optional[Tensor], + ) -> Tuple[Tensor, Optional[Tensor]]: + """ + Args: + x (Tensor): Feature Tensor representing log Mel Spectrogram output. shape ``(B, T, D)``. + lengths (Tensor or None): + Valid length of each input sample. shape: ``(B, )``. + + Returns: + (Tensor, Optional[Tensor]): + Tensor: output sequence after undergoing time reduction and linear projection. + Shape ``(B, T // stride, D * stride). + Optional[Tensor]: output lengths of shape ``(B,)`` if lengths parameter is provided, + otherwise `None`. + """ + if lengths is None: + B, T, D = x.shape + dummy_lengths = torch.full((B,), T) + x, _ = self.time_reduction_layer(x, dummy_lengths) + x = self.linear_layer(x) + return x, None + + x, lengths = self.time_reduction_layer(x, lengths) + x = self.linear_layer(x) + return x, lengths + + +class ConformerEncoder(Module): + """Conformer Encoder class, consisting of feature projection and conformer modules. + + Args: + feature_projection (nn.Module): + Projects feature to encoder dimension. + conformer (nn.ModuleList) + List of Conformer layers. + """ + + def __init__( + self, + feature_projection: Module, + conformer: ModuleList, + ): + super().__init__() + self.feature_projection = feature_projection + self.conformer = conformer + + def _preprocess( + self, + features: Tensor, + lengths: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor]]: + x = self.feature_projection(features) + if lengths is not None: + mask = components._get_padding_mask(x, lengths) + else: + mask = None + return x, mask + + def _get_intermediate_outputs( + self, + x: Tensor, + mask: Optional[Tensor] = None, + num_layers: Optional[int] = None, + ) -> List[Tensor]: + if num_layers is not None: + if not 0 < num_layers <= len(self.conformer): + raise ValueError(f"`num_layers` must be between [1, {len(self.conformer)}]") + + ret: List[Tensor] = [] + + x = x.transpose(0, 1) + for layer in self.conformer: + x = layer(x, mask) + ret.append(x.transpose(0, 1)) + if num_layers is not None and len(ret) >= num_layers: + return ret + return ret + + def forward( + self, + features: Tensor, + lengths: Optional[Tensor] = None, + ) -> Tensor: + """ + Args: + features (Tensor): Tensor of features of shape ``(B, T, D)``. + lengths (Tensor or None, optional): Valid length of each input sample. shape: ``(B, )``. + + Returns: + Tensor: result after applying conformer encoder to features. + """ + x, mask = self._preprocess(features, lengths) + x = x.transpose(0, 1) + for layer in self.conformer: + x = layer(x, mask) + return x.transpose(0, 1) + + def extract_features( + self, + features: Tensor, + lengths: Optional[Tensor] = None, + num_layers: Optional[int] = None, + ) -> List[Tensor]: + """Returns the list of outputs from the intermediate layers of conformer block in the encoder. + + Args: + features (Tensor): Tensor of features of shape ``(B, T, D)``. + lengths (Tensor or None, optional): Valid length of each input sample. shape: ``(B, )``. + + Returns: + List[Tensor]: + Features from requested layers. Each Tensor is of shape: `(batch, time frame, feature dimension)`. + """ + x, masks = self._preprocess(features, lengths) + return self._get_intermediate_outputs(x, mask=masks, num_layers=num_layers) + + +class ConformerWav2Vec2PretrainModel(Module): + """Conformer Wav2Vec2 pre-train model for training from scratch. + + Note: + To build the model, please use one of the factory functions, + :py:func:`conformer_wav2vec2_base` or :py:func:`conformer_wav2vec2_large` + + Args: + wav2vec2 (nn.Module): + Conformer based Wav2Vec2 model, including feature extractor and conformer encoder components. + mask_generator (nn.Module): + Mask generator that generates the mask for masked prediction during training. + negative_sampler (nn.Module): + Negative sampler to apply after masking. + + """ + + def __init__( + self, + wav2vec2: Wav2Vec2Model, + mask_generator: Module, + negative_sampler: Module, + ): + super().__init__() + self.wav2vec2 = wav2vec2 + self.mask_generator = mask_generator + self.negative_sampler = negative_sampler + + def forward( + self, + features: Tensor, + audio_lengths: Optional[Tensor] = None, + ) -> Tuple[Tensor, Optional[Tensor], Tensor, Tensor]: + """ + Args: + features (Tensor): + Tensor of audio features of shape `(batch, frame, dim)`. + audio_lengths (Tensor or None, optional): + Tensor of valid length of each valid auidio in the batch. + shape: `(batch, )` (Default: ``None``) + + Returns: + (Tensor, Optional[Tensor], Tensor, Tensor, Tensor, Tensor): + Tensor + The masked sequences of probability distribution of shape `(batch, frame dim)`. + Tensor or None + If ``lengths`` argument was provided, a Tensor of shape `(batch, )` representing + valid length in time axis is returns. + Tensor + The mask indices. + Tensor + The targets, prior to negative sampling. + Tensor + The negative samples. + Tensor + The indices of the negative samples. + """ + x, lengths = self.wav2vec2.feature_extractor(features, audio_lengths) + + if lengths is not None: + padding_mask = components._get_padding_mask(x, lengths) + else: + padding_mask = None + + x = self.wav2vec2.encoder.feature_projection.layer_norm(x) + x = self.wav2vec2.encoder.feature_projection.dropout(x) + + # Unmasked feature is used to generate positive and negative samples. + unmasked_x = x.clone() + # Apply masking to x before passing it to Conformer layers. + x, mask_idxs = self.mask_generator(x, padding_mask) + # Select the frames from masked indices for negative sampling. + unmasked_x = unmasked_x[mask_idxs].view(x.shape[0], -1, x.shape[-1]) + targets, negs, neg_idxs = self.negative_sampler(unmasked_x) + + x = self.wav2vec2.encoder.feature_projection.projection(x) + x = x.transpose(0, 1) + for conformer_layer in self.wav2vec2.encoder.conformer: + x = conformer_layer(x, padding_mask) + x = x.transpose(0, 1) + + return x, lengths, mask_idxs, targets, negs, neg_idxs + + +################################################################################ +def _get_conformer_feature_extractor( + input_dim: int, + output_dim: int, + stride: int, +) -> FeatureEncoder: + """Construct Feature Extractor + + Args: + input_dim (int): Input dimension of features. + output_dim (int): Output dimension after feature extraction. + stride (int): Stride used in Time Reduction layer of feature extractor. + + Returns: + FeatureEncoder: The resulting feature extraction. + """ + return FeatureEncoder(input_dim, output_dim, stride) + + +def _get_conformer_encoder( + in_features: int, + embed_dim: int, + dropout_input: float, + num_layers: int, + num_heads: int, + ff_interm_features: int, + dropout: float, + depthwise_conv_kernel_size: Union[int, List[int]], + convolution_first: bool, + use_group_norm: bool, +) -> ConformerEncoder: + """Construct Conformer Encoder + + Args: + in_features (int): The number of input features. + embed_dim (int): The dimension of the embedding in the feature projection. + dropout_input (float): The dropout probability applied after the input feature + is projected to ``embed_dim``. + num_layers (int): Number of Conformer layers in the encoder. + num_heads (int): Number of heads in each Conformer layer. + ff_interm_features (int): Hidden layer dimension of the feedforward network in + each Conformer layer. + dropout (float): Dropout probability in each Conformer layer. + depthwise_conv_kernel_size (int or List[int]): List of kernel sizes corresponding + to each of the Conformer layers.If int is provided, all layers will have the + same kernel size. + convolution_first (bool): Whether to apply the convolution module ahead of the + attention module in each Conformer layer. + use_group_norm (bool): Whether to use ``GroupNorm`` rather than ``BatchNorm1d`` in + the convolution module in each Conformer layer. + + Returns: + ConformerEncoder: + The resulting conformer encoder module. + """ + feature_projection = components.FeatureProjection(in_features, embed_dim, dropout_input) + + if type(depthwise_conv_kernel_size) == int: + depthwise_conv_kernel_size = [depthwise_conv_kernel_size] * num_layers + + assert len(depthwise_conv_kernel_size) == num_layers + + conformer_layers = [] + for l in range(num_layers): + layer = ConformerLayer( + input_dim=embed_dim, + ffn_dim=ff_interm_features, + num_attention_heads=num_heads, + depthwise_conv_kernel_size=depthwise_conv_kernel_size[l], + dropout=dropout, + use_group_norm=use_group_norm, + convolution_first=convolution_first, + ) + conformer_layers.append(layer) + + return ConformerEncoder(feature_projection, ModuleList(conformer_layers)) + + +def _get_conformer_negativer_sampler( + input_dim: int, + output_dim: int, + num_negatives: int, + cross_sample_negatives: int, +) -> NegativeSampler: + """Build custom NegativeSampler module, including linear layer and negative sampling. + + Args: + input_dim (int): Dimension of input after feature extraction. + output_dim (int): Dimension of embedding for use in negative sampling. Same as the + embedding in the feature projection. + num_negatives (int): Number of negatives to sample. + cross_sample_negatives (int): Number of cross sampled negatives. + + Returns: + NegativeSampler: + The resulting negative sampler module. + """ + preprocessor = nn.Linear(input_dim, output_dim) + return NegativeSampler(preprocessor, num_negatives, cross_sample_negatives) + + +def conformer_wav2vec2_model( + extractor_input_dim: int, + extractor_output_dim: int, + extractor_stride: int, + encoder_embed_dim: int, + encoder_projection_dropout: float, + encoder_num_layers: int, + encoder_num_heads: int, + encoder_ff_interm_features: int, + encoder_depthwise_conv_kernel_size: Union[int, List[int]], + encoder_dropout: float, + encoder_convolution_first: bool, + encoder_use_group_norm: bool, +) -> Wav2Vec2Model: + """Build a custom Conformer Wav2Vec2Model + + Args: + extractor_input_dim (int): Input dimension of the features. + extractor_output_dim (int): Output dimension after feature extraction. + extractor_stride (int): Stride used in time reduction layer of feature extraction. + encoder_embed_dim (int): The dimension of the embedding in the feature projection. + encoder_projection_dropout (float): + The dropout probability applied after the input feature is projected to ``embed_dim`` + encoder_num_layers (int): Number of Conformer layers in the encoder. + encoder_num_heads (int): Number of heads in each Conformer layer. + encoder_ff_interm_features (int): + Hidden layer dimension of the feedforward network in each Conformer layer. + encoder_depthwise_conv_kernel_size (int or List[int]): + List of kernel sizes corresponding to each of the Conformer layers. + If int is provided, all layers will have the same kernel size. + encoder_dropout (float): Dropout probability in each Conformer layer. + encoder_convolution_first (bool): + Whether to apply the convolution module ahead of the attention module + in each Conformer layer. + encoder_use_group_norm (bool): + Whether to use ``GroupNorm`` rather than ``BatchNorm1d`` in the convolution + module in each Conformer layer. + + Returns: + Wav2Vec2Model: + The resulting wav2vec2 model with a conformer encoder. + """ + feature_extractor = _get_conformer_feature_extractor( + extractor_input_dim, + extractor_output_dim, + extractor_stride, + ) + + encoder = _get_conformer_encoder( + in_features=extractor_output_dim, + embed_dim=encoder_embed_dim, + dropout_input=encoder_projection_dropout, + num_layers=encoder_num_layers, + num_heads=encoder_num_heads, + ff_interm_features=encoder_ff_interm_features, + depthwise_conv_kernel_size=encoder_depthwise_conv_kernel_size, + dropout=encoder_dropout, + convolution_first=encoder_convolution_first, + use_group_norm=encoder_use_group_norm, + ) + + return Wav2Vec2Model(feature_extractor, encoder) + + +def conformer_wav2vec2_base( + extractor_input_dim: int = 64, + extractor_output_dim: int = 256, + encoder_projection_dropout: float = 0.0, +) -> Wav2Vec2Model: + """ + Build Conformer Wav2Vec2 Model with "small" architecture from + *Conformer-Based Slef-Supervised Learning for Non-Speech Audio Tasks* :cite:`9746490` + + Args: + extractor_input_dim (int, optional): Input dimension of feature extractor. (Default: 64) + extractor_output_dim (int, optional): Output dimension of feature extractor. (Default: 256) + encoder_projection_dropout (float, optional): + Dropout probability applied after feature projection. (Default: 0.0) + + Returns: + Wav2Vec2Model: + The resulting wav2vec2 model with a conformer encoder and ``base`` configuration. + """ + return conformer_wav2vec2_model( + extractor_input_dim=extractor_input_dim, + extractor_output_dim=extractor_output_dim, + extractor_stride=4, + encoder_embed_dim=256, + encoder_projection_dropout=encoder_projection_dropout, + encoder_num_layers=12, + encoder_num_heads=8, + encoder_ff_interm_features=1024, + encoder_depthwise_conv_kernel_size=[31] + [15] * 11, + encoder_dropout=0.1, + encoder_convolution_first=True, + encoder_use_group_norm=True, + ) + + +def conformer_wav2vec2_pretrain_model( + extractor_input_dim: int, + extractor_output_dim: int, + extractor_stride: int, + encoder_embed_dim: int, + encoder_projection_dropout: float, + encoder_num_layers: int, + encoder_num_heads: int, + encoder_ff_interm_features: int, + encoder_depthwise_conv_kernel_size: int, + encoder_dropout: float, + encoder_convolution_first: bool, + encoder_use_group_norm: bool, + mask_prob: float, + mask_selection: str, + mask_other: float, + mask_length: int, + no_mask_overlap: bool, + mask_min_space: int, + mask_channel_prob: float, + mask_channel_selection: str, + mask_channel_other: float, + mask_channel_length: int, + no_mask_channel_overlap: bool, + mask_channel_min_space: int, + num_negatives: int, + cross_sample_negatives: int, +) -> ConformerWav2Vec2PretrainModel: + """Build a custom Conformer Wav2Vec2 Model for pre-training + + Args: + extractor_input_dim (int): Input dimension of the features. + extractor_output_dim (int): Output dimension after feature extraction. + extractor_stride (int): + Stride used in time reduction layer of feature extraction. + encoder_embed_dim (int): + The dimension of the embedding in the feature projection. + encoder_projection_dropout (float): + The dropout probability applied after the input feature is projected to + ``embed_dim`` + encoder_num_layers (int): + Number of Conformer layers in the encoder. + encoder_num_heads (int): + Number of heads in each Conformer layer. + encoder_ff_interm_features (int): + Hidden layer dimension of the feedforward network in each Conformer layer. + encoder_depthwise_conv_kernel_size (int or List[int]): + List of kernel sizes corresponding to each of the Conformer layers. + If int is provided, all layers will have the same kernel size. + encoder_dropout (float): + Dropout probability in each Conformer layer. + encoder_convolution_first (bool): + Whether to apply the convolution module ahead of the attention module + in each Conformer layer. + encoder_use_group_norm (bool): + Whether to use ``GroupNorm`` rather than ``BatchNorm1d`` in the convolution + module in each Conformer layer. + mask_prob (float): + Probability for each token to be chosen as start of the span to be masked. + mask_selection (str) + How to choose the mask length. Options: [``static``, ``uniform``, ``normal``, ``poisson``]. + mask_other (float): + Secondary mask argument (used for more complex distributions). + mask_length (int): + The lengths of the mask. + no_mask_overlap (bool): + Whether to allow masks to overlap. + mask_min_space (int): + Minimum space between spans (if no overlap is enabled). + mask_channel_prob: (float): + The probability of replacing a feature with 0. + mask_channel_selection (str): + How to choose the mask length for channel masking. + Options: [``static``, ``uniform``, ``normal``, ``poisson``]. + mask_channel_other (float): + Secondary mask argument for channel masking (used for more complex distributions). + mask_channel_length (int): + Minimum space between spans (if no overlap is enabled) for channel masking. + no_mask_channel_overlap (bool): + Whether to allow channel masks to overlap. + mask_channel_min_space (int): + Minimum space between spans for channel masking (if no overlap is enabled). + num_negatives (int): + Number of negatives to sample. + cross_sample_negatives (int): + Number of cross sampled negatives. + + Returns: + ConformerWav2Vec2PretrainModel: + The resulting model. + """ + wav2vec2 = conformer_wav2vec2_model( + extractor_input_dim, + extractor_output_dim, + extractor_stride, + encoder_embed_dim, + encoder_projection_dropout, + encoder_num_layers, + encoder_num_heads, + encoder_ff_interm_features, + encoder_depthwise_conv_kernel_size, + encoder_dropout, + encoder_convolution_first, + encoder_use_group_norm, + ) + + mask_generator = components.MaskGenerator( + extractor_output_dim, + mask_prob, + mask_selection, + mask_other, + mask_length, + no_mask_overlap, + mask_min_space, + mask_channel_prob, + mask_channel_selection, + mask_channel_other, + mask_channel_length, + no_mask_channel_overlap, + mask_channel_min_space, + ) + + negative_sampler = _get_conformer_negativer_sampler( + extractor_output_dim, + encoder_embed_dim, + num_negatives, + cross_sample_negatives, + ) + + return ConformerWav2Vec2PretrainModel( + wav2vec2=wav2vec2, + mask_generator=mask_generator, + negative_sampler=negative_sampler, + ) + + +def conformer_wav2vec2_pretrain_base( + extractor_input_dim: int = 64, + extractor_output_dim: int = 256, + encoder_projection_dropout: float = 0.0, + mask_prob: float = 0.3, + mask_length: int = 3, + num_negatives: int = 100, + cross_sample_negatives: int = 0, +) -> ConformerWav2Vec2PretrainModel: + """Build Conformer Wav2Vec2 Model for pre-training with "small" architecture from + *Conformer-Based Self-Supervised Learning for Non-Speech Audio Tasks* :cite:`9746490` + + Args: + extractor_input_dim (int, optional): Input dimension of the features. (Default: 64) + extractor_output_dim (int, optional): Output dimension after feature extraction. (Default: 256) + encoder_projection_dropout (float, optional): + The dropout probability applied after the input feature is projected to + ``embed_dim``. (Default: 0.0) + mask_prob (float, optional): + Probability for each token to be chosen as start of the span to be masked. (Default: 0.3) + mask_length (int, optional): + The lengths of the mask. (Default: 3) + num_negatives (int, optional): + Number of sampled negatives. (Default: 0) + cross_sample_negatives (int, optional): + Number of cross sampled negatives. (Default: 0) + + Returns: + ConformerWav2Vec2PretrainModel: + The resulting model. + """ + return conformer_wav2vec2_pretrain_model( + extractor_input_dim=extractor_input_dim, + extractor_output_dim=extractor_output_dim, + extractor_stride=4, + encoder_embed_dim=256, + encoder_projection_dropout=encoder_projection_dropout, + encoder_num_layers=12, + encoder_num_heads=8, + encoder_ff_interm_features=1024, + encoder_depthwise_conv_kernel_size=[31] + [15] * 11, + encoder_dropout=0.1, + encoder_convolution_first=True, + encoder_use_group_norm=True, + mask_prob=mask_prob, + mask_selection="static", + mask_other=0.0, + mask_length=mask_length, + no_mask_overlap=False, + mask_min_space=0, + mask_channel_prob=0, + mask_channel_selection="static", + mask_channel_other=0, + mask_channel_length=10, + no_mask_channel_overlap=False, + mask_channel_min_space=1, + num_negatives=num_negatives, + cross_sample_negatives=cross_sample_negatives, + ) + + +def conformer_wav2vec2_pretrain_large( + extractor_input_dim: int = 64, + extractor_output_dim: int = 256, + encoder_projection_dropout: float = 0.0, + mask_prob: float = 0.3, + mask_length: int = 3, + num_negatives: int = 100, + cross_sample_negatives: int = 0, +) -> ConformerWav2Vec2PretrainModel: + """Build Conformer Wav2Vec2 Model for pre-training with "large" architecture from + *Conformer-Based Slef-Supervised Learning for Non-Speech Audio Tasks* :cite:`9746490` + + Args: + extractor_input_dim (int, optional): Input dimension of the features. (Default: 64) + extractor_output_dim (int, optional): Output dimension after feature extraction. (Default: 256) + encoder_projection_dropout (float, optional): + The dropout probability applied after the input feature is projected to + ``embed_dim``. (Default: 0.0) + mask_prob (float, optional): + Probability for each token to be chosen as start of the span to be masked. (Default: 0.3) + mask_length (int, optional): + The lengths of the mask. (Default: 3) + num_negatives (int, optional): + Number of sampled negatives. (Default: 0) + cross_sample_negatives (int, optional): + Number of cross sampled negatives. (Default: 0) + + Returns: + ConformerWav2Vec2PretrainModel: + The resulting model. + """ + return conformer_wav2vec2_pretrain_model( + extractor_input_dim=extractor_input_dim, + extractor_output_dim=extractor_output_dim, + extractor_stride=4, + encoder_embed_dim=768, + encoder_projection_dropout=encoder_projection_dropout, + encoder_num_layers=12, + encoder_num_heads=12, + encoder_ff_interm_features=1024, + encoder_depthwise_conv_kernel_size=[31] + [15] * 11, + encoder_dropout=0.1, + encoder_convolution_first=True, + encoder_use_group_norm=True, + mask_prob=mask_prob, + mask_selection="static", + mask_other=0.0, + mask_length=mask_length, + no_mask_overlap=False, + mask_min_space=0, + mask_channel_prob=0, + mask_channel_selection="static", + mask_channel_other=0, + mask_channel_length=10, + no_mask_channel_overlap=False, + mask_channel_min_space=1, + num_negatives=num_negatives, + cross_sample_negatives=cross_sample_negatives, + ) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/_emformer_hubert.py b/lib/python3.12/site-packages/torchaudio/prototype/models/_emformer_hubert.py new file mode 100644 index 0000000000000000000000000000000000000000..872c6ce90191a841d7a1387bf17a1803b1689b83 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/_emformer_hubert.py @@ -0,0 +1,333 @@ +from typing import List, Optional, Tuple + +import torch +from torchaudio.models import Wav2Vec2Model +from torchaudio.models.emformer import Emformer +from torchaudio.models.rnnt import _TimeReduction + + +class FeatureEncoder(torch.nn.Module): + """Extract features from log-mel spectrogram input. Consists of linear layer and time reduction layer. + + Args: + input_dim (int): The feature dimension of log-mel spectrogram feature. + output_dim (int): The feature dimension after linear layer. + use_bias (bool): If ``True``, enable bias parameter in the linear layer. + stride (int): Number of frames to merge for the output frame. + """ + + def __init__(self, input_dim: int, output_dim: int, use_bias: bool, stride: int): + super().__init__() + self.linear = torch.nn.Linear(input_dim, output_dim, bias=use_bias) + self.time_reduction = _TimeReduction(stride) + + def forward( + self, input: torch.Tensor, lengths: Optional[torch.Tensor] + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """ + Args: + input (torch.Tensor): The log-mel spectrogram input. + Tensor with dimensions `(batch, time, input_dim)`. + lengths (torch.Tensor or None): Valid length of each input sample. + Tensor with dimension `(batch, )`. + + Returns: + (torch.Tensor, torch.Tensor or None): + torch.Tensor + Returned feature Tensor after linear layer and time reduction layer. + Tensor with dimensions `(batch, time // stride, output_dim)`. + torch.Tensor or None + The reduced lengths Tensor. + """ + output = self.linear(input) + if lengths is None: + B, T, _ = input.shape + dummy_lengths = torch.full((B,), T) + output, _ = self.time_reduction(output, dummy_lengths) + else: + output, lengths = self.time_reduction(output, lengths) + return output, lengths + + +class EmformerEncoder(torch.nn.Module): + """Emformer Encoder class for HuBERT pre-training. Consists of emformer module, + linear layer and layer normalization layer. + + Args: + emformer (torch.nn.Module): + :py:class:`torchaudio.models.Emformer` module that consists of a list of emformer layers. + output_linear (torch.nn.Module): + Linear layer after emformer module. + layer_norm (torch.nn.Module): + Apply layer normalization to the output. + """ + + def __init__( + self, + emformer: torch.nn.Module, + output_linear: torch.nn.Module, + layer_norm: torch.nn.Module, + ): + super().__init__() + self.emformer = emformer + self.output_linear = output_linear + self.layer_norm = layer_norm + + def forward( + self, + input: torch.Tensor, + lengths: Optional[torch.Tensor], + ) -> torch.Tensor: + """ + Args: + input (torch.Tensor): The input feature for emformer encoder. + Tensor with dimensions `(batch, time, feature_dim)`. + lengths (torch.Tensor or None): Valid length of each input sample. + Tensor with dimension `(batch, )`. + + Returns: + torch.Tensor: The feature Tensor after emformer encoder. + """ + if lengths is None: + B, T, _ = input.shape + dummy_lengths = torch.full((B,), T) + output, _ = self.emformer(input, dummy_lengths) + else: + output, lengths = self.emformer(input, lengths) + output = self.output_linear(output) + output = self.layer_norm(output) + return output + + def extract_features( + self, + input: torch.Tensor, + lengths: Optional[torch.Tensor], + num_layers: Optional[int] = None, + ) -> List[torch.Tensor]: + """Extract output Tensors of the emformer layers. + + Args: + input (torch.Tensor): The input feature for emformer encoder. + Tensor with dimensions `(batch, time, feature_dim)`. + lengths (torch.Tensor or None): Valid length of each input sample. + Tensor with dimension `(batch, )`. + num_layers (int or None, optional): If not ``None``, returns the first + `num_layers` layers of Tensors as the output, otherwise returns the + Tensors from all emformer layers. + + Returns: + List[torch.Tensor]: + Output Tensors of selected emformer layers. + """ + if num_layers is not None: + if not 0 < num_layers <= len(self.emformer.emformer_layers): + raise ValueError(f"`num_layers` must be between [1, {len(self.emformer.emformer_layers)}]") + + ret: List[torch.Tensor] = [] + + input = input.permute(1, 0, 2) + right_context = self.emformer._gen_right_context(input) + utterance = input[: input.size(0) - self.emformer.right_context_length] + attention_mask = self.emformer._gen_attention_mask(utterance) + mems = ( + self.emformer.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)[:-1] + if self.emformer.use_mem + else torch.empty(0).to(dtype=input.dtype, device=input.device) + ) + output = utterance + if lengths is None: + B, T, _ = input.shape + lengths = torch.full((B,), T) + for layer in self.emformer.emformer_layers: + output, right_context, mems = layer(output, lengths, right_context, mems, attention_mask) + ret.append(output.permute(1, 0, 2)) + if num_layers is not None and len(ret) >= num_layers: + return ret + return ret + + +def _get_emformer_feature_extractor(input_dim: int, output_dim: int, use_bias: bool, stride: int) -> FeatureEncoder: + """Construct FeatureEncoder for emformer model. + + Args: + input_dim (int): The feature dimension of log-mel spectrogram feature. + output_dim (int): The feature dimension after linear layer. + use_bias (bool): If ``True``, enable bias parameter in the linear layer. + stride (int): Number of frames to merge for the output frame. + + Returns: + FeatureEncoder: The resulting FeatureEncoder module. + """ + return FeatureEncoder(input_dim, output_dim, use_bias, stride) + + +def _get_emformer_encoder( + input_dim: int, + output_dim: int, + num_heads: int, + ffn_dim: int, + num_layers: int, + segment_length: int, + left_context_length: int, + right_context_length: int, + dropout: float, + activation: str, + max_memory_size: int, + weight_init_scale_strategy: Optional[str], + tanh_on_mem: bool, +) -> EmformerEncoder: + """Construct EmformerEncoder for emformer model. + + Args: + input_dim (int): The feature dimension of input Tensor. + output_dim (int): The feature dimension after EmformerEncoder. + num_heads (int): Number of attention heads in each Emformer layer. + ffn_dim: (int): Hidden layer dimension of feedforward network. + num_layers (int): Number of Emformer layers to instantiate. + segment_length (int): Length of each input segment. + left_context_length (int): Length of left context. + right_context_length (int): Length of right context. + dropout (float): Dropout probability. + activation (str): Activation function to use in each Emformer layer's + feedforward network. Must be one of ("relu", "gelu", "silu"). + max_memory_size (int): Maximum number of memory elements to use. + weight_init_scale_strategy (str or None): Per-layer weight initialization scaling + strategy. Must be one of ("depthwise", "constant", ``None``). + tanh_on_mem (bool): If ``True``, applies tanh to memory elements. + + Returns: + EmformerEncoder: The resulting EmformerEncoder module. + """ + emformer = Emformer( + input_dim=input_dim, + num_heads=num_heads, + ffn_dim=ffn_dim, + num_layers=num_layers, + segment_length=segment_length, + left_context_length=left_context_length, + right_context_length=right_context_length, + dropout=dropout, + activation=activation, + max_memory_size=max_memory_size, + weight_init_scale_strategy=weight_init_scale_strategy, + tanh_on_mem=tanh_on_mem, + ) + output_linear = torch.nn.Linear(input_dim, output_dim) + layer_norm = torch.nn.LayerNorm(output_dim) + return EmformerEncoder(emformer, output_linear, layer_norm) + + +def emformer_hubert_model( + extractor_input_dim: int, + extractor_output_dim: int, + extractor_use_bias: bool, + extractor_stride: int, + encoder_input_dim: int, + encoder_output_dim: int, + encoder_num_heads: int, + encoder_ffn_dim: int, + encoder_num_layers: int, + encoder_segment_length: int, + encoder_left_context_length: int, + encoder_right_context_length: int, + encoder_dropout: float, + encoder_activation: str, + encoder_max_memory_size: int, + encoder_weight_init_scale_strategy: Optional[str], + encoder_tanh_on_mem: bool, + aux_num_out: Optional[int], +) -> Wav2Vec2Model: + """Build a custom Emformer HuBERT model. + + Args: + extractor_input_dim (int): The input dimension for feature extractor. + extractor_output_dim (int): The output dimension after feature extractor. + extractor_use_bias (bool): If ``True``, enable bias parameter in the linear layer of feature extractor. + extractor_stride (int): Number of frames to merge for the output frame in feature extractor. + encoder_input_dim (int): The input dimension for Emformer layer. + encoder_output_dim (int): The output dimension after EmformerEncoder. + encoder_num_heads (int): Number of attention heads in each Emformer layer. + encoder_ffn_dim (int): Hidden layer dimension of feedforward network in Emformer. + encoder_num_layers (int): Number of Emformer layers to instantiate. + encoder_segment_length (int): Length of each input segment. + encoder_left_context_length (int): Length of left context. + encoder_right_context_length (int): Length of right context. + encoder_dropout (float): Dropout probability. + encoder_activation (str): Activation function to use in each Emformer layer's + feedforward network. Must be one of ("relu", "gelu", "silu"). + encoder_max_memory_size (int): Maximum number of memory elements to use. + encoder_weight_init_scale_strategy (str or None): Per-layer weight initialization scaling + strategy. Must be one of ("depthwise", "constant", ``None``). + encoder_tanh_on_mem (bool): If ``True``, applies tanh to memory elements. + aux_num_out (int or None): + When provided, attach an extra linear layer on top of encoder, which can be + used for fine-tuning. + + Returns: + Wav2Vec2Model: + The resulting :py:class:`torchaudio.models.Wav2Vec2Model` model + with a :py:class:`torchaudio.models.Emformer` encoder. + """ + feature_extractor = _get_emformer_feature_extractor( + extractor_input_dim, extractor_output_dim, extractor_use_bias, extractor_stride + ) + emformer = _get_emformer_encoder( + encoder_input_dim, + encoder_output_dim, + encoder_num_heads, + encoder_ffn_dim, + encoder_num_layers, + encoder_segment_length, + encoder_left_context_length, + encoder_right_context_length, + encoder_dropout, + encoder_activation, + encoder_max_memory_size, + encoder_weight_init_scale_strategy, + encoder_tanh_on_mem, + ) + aux = None + if aux_num_out is not None: + aux = torch.nn.Linear(in_features=encoder_output_dim, out_features=aux_num_out) + return Wav2Vec2Model(feature_extractor, emformer, aux) + + +def emformer_hubert_base( + extractor_input_dim: int = 80, + extractor_output_dim: int = 128, + encoder_dropout: float = 0.1, + aux_num_out: Optional[int] = None, +) -> Wav2Vec2Model: + """Build Emformer HuBERT Model with 20 Emformer layers. + + Args: + extractor_input_dim (int, optional): The input dimension for feature extractor. (Default: 80) + extractor_output_dim (int, optional): The output dimension after feature extractor. (Default: 128) + encoder_dropout (float, optional): Dropout probability in Emformer. (Default: 0.1) + aux_num_out (int or None, optional): Output dimension of aux layer for fine-tuning. (Default: ``None``) + + Returns: + Wav2Vec2Model: + The resulting :py:class:`torchaudio.models.Wav2Vec2Model` model + with a :py:class:`torchaudio.models.Emformer` encoder. + """ + return emformer_hubert_model( + extractor_input_dim=extractor_input_dim, + extractor_output_dim=extractor_output_dim, + extractor_use_bias=False, + extractor_stride=4, + encoder_input_dim=512, + encoder_output_dim=1024, + encoder_num_heads=8, + encoder_ffn_dim=2048, + encoder_num_layers=20, + encoder_segment_length=4, + encoder_left_context_length=30, + encoder_right_context_length=1, + encoder_dropout=encoder_dropout, + encoder_activation="gelu", + encoder_max_memory_size=0, + encoder_weight_init_scale_strategy="depthwise", + encoder_tanh_on_mem=True, + aux_num_out=aux_num_out, + ) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/conv_emformer.py b/lib/python3.12/site-packages/torchaudio/prototype/models/conv_emformer.py new file mode 100644 index 0000000000000000000000000000000000000000..b5495cfcd3a4753cdf17d426a1b2bcc8479ee3fa --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/conv_emformer.py @@ -0,0 +1,525 @@ +import math +from typing import List, Optional, Tuple + +import torch +from torchaudio.models.emformer import _EmformerAttention, _EmformerImpl, _get_weight_init_gains + + +def _get_activation_module(activation: str) -> torch.nn.Module: + if activation == "relu": + return torch.nn.ReLU() + elif activation == "gelu": + return torch.nn.GELU() + elif activation == "silu": + return torch.nn.SiLU() + else: + raise ValueError(f"Unsupported activation {activation}") + + +class _ResidualContainer(torch.nn.Module): + def __init__(self, module: torch.nn.Module, output_weight: int): + super().__init__() + self.module = module + self.output_weight = output_weight + + def forward(self, input: torch.Tensor): + output = self.module(input) + return output * self.output_weight + input + + +class _ConvolutionModule(torch.nn.Module): + def __init__( + self, + input_dim: int, + segment_length: int, + right_context_length: int, + kernel_size: int, + activation: str = "silu", + dropout: float = 0.0, + ): + super().__init__() + self.input_dim = input_dim + self.segment_length = segment_length + self.right_context_length = right_context_length + self.state_size = kernel_size - 1 + + self.pre_conv = torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), torch.nn.Linear(input_dim, 2 * input_dim, bias=True), torch.nn.GLU() + ) + self.conv = torch.nn.Conv1d( + in_channels=input_dim, + out_channels=input_dim, + kernel_size=kernel_size, + stride=1, + padding=0, + groups=input_dim, + ) + self.post_conv = torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), + _get_activation_module(activation), + torch.nn.Linear(input_dim, input_dim, bias=True), + torch.nn.Dropout(p=dropout), + ) + + def _split_right_context(self, utterance: torch.Tensor, right_context: torch.Tensor) -> torch.Tensor: + T, B, D = right_context.size() + if T % self.right_context_length != 0: + raise ValueError("Tensor length should be divisible by its right context length") + num_segments = T // self.right_context_length + # (num_segments, right context length, B, D) + right_context_segments = right_context.reshape(num_segments, self.right_context_length, B, D) + right_context_segments = right_context_segments.permute(0, 2, 1, 3).reshape( + num_segments * B, self.right_context_length, D + ) + + pad_segments = [] # [(kernel_size - 1, B, D), ...] + for seg_idx in range(num_segments): + end_idx = min(self.state_size + (seg_idx + 1) * self.segment_length, utterance.size(0)) + start_idx = end_idx - self.state_size + pad_segments.append(utterance[start_idx:end_idx, :, :]) + + pad_segments = torch.cat(pad_segments, dim=1).permute(1, 0, 2) # (num_segments * B, kernel_size - 1, D) + return torch.cat([pad_segments, right_context_segments], dim=1).permute(0, 2, 1) + + def _merge_right_context(self, right_context: torch.Tensor, B: int) -> torch.Tensor: + # (num_segments * B, D, right_context_length) + right_context = right_context.reshape(-1, B, self.input_dim, self.right_context_length) + right_context = right_context.permute(0, 3, 1, 2) + return right_context.reshape(-1, B, self.input_dim) # (right_context_length * num_segments, B, D) + + def forward( + self, utterance: torch.Tensor, right_context: torch.Tensor, state: Optional[torch.Tensor] + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input = torch.cat((right_context, utterance)) # input: (T, B, D) + x = self.pre_conv(input) + x_right_context, x_utterance = x[: right_context.size(0), :, :], x[right_context.size(0) :, :, :] + x_utterance = x_utterance.permute(1, 2, 0) # (B, D, T_utterance) + + if state is None: + state = torch.zeros( + input.size(1), + input.size(2), + self.state_size, + device=input.device, + dtype=input.dtype, + ) # (B, D, T) + state_x_utterance = torch.cat([state, x_utterance], dim=2) + + conv_utterance = self.conv(state_x_utterance) # (B, D, T_utterance) + conv_utterance = conv_utterance.permute(2, 0, 1) + + if self.right_context_length > 0: + # (B * num_segments, D, right_context_length + kernel_size - 1) + right_context_block = self._split_right_context(state_x_utterance.permute(2, 0, 1), x_right_context) + conv_right_context_block = self.conv(right_context_block) # (B * num_segments, D, right_context_length) + # (T_right_context, B, D) + conv_right_context = self._merge_right_context(conv_right_context_block, input.size(1)) + y = torch.cat([conv_right_context, conv_utterance], dim=0) + else: + y = conv_utterance + + output = self.post_conv(y) + input + new_state = state_x_utterance[:, :, -self.state_size :] + return output[right_context.size(0) :], output[: right_context.size(0)], new_state + + def infer( + self, utterance: torch.Tensor, right_context: torch.Tensor, state: Optional[torch.Tensor] + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + input = torch.cat((utterance, right_context)) + x = self.pre_conv(input) # (T, B, D) + x = x.permute(1, 2, 0) # (B, D, T) + + if state is None: + state = torch.zeros( + input.size(1), + input.size(2), + self.state_size, + device=input.device, + dtype=input.dtype, + ) # (B, D, T) + state_x = torch.cat([state, x], dim=2) + conv_out = self.conv(state_x) + conv_out = conv_out.permute(2, 0, 1) # T, B, D + output = self.post_conv(conv_out) + input + new_state = state_x[:, :, -self.state_size - right_context.size(0) : -right_context.size(0)] + return output[: utterance.size(0)], output[utterance.size(0) :], new_state + + +class _ConvEmformerLayer(torch.nn.Module): + r"""Convolution-augmented Emformer layer that constitutes ConvEmformer. + + Args: + input_dim (int): input dimension. + num_heads (int): number of attention heads. + ffn_dim: (int): hidden layer dimension of feedforward network. + segment_length (int): length of each input segment. + kernel_size (int): size of kernel to use in convolution module. + dropout (float, optional): dropout probability. (Default: 0.0) + ffn_activation (str, optional): activation function to use in feedforward network. + Must be one of ("relu", "gelu", "silu"). (Default: "relu") + left_context_length (int, optional): length of left context. (Default: 0) + right_context_length (int, optional): length of right context. (Default: 0) + max_memory_size (int, optional): maximum number of memory elements to use. (Default: 0) + weight_init_gain (float or None, optional): scale factor to apply when initializing + attention module parameters. (Default: ``None``) + tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``) + negative_inf (float, optional): value to use for negative infinity in attention weights. (Default: -1e8) + conv_activation (str, optional): activation function to use in convolution module. + Must be one of ("relu", "gelu", "silu"). (Default: "silu") + """ + + def __init__( + self, + input_dim: int, + num_heads: int, + ffn_dim: int, + segment_length: int, + kernel_size: int, + dropout: float = 0.0, + ffn_activation: str = "relu", + left_context_length: int = 0, + right_context_length: int = 0, + max_memory_size: int = 0, + weight_init_gain: Optional[float] = None, + tanh_on_mem: bool = False, + negative_inf: float = -1e8, + conv_activation: str = "silu", + ): + super().__init__() + # TODO: implement talking heads attention. + self.attention = _EmformerAttention( + input_dim=input_dim, + num_heads=num_heads, + dropout=dropout, + weight_init_gain=weight_init_gain, + tanh_on_mem=tanh_on_mem, + negative_inf=negative_inf, + ) + self.dropout = torch.nn.Dropout(dropout) + self.memory_op = torch.nn.AvgPool1d(kernel_size=segment_length, stride=segment_length, ceil_mode=True) + + activation_module = _get_activation_module(ffn_activation) + self.ffn0 = _ResidualContainer( + torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), + torch.nn.Linear(input_dim, ffn_dim), + activation_module, + torch.nn.Dropout(dropout), + torch.nn.Linear(ffn_dim, input_dim), + torch.nn.Dropout(dropout), + ), + 0.5, + ) + self.ffn1 = _ResidualContainer( + torch.nn.Sequential( + torch.nn.LayerNorm(input_dim), + torch.nn.Linear(input_dim, ffn_dim), + activation_module, + torch.nn.Dropout(dropout), + torch.nn.Linear(ffn_dim, input_dim), + torch.nn.Dropout(dropout), + ), + 0.5, + ) + self.layer_norm_input = torch.nn.LayerNorm(input_dim) + self.layer_norm_output = torch.nn.LayerNorm(input_dim) + + self.conv = _ConvolutionModule( + input_dim=input_dim, + kernel_size=kernel_size, + activation=conv_activation, + dropout=dropout, + segment_length=segment_length, + right_context_length=right_context_length, + ) + + self.left_context_length = left_context_length + self.segment_length = segment_length + self.max_memory_size = max_memory_size + self.input_dim = input_dim + self.kernel_size = kernel_size + self.use_mem = max_memory_size > 0 + + def _init_state(self, batch_size: int, device: Optional[torch.device]) -> List[torch.Tensor]: + empty_memory = torch.zeros(self.max_memory_size, batch_size, self.input_dim, device=device) + left_context_key = torch.zeros(self.left_context_length, batch_size, self.input_dim, device=device) + left_context_val = torch.zeros(self.left_context_length, batch_size, self.input_dim, device=device) + past_length = torch.zeros(1, batch_size, dtype=torch.int32, device=device) + conv_cache = torch.zeros( + batch_size, + self.input_dim, + self.kernel_size - 1, + device=device, + ) + return [empty_memory, left_context_key, left_context_val, past_length, conv_cache] + + def _unpack_state(self, state: List[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + past_length = state[3][0][0].item() + past_left_context_length = min(self.left_context_length, past_length) + past_mem_length = min(self.max_memory_size, math.ceil(past_length / self.segment_length)) + pre_mems = state[0][self.max_memory_size - past_mem_length :] + lc_key = state[1][self.left_context_length - past_left_context_length :] + lc_val = state[2][self.left_context_length - past_left_context_length :] + conv_cache = state[4] + return pre_mems, lc_key, lc_val, conv_cache + + def _pack_state( + self, + next_k: torch.Tensor, + next_v: torch.Tensor, + update_length: int, + mems: torch.Tensor, + conv_cache: torch.Tensor, + state: List[torch.Tensor], + ) -> List[torch.Tensor]: + new_k = torch.cat([state[1], next_k]) + new_v = torch.cat([state[2], next_v]) + state[0] = torch.cat([state[0], mems])[-self.max_memory_size :] + state[1] = new_k[new_k.shape[0] - self.left_context_length :] + state[2] = new_v[new_v.shape[0] - self.left_context_length :] + state[3] = state[3] + update_length + state[4] = conv_cache + return state + + def _apply_pre_attention( + self, utterance: torch.Tensor, right_context: torch.Tensor, summary: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + x = torch.cat([right_context, utterance, summary]) + ffn0_out = self.ffn0(x) + layer_norm_input_out = self.layer_norm_input(ffn0_out) + layer_norm_input_right_context, layer_norm_input_utterance, layer_norm_input_summary = ( + layer_norm_input_out[: right_context.size(0)], + layer_norm_input_out[right_context.size(0) : right_context.size(0) + utterance.size(0)], + layer_norm_input_out[right_context.size(0) + utterance.size(0) :], + ) + return ffn0_out, layer_norm_input_right_context, layer_norm_input_utterance, layer_norm_input_summary + + def _apply_post_attention( + self, + rc_output: torch.Tensor, + ffn0_out: torch.Tensor, + conv_cache: Optional[torch.Tensor], + rc_length: int, + utterance_length: int, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + result = self.dropout(rc_output) + ffn0_out[: rc_length + utterance_length] + conv_utterance, conv_right_context, conv_cache = self.conv(result[rc_length:], result[:rc_length], conv_cache) + result = torch.cat([conv_right_context, conv_utterance]) + result = self.ffn1(result) + result = self.layer_norm_output(result) + output_utterance, output_right_context = result[rc_length:], result[:rc_length] + return output_utterance, output_right_context, conv_cache + + def forward( + self, + utterance: torch.Tensor, + lengths: torch.Tensor, + right_context: torch.Tensor, + mems: torch.Tensor, + attention_mask: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + r"""Forward pass for training. + + B: batch size; + D: feature dimension of each frame; + T: number of utterance frames; + R: number of right context frames; + M: number of memory elements. + + Args: + utterance (torch.Tensor): utterance frames, with shape `(T, B, D)`. + lengths (torch.Tensor): with shape `(B,)` and i-th element representing + number of valid frames for i-th batch element in ``utterance``. + right_context (torch.Tensor): right context frames, with shape `(R, B, D)`. + mems (torch.Tensor): memory elements, with shape `(M, B, D)`. + attention_mask (torch.Tensor): attention mask for underlying attention module. + + Returns: + (Tensor, Tensor, Tensor): + Tensor + encoded utterance frames, with shape `(T, B, D)`. + Tensor + updated right context frames, with shape `(R, B, D)`. + Tensor + updated memory elements, with shape `(M, B, D)`. + """ + if self.use_mem: + summary = self.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1) + else: + summary = torch.empty(0).to(dtype=utterance.dtype, device=utterance.device) + + ( + ffn0_out, + layer_norm_input_right_context, + layer_norm_input_utterance, + layer_norm_input_summary, + ) = self._apply_pre_attention(utterance, right_context, summary) + + rc_output, output_mems = self.attention( + utterance=layer_norm_input_utterance, + lengths=lengths, + right_context=layer_norm_input_right_context, + summary=layer_norm_input_summary, + mems=mems, + attention_mask=attention_mask, + ) + + output_utterance, output_right_context, _ = self._apply_post_attention( + rc_output, ffn0_out, None, right_context.size(0), utterance.size(0) + ) + + return output_utterance, output_right_context, output_mems + + @torch.jit.export + def infer( + self, + utterance: torch.Tensor, + lengths: torch.Tensor, + right_context: torch.Tensor, + state: Optional[List[torch.Tensor]], + mems: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor, List[torch.Tensor], torch.Tensor]: + r"""Forward pass for inference. + + B: batch size; + D: feature dimension of each frame; + T: number of utterance frames; + R: number of right context frames; + M: number of memory elements. + + Args: + utterance (torch.Tensor): utterance frames, with shape `(T, B, D)`. + lengths (torch.Tensor): with shape `(B,)` and i-th element representing + number of valid frames for i-th batch element in ``utterance``. + right_context (torch.Tensor): right context frames, with shape `(R, B, D)`. + state (List[torch.Tensor] or None): list of tensors representing layer internal state + generated in preceding invocation of ``infer``. + mems (torch.Tensor): memory elements, with shape `(M, B, D)`. + + Returns: + (Tensor, Tensor, List[torch.Tensor], Tensor): + Tensor + encoded utterance frames, with shape `(T, B, D)`. + Tensor + updated right context frames, with shape `(R, B, D)`. + List[Tensor] + list of tensors representing layer internal state + generated in current invocation of ``infer``. + Tensor + updated memory elements, with shape `(M, B, D)`. + """ + if self.use_mem: + summary = self.memory_op(utterance.permute(1, 2, 0)).permute(2, 0, 1)[:1] + else: + summary = torch.empty(0).to(dtype=utterance.dtype, device=utterance.device) + + ( + ffn0_out, + layer_norm_input_right_context, + layer_norm_input_utterance, + layer_norm_input_summary, + ) = self._apply_pre_attention(utterance, right_context, summary) + + if state is None: + state = self._init_state(layer_norm_input_utterance.size(1), device=layer_norm_input_utterance.device) + pre_mems, lc_key, lc_val, conv_cache = self._unpack_state(state) + + rc_output, next_m, next_k, next_v = self.attention.infer( + utterance=layer_norm_input_utterance, + lengths=lengths, + right_context=layer_norm_input_right_context, + summary=layer_norm_input_summary, + mems=pre_mems, + left_context_key=lc_key, + left_context_val=lc_val, + ) + + output_utterance, output_right_context, conv_cache = self._apply_post_attention( + rc_output, ffn0_out, conv_cache, right_context.size(0), utterance.size(0) + ) + output_state = self._pack_state(next_k, next_v, utterance.size(0), mems, conv_cache, state) + return output_utterance, output_right_context, output_state, next_m + + +class ConvEmformer(_EmformerImpl): + r"""Implements the convolution-augmented streaming transformer architecture introduced in + *Streaming Transformer Transducer based Speech Recognition Using Non-Causal Convolution* + :cite:`9747706`. + + Args: + input_dim (int): input dimension. + num_heads (int): number of attention heads in each ConvEmformer layer. + ffn_dim (int): hidden layer dimension of each ConvEmformer layer's feedforward network. + num_layers (int): number of ConvEmformer layers to instantiate. + segment_length (int): length of each input segment. + kernel_size (int): size of kernel to use in convolution modules. + dropout (float, optional): dropout probability. (Default: 0.0) + ffn_activation (str, optional): activation function to use in feedforward networks. + Must be one of ("relu", "gelu", "silu"). (Default: "relu") + left_context_length (int, optional): length of left context. (Default: 0) + right_context_length (int, optional): length of right context. (Default: 0) + max_memory_size (int, optional): maximum number of memory elements to use. (Default: 0) + weight_init_scale_strategy (str or None, optional): per-layer weight initialization scaling + strategy. Must be one of ("depthwise", "constant", ``None``). (Default: "depthwise") + tanh_on_mem (bool, optional): if ``True``, applies tanh to memory elements. (Default: ``False``) + negative_inf (float, optional): value to use for negative infinity in attention weights. (Default: -1e8) + conv_activation (str, optional): activation function to use in convolution modules. + Must be one of ("relu", "gelu", "silu"). (Default: "silu") + + Examples: + >>> conv_emformer = ConvEmformer(80, 4, 1024, 12, 16, 8, right_context_length=4) + >>> input = torch.rand(10, 200, 80) + >>> lengths = torch.randint(1, 200, (10,)) + >>> output, lengths = conv_emformer(input, lengths) + >>> input = torch.rand(4, 20, 80) + >>> lengths = torch.ones(4) * 20 + >>> output, lengths, states = conv_emformer.infer(input, lengths, None) + """ + + def __init__( + self, + input_dim: int, + num_heads: int, + ffn_dim: int, + num_layers: int, + segment_length: int, + kernel_size: int, + dropout: float = 0.0, + ffn_activation: str = "relu", + left_context_length: int = 0, + right_context_length: int = 0, + max_memory_size: int = 0, + weight_init_scale_strategy: Optional[str] = "depthwise", + tanh_on_mem: bool = False, + negative_inf: float = -1e8, + conv_activation: str = "silu", + ): + weight_init_gains = _get_weight_init_gains(weight_init_scale_strategy, num_layers) + emformer_layers = torch.nn.ModuleList( + [ + _ConvEmformerLayer( + input_dim, + num_heads, + ffn_dim, + segment_length, + kernel_size, + dropout=dropout, + ffn_activation=ffn_activation, + left_context_length=left_context_length, + right_context_length=right_context_length, + max_memory_size=max_memory_size, + weight_init_gain=weight_init_gains[layer_idx], + tanh_on_mem=tanh_on_mem, + negative_inf=negative_inf, + conv_activation=conv_activation, + ) + for layer_idx in range(num_layers) + ] + ) + super().__init__( + emformer_layers, + segment_length, + left_context_length=left_context_length, + right_context_length=right_context_length, + max_memory_size=max_memory_size, + ) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/hifi_gan.py b/lib/python3.12/site-packages/torchaudio/prototype/models/hifi_gan.py new file mode 100644 index 0000000000000000000000000000000000000000..93d92e1854651358367c6388fa3d916941961367 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/hifi_gan.py @@ -0,0 +1,336 @@ +""" +MIT License + +Copyright (c) 2020 Jungil Kong + +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 typing import Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Conv1d, ConvTranspose1d + + +class HiFiGANVocoder(torch.nn.Module): + """Generator part of *HiFi GAN* :cite:`NEURIPS2020_c5d73680`. + Source: https://github.com/jik876/hifi-gan/blob/4769534d45265d52a904b850da5a622601885777/models.py#L75 + + Note: + To build the model, please use one of the factory functions: :py:func:`hifigan_vocoder`, + :py:func:`hifigan_vocoder_v1`, :py:func:`hifigan_vocoder_v2`, :py:func:`hifigan_vocoder_v3`. + + Args: + in_channels (int): Number of channels in the input features. + upsample_rates (tuple of ``int``): Factors by which each upsampling layer increases the time dimension. + upsample_initial_channel (int): Number of channels in the input feature tensor. + upsample_kernel_sizes (tuple of ``int``): Kernel size for each upsampling layer. + resblock_kernel_sizes (tuple of ``int``): Kernel size for each residual block. + resblock_dilation_sizes (tuple of tuples of ``int``): Dilation sizes for each 1D convolutional layer in each + residual block. For resblock type 1 inner tuples should have length 3, because there are 3 + convolutions in each layer. For resblock type 2 they should have length 2. + resblock_type (int, 1 or 2): Determines whether ``ResBlock1`` or ``ResBlock2`` will be used. + lrelu_slope (float): Slope of leaky ReLUs in activations. + """ + + def __init__( + self, + in_channels: int, + upsample_rates: Tuple[int, ...], + upsample_initial_channel: int, + upsample_kernel_sizes: Tuple[int, ...], + resblock_kernel_sizes: Tuple[int, ...], + resblock_dilation_sizes: Tuple[Tuple[int, ...], ...], + resblock_type: int, + lrelu_slope: float, + ): + super(HiFiGANVocoder, self).__init__() + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.conv_pre = Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3) + resblock = ResBlock1 if resblock_type == 1 else ResBlock2 + + self.ups = nn.ModuleList() + for i, (u, k) in enumerate(zip(upsample_rates, upsample_kernel_sizes)): + self.ups.append( + ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + k, + u, + padding=(k - u) // 2, + ) + ) + + self.resblocks = nn.ModuleList() + for i in range(len(self.ups)): + ch = upsample_initial_channel // (2 ** (i + 1)) + for (k, d) in zip(resblock_kernel_sizes, resblock_dilation_sizes): + self.resblocks.append(resblock(ch, k, d, lrelu_slope)) + + self.conv_post = Conv1d(ch, 1, 7, 1, padding=3) + self.lrelu_slope = lrelu_slope + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (Tensor): Feature input tensor of shape `(batch_size, num_channels, time_length)`. + + Returns: + Tensor of shape `(batch_size, 1, time_length * upsample_rate)`, where `upsample_rate` is the product + of upsample rates for all layers. + """ + x = self.conv_pre(x) + for i, upsampling_layer in enumerate(self.ups): + x = F.leaky_relu(x, self.lrelu_slope) + x = upsampling_layer(x) + xs = torch.zeros_like(x) + for j in range(self.num_kernels): + res_block: ResBlockInterface = self.resblocks[i * self.num_kernels + j] + xs += res_block.forward(x) + x = xs / self.num_kernels + + x = F.leaky_relu(x) + x = self.conv_post(x) + x = torch.tanh(x) + + return x + + +@torch.jit.interface +class ResBlockInterface(torch.nn.Module): + """Interface for ResBlock - necessary to make type annotations in ``HiFiGANVocoder.forward`` compatible + with TorchScript + """ + + def forward(self, x: torch.Tensor) -> torch.Tensor: + pass + + +class ResBlock1(torch.nn.Module): + """Residual block of type 1 for HiFiGAN Vocoder :cite:`NEURIPS2020_c5d73680`. + Args: + channels (int): Number of channels in the input features. + kernel_size (int, optional): Kernel size for 1D convolutions. (Default: ``3``) + dilation (tuple of 3 ``int``, optional): Dilations for each 1D convolution. (Default: ``(1, 3, 5)``) + lrelu_slope (float): Slope of leaky ReLUs in activations. + """ + + def __init__( + self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int, int] = (1, 3, 5), lrelu_slope: float = 0.1 + ): + super(ResBlock1, self).__init__() + self.convs1 = nn.ModuleList( + [ + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ), + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ), + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ), + ] + ) + + self.convs2 = nn.ModuleList( + [ + Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), + Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), + Conv1d(channels, channels, kernel_size, 1, dilation=1, padding=get_padding(kernel_size, 1)), + ] + ) + self.lrelu_slope = lrelu_slope + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x (Tensor): input of shape ``(batch_size, channels, time_length)``. + Returns: + Tensor of the same shape as input. + """ + for conv1, conv2 in zip(self.convs1, self.convs2): + xt = F.leaky_relu(x, self.lrelu_slope) + xt = conv1(xt) + xt = F.leaky_relu(xt, self.lrelu_slope) + xt = conv2(xt) + x = xt + x + return x + + +class ResBlock2(torch.nn.Module): + """Residual block of type 2 for HiFiGAN Vocoder :cite:`NEURIPS2020_c5d73680`. + Args: + channels (int): Number of channels in the input features. + kernel_size (int, optional): Kernel size for 1D convolutions. (Default: ``3``) + dilation (tuple of 2 ``int``, optional): Dilations for each 1D convolution. (Default: ``(1, 3)``) + lrelu_slope (float): Slope of leaky ReLUs in activations. + """ + + def __init__( + self, channels: int, kernel_size: int = 3, dilation: Tuple[int, int] = (1, 3), lrelu_slope: float = 0.1 + ): + super(ResBlock2, self).__init__() + self.convs = nn.ModuleList( + [ + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ), + Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ), + ] + ) + self.lrelu_slope = lrelu_slope + + def forward(self, x: torch.Tensor): + """ + Args: + x (Tensor): input of shape ``(batch_size, channels, time_length)``. + Returns: + Tensor of the same shape as input. + """ + for c in self.convs: + xt = F.leaky_relu(x, self.lrelu_slope) + xt = c(xt) + x = xt + x + return x + + +def get_padding(kernel_size, dilation=1): + """Find padding for which 1D convolution preserves the input shape.""" + return int((kernel_size * dilation - dilation) / 2) + + +def hifigan_vocoder( + in_channels: int, + upsample_rates: Tuple[int, ...], + upsample_initial_channel: int, + upsample_kernel_sizes: Tuple[int, ...], + resblock_kernel_sizes: Tuple[int, ...], + resblock_dilation_sizes: Tuple[Tuple[int, ...], ...], + resblock_type: int, + lrelu_slope: float, +) -> HiFiGANVocoder: + r"""Builds HiFi GAN Vocoder :cite:`NEURIPS2020_c5d73680`. + + Args: + in_channels (int): See :py:class:`HiFiGANVocoder`. + upsample_rates (tuple of ``int``): See :py:class:`HiFiGANVocoder`. + upsample_initial_channel (int): See :py:class:`HiFiGANVocoder`. + upsample_kernel_sizes (tuple of ``int``): See :py:class:`HiFiGANVocoder`. + resblock_kernel_sizes (tuple of ``int``): See :py:class:`HiFiGANVocoder`. + resblock_dilation_sizes (tuple of tuples of ``int``): See :py:class:`HiFiGANVocoder`. + resblock_type (int, 1 or 2): See :py:class:`HiFiGANVocoder`. + Returns: + HiFiGANVocoder: generated model. + """ + + return HiFiGANVocoder( + upsample_rates=upsample_rates, + resblock_kernel_sizes=resblock_kernel_sizes, + resblock_dilation_sizes=resblock_dilation_sizes, + resblock_type=resblock_type, + upsample_initial_channel=upsample_initial_channel, + upsample_kernel_sizes=upsample_kernel_sizes, + in_channels=in_channels, + lrelu_slope=lrelu_slope, + ) + + +def hifigan_vocoder_v1() -> HiFiGANVocoder: + r"""Builds HiFiGAN Vocoder with V1 architecture :cite:`NEURIPS2020_c5d73680`. + + Returns: + HiFiGANVocoder: generated model. + """ + return hifigan_vocoder( + upsample_rates=(8, 8, 2, 2), + upsample_kernel_sizes=(16, 16, 4, 4), + upsample_initial_channel=512, + resblock_kernel_sizes=(3, 7, 11), + resblock_dilation_sizes=((1, 3, 5), (1, 3, 5), (1, 3, 5)), + resblock_type=1, + in_channels=80, + lrelu_slope=0.1, + ) + + +def hifigan_vocoder_v2() -> HiFiGANVocoder: + r"""Builds HiFiGAN Vocoder with V2 architecture :cite:`NEURIPS2020_c5d73680`. + + Returns: + HiFiGANVocoder: generated model. + """ + return hifigan_vocoder( + upsample_rates=(8, 8, 2, 2), + upsample_kernel_sizes=(16, 16, 4, 4), + upsample_initial_channel=128, + resblock_kernel_sizes=(3, 7, 11), + resblock_dilation_sizes=((1, 3, 5), (1, 3, 5), (1, 3, 5)), + resblock_type=1, + in_channels=80, + lrelu_slope=0.1, + ) + + +def hifigan_vocoder_v3() -> HiFiGANVocoder: + r"""Builds HiFiGAN Vocoder with V3 architecture :cite:`NEURIPS2020_c5d73680`. + + Returns: + HiFiGANVocoder: generated model. + """ + return hifigan_vocoder( + upsample_rates=(8, 8, 4), + upsample_kernel_sizes=(16, 16, 8), + upsample_initial_channel=256, + resblock_kernel_sizes=(3, 5, 7), + resblock_dilation_sizes=((1, 2), (2, 6), (3, 12)), + resblock_type=2, + in_channels=80, + lrelu_slope=0.1, + ) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/rnnt.py b/lib/python3.12/site-packages/torchaudio/prototype/models/rnnt.py new file mode 100644 index 0000000000000000000000000000000000000000..aa7e32d5b961a80b2637995a3e3be07eaa5d7165 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/rnnt.py @@ -0,0 +1,711 @@ +import math +from typing import Dict, List, Optional, Tuple + +import torch +from torchaudio.models import Conformer, RNNT +from torchaudio.models.rnnt import _Joiner, _Predictor, _TimeReduction, _Transcriber + + +TrieNode = Tuple[Dict[int, "TrieNode"], int, Optional[Tuple[int, int]]] + + +class _ConformerEncoder(torch.nn.Module, _Transcriber): + def __init__( + self, + *, + input_dim: int, + output_dim: int, + time_reduction_stride: int, + conformer_input_dim: int, + conformer_ffn_dim: int, + conformer_num_layers: int, + conformer_num_heads: int, + conformer_depthwise_conv_kernel_size: int, + conformer_dropout: float, + ) -> None: + super().__init__() + self.time_reduction = _TimeReduction(time_reduction_stride) + self.input_linear = torch.nn.Linear(input_dim * time_reduction_stride, conformer_input_dim) + self.conformer = Conformer( + num_layers=conformer_num_layers, + input_dim=conformer_input_dim, + ffn_dim=conformer_ffn_dim, + num_heads=conformer_num_heads, + depthwise_conv_kernel_size=conformer_depthwise_conv_kernel_size, + dropout=conformer_dropout, + use_group_norm=True, + convolution_first=True, + ) + self.output_linear = torch.nn.Linear(conformer_input_dim, output_dim) + self.layer_norm = torch.nn.LayerNorm(output_dim) + + def forward(self, input: torch.Tensor, lengths: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + time_reduction_out, time_reduction_lengths = self.time_reduction(input, lengths) + input_linear_out = self.input_linear(time_reduction_out) + x, lengths = self.conformer(input_linear_out, time_reduction_lengths) + output_linear_out = self.output_linear(x) + layer_norm_out = self.layer_norm(output_linear_out) + return layer_norm_out, lengths + + def infer( + self, + input: torch.Tensor, + lengths: torch.Tensor, + states: Optional[List[List[torch.Tensor]]], + ) -> Tuple[torch.Tensor, torch.Tensor, List[List[torch.Tensor]]]: + raise RuntimeError("Conformer does not support streaming inference.") + + +class _JoinerBiasing(torch.nn.Module): + r"""Recurrent neural network transducer (RNN-T) joint network. + + Args: + input_dim (int): source and target input dimension. + output_dim (int): output dimension. + activation (str, optional): activation function to use in the joiner. + Must be one of ("relu", "tanh"). (Default: "relu") + biasing (bool): perform biasing + deepbiasing (bool): perform deep biasing + attndim (int): dimension of the biasing vector hptr + + """ + + def __init__( + self, + input_dim: int, + output_dim: int, + activation: str = "relu", + biasing: bool = False, + deepbiasing: bool = False, + attndim: int = 1, + ) -> None: + super().__init__() + self.linear = torch.nn.Linear(input_dim, output_dim, bias=True) + self.biasing = biasing + self.deepbiasing = deepbiasing + if self.biasing and self.deepbiasing: + self.biasinglinear = torch.nn.Linear(attndim, input_dim, bias=True) + self.attndim = attndim + if activation == "relu": + self.activation = torch.nn.ReLU() + elif activation == "tanh": + self.activation = torch.nn.Tanh() + else: + raise ValueError(f"Unsupported activation {activation}") + + def forward( + self, + source_encodings: torch.Tensor, + source_lengths: torch.Tensor, + target_encodings: torch.Tensor, + target_lengths: torch.Tensor, + hptr: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + r"""Forward pass for training. + + B: batch size; + T: maximum source sequence length in batch; + U: maximum target sequence length in batch; + D: dimension of each source and target sequence encoding. + + Args: + source_encodings (torch.Tensor): source encoding sequences, with + shape `(B, T, D)`. + source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing + valid sequence length of i-th batch element in ``source_encodings``. + target_encodings (torch.Tensor): target encoding sequences, with shape `(B, U, D)`. + target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing + valid sequence length of i-th batch element in ``target_encodings``. + hptr (torch.Tensor): deep biasing vector with shape `(B, T, U, A)`. + + Returns: + (torch.Tensor, torch.Tensor, torch.Tensor): + torch.Tensor + joint network output, with shape `(B, T, U, output_dim)`. + torch.Tensor + output source lengths, with shape `(B,)` and i-th element representing + number of valid elements along dim 1 for i-th batch element in joint network output. + torch.Tensor + output target lengths, with shape `(B,)` and i-th element representing + number of valid elements along dim 2 for i-th batch element in joint network output. + torch.Tensor + joint network second last layer output (i.e. before self.linear), with shape `(B, T, U, D)`. + """ + joint_encodings = source_encodings.unsqueeze(2).contiguous() + target_encodings.unsqueeze(1).contiguous() + if self.biasing and self.deepbiasing and hptr is not None: + hptr = self.biasinglinear(hptr) + joint_encodings += hptr + elif self.biasing and self.deepbiasing: + # Hack here for unused parameters + joint_encodings += self.biasinglinear(joint_encodings.new_zeros(1, self.attndim)).mean() * 0 + activation_out = self.activation(joint_encodings) + output = self.linear(activation_out) + return output, source_lengths, target_lengths, activation_out + + +class RNNTBiasing(RNNT): + r"""torchaudio.models.RNNT() + + Recurrent neural network transducer (RNN-T) model. + + Note: + To build the model, please use one of the factory functions. + + Args: + transcriber (torch.nn.Module): transcription network. + predictor (torch.nn.Module): prediction network. + joiner (torch.nn.Module): joint network. + attndim (int): TCPGen attention dimension + biasing (bool): If true, use biasing, otherwise use standard RNN-T + deepbiasing (bool): If true, use deep biasing by extracting the biasing vector + embdim (int): dimension of symbol embeddings + jointdim (int): dimension of the joint network joint dimension + charlist (list): The list of word piece tokens in the same order as the output layer + encoutdim (int): dimension of the encoder output vectors + dropout_tcpgen (float): dropout rate for TCPGen + tcpsche (int): The epoch at which TCPGen starts to train + DBaverage (bool): If true, instead of TCPGen, use DBRNNT for biasing + """ + + def __init__( + self, + transcriber: _Transcriber, + predictor: _Predictor, + joiner: _Joiner, + attndim: int, + biasing: bool, + deepbiasing: bool, + embdim: int, + jointdim: int, + charlist: List[str], + encoutdim: int, + dropout_tcpgen: float, + tcpsche: int, + DBaverage: bool, + ) -> None: + super().__init__(transcriber, predictor, joiner) + self.attndim = attndim + self.deepbiasing = deepbiasing + self.jointdim = jointdim + self.embdim = embdim + self.encoutdim = encoutdim + self.char_list = charlist or [] + self.blank_idx = self.char_list.index("") + self.nchars = len(self.char_list) + self.DBaverage = DBaverage + self.biasing = biasing + if self.biasing: + if self.deepbiasing and self.DBaverage: + # Deep biasing without TCPGen + self.biasingemb = torch.nn.Linear(self.nchars, self.attndim, bias=False) + else: + # TCPGen parameters + self.ooKBemb = torch.nn.Embedding(1, self.embdim) + self.Qproj_char = torch.nn.Linear(self.embdim, self.attndim) + self.Qproj_acoustic = torch.nn.Linear(self.encoutdim, self.attndim) + self.Kproj = torch.nn.Linear(self.embdim, self.attndim) + self.pointer_gate = torch.nn.Linear(self.attndim + self.jointdim, 1) + self.dropout_tcpgen = torch.nn.Dropout(dropout_tcpgen) + self.tcpsche = tcpsche + + def forward( + self, + sources: torch.Tensor, + source_lengths: torch.Tensor, + targets: torch.Tensor, + target_lengths: torch.Tensor, + tries: TrieNode, + current_epoch: int, + predictor_state: Optional[List[List[torch.Tensor]]] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, List[List[torch.Tensor]], torch.Tensor, torch.Tensor]: + r"""Forward pass for training. + + B: batch size; + T: maximum source sequence length in batch; + U: maximum target sequence length in batch; + D: feature dimension of each source sequence element. + + Args: + sources (torch.Tensor): source frame sequences right-padded with right context, with + shape `(B, T, D)`. + source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing + number of valid frames for i-th batch element in ``sources``. + targets (torch.Tensor): target sequences, with shape `(B, U)` and each element + mapping to a target symbol. + target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing + number of valid frames for i-th batch element in ``targets``. + tries (TrieNode): wordpiece prefix trees representing the biasing list to be searched + current_epoch (Int): the current epoch number to determine if TCPGen should be trained + at this epoch + predictor_state (List[List[torch.Tensor]] or None, optional): list of lists of tensors + representing prediction network internal state generated in preceding invocation + of ``forward``. (Default: ``None``) + + Returns: + (torch.Tensor, torch.Tensor, torch.Tensor, List[List[torch.Tensor]]): + torch.Tensor + joint network output, with shape + `(B, max output source length, max output target length, output_dim (number of target symbols))`. + torch.Tensor + output source lengths, with shape `(B,)` and i-th element representing + number of valid elements along dim 1 for i-th batch element in joint network output. + torch.Tensor + output target lengths, with shape `(B,)` and i-th element representing + number of valid elements along dim 2 for i-th batch element in joint network output. + List[List[torch.Tensor]] + output states; list of lists of tensors + representing prediction network internal state generated in current invocation + of ``forward``. + torch.Tensor + TCPGen distribution, with shape + `(B, max output source length, max output target length, output_dim (number of target symbols))`. + torch.Tensor + Generation probability (or copy probability), with shape + `(B, max output source length, max output target length, 1)`. + """ + source_encodings, source_lengths = self.transcriber( + input=sources, + lengths=source_lengths, + ) + target_encodings, target_lengths, predictor_state = self.predictor( + input=targets, + lengths=target_lengths, + state=predictor_state, + ) + # Forward TCPGen + hptr = None + tcpgen_dist, p_gen = None, None + if self.biasing and current_epoch >= self.tcpsche and tries != []: + ptrdist_mask, p_gen_mask = self.get_tcpgen_step_masks(targets, tries) + hptr, tcpgen_dist = self.forward_tcpgen(targets, ptrdist_mask, source_encodings) + hptr = self.dropout_tcpgen(hptr) + elif self.biasing: + # Hack here to bypass unused parameters + if self.DBaverage and self.deepbiasing: + dummy = self.biasingemb(source_encodings.new_zeros(1, len(self.char_list))).mean() + else: + dummy = source_encodings.new_zeros(1, self.embdim) + dummy = self.Qproj_char(dummy).mean() + dummy += self.Qproj_acoustic(source_encodings.new_zeros(1, source_encodings.size(-1))).mean() + dummy += self.Kproj(source_encodings.new_zeros(1, self.embdim)).mean() + dummy += self.pointer_gate(source_encodings.new_zeros(1, self.attndim + self.jointdim)).mean() + dummy += self.ooKBemb.weight.mean() + dummy = dummy * 0 + source_encodings += dummy + + output, source_lengths, target_lengths, jointer_activation = self.joiner( + source_encodings=source_encodings, + source_lengths=source_lengths, + target_encodings=target_encodings, + target_lengths=target_lengths, + hptr=hptr, + ) + + # Calculate Generation Probability + if self.biasing and hptr is not None and tcpgen_dist is not None: + p_gen = torch.sigmoid(self.pointer_gate(torch.cat((jointer_activation, hptr), dim=-1))) + # avoid collapsing to ooKB token in the first few updates + # if current_epoch == self.tcpsche: + # p_gen = p_gen * 0.1 + p_gen = p_gen.masked_fill(p_gen_mask.bool().unsqueeze(1).unsqueeze(-1), 0) + + return (output, source_lengths, target_lengths, predictor_state, tcpgen_dist, p_gen) + + def get_tcpgen_distribution(self, query, ptrdist_mask): + # Make use of the predictor embedding matrix + keyvalues = torch.cat([self.predictor.embedding.weight.data, self.ooKBemb.weight], dim=0) + keyvalues = self.dropout_tcpgen(self.Kproj(keyvalues)) + # B * T * U * attndim, nbpe * attndim -> B * T * U * nbpe + tcpgendist = torch.einsum("ntuj,ij->ntui", query, keyvalues) + tcpgendist = tcpgendist / math.sqrt(query.size(-1)) + ptrdist_mask = ptrdist_mask.unsqueeze(1).repeat(1, tcpgendist.size(1), 1, 1) + tcpgendist.masked_fill_(ptrdist_mask.bool(), -1e9) + tcpgendist = torch.nn.functional.softmax(tcpgendist, dim=-1) + # B * T * U * nbpe, nbpe * attndim -> B * T * U * attndim + hptr = torch.einsum("ntui,ij->ntuj", tcpgendist[:, :, :, :-1], keyvalues[:-1, :]) + return hptr, tcpgendist + + def forward_tcpgen(self, targets, ptrdist_mask, source_encodings): + tcpgen_dist = None + if self.DBaverage and self.deepbiasing: + hptr = self.biasingemb(1 - ptrdist_mask[:, :, :-1].float()).unsqueeze(1) + else: + query_char = self.predictor.embedding(targets) + query_char = self.Qproj_char(query_char).unsqueeze(1) # B * 1 * U * attndim + query_acoustic = self.Qproj_acoustic(source_encodings).unsqueeze(2) # B * T * 1 * attndim + query = query_char + query_acoustic # B * T * U * attndim + hptr, tcpgen_dist = self.get_tcpgen_distribution(query, ptrdist_mask) + return hptr, tcpgen_dist + + def get_tcpgen_step_masks(self, yseqs, resettrie): + seqlen = len(yseqs[0]) + batch_masks = yseqs.new_ones(len(yseqs), seqlen, len(self.char_list) + 1) + p_gen_masks = [] + for i, yseq in enumerate(yseqs): + new_tree = resettrie + p_gen_mask = [] + for j, vy in enumerate(yseq): + vy = vy.item() + new_tree = new_tree[0] + if vy in [self.blank_idx]: + new_tree = resettrie + p_gen_mask.append(0) + elif self.char_list[vy].endswith("▁"): + if vy in new_tree and new_tree[vy][0] != {}: + new_tree = new_tree[vy] + else: + new_tree = resettrie + p_gen_mask.append(0) + elif vy not in new_tree: + new_tree = [{}] + p_gen_mask.append(1) + else: + new_tree = new_tree[vy] + p_gen_mask.append(0) + batch_masks[i, j, list(new_tree[0].keys())] = 0 + # In the original paper, ooKB node was not masked + # In this implementation, if not masking ooKB, ooKB probability + # would quickly collapse to 1.0 in the first few updates. + # Haven't found out why this happened. + # batch_masks[i, j, -1] = 0 + p_gen_masks.append(p_gen_mask + [1] * (seqlen - len(p_gen_mask))) + p_gen_masks = torch.Tensor(p_gen_masks).to(yseqs.device).byte() + return batch_masks, p_gen_masks + + def get_tcpgen_step_masks_prefix(self, yseqs, resettrie): + # Implemented for prefix-based wordpieces, not tested yet + seqlen = len(yseqs[0]) + batch_masks = yseqs.new_ones(len(yseqs), seqlen, len(self.char_list) + 1) + p_gen_masks = [] + for i, yseq in enumerate(yseqs): + p_gen_mask = [] + new_tree = resettrie + for j, vy in enumerate(yseq): + vy = vy.item() + new_tree = new_tree[0] + if vy in [self.blank_idx]: + new_tree = resettrie + batch_masks[i, j, list(new_tree[0].keys())] = 0 + elif self.char_list[vy].startswith("▁"): + new_tree = resettrie + if vy not in new_tree[0]: + batch_masks[i, j, list(new_tree[0].keys())] = 0 + else: + new_tree = new_tree[0][vy] + batch_masks[i, j, list(new_tree[0].keys())] = 0 + if new_tree[1] != -1: + batch_masks[i, j, list(resettrie[0].keys())] = 0 + else: + if vy not in new_tree: + new_tree = resettrie + batch_masks[i, j, list(new_tree[0].keys())] = 0 + else: + new_tree = new_tree[vy] + batch_masks[i, j, list(new_tree[0].keys())] = 0 + if new_tree[1] != -1: + batch_masks[i, j, list(resettrie[0].keys())] = 0 + p_gen_mask.append(0) + # batch_masks[i, j, -1] = 0 + p_gen_masks.append(p_gen_mask + [1] * (seqlen - len(p_gen_mask))) + p_gen_masks = torch.Tensor(p_gen_masks).to(yseqs.device).byte() + + return batch_masks, p_gen_masks + + def get_tcpgen_step(self, vy, trie, resettrie): + new_tree = trie[0] + if vy in [self.blank_idx]: + new_tree = resettrie + elif self.char_list[vy].endswith("▁"): + if vy in new_tree and new_tree[vy][0] != {}: + new_tree = new_tree[vy] + else: + new_tree = resettrie + elif vy not in new_tree: + new_tree = [{}] + else: + new_tree = new_tree[vy] + return new_tree + + def join( + self, + source_encodings: torch.Tensor, + source_lengths: torch.Tensor, + target_encodings: torch.Tensor, + target_lengths: torch.Tensor, + hptr: torch.Tensor = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + r"""Applies joint network to source and target encodings. + + B: batch size; + T: maximum source sequence length in batch; + U: maximum target sequence length in batch; + D: dimension of each source and target sequence encoding. + A: TCPGen attention dimension + + Args: + source_encodings (torch.Tensor): source encoding sequences, with + shape `(B, T, D)`. + source_lengths (torch.Tensor): with shape `(B,)` and i-th element representing + valid sequence length of i-th batch element in ``source_encodings``. + target_encodings (torch.Tensor): target encoding sequences, with shape `(B, U, D)`. + target_lengths (torch.Tensor): with shape `(B,)` and i-th element representing + valid sequence length of i-th batch element in ``target_encodings``. + hptr (torch.Tensor): deep biasing vector with shape `(B, T, U, A)`. + + Returns: + (torch.Tensor, torch.Tensor, torch.Tensor): + torch.Tensor + joint network output, with shape `(B, T, U, output_dim)`. + torch.Tensor + output source lengths, with shape `(B,)` and i-th element representing + number of valid elements along dim 1 for i-th batch element in joint network output. + torch.Tensor + joint network second last layer output, with shape `(B, T, U, D)`. + """ + output, source_lengths, target_lengths, jointer_activation = self.joiner( + source_encodings=source_encodings, + source_lengths=source_lengths, + target_encodings=target_encodings, + target_lengths=target_lengths, + hptr=hptr, + ) + return output, source_lengths, jointer_activation + + +def conformer_rnnt_model( + *, + input_dim: int, + encoding_dim: int, + time_reduction_stride: int, + conformer_input_dim: int, + conformer_ffn_dim: int, + conformer_num_layers: int, + conformer_num_heads: int, + conformer_depthwise_conv_kernel_size: int, + conformer_dropout: float, + num_symbols: int, + symbol_embedding_dim: int, + num_lstm_layers: int, + lstm_hidden_dim: int, + lstm_layer_norm: int, + lstm_layer_norm_epsilon: int, + lstm_dropout: int, + joiner_activation: str, +) -> RNNT: + r"""Builds Conformer-based recurrent neural network transducer (RNN-T) model. + + Args: + input_dim (int): dimension of input sequence frames passed to transcription network. + encoding_dim (int): dimension of transcription- and prediction-network-generated encodings + passed to joint network. + time_reduction_stride (int): factor by which to reduce length of input sequence. + conformer_input_dim (int): dimension of Conformer input. + conformer_ffn_dim (int): hidden layer dimension of each Conformer layer's feedforward network. + conformer_num_layers (int): number of Conformer layers to instantiate. + conformer_num_heads (int): number of attention heads in each Conformer layer. + conformer_depthwise_conv_kernel_size (int): kernel size of each Conformer layer's depthwise convolution layer. + conformer_dropout (float): Conformer dropout probability. + num_symbols (int): cardinality of set of target tokens. + symbol_embedding_dim (int): dimension of each target token embedding. + num_lstm_layers (int): number of LSTM layers to instantiate. + lstm_hidden_dim (int): output dimension of each LSTM layer. + lstm_layer_norm (bool): if ``True``, enables layer normalization for LSTM layers. + lstm_layer_norm_epsilon (float): value of epsilon to use in LSTM layer normalization layers. + lstm_dropout (float): LSTM dropout probability. + joiner_activation (str): activation function to use in the joiner. + Must be one of ("relu", "tanh"). (Default: "relu") + + Returns: + RNNT: + Conformer RNN-T model. + """ + encoder = _ConformerEncoder( + input_dim=input_dim, + output_dim=encoding_dim, + time_reduction_stride=time_reduction_stride, + conformer_input_dim=conformer_input_dim, + conformer_ffn_dim=conformer_ffn_dim, + conformer_num_layers=conformer_num_layers, + conformer_num_heads=conformer_num_heads, + conformer_depthwise_conv_kernel_size=conformer_depthwise_conv_kernel_size, + conformer_dropout=conformer_dropout, + ) + predictor = _Predictor( + num_symbols=num_symbols, + output_dim=encoding_dim, + symbol_embedding_dim=symbol_embedding_dim, + num_lstm_layers=num_lstm_layers, + lstm_hidden_dim=lstm_hidden_dim, + lstm_layer_norm=lstm_layer_norm, + lstm_layer_norm_epsilon=lstm_layer_norm_epsilon, + lstm_dropout=lstm_dropout, + ) + joiner = _Joiner(encoding_dim, num_symbols, activation=joiner_activation) + return RNNT(encoder, predictor, joiner) + + +def conformer_rnnt_base() -> RNNT: + r"""Builds basic version of Conformer RNN-T model. + + Returns: + RNNT: + Conformer RNN-T model. + """ + return conformer_rnnt_model( + input_dim=80, + encoding_dim=1024, + time_reduction_stride=4, + conformer_input_dim=256, + conformer_ffn_dim=1024, + conformer_num_layers=16, + conformer_num_heads=4, + conformer_depthwise_conv_kernel_size=31, + conformer_dropout=0.1, + num_symbols=1024, + symbol_embedding_dim=256, + num_lstm_layers=2, + lstm_hidden_dim=512, + lstm_layer_norm=True, + lstm_layer_norm_epsilon=1e-5, + lstm_dropout=0.3, + joiner_activation="tanh", + ) + + +def conformer_rnnt_biasing( + *, + input_dim: int, + encoding_dim: int, + time_reduction_stride: int, + conformer_input_dim: int, + conformer_ffn_dim: int, + conformer_num_layers: int, + conformer_num_heads: int, + conformer_depthwise_conv_kernel_size: int, + conformer_dropout: float, + num_symbols: int, + symbol_embedding_dim: int, + num_lstm_layers: int, + lstm_hidden_dim: int, + lstm_layer_norm: int, + lstm_layer_norm_epsilon: int, + lstm_dropout: int, + joiner_activation: str, + attndim: int, + biasing: bool, + charlist: List[str], + deepbiasing: bool, + tcpsche: int, + DBaverage: bool, +) -> RNNTBiasing: + r"""Builds Conformer-based recurrent neural network transducer (RNN-T) model. + + Args: + input_dim (int): dimension of input sequence frames passed to transcription network. + encoding_dim (int): dimension of transcription- and prediction-network-generated encodings + passed to joint network. + time_reduction_stride (int): factor by which to reduce length of input sequence. + conformer_input_dim (int): dimension of Conformer input. + conformer_ffn_dim (int): hidden layer dimension of each Conformer layer's feedforward network. + conformer_num_layers (int): number of Conformer layers to instantiate. + conformer_num_heads (int): number of attention heads in each Conformer layer. + conformer_depthwise_conv_kernel_size (int): kernel size of each Conformer layer's depthwise convolution layer. + conformer_dropout (float): Conformer dropout probability. + num_symbols (int): cardinality of set of target tokens. + symbol_embedding_dim (int): dimension of each target token embedding. + num_lstm_layers (int): number of LSTM layers to instantiate. + lstm_hidden_dim (int): output dimension of each LSTM layer. + lstm_layer_norm (bool): if ``True``, enables layer normalization for LSTM layers. + lstm_layer_norm_epsilon (float): value of epsilon to use in LSTM layer normalization layers. + lstm_dropout (float): LSTM dropout probability. + joiner_activation (str): activation function to use in the joiner. + Must be one of ("relu", "tanh"). (Default: "relu") + attndim (int): TCPGen attention dimension + biasing (bool): If true, use biasing, otherwise use standard RNN-T + charlist (list): The list of word piece tokens in the same order as the output layer + deepbiasing (bool): If true, use deep biasing by extracting the biasing vector + tcpsche (int): The epoch at which TCPGen starts to train + DBaverage (bool): If true, instead of TCPGen, use DBRNNT for biasing + + Returns: + RNNT: + Conformer RNN-T model with TCPGen-based biasing support. + """ + encoder = _ConformerEncoder( + input_dim=input_dim, + output_dim=encoding_dim, + time_reduction_stride=time_reduction_stride, + conformer_input_dim=conformer_input_dim, + conformer_ffn_dim=conformer_ffn_dim, + conformer_num_layers=conformer_num_layers, + conformer_num_heads=conformer_num_heads, + conformer_depthwise_conv_kernel_size=conformer_depthwise_conv_kernel_size, + conformer_dropout=conformer_dropout, + ) + predictor = _Predictor( + num_symbols=num_symbols, + output_dim=encoding_dim, + symbol_embedding_dim=symbol_embedding_dim, + num_lstm_layers=num_lstm_layers, + lstm_hidden_dim=lstm_hidden_dim, + lstm_layer_norm=lstm_layer_norm, + lstm_layer_norm_epsilon=lstm_layer_norm_epsilon, + lstm_dropout=lstm_dropout, + ) + joiner = _JoinerBiasing( + encoding_dim, + num_symbols, + activation=joiner_activation, + deepbiasing=deepbiasing, + attndim=attndim, + biasing=biasing, + ) + return RNNTBiasing( + encoder, + predictor, + joiner, + attndim, + biasing, + deepbiasing, + symbol_embedding_dim, + encoding_dim, + charlist, + encoding_dim, + conformer_dropout, + tcpsche, + DBaverage, + ) + + +def conformer_rnnt_biasing_base(charlist=None, biasing=True) -> RNNT: + r"""Builds basic version of Conformer RNN-T model with TCPGen. + + Returns: + RNNT: + Conformer RNN-T model with TCPGen-based biasing support. + """ + return conformer_rnnt_biasing( + input_dim=80, + encoding_dim=576, + time_reduction_stride=4, + conformer_input_dim=144, + conformer_ffn_dim=576, + conformer_num_layers=16, + conformer_num_heads=4, + conformer_depthwise_conv_kernel_size=31, + conformer_dropout=0.1, + num_symbols=601, + symbol_embedding_dim=256, + num_lstm_layers=1, + lstm_hidden_dim=320, + lstm_layer_norm=True, + lstm_layer_norm_epsilon=1e-5, + lstm_dropout=0.3, + joiner_activation="tanh", + attndim=256, + biasing=biasing, + charlist=charlist, + deepbiasing=True, + tcpsche=30, + DBaverage=False, + ) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/models/rnnt_decoder.py b/lib/python3.12/site-packages/torchaudio/prototype/models/rnnt_decoder.py new file mode 100644 index 0000000000000000000000000000000000000000..4e7c4ac661e29da7d104f80541dc0c9919b98ea0 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/models/rnnt_decoder.py @@ -0,0 +1,399 @@ +from typing import Callable, Dict, List, Optional, Tuple + +import torch +from torchaudio.models import RNNT +from torchaudio.prototype.models.rnnt import TrieNode + +__all__ = ["Hypothesis", "RNNTBeamSearchBiasing"] + + +Hypothesis = Tuple[List[int], torch.Tensor, List[List[torch.Tensor]], float, list] +Hypothesis.__doc__ = """Hypothesis generated by RNN-T beam search decoder, + represented as tuple of (tokens, prediction network output, prediction network state, score). + """ + + +def _get_hypo_tokens(hypo: Hypothesis) -> List[int]: + return hypo[0] + + +def _get_hypo_predictor_out(hypo: Hypothesis) -> torch.Tensor: + return hypo[1] + + +def _get_hypo_state(hypo: Hypothesis) -> List[List[torch.Tensor]]: + return hypo[2] + + +def _get_hypo_score(hypo: Hypothesis) -> float: + return hypo[3] + + +def _get_hypo_trie(hypo: Hypothesis) -> TrieNode: + return hypo[4] + + +def _set_hypo_trie(hypo: Hypothesis, trie: TrieNode) -> None: + hypo[4] = trie + + +def _get_hypo_key(hypo: Hypothesis) -> str: + return str(hypo[0]) + + +def _batch_state(hypos: List[Hypothesis]) -> List[List[torch.Tensor]]: + states: List[List[torch.Tensor]] = [] + for i in range(len(_get_hypo_state(hypos[0]))): + batched_state_components: List[torch.Tensor] = [] + for j in range(len(_get_hypo_state(hypos[0])[i])): + batched_state_components.append(torch.cat([_get_hypo_state(hypo)[i][j] for hypo in hypos])) + states.append(batched_state_components) + return states + + +def _slice_state(states: List[List[torch.Tensor]], idx: int, device: torch.device) -> List[List[torch.Tensor]]: + idx_tensor = torch.tensor([idx], device=device) + return [[state.index_select(0, idx_tensor) for state in state_tuple] for state_tuple in states] + + +def _default_hypo_sort_key(hypo: Hypothesis) -> float: + return _get_hypo_score(hypo) / (len(_get_hypo_tokens(hypo)) + 1) + + +def _compute_updated_scores( + hypos: List[Hypothesis], + next_token_probs: torch.Tensor, + beam_width: int, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hypo_scores = torch.tensor([_get_hypo_score(h) for h in hypos]).unsqueeze(1) + nonblank_scores = hypo_scores + next_token_probs[:, :-1] # [beam_width, num_tokens - 1] + nonblank_nbest_scores, nonblank_nbest_idx = nonblank_scores.reshape(-1).topk(beam_width) + nonblank_nbest_hypo_idx = nonblank_nbest_idx.div(nonblank_scores.shape[1], rounding_mode="trunc") + nonblank_nbest_token = nonblank_nbest_idx % nonblank_scores.shape[1] + return nonblank_nbest_scores, nonblank_nbest_hypo_idx, nonblank_nbest_token + + +def _remove_hypo(hypo: Hypothesis, hypo_list: List[Hypothesis]) -> None: + for i, elem in enumerate(hypo_list): + if _get_hypo_key(hypo) == _get_hypo_key(elem): + del hypo_list[i] + break + + +class RNNTBeamSearchBiasing(torch.nn.Module): + r"""Beam search decoder for RNN-T model with biasing support. + + Args: + model (RNNT): RNN-T model to use. + blank (int): index of blank token in vocabulary. + temperature (float, optional): temperature to apply to joint network output. + Larger values yield more uniform samples. (Default: 1.0) + hypo_sort_key (Callable[[Hypothesis], float] or None, optional): callable that computes a score + for a given hypothesis to rank hypotheses by. If ``None``, defaults to callable that returns + hypothesis score normalized by token sequence length. (Default: None) + step_max_tokens (int, optional): maximum number of tokens to emit per input time step. (Default: 100) + trie (list, optional): the prefix tree for TCPGen biasing + biasing (bool, optional): If true, do biasing, otherwise use standard RNN-T support + """ + + def __init__( + self, + model: RNNT, + blank: int, + temperature: float = 1.0, + hypo_sort_key: Optional[Callable[[Hypothesis], float]] = None, + step_max_tokens: int = 100, + trie: TrieNode = None, + biasing: bool = False, + ) -> None: + super().__init__() + self.model = model + self.blank = blank + self.temperature = temperature + self.resettrie = trie or [] + self.dobiasing = biasing + + if hypo_sort_key is None: + self.hypo_sort_key = _default_hypo_sort_key + else: + self.hypo_sort_key = hypo_sort_key + + self.step_max_tokens = step_max_tokens + + def _init_b_hypos(self, hypo: Optional[Hypothesis], device: torch.device) -> List[Hypothesis]: + if hypo is not None: + token = _get_hypo_tokens(hypo)[-1] + state = _get_hypo_state(hypo) + else: + token = self.blank + state = None + + one_tensor = torch.tensor([1], device=device) + pred_out, _, pred_state = self.model.predict(torch.tensor([[token]], device=device), one_tensor, state) + init_hypo = ([token], pred_out[0].detach(), pred_state, 0.0, self.resettrie) + return [init_hypo] + + def _get_trie_mask(self, trie): + step_mask = torch.ones(len(self.model.char_list) + 1) + step_mask[list(trie[0].keys())] = 0 + # step_mask[-1] = 0 + return step_mask + + def _get_generation_prob(self, trie): + if len(trie[0].keys()) == 0: + return True + else: + return False + + def _gen_next_token_probs( + self, enc_out: torch.Tensor, hypos: List[Hypothesis], device: torch.device + ) -> torch.Tensor: + one_tensor = torch.tensor([1], device=device) + predictor_out = torch.stack([_get_hypo_predictor_out(h) for h in hypos], dim=0) + if self.dobiasing: + # Get valid subset of wordpieces + trie_masks = torch.stack([self._get_trie_mask(_get_hypo_trie(h)) for h in hypos], dim=0) + trie_masks = trie_masks.to(enc_out.device).unsqueeze(1) # beam_width, 1, nchars + # Determine if there is any paths on the trie + genprob_masks = torch.tensor([self._get_generation_prob(_get_hypo_trie(h)) for h in hypos]) # beam_width + genprob_masks = genprob_masks.to(enc_out.device) + # Forward TCPGen component + last_tokens = torch.tensor([_get_hypo_tokens(h)[-1] for h in hypos]).unsqueeze(-1).to(enc_out.device) + hptr, tcpgen_dist = self.model.forward_tcpgen(last_tokens, trie_masks, enc_out) + else: + hptr = None + # hptr sent to joiner, if deepbiasing is True joiner will use it + joined_out, _, joined_activation = self.model.join( + enc_out, + one_tensor, + predictor_out, + torch.tensor([1] * len(hypos), device=device), + hptr=hptr, + ) # [beam_width, 1, 1, num_tokens] + if self.dobiasing: + p_gen = torch.sigmoid(self.model.pointer_gate(torch.cat((joined_activation, hptr), dim=-1))) + p_gen = p_gen.masked_fill(genprob_masks.view(p_gen.size(0), 1, 1, 1), 0) + model_tu = torch.softmax(joined_out / self.temperature, dim=3) + # assuming last token is blank + p_not_null = 1.0 - model_tu[:, :, :, -1:] + ptr_dist_fact = torch.cat([tcpgen_dist[:, :, :, :-2], tcpgen_dist[:, :, :, -1:]], dim=-1) * p_not_null + ptr_gen_complement = tcpgen_dist[:, :, :, -1:] * p_gen + p_partial = ptr_dist_fact[:, :, :, :-1] * p_gen + model_tu[:, :, :, :-1] * (1 - p_gen + ptr_gen_complement) + p_final = torch.cat([p_partial, model_tu[:, :, :, -1:]], dim=-1) + joined_out = torch.log(p_final) + else: + joined_out = torch.nn.functional.log_softmax(joined_out / self.temperature, dim=3) + return joined_out[:, 0, 0] + + def _gen_b_hypos( + self, + b_hypos: List[Hypothesis], + a_hypos: List[Hypothesis], + next_token_probs: torch.Tensor, + key_to_b_hypo: Dict[str, Hypothesis], + ) -> List[Hypothesis]: + for i in range(len(a_hypos)): + h_a = a_hypos[i] + append_blank_score = _get_hypo_score(h_a) + next_token_probs[i, -1] + if _get_hypo_key(h_a) in key_to_b_hypo: + h_b = key_to_b_hypo[_get_hypo_key(h_a)] + _remove_hypo(h_b, b_hypos) + score = float(torch.tensor(_get_hypo_score(h_b)).logaddexp(append_blank_score)) + else: + score = float(append_blank_score) + h_b = ( + _get_hypo_tokens(h_a), + _get_hypo_predictor_out(h_a), + _get_hypo_state(h_a), + score, + _get_hypo_trie(h_a), + ) + b_hypos.append(h_b) + key_to_b_hypo[_get_hypo_key(h_b)] = h_b + _, sorted_idx = torch.tensor([_get_hypo_score(hypo) for hypo in b_hypos]).sort() + return [b_hypos[idx] for idx in sorted_idx] + + def _gen_a_hypos( + self, + a_hypos: List[Hypothesis], + b_hypos: List[Hypothesis], + next_token_probs: torch.Tensor, + t: int, + beam_width: int, + device: torch.device, + ) -> List[Hypothesis]: + ( + nonblank_nbest_scores, + nonblank_nbest_hypo_idx, + nonblank_nbest_token, + ) = _compute_updated_scores(a_hypos, next_token_probs, beam_width) + + if len(b_hypos) < beam_width: + b_nbest_score = -float("inf") + else: + b_nbest_score = _get_hypo_score(b_hypos[-beam_width]) + + base_hypos: List[Hypothesis] = [] + new_tokens: List[int] = [] + new_scores: List[float] = [] + for i in range(beam_width): + score = float(nonblank_nbest_scores[i]) + if score > b_nbest_score: + a_hypo_idx = int(nonblank_nbest_hypo_idx[i]) + base_hypos.append(a_hypos[a_hypo_idx]) + new_tokens.append(int(nonblank_nbest_token[i])) + new_scores.append(score) + + if base_hypos: + new_hypos = self._gen_new_hypos(base_hypos, new_tokens, new_scores, t, device) + else: + new_hypos: List[Hypothesis] = [] + + return new_hypos + + def _gen_new_hypos( + self, + base_hypos: List[Hypothesis], + tokens: List[int], + scores: List[float], + t: int, + device: torch.device, + ) -> List[Hypothesis]: + tgt_tokens = torch.tensor([[token] for token in tokens], device=device) + states = _batch_state(base_hypos) + pred_out, _, pred_states = self.model.predict( + tgt_tokens, + torch.tensor([1] * len(base_hypos), device=device), + states, + ) + new_hypos: List[Hypothesis] = [] + for i, h_a in enumerate(base_hypos): + new_tokens = _get_hypo_tokens(h_a) + [tokens[i]] + if self.dobiasing: + new_trie = self.model.get_tcpgen_step(tokens[i], _get_hypo_trie(h_a), self.resettrie) + else: + new_trie = self.resettrie + new_hypos.append( + (new_tokens, pred_out[i].detach(), _slice_state(pred_states, i, device), scores[i], new_trie) + ) + return new_hypos + + def _search( + self, + enc_out: torch.Tensor, + hypo: Optional[Hypothesis], + beam_width: int, + ) -> List[Hypothesis]: + n_time_steps = enc_out.shape[1] + device = enc_out.device + + a_hypos: List[Hypothesis] = [] + b_hypos = self._init_b_hypos(hypo, device) + for t in range(n_time_steps): + a_hypos = b_hypos + b_hypos = torch.jit.annotate(List[Hypothesis], []) + key_to_b_hypo: Dict[str, Hypothesis] = {} + symbols_current_t = 0 + + while a_hypos: + next_token_probs = self._gen_next_token_probs(enc_out[:, t : t + 1], a_hypos, device) + next_token_probs = next_token_probs.cpu() + b_hypos = self._gen_b_hypos(b_hypos, a_hypos, next_token_probs, key_to_b_hypo) + + if symbols_current_t == self.step_max_tokens: + break + + a_hypos = self._gen_a_hypos( + a_hypos, + b_hypos, + next_token_probs, + t, + beam_width, + device, + ) + if a_hypos: + symbols_current_t += 1 + + _, sorted_idx = torch.tensor([self.hypo_sort_key(hypo) for hypo in b_hypos]).topk(beam_width) + b_hypos = [b_hypos[idx] for idx in sorted_idx] + + return b_hypos + + def forward( + self, + input: torch.Tensor, + length: torch.Tensor, + beam_width: int, + ) -> List[Hypothesis]: + r"""Performs beam search for the given input sequence. + + T: number of frames; + D: feature dimension of each frame. + + Args: + input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D). + length (torch.Tensor): number of valid frames in input + sequence, with shape () or (1,). + beam_width (int): beam size to use during search. + + Returns: + List[Hypothesis]: top-``beam_width`` hypotheses found by beam search. + """ + if input.dim() != 2 and not (input.dim() == 3 and input.shape[0] == 1): + raise ValueError("input must be of shape (T, D) or (1, T, D)") + if input.dim() == 2: + input = input.unsqueeze(0) + + if length.shape != () and length.shape != (1,): + raise ValueError("length must be of shape () or (1,)") + if input.dim() == 0: + input = input.unsqueeze(0) + + enc_out, _ = self.model.transcribe(input, length) + return self._search(enc_out, None, beam_width) + + @torch.jit.export + def infer( + self, + input: torch.Tensor, + length: torch.Tensor, + beam_width: int, + state: Optional[List[List[torch.Tensor]]] = None, + hypothesis: Optional[Hypothesis] = None, + ) -> Tuple[List[Hypothesis], List[List[torch.Tensor]]]: + r"""Performs beam search for the given input sequence in streaming mode. + + T: number of frames; + D: feature dimension of each frame. + + Args: + input (torch.Tensor): sequence of input frames, with shape (T, D) or (1, T, D). + length (torch.Tensor): number of valid frames in input + sequence, with shape () or (1,). + beam_width (int): beam size to use during search. + state (List[List[torch.Tensor]] or None, optional): list of lists of tensors + representing transcription network internal state generated in preceding + invocation. (Default: ``None``) + hypothesis (Hypothesis or None): hypothesis from preceding invocation to seed + search with. (Default: ``None``) + + Returns: + (List[Hypothesis], List[List[torch.Tensor]]): + List[Hypothesis] + top-``beam_width`` hypotheses found by beam search. + List[List[torch.Tensor]] + list of lists of tensors representing transcription network + internal state generated in current invocation. + """ + if input.dim() != 2 and not (input.dim() == 3 and input.shape[0] == 1): + raise ValueError("input must be of shape (T, D) or (1, T, D)") + if input.dim() == 2: + input = input.unsqueeze(0) + + if length.shape != () and length.shape != (1,): + raise ValueError("length must be of shape () or (1,)") + if length.dim() == 0: + length = length.unsqueeze(0) + + enc_out, _, state = self.model.transcribe_streaming(input, length, state) + return self._search(enc_out, hypothesis, beam_width), state diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..bf7bb4b86e58121241ca4ffd176cac92d4342e6a --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__init__.py @@ -0,0 +1,12 @@ +from ._vggish import VGGISH, VGGishBundle +from .hifigan_pipeline import HIFIGAN_VOCODER_V3_LJSPEECH, HiFiGANVocoderBundle +from .rnnt_pipeline import EMFORMER_RNNT_BASE_MUSTC, EMFORMER_RNNT_BASE_TEDLIUM3 + +__all__ = [ + "EMFORMER_RNNT_BASE_MUSTC", + "EMFORMER_RNNT_BASE_TEDLIUM3", + "HIFIGAN_VOCODER_V3_LJSPEECH", + "HiFiGANVocoderBundle", + "VGGISH", + "VGGishBundle", +] diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..948cfa41a7d917541568721802eff627de70643e Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/hifigan_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/hifigan_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..94d77f20a2746008e4ec3ccc877da0cbd8bc7938 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/hifigan_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..62a4c155cf23537ea5b4a81ad91c69e58c74d4b5 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/__pycache__/rnnt_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..8cd4774f56a300d099b24f3a9e905224967da522 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__init__.py @@ -0,0 +1,3 @@ +from ._vggish_pipeline import VGGISH, VGGishBundle + +__all__ = ["VGGISH", "VGGishBundle"] diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..e429665d3a8b74a7f530264b3049886a8d3147e3 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_impl.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_impl.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..25096426d0d7be2586cbf378b4e87fb106c71080 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_impl.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_pipeline.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_pipeline.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..54ec6f3fed3064b2afe5e3a45596f87defbcaa11 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/__pycache__/_vggish_pipeline.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/_vggish_impl.py b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/_vggish_impl.py new file mode 100644 index 0000000000000000000000000000000000000000..6eb6ea8f59490eab777f2ba699f128d7c7876adc --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/_vggish_impl.py @@ -0,0 +1,233 @@ +# Derived from torchvggish (https://github.com/harritaylor/torchvggish). +# Copyright 2017 The TensorFlow Authors All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +import math + +import torch + + +_MEL_BREAK_FREQUENCY_HERTZ = 700.0 +_MEL_HIGH_FREQUENCY_Q = 1127.0 + + +_SAMPLE_RATE = 16000 +_STFT_WINDOW_LENGTH_SECONDS = 0.025 +_STFT_HOP_LENGTH_SECONDS = 0.010 +_MEL_MIN_HZ = 125 +_MEL_MAX_HZ = 7500 +_NUM_BANDS = 64 +_LOG_OFFSET = 0.01 +_EXAMPLE_WINDOW_SECONDS = 0.96 # Each example contains 96 10ms frames +_EXAMPLE_HOP_SECONDS = 0.96 # with zero overlap. + + +def _build_features_network(): + layers = [] + + for input_dim, output_dim in [(1, 64), (64, 128)]: + layers += [ + torch.nn.Conv2d(input_dim, output_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), + torch.nn.ReLU(inplace=True), + torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False), + ] + + for input_dim, output_dim in [(128, 256), (256, 512)]: + layers += [ + torch.nn.Conv2d(input_dim, output_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), + torch.nn.ReLU(inplace=True), + torch.nn.Conv2d( + output_dim, + output_dim, + kernel_size=(3, 3), + stride=(1, 1), + padding=(1, 1), + ), + torch.nn.ReLU(inplace=True), + torch.nn.MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False), + ] + + return torch.nn.Sequential(*layers) + + +def _build_embedding_network(): + return torch.nn.Sequential( + torch.nn.Linear(512 * 4 * 6, 4096), + torch.nn.ReLU(True), + torch.nn.Linear(4096, 4096), + torch.nn.ReLU(True), + torch.nn.Linear(4096, 128), + torch.nn.ReLU(True), + ) + + +def _frame(data, window_length, hop_length): + num_samples = data.shape[0] + num_frames = 1 + int(math.floor((num_samples - window_length) / hop_length)) + shape = (num_frames, window_length) + data.shape[1:] + strides = (data.stride()[0] * hop_length,) + data.stride() + return torch.as_strided(data, shape, strides) + + +def _stft_magnitude(signal, fft_length, hop_length=None, window_length=None): + frames = _frame(signal, window_length, hop_length) + window = torch.hann_window(window_length, periodic=True).to(signal.device) + windowed_frames = frames * window + return torch.abs(torch.fft.rfft(windowed_frames, int(fft_length))) + + +def _hertz_to_mel(frequencies_hertz): + return _MEL_HIGH_FREQUENCY_Q * torch.log(1.0 + (frequencies_hertz / _MEL_BREAK_FREQUENCY_HERTZ)) + + +def _spectrogram_to_mel_matrix( + num_mel_bins=20, + num_spectrogram_bins=129, + audio_sample_rate=8000, + lower_edge_hertz=125.0, + upper_edge_hertz=3800.0, +): + nyquist_hertz = audio_sample_rate / 2.0 + if lower_edge_hertz < 0.0: + raise ValueError("lower_edge_hertz %.1f must be >= 0" % lower_edge_hertz) + if lower_edge_hertz >= upper_edge_hertz: + raise ValueError("lower_edge_hertz %.1f >= upper_edge_hertz %.1f" % (lower_edge_hertz, upper_edge_hertz)) + + if upper_edge_hertz > nyquist_hertz: + raise ValueError("upper_edge_hertz %.1f is greater than Nyquist %.1f" % (upper_edge_hertz, nyquist_hertz)) + spectrogram_bins_hertz = torch.linspace(0.0, nyquist_hertz, num_spectrogram_bins) + + spectrogram_bins_mel = _hertz_to_mel(spectrogram_bins_hertz) + # The i'th mel band (starting from i=1) has center frequency + # band_edges_mel[i], lower edge band_edges_mel[i-1], and higher edge + # band_edges_mel[i+1]. Thus, we need num_mel_bins + 2 values in + # the band_edges_mel arrays. + band_edges_mel = torch.linspace( + _hertz_to_mel(torch.tensor(lower_edge_hertz)), + _hertz_to_mel(torch.tensor(upper_edge_hertz)), + num_mel_bins + 2, + ) + # Matrix to post-multiply feature arrays whose rows are num_spectrogram_bins + # of spectrogram values. + mel_weights_matrix = torch.empty((num_spectrogram_bins, num_mel_bins)) + for i in range(num_mel_bins): + lower_edge_mel, center_mel, upper_edge_mel = band_edges_mel[i : i + 3] + # Calculate lower and upper slopes for every spectrogram bin. + # Line segments are linear in the *mel* domain, not hertz. + lower_slope = (spectrogram_bins_mel - lower_edge_mel) / (center_mel - lower_edge_mel) + upper_slope = (upper_edge_mel - spectrogram_bins_mel) / (upper_edge_mel - center_mel) + + # .. then intersect them with each other and zero. + mel_weights_matrix[:, i] = torch.maximum(torch.tensor(0.0), torch.minimum(lower_slope, upper_slope)) + + # HTK excludes the spectrogram DC bin; make sure it always gets a zero + # coefficient. + mel_weights_matrix[0, :] = 0.0 + return mel_weights_matrix + + +def _log_mel_spectrogram( + data, + audio_sample_rate=8000, + log_offset=0.0, + window_length_secs=0.025, + hop_length_secs=0.010, + **kwargs, +): + window_length_samples = int(round(audio_sample_rate * window_length_secs)) + hop_length_samples = int(round(audio_sample_rate * hop_length_secs)) + fft_length = 2 ** int(math.ceil(math.log(window_length_samples) / math.log(2.0))) + + spectrogram = _stft_magnitude( + data, + fft_length=fft_length, + hop_length=hop_length_samples, + window_length=window_length_samples, + ) + mel_spectrogram = torch.matmul( + spectrogram, + _spectrogram_to_mel_matrix( + num_spectrogram_bins=spectrogram.shape[1], + audio_sample_rate=audio_sample_rate, + **kwargs, + ).to(spectrogram), + ) + return torch.log(mel_spectrogram + log_offset) + + +def _waveform_to_examples(data): + # Compute log mel spectrogram features, with shape (n_frame, n_mel) + log_mel = _log_mel_spectrogram( + data, + audio_sample_rate=_SAMPLE_RATE, + log_offset=_LOG_OFFSET, + window_length_secs=_STFT_WINDOW_LENGTH_SECONDS, + hop_length_secs=_STFT_HOP_LENGTH_SECONDS, + num_mel_bins=_NUM_BANDS, + lower_edge_hertz=_MEL_MIN_HZ, + upper_edge_hertz=_MEL_MAX_HZ, + ) + + # Frame features into examples, with shape (n_example, n_frame, n_mel) + features_sample_rate = 1.0 / _STFT_HOP_LENGTH_SECONDS + example_window_length = int(round(_EXAMPLE_WINDOW_SECONDS * features_sample_rate)) + + example_hop_length = int(round(_EXAMPLE_HOP_SECONDS * features_sample_rate)) + log_mel_examples = _frame(log_mel, window_length=example_window_length, hop_length=example_hop_length) + + # (n_example, 1, n_frame, n_mel) + return log_mel_examples.unsqueeze(1) + + +class VGGish(torch.nn.Module): + """Implementation of VGGish model :cite:`45611`.""" + + def __init__(self): + super().__init__() + + self.features_network = _build_features_network() + self.embedding_network = _build_embedding_network() + + def forward(self, input: torch.Tensor) -> torch.Tensor: + """ + Args: + input (torch.Tensor): batch of spectrograms, with shape `(n_example, 1, n_frame, 64)`. + + Returns: + torch.Tensor: model output, with shape `(n_example, 128)`. + """ + x = self.features_network(input) + + x = x.permute(0, 2, 3, 1) + x = x.reshape(x.size(0), -1) + + return self.embedding_network(x) + + +class VGGishInputProcessor: + """Converts raw waveforms to batches of examples to use as inputs to VGGish.""" + + def __call__(self, input: torch.Tensor) -> torch.Tensor: + """ + Args: + input (torch.Tensor): waveform, with shape `(T,)`. + sample_rate (int): sample rate of waveform in hertz. + + Returns: + torch.Tensor: batch of examples to pass to VGGish, with shape `(n_example, 1, n_frame, 64)`. + """ + if len(input.shape) != 1: + raise ValueError("input waveform must have dimension of 1.") + return _waveform_to_examples(input) diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/_vggish_pipeline.py b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/_vggish_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..7b0f527e73eb61fe4a9fa7d7d86ea467f9ae8a9e --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/_vggish/_vggish_pipeline.py @@ -0,0 +1,82 @@ +from dataclasses import dataclass +from typing import Callable, Dict + +import torch +import torchaudio + +from ._vggish_impl import _SAMPLE_RATE, VGGish as _VGGish, VGGishInputProcessor as _VGGishInputProcessor + + +def _get_state_dict(): + path = torchaudio.utils.download_asset("models/vggish.pt") + return torch.load(path) + + +@dataclass +class VGGishBundle: + """VGGish :cite:`45611` inference pipeline ported from + `torchvggish `__ + and `tensorflow-models `__. + + Example: + >>> import torchaudio + >>> from torchaudio.prototype.pipelines import VGGISH + >>> + >>> input_sr = VGGISH.sample_rate + >>> input_proc = VGGISH.get_input_processor() + >>> model = VGGISH.get_model() + >>> + >>> waveform, sr = torchaudio.load( + >>> "Chopin_Ballade_-1_In_G_Minor,_Op._23.mp3", + >>> ) + >>> waveform = waveform.squeeze(0) + >>> waveform = torchaudio.functional.resample(waveform, sr, input_sr) + >>> mono_output = model(input_proc(waveform)) + """ + + class VGGish(_VGGish): + __doc__ = _VGGish.__doc__ + + class VGGishInputProcessor(_VGGishInputProcessor): + __doc__ = _VGGishInputProcessor.__doc__ + + _state_dict_func: Callable[[], Dict] + + @property + def sample_rate(self) -> int: + """Sample rate of input waveform expected by input processor and model. + + :type: int + """ + return _SAMPLE_RATE + + def get_model(self) -> VGGish: + """Constructs pre-trained VGGish model. Downloads and caches weights as necessary. + + Returns: + VGGish: VGGish model with pre-trained weights loaded. + """ + model = self.VGGish() + state_dict = self._state_dict_func() + model.load_state_dict(state_dict) + model.eval() + return model + + def get_input_processor(self) -> VGGishInputProcessor: + """Constructs input processor for VGGish. + + Returns: + VGGishInputProcessor: input processor for VGGish. + """ + return self.VGGishInputProcessor() + + +VGGISH = VGGishBundle(_get_state_dict) +VGGISH.__doc__ = """Pre-trained VGGish :cite:`45611` inference pipeline ported from + `torchvggish `__ + and `tensorflow-models `__. + + Per the `documentation `__ + for the original model, the model is "trained on a large YouTube dataset (a preliminary version of + what later became YouTube-8M)". + """ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/hifigan_pipeline.py b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/hifigan_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..9aa383deb010c872ef5817962396cb122281b382 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/hifigan_pipeline.py @@ -0,0 +1,228 @@ +from dataclasses import dataclass +from typing import Any, Dict, Optional + +import torch +import torch.nn.functional as F +from torch.nn import Module +from torchaudio._internal import load_state_dict_from_url + +from torchaudio.prototype.models.hifi_gan import hifigan_vocoder, HiFiGANVocoder +from torchaudio.transforms import MelSpectrogram + + +@dataclass +class HiFiGANVocoderBundle: + """Data class that bundles associated information to use pretrained + :py:class:`~torchaudio.prototype.models.HiFiGANVocoder`. + + This class provides interfaces for instantiating the pretrained model along with + the information necessary to retrieve pretrained weights and additional data + to be used with the model. + + Torchaudio library instantiates objects of this class, each of which represents + a different pretrained model. Client code should access pretrained models via these + instances. + + This bundle can convert mel spectrorgam to waveforms and vice versa. A typical use case would be a flow like + `text -> mel spectrogram -> waveform`, where one can use an external component, e.g. Tacotron2, + to generate mel spectrogram from text. Please see below for the code example. + + Example: Transform synthetic mel spectrogram to audio. + >>> import torch + >>> import torchaudio + >>> # Since HiFiGAN bundle is in prototypes, it needs to be exported explicitly + >>> from torchaudio.prototype.pipelines import HIFIGAN_VOCODER_V3_LJSPEECH as bundle + >>> + >>> # Load the HiFiGAN bundle + >>> vocoder = bundle.get_vocoder() + Downloading: "https://download.pytorch.org/torchaudio/models/hifigan_vocoder_v3_ljspeech.pth" + 100%|████████████| 5.59M/5.59M [00:00<00:00, 18.7MB/s] + >>> + >>> # Generate synthetic mel spectrogram + >>> specgram = torch.sin(0.5 * torch.arange(start=0, end=100)).expand(bundle._vocoder_params["in_channels"], 100) + >>> + >>> # Transform mel spectrogram into audio + >>> waveform = vocoder(specgram) + >>> torchaudio.save('sample.wav', waveform, bundle.sample_rate) + + Example: Usage together with Tacotron2, text to audio. + >>> import torch + >>> import torchaudio + >>> # Since HiFiGAN bundle is in prototypes, it needs to be exported explicitly + >>> from torchaudio.prototype.pipelines import HIFIGAN_VOCODER_V3_LJSPEECH as bundle_hifigan + >>> + >>> # Load Tacotron2 bundle + >>> bundle_tactron2 = torchaudio.pipelines.TACOTRON2_WAVERNN_CHAR_LJSPEECH + >>> processor = bundle_tactron2.get_text_processor() + >>> tacotron2 = bundle_tactron2.get_tacotron2() + >>> + >>> # Use Tacotron2 to convert text to mel spectrogram + >>> text = "A quick brown fox jumped over a lazy dog" + >>> input, lengths = processor(text) + >>> specgram, lengths, _ = tacotron2.infer(input, lengths) + >>> + >>> # Load HiFiGAN bundle + >>> vocoder = bundle_hifigan.get_vocoder() + Downloading: "https://download.pytorch.org/torchaudio/models/hifigan_vocoder_v3_ljspeech.pth" + 100%|████████████| 5.59M/5.59M [00:03<00:00, 1.55MB/s] + >>> + >>> # Use HiFiGAN to convert mel spectrogram to audio + >>> waveform = vocoder(specgram).squeeze(0) + >>> torchaudio.save('sample.wav', waveform, bundle_hifigan.sample_rate) + """ # noqa: E501 + + _path: str + _vocoder_params: Dict[str, Any] # Vocoder parameters + _mel_params: Dict[str, Any] # Mel transformation parameters + _sample_rate: float + + def _get_state_dict(self, dl_kwargs): + url = f"https://download.pytorch.org/torchaudio/models/{self._path}" + dl_kwargs = {} if dl_kwargs is None else dl_kwargs + state_dict = load_state_dict_from_url(url, **dl_kwargs) + return state_dict + + def get_vocoder(self, *, dl_kwargs=None) -> HiFiGANVocoder: + """Construct the HiFiGAN Generator model, which can be used a vocoder, and load the pretrained weight. + + The weight file is downloaded from the internet and cached with + :func:`torch.hub.load_state_dict_from_url` + + Args: + dl_kwargs (dictionary of keyword arguments): Passed to :func:`torch.hub.load_state_dict_from_url`. + + Returns: + Variation of :py:class:`~torchaudio.prototype.models.HiFiGANVocoder`. + """ + model = hifigan_vocoder(**self._vocoder_params) + model.load_state_dict(self._get_state_dict(dl_kwargs)) + model.eval() + return model + + def get_mel_transform(self) -> Module: + """Construct an object which transforms waveforms into mel spectrograms.""" + return _HiFiGANMelSpectrogram( + n_mels=self._vocoder_params["in_channels"], + sample_rate=self._sample_rate, + **self._mel_params, + ) + + @property + def sample_rate(self): + """Sample rate of the audio that the model is trained on. + + :type: float + """ + return self._sample_rate + + +class _HiFiGANMelSpectrogram(torch.nn.Module): + """ + Generate mel spectrogram in a way equivalent to the original HiFiGAN implementation: + https://github.com/jik876/hifi-gan/blob/4769534d45265d52a904b850da5a622601885777/meldataset.py#L49-L72 + + This class wraps around :py:class:`torchaudio.transforms.MelSpectrogram`, but performs extra steps to achive + equivalence with the HiFiGAN implementation. + + Args: + hop_size (int): Length of hop between STFT windows. + n_fft (int): Size of FFT, creates ``n_fft // 2 + 1`` bins. + win_length (int): Window size. + f_min (float or None): Minimum frequency. + f_max (float or None): Maximum frequency. + sample_rate (int): Sample rate of audio signal. + n_mels (int): Number of mel filterbanks. + """ + + def __init__( + self, + hop_size: int, + n_fft: int, + win_length: int, + f_min: Optional[float], + f_max: Optional[float], + sample_rate: float, + n_mels: int, + ): + super(_HiFiGANMelSpectrogram, self).__init__() + self.mel_transform = MelSpectrogram( + sample_rate=sample_rate, + n_fft=n_fft, + win_length=win_length, + hop_length=hop_size, + f_min=f_min, + f_max=f_max, + n_mels=n_mels, + normalized=False, + pad=0, + mel_scale="slaney", + norm="slaney", + center=False, + ) + self.sample_rate = sample_rate + self.hop_size = hop_size + self.n_fft = n_fft + self.win_length = win_length + self.f_min = f_min + self.f_max = f_max + self.n_mels = n_mels + self.pad_size = int((n_fft - hop_size) / 2) + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + """Generate mel spectrogram from a waveform. Should have same sample rate as ``self.sample_rate``. + + Args: + waveform (Tensor): waveform of shape ``(batch_size, time_length)``. + Returns: + Tensor of shape ``(batch_size, n_mel, time_length)`` + """ + ref_waveform = F.pad(waveform.unsqueeze(1), (self.pad_size, self.pad_size), mode="reflect") + ref_waveform = ref_waveform.squeeze(1) + + spectr = (self.mel_transform.spectrogram(ref_waveform) + 1e-9) ** 0.5 + mel_spectrogram = self.mel_transform.mel_scale(spectr) + mel_spectrogram = torch.log(torch.clamp(mel_spectrogram, min=1e-5)) + return mel_spectrogram + + +HIFIGAN_VOCODER_V3_LJSPEECH = HiFiGANVocoderBundle( + "hifigan_vocoder_v3_ljspeech.pth", + _vocoder_params={ + "upsample_rates": (8, 8, 4), + "upsample_kernel_sizes": (16, 16, 8), + "upsample_initial_channel": 256, + "resblock_kernel_sizes": (3, 5, 7), + "resblock_dilation_sizes": ((1, 2), (2, 6), (3, 12)), + "resblock_type": 2, + "in_channels": 80, + "lrelu_slope": 0.1, + }, + _mel_params={ + "hop_size": 256, + "n_fft": 1024, + "win_length": 1024, + "f_min": 0, + "f_max": 8000, + }, + _sample_rate=22050, +) +HIFIGAN_VOCODER_V3_LJSPEECH.__doc__ = """HiFiGAN Vocoder pipeline, trained on *The LJ Speech Dataset* + :cite:`ljspeech17`. + + This pipeine can be used with an external component which generates mel spectrograms from text, for example, + Tacotron2 - see examples in :py:class:`HiFiGANVocoderBundle`. + Although this works with the existing Tacotron2 bundles, for the best results one needs to retrain Tacotron2 + using the same data preprocessing pipeline which was used for training HiFiGAN. In particular, the original + HiFiGAN implementation uses a custom method of generating mel spectrograms from waveforms, different from + :py:class:`torchaudio.transforms.MelSpectrogram`. We reimplemented this transform as + :py:meth:`HiFiGANVocoderBundle.get_mel_transform`, making sure it is equivalent to the original HiFiGAN code `here + `_. + + The underlying vocoder is constructed by + :py:func:`torchaudio.prototype.models.hifigan_vocoder`. The weights are converted from the ones published + with the original paper :cite:`NEURIPS2020_c5d73680` under `MIT License + `__. See links to + pre-trained models on `GitHub `__. + + Please refer to :py:class:`HiFiGANVocoderBundle` for usage instructions. + """ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/pipelines/rnnt_pipeline.py b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/rnnt_pipeline.py new file mode 100644 index 0000000000000000000000000000000000000000..c82e2f83a2a4a1dc241e5b1cf15fad0690446d72 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/pipelines/rnnt_pipeline.py @@ -0,0 +1,58 @@ +from functools import partial + +from torchaudio.models import emformer_rnnt_base +from torchaudio.pipelines import RNNTBundle + + +EMFORMER_RNNT_BASE_MUSTC = RNNTBundle( + _rnnt_path="models/emformer_rnnt_base_mustc.pt", + _rnnt_factory_func=partial(emformer_rnnt_base, num_symbols=501), + _global_stats_path="pipeline-assets/global_stats_rnnt_mustc.json", + _sp_model_path="pipeline-assets/spm_bpe_500_mustc.model", + _right_padding=4, + _blank=500, + _sample_rate=16000, + _n_fft=400, + _n_mels=80, + _hop_length=160, + _segment_length=16, + _right_context_length=4, +) +EMFORMER_RNNT_BASE_MUSTC.__doc__ = """Pre-trained Emformer-RNNT-based ASR pipeline capable of performing both +streaming and non-streaming inference. + +The underlying model is constructed by :py:func:`torchaudio.models.emformer_rnnt_base` +and utilizes weights trained on *MuST-C release v2.0* :cite:`CATTONI2021101155` dataset +using training script ``train.py`` +`here `__ +with ``num_symbols=501``. + +Please refer to :py:class:`torchaudio.pipelines.RNNTBundle` for usage instructions. +""" + + +EMFORMER_RNNT_BASE_TEDLIUM3 = RNNTBundle( + _rnnt_path="models/emformer_rnnt_base_tedlium3.pt", + _rnnt_factory_func=partial(emformer_rnnt_base, num_symbols=501), + _global_stats_path="pipeline-assets/global_stats_rnnt_tedlium3.json", + _sp_model_path="pipeline-assets/spm_bpe_500_tedlium3.model", + _right_padding=4, + _blank=500, + _sample_rate=16000, + _n_fft=400, + _n_mels=80, + _hop_length=160, + _segment_length=16, + _right_context_length=4, +) +EMFORMER_RNNT_BASE_TEDLIUM3.__doc__ = """Pre-trained Emformer-RNNT-based ASR pipeline capable of performing both +streaming and non-streaming inference. + +The underlying model is constructed by :py:func:`torchaudio.models.emformer_rnnt_base` +and utilizes weights trained on *TED-LIUM Release 3* :cite:`rousseau2012tedlium` dataset +using training script ``train.py`` +`here `__ +with ``num_symbols=501``. + +Please refer to :py:class:`torchaudio.pipelines.RNNTBundle` for usage instructions. +""" diff --git a/lib/python3.12/site-packages/torchaudio/prototype/transforms/__init__.py b/lib/python3.12/site-packages/torchaudio/prototype/transforms/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..457f20e119a0640336ff91eb92ff68dd42fd23f6 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/transforms/__init__.py @@ -0,0 +1,9 @@ +from ._transforms import BarkScale, BarkSpectrogram, ChromaScale, ChromaSpectrogram, InverseBarkScale + +__all__ = [ + "BarkScale", + "BarkSpectrogram", + "ChromaScale", + "ChromaSpectrogram", + "InverseBarkScale", +] diff --git a/lib/python3.12/site-packages/torchaudio/prototype/transforms/__pycache__/__init__.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/transforms/__pycache__/__init__.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..8f84397b0f800a4e3d55da069c795b8499f03c82 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/transforms/__pycache__/__init__.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/transforms/__pycache__/_transforms.cpython-312.pyc b/lib/python3.12/site-packages/torchaudio/prototype/transforms/__pycache__/_transforms.cpython-312.pyc new file mode 100644 index 0000000000000000000000000000000000000000..894c4770d3f1fe8be9d96b9241fb641b44ad8613 Binary files /dev/null and b/lib/python3.12/site-packages/torchaudio/prototype/transforms/__pycache__/_transforms.cpython-312.pyc differ diff --git a/lib/python3.12/site-packages/torchaudio/prototype/transforms/_transforms.py b/lib/python3.12/site-packages/torchaudio/prototype/transforms/_transforms.py new file mode 100644 index 0000000000000000000000000000000000000000..9d89cc5339c84d927c5a4d91a014026a9242f675 --- /dev/null +++ b/lib/python3.12/site-packages/torchaudio/prototype/transforms/_transforms.py @@ -0,0 +1,456 @@ +from typing import Callable, Optional + +import torch +from torchaudio.prototype.functional import barkscale_fbanks, chroma_filterbank +from torchaudio.transforms import Spectrogram + + +class BarkScale(torch.nn.Module): + r"""Turn a normal STFT into a bark frequency STFT with triangular filter banks. + + .. devices:: CPU CUDA + + .. properties:: Autograd TorchScript + + Args: + n_barks (int, optional): Number of bark filterbanks. (Default: ``128``) + sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``) + f_min (float, optional): Minimum frequency. (Default: ``0.``) + f_max (float or None, optional): Maximum frequency. (Default: ``sample_rate // 2``) + n_stft (int, optional): Number of bins in STFT. See ``n_fft`` in :class:`Spectrogram`. (Default: ``201``) + norm (str or None, optional): If ``"slaney"``, divide the triangular bark weights by the width of the bark band + (area normalization). (Default: ``None``) + bark_scale (str, optional): Scale to use: ``traunmuller``, ``schroeder`` or ``wang``. (Default: ``traunmuller``) + + Example + >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True) + >>> spectrogram_transform = transforms.Spectrogram(n_fft=1024) + >>> spectrogram = spectrogram_transform(waveform) + >>> barkscale_transform = transforms.BarkScale(sample_rate=sample_rate, n_stft=1024 // 2 + 1) + >>> barkscale_spectrogram = barkscale_transform(spectrogram) + + See also: + :py:func:`torchaudio.prototype.functional.barkscale_fbanks` - The function used to + generate the filter banks. + """ + __constants__ = ["n_barks", "sample_rate", "f_min", "f_max"] + + def __init__( + self, + n_barks: int = 128, + sample_rate: int = 16000, + f_min: float = 0.0, + f_max: Optional[float] = None, + n_stft: int = 201, + bark_scale: str = "traunmuller", + ) -> None: + super(BarkScale, self).__init__() + self.n_barks = n_barks + self.sample_rate = sample_rate + self.f_max = f_max if f_max is not None else float(sample_rate // 2) + self.f_min = f_min + self.bark_scale = bark_scale + + if f_min > self.f_max: + raise ValueError("Require f_min: {} <= f_max: {}".format(f_min, self.f_max)) + + fb = barkscale_fbanks(n_stft, self.f_min, self.f_max, self.n_barks, self.sample_rate, self.bark_scale) + self.register_buffer("fb", fb) + + def forward(self, specgram: torch.Tensor) -> torch.Tensor: + r""" + Args: + specgram (torch.Tensor): A spectrogram STFT of dimension (..., freq, time). + + Returns: + torch.Tensor: Bark frequency spectrogram of size (..., ``n_barks``, time). + """ + + # (..., time, freq) dot (freq, n_mels) -> (..., n_mels, time) + bark_specgram = torch.matmul(specgram.transpose(-1, -2), self.fb).transpose(-1, -2) + + return bark_specgram + + +class InverseBarkScale(torch.nn.Module): + r"""Estimate a STFT in normal frequency domain from bark frequency domain. + + .. devices:: CPU CUDA + + It minimizes the euclidian norm between the input bark-spectrogram and the product between + the estimated spectrogram and the filter banks using SGD. + + Args: + n_stft (int): Number of bins in STFT. See ``n_fft`` in :class:`Spectrogram`. + n_barks (int, optional): Number of bark filterbanks. (Default: ``128``) + sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``) + f_min (float, optional): Minimum frequency. (Default: ``0.``) + f_max (float or None, optional): Maximum frequency. (Default: ``sample_rate // 2``) + max_iter (int, optional): Maximum number of optimization iterations. (Default: ``100000``) + tolerance_loss (float, optional): Value of loss to stop optimization at. (Default: ``1e-5``) + tolerance_change (float, optional): Difference in losses to stop optimization at. (Default: ``1e-8``) + sgdargs (dict or None, optional): Arguments for the SGD optimizer. (Default: ``None``) + bark_scale (str, optional): Scale to use: ``traunmuller``, ``schroeder`` or ``wang``. (Default: ``traunmuller``) + + Example + >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True) + >>> mel_spectrogram_transform = transforms.BarkSpectrogram(sample_rate, n_fft=1024) + >>> mel_spectrogram = bark_spectrogram_transform(waveform) + >>> inverse_barkscale_transform = transforms.InverseBarkScale(n_stft=1024 // 2 + 1) + >>> spectrogram = inverse_barkscale_transform(mel_spectrogram) + """ + __constants__ = [ + "n_stft", + "n_barks", + "sample_rate", + "f_min", + "f_max", + "max_iter", + "tolerance_loss", + "tolerance_change", + "sgdargs", + ] + + def __init__( + self, + n_stft: int, + n_barks: int = 128, + sample_rate: int = 16000, + f_min: float = 0.0, + f_max: Optional[float] = None, + max_iter: int = 100000, + tolerance_loss: float = 1e-5, + tolerance_change: float = 1e-8, + sgdargs: Optional[dict] = None, + bark_scale: str = "traunmuller", + ) -> None: + super(InverseBarkScale, self).__init__() + self.n_barks = n_barks + self.sample_rate = sample_rate + self.f_max = f_max or float(sample_rate // 2) + self.f_min = f_min + self.max_iter = max_iter + self.tolerance_loss = tolerance_loss + self.tolerance_change = tolerance_change + self.sgdargs = sgdargs or {"lr": 0.1, "momentum": 0.9} + + if f_min > self.f_max: + raise ValueError("Require f_min: {} <= f_max: {}".format(f_min, self.f_max)) + + fb = barkscale_fbanks(n_stft, self.f_min, self.f_max, self.n_barks, self.sample_rate, bark_scale) + self.register_buffer("fb", fb) + + def forward(self, barkspec: torch.Tensor) -> torch.Tensor: + r""" + Args: + barkspec (torch.Tensor): A Bark frequency spectrogram of dimension (..., ``n_barks``, time) + + Returns: + torch.Tensor: Linear scale spectrogram of size (..., freq, time) + """ + # pack batch + shape = barkspec.size() + barkspec = barkspec.view(-1, shape[-2], shape[-1]) + + n_barks, time = shape[-2], shape[-1] + freq, _ = self.fb.size() # (freq, n_mels) + barkspec = barkspec.transpose(-1, -2) + if self.n_barks != n_barks: + raise ValueError("Expected an input with {} bark bins. Found: {}".format(self.n_barks, n_barks)) + + specgram = torch.rand( + barkspec.size()[0], time, freq, requires_grad=True, dtype=barkspec.dtype, device=barkspec.device + ) + + optim = torch.optim.SGD([specgram], **self.sgdargs) + + loss = float("inf") + for _ in range(self.max_iter): + optim.zero_grad() + diff = barkspec - specgram.matmul(self.fb) + new_loss = diff.pow(2).sum(axis=-1).mean() + # take sum over bark-frequency then average over other dimensions + # so that loss threshold is applied par unit timeframe + new_loss.backward() + optim.step() + specgram.data = specgram.data.clamp(min=0) + + new_loss = new_loss.item() + if new_loss < self.tolerance_loss or abs(loss - new_loss) < self.tolerance_change: + break + loss = new_loss + + specgram.requires_grad_(False) + specgram = specgram.clamp(min=0).transpose(-1, -2) + + # unpack batch + specgram = specgram.view(shape[:-2] + (freq, time)) + return specgram + + +class BarkSpectrogram(torch.nn.Module): + r"""Create BarkSpectrogram for a raw audio signal. + + .. devices:: CPU CUDA + + .. properties:: Autograd TorchScript + + This is a composition of :py:func:`torchaudio.transforms.Spectrogram` and + and :py:func:`torchaudio.transforms.BarkScale`. + + Sources + * https://www.fon.hum.uva.nl/praat/manual/BarkSpectrogram.html + * Traunmüller, Hartmut. "Analytical Expressions for the Tonotopic Sensory Scale." Journal of the Acoustical + * Society of America. Vol. 88, Issue 1, 1990, pp. 97–100. + * https://ccrma.stanford.edu/courses/120-fall-2003/lecture-5.html + + Args: + sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``) + n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. (Default: ``400``) + win_length (int or None, optional): Window size. (Default: ``n_fft``) + hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``) + f_min (float, optional): Minimum frequency. (Default: ``0.``) + f_max (float or None, optional): Maximum frequency. (Default: ``None``) + pad (int, optional): Two sided padding of signal. (Default: ``0``) + n_mels (int, optional): Number of mel filterbanks. (Default: ``128``) + window_fn (Callable[..., torch.Tensor], optional): A function to create a window tensor + that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``) + power (float, optional): Exponent for the magnitude spectrogram, + (must be > 0) e.g., 1 for energy, 2 for power, etc. (Default: ``2``) + normalized (bool, optional): Whether to normalize by magnitude after stft. (Default: ``False``) + wkwargs (Dict[..., ...] or None, optional): Arguments for window function. (Default: ``None``) + center (bool, optional): whether to pad :attr:`waveform` on both sides so + that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`. + (Default: ``True``) + pad_mode (string, optional): controls the padding method used when + :attr:`center` is ``True``. (Default: ``"reflect"``) + bark_scale (str, optional): Scale to use: ``traunmuller``, ``schroeder`` or ``wang``. (Default: ``traunmuller``) + + Example + >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True) + >>> transform = transforms.BarkSpectrogram(sample_rate) + >>> bark_specgram = transform(waveform) # (channel, n_barks, time) + + See also: + :py:func:`torchaudio.functional.melscale_fbanks` - The function used to + generate the filter banks. + """ + __constants__ = ["sample_rate", "n_fft", "win_length", "hop_length", "pad", "n_barks", "f_min"] + + def __init__( + self, + sample_rate: int = 16000, + n_fft: int = 400, + win_length: Optional[int] = None, + hop_length: Optional[int] = None, + f_min: float = 0.0, + f_max: Optional[float] = None, + pad: int = 0, + n_barks: int = 128, + window_fn: Callable[..., torch.Tensor] = torch.hann_window, + power: float = 2.0, + normalized: bool = False, + wkwargs: Optional[dict] = None, + center: bool = True, + pad_mode: str = "reflect", + bark_scale: str = "traunmuller", + ) -> None: + super(BarkSpectrogram, self).__init__() + + self.sample_rate = sample_rate + self.n_fft = n_fft + self.win_length = win_length if win_length is not None else n_fft + self.hop_length = hop_length if hop_length is not None else self.win_length // 2 + self.pad = pad + self.power = power + self.normalized = normalized + self.n_barks = n_barks # number of bark frequency bins + self.f_max = f_max + self.f_min = f_min + self.spectrogram = Spectrogram( + n_fft=self.n_fft, + win_length=self.win_length, + hop_length=self.hop_length, + pad=self.pad, + window_fn=window_fn, + power=self.power, + normalized=self.normalized, + wkwargs=wkwargs, + center=center, + pad_mode=pad_mode, + onesided=True, + ) + self.bark_scale = BarkScale( + self.n_barks, self.sample_rate, self.f_min, self.f_max, self.n_fft // 2 + 1, bark_scale + ) + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + r""" + Args: + waveform (torch.Tensor): torch.Tensor of audio of dimension (..., time). + + Returns: + torch.Tensor: Bark frequency spectrogram of size (..., ``n_barks``, time). + """ + specgram = self.spectrogram(waveform) + bark_specgram = self.bark_scale(specgram) + return bark_specgram + + +class ChromaScale(torch.nn.Module): + r"""Converts spectrogram to chromagram. + + .. devices:: CPU CUDA + + .. properties:: Autograd + + Args: + sample_rate (int): Sample rate of audio signal. + n_freqs (int): Number of frequency bins in STFT. See ``n_fft`` in :class:`Spectrogram`. + n_chroma (int, optional): Number of chroma. (Default: ``12``) + tuning (float, optional): Tuning deviation from A440 in fractions of a chroma bin. (Default: 0.0) + ctroct (float, optional): Center of Gaussian dominance window to weight filters by, in octaves. (Default: 5.0) + octwidth (float or None, optional): Width of Gaussian dominance window to weight filters by, in octaves. + If ``None``, then disable weighting altogether. (Default: 2.0) + norm (int, optional): order of norm to normalize filter bank by. (Default: 2) + base_c (bool, optional): If True, then start filter bank at C. Otherwise, start at A. (Default: True) + + Example + >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True) + >>> spectrogram_transform = transforms.Spectrogram(n_fft=1024) + >>> spectrogram = spectrogram_transform(waveform) + >>> chroma_transform = transforms.ChromaScale(sample_rate=sample_rate, n_freqs=1024 // 2 + 1) + >>> chroma_spectrogram = chroma_transform(spectrogram) + + See also: + :py:func:`torchaudio.prototype.functional.chroma_filterbank` — function used to + generate the filter bank. + """ + + def __init__( + self, + sample_rate: int, + n_freqs: int, + *, + n_chroma: int = 12, + tuning: float = 0.0, + ctroct: float = 5.0, + octwidth: Optional[float] = 2.0, + norm: int = 2, + base_c: bool = True, + ): + super().__init__() + fb = chroma_filterbank( + sample_rate, n_freqs, n_chroma, tuning=tuning, ctroct=ctroct, octwidth=octwidth, norm=norm, base_c=base_c + ) + self.register_buffer("fb", fb) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + r""" + Args: + specgram (torch.Tensor): Spectrogram of dimension (..., ``n_freqs``, time). + + Returns: + torch.Tensor: Chroma spectrogram of size (..., ``n_chroma``, time). + """ + return torch.matmul(x.transpose(-1, -2), self.fb).transpose(-1, -2) + + +class ChromaSpectrogram(torch.nn.Module): + r"""Generates chromagram for audio signal. + + .. devices:: CPU CUDA + + .. properties:: Autograd + + Composes :py:func:`torchaudio.transforms.Spectrogram` and + and :py:func:`torchaudio.prototype.transforms.ChromaScale`. + + Args: + sample_rate (int): Sample rate of audio signal. + n_fft (int, optional): Size of FFT, creates ``n_fft // 2 + 1`` bins. + win_length (int or None, optional): Window size. (Default: ``n_fft``) + hop_length (int or None, optional): Length of hop between STFT windows. (Default: ``win_length // 2``) + pad (int, optional): Two sided padding of signal. (Default: ``0``) + window_fn (Callable[..., torch.Tensor], optional): A function to create a window tensor + that is applied/multiplied to each frame/window. (Default: ``torch.hann_window``) + power (float, optional): Exponent for the magnitude spectrogram, + (must be > 0) e.g., 1 for energy, 2 for power, etc. (Default: ``2``) + normalized (bool, optional): Whether to normalize by magnitude after stft. (Default: ``False``) + wkwargs (Dict[..., ...] or None, optional): Arguments for window function. (Default: ``None``) + center (bool, optional): whether to pad :attr:`waveform` on both sides so + that the :math:`t`-th frame is centered at time :math:`t \times \text{hop\_length}`. + (Default: ``True``) + pad_mode (string, optional): controls the padding method used when + :attr:`center` is ``True``. (Default: ``"reflect"``) + n_chroma (int, optional): Number of chroma. (Default: ``12``) + tuning (float, optional): Tuning deviation from A440 in fractions of a chroma bin. (Default: 0.0) + ctroct (float, optional): Center of Gaussian dominance window to weight filters by, in octaves. (Default: 5.0) + octwidth (float or None, optional): Width of Gaussian dominance window to weight filters by, in octaves. + If ``None``, then disable weighting altogether. (Default: 2.0) + norm (int, optional): order of norm to normalize filter bank by. (Default: 2) + base_c (bool, optional): If True, then start filter bank at C. Otherwise, start at A. (Default: True) + + Example + >>> waveform, sample_rate = torchaudio.load("test.wav", normalize=True) + >>> transform = transforms.ChromaSpectrogram(sample_rate=sample_rate, n_fft=400) + >>> chromagram = transform(waveform) # (channel, n_chroma, time) + """ + + def __init__( + self, + sample_rate: int, + n_fft: int, + *, + win_length: Optional[int] = None, + hop_length: Optional[int] = None, + pad: int = 0, + window_fn: Callable[..., torch.Tensor] = torch.hann_window, + power: float = 2.0, + normalized: bool = False, + wkwargs: Optional[dict] = None, + center: bool = True, + pad_mode: str = "reflect", + n_chroma: int = 12, + tuning: float = 0.0, + ctroct: float = 5.0, + octwidth: Optional[float] = 2.0, + norm: int = 2, + base_c: bool = True, + ): + super().__init__() + self.spectrogram = Spectrogram( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + pad=pad, + window_fn=window_fn, + power=power, + normalized=normalized, + wkwargs=wkwargs, + center=center, + pad_mode=pad_mode, + onesided=True, + ) + self.chroma_scale = ChromaScale( + sample_rate, + n_fft // 2 + 1, + n_chroma=n_chroma, + tuning=tuning, + base_c=base_c, + ctroct=ctroct, + octwidth=octwidth, + norm=norm, + ) + + def forward(self, waveform: torch.Tensor) -> torch.Tensor: + r""" + Args: + waveform (Tensor): Tensor of audio of dimension (..., time). + + Returns: + Tensor: Chromagram of size (..., ``n_chroma``, time). + """ + spectrogram = self.spectrogram(waveform) + chroma_spectrogram = self.chroma_scale(spectrogram) + return chroma_spectrogram