diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 0000000000..0f84aa1211 --- /dev/null +++ b/.dockerignore @@ -0,0 +1,14 @@ +*secret*sh +*~ +.env +.env-secrets +.git +.idea/ +__pycache__/ +_ignored_schema.rs +models/Wav2Lip/python/ +models/tacotron2/python/ +production_database_secret_connect.sh +secrets.toml +secrets.yaml +target/ diff --git a/.env b/.env new file mode 100644 index 0000000000..7545992311 --- /dev/null +++ b/.env @@ -0,0 +1,62 @@ +DATABASE_URL=mysql://storyteller:password@localhost/storyteller +#DATABASE_URL=mysql://root:root@localhost/storyteller +#DATABASE_URL=mysql://root@localhost/storyteller + +# SQLX normally requires a live database to reflect upon at build time. +# This setting turns that on and requires a `sql-data.json` cache file instead. +# See https://github.com/launchbadge/sqlx/tree/master/sqlx-cli#readme +#SQLX_OFFLINE=true + +# We can't serve HTTPS locally without self-signed certificate nonsense. +COOKIE_SECURE=false +COOKIE_DOMAIN=.jungle.horse +WEBSITE_HOMEPAGE_REDIRECT=http://jungle.horse + +# Bucket config (non-secret) +REGION_NAME=us-east1 + +TTS_DOWNLOAD_BUCKET_NAME=dev-vocodes-private-uploads +TTS_DOWNLOAD_BUCKET_ROOT=/user_uploaded_tts_models + +W2L_PRIVATE_DOWNLOAD_BUCKET_NAME=dev-vocodes-private-uploads +W2L_PUBLIC_DOWNLOAD_BUCKET_NAME=dev-vocodes-public + +PRIVATE_BUCKET_NAME=dev-vocodes-private-uploads +PUBLIC_BUCKET_NAME=dev-vocodes-public + +# Bucket roots +AUDIO_UPLOADS_BUCKET_ROOT=/user_uploaded_w2l_audio +W2L_DOWNLOAD_BUCKET_ROOT=/user_uploaded_w2l_templates +W2L_INFERENCE_OUTPUT_BUCKET_ROOT=/w2l_output + +W2L_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/Wav2Lip +W2L_INFERENCE_SCRIPT_NAME=vocodes_process_inference.py +W2L_MODEL_FILENAME=wav2lip_gan.pth + +#W2L_MODEL_CHECKPOINT=/home/bt/models/wav2lip/s3fd-619a316812.pth +W2L_MODEL_CHECKPOINT=/home/bt/models/wav2lip/wav2lip_gan.pth +W2L_END_BUMP_FILENAME=vocodes-short-end-bump.mp4 + +W2L_DOWNLOAD_SCRIPT_NAME=vocodes_process_upload.py + +TTS_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/tts +TTS_INFERENCE_SCRIPT_NAME=vocodes_inference.py + +TACOTRON_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/tts +TACOTRON_VIRTUAL_ENV_ACTIVATION_COMMAND="source python-tacotron/bin/activate" +TACOTRON_MODEL_CHECK_SCRIPT_NAME="vocodes_model_check_tacotron.py" + +TALKNET_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/tts +TALKNET_MODEL_CHECK_SCRIPT_NAME="vocodes_model_check_talknet.py" + +TTS_INFERENCE_SIDECAR_HOSTNAME=localhost:8000 + +#TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_universal_v5.pt # THIS SOUNDS HORRID +#TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_v4_uberduck.pt # THIS IS TOO SMALL! 300MB vs 1.2 GB! +TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_v4.pt + +TTS_HIFIGAN_VOCODER_MODEL_FILENAME=hifigan_universal_g_02500000.pt +TTS_HIFIGAN_SUPERRES_VOCODER_MODEL_FILENAME=hifigan_superres_twilight_33000.pt + +# Development key. Permissible to check in. +NEWRELIC_API_KEY=c402be1c4c4216ec83d0cfc0851520cfFFFFNRAL diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml new file mode 100644 index 0000000000..95e3e4a398 --- /dev/null +++ b/.github/workflows/docker-publish.yml @@ -0,0 +1,65 @@ +# Adapted from +# https://github.com/actions/starter-workflows/ci/docker-publish.yml +# If using base images, don't forget to add READ access, eg. from a package settings page like: +# https://github.com/orgs/storytold/packages/container/docker-base-images-rust-ssl/settings/actions_access +name: Publish Docker Container to GHCR +on: + push: + branches: + # NB: Default-branch doesn't work, despite Github's documentation. + #- $default-branch + - master +env: + IMAGE_NAME: storyteller-web + +jobs: + # Push image to GitHub Packages. + # See also https://docs.docker.com/docker-hub/builds/ + push: + runs-on: ubuntu-latest + if: github.event_name == 'push' + permissions: + packages: write + contents: read + + steps: + - uses: actions/checkout@v2 + + - name: Login to GitHub Container Registry + uses: docker/login-action@v1 + with: + registry: ghcr.io + #username: ${{ github.repository_owner }} + username: ${{ github.actor }} + password: ${{ secrets.GITHUB_TOKEN }} + + - name: Build image + run: docker build . --file Dockerfile --tag $IMAGE_NAME --label "runnumber=${GITHUB_RUN_ID}" + + - name: Push image to GitHub Container Registry + run: | + IMAGE_ID=ghcr.io/${{ github.repository_owner }}/$IMAGE_NAME + + # Change all uppercase to lowercase + IMAGE_ID=$(echo $IMAGE_ID | tr '[A-Z]' '[a-z]') + + # Strip git ref prefix from version + VERSION=$(echo "${{ github.ref }}" | sed -e 's,.*/\(.*\),\1,') + + # Strip "v" prefix from tag name + [[ "${{ github.ref }}" == "refs/tags/"* ]] && VERSION=$(echo $VERSION | sed -e 's/^v//') + + # Use Docker `latest` tag convention + [ "$VERSION" == "$default-branch" ] && VERSION=latest + + echo IMAGE_ID=$IMAGE_ID + echo VERSION=$VERSION + + docker tag $IMAGE_NAME $IMAGE_ID:$VERSION + docker push $IMAGE_ID:$VERSION + + SHORT_SHA=$(echo ${GITHUB_SHA} | cut -c1-12) + + docker tag $IMAGE_NAME $IMAGE_ID:$SHORT_SHA + docker push $IMAGE_ID:$SHORT_SHA + diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000000..e32e3775cb --- /dev/null +++ b/.gitignore @@ -0,0 +1,17 @@ +*.swo +*.swp +*secret*sh +*~ +.env +.env-secrets +.idea/ +__pycache__/ +_ignored_schema.rs +models/Wav2Lip/python/ +models/tacotron2/python/ +production_database_secret_connect.sh +production_database_secret_query.sh +secrets.toml +secrets.yaml +target/ +twitch_secrets.toml diff --git a/Cargo.lock b/Cargo.lock new file mode 100644 index 0000000000..cf28212e13 --- /dev/null +++ b/Cargo.lock @@ -0,0 +1,4937 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 3 + +[[package]] +name = "actix" +version = "0.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3720d0064a0ce5c0de7bd93bdb0a6caebab2a9b5668746145d7b3b0c5da02914" +dependencies = [ + "actix-rt 2.4.0", + "actix_derive", + "bitflags", + "bytes 1.0.1", + "crossbeam-channel", + "futures-core", + "futures-sink", + "futures-task", + "futures-util", + "log", + "once_cell", + "parking_lot 0.11.1", + "pin-project-lite 0.2.6", + "smallvec 1.6.1", + "tokio 1.13.0", + "tokio-util 0.6.9", +] + +[[package]] +name = "actix-codec" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78d1833b3838dbe990df0f1f87baf640cf6146e898166afe401839d1b001e570" +dependencies = [ + "bitflags", + "bytes 0.5.6", + "futures-core", + "futures-sink", + "log", + "pin-project 0.4.28", + "tokio 0.2.25", + "tokio-util 0.3.1", +] + +[[package]] +name = "actix-codec" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13895df506faee81e423febbae3a33b27fca71831b96bb3d60adf16ebcfea952" +dependencies = [ + "bitflags", + "bytes 1.0.1", + "futures-core", + "futures-sink", + "log", + "memchr", + "pin-project-lite 0.2.6", + "tokio 1.13.0", + "tokio-util 0.6.9", +] + +[[package]] +name = "actix-connect" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "177837a10863f15ba8d3ae3ec12fac1099099529ed20083a27fdfe247381d0dc" +dependencies = [ + "actix-codec 0.3.0", + "actix-rt 1.1.1", + "actix-service 1.0.6", + "actix-utils 2.0.0", + "derive_more", + "either", + "futures-util", + "http", + "log", + "trust-dns-proto", + "trust-dns-resolver", +] + +[[package]] +name = "actix-cors" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36b133d8026a9f209a9aeeeacd028e7451bcca975f592881b305d37983f303d7" +dependencies = [ + "actix-web 3.3.2", + "derive_more", + "futures-util", + "log", + "once_cell", + "tinyvec", +] + +[[package]] +name = "actix-cors" +version = "0.6.0-beta.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa59d36d7ad063401d94ada05264a303791796ad5222d0954cc21f2e1052e99b" +dependencies = [ + "actix-service 2.0.1", + "actix-web 4.0.0-beta.10", + "derive_more", + "futures-util", + "log", + "once_cell", + "smallvec 1.6.1", +] + +[[package]] +name = "actix-files" +version = "0.6.0-beta.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf3754d4632de16dd68f6a5d379cf9d7200ec26020482e26aa16560a258e92f3" +dependencies = [ + "actix-http 3.0.0-beta.11", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "actix-web 4.0.0-beta.10", + "askama_escape", + "bitflags", + "bytes 1.0.1", + "derive_more", + "futures-core", + "http-range", + "log", + "mime", + "mime_guess", + "percent-encoding 2.1.0", +] + +[[package]] +name = "actix-http" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "452299e87817ae5673910e53c243484ca38be3828db819b6011736fc6982e874" +dependencies = [ + "actix-codec 0.3.0", + "actix-connect", + "actix-rt 1.1.1", + "actix-service 1.0.6", + "actix-threadpool", + "actix-utils 2.0.0", + "base64 0.13.0", + "bitflags", + "bytes 0.5.6", + "cookie 0.14.4", + "copyless", + "derive_more", + "either", + "encoding_rs", + "futures-channel", + "futures-core", + "futures-util", + "fxhash", + "h2 0.2.7", + "http", + "httparse", + "indexmap", + "itoa", + "language-tags 0.2.2", + "lazy_static", + "log", + "mime", + "percent-encoding 2.1.0", + "pin-project 1.0.7", + "rand 0.7.3", + "regex", + "serde", + "serde_json", + "serde_urlencoded", + "sha-1", + "slab", + "time 0.2.26", +] + +[[package]] +name = "actix-http" +version = "3.0.0-beta.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9b510d35f13987537289f38bf136e7e702a5c87cc28760310cc459544f40afd" +dependencies = [ + "actix-codec 0.4.1", + "actix-rt 2.4.0", + "actix-service 2.0.1", + "actix-tls 3.0.0-beta.7", + "actix-utils 3.0.0", + "ahash 0.7.4", + "base64 0.13.0", + "bitflags", + "brotli2", + "bytes 1.0.1", + "bytestring", + "derive_more", + "encoding_rs", + "flate2", + "futures-core", + "futures-util", + "h2 0.3.7", + "http", + "httparse", + "httpdate 1.0.1", + "itoa", + "language-tags 0.3.2", + "local-channel", + "log", + "mime", + "once_cell", + "percent-encoding 2.1.0", + "pin-project 1.0.7", + "pin-project-lite 0.2.6", + "rand 0.8.4", + "sha-1", + "smallvec 1.6.1", + "tokio 1.13.0", + "zstd", +] + +[[package]] +name = "actix-macros" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ca8ce00b267af8ccebbd647de0d61e0674b6e61185cc7a592ff88772bed655" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "actix-macros" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "465a6172cf69b960917811022d8f29bc0b7fa1398bc4f78b3c466673db1213b6" +dependencies = [ + "quote", + "syn", +] + +[[package]] +name = "actix-multipart" +version = "0.4.0-beta.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee3878dadc5d0e2afa0439b4e1ef3b03fa2feb72c208d8c6f196ed6242979f6d" +dependencies = [ + "actix-utils 3.0.0", + "actix-web 4.0.0-beta.10", + "bytes 1.0.1", + "derive_more", + "futures-core", + "futures-util", + "httparse", + "local-waker", + "log", + "mime", + "twoway", +] + +[[package]] +name = "actix-router" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ad299af73649e1fc893e333ccf86f377751eb95ff875d095131574c6f43452c" +dependencies = [ + "bytestring", + "http", + "log", + "regex", + "serde", +] + +[[package]] +name = "actix-router" +version = "0.5.0-beta.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36b95ce0d76d1aa2f98b681702807475ade0f99bd4552546a6843a966d42ea3d" +dependencies = [ + "bytestring", + "firestorm", + "http", + "log", + "regex", + "serde", +] + +[[package]] +name = "actix-rt" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "143fcc2912e0d1de2bcf4e2f720d2a60c28652ab4179685a1ee159e0fb3db227" +dependencies = [ + "actix-macros 0.1.3", + "actix-threadpool", + "copyless", + "futures-channel", + "futures-util", + "smallvec 1.6.1", + "tokio 0.2.25", +] + +[[package]] +name = "actix-rt" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a0c218d0a17c120f10ee0c69c9f0c45d87319e8f66b1f065e8412b612fc3e24" +dependencies = [ + "actix-macros 0.2.3", + "futures-core", + "tokio 1.13.0", +] + +[[package]] +name = "actix-server" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45407e6e672ca24784baa667c5d32ef109ccdd8d5e0b5ebb9ef8a67f4dfb708e" +dependencies = [ + "actix-codec 0.3.0", + "actix-rt 1.1.1", + "actix-service 1.0.6", + "actix-utils 2.0.0", + "futures-channel", + "futures-util", + "log", + "mio 0.6.23", + "mio-uds", + "num_cpus", + "slab", + "socket2 0.3.19", +] + +[[package]] +name = "actix-server" +version = "2.0.0-beta.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7367665785765b066ad16e1086d26a087f696bc7c42b6f93004ced6cfcf1eeca" +dependencies = [ + "actix-rt 2.4.0", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "futures-core", + "log", + "mio 0.7.11", + "num_cpus", + "tokio 1.13.0", +] + +[[package]] +name = "actix-service" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0052435d581b5be835d11f4eb3bce417c8af18d87ddf8ace99f8e67e595882bb" +dependencies = [ + "futures-util", + "pin-project 0.4.28", +] + +[[package]] +name = "actix-service" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d3dc6a618b082974a08d7a4781d24d4691cba51500059bfebe6656a61ebfe1e" +dependencies = [ + "futures-core", + "paste", + "pin-project-lite 0.2.6", +] + +[[package]] +name = "actix-testing" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47239ca38799ab74ee6a8a94d1ce857014b2ac36f242f70f3f75a66f691e791c" +dependencies = [ + "actix-macros 0.1.3", + "actix-rt 1.1.1", + "actix-server 1.0.4", + "actix-service 1.0.6", + "log", + "socket2 0.3.19", +] + +[[package]] +name = "actix-threadpool" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d209f04d002854b9afd3743032a27b066158817965bf5d036824d19ac2cc0e30" +dependencies = [ + "derive_more", + "futures-channel", + "lazy_static", + "log", + "num_cpus", + "parking_lot 0.11.1", + "threadpool", +] + +[[package]] +name = "actix-tls" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24789b7d7361cf5503a504ebe1c10806896f61e96eca9a7350e23001aca715fb" +dependencies = [ + "actix-codec 0.3.0", + "actix-service 1.0.6", + "actix-utils 2.0.0", + "futures-util", +] + +[[package]] +name = "actix-tls" +version = "3.0.0-beta.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4af84e13e4600829858a3e68079be710d1ada461431e1e4c5ae663479ea0a3c" +dependencies = [ + "actix-codec 0.4.1", + "actix-rt 2.4.0", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "derive_more", + "futures-core", + "http", + "log", + "tokio-util 0.6.9", +] + +[[package]] +name = "actix-utils" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e9022dec56632d1d7979e59af14f0597a28a830a9c1c7fec8b2327eb9f16b5a" +dependencies = [ + "actix-codec 0.3.0", + "actix-rt 1.1.1", + "actix-service 1.0.6", + "bitflags", + "bytes 0.5.6", + "either", + "futures-channel", + "futures-sink", + "futures-util", + "log", + "pin-project 0.4.28", + "slab", +] + +[[package]] +name = "actix-utils" +version = "3.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e491cbaac2e7fc788dfff99ff48ef317e23b3cf63dbaf7aaab6418f40f92aa94" +dependencies = [ + "local-waker", + "pin-project-lite 0.2.6", +] + +[[package]] +name = "actix-web" +version = "3.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e641d4a172e7faa0862241a20ff4f1f5ab0ab7c279f00c2d4587b77483477b86" +dependencies = [ + "actix-codec 0.3.0", + "actix-http 2.2.0", + "actix-macros 0.1.3", + "actix-router 0.2.7", + "actix-rt 1.1.1", + "actix-server 1.0.4", + "actix-service 1.0.6", + "actix-testing", + "actix-threadpool", + "actix-tls 2.0.0", + "actix-utils 2.0.0", + "actix-web-codegen 0.4.0", + "awc", + "bytes 0.5.6", + "derive_more", + "encoding_rs", + "futures-channel", + "futures-core", + "futures-util", + "fxhash", + "log", + "mime", + "pin-project 1.0.7", + "regex", + "serde", + "serde_json", + "serde_urlencoded", + "socket2 0.3.19", + "time 0.2.26", + "tinyvec", + "url 2.2.2", +] + +[[package]] +name = "actix-web" +version = "4.0.0-beta.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8a4b9d00991d8da308070a5cea7f1bbaa153a91c3fb5567937d99b9f46d601e" +dependencies = [ + "actix-codec 0.4.1", + "actix-http 3.0.0-beta.11", + "actix-macros 0.2.3", + "actix-router 0.5.0-beta.2", + "actix-rt 2.4.0", + "actix-server 2.0.0-beta.6", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "actix-web-codegen 0.5.0-beta.5", + "ahash 0.7.4", + "bytes 1.0.1", + "cfg-if 1.0.0", + "cookie 0.15.1", + "derive_more", + "either", + "encoding_rs", + "futures-core", + "futures-util", + "itoa", + "language-tags 0.3.2", + "log", + "mime", + "once_cell", + "paste", + "pin-project 1.0.7", + "regex", + "serde", + "serde_json", + "serde_urlencoded", + "smallvec 1.6.1", + "socket2 0.4.0", + "time 0.3.4", + "url 2.2.2", +] + +[[package]] +name = "actix-web-actors" +version = "4.0.0-beta.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b43c8e03f1877bc010316de455ee479aa0de0ae0717d49041db5f38fa192de" +dependencies = [ + "actix", + "actix-codec 0.4.1", + "actix-http 3.0.0-beta.11", + "actix-web 4.0.0-beta.10", + "bytes 1.0.1", + "bytestring", + "futures-core", + "pin-project 1.0.7", + "tokio 1.13.0", +] + +[[package]] +name = "actix-web-codegen" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ad26f77093333e0e7c6ffe54ebe3582d908a104e448723eec6d43d08b07143fb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "actix-web-codegen" +version = "0.5.0-beta.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfe80a8828fa88a0420dc8fdd4c16b8207326c917f17701881b063eadc2a8d3b" +dependencies = [ + "actix-router 0.5.0-beta.2", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "actix_derive" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d44b8fee1ced9671ba043476deddef739dd0959bf77030b26b738cc591737a7" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "adler" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f26201604c87b1e01bd3d98f8d5d9a8fcbb815e8cedb41ffccbeb4bf593a35fe" + +[[package]] +name = "aes-soft" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be14c7498ea50828a38d0e24a765ed2effe92a705885b57d029cd67d45744072" +dependencies = [ + "cipher", + "opaque-debug", +] + +[[package]] +name = "ahash" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "739f4a8db6605981345c5654f3a85b056ce52f37a39d34da03f25bf2151ea16e" + +[[package]] +name = "ahash" +version = "0.7.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43bb833f0bf979d8475d38fbf09ed3b8a55e1885fe93ad3f93239fc6a4f17b98" +dependencies = [ + "getrandom 0.2.2", + "once_cell", + "version_check", +] + +[[package]] +name = "aho-corasick" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e37cfd5e7657ada45f742d6e99ca5788580b5c529dc78faf11ece6dc702656f" +dependencies = [ + "memchr", +] + +[[package]] +name = "aliri_braid" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4961918d91b2182cf3adc6e3cd8d6f44452e0f42db367691447af31afdd1ad7d" +dependencies = [ + "aliri_braid_impl", +] + +[[package]] +name = "aliri_braid_impl" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e21d6298c933984056b381a70c7248a97370b76dfdbc95dc3da0fcbcd994342" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "any_ascii" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70033777eb8b5124a81a1889416543dddef2de240019b674c81285a2635a7e1e" + +[[package]] +name = "anyhow" +version = "1.0.40" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28b2cd92db5cbd74e8e5028f7e27dd7aa3090e89e4f2a197cc7c8dfb69c7063b" + +[[package]] +name = "arrayvec" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "23b62fc65de8e4e7f52534fb52b0f3ed04746ae267519eef2a83941e8085068b" + +[[package]] +name = "ascii" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eab1c04a571841102f5345a8fc0f6bb3d31c315dec879b5c6e42e40ce7ffa34e" + +[[package]] +name = "askama_escape" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90c108c1a94380c89d2215d0ac54ce09796823cca0fd91b299cfff3b33e346fb" + +[[package]] +name = "async-trait" +version = "0.1.51" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44318e776df68115a881de9a8fd1b9e53368d7a4a5ce4cc48517da3393233a5e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "atoi" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "616896e05fc0e2649463a93a15183c6a16bf03413a7af88ef1285ddedfa9cda5" +dependencies = [ + "num-traits", +] + +[[package]] +name = "attohttpc" +version = "0.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e69e13a99a7e6e070bb114f7ff381e58c7ccc188630121fc4c2fe4bcf24cd072" +dependencies = [ + "http", + "log", + "native-tls", + "openssl", + "serde", + "serde_json", + "url 2.2.2", + "wildmatch", +] + +[[package]] +name = "atty" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b39be18770d11421cdb1b9947a45dd3f37e93092cbf377614828a319d5fee8" +dependencies = [ + "hermit-abi", + "libc", + "winapi 0.3.9", +] + +[[package]] +name = "autocfg" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d49d90015b3c36167a20fe2810c5cd875ad504b39cff3d4eae7977e6b7c1cb2" + +[[package]] +name = "autocfg" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdb031dd78e28731d87d56cc8ffef4a8f36ca26c38fe2de700543e627f8a464a" + +[[package]] +name = "awc" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b381e490e7b0cfc37ebc54079b0413d8093ef43d14a4e4747083f7fa47a9e691" +dependencies = [ + "actix-codec 0.3.0", + "actix-http 2.2.0", + "actix-rt 1.1.1", + "actix-service 1.0.6", + "base64 0.13.0", + "bytes 0.5.6", + "cfg-if 1.0.0", + "derive_more", + "futures-core", + "log", + "mime", + "percent-encoding 2.1.0", + "rand 0.7.3", + "serde", + "serde_json", + "serde_urlencoded", +] + +[[package]] +name = "aws-creds" +version = "0.27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460a75eac8f3cb7683e0a9a588a83c3ff039331ea7bfbfbfcecf1dacab276e11" +dependencies = [ + "anyhow", + "attohttpc", + "dirs", + "rust-ini", + "serde", + "serde-xml-rs", + "serde_derive", + "url 2.2.2", +] + +[[package]] +name = "aws-region" +version = "0.23.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1efb67b8f201dd0deea4e1240bce7da0366f8df90509a7df054973604b20b34" +dependencies = [ + "anyhow", +] + +[[package]] +name = "base-x" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4521f3e3d031370679b3b140beb36dfe4801b09ac77e30c61941f97df3ef28b" + +[[package]] +name = "base64" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3441f0f7b02788e948e47f457ca01f1d7e6d92c693bc132c22b087d3141c03ff" + +[[package]] +name = "base64" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "904dfeac50f3cdaba28fc6f57fdcddb75f49ed61346676a78c4ffe55877802fd" + +[[package]] +name = "bcrypt" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2cab630912253fb9dc92c0e2fabd0a7b51f5a5a4007177cfa31e517015b7204" +dependencies = [ + "base64 0.12.3", + "blowfish", + "byteorder", + "getrandom 0.1.16", +] + +[[package]] +name = "bitflags" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf1de2fe8c75bc145a2f577add951f8134889b4795d47466a54a5c846d691693" + +[[package]] +name = "bitvec" +version = "0.19.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8942c8d352ae1838c9dda0b0ca2ab657696ef2232a20147cf1b30ae1a9cb4321" +dependencies = [ + "funty", + "radium", + "tap", + "wyz", +] + +[[package]] +name = "block-buffer" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4152116fd6e9dadb291ae18fc1ec3575ed6d84c29642d97890f4b4a3417297e4" +dependencies = [ + "block-padding", + "generic-array", +] + +[[package]] +name = "block-cipher" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f337a3e6da609650eb74e02bc9fac7b735049f7623ab12f2e4c719316fcc7e80" +dependencies = [ + "generic-array", +] + +[[package]] +name = "block-modes" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57a0e8073e8baa88212fb5823574c02ebccb395136ba9a164ab89379ec6072f0" +dependencies = [ + "block-padding", + "cipher", +] + +[[package]] +name = "block-padding" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d696c370c750c948ada61c69a0ee2cbbb9c50b1019ddb86d9317157a99c2cae" + +[[package]] +name = "blowfish" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0f06850ba969bc59388b2cc0a4f186fc6d9d37208863b15b84ae3866ac90ac06" +dependencies = [ + "block-cipher", + "byteorder", + "opaque-debug", +] + +[[package]] +name = "brotli-sys" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4445dea95f4c2b41cde57cc9fee236ae4dbae88d8fcbdb4750fc1bb5d86aaecd" +dependencies = [ + "cc", + "libc", +] + +[[package]] +name = "brotli2" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cb036c3eade309815c15ddbacec5b22c4d1f3983a774ab2eac2e3e9ea85568e" +dependencies = [ + "brotli-sys", + "libc", +] + +[[package]] +name = "bumpalo" +version = "3.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63396b8a4b9de3f4fdfb320ab6080762242f66a8ef174c49d8e19b674db4cdbe" + +[[package]] +name = "byteorder" +version = "1.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14c189c53d098945499cdfa7ecc63567cf3886b3332b312a5b4585d8d3a6a610" + +[[package]] +name = "bytes" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "206fdffcfa2df7cbe15601ef46c813fce0965eb3286db6b56c583b814b51c81c" +dependencies = [ + "byteorder", + "iovec", +] + +[[package]] +name = "bytes" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e4cec68f03f32e44924783795810fa50a7035d8c8ebe78580ad7e6c703fba38" + +[[package]] +name = "bytes" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b700ce4376041dcd0a327fd0097c41095743c4c8af8887265942faf1100bd040" + +[[package]] +name = "bytestring" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90706ba19e97b90786e19dc0d5e2abd80008d99d4c0c5d1ad0b5e72cec7c494d" +dependencies = [ + "bytes 1.0.1", +] + +[[package]] +name = "cc" +version = "1.0.67" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3c69b077ad434294d3ce9f1f6143a2a4b89a8a2d54ef813d85003a4fd1137fd" +dependencies = [ + "jobserver", +] + +[[package]] +name = "cfb" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca453e8624711b2f0f4eb47076a318feda166252a827ee25d067b43de83dcba0" +dependencies = [ + "byteorder", + "uuid", +] + +[[package]] +name = "cfg-if" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4785bdd1c96b2a846b2bd7cc02e86b6b3dbf14e7e53446c4f54c92a361040822" + +[[package]] +name = "cfg-if" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd" + +[[package]] +name = "chrono" +version = "0.4.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670ad68c9088c2a963aaa298cb369688cf3f9465ce5e2d4ca10e6e0098a1ce73" +dependencies = [ + "libc", + "num-integer", + "num-traits", + "serde", + "time 0.1.44", + "winapi 0.3.9", +] + +[[package]] +name = "cipher" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "12f8e7987cbd042a63249497f41aed09f8e65add917ea6566effbc56578d6801" +dependencies = [ + "generic-array", +] + +[[package]] +name = "cloudabi" +version = "0.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddfc5b9aa5d4507acaf872de71051dfd0e309860e88966e1051e462a077aac4f" +dependencies = [ + "bitflags", +] + +[[package]] +name = "combine" +version = "3.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da3da6baa321ec19e1cc41d31bf599f00c783d0517095cdaf0332e3fe8d20680" +dependencies = [ + "ascii", + "byteorder", + "either", + "memchr", + "unreachable", +] + +[[package]] +name = "combine" +version = "4.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc4369b5e4c0cddf64ad8981c0111e7df4f7078f4d6ba98fb31f2e17c4c57b7e" +dependencies = [ + "bytes 1.0.1", + "memchr", +] + +[[package]] +name = "config" +version = "0.0.1" + +[[package]] +name = "const_fn" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "402da840495de3f976eaefc3485b7f5eb5b0bf9761f9a47be27fe975b3b8c2ec" + +[[package]] +name = "container_common" +version = "0.0.1" +dependencies = [ + "anyhow", + "once_cell", + "rand 0.8.4", + "uuid", +] + +[[package]] +name = "convert_case" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6245d59a3e82a7fc217c5828a6692dbc6dfb63a0c8c90495621f7b9d79704a0e" + +[[package]] +name = "cookie" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03a5d7b21829bc7b4bf4754a978a241ae54ea55a40f92bb20216e54096f4b951" +dependencies = [ + "percent-encoding 2.1.0", + "time 0.2.26", + "version_check", +] + +[[package]] +name = "cookie" +version = "0.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d5f1c7727e460397e56abc4bddc1d49e07a1ad78fc98eb2e1c8f032a58a2f80d" +dependencies = [ + "percent-encoding 2.1.0", + "time 0.2.26", + "version_check", +] + +[[package]] +name = "copyless" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2df960f5d869b2dd8532793fde43eb5427cceb126c929747a26823ab0eeb536" + +[[package]] +name = "core-foundation" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a89e2ae426ea83155dccf10c0fa6b1463ef6d5fcb44cee0b224a408fa640a62" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea221b5284a47e40033bf9b66f35f984ec0ea2931eb03505246cd27a963f981b" + +[[package]] +name = "cpufeatures" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec1028182c380cc45a2e2c5ec841134f2dfd0f8f5f0a5bcd68004f81b5efdf4" +dependencies = [ + "libc", +] + +[[package]] +name = "crc" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "49fc9a695bca7f35f5f4c15cddc84415f66a74ea78eef08e90c5024f2b540e23" +dependencies = [ + "crc-catalog", +] + +[[package]] +name = "crc-any" +version = "2.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9950e91c5c444b0729f5f1b0aec76c523e01920ce828e37bccfa27803ff34e1" +dependencies = [ + "debug-helper", +] + +[[package]] +name = "crc-catalog" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccaeedb56da03b09f598226e25e80088cb4cd25f316e6e4df7d695f0feeb1403" + +[[package]] +name = "crc32fast" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81156fece84ab6a9f2afdb109ce3ae577e42b1228441eded99bd77f627953b1a" +dependencies = [ + "cfg-if 1.0.0", +] + +[[package]] +name = "crossbeam" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ae5588f6b3c3cb05239e90bd110f257254aecd01e4635400391aeae07497845" +dependencies = [ + "cfg-if 1.0.0", + "crossbeam-channel", + "crossbeam-deque", + "crossbeam-epoch", + "crossbeam-queue", + "crossbeam-utils 0.8.5", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06ed27e177f16d65f0f0c22a213e17c696ace5dd64b14258b52f9417ccb52db4" +dependencies = [ + "cfg-if 1.0.0", + "crossbeam-utils 0.8.5", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6455c0ca19f0d2fbf751b908d5c55c1f5cbc65e03c4225427254b46890bdde1e" +dependencies = [ + "cfg-if 1.0.0", + "crossbeam-epoch", + "crossbeam-utils 0.8.5", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec02e091aa634e2c3ada4a392989e7c3116673ef0ac5b72232439094d73b7fd" +dependencies = [ + "cfg-if 1.0.0", + "crossbeam-utils 0.8.5", + "lazy_static", + "memoffset", + "scopeguard", +] + +[[package]] +name = "crossbeam-queue" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b10ddc024425c88c2ad148c1b0fd53f4c6d38db9697c9f1588381212fa657c9" +dependencies = [ + "cfg-if 1.0.0", + "crossbeam-utils 0.8.5", +] + +[[package]] +name = "crossbeam-utils" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3c7c73a2d1e9fc0886a08b93e98eb643461230d5f1925e4036204d5f2e261a8" +dependencies = [ + "autocfg 1.0.1", + "cfg-if 0.1.10", + "lazy_static", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d82cfc11ce7f2c3faef78d8a684447b40d503d9681acebed6cb728d45940c4db" +dependencies = [ + "cfg-if 1.0.0", + "lazy_static", +] + +[[package]] +name = "crypto-mac" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4857fd85a0c34b3c3297875b747c1e02e06b6a0ea32dd892d8192b9ce0813ea6" +dependencies = [ + "generic-array", + "subtle", +] + +[[package]] +name = "crypto-mac" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b1d1a86f49236c215f271d40892d5fc950490551400b02ef360692c29815c714" +dependencies = [ + "generic-array", + "subtle", +] + +[[package]] +name = "data-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3ee2393c4a91429dffb4bedf19f4d6abf27d8a732c8ce4980305d782e5426d57" + +[[package]] +name = "database_queries" +version = "0.0.1" +dependencies = [ + "anyhow", + "base64 0.13.0", + "chrono", + "config", + "container_common", + "dotenv", + "log", + "once_cell", + "regex", + "serde", + "serde_derive", + "serde_json", + "sqlx", + "toml", +] + +[[package]] +name = "debug-helper" +version = "0.3.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4460596867846f73bddca51f7403b6a29f5315125be10a1640259b4db5b9494c" + +[[package]] +name = "derive_more" +version = "0.99.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5cc7b9cef1e351660e5443924e4f43ab25fbbed3e9a5f052df3677deb4d6b320" +dependencies = [ + "convert_case", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "des" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b24e7c748888aa2fa8bce21d8c64a52efc810663285315ac7476f7197a982fae" +dependencies = [ + "byteorder", + "cipher", + "opaque-debug", +] + +[[package]] +name = "diesel" +version = "1.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "047bfc4d5c3bd2ef6ca6f981941046113524b9a9f9a7cbdfdd7ff40f58e6f542" +dependencies = [ + "byteorder", + "diesel_derives", + "mysqlclient-sys", + "url 1.7.2", +] + +[[package]] +name = "diesel_derives" +version = "1.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45f5098f628d02a7a0f68ddba586fb61e80edec3bdc1be3b921f4ceec60858d3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "digest" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3dd60d1080a57a05ab032377049e0591415d2b31afd7028356dbf3cc6dcb066" +dependencies = [ + "generic-array", +] + +[[package]] +name = "dirs" +version = "4.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3aa72a6f96ea37bbc5aa912f6788242832f75369bdfdadcb0e38423f100059" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03d86534ed367a67548dc68113a0f5db55432fdfbb6e6f9d77704397d95d5780" +dependencies = [ + "libc", + "redox_users", + "winapi 0.3.9", +] + +[[package]] +name = "discard" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "212d0f5754cb6769937f4501cc0e67f4f4483c8d2c3e1e922ee9edbe4ab4c7c0" + +[[package]] +name = "displaydoc" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3bf95dc3f046b9da4f2d51833c0d3547d8564ef6910f5c1ed130306a75b92886" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dlv-list" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68df3f2b690c1b86e65ef7830956aededf3cb0a16f898f79b9a6f421a7b6211b" +dependencies = [ + "rand 0.8.4", +] + +[[package]] +name = "dotenv" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77c90badedccf4105eca100756a0b1289e191f6fcbdadd3cee1d2f614f97da8f" + +[[package]] +name = "dtoa" +version = "0.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "56899898ce76aaf4a0f24d914c97ea6ed976d42fec6ad33fcbb0a1103e07b2b0" + +[[package]] +name = "easyenv" +version = "0.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8044ea5153408cb704ee2a9dafbd33aa5d14c8a34c8405910d90e8c8ef45944" +dependencies = [ + "dotenv", + "env_logger", + "log", +] + +[[package]] +name = "either" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e78d4f1cc4ae33bbfc157ed5d5a5ef3bc29227303d595861deb238fcec4e9457" +dependencies = [ + "serde", +] + +[[package]] +name = "encoding_rs" +version = "0.8.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "80df024fbc5ac80f87dfef0d9f5209a252f2a497f7f42944cff24d8253cac065" +dependencies = [ + "cfg-if 1.0.0", +] + +[[package]] +name = "enum-as-inner" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c5f0096a91d210159eceb2ff5e1c4da18388a170e1e3ce948aac9c8fdbbf595" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "env_logger" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44533bbbb3bb3c1fa17d9f2e4e38bbbaf8396ba82193c4cb1b6445d711445d36" +dependencies = [ + "atty", + "humantime", + "log", + "regex", + "termcolor", +] + +[[package]] +name = "firestorm" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31586bda1b136406162e381a3185a506cdfc1631708dd40cba2f6628d8634499" + +[[package]] +name = "flate2" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cd3aec53de10fe96d7d8c565eb17f2c687bb5518a2ec453b5b1252964526abe0" +dependencies = [ + "cfg-if 1.0.0", + "crc32fast", + "libc", + "miniz_oxide", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + +[[package]] +name = "form_urlencoded" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fc25a87fa4fd2094bffb06925852034d90a17f0d1e05197d4956d3555752191" +dependencies = [ + "matches", + "percent-encoding 2.1.0", +] + +[[package]] +name = "fuchsia-cprng" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a06f77d526c1a601b7c4cdd98f54b5eaabffc14d5f2f0296febdc7f357c6d3ba" + +[[package]] +name = "fuchsia-zircon" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e9763c69ebaae630ba35f74888db465e49e259ba1bc0eda7d06f4a067615d82" +dependencies = [ + "bitflags", + "fuchsia-zircon-sys", +] + +[[package]] +name = "fuchsia-zircon-sys" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3dcaa9ae7725d12cdb85b3ad99a434db70b468c09ded17e012d86b5c1010f7a7" + +[[package]] +name = "funty" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fed34cd105917e91daa4da6b3728c47b068749d6a62c59811f06ed2ac71d9da7" + +[[package]] +name = "futures" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a471a38ef8ed83cd6e40aa59c1ffe17db6855c18e3604d9c4ed8c08ebc28678" + +[[package]] +name = "futures" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a12aa0eb539080d55c3f2d45a67c3b58b6b0773c1a3ca2dfec66d58c97fd66ca" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5da6ba8c3bb3c165d3c7319fc1cc8304facf1fb8db99c5de877183c08a273888" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88d1c26957f23603395cd326b0ffe64124b818f4449552f960d815cfba83a53d" + +[[package]] +name = "futures-executor" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45025be030969d763025784f7f355043dc6bc74093e4ecc5000ca4dc50d8745c" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", + "num_cpus", +] + +[[package]] +name = "futures-intrusive" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "62007592ac46aa7c2b6416f7deb9a8a8f63a01e0f1d6e1787d5630170db2b63e" +dependencies = [ + "futures-core", + "lock_api 0.4.4", + "parking_lot 0.11.1", +] + +[[package]] +name = "futures-io" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "522de2a0fe3e380f1bc577ba0474108faf3f6b18321dbf60b3b9c39a75073377" + +[[package]] +name = "futures-macro" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18e4a4b95cea4b4ccbcf1c5675ca7c4ee4e9e75eb79944d07defde18068f79bb" +dependencies = [ + "autocfg 1.0.1", + "proc-macro-hack", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36ea153c13024fe480590b3e3d4cad89a0cfacecc24577b68f86c6ced9c2bc11" + +[[package]] +name = "futures-task" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d3d00f4eddb73e498a54394f228cd55853bdf059259e8e7bc6e69d408892e99" + +[[package]] +name = "futures-timer" +version = "3.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e64b03909df88034c26dc1547e8970b91f98bdb65165d6a4e9110d94263dbb2c" + +[[package]] +name = "futures-util" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "36568465210a3a6ee45e1f165136d68671471a501e632e9a98d96872222b5481" +dependencies = [ + "autocfg 1.0.1", + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite 0.2.6", + "pin-utils", + "proc-macro-hack", + "proc-macro-nested", + "slab", +] + +[[package]] +name = "fxhash" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c31b6d751ae2c7f11320402d34e41349dd1016f8d5d45e48c4312bc8625af50c" +dependencies = [ + "byteorder", +] + +[[package]] +name = "generic-array" +version = "0.14.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "501466ecc8a30d1d3b7fc9229b122b2ce8ed6e9d9223f1138d4babb253e51817" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getopts" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "14dbbfd5c71d70241ecf9e6f13737f7b5ce823821063188d7e46c41d371eebd5" +dependencies = [ + "unicode-width", +] + +[[package]] +name = "getrandom" +version = "0.1.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fc3cb4d91f53b50155bdcfd23f6a4c39ae1969c2ae85982b135750cccaf5fce" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "wasi 0.9.0+wasi-snapshot-preview1", +] + +[[package]] +name = "getrandom" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9495705279e7140bf035dde1f6e750c162df8b625267cd52cc44e0b156732c8" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "wasi 0.10.0+wasi-snapshot-preview1", +] + +[[package]] +name = "h2" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e4728fd124914ad25e99e3d15a9361a879f6620f63cb56bbb08f95abb97a535" +dependencies = [ + "bytes 0.5.6", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http", + "indexmap", + "slab", + "tokio 0.2.25", + "tokio-util 0.3.1", + "tracing", + "tracing-futures", +] + +[[package]] +name = "h2" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fd819562fcebdac5afc5c113c3ec36f902840b70fd4fc458799c8ce4607ae55" +dependencies = [ + "bytes 1.0.1", + "fnv", + "futures-core", + "futures-sink", + "futures-util", + "http", + "indexmap", + "slab", + "tokio 1.13.0", + "tokio-util 0.6.9", + "tracing", +] + +[[package]] +name = "hashbrown" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7afe4a420e3fe79967a00898cc1f4db7c8a49a9333a29f8a4bd76a253d5cd04" +dependencies = [ + "ahash 0.4.7", +] + +[[package]] +name = "hashbrown" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab5ef0d4909ef3724cc8cce6ccc8572c5c817592e9285f5464f8e86f8bd3726e" +dependencies = [ + "ahash 0.7.4", +] + +[[package]] +name = "hashlink" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7249a3129cbc1ffccd74857f81464a323a152173cdb134e0fd81bc803b29facf" +dependencies = [ + "hashbrown 0.11.2", +] + +[[package]] +name = "heck" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "87cbf45460356b7deeb5e3415b5563308c0a9b057c85e12b06ad551f98d0a6ac" +dependencies = [ + "unicode-segmentation", +] + +[[package]] +name = "hermit-abi" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "322f4de77956e22ed0e5032c359a0f1273f1f7f0d79bfa3b8ffbc730d7fbcc5c" +dependencies = [ + "libc", +] + +[[package]] +name = "hex" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" + +[[package]] +name = "hmac" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c1441c6b1e930e2817404b5046f1f989899143a12bf92de603b69f4e0aee1e15" +dependencies = [ + "crypto-mac 0.10.0", + "digest", +] + +[[package]] +name = "hmac" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a2a2320eb7ec0ebe8da8f744d7812d9fc4cb4d09344ac01898dbcb6a20ae69b" +dependencies = [ + "crypto-mac 0.11.1", + "digest", +] + +[[package]] +name = "hostname" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867" +dependencies = [ + "libc", + "match_cfg", + "winapi 0.3.9", +] + +[[package]] +name = "http" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1323096b05d41827dadeaee54c9981958c0f94e670bc94ed80037d1a7b8b186b" +dependencies = [ + "bytes 1.0.1", + "fnv", + "itoa", +] + +[[package]] +name = "http-body" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13d5ff830006f7646652e057693569bfe0d51760c0085a071769d142a205111b" +dependencies = [ + "bytes 0.5.6", + "http", +] + +[[package]] +name = "http-body" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ff4f84919677303da5f147645dbea6b1881f368d03ac84e1dc09031ebd7b2c6" +dependencies = [ + "bytes 1.0.1", + "http", + "pin-project-lite 0.2.6", +] + +[[package]] +name = "http-range" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eee9694f83d9b7c09682fdb32213682939507884e5bcf227be9aff5d644b90dc" + +[[package]] +name = "http_server_common" +version = "0.0.1" +dependencies = [ + "actix-cors 0.6.0-beta.3", + "actix-http 3.0.0-beta.11", + "actix-web 4.0.0-beta.10", + "anyhow", + "base64 0.13.0", + "chrono", + "container_common", + "dotenv", + "language-tags 0.3.2", + "log", + "serde", + "serde_derive", + "serde_json", + "toml", +] + +[[package]] +name = "httparse" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "acd94fdbe1d4ff688b67b04eee2e17bd50995534a61539e45adfefb45e5e5503" + +[[package]] +name = "httpdate" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "494b4d60369511e7dea41cf646832512a94e542f68bb9c49e54518e0f468eb47" + +[[package]] +name = "httpdate" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6456b8a6c8f33fee7d958fcd1b60d55b11940a79e63ae87013e6d22e26034440" + +[[package]] +name = "humantime" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df004cfca50ef23c36850aaaa59ad52cc70d0e90243c3c7737a4dd32dc7a3c4f" +dependencies = [ + "quick-error", +] + +[[package]] +name = "hyper" +version = "0.13.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a6f157065790a3ed2f88679250419b5cdd96e714a0d65f7797fd337186e96bb" +dependencies = [ + "bytes 0.5.6", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.2.7", + "http", + "http-body 0.3.1", + "httparse", + "httpdate 0.3.2", + "itoa", + "pin-project 1.0.7", + "socket2 0.3.19", + "tokio 0.2.25", + "tower-service", + "tracing", + "want", +] + +[[package]] +name = "hyper" +version = "0.14.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b91bb1f221b6ea1f1e4371216b70f40748774c2fb5971b450c07773fb92d26b" +dependencies = [ + "bytes 1.0.1", + "futures-channel", + "futures-core", + "futures-util", + "h2 0.3.7", + "http", + "http-body 0.4.4", + "httparse", + "httpdate 1.0.1", + "itoa", + "pin-project-lite 0.2.6", + "socket2 0.4.0", + "tokio 1.13.0", + "tower-service", + "tracing", + "want", +] + +[[package]] +name = "hyper-tls" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6183ddfa99b85da61a140bea0efc93fdf56ceaa041b37d553518030827f9905" +dependencies = [ + "bytes 1.0.1", + "hyper 0.14.14", + "native-tls", + "tokio 1.13.0", + "tokio-native-tls", +] + +[[package]] +name = "idna" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38f09e0f0b1fb55fdee1f17470ad800da77af5186a1a76c026b679358b7e844e" +dependencies = [ + "matches", + "unicode-bidi", + "unicode-normalization", +] + +[[package]] +name = "idna" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "418a0a6fab821475f634efe3ccc45c013f742efe03d853e8d3355d5cb850ecf8" +dependencies = [ + "matches", + "unicode-bidi", + "unicode-normalization", +] + +[[package]] +name = "indexmap" +version = "1.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824845a0bf897a9042383849b02c1bc219c2383772efcd5c6f9766fa4b81aef3" +dependencies = [ + "autocfg 1.0.1", + "hashbrown 0.9.1", +] + +[[package]] +name = "infer" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea70330449622910e0edebab230734569516269fb32342fb0a8956340fa48c6c" +dependencies = [ + "cfb", +] + +[[package]] +name = "instant" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "61124eeebbd69b8190558df225adf7e4caafce0d743919e5d6b19652314ec5ec" +dependencies = [ + "cfg-if 1.0.0", +] + +[[package]] +name = "iovec" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2b3ea6ff95e175473f8ffe6a7eb7c00d054240321b84c57051175fe3c1e075e" +dependencies = [ + "libc", +] + +[[package]] +name = "ipconfig" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7e2f18aece9709094573a9f24f483c4f65caa4298e2f7ae1b71cc65d853fad7" +dependencies = [ + "socket2 0.3.19", + "widestring", + "winapi 0.3.9", + "winreg 0.6.2", +] + +[[package]] +name = "ipnet" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47be2f14c678be2fdcab04ab1171db51b2762ce6f0a8ee87c8dd4a04ed216135" + +[[package]] +name = "itoa" +version = "0.4.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd25036021b0de88a0aff6b850051563c6516d0bf53f8638938edbb9de732736" + +[[package]] +name = "jobserver" +version = "0.1.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "af25a77299a7f711a01975c35a6a424eb6862092cc2d6c72c4ed6cbc56dfc1fa" +dependencies = [ + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.51" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83bdfbace3a0e81a4253f73b49e960b053e396a11012cbd49b9b74d6a2b67062" +dependencies = [ + "wasm-bindgen", +] + +[[package]] +name = "jwt" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "86e46349d67dc03bdbdb28da0337a355a53ca1d5156452722c36fe21d0e6389b" +dependencies = [ + "base64 0.13.0", + "crypto-mac 0.10.0", + "digest", + "hmac 0.10.1", + "serde", + "serde_json", + "sha2", +] + +[[package]] +name = "kernel32-sys" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7507624b29483431c0ba2d82aece8ca6cdba9382bff4ddd0f7490560c056098d" +dependencies = [ + "winapi 0.2.8", + "winapi-build", +] + +[[package]] +name = "language-tags" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a91d884b6667cd606bb5a69aa0c99ba811a115fc68915e7056ec08a46e93199a" + +[[package]] +name = "language-tags" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4345964bb142484797b161f473a503a434de77149dd8c7427788c6e13379388" + +[[package]] +name = "lazy_static" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e2abad23fbc42b3700f2f279844dc832adb2b2eb069b2df918f455c4e18cc646" +dependencies = [ + "spin", +] + +[[package]] +name = "lexical-core" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6607c62aa161d23d17a9072cc5da0be67cdfc89d3afb1e8d9c842bebc2525ffe" +dependencies = [ + "arrayvec", + "bitflags", + "cfg-if 1.0.0", + "ryu", + "static_assertions", +] + +[[package]] +name = "lexical-sort" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c09e4591611e231daf4d4c685a66cb0410cc1e502027a20ae55f2bb9e997207a" +dependencies = [ + "any_ascii", +] + +[[package]] +name = "lfu" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "256290473df5a92218b86e619ab5fe8a831a167b70af720726b565c68b5360d3" +dependencies = [ + "linked_hash_set", +] + +[[package]] +name = "libc" +version = "0.2.107" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fbe5e23404da5b4f555ef85ebed98fb4083e55a00c317800bc2a50ede9f3d219" + +[[package]] +name = "libm" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7d73b3f436185384286bd8098d17ec07c9a7d2388a6599f824d8502b529702a" + +[[package]] +name = "limitation" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e415413963930f7050f21994bfd336dc75ef8b5b3b142e99c55b72b61f726d5" +dependencies = [ + "chrono", + "futures 0.1.31", + "redis 0.13.0", + "time 0.1.44", +] + +[[package]] +name = "linked-hash-map" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7fb9b38af92608140b86b693604b9ffcc5824240a484d1ecd4795bacb2fe88f3" + +[[package]] +name = "linked_hash_set" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47186c6da4d81ca383c7c47c1bfc80f4b95f4720514d860a5407aaf4233f9588" +dependencies = [ + "linked-hash-map", +] + +[[package]] +name = "local-channel" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6246c68cf195087205a0512559c97e15eaf95198bf0e206d662092cdcb03fe9f" +dependencies = [ + "futures-core", + "futures-sink", + "futures-util", + "local-waker", +] + +[[package]] +name = "local-waker" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "84f9a2d3e27ce99ce2c3aad0b09b1a7b916293ea9b2bf624c13fe646fadd8da4" + +[[package]] +name = "lock_api" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4da24a77a3d8a6d4862d95f72e6fdb9c09a643ecdb402d754004a557f2bec75" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "lock_api" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0382880606dff6d15c9476c416d18690b72742aa7b605bb6dd6ec9030fbf07eb" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51b9bbe6c47d51fc3e1a9b945965946b4c44142ab8792c50835a980d362c2710" +dependencies = [ + "cfg-if 1.0.0", +] + +[[package]] +name = "lru-cache" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31e24f1ad8321ca0e8a1e0ac13f23cb668e6f5466c2c57319f6a5cf1cc8e3b1c" +dependencies = [ + "linked-hash-map", +] + +[[package]] +name = "lru_time_cache" +version = "0.11.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9106e1d747ffd48e6be5bb2d97fa706ed25b144fbee4d5c02eae110cd8d6badd" + +[[package]] +name = "magic-crypt" +version = "3.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3c94f1281833c690f81e6a00c545063b4f034509d3af6d29b58d48e39aa64c9" +dependencies = [ + "aes-soft", + "base64 0.13.0", + "block-modes", + "crc-any", + "des", + "digest", + "md-5", + "sha2", + "tiger", +] + +[[package]] +name = "maplit" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e2e65a1a2e43cfcb47a895c4c8b10d1f4a61097f9f254f183aee60cad9c651d" + +[[package]] +name = "match_cfg" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4" + +[[package]] +name = "matches" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ffc5c5338469d4d3ea17d269fa8ea3512ad247247c30bd2df69e68309ed0a08" + +[[package]] +name = "maybe-async" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6007f9dad048e0a224f27ca599d669fca8cfa0dac804725aab542b2eb032bce6" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "maybe-uninit" +version = "2.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "60302e4db3a61da70c0cb7991976248362f30319e88850c487b9b95bbf059e00" + +[[package]] +name = "md-5" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5a279bb9607f9f53c22d496eade00d138d1bdcccd07d74650387cf94942a15" +dependencies = [ + "block-buffer", + "digest", + "opaque-debug", +] + +[[package]] +name = "md5" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "490cc448043f947bae3cbee9c203358d62dbee0db12107a74be5c30ccfd09771" + +[[package]] +name = "memchr" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b16bd47d9e329435e309c58469fe0791c2d0d1ba96ec0954152a5ae2b04387dc" + +[[package]] +name = "memoffset" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59accc507f1338036a0477ef61afdae33cde60840f4dfe481319ce3ad116ddf9" +dependencies = [ + "autocfg 1.0.1", +] + +[[package]] +name = "mime" +version = "0.3.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a60c7ce501c71e03a9c9c0d35b861413ae925bd979cc7a4e30d060069aaac8d" + +[[package]] +name = "mime_guess" +version = "2.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2684d4c2e97d99848d30b324b00c8fcc7e5c897b7cbb5819b09e7c90e8baf212" +dependencies = [ + "mime", + "unicase", +] + +[[package]] +name = "minidom" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "332592c2149fc7dd40a64fc9ef6f0d65607284b474cef9817d1fc8c7e7b3608e" +dependencies = [ + "quick-xml", +] + +[[package]] +name = "miniz_oxide" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a92518e98c078586bc6c934028adcca4c92a53d6a958196de835170a01d84e4b" +dependencies = [ + "adler", + "autocfg 1.0.1", +] + +[[package]] +name = "mio" +version = "0.6.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4afd66f5b91bf2a3bc13fad0e21caedac168ca4c707504e75585648ae80e4cc4" +dependencies = [ + "cfg-if 0.1.10", + "fuchsia-zircon", + "fuchsia-zircon-sys", + "iovec", + "kernel32-sys", + "libc", + "log", + "miow 0.2.2", + "net2", + "slab", + "winapi 0.2.8", +] + +[[package]] +name = "mio" +version = "0.7.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf80d3e903b34e0bd7282b218398aec54e082c840d9baf8339e0080a0c542956" +dependencies = [ + "libc", + "log", + "miow 0.3.7", + "ntapi", + "winapi 0.3.9", +] + +[[package]] +name = "mio-uds" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "afcb699eb26d4332647cc848492bbc15eafb26f08d0304550d5aa1f612e066f0" +dependencies = [ + "iovec", + "libc", + "mio 0.6.23", +] + +[[package]] +name = "miow" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebd808424166322d4a38da87083bfddd3ac4c131334ed55856112eb06d46944d" +dependencies = [ + "kernel32-sys", + "net2", + "winapi 0.2.8", + "ws2_32-sys", +] + +[[package]] +name = "miow" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9f1c5b025cda876f66ef43a113f91ebc9f4ccef34843000e0adf6ebbab84e21" +dependencies = [ + "winapi 0.3.9", +] + +[[package]] +name = "mysqlclient-sys" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e9637d93448044078aaafea7419aed69d301b4a12bcc4aa0ae856eb169bef85" +dependencies = [ + "pkg-config", + "vcpkg", +] + +[[package]] +name = "native-tls" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48ba9f7719b5a0f42f338907614285fb5fd70e53858141f69898a1fb7203b24d" +dependencies = [ + "lazy_static", + "libc", + "log", + "openssl", + "openssl-probe", + "openssl-sys", + "schannel", + "security-framework", + "security-framework-sys", + "tempfile", +] + +[[package]] +name = "net2" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "391630d12b68002ae1e25e8f974306474966550ad82dac6886fb8910c19568ae" +dependencies = [ + "cfg-if 0.1.10", + "libc", + "winapi 0.3.9", +] + +[[package]] +name = "newrelic-telemetry" +version = "0.1.0" +dependencies = [ + "anyhow", + "env_logger", + "flate2", + "futures 0.3.17", + "hyper 0.14.14", + "hyper-tls", + "log", + "rand 0.7.3", + "serde", + "serde_json", + "tokio 1.13.0", + "uuid", +] + +[[package]] +name = "nom" +version = "6.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7413f999671bd4745a7b624bd370a569fb6bc574b23c83a3c5ed2e453f3d5e2" +dependencies = [ + "bitvec", + "funty", + "lexical-core", + "memchr", + "version_check", +] + +[[package]] +name = "ntapi" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f6bb902e437b6d86e03cce10a7e2af662292c5dfef23b65899ea3ac9354ad44" +dependencies = [ + "winapi 0.3.9", +] + +[[package]] +name = "num-bigint" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7d0a3d5e207573f948a9e5376662aa743a2ea13f7c50a554d7af443a73fbfeba" +dependencies = [ + "autocfg 1.0.1", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-bigint" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f93ab6289c7b344a8a9f60f88d80aa20032336fe78da341afc91c8a2341fc75f" +dependencies = [ + "autocfg 1.0.1", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-bigint-dig" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4547ee5541c18742396ae2c895d0717d0f886d8823b8399cdaf7b07d63ad0480" +dependencies = [ + "autocfg 0.1.7", + "byteorder", + "lazy_static", + "libm", + "num-integer", + "num-iter", + "num-traits", + "rand 0.8.4", + "smallvec 1.6.1", + "zeroize", +] + +[[package]] +name = "num-integer" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d2cc698a63b549a70bc047073d2949cce27cd1c7b0a4a862d08a8031bc2801db" +dependencies = [ + "autocfg 1.0.1", + "num-traits", +] + +[[package]] +name = "num-iter" +version = "0.1.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2021c8337a54d21aca0d59a92577a029af9431cb59b909b03252b9c164fad59" +dependencies = [ + "autocfg 1.0.1", + "num-integer", + "num-traits", +] + +[[package]] +name = "num-traits" +version = "0.2.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a64b1ec5cda2586e284722486d802acf1f7dbdc623e2bfc57e65ca1cd099290" +dependencies = [ + "autocfg 1.0.1", + "libm", +] + +[[package]] +name = "num_cpus" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05499f3756671c15885fee9034446956fff3f243d6077b91e5767df161f766b3" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "once_cell" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da32515d9f6e6e489d7bc9d84c71b060db7247dc035bbe44eac88cf87486d8d5" + +[[package]] +name = "opaque-debug" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "624a8340c38c1b80fd549087862da4ba43e08858af025b236e509b6649fc13d5" + +[[package]] +name = "openssl" +version = "0.10.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d7830286ad6a3973c0f1d9b73738f69c76b739301d0229c4b96501695cbe4c8" +dependencies = [ + "bitflags", + "cfg-if 1.0.0", + "foreign-types", + "libc", + "once_cell", + "openssl-sys", +] + +[[package]] +name = "openssl-probe" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28988d872ab76095a6e6ac88d99b54fd267702734fd7ffe610ca27f533ddb95a" + +[[package]] +name = "openssl-sys" +version = "0.9.63" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6b0d6fb7d80f877617dfcb014e605e2b5ab2fb0afdf27935219bb6bd984cb98" +dependencies = [ + "autocfg 1.0.1", + "cc", + "libc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "ordered-multimap" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1c672c7ad9ec066e428c00eb917124a06f08db19e2584de982cc34b1f4c12485" +dependencies = [ + "dlv-list", + "hashbrown 0.9.1", +] + +[[package]] +name = "parking_lot" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f842b1982eb6c2fe34036a4fbfb06dd185a3f5c8edfaacdf7d1ea10b07de6252" +dependencies = [ + "lock_api 0.3.4", + "parking_lot_core 0.6.2", + "rustc_version", +] + +[[package]] +name = "parking_lot" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d7744ac029df22dca6284efe4e898991d28e3085c706c972bcd7da4a27a15eb" +dependencies = [ + "instant", + "lock_api 0.4.4", + "parking_lot_core 0.8.3", +] + +[[package]] +name = "parking_lot_core" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b876b1b9e7ac6e1a74a6da34d25c42e17e8862aa409cbbbdcfc8d86c6f3bc62b" +dependencies = [ + "cfg-if 0.1.10", + "cloudabi", + "libc", + "redox_syscall 0.1.57", + "rustc_version", + "smallvec 0.6.14", + "winapi 0.3.9", +] + +[[package]] +name = "parking_lot_core" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7a782938e745763fe6907fc6ba86946d72f49fe7e21de074e08128a99fb018" +dependencies = [ + "cfg-if 1.0.0", + "instant", + "libc", + "redox_syscall 0.2.8", + "smallvec 1.6.1", + "winapi 0.3.9", +] + +[[package]] +name = "paste" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0744126afe1a6dd7f394cb50a716dbe086cb06e255e53d8d0185d82828358fb5" + +[[package]] +name = "pem" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fd56cbd21fea48d0c440b41cd69c589faacade08c992d9a54e471b79d0fd13eb" +dependencies = [ + "base64 0.13.0", + "once_cell", + "regex", +] + +[[package]] +name = "percent-encoding" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "31010dd2e1ac33d5b46a5b413495239882813e0369f8ed8a5e266f173602f831" + +[[package]] +name = "percent-encoding" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4fd5641d01c8f18a23da7b6fe29298ff4b55afcccdf78973b24cf3175fee32e" + +[[package]] +name = "pin-project" +version = "0.4.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "918192b5c59119d51e0cd221f4d49dde9112824ba717369e903c97d076083d0f" +dependencies = [ + "pin-project-internal 0.4.28", +] + +[[package]] +name = "pin-project" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7509cc106041c40a4518d2af7a61530e1eed0e6285296a3d8c5472806ccc4a4" +dependencies = [ + "pin-project-internal 1.0.7", +] + +[[package]] +name = "pin-project-internal" +version = "0.4.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3be26700300be6d9d23264c73211d8190e755b6b5ca7a1b28230025511b52a5e" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-internal" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c950132583b500556b1efd71d45b319029f2b71518d979fcc208e16b42426f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-lite" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "257b64915a082f7811703966789728173279bdebb956b143dbcd23f6f970a777" + +[[package]] +name = "pin-project-lite" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc0e1f259c92177c30a4c9d177246edd0a3568b25756a977d0632cf8fa37e905" + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "pkg-config" +version = "0.3.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3831453b3449ceb48b6d9c7ad7c96d5ea673e9b470a1dc578c2ce6521230884c" + +[[package]] +name = "ppv-lite86" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac74c624d6b2d21f425f752262f42188365d7b8ff1aff74c82e45136510a4857" + +[[package]] +name = "proc-macro-hack" +version = "0.5.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dbf0c48bc1d91375ae5c3cd81e3722dff1abcf81a30960240640d223f59fe0e5" + +[[package]] +name = "proc-macro-nested" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc881b2c22681370c6a780e47af9840ef841837bc98118431d4e1868bd0c1086" + +[[package]] +name = "proc-macro2" +version = "1.0.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a152013215dca273577e18d2bf00fa862b89b24169fb78c4c95aeb07992c9cec" +dependencies = [ + "unicode-xid", +] + +[[package]] +name = "pulldown-cmark" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffade02495f22453cd593159ea2f59827aae7f53fa8323f756799b670881dcf8" +dependencies = [ + "bitflags", + "getopts", + "memchr", + "unicase", +] + +[[package]] +name = "quick-error" +version = "1.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" + +[[package]] +name = "quick-xml" +version = "0.20.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26aab6b48e2590e4a64d1ed808749ba06257882b461d01ca71baeb747074a6dd" +dependencies = [ + "memchr", +] + +[[package]] +name = "quote" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3d0b9745dc2debf507c8422de05d7226cc1f0644216dfdfead988f9b1ab32a7" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r2d2" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "545c5bc2b880973c9c10e4067418407a0ccaa3091781d1671d46eb35107cb26f" +dependencies = [ + "log", + "parking_lot 0.11.1", + "scheduled-thread-pool", +] + +[[package]] +name = "r2d2_redis" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "182473b876b0b93e353682ec58e207dd1cb4a62278bbe0045fe52b86b74363bb" +dependencies = [ + "r2d2", + "redis 0.20.1", +] + +[[package]] +name = "radium" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "941ba9d78d8e2f7ce474c015eea4d9c6d25b6a3327f9832ee29a4de27f91bbb8" + +[[package]] +name = "rand" +version = "0.4.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "552840b97013b1a26992c11eac34bdd778e464601a4c2054b5f0bff7c6761293" +dependencies = [ + "fuchsia-cprng", + "libc", + "rand_core 0.3.1", + "rdrand", + "winapi 0.3.9", +] + +[[package]] +name = "rand" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a6b1679d49b24bbfe0c803429aa1874472f50d9b363131f0e89fc356b544d03" +dependencies = [ + "getrandom 0.1.16", + "libc", + "rand_chacha 0.2.2", + "rand_core 0.5.1", + "rand_hc 0.2.0", +] + +[[package]] +name = "rand" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e7573632e6454cf6b99d7aac4ccca54be06da05aca2ef7423d22d27d4d4bcd8" +dependencies = [ + "libc", + "rand_chacha 0.3.0", + "rand_core 0.6.2", + "rand_hc 0.3.0", +] + +[[package]] +name = "rand_chacha" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f4c8ed856279c9737206bf725bf36935d8666ead7aa69b52be55af369d193402" +dependencies = [ + "ppv-lite86", + "rand_core 0.5.1", +] + +[[package]] +name = "rand_chacha" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e12735cf05c9e10bf21534da50a147b924d555dc7a547c42e6bb2d5b6017ae0d" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.2", +] + +[[package]] +name = "rand_core" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a6fdeb83b075e8266dcc8762c22776f6877a63111121f5f8c7411e5be7eed4b" +dependencies = [ + "rand_core 0.4.2", +] + +[[package]] +name = "rand_core" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c33a3c44ca05fa6f1807d8e6743f3824e8509beca625669633be0acbdf509dc" + +[[package]] +name = "rand_core" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90bde5296fc891b0cef12a6d03ddccc162ce7b2aff54160af9338f8d40df6d19" +dependencies = [ + "getrandom 0.1.16", +] + +[[package]] +name = "rand_core" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34cf66eb183df1c5876e2dcf6b13d57340741e8dc255b48e40a26de954d06ae7" +dependencies = [ + "getrandom 0.2.2", +] + +[[package]] +name = "rand_hc" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ca3129af7b92a17112d59ad498c6f81eaf463253766b90396d39ea7a39d6613c" +dependencies = [ + "rand_core 0.5.1", +] + +[[package]] +name = "rand_hc" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3190ef7066a446f2e7f42e239d161e905420ccab01eb967c9eb27d21b2322a73" +dependencies = [ + "rand_core 0.6.2", +] + +[[package]] +name = "rdrand" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "678054eb77286b51581ba43620cc911abf02758c91f93f479767aed0f90458b2" +dependencies = [ + "rand_core 0.3.1", +] + +[[package]] +name = "redis" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9bb1079bc5692c03e1479bb7086ac870ac2aaf7115b5e72d314c01368f1a747e" +dependencies = [ + "bytes 0.4.12", + "combine 3.8.1", + "dtoa", + "futures 0.1.31", + "itoa", + "percent-encoding 2.1.0", + "sha1", + "tokio-codec", + "tokio-executor", + "tokio-io", + "tokio-sync", + "tokio-tcp", + "tokio-uds", + "url 2.2.2", +] + +[[package]] +name = "redis" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0a32cb439c4e89c1e6415e5b3b23d9d8cc6dc1bf5a8cade19adbd5418de803be" +dependencies = [ + "async-trait", + "combine 4.5.2", + "dtoa", + "itoa", + "percent-encoding 2.1.0", + "sha1", + "url 2.2.2", +] + +[[package]] +name = "redis_common" +version = "0.0.1" +dependencies = [ + "anyhow", + "container_common", + "serde", + "serde_derive", + "serde_json", + "time 0.2.26", +] + +[[package]] +name = "redox_syscall" +version = "0.1.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41cc0f7e4d5d4544e8861606a285bb08d3e70712ccc7d2b84d7c0ccfaf4b05ce" + +[[package]] +name = "redox_syscall" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "742739e41cd49414de871ea5e549afb7e2a3ac77b589bcbebe8c82fab37147fc" +dependencies = [ + "bitflags", +] + +[[package]] +name = "redox_users" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "528532f3d801c87aec9def2add9ca802fe569e44a544afe633765267840abe64" +dependencies = [ + "getrandom 0.2.2", + "redox_syscall 0.2.8", +] + +[[package]] +name = "regex" +version = "1.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d07a8629359eb56f1e2fb1652bb04212c072a87ba68546a04065d525673ac461" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.6.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f497285884f3fcff424ffc933e56d7cbca511def0c9831a7f9b5f6153e3cc89b" + +[[package]] +name = "remove_dir_all" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3acd125665422973a33ac9d3dd2df85edad0f4ae9b00dafb1a05e43a9f5ef8e7" +dependencies = [ + "winapi 0.3.9", +] + +[[package]] +name = "reqwest" +version = "0.11.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "66d2927ca2f685faf0fc620ac4834690d29e7abb153add10f5812eef20b5e280" +dependencies = [ + "base64 0.13.0", + "bytes 1.0.1", + "encoding_rs", + "futures-core", + "futures-util", + "http", + "http-body 0.4.4", + "hyper 0.14.14", + "hyper-tls", + "ipnet", + "js-sys", + "lazy_static", + "log", + "mime", + "native-tls", + "percent-encoding 2.1.0", + "pin-project-lite 0.2.6", + "serde", + "serde_json", + "serde_urlencoded", + "tokio 1.13.0", + "tokio-native-tls", + "url 2.2.2", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", + "winreg 0.7.0", +] + +[[package]] +name = "resolv-conf" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52e44394d2086d010551b14b53b1f24e31647570cd1deb0379e2c21b329aba00" +dependencies = [ + "hostname", + "quick-error", +] + +[[package]] +name = "ring" +version = "0.16.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3053cf52e236a3ed746dfc745aa9cacf1b791d846bdaf412f60a8d7d6e17c8fc" +dependencies = [ + "cc", + "libc", + "once_cell", + "spin", + "untrusted", + "web-sys", + "winapi 0.3.9", +] + +[[package]] +name = "ring" +version = "0.17.0-alpha.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55ae9a4d2975bdd8254d7bcdd2261be62b42b8920f903ec682b08cdc11b87af2" +dependencies = [ + "cc", + "libc", + "once_cell", + "spin", + "untrusted", + "web-sys", + "winapi 0.3.9", +] + +[[package]] +name = "rpassword" +version = "5.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffc936cf8a7ea60c58f030fd36a612a48f440610214dc54bc36431f9ea0c3efb" +dependencies = [ + "libc", + "winapi 0.3.9", +] + +[[package]] +name = "rsa" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b0aeddcca1082112a6eeb43bf25fd7820b066aaf6eaef776e19d0a1febe38fe" +dependencies = [ + "byteorder", + "digest", + "lazy_static", + "num-bigint-dig", + "num-integer", + "num-iter", + "num-traits", + "pem", + "rand 0.8.4", + "simple_asn1", + "subtle", + "zeroize", +] + +[[package]] +name = "rust-ini" +version = "0.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63471c4aa97a1cf8332a5f97709a79a4234698de6a1f5087faf66f2dae810e22" +dependencies = [ + "cfg-if 1.0.0", + "ordered-multimap", +] + +[[package]] +name = "rust-s3" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18c58d4682844a5d6301efbf915dd7a9d3d638006f9bb821527a0bbbf2a4cfc2" +dependencies = [ + "anyhow", + "async-trait", + "aws-creds", + "aws-region", + "base64 0.13.0", + "cfg-if 1.0.0", + "chrono", + "hex", + "hmac 0.11.0", + "http", + "log", + "maybe-async", + "md5", + "minidom", + "percent-encoding 2.1.0", + "reqwest", + "serde", + "serde-xml-rs", + "serde_derive", + "sha2", + "tokio 1.13.0", + "tokio-stream", + "url 2.2.2", +] + +[[package]] +name = "rustc_version" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "138e3e0acb6c9fb258b19b67cb8abd63c00679d2851805ea151465464fe9030a" +dependencies = [ + "semver", +] + +[[package]] +name = "rustls" +version = "0.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "35edb675feee39aec9c99fa5ff985081995a06d594114ae14cbe797ad7b7a6d7" +dependencies = [ + "base64 0.13.0", + "log", + "ring 0.16.20", + "sct", + "webpki", +] + +[[package]] +name = "rustls-native-certs" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a07b7c1885bd8ed3831c289b7870b13ef46fe0e856d288c30d9cc17d75a2092" +dependencies = [ + "openssl-probe", + "rustls", + "schannel", + "security-framework", +] + +[[package]] +name = "ryu" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "71d301d4193d031abdd79ff7e3dd721168a9572ef3fe51a1517aba235bd8f86e" + +[[package]] +name = "schannel" +version = "0.1.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f05ba609c234e60bee0d547fe94a4c7e9da733d1c962cf6e59efa4cd9c8bc75" +dependencies = [ + "lazy_static", + "winapi 0.3.9", +] + +[[package]] +name = "scheduled-thread-pool" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc6f74fd1204073fa02d5d5d68bec8021be4c38690b61264b2fdb48083d0e7d7" +dependencies = [ + "parking_lot 0.11.1", +] + +[[package]] +name = "scopeguard" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd" + +[[package]] +name = "sct" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b362b83898e0e69f38515b82ee15aa80636befe47c3b6d3d89a911e78fc228ce" +dependencies = [ + "ring 0.16.20", + "untrusted", +] + +[[package]] +name = "security-framework" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3670b1d2fdf6084d192bc71ead7aabe6c06aa2ea3fbd9cc3ac111fa5c2b1bd84" +dependencies = [ + "bitflags", + "core-foundation", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3676258fd3cfe2c9a0ec99ce3038798d847ce3e4bb17746373eb9f0f1ac16339" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d7eb9ef2c18661902cc47e535f9bc51b78acd254da71d375c2f6720d9a40403" +dependencies = [ + "semver-parser", +] + +[[package]] +name = "semver-parser" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "388a1df253eca08550bef6c72392cfe7c30914bf41df5269b68cbd6ff8f570a3" + +[[package]] +name = "serde" +version = "1.0.130" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f12d06de37cf59146fbdecab66aa99f9fe4f78722e3607577a5375d66bd0c913" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde-xml-rs" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65162e9059be2f6a3421ebbb4fef3e74b7d9e7c60c50a0e292c6239f19f1edfa" +dependencies = [ + "log", + "serde", + "thiserror", + "xml-rs", +] + +[[package]] +name = "serde_derive" +version = "1.0.130" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7bc1a1ab1961464eae040d96713baa5a724a8152c1222492465b54322ec508b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e466864e431129c7e0d3476b92f20458e5879919a0596c6472738d9fa2d342f8" +dependencies = [ + "indexmap", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0421d4f173fab82d72d6babf36d57fae38b994ca5c2d78e704260ba6d12118b" +dependencies = [ + "serde", +] + +[[package]] +name = "serde_repr" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98d0516900518c29efa217c298fa1f4e6c6ffc85ae29fd7f4ee48f176e1a9ed5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edfa57a7f8d9c1d260a549e7224100f6c43d43f9103e06dd8b4095a9b2b43ce9" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "sha-1" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8c4cfa741c5832d0ef7fab46cabed29c2aae926db0b11bb2069edd8db5e64e16" +dependencies = [ + "block-buffer", + "cfg-if 1.0.0", + "cpufeatures", + "digest", + "opaque-debug", +] + +[[package]] +name = "sha1" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2579985fda508104f7587689507983eadd6a6e84dd35d6d115361f530916fa0d" + +[[package]] +name = "sha2" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b362ae5752fd2137731f9fa25fd4d9058af34666ca1966fb969119cc35719f12" +dependencies = [ + "block-buffer", + "cfg-if 1.0.0", + "cpufeatures", + "digest", + "opaque-debug", +] + +[[package]] +name = "signal-hook-registry" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "16f1d0fef1604ba8f7a073c7e701f213e056707210e9020af4528e0101ce11a6" +dependencies = [ + "libc", +] + +[[package]] +name = "simple_asn1" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8eb4ea60fb301dc81dfc113df680571045d375ab7345d171c5dc7d7e13107a80" +dependencies = [ + "chrono", + "num-bigint 0.4.3", + "num-traits", + "thiserror", +] + +[[package]] +name = "slab" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f173ac3d1a7e3b28003f40de0b5ce7fe2710f9b9dc3fc38664cebee46b3b6527" + +[[package]] +name = "smallvec" +version = "0.6.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97fcaeba89edba30f044a10c6a3cc39df9c3f17d7cd829dd1446cab35f890e0" +dependencies = [ + "maybe-uninit", +] + +[[package]] +name = "smallvec" +version = "1.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe0f37c9e8f3c5a4a66ad655a93c74daac4ad00c441533bf5c6e7990bb42604e" + +[[package]] +name = "socket2" +version = "0.3.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "122e570113d28d773067fab24266b66753f6ea915758651696b6e35e49f88d6e" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "winapi 0.3.9", +] + +[[package]] +name = "socket2" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e3dfc207c526015c632472a77be09cf1b6e46866581aecae5cc38fb4235dea2" +dependencies = [ + "libc", + "winapi 0.3.9", +] + +[[package]] +name = "spin" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e63cff320ae2c57904679ba7cb63280a3dc4613885beafb148ee7bf9aa9042d" + +[[package]] +name = "sqlformat" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d86e3c77ff882a828346ba401a7ef4b8e440df804491c6064fe8295765de71c" +dependencies = [ + "lazy_static", + "maplit", + "nom", + "regex", + "unicode_categories", +] + +[[package]] +name = "sqlx" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e4b94ab0f8c21ee4899b93b06451ef5d965f1a355982ee73684338228498440" +dependencies = [ + "sqlx-core", + "sqlx-macros", +] + +[[package]] +name = "sqlx-core" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec28b91a01e1fe286d6ba66f68289a2286df023fc97444e1fd86c2fd6d5dc026" +dependencies = [ + "ahash 0.7.4", + "atoi", + "base64 0.13.0", + "bitflags", + "byteorder", + "bytes 1.0.1", + "chrono", + "crc", + "crossbeam-channel", + "crossbeam-queue", + "crossbeam-utils 0.8.5", + "digest", + "either", + "futures-channel", + "futures-core", + "futures-intrusive", + "futures-util", + "generic-array", + "hashlink", + "hex", + "itoa", + "libc", + "log", + "memchr", + "num-bigint 0.3.2", + "once_cell", + "parking_lot 0.11.1", + "percent-encoding 2.1.0", + "rand 0.8.4", + "rsa", + "serde", + "sha-1", + "sha2", + "smallvec 1.6.1", + "sqlformat", + "sqlx-rt", + "stringprep", + "thiserror", + "tokio-stream", + "url 2.2.2", + "whoami", +] + +[[package]] +name = "sqlx-macros" +version = "0.5.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4dc33c35d54774eed73d54568d47a6ac099aed8af5e1556a017c131be88217d5" +dependencies = [ + "dotenv", + "either", + "futures 0.3.17", + "heck", + "hex", + "once_cell", + "proc-macro2", + "quote", + "serde", + "serde_json", + "sha2", + "sqlx-core", + "sqlx-rt", + "syn", + "url 2.2.2", +] + +[[package]] +name = "sqlx-rt" +version = "0.5.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d1bd069de53442e7a320f525a6d4deb8bb0621ac7a55f7eccbc2b58b57f43d0" +dependencies = [ + "actix-rt 2.4.0", + "native-tls", + "once_cell", + "tokio 1.13.0", + "tokio-native-tls", +] + +[[package]] +name = "standback" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e113fb6f3de07a243d434a56ec6f186dfd51cb08448239fe7bcae73f87ff28ff" +dependencies = [ + "version_check", +] + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "stdweb" +version = "0.4.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d022496b16281348b52d0e30ae99e01a73d737b2f45d38fed4edf79f9325a1d5" +dependencies = [ + "discard", + "rustc_version", + "stdweb-derive", + "stdweb-internal-macros", + "stdweb-internal-runtime", + "wasm-bindgen", +] + +[[package]] +name = "stdweb-derive" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c87a60a40fccc84bef0652345bbbbbe20a605bf5d0ce81719fc476f5c03b50ef" +dependencies = [ + "proc-macro2", + "quote", + "serde", + "serde_derive", + "syn", +] + +[[package]] +name = "stdweb-internal-macros" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "58fa5ff6ad0d98d1ffa8cb115892b6e69d67799f6763e162a1c9db421dc22e11" +dependencies = [ + "base-x", + "proc-macro2", + "quote", + "serde", + "serde_derive", + "serde_json", + "sha1", + "syn", +] + +[[package]] +name = "stdweb-internal-runtime" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "213701ba3370744dcd1a12960caa4843b3d68b4d1c0a5d575e0d65b2ee9d16c0" + +[[package]] +name = "storyteller-web" +version = "0.0.1" +dependencies = [ + "actix", + "actix-cors 0.6.0-beta.3", + "actix-files", + "actix-http 3.0.0-beta.11", + "actix-multipart", + "actix-rt 2.4.0", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "actix-web 4.0.0-beta.10", + "actix-web-actors", + "anyhow", + "base64 0.13.0", + "bcrypt", + "chrono", + "config", + "container_common", + "data-encoding", + "database_queries", + "derive_more", + "diesel", + "dotenv", + "easyenv", + "futures 0.1.31", + "futures 0.3.17", + "futures-util", + "hex", + "hmac 0.10.1", + "hostname", + "http_server_common", + "hyper 0.14.14", + "infer", + "jwt", + "language-tags 0.3.2", + "lazy_static", + "lexical-sort", + "lfu", + "limitation", + "log", + "lru_time_cache", + "magic-crypt", + "md-5", + "native-tls", + "newrelic-telemetry", + "pulldown-cmark", + "r2d2_redis", + "rand 0.8.4", + "redis_common", + "regex", + "reqwest", + "ring 0.17.0-alpha.10", + "rust-s3", + "serde", + "serde_derive", + "serde_json", + "sha2", + "sqlx", + "subprocess", + "tempdir", + "time 0.2.26", + "tokio 1.13.0", + "toml", + "tts_common", + "twitch_api2", + "twitch_common", + "twitch_oauth2", + "user_input_common", + "uuid", +] + +[[package]] +name = "stringprep" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ee348cb74b87454fff4b551cbf727025810a004f88aeacae7f85b87f4e9a1c1" +dependencies = [ + "unicode-bidi", + "unicode-normalization", +] + +[[package]] +name = "subprocess" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "334b801f9ca2529ba9605109f1d2835f15aff94cd6cfe5afe6ce8cf71e5fd3c4" +dependencies = [ + "libc", + "winapi 0.3.9", +] + +[[package]] +name = "subtle" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e81da0851ada1f3e9d4312c704aa4f8806f0f9d69faaf8df2f3464b4a9437c2" + +[[package]] +name = "syn" +version = "1.0.72" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a1e8cdbefb79a9a5a65e0db8b47b723ee907b7c7f8496c76a1770b5c310bab82" +dependencies = [ + "proc-macro2", + "quote", + "unicode-xid", +] + +[[package]] +name = "synstructure" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b834f2d66f734cb897113e34aaff2f1ab4719ca946f9a7358dba8f8064148701" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "unicode-xid", +] + +[[package]] +name = "tap" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "55937e1799185b12863d447f42597ed69d9928686b8d88a1df17376a097d8369" + +[[package]] +name = "tempdir" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15f2b5fb00ccdf689e0149d1b1b3c03fead81c2b37735d812fa8bddbbf41b6d8" +dependencies = [ + "rand 0.4.6", + "remove_dir_all", +] + +[[package]] +name = "tempfile" +version = "3.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dac1c663cfc93810f88aed9b8941d48cabf856a1b111c29a40439018d870eb22" +dependencies = [ + "cfg-if 1.0.0", + "libc", + "rand 0.8.4", + "redox_syscall 0.2.8", + "remove_dir_all", + "winapi 0.3.9", +] + +[[package]] +name = "termcolor" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2dfed899f0eb03f32ee8c6a0aabdb8a7949659e3466561fc0adf54e26d88c5f4" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "thiserror" +version = "1.0.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "854babe52e4df1653706b98fcfc05843010039b406875930a70e4d9644e5c417" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.30" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa32fd3f627f367fe16f893e2597ae3c05020f8bba2666a4e6ea73d377e5714b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "threadpool" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d050e60b33d41c19108b32cea32164033a9013fe3b46cbd4457559bfbf77afaa" +dependencies = [ + "num_cpus", +] + +[[package]] +name = "tiger" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "443e531cbcf9de83258cfef70bcd56c91188de5819ebd4b19c85f589e0617005" +dependencies = [ + "block-buffer", + "byteorder", + "digest", +] + +[[package]] +name = "time" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db9e6914ab8b1ae1c260a4ae7a49b6c5611b40328a735b21862567685e73255" +dependencies = [ + "libc", + "wasi 0.10.0+wasi-snapshot-preview1", + "winapi 0.3.9", +] + +[[package]] +name = "time" +version = "0.2.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08a8cbfbf47955132d0202d1662f49b2423ae35862aee471f3ba4b133358f372" +dependencies = [ + "const_fn", + "libc", + "standback", + "stdweb", + "time-macros", + "version_check", + "winapi 0.3.9", +] + +[[package]] +name = "time" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "99beeb0daeac2bd1e86ac2c21caddecb244b39a093594da1a661ec2060c7aedd" +dependencies = [ + "itoa", + "libc", +] + +[[package]] +name = "time-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "957e9c6e26f12cb6d0dd7fc776bb67a706312e7299aed74c8dd5b17ebb27e2f1" +dependencies = [ + "proc-macro-hack", + "time-macros-impl", +] + +[[package]] +name = "time-macros-impl" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5c3be1edfad6027c69f5491cf4cb310d1a71ecd6af742788c6ff8bced86b8fa" +dependencies = [ + "proc-macro-hack", + "proc-macro2", + "quote", + "standback", + "syn", +] + +[[package]] +name = "tinyvec" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b5220f05bb7de7f3f53c7c065e1199b3172696fe2db9f9c4d8ad9b4ee74c342" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cda74da7e1a664f795bb1f8a87ec406fb89a02522cf6e50620d016add6dbbf5c" + +[[package]] +name = "tokio" +version = "0.2.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6703a273949a90131b290be1fe7b039d0fc884aa1935860dfcbe056f28cd8092" +dependencies = [ + "bytes 0.5.6", + "fnv", + "futures-core", + "iovec", + "lazy_static", + "libc", + "memchr", + "mio 0.6.23", + "mio-uds", + "pin-project-lite 0.1.12", + "signal-hook-registry", + "slab", + "winapi 0.3.9", +] + +[[package]] +name = "tokio" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "588b2d10a336da58d877567cd8fb8a14b463e2104910f8132cd054b4b96e29ee" +dependencies = [ + "autocfg 1.0.1", + "bytes 1.0.1", + "libc", + "memchr", + "mio 0.7.11", + "num_cpus", + "once_cell", + "parking_lot 0.11.1", + "pin-project-lite 0.2.6", + "signal-hook-registry", + "tokio-macros", + "winapi 0.3.9", +] + +[[package]] +name = "tokio-codec" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25b2998660ba0e70d18684de5d06b70b70a3a747469af9dea7618cc59e75976b" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "tokio-io", +] + +[[package]] +name = "tokio-executor" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb2d1b8f4548dbf5e1f7818512e9c406860678f29c300cdf0ebac72d1a3a1671" +dependencies = [ + "crossbeam-utils 0.7.2", + "futures 0.1.31", +] + +[[package]] +name = "tokio-io" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57fc868aae093479e3131e3d165c93b1c7474109d13c90ec0dda2a1bbfff0674" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "log", +] + +[[package]] +name = "tokio-macros" +version = "1.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "114383b041aa6212c579467afa0075fbbdd0718de036100bc0ba7961d8cb9095" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-native-tls" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f7d995660bd2b7f8c1568414c1126076c13fbb725c40112dc0120b78eb9b717b" +dependencies = [ + "native-tls", + "tokio 1.13.0", +] + +[[package]] +name = "tokio-reactor" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09bc590ec4ba8ba87652da2068d150dcada2cfa2e07faae270a5e0409aa51351" +dependencies = [ + "crossbeam-utils 0.7.2", + "futures 0.1.31", + "lazy_static", + "log", + "mio 0.6.23", + "num_cpus", + "parking_lot 0.9.0", + "slab", + "tokio-executor", + "tokio-io", + "tokio-sync", +] + +[[package]] +name = "tokio-rustls" +version = "0.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc6844de72e57df1980054b38be3a9f4702aba4858be64dd700181a8a6d0e1b6" +dependencies = [ + "rustls", + "tokio 1.13.0", + "webpki", +] + +[[package]] +name = "tokio-stream" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50145484efff8818b5ccd256697f36863f587da82cf8b409c53adf1e840798e3" +dependencies = [ + "futures-core", + "pin-project-lite 0.2.6", + "tokio 1.13.0", +] + +[[package]] +name = "tokio-sync" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edfe50152bc8164fcc456dab7891fa9bf8beaf01c5ee7e1dd43a397c3cf87dee" +dependencies = [ + "fnv", + "futures 0.1.31", +] + +[[package]] +name = "tokio-tcp" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "98df18ed66e3b72e742f185882a9e201892407957e45fbff8da17ae7a7c51f72" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "iovec", + "mio 0.6.23", + "tokio-io", + "tokio-reactor", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "511de3f85caf1c98983545490c3d09685fa8eb634e57eec22bb4db271f46cbd8" +dependencies = [ + "futures-util", + "log", + "pin-project 1.0.7", + "rustls", + "tokio 1.13.0", + "tokio-rustls", + "tungstenite", + "webpki", + "webpki-roots", +] + +[[package]] +name = "tokio-uds" +version = "0.2.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab57a4ac4111c8c9dbcf70779f6fc8bc35ae4b2454809febac840ad19bd7e4e0" +dependencies = [ + "bytes 0.4.12", + "futures 0.1.31", + "iovec", + "libc", + "log", + "mio 0.6.23", + "mio-uds", + "tokio-codec", + "tokio-io", + "tokio-reactor", +] + +[[package]] +name = "tokio-util" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be8242891f2b6cbef26a2d7e8605133c2c554cd35b3e4948ea892d6d68436499" +dependencies = [ + "bytes 0.5.6", + "futures-core", + "futures-sink", + "log", + "pin-project-lite 0.1.12", + "tokio 0.2.25", +] + +[[package]] +name = "tokio-util" +version = "0.6.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e99e1983e5d376cd8eb4b66604d2e99e79f5bd988c3055891dcd8c9e2604cc0" +dependencies = [ + "bytes 1.0.1", + "futures-core", + "futures-sink", + "log", + "pin-project-lite 0.2.6", + "tokio 1.13.0", +] + +[[package]] +name = "toml" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a31142970826733df8241ef35dc040ef98c679ab14d7c3e54d827099b3acecaa" +dependencies = [ + "serde", +] + +[[package]] +name = "tower-service" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "360dfd1d6d30e05fda32ace2c8c70e9c0a9da713275777f5a4dbb8a1893930c6" + +[[package]] +name = "tracing" +version = "0.1.26" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09adeb8c97449311ccd28a427f96fb563e7fd31aabf994189879d9da2394b89d" +dependencies = [ + "cfg-if 1.0.0", + "log", + "pin-project-lite 0.2.6", + "tracing-core", +] + +[[package]] +name = "tracing-core" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9ff14f98b1a4b289c6248a023c1c2fa1491062964e9fed67ab29c4e4da4a052" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "tracing-futures" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97d095ae15e245a057c8e8451bab9b3ee1e1f68e9ba2b4fbc18d0ac5237835f2" +dependencies = [ + "pin-project 1.0.7", + "tracing", +] + +[[package]] +name = "trust-dns-proto" +version = "0.19.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1cad71a0c0d68ab9941d2fb6e82f8fb2e86d9945b94e1661dd0aaea2b88215a9" +dependencies = [ + "async-trait", + "cfg-if 1.0.0", + "enum-as-inner", + "futures 0.3.17", + "idna 0.2.3", + "lazy_static", + "log", + "rand 0.7.3", + "smallvec 1.6.1", + "thiserror", + "tokio 0.2.25", + "url 2.2.2", +] + +[[package]] +name = "trust-dns-resolver" +version = "0.19.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "710f593b371175db53a26d0b38ed2978fafb9e9e8d3868b1acd753ea18df0ceb" +dependencies = [ + "cfg-if 0.1.10", + "futures 0.3.17", + "ipconfig", + "lazy_static", + "log", + "lru-cache", + "resolv-conf", + "smallvec 1.6.1", + "thiserror", + "tokio 0.2.25", + "trust-dns-proto", +] + +[[package]] +name = "try-lock" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59547bce71d9c38b83d9c0e92b6066c4253371f15005def0c30d9657f50c7642" + +[[package]] +name = "tts_common" +version = "0.0.1" +dependencies = [ + "anyhow", + "container_common", + "log", + "once_cell", + "regex", + "unicode-segmentation", + "user_input_common", +] + +[[package]] +name = "tungstenite" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a0b2d8558abd2e276b0a8df5c05a2ec762609344191e5fd23e292c910e9165b5" +dependencies = [ + "base64 0.13.0", + "byteorder", + "bytes 1.0.1", + "http", + "httparse", + "log", + "rand 0.8.4", + "rustls", + "rustls-native-certs", + "sha-1", + "thiserror", + "url 2.2.2", + "utf-8", + "webpki", +] + +[[package]] +name = "twitch-pubsub-subscriber" +version = "0.0.1" +dependencies = [ + "actix", + "actix-cors 0.5.4", + "actix-files", + "actix-http 3.0.0-beta.11", + "actix-rt 2.4.0", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "actix-web 4.0.0-beta.10", + "actix-web-actors", + "anyhow", + "base64 0.13.0", + "bcrypt", + "chrono", + "config", + "container_common", + "crossbeam", + "database_queries", + "derive_more", + "diesel", + "dotenv", + "easyenv", + "futures 0.1.31", + "futures 0.3.17", + "futures-timer", + "futures-util", + "hex", + "hmac 0.10.1", + "hostname", + "http_server_common", + "hyper 0.13.10", + "jwt", + "lazy_static", + "lfu", + "limitation", + "log", + "lru_time_cache", + "magic-crypt", + "md-5", + "native-tls", + "pulldown-cmark", + "r2d2_redis", + "rand 0.8.4", + "redis_common", + "regex", + "reqwest", + "rpassword", + "serde", + "serde_derive", + "serde_json", + "sha2", + "sqlx", + "time 0.2.26", + "tokio 1.13.0", + "tokio-tungstenite", + "toml", + "twitch_api2", + "twitch_common", + "twitch_oauth2", + "uuid", +] + +[[package]] +name = "twitch_api2" +version = "0.6.0-rc.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "600ce224d5eafc176688f05836bbcfec3d314956b5f4a5a360218b4289923dd5" +dependencies = [ + "aliri_braid", + "async-trait", + "chrono", + "displaydoc", + "futures 0.3.17", + "http", + "once_cell", + "reqwest", + "serde", + "serde_derive", + "serde_json", + "serde_path_to_error", + "serde_repr", + "thiserror", + "twitch_oauth2", + "typed-builder", + "url 2.2.2", + "version_check", +] + +[[package]] +name = "twitch_common" +version = "0.0.1" +dependencies = [ + "anyhow", + "container_common", + "easyenv", + "once_cell", + "regex", + "serde", + "serde_derive", + "twitch_oauth2", +] + +[[package]] +name = "twitch_oauth2" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f454ca1e9b8dc97a7255ff50e13a8a1d6ffab8ce6b9445cfc331d6420d40a953" +dependencies = [ + "aliri_braid", + "async-trait", + "base64 0.13.0", + "displaydoc", + "http", + "once_cell", + "rand 0.8.4", + "reqwest", + "serde", + "serde_json", + "thiserror", + "url 2.2.2", +] + +[[package]] +name = "twoway" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c57ffb460d7c24cd6eda43694110189030a3d1dfe418416d9468fd1c1d290b47" +dependencies = [ + "memchr", + "unchecked-index", +] + +[[package]] +name = "typed-builder" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a46ee5bd706ff79131be9c94e7edcb82b703c487766a114434e5790361cf08c5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "typenum" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f6906492a7cd215bfa4cf595b600146ccfac0c79bcbd1f3000162af5e8b06" + +[[package]] +name = "unchecked-index" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eeba86d422ce181a719445e51872fa30f1f7413b62becb52e95ec91aa262d85c" + +[[package]] +name = "unicase" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50f37be617794602aabbeee0be4f259dc1778fabe05e2d67ee8f79326d5cb4f6" +dependencies = [ + "version_check", +] + +[[package]] +name = "unicode-bidi" +version = "0.3.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eeb8be209bb1c96b7c177c7420d26e04eccacb0eeae6b980e35fcb74678107e0" +dependencies = [ + "matches", +] + +[[package]] +name = "unicode-normalization" +version = "0.1.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07fbfce1c8a97d547e8b5334978438d9d6ec8c20e38f56d4a4374d181493eaef" +dependencies = [ + "tinyvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0d2e7be6ae3a5fa87eed5fb451aff96f2573d2694942e40543ae0bbe19c796" + +[[package]] +name = "unicode-width" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9337591893a19b88d8d87f2cec1e73fad5cdfd10e5a6f349f498ad6ea2ffb1e3" + +[[package]] +name = "unicode-xid" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ccb82d61f80a663efe1f787a51b16b5a51e3314d6ac365b08639f52387b33f3" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unreachable" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "382810877fe448991dfc7f0dd6e3ae5d58088fd0ea5e35189655f84e6814fa56" +dependencies = [ + "void", +] + +[[package]] +name = "untrusted" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a156c684c91ea7d62626509bce3cb4e1d9ed5c4d978f7b4352658f96a4c26b4a" + +[[package]] +name = "url" +version = "1.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dd4e7c0d531266369519a4aa4f399d748bd37043b00bde1e4ff1f60a120b355a" +dependencies = [ + "idna 0.1.5", + "matches", + "percent-encoding 1.0.1", +] + +[[package]] +name = "url" +version = "2.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a507c383b2d33b5fc35d1861e77e6b383d158b2da5e14fe51b83dfedf6fd578c" +dependencies = [ + "form_urlencoded", + "idna 0.2.3", + "matches", + "percent-encoding 2.1.0", +] + +[[package]] +name = "user_input_common" +version = "0.0.1" +dependencies = [ + "anyhow", + "container_common", + "log", + "once_cell", + "regex", + "unicode-segmentation", +] + +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + +[[package]] +name = "uuid" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bc5cf98d8186244414c848017f0e2676b3fcb46807f6668a97dfe67359a3c4b7" +dependencies = [ + "getrandom 0.2.2", +] + +[[package]] +name = "vcpkg" +version = "0.2.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cbdbff6266a24120518560b5dc983096efb98462e51d0d68169895b237be3e5d" + +[[package]] +name = "version_check" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fecdca9a5291cc2b8dcf7dc02453fee791a280f3743cb0905f8822ae463b3fe" + +[[package]] +name = "void" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a02e4885ed3bc0f2de90ea6dd45ebcbb66dacffe03547fadbb0eeae2770887d" + +[[package]] +name = "want" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1ce8a968cb1cd110d136ff8b819a556d6fb6d919363c61534f6860c7eb172ba0" +dependencies = [ + "log", + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.9.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cccddf32554fecc6acb585f82a32a72e28b48f8c4c1883ddfeeeaa96f7d8e519" + +[[package]] +name = "wasi" +version = "0.10.0+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a143597ca7c7793eff794def352d41792a93c481eb1042423ff7ff72ba2c31f" + +[[package]] +name = "wasm-bindgen" +version = "0.2.74" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d54ee1d4ed486f78874278e63e4069fc1ab9f6a18ca492076ffb90c5eb2997fd" +dependencies = [ + "cfg-if 1.0.0", + "wasm-bindgen-macro", +] + +[[package]] +name = "wasm-bindgen-backend" +version = "0.2.74" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b33f6a0694ccfea53d94db8b2ed1c3a8a4c86dd936b13b9f0a15ec4a451b900" +dependencies = [ + "bumpalo", + "lazy_static", + "log", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5fba7978c679d53ce2d0ac80c8c175840feb849a161664365d1287b41f2e67f1" +dependencies = [ + "cfg-if 1.0.0", + "js-sys", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.74" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "088169ca61430fe1e58b8096c24975251700e7b1f6fd91cc9d59b04fb9b18bd4" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.74" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be2241542ff3d9f241f5e2cb6dd09b37efe786df8851c54957683a49f0987a97" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-backend", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.74" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7cff876b8f18eed75a66cf49b65e7f967cb354a7aa16003fb55dbfd25b44b4f" + +[[package]] +name = "web-sys" +version = "0.3.51" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e828417b379f3df7111d3a2a9e5753706cae29c41f7c4029ee9fd77f3e09e582" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki" +version = "0.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e38c0608262c46d4a56202ebabdeb094cef7e560ca7a226c6bf055188aa4ea" +dependencies = [ + "ring 0.16.20", + "untrusted", +] + +[[package]] +name = "webpki-roots" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aabe153544e473b775453675851ecc86863d2a81d786d741f6b76778f2a48940" +dependencies = [ + "webpki", +] + +[[package]] +name = "websocket-gateway" +version = "0.0.1" +dependencies = [ + "actix", + "actix-cors 0.5.4", + "actix-files", + "actix-http 3.0.0-beta.11", + "actix-rt 2.4.0", + "actix-service 2.0.1", + "actix-utils 3.0.0", + "actix-web 4.0.0-beta.10", + "actix-web-actors", + "anyhow", + "base64 0.13.0", + "bcrypt", + "chrono", + "config", + "container_common", + "crossbeam", + "database_queries", + "derive_more", + "diesel", + "dotenv", + "easyenv", + "futures 0.1.31", + "futures 0.3.17", + "futures-timer", + "futures-util", + "hex", + "hmac 0.10.1", + "hostname", + "http_server_common", + "hyper 0.13.10", + "jwt", + "lazy_static", + "lfu", + "limitation", + "log", + "lru_time_cache", + "magic-crypt", + "md-5", + "native-tls", + "pulldown-cmark", + "r2d2_redis", + "rand 0.8.4", + "redis_common", + "regex", + "reqwest", + "rpassword", + "serde", + "serde_derive", + "serde_json", + "sha2", + "sqlx", + "time 0.2.26", + "tokio 1.13.0", + "tokio-tungstenite", + "toml", + "twitch_api2", + "twitch_common", + "twitch_oauth2", + "uuid", +] + +[[package]] +name = "whoami" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4abacf325c958dfeaf1046931d37f2a901b6dfe0968ee965a29e94c6766b2af6" +dependencies = [ + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "widestring" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c168940144dd21fd8046987c16a46a33d5fc84eec29ef9dcddc2ac9e31526b7c" + +[[package]] +name = "wildmatch" +version = "2.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6c48bd20df7e4ced539c12f570f937c6b4884928a87fee70a479d72f031d4e0" + +[[package]] +name = "winapi" +version = "0.2.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "167dc9d6949a9b857f3451275e911c3f44255842c1f7a76f33c55103a909087a" + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-build" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d315eee3b34aca4797b2da6b13ed88266e6d612562a0c46390af8299fc699bc" + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70ec6ce85bb158151cae5e5c87f95a8e97d2c0c4b001223f33a334e3ce5de178" +dependencies = [ + "winapi 0.3.9", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "winreg" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2986deb581c4fe11b621998a5e53361efe6b48a151178d0cd9eeffa4dc6acc9" +dependencies = [ + "winapi 0.3.9", +] + +[[package]] +name = "winreg" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0120db82e8a1e0b9fb3345a539c478767c0048d842860994d96113d5b667bd69" +dependencies = [ + "winapi 0.3.9", +] + +[[package]] +name = "ws2_32-sys" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d59cefebd0c892fa2dd6de581e937301d8552cb44489cdff035c6187cb63fa5e" +dependencies = [ + "winapi 0.2.8", + "winapi-build", +] + +[[package]] +name = "wyz" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85e60b0d1b5f99db2556934e21937020776a5d31520bf169e851ac44e6420214" + +[[package]] +name = "xml-rs" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b07db065a5cf61a7e4ba64f29e67db906fb1787316516c4e6e5ff0fea1efcd8a" + +[[package]] +name = "zeroize" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4756f7db3f7b5574938c3eb1c117038b8e07f95ee6718c0efad4ac21508f1efd" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2c1e130bebaeab2f23886bf9acbaca14b092408c452543c857f66399cd6dab1" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zstd" +version = "0.7.0+zstd.1.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9428752481d8372e15b1bf779ea518a179ad6c771cca2d2c60e4fbff3cc2cd52" +dependencies = [ + "zstd-safe", +] + +[[package]] +name = "zstd-safe" +version = "3.1.0+zstd.1.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aa1926623ad7fe406e090555387daf73db555b948134b4d73eac5eb08fb666d" +dependencies = [ + "libc", + "zstd-sys", +] + +[[package]] +name = "zstd-sys" +version = "1.5.0+zstd.1.4.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4e6c094340240369025fc6b731b054ee2a834328fa584310ac96aa4baebdc465" +dependencies = [ + "cc", + "libc", +] diff --git a/Cargo.toml b/Cargo.toml new file mode 100644 index 0000000000..51af671f1b --- /dev/null +++ b/Cargo.toml @@ -0,0 +1,19 @@ +[workspace] + +members = [ + # Shared libraries + "crates/lib/config", + "crates/lib/container_common", + "crates/lib/database_queries", + "crates/lib/http_server_common", + "crates/lib/redis_common", + "crates/lib/tts_common", + "crates/lib/twitch_common", + "crates/lib/user_input_common", + + # Application + "crates/service/storyteller_web", + "crates/service/twitch_pubsub_subscriber", + "crates/service/websocket_gateway", +] + diff --git a/Dockerfile b/Dockerfile new file mode 100644 index 0000000000..91d0a45769 --- /dev/null +++ b/Dockerfile @@ -0,0 +1,318 @@ +# ==================== Python Build Base ==================== + +FROM ghcr.io/storytold/docker-base-images-nvidia-cuda-experimental:080f96bc7087 as pybuild-base + +## Authenticating with GHCR locally: +## https://docs.github.com/en/packages/working-with-a-github-packages-registry/migrating-to-the-container-registry-from-the-docker-registry +## ``` +## export CR_PAT=YOUR_TOKEN +## echo $CR_PAT | docker login ghcr.io -u USERNAME --password-stdin +## ``` +# +##FROM nvidia/cuda:10.1-devel-ubuntu18.04 as pybuild-base +#FROM nvidia/cuda:11.3.0-devel-ubuntu20.04 as pybuild-base +# +## tzdata install hangs without presetting the time zone +#ENV TZ=UTC +#RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone +# +## NB(2021-05-30): This is the image that Tacotron 2's Dockerfile specifies. +## I have not tried this, but it might be necessary. +##FROM pytorch/pytorch:nightly-devel-cuda10.0-cudnn7 +##ENV PATH /usr/local/nvidia/bin:/usr/local/cuda/bin:${PATH} +# +## NB: https://github.com/NVIDIA/nvidia-docker/issues/864#issuecomment-439848887 +## NB: We do not install ffmpeg, since the version is 3.8.* series and we need 4.2.* +## +## NB(2021-05-30): We need to install `cython3` for building modern Tacotron 2 +## dependencies (ie. protoc) on old Ubuntu +## RUN apt-get update \ +## && apt-get install -y \ +#RUN apt-get update +##RUN apt-cache search python3.8 +# +## NB: Here we install nvidia drivers +## NB: Avoid keyboard prompt +#RUN DEBIAN_FRONTEND=noninteractive apt-get install -y nvidia-384 --no-install-recommends +# +## ffmpeg version 4.2.4-1ubuntu0.1 Copyright (c) 2000-2020 the FFmpeg developers +##RUN apt-get install -y ffmpeg +##RUN ffmpeg -version +##RUN sleep 100 +# +## Notes on packages: +## * software-properties-common - needed for `apt-add-repository` +## * pciutils - needed for `lspci | grep -i nvidia` +#RUN apt-get install -y \ +# build-essential \ +# curl \ +# cython3 \ +# ffmpeg \ +# g++-7 \ +# gcc-7 \ +# git \ +# htop \ +# imagemagick \ +# libc++-7-dev \ +# libffi-dev \ +# libgcc-7-dev \ +# libmagic1 \ +# libsm6 \ +# libsndfile1 \ +# libssl-dev \ +# libxext6 \ +# libxrender-dev \ +# netcat \ +# pciutils \ +# python-dev \ +# python3-pip \ +# python3.8 \ +# python3.8-dev \ +# python3.8-venv\ +# redis-tools \ +# silversearcher-ag \ +# software-properties-common \ +# sox \ +# tmux \ +# vim \ +# wget +# +# +## Wav2Lip needs Python3.6 +## NB: we install python3-virtualenv as `python3.6 -m venv` won't work anymore +#RUN add-apt-repository ppa:deadsnakes/ppa +#RUN apt-get update +#RUN apt-get install -y python3.6 +#RUN apt-get install -y python3-virtualenv +#RUN apt-get install -y python3.6-venv +#RUN apt-get install -y python3.6-dev +##RUN apt-get install python3-virtualenv +# +# +#RUN python3.6 --version +#RUN python3.8 --version + +# We need ffmpeg version >= 4.2 (https://superuser.com/a/579110) +# http://ubuntuhandbook.org/index.php/2020/06/install-ffmpeg-4-3-via-ppa-ubuntu-18-04-16-04/ + +#RUN apt-get install -y software-properties-common \ +# && add-apt-repository ppa:jonathonf/ffmpeg-4 \ +# && apt-get update \ +# && apt-get install -y ffmpeg +# + +# The downloader needs youtube-dl as a python library +# NB: No longer maintained: https://news.ycombinator.com/item?id=28319624 +#RUN pip3 install --upgrade youtube-dl + +# Used by the downloader script. +# And it looks like youtube-dl is dead... +# Per the README, it looks like pip will always pull the latest version. +# https://github.com/yt-dlp/yt-dlp#update +RUN python3 -m pip install --upgrade yt-dlp + +# ==================== Python Build Step 2: Wav2Lip Requirements ==================== + +FROM pybuild-base as pybuild-requirements + +WORKDIR / + +COPY models/Wav2Lip ./models/Wav2Lip +WORKDIR models/Wav2Lip + +# NB: We need Python3.6 +RUN python3.6 -m venv python +#RUN virtualenv -p /usr/bin/python3.6 ... wait +RUN . python/bin/activate \ + && pip install --upgrade pip \ + && pip install -r requirements.txt \ + && deactivate + +WORKDIR / + +COPY models/tts ./models/tts +WORKDIR models/tts + +# NB: Not sure if we need Python3.6 for Tacotron2, but Python3.8 gave me trouble. +# NB: Setuptools fix: https://github.com/tensorflow/tensorflow/issues/34302#issuecomment-554450289 +RUN python3.6 -m venv python +RUN . python/bin/activate \ + && pip install --upgrade pip setuptools \ + && pip install -r requirements-tacotron-python36.txt \ + && deactivate + +# ==================== Rust Build Base ==================== + +# Custom base image +# Make sure to add this repository so it has read acces to the base image: +# https://github.com/orgs/storytold/packages/container/docker-base-images-rust-ssl/settings/actions_access +FROM ghcr.io/storytold/docker-base-images-rust-ssl:latest as rust-build +WORKDIR /tmp + +COPY Cargo.lock . +COPY Cargo.toml . +COPY sqlx-data.json . +COPY crates/ ./crates +COPY db/ ./db +# TODO(2022-01-16): Might not be necessary: +# COPY _migrations/ ./_migrations + +RUN $HOME/.cargo/bin/cargo fetch + +# Build all the binaries. +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin storyteller-web + +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin tts-download-job + +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin w2l-download-job + +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin tts-inference-job + +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin w2l-inference-job + +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin websocket-gateway + +RUN SQLX_OFFLINE=true \ + LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} \ + $HOME/.cargo/bin/cargo build \ + --release \ + --bin twitch-pubsub-subscriber + +# Final image +# FROM ubuntu:xenial +FROM pybuild-requirements as final +#RUN mkdir /storyteller +#WORKDIR /storyteller +WORKDIR / + + +COPY scripts/ ./scripts + +# See: https://github.com/opencontainers/image-spec/blob/master/annotations.md +LABEL org.opencontainers.image.authors='bt@brand.io, echelon@gmail.com' +LABEL org.opencontainers.image.url='https://github.com/storytold/storyteller-web' +LABEL org.opencontainers.image.documentation='https://github.com/storytold/storyteller-web' +LABEL org.opencontainers.image.source='https://github.com/storytold/storyteller-web' + +# NB: Comment this out for non-debug images +# TODO: Figure out how this is done elsewhere with just the static binaries +# Others: `mysql-client` +RUN apt-get update \ + && apt-get install -y \ + curl \ + python3-pip \ + python3 \ + wget + +RUN python3 --version +#RUN python3.7 -m venv python +#RUN . python/bin/activate \ +# && pip install --upgrade pip \ +# && pip install -r requirements.txt + + +# TODO: Use venv to do this instead. +RUN pip3 install gdown + +# Copy all the binaries. +COPY --from=rust-build /tmp/target/release/storyteller-web / +COPY --from=rust-build /tmp/target/release/tts-download-job / +COPY --from=rust-build /tmp/target/release/tts-inference-job / +COPY --from=rust-build /tmp/target/release/w2l-download-job / +COPY --from=rust-build /tmp/target/release/w2l-inference-job / +COPY --from=rust-build /tmp/target/release/websocket-gateway / +COPY --from=rust-build /tmp/target/release/twitch-pubsub-subscriber / + +# SSL certs are required for crypto +COPY --from=rust-build /etc/ssl /etc/ssl + +# Required dynamically linked libraries +COPY --from=rust-build /lib/x86_64-linux-gnu/libssl.* /lib/x86_64-linux-gnu/ +COPY --from=rust-build /lib/x86_64-linux-gnu/libcrypto.* /lib/x86_64-linux-gnu/ + +# Make sure all the links resolve +RUN ldd storyteller-web + +# Without a .env file, Rust crashes "mysteriously" (ugh) +RUN touch .env +RUN touch .env-secrets + +EXPOSE 8080 +CMD LD_LIBRARY_PATH=/usr/lib /storyteller-web + +# =================================================================================================== +# WAV2LIP WORKER + +# ==================== Rust Build Base ==================== + +# #FROM ubuntu:xenial as rustbuild-base +# WORKDIR /tmp +# RUN apt-get update \ +# && apt-get install -y \ +# build-essential \ +# curl \ +# libssl-dev \ +# libssl1.0.0 \ +# pkg-config \ +# unzip \ +# wget +# RUN curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs \ +# | sh -s -- --default-toolchain stable -y + + + +# ==================== Rust Build Binary ==================== + +# FROM rustbuild-base as rustbuild-binary +# +# COPY Cargo.lock . +# COPY Cargo.toml . +# COPY w2l_server ./w2l_server +# COPY w2l_shared ./w2l_shared +# COPY w2l_worker ./w2l_worker +# +# RUN $HOME/.cargo/bin/cargo fetch +# RUN LD_LIBRARY_PATH=/usr/lib:${LD_LIBRARY_PATH} $HOME/.cargo/bin/cargo build --release --bin wav2lip_worker + +# ==================== Final Image ==================== + +# FROM pybuild-requirements as final +# WORKDIR / +# +# COPY --from=rustbuild-binary /tmp/target/release/wav2lip_worker / +# COPY --from=rustbuild-binary /etc/ssl /etc/ssl +# COPY --from=rustbuild-binary /lib/x86_64-linux-gnu/libssl.* /lib/x86_64-linux-gnu/ +# COPY --from=rustbuild-binary /lib/x86_64-linux-gnu/libcrypto.* /lib/x86_64-linux-gnu/ +# +# # Without a .env file, Rust crashes "mysteriously" (ugh) +# RUN touch .env +# +# RUN ldd wav2lip_worker +# +# EXPOSE 8080 +# CMD LD_LIBRARY_PATH=/usr/lib /wav2lip_worker + diff --git a/README.md b/README.md new file mode 100644 index 0000000000..61f9e71907 --- /dev/null +++ b/README.md @@ -0,0 +1,75 @@ +storyteller-web +=============== + +This is the main user account monolith that we'll also bake other pieces into. + +Application overview +-------------------- + +* storyteller_web + * `storyteller-web` - HTTP + database monolith for FakeYou + * `tts-download-job` - Async download of models + * `tts-inference-job` - Async TTS inference + * `w2l-download-job` - Async download of videos and images + * `w2l-inference-job` - Async W2L inference +* obs_gateway + * `obs-gateway-server` - websocket for hosting OBS, twitch pieces (move!) + * `twitch-pubsub-job` - Subscribes to Twitch PubSub + * `twitch-chat-job` (*TODO*) - Subscribes to Twitch chat + * `reddit-chat-job` (*TODO*) - Subscribes to Reddit RPAN chat +* social + * `discord-chat-job` (*TODO*) - Subscribes to Reddit RPAN chat + * `twitter-feed-job` (*TODO*) - Subscribes to Reddit RPAN chat + +Schema and API notes +-------------------- + +* [Redis Schema Notes README](_docs/redis_schema.md) + +Local development +----------------- + +[Local development setup README](_docs/local_development_setup.md) + +Production and Deployment +------------------------- + +[Production builds and deployment README](_docs/production_builds_and_deployment.md) + +Development notes +----------------- + +[Development notes README](_docs/development_notes.md) + +Dump a MySQL report +------------------- + +``` +mysql -u storyteller \ + -pPasswordHere \ + -h IP_HERE \ + -P 3306 \ + -D storyteller \ + -B \ + -e "select * from voice_clone_requests;" | sed "s/'/\'/;s/\t/\",\"/g;s/^/\"/;s/$/\"/;s/\n//g" > voice_clone_report.csv +``` + +TODO +---- + +* Revise this README to be more useful. + +Notes / TODOs: + +* Examples for good Actix+Sqlx Tests: + https://stackoverflow.com/questions/65370752/how-do-i-create-an-actix-web-server-that-accepts-both-sqlx-database-pools-and-tr + +* Actix/sqlx runtime compat: + https://github.com/launchbadge/sqlx/issues/1117#issuecomment-801237734 + +* Redis caching + +* Jobs for analytics queries + +* Pull jobs into own repo + diff --git a/_docs/database_schema.md b/_docs/database_schema.md new file mode 100644 index 0000000000..ef0d3eaa9f --- /dev/null +++ b/_docs/database_schema.md @@ -0,0 +1,19 @@ +database schema +================ + +Twitch Model +------------ + +- `twitch_user_id` + Though numeric, this should always be a string per + [https://discuss.dev.twitch.tv/t/bug-v5-api-returns--id-as-string-for-featured-channels/10310] + and + [https://discuss.dev.twitch.tv/t/type-of-user-id-in-api-responses/10205] + +- `twitch_channel_id` + This is the same as `twitch_user_id`. They're exactly 1:1. Prefer to + use 'user_id' terminology. + +- `twitch_username` + Between 4 and 25 characters. + diff --git a/_docs/engineering_notes.md b/_docs/engineering_notes.md new file mode 100644 index 0000000000..145effc265 --- /dev/null +++ b/_docs/engineering_notes.md @@ -0,0 +1,21 @@ +development notes +================= + +Actix notes +----------- + +json request +``` +async fn handler(request: web::Json) -> impl Responder { "whatever" } +``` + +form-multipart request +``` +async fn handler(request: web::Form) -> impl Responder { "whatever" } +``` + +route parameters +``` +#[get("/{name}")] +async fn hello(name: web::Path) -> impl Responder { format!("Hi {}", name) } +``` diff --git a/_docs/local_development_setup.md b/_docs/local_development_setup.md new file mode 100644 index 0000000000..57dce2a3cc --- /dev/null +++ b/_docs/local_development_setup.md @@ -0,0 +1,207 @@ +dev setup +========= + +[Back to main README](../README.md) + +### Debugging and Fixing CUDA/PyTorch + +See this spreadsheet for previous battles with version incompatibility: +https://docs.google.com/spreadsheets/d/1BEdLmwOzo3r83-iJn9sj6co1VT92t3lIeJ-lNH25bdQ/edit#gid=0 + +Ubuntu might upgrade the driver by accident, and it might need reinstallation. Who knows. +This stuff is a nightmare. + +### Database Setup + +Install the following libraries, and see the notes further below about MySQL on Ubuntu 20.04. + +``` +mysql-server +imagemagick +``` + +To manage the database and perform migrations, install the Rust tools diesel and sqlx. + +We'll be using diesel to manage the migrations, but sqlx within the app to actually perform queries. +Diesel is an ORM, which is dumb, so we use sqlx as at-compile-time typesafe SQL. + +``` +sudo apt-get install libmysqlclient-dev +cargo install sqlx-cli --no-default-features --features mysql [2022-01-16: is this needed?] +``` + +To install a more modern version of Diesel (unreleased) that +supports the "migrations_dir" toml config option: + +``` +cargo install diesel_cli \ + --git https://github.com/diesel-rs/diesel.git \ + --rev a213fe232a122f35a812b0ce0269708a1845a4c9 \ + --no-default-features \ + --features mysql +``` + +#### Sqlx Error on Linux when Performing Schema Migrations + +You might get this error message during migration, + +``` +Encountered unknown type for Mysql: enum +thread 'main' panicked at 'internal error: entered unreachable code: Mysql only supports a closed set of types. + If you ever see this error message please open an issue at https://github.com/diesel-rs/diesel containing a dump of your schema definition.', diesel_cli/src/print_schema.rs:310:17 +note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace +``` + +I haven't found the cause (it doesn't happen on Mac), but the migrations appear to work regardless of the error message. + +#### Linux database notes + +If MySql in local dev can't be connected to, reset the accounts: + +https://linuxconfig.org/how-to-reset-root-mysql-mariadb-password-on-ubuntu-20-04-focal-fossa-linux + +#### Fixing dev MySql on Ubuntu 20.04 + +For some reason, the MySql default install on 20.04 gave me a bunch of trouble. + +In retrospect, I _think_ this is because 'root@localhost' requires sudo to access, but if this +gives any trouble in the future, here's how I got around it (two full hours of distraction!) + +``` +sudo apt-get install mysql-server + +# But for some reason the default password doesn't work and diesel/sqlx can't connect? + +# Kill everything + +sudo apt purge mysql-server mysql-client mysql-common +sudo apt autoremove +sudo mv -iv /var/lib/mysql /var/tmp/mysql-backup +sudo rm -rf /var/lib/mysql* + +sudo /usr/bin/mysql_secure_installation + +# still no work... + +sudo systemctl stop mysql.service + +# ugh, it wasn't chowned or created (both states observed in different installs)… +sudo mkdir -p /var/run/mysqld +sudo chown mysql:mysql /var/run/mysqld + +sudo mysqld_safe --skip-grant-tables --skip-networking & + +mysql -u root + +flush privileges; +USE mysql; +ALTER USER 'root'@'localhost' IDENTIFIED BY 'root'; + +sudo killall -u mysql +sudo systemctl restart mysql.service + +# now it works +sudo mysql -u root -p + +# but now I have to use "sudo" !?!? + +use mysql; + +CREATE USER 'storyteller'@'localhost' IDENTIFIED BY 'password'; +GRANT ALL PRIVILEGES ON storyteller.* TO 'storyteller'@'localhost'; +``` + +### Install Diesel (migrations only) + +We use Diesel to manage migrations, but we don't use it in server code. +Actual server code uses SQLx. To install the CLI tool, run the following: + +``` +cargo install diesel_cli +``` + +Mac [has issues](https://github.com/diesel-rs/diesel/issues/2605) and requires a few dependencies: + +``` +brew install libpq +``` + +### Database migrations + +To reset the entire database (drop, migrate), run: + +``` +diesel database reset +``` + +To migrate at the current step and beyond: + +``` +diesel migration run +``` + +To undo migrations at the current step: + +``` +diesel migration redo +``` + +### Server Query Codegen + +We use SQLx instead of Diesel in the production server. It's typesafe +SQL instead of an ORM like Diesel. + +SQLx connects to a database to derive type information, but obviously +cannot do this for builds in CI. In order to cache the types, we build +and check in a cache file (necessary for builds): + +``` +SQLX_OFFLINE=true cargo sqlx prepare +``` + +Now that we have multiple binaries, it's required to include all the queries in the main +binary so we can generate the cached queries as a single target. That's then executed +with: + +``` +SQLX_OFFLINE=true cargo sqlx prepare -- --bin storyteller-web +``` + +### Setting Up a Local Nginx Proxy + +Set up a local nginx to proxy to the frontend and backend so cookie issues aren't annoying + +Configure Nginx per the checked in Nginx configs (and instructions) in `localdev/nginx-http-config`. + +``` + +And in /etc/hosts, + +``` +127.0.0.1 jungle.horse +127.0.0.1 api.jungle.horse + +``` + +### Python 3.6 on Apple M1 Mac + +Python3.6 isn't supported on Apple silicon, and it's not in homebrew. It can be installed with +[nix using Rosetta](https://stackoverflow.com/a/65980989): + +Download: https://nixos.org/download.html#nix-quick-install + +``` +nix run nixpkgs.python36 -c python +``` + +Install venv: + +``` +nix run nixpkgs.python36 -c python -m venv python +``` + +Install other packages on Mac that aren't used in venv: + +``` +python3 -m pip install --user requests gdown youtube_dl +``` diff --git a/_docs/production_builds_and_deployment.md b/_docs/production_builds_and_deployment.md new file mode 100644 index 0000000000..855832be1f --- /dev/null +++ b/_docs/production_builds_and_deployment.md @@ -0,0 +1,36 @@ +production builds and deployment +================================ + +### Github Docker builds + +The repository needs to be given read access to the base docker image: + +https://github.com/orgs/storytold/packages/container/docker-base-images-rust-ssl/settings + +These instructions assume running on GCP. + +### Database migrations + +1. Set `DATABASE_URL` in `.env` to the production secrets (DO NOT COMMIT!) +2. Run `diesel migration run` + +### Setting up public buckets without list permission + +Public buckets that deny the `list` action should use the following Role: + +`roles/storage.legacyObjectReader` + +See: + +* https://stackoverflow.com/a/56354633 +* https://cloud.google.com/storage/docs/access-control/making-data-public#buckets + +### Generating Bucket Access Key and Secret Key + +https://cloud.telestream.net/tutorials/how-to-setting-up-google-cloud-storage/ + +1. Go to the GCS page +2. Click "settings" +3. Click "interoperability" tab +4. (enable interoperable access if not already set) +5. Click "create new key" diff --git a/_docs/redis_schema.md b/_docs/redis_schema.md new file mode 100644 index 0000000000..4b3ed5420f --- /dev/null +++ b/_docs/redis_schema.md @@ -0,0 +1,27 @@ +redis schema +============ + +Database 0: TTS / RateLimiter +----------------------------- + +TODO: Document rate limiter + +Public statistics on usage displayed on the website: + +* `ttsUseCount:{model_token}`: INCR, GET +* `w2lUseCount:{template_token}`: INCR, GET + +Temporary job progress reports to the user: + +These are written as a single string, such as "downloading model", "uploading model", +"generating preview", "executing", etc. + +* `ttsDownloadExtraStatus:{inference_job_token}`: SETEX, GET +* `ttsInferenceExtraStatus:{inference_job_token}`: SETEX, GET +* `w2lDownloadExtraStatus:{inference_job_token}`: SETEX, GET +* `w2lInferenceExtraStatus:{inference_job_token}`: SETEX, GET + +Database 1: Obs / Twitch +------------------------ + +* WIP diff --git a/_migrations/.gitkeep b/_migrations/.gitkeep new file mode 100644 index 0000000000..e69de29bb2 diff --git a/_migrations/2021-05-15-071534_users/down.sql b/_migrations/2021-05-15-071534_users/down.sql new file mode 100644 index 0000000000..262c9cf4d3 --- /dev/null +++ b/_migrations/2021-05-15-071534_users/down.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table user_roles; +drop table user_sessions; +drop table users; diff --git a/_migrations/2021-05-15-071534_users/up.sql b/_migrations/2021-05-15-071534_users/up.sql new file mode 100644 index 0000000000..fb511b3835 --- /dev/null +++ b/_migrations/2021-05-15-071534_users/up.sql @@ -0,0 +1,235 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE users ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Visible "primary key" + token VARCHAR(32) NOT NULL, + + -- Username is a lookup key; display_name allows the user to add custom case. + username VARCHAR(20) NOT NULL, + display_name VARCHAR(20) NOT NULL, + + email_address VARCHAR(255) NOT NULL, + email_confirmed BOOLEAN NOT NULL DEFAULT false, + + -- The role assigned to the user confers permissions. + user_role_slug VARCHAR(16) NOT NULL, + + -- ========== CREDENTIALS ========== + + -- Bcrypt password hash. Granted, there are newer methods: + -- https://cheatsheetseries.owasp.org/cheatsheets/Password_Storage_Cheat_Sheet.html + -- Passwords may be a max of 64 characters per BCrypt. + password_hash BINARY(60) NOT NULL, + + -- Incremented with every update to the password. + password_version INT NOT NULL DEFAULT 0, + + -- ========== ABUSE TRACKING ========== + + -- For abuse tracking. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + ip_address_last_login VARCHAR(40) NOT NULL, + ip_address_last_update VARCHAR(40) NOT NULL, + + -- ========== DISPLAY / PROFILE ========== + + -- Gravatar image hashes are precomputed. + email_gravatar_hash CHAR(32) NOT NULL, + + -- The profile of the user in markdown (user editable). + profile_markdown TEXT NOT NULL, + + -- Generated HTML (not user-editable). + profile_rendered_html TEXT NOT NULL, + + -- An uploaded avatar. Public hash in our bucket. + -- avatar_public_bucket_hash CHAR(32) DEFAULT NULL, + + -- ========== SOCIAL MEDIA ========== + + -- Social media usernames + -- These are not confirmed. We'll need to build an OAuth system to handle that. + discord_username VARCHAR(36) DEFAULT NULL, + twitter_username VARCHAR(36) DEFAULT NULL, + twitch_username VARCHAR(36) DEFAULT NULL, + patreon_username VARCHAR(36) DEFAULT NULL, + github_username VARCHAR(36) DEFAULT NULL, + cashapp_username VARCHAR(36) DEFAULT NULL, + website_url VARCHAR(255) DEFAULT NULL, + + -- ========== USER PREFERENCES ========== + + -- If the user doesn't want to use gravatar and doesn't have an uploaded avatar. + disable_gravatar BOOLEAN NOT NULL DEFAULT false, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + -- Visibility preference used at TTS inference time (which can be changed on the result itself). + -- Moderators will still see them. + preferred_tts_result_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + -- Visibility preference used at TTS inference time (which can be changed on the result itself). + -- Moderators will still see them. + preferred_w2l_result_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- Auto play preferences + -- Tri-state because we want to potentially change the default behavior. + auto_play_audio_preference BOOLEAN DEFAULT NULL, + auto_play_video_preference BOOLEAN DEFAULT NULL, + + -- Favorite TTS voice to use by default + -- maybe_preferred_tts_model_token VARCHAR(32) DEFAULT NULL, + + -- Favorite W2L model to use by default + -- maybe_preferred_w2l_template_token VARCHAR(32) DEFAULT NULL, + + -- Settings + -- DO NOT REORDER. + -- dark_mode_preference ENUM( + -- 'light-mode', + -- 'dark-mode', + -- 'use-clock' + -- ) NOT NULL DEFAULT 'light-mode', + + -- ========== STATS ========== + + -- For tracking stats. + -- The "cached" values are updated by a background job. + -- cached_tts_rendered_counter INT(10) NOT NULL DEFAULT 0, + -- cached_w2l_rendered_counter INT(10) NOT NULL DEFAULT 0, + + -- ========== MODERATION DETAILS ========== + + -- Different than deleted. + -- Users still show up, but can't do anything. + is_banned BOOLEAN NOT NULL DEFAULT false, + + -- If a moderator has comments. + maybe_mod_comments VARCHAR(255) DEFAULT NULL, + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If the user is deleted, we set this. + -- This is different than banned. These users won't show up at all. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (username), + UNIQUE KEY (email_address), + KEY fk_user_role_slug (user_role_slug) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +CREATE TABLE user_roles ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" + slug CHAR(16) NOT NULL, + + name VARCHAR(255) NOT NULL, + + -- Usage + can_use_tts BOOLEAN NOT NULL DEFAULT FALSE, + can_use_w2l BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_own_tts_results BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_own_w2l_results BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_own_account BOOLEAN NOT NULL DEFAULT FALSE, + + -- Contribution + can_upload_tts_models BOOLEAN NOT NULL DEFAULT FALSE, + can_upload_w2l_templates BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_own_tts_models BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_own_w2l_templates BOOLEAN NOT NULL DEFAULT FALSE, + + -- Moderation + can_approve_w2l_templates BOOLEAN NOT NULL DEFAULT FALSE, + -- Edit includes "lock_from_use" and "lock_from_modification" + can_edit_other_users_profiles BOOLEAN NOT NULL DEFAULT FALSE, + can_edit_other_users_tts_models BOOLEAN NOT NULL DEFAULT FALSE, + can_edit_other_users_w2l_templates BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_other_users_tts_models BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_other_users_tts_results BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_other_users_w2l_templates BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_other_users_w2l_results BOOLEAN NOT NULL DEFAULT FALSE, + can_ban_users BOOLEAN NOT NULL DEFAULT FALSE, + can_delete_users BOOLEAN NOT NULL DEFAULT FALSE, + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (slug) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +CREATE TABLE user_sessions ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Session entropy + token VARCHAR(32) NOT NULL, + + -- Foreign key to user + user_token VARCHAR(32) NOT NULL, + + -- Track each session's creation IP. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Session termination time. + -- This must be set by the server code, or the session is invalid + expires_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + -- deletion = session termination + -- Typically these are destroyed by users, but if in the future we allow mods to + -- delete them, it doesn't really matter who did the deletion: sessions are not + -- designed to be recoverable. + deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_user_token (user_token), + KEY index_ip_address_creation (ip_address_creation) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-05-15-071545_asset_upload_jobs/down.sql b/_migrations/2021-05-15-071545_asset_upload_jobs/down.sql new file mode 100644 index 0000000000..117aee63d0 --- /dev/null +++ b/_migrations/2021-05-15-071545_asset_upload_jobs/down.sql @@ -0,0 +1,5 @@ +# noinspection SqlNoDataSourceInspectionForFile +# noinspection SqlResolveForFile + +drop table tts_model_upload_jobs; +drop table w2l_template_upload_jobs; diff --git a/_migrations/2021-05-15-071545_asset_upload_jobs/up.sql b/_migrations/2021-05-15-071545_asset_upload_jobs/up.sql new file mode 100644 index 0000000000..5a377f2be5 --- /dev/null +++ b/_migrations/2021-05-15-071545_asset_upload_jobs/up.sql @@ -0,0 +1,231 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- NB: See the "squashed" copy of this for better documentation +-- and the current state of the table. + +-- TTS MODELS +CREATE TABLE tts_model_upload_jobs ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== UPLOAD DETAILS ========== + + -- The name of the voice + -- The "name" of the voice model, which might be complicated. + -- If maybe_subject_token (etc.) is set, then it's authoritative instead. + title VARCHAR(255) NOT NULL, + + -- NB: THIS ENUM IS OUT OF DATE! + -- Check the squashed sql file as an updated reference! + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + tts_model_type ENUM( + 'not-set', + 'tacotron2', + 'glowtts', + 'glowtts-vocodes' + ) NOT NULL DEFAULT 'not-set', + + -- Can be linked to a well-known voice + -- maybe_subject_token VARCHAR(32) DEFAULT NULL, + -- maybe_actor_subject_token VARCHAR(32) DEFAULT NULL, + + -- If we need to download the file from Google Drive. + download_url VARCHAR(512) DEFAULT NULL, + + -- NB: DO NOT SORT! + -- download_url_type ENUM( + -- 'google-drive', + -- 'web' + -- ) NOT NULL DEFAULT 'web', + + -- ========== CREATOR DETAILS AND PREFERENCES ========== + + -- Foreign key to user + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_success_result_token (on_success_result_token), + KEY fk_creator_user_token (creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +-- WAV2LIP TEMPLATES +CREATE TABLE w2l_template_upload_jobs ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== UPLOAD DETAILS ========== + + -- The title of the template + title VARCHAR(255) NOT NULL, + + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + template_type ENUM( + 'unknown', + 'video', + 'image' + ) NOT NULL DEFAULT 'unknown', + + -- Can be linked to a well-known subject + -- maybe_subject_token VARCHAR(32) DEFAULT NULL, + -- maybe_actor_subject_token VARCHAR(32) DEFAULT NULL, + + -- If we need to download the file from Google Drive. + download_url VARCHAR(512) DEFAULT NULL, + + -- NB: DO NOT SORT! + -- download_url_type ENUM( + -- 'google-drive', + -- 'web' + -- ) NOT NULL DEFAULT 'web', + + -- ========== CREATOR DETAILS AND PREFERENCES ========== + + -- Foreign key to user + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_success_result_token (on_success_result_token), + KEY fk_creator_user_token (creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-05-15-071601_assets/down.sql b/_migrations/2021-05-15-071601_assets/down.sql new file mode 100644 index 0000000000..872ae266d7 --- /dev/null +++ b/_migrations/2021-05-15-071601_assets/down.sql @@ -0,0 +1,5 @@ +# noinspection SqlNoDataSourceInspectionForFile +# noinspection SqlResolveForFile + +drop table tts_models; +drop table w2l_templates; \ No newline at end of file diff --git a/_migrations/2021-05-15-071601_assets/up.sql b/_migrations/2021-05-15-071601_assets/up.sql new file mode 100644 index 0000000000..6a9a8ee855 --- /dev/null +++ b/_migrations/2021-05-15-071601_assets/up.sql @@ -0,0 +1,328 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- NB: See the "squashed" copy of this for better documentation +-- and the current state of the table. + +CREATE TABLE tts_models ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- NB: THIS ENUM IS OUT OF DATE! + -- Check the squashed sql file as an updated reference! + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + tts_model_type ENUM( + 'not-set', + 'tacotron2', + 'glowtts', + 'glowtts-vocodes' + ) NOT NULL DEFAULT 'not-set', + + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- How text should be handled. + text_preprocessing_algorithm ENUM( + 'basic', + 'arpabet' + ) NOT NULL DEFAULT 'basic', + + -- NB: ADDED BY ALTER + -- maybe_default_pretrained_vocoder VARCHAR(64) DEFAULT NULL, + + -- Optional Pointer to a newer version of the voice + -- If there's a newer version, we can disable this one. + -- maybe_updated_model_token VARCHAR(32) DEFAULT NULL, + + -- We can set as a combination of ['username' + 'voice-name'] + -- There can be public aliases for voices, eg. a voice's default model. + -- A user can change this. + -- As such, these should not be foreign keys. + -- updatable_slug VARCHAR(64) NOT NULL, + + -- Can be linked to a well-known voice + -- maybe_subject_token VARCHAR(32) DEFAULT NULL, + -- maybe_actor_subject_token VARCHAR(32) DEFAULT NULL, + + -- ========== UNSTRUCTURED METADATA ========== + + -- The "name" of the voice model, which might be complicated. + -- If maybe_subject_token (etc.) is set, then it's authoritative instead. + title VARCHAR(255) NOT NULL, + + -- The description of the model in markdown. + description_markdown TEXT NOT NULL, + + -- Generated HTML (not user-editable). + description_rendered_html TEXT NOT NULL, + + -- eg. Bill Clinton, Batman, or Goku. + voice_name VARCHAR(255) DEFAULT NULL, + + -- Optional; eg. Sean Schemmel + actor_name VARCHAR(255) DEFAULT NULL, + + -- Optional; eg. if the voice is "happy" or a singer "a-capella", etc. + characteristic VARCHAR(255) DEFAULT NULL, + + -- Optional; eg. if you train the model more than once. + version_string VARCHAR(255) DEFAULT NULL, + + -- ========== IMMUTABLE PROVENANCE AND METADATA ========== + + -- Where the file was originally downloaded (if it was downloaded) + original_download_url VARCHAR(512) DEFAULT NULL, + + -- The filename that was used at upload time (if available) + original_filename VARCHAR(255) DEFAULT NULL, + + -- File characteristics + file_size_bytes INT(10) NOT NULL DEFAULT 0, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address_creation VARCHAR(40) NOT NULL, + creator_ip_address_last_update VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== BUCKET STORAGE ========== + + -- The pytorch model + -- For now, this will be a hash of the file contents. + -- NB: NOT UNIQUE! We can allow duplicate uploads. + private_bucket_hash CHAR(64) NOT NULL, + -- The "full url" version of the path + private_bucket_object_name VARCHAR(255) NOT NULL, + + -- ========== RATING AND RANKING ========== + + -- Calculated average, on a scale of 0-100 + -- Null with zero ratings. + -- calculated_average_score INT(3) DEFAULT NULL, + + -- Other metrics + -- calculated_total_ratings_submitted_count INT(10) NOT NULL DEFAULT 0, + -- calculated_total_uses_count BIGINT(10) NOT NULL DEFAULT 0, + + -- ========== MODERATION DETAILS ========== + + -- Mods have to approve of tts models for them to show up in a public index. + -- Models can be used by the author (or other parties that know the URL) before + -- they're approved, but unapproved models won't show up in public indices. + -- This is independent of other visibility controls (is_mod_disabled, deleted_at, + -- and creator_set_visibility) + is_public_listing_approved BOOLEAN DEFAULT NULL, + + -- Extremely popular models may be locked from deletion or modification by users. + is_locked_from_user_modification BOOLEAN NOT NULL DEFAULT FALSE, + + -- In this case, a moderator disables it from being used for inference. + -- This also disables it for the creator. + -- Unlike deletion, it remains "visible" to those that have access. + is_locked_from_use BOOLEAN NOT NULL DEFAULT FALSE, + + -- If a moderator has comments. + maybe_mod_comments VARCHAR(255) DEFAULT NULL, + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod. + -- It completely disappears from the system. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_creator_user_token (creator_user_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_creator_ip_address_creation (creator_ip_address_creation), + KEY index_creator_ip_address_last_update (creator_ip_address_last_update), + KEY index_creator_set_visibility (creator_set_visibility), + KEY index_private_bucket_hash (private_bucket_hash), + KEY index_is_public_listing_approved (is_public_listing_approved), + KEY index_is_locked_from_user_modification (is_locked_from_user_modification), + KEY index_is_locked_from_use (is_locked_from_use) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +CREATE TABLE w2l_templates ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + template_type ENUM( + 'not-set', + 'image', + 'video' + ) NOT NULL DEFAULT 'not-set', + + -- Optional Pointer to a newer version of the template + -- If there's a newer version, we can disable this one. + -- maybe_updated_template_token VARCHAR(32) DEFAULT NULL, + + -- A combination of ['username' + 'template-name'] + -- There can be public aliases for voices, eg. a voice's default model. + -- A user can change this. + -- As such, these should not be foreign keys. + -- updatable_slug VARCHAR(64) NOT NULL, + + -- Can be linked to a well-known voice + -- maybe_subject_token VARCHAR(32) DEFAULT NULL, + -- maybe_actor_subject_token VARCHAR(32) DEFAULT NULL, + + -- ========== UNSTRUCTURED METADATA ========== + + -- The title of the template. + title CHAR(255) NOT NULL, + + -- The description of the template in markdown. + description_markdown TEXT NOT NULL, + + -- Generated HTML (not user-editable). + description_rendered_html TEXT NOT NULL, + + -- ========== IMMUTABLE PROVENANCE AND METADATA ========== + + -- Where the file was originally downloaded (if it was downloaded) + original_download_url VARCHAR(512) DEFAULT NULL, + + -- The filename that was used at upload time (if available) + original_filename CHAR(255) DEFAULT NULL, + + -- Dimensions and other characteristics + file_size_bytes INT(10) NOT NULL DEFAULT 0, + mime_type VARCHAR(32) NOT NULL DEFAULT '', + frame_width INT(10) NOT NULL DEFAULT 0, + frame_height INT(10) NOT NULL DEFAULT 0, + frame_count INT(10) NOT NULL DEFAULT 0, + fps FLOAT NOT NULL DEFAULT 0.0, + duration_millis INT(10) NOT NULL DEFAULT 0, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address_creation VARCHAR(40) NOT NULL, + creator_ip_address_last_update VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== BUCKET STORAGE ========== + + -- The hash of the original source image/video + private_bucket_hash CHAR(64) NOT NULL, + + -- The "full path" version of the path. Contains the hash. + private_bucket_object_name VARCHAR(255) NOT NULL, + -- The "full path" of the cached faces object. Contains the hash. + private_bucket_cached_faces_object_name VARCHAR(255) NOT NULL, + + -- The "full url" of the preview image and/or video. Contains the hash. + maybe_public_bucket_preview_image_object_name VARCHAR(255) DEFAULT NULL, + maybe_public_bucket_preview_video_object_name VARCHAR(255) DEFAULT NULL, + + -- ========== RATING AND RANKING ========== + + -- Calculated average, on a scale of 0-100 + -- Null with zero ratings. + -- calculated_average_score INT(3) DEFAULT NULL, + + -- Other metrics + -- calculated_total_ratings_submitted_count INT(10) NOT NULL DEFAULT 0, + -- calculated_total_uses_count BIGINT(10) NOT NULL DEFAULT 0, + + -- ========== MODERATION DETAILS ========== + + -- Mods have to approve of w2l templates for them to show up in a public index. + -- Templates can be used by the author (or other parties that know the URL) before + -- they're approved, but unapproved templates won't show up in public indices. + -- This is independent of other visibility controls (is_mod_disabled, deleted_at, + -- and creator_set_visibility) + is_public_listing_approved BOOLEAN DEFAULT NULL, + + -- Extremely popular models may be locked from deletion or modification by users. + is_locked_from_user_modification BOOLEAN NOT NULL DEFAULT FALSE, + + -- In this case, a moderator disables it from being used for inference. + -- This also disables it for the creator. + -- Unlike deletion, it remains "visible" to those that have access. + is_locked_from_use BOOLEAN NOT NULL DEFAULT FALSE, + + -- If a moderator has comments. + maybe_mod_comments VARCHAR(255) DEFAULT NULL, + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod. + -- It shows up nowhere if deleted. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_creator_user_token (creator_user_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_template_type (template_type), + KEY index_creator_ip_address_creation (creator_ip_address_creation), + KEY index_creator_ip_address_last_update (creator_ip_address_last_update), + KEY index_creator_set_visibility (creator_set_visibility), + KEY index_private_bucket_hash (private_bucket_hash), + KEY index_is_public_listing_approved (is_public_listing_approved), + KEY index_is_locked_from_user_modification (is_locked_from_user_modification), + KEY index_is_locked_from_use (is_locked_from_use) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-05-15-212913_inference_jobs/down.sql b/_migrations/2021-05-15-212913_inference_jobs/down.sql new file mode 100644 index 0000000000..cad1d8e8ff --- /dev/null +++ b/_migrations/2021-05-15-212913_inference_jobs/down.sql @@ -0,0 +1,5 @@ +# noinspection SqlResolveForFile +# noinspection SqlNoDataSourceInspectionForFile + +drop table tts_inference_jobs; +drop table w2l_inference_jobs; diff --git a/_migrations/2021-05-15-212913_inference_jobs/up.sql b/_migrations/2021-05-15-212913_inference_jobs/up.sql new file mode 100644 index 0000000000..28a68313ad --- /dev/null +++ b/_migrations/2021-05-15-212913_inference_jobs/up.sql @@ -0,0 +1,229 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- NB: See the "squashed" copy of this for better documentation +-- and the current state of the table. + +CREATE TABLE tts_inference_jobs ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token from client + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== INFERENCE DETAILS ========== + + -- The model to use. + -- This also determines which architecture we're using. + model_token VARCHAR(32) NOT NULL, + + -- NB: ADDED BY ALTER + -- maybe_override_pretrained_vocoder VARCHAR(64) DEFAULT NULL, + + -- The raw, unprocessed user input. + raw_inference_text TEXT NOT NULL, + + -- ========== CREATOR DETAILS ========== + + -- Foreign key to user + -- If no user is logged in, this is null. + maybe_creator_user_token VARCHAR(32) DEFAULT NULL, + + -- Based on a cookie sent by the frontend. + -- We'll save this even if the user is logged in. + maybe_creator_anonymous_visitor_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_success_result_token (on_success_result_token), + KEY fk_model_token (model_token), + KEY fk_maybe_creator_user_token (maybe_creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +CREATE TABLE w2l_inference_jobs ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token from client + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== INFERENCE DETAILS : FACE TEMPLATE ========== + + -- The W2L template to use + -- Can be an image or video. + -- This is null if we're using a custom uploaded image. + maybe_w2l_template_token VARCHAR(32) DEFAULT NULL, + + -- If we're using a custom uploaded image, this will be present. + -- NOTE: This may never be supported. + maybe_public_image_bucket_location VARCHAR(255) DEFAULT NULL, + + -- ========== INFERENCE DETAILS : AUDIO SOURCE ========== + + -- If we're using TTS results, this will be present + maybe_tts_inference_result_token VARCHAR(32) DEFAULT NULL, + + -- If we're using custom uploaded audio, this will be present. + maybe_public_audio_bucket_hash VARCHAR(64) DEFAULT NULL, + maybe_public_audio_bucket_location VARCHAR(255) DEFAULT NULL, + + -- The filename that was used at upload time (if available) + maybe_original_audio_filename CHAR(255) DEFAULT NULL, + + -- Where the audio file was originally downloaded (if it was downloaded) + maybe_original_audio_download_url VARCHAR(512) DEFAULT NULL, + + maybe_audio_mime_type VARCHAR(32) DEFAULT NULL, + + -- ========== CREATOR DETAILS AND PREFERENCES ========== + + -- Foreign key to user + -- If no user is logged in, this is null. + maybe_creator_user_token VARCHAR(32) DEFAULT NULL, + + -- Based on a cookie sent by the frontend. + -- We'll save this even if the user is logged in. + maybe_creator_anonymous_visitor_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- Flags to disable branding + disable_end_bump BOOL NOT NULL DEFAULT false, + disable_watermark BOOL NOT NULL DEFAULT false, + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_success_result_token (on_success_result_token), + KEY fk_maybe_w2l_template_token (maybe_w2l_template_token), + KEY fk_maybe_tts_inference_result_token (maybe_tts_inference_result_token), + KEY fk_maybe_creator_user_token (maybe_creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-05-15-212918_inference_results/down.sql b/_migrations/2021-05-15-212918_inference_results/down.sql new file mode 100644 index 0000000000..261d58e669 --- /dev/null +++ b/_migrations/2021-05-15-212918_inference_results/down.sql @@ -0,0 +1,5 @@ +# noinspection SqlResolveForFile +# noinspection SqlNoDataSourceInspectionForFile + +drop table tts_results; +drop table w2l_results; diff --git a/_migrations/2021-05-15-212918_inference_results/up.sql b/_migrations/2021-05-15-212918_inference_results/up.sql new file mode 100644 index 0000000000..f6aa7ab2ae --- /dev/null +++ b/_migrations/2021-05-15-212918_inference_results/up.sql @@ -0,0 +1,184 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- NB: See the "squashed" copy of this for better documentation +-- and the current state of the table. + +CREATE TABLE tts_results ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- ========== INFERENCE DETAILS ========== + + -- The TTS model that was used + model_token VARCHAR(32) NOT NULL, + + -- NB: ADDED BY ALTER + -- maybe_pretrained_vocoder_used VARCHAR(64) DEFAULT NULL, + + -- The original raw, unprocessed user input. + raw_inference_text TEXT NOT NULL, + + -- SHA2 hash of the text [SHA2 = CHAR(64), SHA1 = CHAR(40), MD5 = CHAR(32)] + raw_inference_text_hash_sha2 CHAR(64) NOT NULL, + + -- First pass at text normalization. + -- eg. 14th -> fourteenth, etc. + normalized_inference_text TEXT NOT NULL, + + -- If the model uses arpabet, we'll save the arpabet we used. + maybe_arpabet_text TEXT DEFAULT NULL, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + -- If the user wasn't logged in, this is null + maybe_creator_user_token VARCHAR(32) DEFAULT NULL, + + -- Based on a cookie sent by the frontend. + -- We'll save this even if the user is logged in. + maybe_creator_anonymous_visitor_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- The synthetic id associated with this result. + -- These ids are incremented on a per-user basis to help users + -- sequence their own work. They serve no other purpose. + maybe_creator_synthetic_id BIGINT(20) DEFAULT NULL, + + -- ========== BUCKET STORAGE ========== + + -- Where the wav, spectrogram, and etc. are located. + public_bucket_wav_audio_path VARCHAR(255) NOT NULL, + public_bucket_spectrogram_path VARCHAR(255) NOT NULL, + + -- ========== METADATA ========== + + file_size_bytes INT(10) NOT NULL DEFAULT 0, + duration_millis INT(10) NOT NULL DEFAULT 0, + + -- ========== MODERATION DETAILS ========== + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod or the creator. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_model_token (model_token), + KEY fk_maybe_creator_user_token (maybe_creator_user_token), + KEY fk_maybe_creator_anonymous_visitor_token (maybe_creator_anonymous_visitor_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_creator_ip_address (creator_ip_address), + KEY index_user_deleted_at(user_deleted_at), + KEY index_mod_deleted_at(mod_deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +CREATE TABLE w2l_results ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- ========== INFERENCE DETAILS ========== + + -- The w2l template that was used (if set) + maybe_w2l_template_token VARCHAR(32) NOT NULL, + + -- The inference result, if we're using them. + maybe_tts_inference_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + -- If the user wasn't logged in, this is null + maybe_creator_user_token VARCHAR(32) DEFAULT NULL, + + -- Based on a cookie sent by the frontend. + -- We'll save this even if the user is logged in. + maybe_creator_anonymous_visitor_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- The synthetic id associated with this result. + -- These ids are incremented on a per-user basis to help users + -- sequence their own work. They serve no other purpose. + maybe_creator_synthetic_id BIGINT(20) DEFAULT NULL, + + -- ========== BUCKET STORAGE ========== + + -- Where the wav, spectrogram, and etc. are located. + public_bucket_video_path VARCHAR(255) NOT NULL, + + -- ========== METADATA ========== + + file_size_bytes INT(10) NOT NULL DEFAULT 0, + mime_type VARCHAR(32) NOT NULL DEFAULT '', + frame_width INT(10) NOT NULL DEFAULT 0, + frame_height INT(10) NOT NULL DEFAULT 0, + duration_millis INT(10) NOT NULL DEFAULT 0, + + -- ========== MODERATION DETAILS ========== + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod or the creator. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_maybe_creator_user_token (maybe_creator_user_token), + KEY fk_maybe_creator_anonymous_visitor_token (maybe_creator_anonymous_visitor_token), + KEY fk_maybe_w2l_template_token (maybe_w2l_template_token), + KEY fk_maybe_tts_inference_result_token (maybe_tts_inference_result_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_creator_ip_address (creator_ip_address), + KEY index_user_deleted_at(user_deleted_at), + KEY index_mod_deleted_at(mod_deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-05-15-212920_badges/down.sql b/_migrations/2021-05-15-212920_badges/down.sql new file mode 100644 index 0000000000..97075e079a --- /dev/null +++ b/_migrations/2021-05-15-212920_badges/down.sql @@ -0,0 +1,6 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table badges; +drop table user_badges; diff --git a/_migrations/2021-05-15-212920_badges/up.sql b/_migrations/2021-05-15-212920_badges/up.sql new file mode 100644 index 0000000000..0b6e499cea --- /dev/null +++ b/_migrations/2021-05-15-212920_badges/up.sql @@ -0,0 +1,42 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Badges are rewards for loyalty +CREATE TABLE badges ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" + slug VARCHAR(32) NOT NULL, + + -- Description + title VARCHAR(255) NOT NULL, + description VARCHAR(512) NOT NULL, + image_url VARCHAR(512) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (slug) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +-- Join table +CREATE TABLE user_badges ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + user_token VARCHAR(32) NOT NULL, + badge_slug VARCHAR(32) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + KEY fk_user_token (user_token), + KEY fk_badge_slug (badge_slug) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-05-30-010200_firehose/down.sql b/_migrations/2021-05-30-010200_firehose/down.sql new file mode 100644 index 0000000000..6bc740efcd --- /dev/null +++ b/_migrations/2021-05-30-010200_firehose/down.sql @@ -0,0 +1,4 @@ +# noinspection SqlResolveForFile +# noinspection SqlNoDataSourceInspectionForFile + +drop table firehose_entries; diff --git a/_migrations/2021-05-30-010200_firehose/up.sql b/_migrations/2021-05-30-010200_firehose/up.sql new file mode 100644 index 0000000000..40302a1e21 --- /dev/null +++ b/_migrations/2021-05-30-010200_firehose/up.sql @@ -0,0 +1,37 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE firehose_entries ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Visible "primary key" + token VARCHAR(32) NOT NULL, + + -- The type of the event + event_type VARCHAR(32) NOT NULL, + + -- The target user + maybe_target_user_token VARCHAR(32) DEFAULT NULL, + + -- The target "entity", which varies by event type. + -- eg. a model token or template token + maybe_target_entity_token VARCHAR(32) DEFAULT NULL, + + -- If something was created, this might be the token (if we know it) + -- eg. an inference job's token + maybe_created_entity_token VARCHAR(32) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_maybe_target_user_token (maybe_target_user_token), + KEY fk_maybe_target_entity_token (maybe_target_entity_token), + KEY fk_maybe_created_entity_token (maybe_created_entity_token), + KEY index_event_type (event_type) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-06-09-010200_ip_bans/down.sql b/_migrations/2021-06-09-010200_ip_bans/down.sql new file mode 100644 index 0000000000..05e576cc5d --- /dev/null +++ b/_migrations/2021-06-09-010200_ip_bans/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table ip_address_bans; diff --git a/_migrations/2021-06-09-010200_ip_bans/up.sql b/_migrations/2021-06-09-010200_ip_bans/up.sql new file mode 100644 index 0000000000..6c71635b98 --- /dev/null +++ b/_migrations/2021-06-09-010200_ip_bans/up.sql @@ -0,0 +1,37 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE ip_address_bans ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- The IP address to ban + ip_address VARCHAR(40) NOT NULL, + + -- Maybe the target user. + maybe_target_user_token VARCHAR(32) DEFAULT NULL, + + -- Mod who created or last edited the ban. + mod_user_token VARCHAR(32) NOT NULL, + + -- Optional notes about the ban. + mod_notes TEXT NOT NULL, + + -- Moderator can set an expire time (optional). + -- If not set, the ban is indefinite + expires_at TIMESTAMP DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If deleted, the ban is removed. + deleted_at TIMESTAMP DEFAULT NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (ip_address), + KEY fk_mod_user_token (mod_user_token), + KEY fk_maybe_target_user_token (maybe_target_user_token) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-06-09-010300_mod_audit_log/down.sql b/_migrations/2021-06-09-010300_mod_audit_log/down.sql new file mode 100644 index 0000000000..72523365cb --- /dev/null +++ b/_migrations/2021-06-09-010300_mod_audit_log/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table moderator_audit_log_entries; \ No newline at end of file diff --git a/_migrations/2021-06-09-010300_mod_audit_log/up.sql b/_migrations/2021-06-09-010300_mod_audit_log/up.sql new file mode 100644 index 0000000000..a131f54031 --- /dev/null +++ b/_migrations/2021-06-09-010300_mod_audit_log/up.sql @@ -0,0 +1,38 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE moderator_audit_log_entries ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Visible "primary key" + token VARCHAR(32) NOT NULL, + + -- The type of the event (ban, delete, etc.) + event_type VARCHAR(32) NOT NULL, + + -- The target user + maybe_target_user_token VARCHAR(32) DEFAULT NULL, + + -- The target "entity", which varies by event type. + -- eg. a model token or template token + maybe_target_entity_token VARCHAR(32) DEFAULT NULL, + + -- If an IP was banned, we record it here. + maybe_target_ip_address VARCHAR(40) DEFAULT NULL, + + -- Optional mod-specified notes. + optional_notes TEXT DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_maybe_target_user_token (maybe_target_user_token), + KEY fk_maybe_target_entity_token (maybe_target_entity_token), + KEY index_event_type (event_type) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-06-13-010101_synthetic_ids/down.sql b/_migrations/2021-06-13-010101_synthetic_ids/down.sql new file mode 100644 index 0000000000..74c48f3c51 --- /dev/null +++ b/_migrations/2021-06-13-010101_synthetic_ids/down.sql @@ -0,0 +1,6 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table tts_result_synthetic_ids; +drop table w2l_result_synthetic_ids; diff --git a/_migrations/2021-06-13-010101_synthetic_ids/up.sql b/_migrations/2021-06-13-010101_synthetic_ids/up.sql new file mode 100644 index 0000000000..f96901a2ff --- /dev/null +++ b/_migrations/2021-06-13-010101_synthetic_ids/up.sql @@ -0,0 +1,47 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- It gets unwieldy to show results without a numeric sequence. +-- We create a synthetic sequence for each user to make revisiting results easier. +CREATE TABLE tts_result_synthetic_ids ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + user_token VARCHAR(32) NOT NULL, + + -- The next result id. + -- After we use an id, we increment this counter. + next_id BIGINT(20) NOT NULL DEFAULT 0, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (user_token) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +-- It gets unwieldy to show results without a numeric sequence. +-- We create a synthetic sequence for each user to make revisiting results easier. +CREATE TABLE w2l_result_synthetic_ids ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + user_token VARCHAR(32) NOT NULL, + + -- The next result id. + -- After we use an id, we increment this counter. + next_id BIGINT(20) NOT NULL DEFAULT 0, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (user_token) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2021-08-25-010101_alter_tables_add_vocoder_type/down.sql b/_migrations/2021-08-25-010101_alter_tables_add_vocoder_type/down.sql new file mode 100644 index 0000000000..16dcc8614f --- /dev/null +++ b/_migrations/2021-08-25-010101_alter_tables_add_vocoder_type/down.sql @@ -0,0 +1,12 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models + DROP COLUMN maybe_default_pretrained_vocoder; + +ALTER TABLE tts_inference_jobs + DROP COLUMN maybe_override_pretrained_vocoder; + +ALTER TABLE tts_results + DROP COLUMN maybe_pretrained_vocoder_used; diff --git a/_migrations/2021-08-25-010101_alter_tables_add_vocoder_type/up.sql b/_migrations/2021-08-25-010101_alter_tables_add_vocoder_type/up.sql new file mode 100644 index 0000000000..3ce73ad8f4 --- /dev/null +++ b/_migrations/2021-08-25-010101_alter_tables_add_vocoder_type/up.sql @@ -0,0 +1,21 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- TODO/FIXME: I originally created this with the wrong directory month. +-- I'll have to repair the migrations table. +-- version: 20210925010101 | run_on: 2021-08-25 03:01:15 (development) +-- UPDATE __diesel_schema_migrations SET version=20210825010101 WHERE version=20210925010101 LIMIT 1; +-- NB: Fixed in production on 2022-01-06 + +ALTER TABLE tts_models +ADD COLUMN maybe_default_pretrained_vocoder VARCHAR(64) DEFAULT NULL +AFTER text_preprocessing_algorithm; + +ALTER TABLE tts_inference_jobs +ADD COLUMN maybe_override_pretrained_vocoder VARCHAR(64) DEFAULT NULL +AFTER model_token; + +ALTER TABLE tts_results +ADD COLUMN maybe_pretrained_vocoder_used VARCHAR(64) DEFAULT NULL +AFTER model_token; diff --git a/_migrations/2021-09-10-053520_alter_tables_add_talknet_model_type/down.sql b/_migrations/2021-09-10-053520_alter_tables_add_talknet_model_type/down.sql new file mode 100644 index 0000000000..291a97c5ce --- /dev/null +++ b/_migrations/2021-09-10-053520_alter_tables_add_talknet_model_type/down.sql @@ -0,0 +1 @@ +-- This file should undo anything in `up.sql` \ No newline at end of file diff --git a/_migrations/2021-09-10-053520_alter_tables_add_talknet_model_type/up.sql b/_migrations/2021-09-10-053520_alter_tables_add_talknet_model_type/up.sql new file mode 100644 index 0000000000..5743c05bee --- /dev/null +++ b/_migrations/2021-09-10-053520_alter_tables_add_talknet_model_type/up.sql @@ -0,0 +1,23 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Add talknet to the tts_model_type. + +ALTER TABLE tts_models +MODIFY COLUMN tts_model_type ENUM( + 'not-set', + 'tacotron2', + 'glowtts', + 'glowtts-vocodes', + 'talknet' +) NOT NULL DEFAULT 'not-set'; + +ALTER TABLE tts_model_upload_jobs +MODIFY COLUMN tts_model_type ENUM( + 'not-set', + 'tacotron2', + 'glowtts', + 'glowtts-vocodes', + 'talknet' +) NOT NULL DEFAULT 'not-set'; diff --git a/_migrations/2021-09-10-074547_alter_tables_talknet_metadata/down.sql b/_migrations/2021-09-10-074547_alter_tables_talknet_metadata/down.sql new file mode 100644 index 0000000000..1da108803b --- /dev/null +++ b/_migrations/2021-09-10-074547_alter_tables_talknet_metadata/down.sql @@ -0,0 +1,15 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models +DROP COLUMN private_bucket_object_is_archive; + +ALTER TABLE tts_models +DROP COLUMN has_self_contained_vocoder; + +ALTER TABLE tts_models +DROP COLUMN has_self_contained_duration_model; + +ALTER TABLE tts_models +DROP COLUMN has_self_contained_pitch_model; diff --git a/_migrations/2021-09-10-074547_alter_tables_talknet_metadata/up.sql b/_migrations/2021-09-10-074547_alter_tables_talknet_metadata/up.sql new file mode 100644 index 0000000000..e0715f0878 --- /dev/null +++ b/_migrations/2021-09-10-074547_alter_tables_talknet_metadata/up.sql @@ -0,0 +1,19 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models +ADD COLUMN private_bucket_object_is_archive BOOLEAN NOT NULL DEFAULT FALSE +AFTER private_bucket_object_name; + +ALTER TABLE tts_models +ADD COLUMN has_self_contained_vocoder BOOLEAN NOT NULL DEFAULT FALSE +AFTER tts_model_type; + +ALTER TABLE tts_models +ADD COLUMN has_self_contained_duration_model BOOLEAN NOT NULL DEFAULT FALSE +AFTER has_self_contained_vocoder; + +ALTER TABLE tts_models +ADD COLUMN has_self_contained_pitch_model BOOLEAN NOT NULL DEFAULT FALSE +AFTER has_self_contained_duration_model; diff --git a/_migrations/2021-12-20-230000_model_categories/down.sql b/_migrations/2021-12-20-230000_model_categories/down.sql new file mode 100644 index 0000000000..986333f646 --- /dev/null +++ b/_migrations/2021-12-20-230000_model_categories/down.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE model_categories; +DROP TABLE tts_category_assignments; +DROP TABLE w2l_category_assignments; diff --git a/_migrations/2021-12-20-230000_model_categories/up.sql b/_migrations/2021-12-20-230000_model_categories/up.sql new file mode 100644 index 0000000000..8c8ab200ad --- /dev/null +++ b/_migrations/2021-12-20-230000_model_categories/up.sql @@ -0,0 +1,174 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE model_categories ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Visible "primary key" + token VARCHAR(32) NOT NULL, + + -- Idempotency token + -- This is so the frontend client doesn't submit duplicate requests. + -- Only used on create. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== MODEL TYPE ========== + + model_type ENUM( + 'not_set', + 'tts', + 'w2l' + ) NOT NULL DEFAULT 'not_set', + + -- ========== CATEGORY TOPOLOGY ========== + + -- Optional foreign key link to super category. + -- If absent, this is a top-level category or top-level super category. + maybe_super_category_token VARCHAR(32) DEFAULT NULL, + + -- If this category can be applied directly to models. + can_directly_have_models BOOLEAN NOT NULL DEFAULT false, + + -- If this category can have subcategories. + can_have_subcategories BOOLEAN NOT NULL DEFAULT false, + + -- If only mods are allowed to tag (or untag) models with this. + can_only_mods_apply BOOLEAN NOT NULL DEFAULT false, + + -- ========== CATEGORY NAME ========== + + -- The name of the category, which might be complicated. + name VARCHAR(255) NOT NULL, + + -- The name displayed in the dropdown (if it should be different), otherwise NULL. + -- eg. "Gender" (name) and "By Gender" (dropdown name) + maybe_dropdown_name VARCHAR(20) DEFAULT NULL, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the category. + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address_creation VARCHAR(40) NOT NULL, + -- Technically we don't allow editing, but just in case. + creator_ip_address_last_update VARCHAR(40) NOT NULL, + + -- ========== MODERATION DETAILS ========== + + -- Mods have to approve of categories. They remain hidden and unusable until approved. + is_mod_approved BOOLEAN DEFAULT NULL, + + -- If a moderator has comments. + maybe_mod_comments VARCHAR(255) DEFAULT NULL, + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If the category is deleted and shouldn't be displayed. + -- Only mods can delete categories. + deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_maybe_super_category_token (maybe_super_category_token), + KEY fk_creator_user_token (creator_user_token), + KEY index_model_type (model_type), + KEY index_can_directly_have_models (can_directly_have_models), + KEY index_can_have_subcategories (can_have_subcategories), + KEY index_is_mod_approved (is_mod_approved) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +-- Join table assigning models to zero or more categories. +CREATE TABLE tts_category_assignments ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + model_token VARCHAR(32) NOT NULL, + category_token VARCHAR(32) NOT NULL, + + -- ========== USER DETAILS ========== + + -- The person that created the assignment. + -- This can be the model owner or a moderator. + category_addition_user_token VARCHAR(32) DEFAULT NULL, + + -- The person that removed the assignment. + -- This can be the model owner or a moderator. + category_removal_user_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- THIS CAN BE A MODERATOR FOR SIMPLICITY. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + ip_address_last_update VARCHAR(40) NOT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (model_token, category_token), + KEY fk_model_token (model_token), + KEY fk_category_token (category_token), + KEY index_deleted_at (deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + +-- Join table assigning models to zero or more categories. +CREATE TABLE w2l_category_assignments ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + model_token VARCHAR(32) NOT NULL, + category_token VARCHAR(32) NOT NULL, + + -- ========== USER DETAILS ========== + + -- The person that created the assignment. + -- This can be the model owner or a moderator. + category_addition_user_token VARCHAR(32) DEFAULT NULL, + + -- The person that removed the assignment. + -- This can be the model owner or a moderator. + category_removal_user_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- THIS CAN BE A MODERATOR FOR SIMPLICITY. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + ip_address_last_update VARCHAR(40) NOT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (model_token, category_token), + KEY fk_model_token (model_token), + KEY fk_category_token (category_token), + KEY index_deleted_at (deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-01-11-042235_twitch_oauth_tokens/down.sql b/_migrations/2022-01-11-042235_twitch_oauth_tokens/down.sql new file mode 100644 index 0000000000..0e3e1baa53 --- /dev/null +++ b/_migrations/2022-01-11-042235_twitch_oauth_tokens/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE twitch_oauth_tokens; diff --git a/_migrations/2022-01-11-042235_twitch_oauth_tokens/up.sql b/_migrations/2022-01-11-042235_twitch_oauth_tokens/up.sql new file mode 100644 index 0000000000..bcaba96033 --- /dev/null +++ b/_migrations/2022-01-11-042235_twitch_oauth_tokens/up.sql @@ -0,0 +1,264 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- **NOTE ABOUT DESIGN**: +-- This table contains _many_ OAuth tokens for any single user. (many users : many tokens) +-- An entire history of tokens is kept in the table. +-- Apart from "deleted_at", etc., these records should largely be IMMUTABLE. + +CREATE TABLE twitch_oauth_tokens( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- Minted on first request and copied on each refresh. + -- This allows us to group access tokens all the way back up the + -- refresh chain. + oauth_refresh_grouping_token VARCHAR(32) NOT NULL, + + -- ========== STORYTELLER/FAKEYOU USER ========== + + -- Foreign key to user (Storyteller/FakeYou user) + -- If no user is associated (anonymous Twitch user), then this is null. + maybe_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== TWITCH USER ========== + + -- The user ID / channel ID are the same + -- https://discuss.dev.twitch.tv/t/what-is-the-difference-between-the-stream--id-and-channel--id/4423 + -- Several suggest this should be a string and not an integer + -- https://discuss.dev.twitch.tv/t/type-of-user-id-in-api-responses/10205 + -- Yep, strings + -- https://discuss.dev.twitch.tv/t/bug-v5-api-returns--id-as-string-for-featured-channels/10310 + twitch_user_id VARCHAR(64) NOT NULL, + + -- Twitch usernames are between 4 and 25 characters. + -- This is returned in the oauth flow + twitch_username VARCHAR(32) NOT NULL, + + -- Twitch username, but lowercase for lookup. + twitch_username_lowercase VARCHAR(32) NOT NULL, + + -- ========== OAUTH TOKEN DETAILS ========== + + -- The ever important OAuth access token. + access_token VARCHAR(128) NOT NULL, + + -- If we can refresh the token, this is the thing to use. + maybe_refresh_token VARCHAR(128) DEFAULT NULL, + + -- Should be "bearer". Included in the OAuth redemption. + token_type VARCHAR(32) DEFAULT NULL, + + -- When the token expires from time of first grant. + -- (We don't update this field.) + -- Null if it does not expire or we were not informed + expires_in_seconds INT(10) UNSIGNED DEFAULT NULL, + + -- Number of times this token has been refreshed + -- The original token creation will be "0", and each + -- subsequent time will bump by one. Mostly for debugging. + refresh_count INT(10) UNSIGNED DEFAULT 0, + + -- ====================================================== + -- ==================== OAUTH SCOPES ==================== + -- ====================================================== + + -- TODO: Probably could have stored this in a bitfield. Ugh. + + -- "analytics:read:extensions" + -- "View analytics data for the Twitch Extensions owned by the authenticated account." + has_analytics_read_extensions BOOLEAN NOT NULL DEFAULT FALSE, + + -- scope: "analytics:read:games" + -- "View analytics data for the games owned by the authenticated account." + has_analytics_read_games BOOLEAN NOT NULL DEFAULT FALSE, + + -- "bits:read" + -- "View Bits information for a channel." + has_bits_read BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:edit:commercial" + -- "Run commercials on a channel." + has_channel_edit_commercial BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:broadcast" + -- "Manage a channel's broadcast configuration, including updating channel configuration and managing stream markers and stream tags." + has_channel_manage_broadcast BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:extensions" + -- "Manage a channel's Extension configuration, including activating Extensions." + has_channel_manage_extensions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:polls" + -- "Manage a channel's polls." + has_channel_manage_polls BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:predictions" + -- "Manage of channel's Channel Points Predictions" + has_channel_manage_predictions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:redemptions" + -- "Manage Channel Points custom rewards and their redemptions on a channel." + has_channel_manage_redemptions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:schedule" + -- "Manage a channel's stream schedule." + has_channel_manage_schedule BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:manage:videos" + -- "Manage a channel's videos, including deleting videos." + has_channel_manage_videos BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:moderate" + -- "Perform moderation actions in a channel. The user requesting the scope must be a moderator in the channel." + has_channel_moderate BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:editors" + -- "View a list of users with the editor role for a channel." + has_channel_read_editors BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:goals" + -- "View Creator Goals for a channel." + has_channel_read_goals BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:hype_train" + -- "View Hype Train information for a channel." + has_channel_read_hype_train BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:polls" + -- "View a channel's polls." + has_channel_read_polls BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:predictions" + -- "View a channel's Channel Points Predictions." + has_channel_read_predictions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:redemptions" + -- "View Channel Points custom rewards and their redemptions on a channel." + -- eg. PubSub subscribe to `channel-points-channel-v1.` (channel points spends) + has_channel_read_redemptions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:stream_key" + -- "View an authorized user's stream key." + has_channel_read_stream_key BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel:read:subscriptions" + -- "View a list of all subscribers to a channel and check if a user is subscribed to a channel." + -- eg. Enumerate list of subscribers for a channel + -- eg. PubSub subscribe to `channel-subscribe-events-v1.` (subscribe, resubscribe, gift) + has_channel_read_subscriptions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "channel_subscriptions" + -- "\\[DEPRECATED\\] Read all subscribers to your channel." + has_channel_subscriptions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "chat:edit" + -- "Send live stream chat and rooms messages." + has_chat_edit BOOLEAN NOT NULL DEFAULT FALSE, + + -- "chat:read" + -- "View live stream chat and rooms messages." + has_chat_read BOOLEAN NOT NULL DEFAULT FALSE, + + -- "clips:edit" + -- "Manage Clips for a channel." + has_clips_edit BOOLEAN NOT NULL DEFAULT FALSE, + + -- "moderation:read" + -- "View a channel's moderation data including Moderators, Bans, Timeouts, and Automod settings." + has_moderation_read BOOLEAN NOT NULL DEFAULT FALSE, + + -- "moderator:manage:automod" + -- "Manage messages held for review by AutoMod in channels where you are a moderator." + has_moderator_manage_auto_mod BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:edit" + -- "Manage a user object." + has_user_edit BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:edit:broadcast" + -- "Edit your channel's broadcast configuration, including extension configuration. (This scope implies user:read:broadcast capability.)" + has_user_edit_broadcast BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:edit:follows" + -- "Edit a user's follows." + has_user_edit_follows BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:manage:blocked_users" + -- "Manage the block list of a user." + has_user_manage_blocked_users BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:read:blocked_users" + -- "View the block list of a user." + has_user_read_blocked_users BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:read:broadcast" + -- "View a user's broadcasting configuration, including Extension configurations." + has_user_read_broadcast BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:read:email" + -- "Read an authorized user's email address." + has_user_read_email BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:read:follows" + -- "View the list of channels a user follows." + -- eg. We can use this to see if a Twitch user follows us. + has_user_read_follows BOOLEAN NOT NULL DEFAULT FALSE, + + -- "user:read:subscriptions" + -- "View if an authorized user is subscribed to specific channels." + has_user_read_subscriptions BOOLEAN NOT NULL DEFAULT FALSE, + + -- "whispers:edit" + -- "Send whisper messages." + has_whispers_edit BOOLEAN NOT NULL DEFAULT FALSE, + + -- "whispers:read" + -- "View your whisper messages." + has_whispers_read BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== SECURITY ========== + + -- For abuse tracking. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT UNSIGNED NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- This is when the OAuth token will expire. + -- While we need to refresh before this time, Twitch recommends against eager + -- refreshes and instead wants the code to lazily refresh. + -- This also has no bearing on the "refresh_token". + -- We'll just use this for bookkeeping. + expires_at TIMESTAMP DEFAULT NULL, + + -- If the OAuth token is deleted, we set this. + -- We'll need to set it for all of a user's OAuth tokens since this table stores many tokens. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + KEY index_oauth_refresh_grouping_token (oauth_refresh_grouping_token), + KEY fk_maybe_user_token (maybe_user_token), + KEY index_twitch_user_id (twitch_user_id), + -- KEY index_twitch_username (twitch_username), + KEY index_twitch_username_lowercase (twitch_username_lowercase), + KEY index_expires_at (expires_at), + KEY index_user_deleted_at (user_deleted_at), + KEY index_mod_deleted_at (mod_deleted_at) + + -- KEY index_has_bits_read (has_bits_read), + -- KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; + diff --git a/_migrations/2022-01-26-045721_twitch_bits/down.sql b/_migrations/2022-01-26-045721_twitch_bits/down.sql new file mode 100644 index 0000000000..fb1e062e72 --- /dev/null +++ b/_migrations/2022-01-26-045721_twitch_bits/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE twitch_bits_events; diff --git a/_migrations/2022-01-26-045721_twitch_bits/up.sql b/_migrations/2022-01-26-045721_twitch_bits/up.sql new file mode 100644 index 0000000000..83564535da --- /dev/null +++ b/_migrations/2022-01-26-045721_twitch_bits/up.sql @@ -0,0 +1,44 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- This is primarily meant to be for analysis. +-- NB: See [_docs/database_schema.md] for details on Twitch-specifics. + +CREATE TABLE twitch_bits_events( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- ========== USER / CHANNEL ========== + + -- The user sending the bits. + -- We'll need to keep the name and lowercase for lookup. + sender_twitch_user_id VARCHAR(64) NOT NULL, + sender_twitch_username VARCHAR(32) NOT NULL, + sender_twitch_username_lowercase VARCHAR(32) NOT NULL, + + -- The channel the bits were being sent to. + -- We technically only need the ID since we can JOIN, but we'll keep + -- the name for convenience. + destination_channel_id VARCHAR(64) NOT NULL, + destination_channel_name VARCHAR(32) NOT NULL, + + -- ========== BITS DETAILS ========== + + is_anonymous BOOLEAN NOT NULL DEFAULT FALSE, + + -- The number of bits just spent + bits_used INT(10) UNSIGNED DEFAULT NULL, + + -- The lifetime(?) number of bits spent by the user in this channel. + total_bits_used INT(10) UNSIGNED DEFAULT NULL, + + chat_message TEXT NOT NULL, + + -- INDICES -- + PRIMARY KEY (id), + KEY index_sender_twitch_user_id (sender_twitch_user_id), + KEY index_sender_twitch_username_lowercase (sender_twitch_username_lowercase), + KEY index_destination_channel_id (destination_channel_id) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-01-26-045730_twitch_channel_points/down.sql b/_migrations/2022-01-26-045730_twitch_channel_points/down.sql new file mode 100644 index 0000000000..9cc5937547 --- /dev/null +++ b/_migrations/2022-01-26-045730_twitch_channel_points/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE twitch_channel_point_events; diff --git a/_migrations/2022-01-26-045730_twitch_channel_points/up.sql b/_migrations/2022-01-26-045730_twitch_channel_points/up.sql new file mode 100644 index 0000000000..a8a23e89ab --- /dev/null +++ b/_migrations/2022-01-26-045730_twitch_channel_points/up.sql @@ -0,0 +1,57 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- This is primarily meant to be for analysis. +-- NB: See [_docs/database_schema.md] for details on Twitch-specifics. + +CREATE TABLE twitch_channel_point_events ( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- ========== USER / CHANNEL ========== + + -- The user sending the bits. + -- We'll need to keep the name and lowercase for lookup. + sender_twitch_user_id VARCHAR(64) NOT NULL, + sender_twitch_username VARCHAR(32) NOT NULL, + sender_twitch_username_lowercase VARCHAR(32) NOT NULL, + + -- The channel the bits were being sent to. + -- We technically only need the ID since we can JOIN, but we'll keep + -- the name for convenience. + destination_channel_id VARCHAR(64) NOT NULL, + destination_channel_name VARCHAR(32) NOT NULL, + + -- ========== CHANNEL POINTS DETAILS ========== + + -- This is the title used for the reward + -- We'll match on this to do TTS + title VARCHAR(256) NOT NULL, + + -- This is the description used for the reward + prompt VARCHAR(256) NOT NULL, + + -- This is the text that the user sends + -- This is optional, but we'll coerce null to empty + user_text_input TEXT NOT NULL, + + -- A UUID for this instance of reward redemption + redemption_id VARCHAR(36) NOT NULL, + + -- A UUID for the reward + reward_id VARCHAR(36) NOT NULL, + + reward_cost INT(10) UNSIGNED DEFAULT NULL, + + is_sub_only BOOLEAN NOT NULL DEFAULT FALSE, + max_per_stream INT(10) UNSIGNED DEFAULT NULL, + max_per_user_per_stream INT(10) UNSIGNED DEFAULT NULL, + + -- INDICES -- + PRIMARY KEY (id), + KEY index_sender_twitch_user_id (sender_twitch_user_id), + KEY index_sender_twitch_username_lowercase (sender_twitch_username_lowercase), + KEY index_destination_channel_id (destination_channel_id) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-02-09-041205_alter_tts_add_twitch/down.sql b/_migrations/2022-02-09-041205_alter_tts_add_twitch/down.sql new file mode 100644 index 0000000000..6f51d20bd7 --- /dev/null +++ b/_migrations/2022-02-09-041205_alter_tts_add_twitch/down.sql @@ -0,0 +1,10 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_inference_jobs DROP COLUMN is_for_twitch; +ALTER TABLE tts_inference_jobs DROP COLUMN is_from_api; +ALTER TABLE tts_inference_jobs DROP COLUMN priority_level; + +ALTER TABLE tts_results DROP COLUMN is_from_api; +ALTER TABLE tts_results DROP COLUMN is_for_twitch; diff --git a/_migrations/2022-02-09-041205_alter_tts_add_twitch/up.sql b/_migrations/2022-02-09-041205_alter_tts_add_twitch/up.sql new file mode 100644 index 0000000000..efda373105 --- /dev/null +++ b/_migrations/2022-02-09-041205_alter_tts_add_twitch/up.sql @@ -0,0 +1,23 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_inference_jobs + ADD COLUMN is_from_api BOOLEAN NOT NULL DEFAULT FALSE + AFTER creator_set_visibility, + + ADD COLUMN is_for_twitch BOOLEAN NOT NULL DEFAULT FALSE + AFTER is_from_api, + + -- Priority *increases*, so a level of 2 will be higher than 1. + -- By default, all jobs from FakeYou have level 0. + -- Twitch TTS has level 1. + -- Paid Twitch and Paid API have level 2. + ADD COLUMN priority_level TINYINT UNSIGNED NOT NULL DEFAULT 0 + AFTER status; + +ALTER TABLE tts_results + ADD COLUMN is_from_api BOOLEAN NOT NULL DEFAULT FALSE + AFTER duration_millis, + ADD COLUMN is_for_twitch BOOLEAN NOT NULL DEFAULT FALSE + AFTER is_from_api; diff --git a/_migrations/2022-02-11-145145_alter_tts_model_add_featured/down.sql b/_migrations/2022-02-11-145145_alter_tts_model_add_featured/down.sql new file mode 100644 index 0000000000..3de7f7caf0 --- /dev/null +++ b/_migrations/2022-02-11-145145_alter_tts_model_add_featured/down.sql @@ -0,0 +1,6 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models DROP COLUMN is_front_page_featured; +ALTER TABLE tts_models DROP COLUMN is_twitch_featured; diff --git a/_migrations/2022-02-11-145145_alter_tts_model_add_featured/up.sql b/_migrations/2022-02-11-145145_alter_tts_model_add_featured/up.sql new file mode 100644 index 0000000000..9eb3c98d56 --- /dev/null +++ b/_migrations/2022-02-11-145145_alter_tts_model_add_featured/up.sql @@ -0,0 +1,10 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models + ADD COLUMN is_front_page_featured BOOLEAN NOT NULL DEFAULT FALSE + AFTER version_string, + + ADD COLUMN is_twitch_featured BOOLEAN NOT NULL DEFAULT FALSE + AFTER is_front_page_featured; diff --git a/_migrations/2022-02-13-223514_twitch_event_rules/down.sql b/_migrations/2022-02-13-223514_twitch_event_rules/down.sql new file mode 100644 index 0000000000..a2a3be1af1 --- /dev/null +++ b/_migrations/2022-02-13-223514_twitch_event_rules/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE twitch_event_rules; diff --git a/_migrations/2022-02-13-223514_twitch_event_rules/up.sql b/_migrations/2022-02-13-223514_twitch_event_rules/up.sql new file mode 100644 index 0000000000..bd6f18513f --- /dev/null +++ b/_migrations/2022-02-13-223514_twitch_event_rules/up.sql @@ -0,0 +1,70 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- How to respond to different types of Twitch events. +-- Multiple rows per user +CREATE TABLE twitch_event_rules( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- Idempotency token from client + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- Foreign key to user + -- This is a FakeYou/Storyteller account, *NOT* a twitch user id. + user_token VARCHAR(32) NOT NULL, + + -- ========== EVENT MATCHING AND BEHAVIOR ========== + + -- What type of event we'll be responding to. + event_category ENUM( + 'bits', + 'channel_points', + 'chat_command' + ) NOT NULL, + + -- A JSON payload containing any predicates we wish to apply to the match. + -- eg. bits_value > 100 + -- eg. channel_points_award_name == "Mario Voice" + event_match_predicate MEDIUMTEXT NOT NULL, + + -- A JSON payload containing how we wish to respond to the event. + -- eg. tts M:model + event_response MEDIUMTEXT NOT NULL, + + -- ========== RULE ORGANIZATION AND MANAGEMENT ========== + + -- The user can rearrange the rules in the UI. + -- This will be the order they apply in if matched. + -- NB: There is nothing in MySQL to guarantee unique ordering. + user_specified_rule_order INT(10) UNSIGNED NOT NULL DEFAULT 0, + + -- Whether or not the rule is enabled. + -- This is different than deleted. It still shows up in the UI. + rule_is_disabled BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== ABUSE TRACKING, ETC. ========== + + -- For abuse tracking. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + ip_address_last_update VARCHAR(40) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + deleted_at TIMESTAMP DEFAULT NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (uuid_idempotency_token), + UNIQUE KEY (token), + KEY fk_user_token (user_token), + KEY index_event_category (event_category), + KEY index_deleted_at (deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-02-13-223523_twitch_general_settings/down.sql b/_migrations/2022-02-13-223523_twitch_general_settings/down.sql new file mode 100644 index 0000000000..bc9ea84a15 --- /dev/null +++ b/_migrations/2022-02-13-223523_twitch_general_settings/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE twitch_general_settings; diff --git a/_migrations/2022-02-13-223523_twitch_general_settings/up.sql b/_migrations/2022-02-13-223523_twitch_general_settings/up.sql new file mode 100644 index 0000000000..956dcbeb29 --- /dev/null +++ b/_migrations/2022-02-13-223523_twitch_general_settings/up.sql @@ -0,0 +1,38 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- One row per user. Contains all of the general (non-repeated) configs. +CREATE TABLE twitch_general_settings ( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- Foreign key to user + -- This is a FakeYou/Storyteller account, *NOT* a twitch user id. + user_token VARCHAR(32) NOT NULL, + + -- Kill switch for the entire service. + storyteller_is_enabled BOOLEAN NOT NULL DEFAULT FALSE, + + -- Kill switch for TTS. + tts_is_enabled BOOLEAN NOT NULL DEFAULT FALSE, + + -- Kill switch for W2L. + w2l_is_enabled BOOLEAN NOT NULL DEFAULT FALSE, + + -- Can set a fallback voice if a TTS voice goes missing. + maybe_fallback_tts_model_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + ip_address_last_update VARCHAR(40) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (user_token) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-02-16-072149_api_tokens/down.sql b/_migrations/2022-02-16-072149_api_tokens/down.sql new file mode 100644 index 0000000000..f8dce0d9c1 --- /dev/null +++ b/_migrations/2022-02-16-072149_api_tokens/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE api_tokens; diff --git a/_migrations/2022-02-16-072149_api_tokens/up.sql b/_migrations/2022-02-16-072149_api_tokens/up.sql new file mode 100644 index 0000000000..7873bb4889 --- /dev/null +++ b/_migrations/2022-02-16-072149_api_tokens/up.sql @@ -0,0 +1,47 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- We'll use application logic to kill any API keys beyond the most recent 5 "active" a user creates. +CREATE TABLE api_tokens ( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- Idempotency token from client + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- This is a private internal reference to the API token. + -- Not shared with consumers / users. + -- This should be used as the foreign key for the token. + internal_token VARCHAR(32) NOT NULL, + + -- This is the secret API token value given to consumers + api_token VARCHAR(32) NOT NULL, + + -- Foreign key to user that owns the API token. + -- A user can have several. + user_token VARCHAR(32) NOT NULL, + + -- A user-defined short description (optional) + maybe_short_description VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + ip_address_last_update VARCHAR(40) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + deleted_at TIMESTAMP DEFAULT NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (uuid_idempotency_token), + UNIQUE KEY (internal_token), + UNIQUE KEY (api_token), + KEY fk_user_token (user_token), + KEY index_created_at (created_at), + KEY index_deleted_at (deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-02-19-004032_voice_clone_requests/down.sql b/_migrations/2022-02-19-004032_voice_clone_requests/down.sql new file mode 100644 index 0000000000..83299ec282 --- /dev/null +++ b/_migrations/2022-02-19-004032_voice_clone_requests/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE voice_clone_requests; diff --git a/_migrations/2022-02-19-004032_voice_clone_requests/up.sql b/_migrations/2022-02-19-004032_voice_clone_requests/up.sql new file mode 100644 index 0000000000..e59d133c09 --- /dev/null +++ b/_migrations/2022-02-19-004032_voice_clone_requests/up.sql @@ -0,0 +1,77 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE voice_clone_requests( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + -- Idempotency token from client + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- ========== CONTACT ========== + + -- Foreign key to user + -- This is a FakeYou/Storyteller account, *NOT* a twitch user id. + maybe_user_token VARCHAR(32) DEFAULT NULL, + + email_address VARCHAR(255) NOT NULL, + + discord_username VARCHAR(255) NOT NULL, + + -- ========== REASON ========== + + -- Visibility + is_for_private_use BOOLEAN NOT NULL DEFAULT FALSE, + is_for_public_use BOOLEAN NOT NULL DEFAULT FALSE, + + -- How they'll use + is_for_studio BOOLEAN NOT NULL DEFAULT FALSE, + is_for_twitch_tts BOOLEAN NOT NULL DEFAULT FALSE, + is_for_api_use BOOLEAN NOT NULL DEFAULT FALSE, + is_for_music BOOLEAN NOT NULL DEFAULT FALSE, + is_for_games BOOLEAN NOT NULL DEFAULT FALSE, + is_for_other BOOLEAN NOT NULL DEFAULT FALSE, + optional_notes_on_use MEDIUMTEXT DEFAULT NULL, + + -- Subject/Ownership + is_own_voice BOOLEAN NOT NULL DEFAULT FALSE, + is_third_party_voice BOOLEAN NOT NULL DEFAULT FALSE, + optional_notes_on_subject VARCHAR(255) DEFAULT NULL, + + -- ========== RECORDING QUALITY / EQUIPMENT ========== + + has_clean_audio_recordings BOOLEAN NOT NULL DEFAULT FALSE, + has_good_microphone BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== COMMENTS ========== + + optional_questions MEDIUMTEXT DEFAULT NULL, + optional_extra_comments MEDIUMTEXT DEFAULT NULL, + + -- ========== BUSINESS END / INTERNAL USE ========== + + internal_has_been_handled BOOLEAN NOT NULL DEFAULT FALSE, + + internal_notes MEDIUMTEXT DEFAULT NULL, + + -- ========== ABUSE TRACKING, ETC. ========== + + -- For abuse tracking. + -- Wide enough for IPv4/6 + ip_address_creation VARCHAR(40) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (uuid_idempotency_token), + UNIQUE KEY (token), + KEY fk_maybe_user_token (maybe_user_token) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-03-01-085243_alter_table_tts_models_add_language/down.sql b/_migrations/2022-03-01-085243_alter_table_tts_models_add_language/down.sql new file mode 100644 index 0000000000..7c3a1d738d --- /dev/null +++ b/_migrations/2022-03-01-085243_alter_table_tts_models_add_language/down.sql @@ -0,0 +1,9 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models + DROP COLUMN ietf_language_tag; + +ALTER TABLE tts_models + DROP COLUMN ietf_primary_language_subtag; diff --git a/_migrations/2022-03-01-085243_alter_table_tts_models_add_language/up.sql b/_migrations/2022-03-01-085243_alter_table_tts_models_add_language/up.sql new file mode 100644 index 0000000000..878a8a62c6 --- /dev/null +++ b/_migrations/2022-03-01-085243_alter_table_tts_models_add_language/up.sql @@ -0,0 +1,13 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- The full IETF BCP47 language tag (eg. en, en-US, es-419, ja-JP, pt, etc.) +ALTER TABLE tts_models + ADD COLUMN ietf_language_tag VARCHAR(64) NOT NULL DEFAULT 'en' + AFTER maybe_default_pretrained_vocoder; + +-- The IETF BCP47 language tag's primary language subtag (eg. "es-419" becomes "es") +ALTER TABLE tts_models + ADD COLUMN ietf_primary_language_subtag VARCHAR(12) NOT NULL DEFAULT 'en' + AFTER ietf_language_tag; diff --git a/_migrations/2022-03-01-085254_alter_table_tts_models_add_bot_command/down.sql b/_migrations/2022-03-01-085254_alter_table_tts_models_add_bot_command/down.sql new file mode 100644 index 0000000000..5b97b7b587 --- /dev/null +++ b/_migrations/2022-03-01-085254_alter_table_tts_models_add_bot_command/down.sql @@ -0,0 +1,6 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE tts_models + DROP COLUMN maybe_suggested_unique_bot_command; diff --git a/_migrations/2022-03-01-085254_alter_table_tts_models_add_bot_command/up.sql b/_migrations/2022-03-01-085254_alter_table_tts_models_add_bot_command/up.sql new file mode 100644 index 0000000000..e617b9b0ed --- /dev/null +++ b/_migrations/2022-03-01-085254_alter_table_tts_models_add_bot_command/up.sql @@ -0,0 +1,14 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- A bot command prefix, eg. 'sonic' that could be interpreted as '/sonic' by a TTS system. +-- These are unique strings and can only be applied to one model at a time. +-- Only moderators can set this, and it's only for the best voices. +ALTER TABLE tts_models + ADD COLUMN maybe_suggested_unique_bot_command VARCHAR(255) DEFAULT NULL + AFTER is_twitch_featured; + +-- And the unique index. +ALTER TABLE tts_models + ADD UNIQUE INDEX unique_index_maybe_suggested_unique_bot_command (maybe_suggested_unique_bot_command); diff --git a/_migrations/2022-03-02-100137_alter_table_twitch_oauth_tokens_add_token_pk/down.sql b/_migrations/2022-03-02-100137_alter_table_twitch_oauth_tokens_add_token_pk/down.sql new file mode 100644 index 0000000000..2768e2a632 --- /dev/null +++ b/_migrations/2022-03-02-100137_alter_table_twitch_oauth_tokens_add_token_pk/down.sql @@ -0,0 +1,6 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE twitch_oauth_tokens + DROP COLUMN internal_token; diff --git a/_migrations/2022-03-02-100137_alter_table_twitch_oauth_tokens_add_token_pk/up.sql b/_migrations/2022-03-02-100137_alter_table_twitch_oauth_tokens_add_token_pk/up.sql new file mode 100644 index 0000000000..39a519c385 --- /dev/null +++ b/_migrations/2022-03-02-100137_alter_table_twitch_oauth_tokens_add_token_pk/up.sql @@ -0,0 +1,12 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- In case we ever want to query these via an admin interface, it'd be nice to not use IDs +ALTER TABLE twitch_oauth_tokens + ADD COLUMN internal_token VARCHAR(32) NOT NULL + AFTER id; + +-- And the unique index. +ALTER TABLE twitch_oauth_tokens + ADD UNIQUE INDEX unique_index_internal_token (internal_token); diff --git a/_migrations/2022-03-14-035745_twitch_event_match_logs/down.sql b/_migrations/2022-03-14-035745_twitch_event_match_logs/down.sql new file mode 100644 index 0000000000..ff2b203e89 --- /dev/null +++ b/_migrations/2022-03-14-035745_twitch_event_match_logs/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE twitch_event_match_logs; diff --git a/_migrations/2022-03-14-035745_twitch_event_match_logs/up.sql b/_migrations/2022-03-14-035745_twitch_event_match_logs/up.sql new file mode 100644 index 0000000000..25e4fee799 --- /dev/null +++ b/_migrations/2022-03-14-035745_twitch_event_match_logs/up.sql @@ -0,0 +1,23 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Log any event matching (bits, channel points, etc.) +CREATE TABLE twitch_event_match_logs( + -- Not used for anything except replication. + id BIGINT(20) UNSIGNED NOT NULL AUTO_INCREMENT, + + token VARCHAR(32) NOT NULL, + + -- Type of event, eg. + event_type VARCHAR(32) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + deleted_at TIMESTAMP DEFAULT NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations_squashed/tts_tables/tts_inference_jobs.sql b/_migrations_squashed/tts_tables/tts_inference_jobs.sql new file mode 100644 index 0000000000..2d3dd610c5 --- /dev/null +++ b/_migrations_squashed/tts_tables/tts_inference_jobs.sql @@ -0,0 +1,120 @@ +-- NB: This is a manually squashed view of all the CREATE and ALTER statements, +-- with comments attached to the fields for centralized documentation. + +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE tts_inference_jobs ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token from client + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== INFERENCE DETAILS ========== + + -- The model to use. + -- This also determines which architecture we're using. + model_token VARCHAR(32) NOT NULL, + + -- If set, use a different pretrained vocoder. + maybe_override_pretrained_vocoder VARCHAR(64) DEFAULT NULL, + + -- The raw, unprocessed user input. + raw_inference_text TEXT NOT NULL, + + -- ========== CREATOR DETAILS ========== + + -- Foreign key to user + -- If no user is logged in, this is null. + maybe_creator_user_token VARCHAR(32) DEFAULT NULL, + + -- Based on a cookie sent by the frontend. + -- We'll save this even if the user is logged in. + maybe_creator_anonymous_visitor_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== SOURCE METADATA ========== + + is_from_api BOOLEAN NOT NULL DEFAULT FALSE, + is_for_twitch BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- Priority *increases*, so a level of 2 will be higher than 1. + -- These are the level semantics currently: + -- - All jobs from anonymous FakeYou users have level 0. + -- - All jobs from logged in FakeYou users have level 1. + -- - All jobs from Twitch TTS (unpaid) have level 10 (ten). + -- - (There will be future levels for paid Twitch and social FakeYou rewards.) + priority_level TINYINT UNSIGNED NOT NULL DEFAULT 0, + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_success_result_token (on_success_result_token), + KEY fk_model_token (model_token), + KEY fk_maybe_creator_user_token (maybe_creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations_squashed/tts_tables/tts_model_upload_jobs.sql b/_migrations_squashed/tts_tables/tts_model_upload_jobs.sql new file mode 100644 index 0000000000..084449a325 --- /dev/null +++ b/_migrations_squashed/tts_tables/tts_model_upload_jobs.sql @@ -0,0 +1,120 @@ +-- NB: This is a manually squashed view of all the CREATE and ALTER statements, +-- with comments attached to the fields for centralized documentation. + +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- TTS MODELS +CREATE TABLE tts_model_upload_jobs ( + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_result_token VARCHAR(32) DEFAULT NULL, + + -- ========== UPLOAD DETAILS ========== + + -- The name of the voice + -- The "name" of the voice model, which might be complicated. + -- If maybe_subject_token (etc.) is set, then it's authoritative instead. + title VARCHAR(255) NOT NULL, + + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + tts_model_type ENUM( + 'not-set', + 'tacotron2', + 'glowtts', + 'glowtts-vocodes', + 'talknet' + ) NOT NULL DEFAULT 'not-set', + + -- Can be linked to a well-known voice + -- maybe_subject_token VARCHAR(32) DEFAULT NULL, + -- maybe_actor_subject_token VARCHAR(32) DEFAULT NULL, + + -- If we need to download the file from Google Drive. + download_url VARCHAR(512) DEFAULT NULL, + + -- NB: DO NOT SORT! + -- download_url_type ENUM( + -- 'google-drive', + -- 'web' + -- ) NOT NULL DEFAULT 'web', + + -- ========== CREATOR DETAILS AND PREFERENCES ========== + + -- Foreign key to user + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_success_result_token (on_success_result_token), + KEY fk_creator_user_token (creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations_squashed/tts_tables/tts_models.sql b/_migrations_squashed/tts_tables/tts_models.sql new file mode 100644 index 0000000000..d31242ce0f --- /dev/null +++ b/_migrations_squashed/tts_tables/tts_models.sql @@ -0,0 +1,204 @@ +-- NB: This is a manually squashed view of all the CREATE and ALTER statements, +-- with comments attached to the fields for centralized documentation. + +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE tts_models ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + tts_model_type ENUM( + 'not-set', + 'tacotron2', + 'glowtts', + 'glowtts-vocodes', + 'talknet' + ) NOT NULL DEFAULT 'not-set', + + -- For models distributed as archives/bundles (eg. TalkNet), we can pack + -- in a vocoder. Note that this vocoder does not have to be used. + has_self_contained_vocoder BOOLEAN NOT NULL DEFAULT FALSE, + + -- For models distributed as archives/bundles (eg. TalkNet), we can pack + -- in a duration model. + has_self_contained_duration_model BOOLEAN NOT NULL DEFAULT FALSE, + + -- For models distributed as archives/bundles (eg. TalkNet), we can pack + -- in a pitch model. + has_self_contained_pitch_model BOOLEAN NOT NULL DEFAULT FALSE, + + -- NB: DO NOT CHANGE ORDER; APPEND ONLY! + -- How text should be handled. + text_preprocessing_algorithm ENUM( + 'basic', + 'arpabet' + ) NOT NULL DEFAULT 'basic', + + -- If set, use a different pretrained vocoder. + -- If not set, use the website default. (Currently HifiGan-SS) + maybe_default_pretrained_vocoder VARCHAR(64) DEFAULT NULL, + + -- Optional Pointer to a newer version of the voice + -- If there's a newer version, we can disable this one. + -- maybe_updated_model_token VARCHAR(32) DEFAULT NULL, + + -- We can set as a combination of ['username' + 'voice-name'] + -- There can be public aliases for voices, eg. a voice's default model. + -- A user can change this. + -- As such, these should not be foreign keys. + -- updatable_slug VARCHAR(64) NOT NULL, + + -- Can be linked to a well-known voice + -- maybe_subject_token VARCHAR(32) DEFAULT NULL, + -- maybe_actor_subject_token VARCHAR(32) DEFAULT NULL, + + -- ========== UNSTRUCTURED METADATA ========== + + -- The full IETF BCP47 language tag (eg. en, en-US, es-419, ja-JP, pt, etc.) + ietf_language_tag VARCHAR(64) NOT NULL DEFAULT 'en', + + -- The IETF BCP47 language tag's primary language subtag (eg. "es-419" becomes "es") + ietf_primary_language_subtag VARCHAR(12) NOT NULL DEFAULT 'en', + + -- The "name" of the voice model, which might be complicated. + -- If maybe_subject_token (etc.) is set, then it's authoritative instead. + title VARCHAR(255) NOT NULL, + + -- The description of the model in markdown. + description_markdown TEXT NOT NULL, + + -- Generated HTML (not user-editable). + description_rendered_html TEXT NOT NULL, + + -- eg. Bill Clinton, Batman, or Goku. + voice_name VARCHAR(255) DEFAULT NULL, + + -- Optional; eg. Sean Schemmel + actor_name VARCHAR(255) DEFAULT NULL, + + -- Optional; eg. if the voice is "happy" or a singer "a-capella", etc. + characteristic VARCHAR(255) DEFAULT NULL, + + -- Optional; eg. if you train the model more than once. + version_string VARCHAR(255) DEFAULT NULL, + + -- Whether the voice features on the front of FakeYou or Storyteller Twitch TTS + is_front_page_featured BOOLEAN NOT NULL DEFAULT FALSE, + is_twitch_featured BOOLEAN NOT NULL DEFAULT FALSE, + + -- A suggested bot command prefix, eg. 'sonic' that could be interpreted as '/sonic' + -- by a TTS system. These are unique strings that can only be applied to one model at + -- a time at most. Only moderators can set this, and it's only for the best voices. + maybe_suggested_unique_bot_command VARCHAR(255) DEFAULT NULL, + + -- ========== IMMUTABLE PROVENANCE AND METADATA ========== + + -- Where the file was originally downloaded (if it was downloaded) + original_download_url VARCHAR(512) DEFAULT NULL, + + -- The filename that was used at upload time (if available) + original_filename VARCHAR(255) DEFAULT NULL, + + -- File characteristics + file_size_bytes INT(10) NOT NULL DEFAULT 0, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address_creation VARCHAR(40) NOT NULL, + creator_ip_address_last_update VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== BUCKET STORAGE ========== + + -- The pytorch model + -- For now, this will be a hash of the file contents. + -- NB: NOT UNIQUE! We can allow duplicate uploads. + private_bucket_hash CHAR(64) NOT NULL, + -- The "full url" version of the path + private_bucket_object_name VARCHAR(255) NOT NULL, + + -- Whether the file is a zip archive (rather than a model file). + private_bucket_object_is_archive BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== RATING AND RANKING ========== + + -- Calculated average, on a scale of 0-100 + -- Null with zero ratings. + -- calculated_average_score INT(3) DEFAULT NULL, + + -- Other metrics + -- calculated_total_ratings_submitted_count INT(10) NOT NULL DEFAULT 0, + -- calculated_total_uses_count BIGINT(10) NOT NULL DEFAULT 0, + + -- ========== MODERATION DETAILS ========== + + -- Mods have to approve of tts models for them to show up in a public index. + -- Models can be used by the author (or other parties that know the URL) before + -- they're approved, but unapproved models won't show up in public indices. + -- This is independent of other visibility controls (is_mod_disabled, deleted_at, + -- and creator_set_visibility) + is_public_listing_approved BOOLEAN DEFAULT NULL, + + -- Extremely popular models may be locked from deletion or modification by users. + is_locked_from_user_modification BOOLEAN NOT NULL DEFAULT FALSE, + + -- In this case, a moderator disables it from being used for inference. + -- This also disables it for the creator. + -- Unlike deletion, it remains "visible" to those that have access. + is_locked_from_use BOOLEAN NOT NULL DEFAULT FALSE, + + -- If a moderator has comments. + maybe_mod_comments VARCHAR(255) DEFAULT NULL, + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod. + -- It completely disappears from the system. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_creator_user_token (creator_user_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_creator_ip_address_creation (creator_ip_address_creation), + KEY index_creator_ip_address_last_update (creator_ip_address_last_update), + KEY index_creator_set_visibility (creator_set_visibility), + KEY index_private_bucket_hash (private_bucket_hash), + KEY index_is_public_listing_approved (is_public_listing_approved), + KEY index_is_locked_from_user_modification (is_locked_from_user_modification), + KEY index_is_locked_from_use (is_locked_from_use) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations_squashed/tts_tables/tts_results.sql b/_migrations_squashed/tts_tables/tts_results.sql new file mode 100644 index 0000000000..72756ab9d4 --- /dev/null +++ b/_migrations_squashed/tts_tables/tts_results.sql @@ -0,0 +1,105 @@ +-- NB: This is a manually squashed view of all the CREATE and ALTER statements, +-- with comments attached to the fields for centralized documentation. + +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE tts_results ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- ========== INFERENCE DETAILS ========== + + -- The TTS model that was used + model_token VARCHAR(32) NOT NULL, + + -- If set, specifies which pretrained vocoder was used. + maybe_pretrained_vocoder_used VARCHAR(64) DEFAULT NULL, + + -- The original raw, unprocessed user input. + raw_inference_text TEXT NOT NULL, + + -- SHA2 hash of the text [SHA2 = CHAR(64), SHA1 = CHAR(40), MD5 = CHAR(32)] + raw_inference_text_hash_sha2 CHAR(64) NOT NULL, + + -- First pass at text normalization. + -- eg. 14th -> fourteenth, etc. + normalized_inference_text TEXT NOT NULL, + + -- If the model uses arpabet, we'll save the arpabet we used. + maybe_arpabet_text TEXT DEFAULT NULL, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + -- If the user wasn't logged in, this is null + maybe_creator_user_token VARCHAR(32) DEFAULT NULL, + + -- Based on a cookie sent by the frontend. + -- We'll save this even if the user is logged in. + maybe_creator_anonymous_visitor_token VARCHAR(32) DEFAULT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- The synthetic id associated with this result. + -- These ids are incremented on a per-user basis to help users + -- sequence their own work. They serve no other purpose. + maybe_creator_synthetic_id BIGINT(20) DEFAULT NULL, + + -- ========== BUCKET STORAGE ========== + + -- Where the wav, spectrogram, and etc. are located. + public_bucket_wav_audio_path VARCHAR(255) NOT NULL, + public_bucket_spectrogram_path VARCHAR(255) NOT NULL, + + -- ========== METADATA ========== + + file_size_bytes INT(10) NOT NULL DEFAULT 0, + duration_millis INT(10) NOT NULL DEFAULT 0, + + -- ========== SOURCE METADATA ========== + + is_from_api BOOLEAN NOT NULL DEFAULT FALSE, + is_for_twitch BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== MODERATION DETAILS ========== + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod or the creator. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_model_token (model_token), + KEY fk_maybe_creator_user_token (maybe_creator_user_token), + KEY fk_maybe_creator_anonymous_visitor_token (maybe_creator_anonymous_visitor_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_creator_ip_address (creator_ip_address), + KEY index_user_deleted_at(user_deleted_at), + KEY index_mod_deleted_at(mod_deleted_at) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/analytics/user_retention.sql b/analytics/user_retention.sql new file mode 100644 index 0000000000..da1362f17c --- /dev/null +++ b/analytics/user_retention.sql @@ -0,0 +1,75 @@ + +-- All logged in users with results in the last X days +select u.created_at from users as u +where u.token IN +(select distinct u.token from tts_inference_jobs as j join users u on u.token = j.maybe_creator_user_token where j.created_at > now() - interval 7 day) +order by u.created_at ASC; + + +-- All logged in users with results in the last X days that are older than Y +select u.created_at from users as u +where u.token IN +(select distinct u.token from tts_inference_jobs as j join users u on u.token = j.maybe_creator_user_token where j.created_at > now() - interval 30 day) +AND u.created_at < now() - interval 30 day +order by u.created_at ASC; + +-- ====================== FEBRUARY COHORT =================================== + +-- [1] All users created between X and Y +select count(*) from users +where created_at > now() - interval 60 day +and created_at < now() - interval 30 day; + +-- [2] Active users from that cohort ^ +select count(*) from users as u +where u.token IN +(select distinct u.token from tts_inference_jobs as j join users u on u.token = j.maybe_creator_user_token where j.created_at > now() - interval 30 day) +AND created_at > now() - interval 60 day +AND created_at < now() - interval 30 day; + + + +-- ====================== JANUARY COHORT =================================== + +-- [1] All users created between X and Y +select count(*) from users +where created_at > now() - interval 90 day +and created_at < now() - interval 60 day; + +-- [2] Active users from that cohort ^ +select count(*) from users as u +where u.token IN +(select distinct u.token from tts_inference_jobs as j join users u on u.token = j.maybe_creator_user_token where j.created_at > now() - interval 30 day) +AND created_at > now() - interval 90 day +AND created_at < now() - interval 60 day; + + + +(select distinct u.token from tts_inference_jobs as j join users u on u.token = j.maybe_creator_user_token where j.created_at > now() - interval 30 day) + + +select j.creator_ip_address, count(*) as number +from tts_inference_jobs as j +where j.created_at > now() - interval 1 day +group by j.creator_ip_address +order by number desc +limit 100; + + +select count(*) from tts_inference_jobs +where created_at > now() - interval 1 day ; + +select count(distinct creator_ip_address) from tts_inference_jobs +where created_at > now() - interval 1 day ; + + + +select j.creator_ip_address, count(*) as number +from tts_inference_jobs as j +where j.created_at > now() - interval 1 day +group by j.creator_ip_address +order by number desc +limit 100; + + + diff --git a/codegen_database.sh b/codegen_database.sh new file mode 100755 index 0000000000..dc9d904b7f --- /dev/null +++ b/codegen_database.sh @@ -0,0 +1,46 @@ +#!/bin/bash + +# This is complicated since we're now in a more complex Rust workspace with +# multiple projects and targets. + +# NB(2022-01-31): For some reason, 'storyteller-web', despite importing the obs queries +# via the `database_queries` shared library, started to exclude these queries. It was +# working, but mysteriously stopped. I'm using this more complicated setup with jq-based +# json merge as a stop gap. + +# NB: This format can be used without changing directory: +# SQLX_OFFLINE=true cargo sqlx prepare -- --bin storyteller-web --manifest-path crates/service/storyteller_web/Cargo.toml + +build_shared_database_library() { + pushd crates/lib/database_queries + SQLX_OFFLINE=true cargo sqlx prepare + popd +} + +build_storyteller_web_app() { + # NB: This imports the inference/upload job queries + # It should also import the shared database lib queries, but something(???) broke. + pushd crates/service/storyteller_web + SQLX_OFFLINE=true cargo sqlx prepare -- --bin storyteller-web # NB: Because multiple binary targets + popd +} + +combine_sqlx_queries() { + # Merge multiple JSON files into a single dictionary. + # https://stackoverflow.com/a/24904276 + jq -s '.[0] * .[1]' \ + crates/lib/database_queries/sqlx-data.json \ + crates/service/storyteller_web/sqlx-data.json \ + > sqlx-data.json +} + +cleanup_temp_files() { + rm crates/lib/database_queries/sqlx-data.json \ + crates/service/storyteller_web/sqlx-data.json +} + +build_shared_database_library +build_storyteller_web_app +combine_sqlx_queries +cleanup_temp_files + diff --git a/crates/lib/config/Cargo.toml b/crates/lib/config/Cargo.toml new file mode 100644 index 0000000000..fc10d38795 --- /dev/null +++ b/crates/lib/config/Cargo.toml @@ -0,0 +1,13 @@ +[package] +name = "config" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "config" +path = "src/lib.rs" diff --git a/crates/lib/config/src/lib.rs b/crates/lib/config/src/lib.rs new file mode 100644 index 0000000000..f78455f59a --- /dev/null +++ b/crates/lib/config/src/lib.rs @@ -0,0 +1 @@ +pub mod shared_constants; diff --git a/crates/lib/config/src/shared_constants.rs b/crates/lib/config/src/shared_constants.rs new file mode 100644 index 0000000000..b142962b30 --- /dev/null +++ b/crates/lib/config/src/shared_constants.rs @@ -0,0 +1,20 @@ +//! Shared constants that all binaries use. +//! Or at least that deserve extremely high visibility. + +/// NB: `sqlx::query` is spammy and dumps all queries as "info"-level log lines. +/// NB: `hyper::proto::h1::io` is incredibly spammy and logs every chunk of bytes in very large files being downloaded. +pub const DEFAULT_RUST_LOG: &'static str = "debug,actix_web=info,sqlx::query=warn,hyper::proto::h1::io=warn,storyteller_web::threads::db_health_checker_thread::db_health_checker_thread=warn"; +//const DEFAULT_RUST_LOG: &'static str = "debug,actix_web=info"; // but sometimes we want to debug + +/// The default Redis connection string for use in development +pub const DEFAULT_REDIS_DATABASE_0_CONNECTION_STRING: &'static str = "redis://localhost/0"; + +/// The default Redis connection string for use in development +/// Database 1 +pub const DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING: &'static str = "redis://localhost/1"; + +/// The default MySql connection string for use in development +pub const DEFAULT_MYSQL_CONNECTION_STRING: &'static str = "mysql://storyteller:password@localhost/storyteller"; + +/// The number of results we return by default for paginated queries. +pub const DEFAULT_MYSQL_QUERY_RESULT_PAGE_SIZE: u16 = 25; diff --git a/crates/lib/container_common/Cargo.toml b/crates/lib/container_common/Cargo.toml new file mode 100644 index 0000000000..6976629abf --- /dev/null +++ b/crates/lib/container_common/Cargo.toml @@ -0,0 +1,21 @@ +[package] +name = "container_common" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "container_common" +path = "src/lib.rs" + +[dependencies] + +# External +anyhow = "1.0.40" +once_cell = "1.9.0" +rand = "0.8.3" +uuid = { version = "0.8.2", features = ["v4"] } diff --git a/crates/lib/container_common/src/anyhow_result.rs b/crates/lib/container_common/src/anyhow_result.rs new file mode 100644 index 0000000000..7784374637 --- /dev/null +++ b/crates/lib/container_common/src/anyhow_result.rs @@ -0,0 +1,2 @@ +/// Easier to import than anyhow::Result. +pub type AnyhowResult = anyhow::Result; diff --git a/crates/lib/container_common/src/cache/mod.rs b/crates/lib/container_common/src/cache/mod.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/lib/container_common/src/collections/mod.rs b/crates/lib/container_common/src/collections/mod.rs new file mode 100644 index 0000000000..dc0324cd61 --- /dev/null +++ b/crates/lib/container_common/src/collections/mod.rs @@ -0,0 +1 @@ +pub mod random_from_vec; \ No newline at end of file diff --git a/crates/lib/container_common/src/collections/random_from_vec.rs b/crates/lib/container_common/src/collections/random_from_vec.rs new file mode 100644 index 0000000000..652da76e83 --- /dev/null +++ b/crates/lib/container_common/src/collections/random_from_vec.rs @@ -0,0 +1,31 @@ +use rand::seq::SliceRandom; + +/// Return a random item from a vec. +pub fn random_from_vec(vec: &Vec) -> Option<&T> { + vec.choose(&mut rand::thread_rng()) +} + +#[cfg(test)] +mod tests { + use crate::collections::random_from_vec::random_from_vec; + + #[test] + fn pulls_item_from_vec() { + let items = vec!["one", "two", "three"]; + let selected = random_from_vec(&items).unwrap(); + assert!(items.contains(selected)); + + let items = vec!["one", "two", "three"].into_iter() + .map(|s| s.to_string()) + .collect(); + let selected = random_from_vec(&items).unwrap(); + assert!(items.contains(selected)); + } + + #[test] + fn behavior_on_empty() { + let items : Vec = vec![]; + let selected = random_from_vec(&items); + assert!(selected.is_none()); + } +} \ No newline at end of file diff --git a/crates/lib/container_common/src/i18n/mod.rs b/crates/lib/container_common/src/i18n/mod.rs new file mode 100644 index 0000000000..5530078424 --- /dev/null +++ b/crates/lib/container_common/src/i18n/mod.rs @@ -0,0 +1 @@ +pub mod supported_languages_for_models; \ No newline at end of file diff --git a/crates/lib/container_common/src/i18n/supported_languages_for_models.rs b/crates/lib/container_common/src/i18n/supported_languages_for_models.rs new file mode 100644 index 0000000000..6f7fd7a764 --- /dev/null +++ b/crates/lib/container_common/src/i18n/supported_languages_for_models.rs @@ -0,0 +1,112 @@ +use once_cell::sync::Lazy; +use std::collections::{HashSet, HashMap}; + +/// These are language tags we try to support. +static SUPPORTED_LANGUAGES_FOR_MODELS : Lazy> = Lazy::new(|| { + let language_tags = [ + + // ========== Technically, we only support these so far ========== + + // English + "en", + "en-AU", + "en-CA", + "en-GB", + "en-US", + // Spanish + "es", + "es-419", + "es-AR", + "es-CL", + "es-CO", + "es-ES", + "es-MX", + "es-PE", + "es-US", + // Portuguese + "pt", + "pt-BR", + + // ========== But these are on the horizon ========== + + // French + "fr", + "fr-FR", + // German + "de", + "de-DE", + // Japanese + "ja", + "ja-JP", + // Misc + "id", + "id-ID", + "it", + "it-IT", + "ru", + "ru-RU", + "th-TH", + "tr", + "tr-TR", + "zh-CN", + "zh-HK", + ]; + language_tags.iter() + .map(|tag| tag.to_string()) + .collect::>() +}); + +/// Convert lower case tags to canonical form +static SUPPORTED_LANGUAGES_FOR_MODELS_CANONICAL_MAP : Lazy> = Lazy::new(|| { + SUPPORTED_LANGUAGES_FOR_MODELS.iter() + .map(|tag| (tag.to_lowercase(), tag.clone())) + .collect::>() +}); + +/// Take a full IETF BCP 47 language tag and return the canonicalized form *IF* we support it +pub fn get_canonicalized_language_tag_for_model(language_tag: &str) -> Option<&'static str> { + let lowercase = language_tag.to_lowercase(); + SUPPORTED_LANGUAGES_FOR_MODELS_CANONICAL_MAP.get(&lowercase).map(|v| v.as_str()) +} + +/// Take a full IETF BCP 47 language tag and return true *IF* we support it +/// It must be the correct case. +pub fn is_valid_language_for_models(language_tag: &str) -> bool { + SUPPORTED_LANGUAGES_FOR_MODELS.contains(language_tag) +} + +#[cfg(test)] +mod tests { + use crate::i18n::supported_languages_for_models::{get_canonicalized_language_tag_for_model, is_valid_language_for_models}; + + #[test] + fn get_canonical_language_for_model_success() { + assert_eq!("en", get_canonicalized_language_tag_for_model("EN").unwrap()); + assert_eq!("es-419", get_canonicalized_language_tag_for_model("ES-419").unwrap()); + assert_eq!("en-US", get_canonicalized_language_tag_for_model("En-Us").unwrap()); + } + + #[test] + fn get_canonical_language_for_model_failure() { + assert!(get_canonicalized_language_tag_for_model("").is_none()); + assert!(get_canonicalized_language_tag_for_model("en-JP").is_none()); + assert!(get_canonicalized_language_tag_for_model("es-Lots-Of-Tags").is_none()); + } + + #[test] + fn test_valid_language_for_model() { + // Valid + assert!(is_valid_language_for_models("en")); + assert!(is_valid_language_for_models("es-419")); + assert!(is_valid_language_for_models("ja-JP")); + // Invalid due to case + assert!(!is_valid_language_for_models("EN")); + assert!(!is_valid_language_for_models("ES-419")); + assert!(!is_valid_language_for_models("JA-jp")); + // Invalid misc + assert!(!is_valid_language_for_models("")); + assert!(!is_valid_language_for_models("foo")); + assert!(!is_valid_language_for_models("en-FOO")); + } +} + diff --git a/crates/lib/container_common/src/lib.rs b/crates/lib/container_common/src/lib.rs new file mode 100644 index 0000000000..2e45d4f08b --- /dev/null +++ b/crates/lib/container_common/src/lib.rs @@ -0,0 +1,6 @@ +pub mod anyhow_result; +pub mod collections; +pub mod i18n; +pub mod numerics; +pub mod thread; +pub mod token; \ No newline at end of file diff --git a/crates/lib/container_common/src/numerics/mod.rs b/crates/lib/container_common/src/numerics/mod.rs new file mode 100644 index 0000000000..9a87330c06 --- /dev/null +++ b/crates/lib/container_common/src/numerics/mod.rs @@ -0,0 +1 @@ +pub mod signed_to_unsigned; \ No newline at end of file diff --git a/crates/lib/container_common/src/numerics/signed_to_unsigned.rs b/crates/lib/container_common/src/numerics/signed_to_unsigned.rs new file mode 100644 index 0000000000..205bb62b57 --- /dev/null +++ b/crates/lib/container_common/src/numerics/signed_to_unsigned.rs @@ -0,0 +1,127 @@ + +// TODO: Use macros to define these (or use a better numeric crate) + +/// Negative numbers coalesce to 0 +pub fn i8_to_unsigned_zeroing_negatives(num: i8) -> u8 { + if num.is_negative() { + 0 + } else { + num.unsigned_abs() + } +} + +/// Negative numbers coalesce to 0 +pub fn i16_to_unsigned_zeroing_negatives(num: i16) -> u16 { + if num.is_negative() { + 0 + } else { + num.unsigned_abs() + } +} + +/// Negative numbers coalesce to 0 +pub fn i32_to_unsigned_zeroing_negatives(num: i32) -> u32 { + if num.is_negative() { + 0 + } else { + num.unsigned_abs() + } +} + +/// Negative numbers coalesce to 0 +pub fn i64_to_unsigned_zeroing_negatives(num: i64) -> u64 { + if num.is_negative() { + 0 + } else { + num.unsigned_abs() + } +} + +/// Negative numbers coalesce to 0 +pub fn i128_to_unsigned_zeroing_negatives(num: i128) -> u128 { + if num.is_negative() { + 0 + } else { + num.unsigned_abs() + } +} + +/// Negative numbers coalesce to 0 +pub fn isize_to_unsigned_zeroing_negatives(num: isize) -> usize { + if num.is_negative() { + 0 + } else { + num.unsigned_abs() + } +} + +#[cfg(test)] +mod tests { + use crate::numerics::signed_to_unsigned::i8_to_unsigned_zeroing_negatives; + use crate::numerics::signed_to_unsigned::i16_to_unsigned_zeroing_negatives; + use crate::numerics::signed_to_unsigned::i32_to_unsigned_zeroing_negatives; + use crate::numerics::signed_to_unsigned::i64_to_unsigned_zeroing_negatives; + use crate::numerics::signed_to_unsigned::i128_to_unsigned_zeroing_negatives; + use crate::numerics::signed_to_unsigned::isize_to_unsigned_zeroing_negatives; + + #[test] + fn i8_to_unsigned() { + assert_eq!(i8_to_unsigned_zeroing_negatives(-127), 0); + assert_eq!(i8_to_unsigned_zeroing_negatives(-1), 0); + assert_eq!(i8_to_unsigned_zeroing_negatives(0), 0); + assert_eq!(i8_to_unsigned_zeroing_negatives(1), 1); + assert_eq!(i8_to_unsigned_zeroing_negatives(100), 100); + assert_eq!(i8_to_unsigned_zeroing_negatives(127), 127); + } + + #[test] + fn i16_to_unsigned() { + assert_eq!(i16_to_unsigned_zeroing_negatives(-127), 0); + assert_eq!(i16_to_unsigned_zeroing_negatives(-1), 0); + assert_eq!(i16_to_unsigned_zeroing_negatives(0), 0); + assert_eq!(i16_to_unsigned_zeroing_negatives(1), 1); + assert_eq!(i16_to_unsigned_zeroing_negatives(100), 100); + assert_eq!(i16_to_unsigned_zeroing_negatives(127), 127); + } + + #[test] + fn i32_to_unsigned() { + assert_eq!(i32_to_unsigned_zeroing_negatives(-127), 0); + assert_eq!(i32_to_unsigned_zeroing_negatives(-1), 0); + assert_eq!(i32_to_unsigned_zeroing_negatives(0), 0); + assert_eq!(i32_to_unsigned_zeroing_negatives(1), 1); + assert_eq!(i32_to_unsigned_zeroing_negatives(100), 100); + assert_eq!(i32_to_unsigned_zeroing_negatives(127), 127); + } + + #[test] + fn i64_to_unsigned() { + assert_eq!(i64_to_unsigned_zeroing_negatives(-127), 0); + assert_eq!(i64_to_unsigned_zeroing_negatives(-1), 0); + assert_eq!(i64_to_unsigned_zeroing_negatives(0), 0); + assert_eq!(i64_to_unsigned_zeroing_negatives(1), 1); + assert_eq!(i64_to_unsigned_zeroing_negatives(100), 100); + assert_eq!(i64_to_unsigned_zeroing_negatives(127), 127); + } + + + #[test] + fn i128_to_unsigned() { + assert_eq!(i128_to_unsigned_zeroing_negatives(-127), 0); + assert_eq!(i128_to_unsigned_zeroing_negatives(-1), 0); + assert_eq!(i128_to_unsigned_zeroing_negatives(0), 0); + assert_eq!(i128_to_unsigned_zeroing_negatives(1), 1); + assert_eq!(i128_to_unsigned_zeroing_negatives(100), 100); + assert_eq!(i128_to_unsigned_zeroing_negatives(127), 127); + } + + #[test] + fn isize_to_unsigned() { + assert_eq!(isize_to_unsigned_zeroing_negatives(-127), 0); + assert_eq!(isize_to_unsigned_zeroing_negatives(-1), 0); + assert_eq!(isize_to_unsigned_zeroing_negatives(0), 0); + assert_eq!(isize_to_unsigned_zeroing_negatives(1), 1); + assert_eq!(isize_to_unsigned_zeroing_negatives(100), 100); + assert_eq!(isize_to_unsigned_zeroing_negatives(127), 127); + } +} \ No newline at end of file diff --git a/crates/lib/container_common/src/thread/mod.rs b/crates/lib/container_common/src/thread/mod.rs new file mode 100644 index 0000000000..4a3d1cc9ff --- /dev/null +++ b/crates/lib/container_common/src/thread/mod.rs @@ -0,0 +1 @@ +pub mod thread_id; \ No newline at end of file diff --git a/crates/lib/container_common/src/thread/thread_id.rs b/crates/lib/container_common/src/thread/thread_id.rs new file mode 100644 index 0000000000..8b33d508e6 --- /dev/null +++ b/crates/lib/container_common/src/thread/thread_id.rs @@ -0,0 +1,59 @@ +use crate::token::random_crockford_token::random_crockford_token; + +/// Thread string identifiers +/// Used for leases, logging, etc. +/// This has nothing to do with actual OS thread names or IDs. +/// The `std::thread::ThreadId` type might be better for some cases. +#[derive(Eq, PartialEq, Debug, Clone)] +pub struct ThreadId { + id: String, +} + +impl ThreadId { + pub fn random_id() -> Self { + Self { + id: random_crockford_token(10), + } + } + + pub fn with_id(id: &str) -> Self { + Self { + id: id.to_string(), + } + } + + pub fn get_id(&self) -> &str { + &self.id + } +} + +#[cfg(test)] +mod tests { + use crate::thread::thread_id::ThreadId; + + #[test] + fn equals() { + let a = ThreadId::with_id("foo"); + let b = ThreadId::with_id("foo"); + assert_eq!(a, b); + } + + #[test] + fn not_equals() { + let a = ThreadId::with_id("foo"); + let b = ThreadId::with_id("bar"); + assert_ne!(a, b); + } + + #[test] + fn manual_id() { + let t = ThreadId::with_id("foo"); + assert_eq!("foo", t.get_id()); + } + + #[test] + fn random_id() { + let t = ThreadId::random_id(); + assert_eq!(t.get_id().len(), 10); + } +} \ No newline at end of file diff --git a/crates/lib/container_common/src/token/mod.rs b/crates/lib/container_common/src/token/mod.rs new file mode 100644 index 0000000000..023a5ee511 --- /dev/null +++ b/crates/lib/container_common/src/token/mod.rs @@ -0,0 +1,3 @@ +pub mod random_crockford_token; +pub mod random_prefix_crockford_token; +pub mod random_uuid; \ No newline at end of file diff --git a/crates/lib/container_common/src/token/random_crockford_token.rs b/crates/lib/container_common/src/token/random_crockford_token.rs new file mode 100644 index 0000000000..a7cdd5346c --- /dev/null +++ b/crates/lib/container_common/src/token/random_crockford_token.rs @@ -0,0 +1,30 @@ +use rand::Rng; + +// Crockford characters +const CROCKFORD_CHARSET: &[u8] = b"0123456789ABCDEFGHJKMNPQRSTVWXYZ"; + +pub fn random_crockford_token(length: usize) -> String { + let mut rng = rand::thread_rng(); + + let token: String = (0..length) + .map(|_| { + let idx = rng.gen_range(0..CROCKFORD_CHARSET.len()); + CROCKFORD_CHARSET[idx] as char + }) + .collect(); + + token +} + +#[cfg(test)] +mod tests { + use crate::token::random_crockford_token::random_crockford_token; + + #[test] + fn random_token_length() { + assert_eq!(random_crockford_token(0), "".to_string()); + assert_eq!(random_crockford_token(1).len(), 1); + assert_eq!(random_crockford_token(10).len(), 10); + assert_eq!(random_crockford_token(32).len(), 32); + } +} diff --git a/crates/lib/container_common/src/token/random_prefix_crockford_token.rs b/crates/lib/container_common/src/token/random_prefix_crockford_token.rs new file mode 100644 index 0000000000..0d9508ce65 --- /dev/null +++ b/crates/lib/container_common/src/token/random_prefix_crockford_token.rs @@ -0,0 +1,98 @@ +use anyhow::anyhow; +use crate::anyhow_result::AnyhowResult; +use rand::Rng; + +// Crockford characters +const CROCKFORD_UPPERCASE_CHARSET: &[u8] = b"0123456789ABCDEFGHJKMNPQRSTVWXYZ"; + +const CROCKFORD_LOWERCASE_CHARSET: &[u8] = b"0123456789abcdefghjkmnpqrstvwxyz"; + +// Random part can't be any less than this. +const MIN_ENTROPY_LENGTH : usize = 8; + +pub fn random_prefix_crockford_token(prefix: &str, length: usize, uppercase: bool) -> AnyhowResult { + let mut prefix = prefix.trim().to_string(); + + if prefix.is_empty() { + return Err(anyhow!("prefix not set; use other crockford functions")); + } + + if !prefix.ends_with(":") { + prefix = format!("{}:", prefix); + } + + let entropy_length : usize = length.saturating_sub(prefix.len()); + + if entropy_length < MIN_ENTROPY_LENGTH { + return Err(anyhow!("not enough entropy in token")); + } + + let mut rng = rand::thread_rng(); + + let charset = if uppercase { CROCKFORD_UPPERCASE_CHARSET } else { CROCKFORD_LOWERCASE_CHARSET }; + + let entropy_part: String = (0..entropy_length) + .map(|_| { + let idx = rng.gen_range(0..charset.len()); + charset[idx] as char + }) + .collect(); + + let token = format!("{}{}", prefix, entropy_part); + Ok(token) +} + +#[cfg(test)] +mod tests { + use crate::token::random_prefix_crockford_token::random_prefix_crockford_token; + + #[test] + fn no_prefix_error() { + assert!(random_prefix_crockford_token("", 0, false).is_err()); + assert!(random_prefix_crockford_token("", 10, false).is_err()); + assert!(random_prefix_crockford_token("", 32, false).is_err()); + assert!(random_prefix_crockford_token("", 0, true).is_err()); + assert!(random_prefix_crockford_token("", 10, true).is_err()); + assert!(random_prefix_crockford_token("", 32, true).is_err()); + } + + #[test] + fn not_enough_entropy_error() { + assert!(random_prefix_crockford_token("foo", 0, false).is_err()); + assert!(random_prefix_crockford_token("foo", 5, false).is_err()); + assert!(random_prefix_crockford_token("foo", 11, false).is_err()); + assert!(random_prefix_crockford_token("foo:", 11, false).is_err()); + assert!(random_prefix_crockford_token("foo", 0, true).is_err()); + assert!(random_prefix_crockford_token("foo", 5, true).is_err()); + assert!(random_prefix_crockford_token("foo", 11, true).is_err()); + assert!(random_prefix_crockford_token("foo:", 11, true).is_err()); + } + + #[test] + fn starts_with() { + assert!(random_prefix_crockford_token("w", 10, false).unwrap().starts_with("w:")); + assert!(random_prefix_crockford_token("w", 10, false).unwrap().starts_with("w:")); + assert!(random_prefix_crockford_token("prefix", 16, false).unwrap().starts_with("prefix:")); + assert!(random_prefix_crockford_token("prefix:", 16, false).unwrap().starts_with("prefix:")); + assert!(random_prefix_crockford_token("w", 10, true).unwrap().starts_with("w:")); + assert!(random_prefix_crockford_token("w", 10, true).unwrap().starts_with("w:")); + assert!(random_prefix_crockford_token("prefix", 16, true).unwrap().starts_with("prefix:")); + assert!(random_prefix_crockford_token("prefix:", 16, true).unwrap().starts_with("prefix:")); + } + + #[test] + fn length() { + assert_eq!(random_prefix_crockford_token("w", 10, false).unwrap().len(), 10); + assert_eq!(random_prefix_crockford_token("w:", 10, false).unwrap().len(), 10); + assert_eq!(random_prefix_crockford_token("prefix", 16, false).unwrap().len(), 16); + assert_eq!(random_prefix_crockford_token("prefix:", 16, false).unwrap().len(), 16); + assert_eq!(random_prefix_crockford_token("asdf:", 32, false).unwrap().len(), 32); + assert_eq!(random_prefix_crockford_token("asdf", 32, false).unwrap().len(), 32); + assert_eq!(random_prefix_crockford_token("w", 10, true).unwrap().len(), 10); + assert_eq!(random_prefix_crockford_token("w:", 10, true).unwrap().len(), 10); + assert_eq!(random_prefix_crockford_token("prefix", 16, true).unwrap().len(), 16); + assert_eq!(random_prefix_crockford_token("prefix:", 16, true).unwrap().len(), 16); + assert_eq!(random_prefix_crockford_token("asdf:", 32, true).unwrap().len(), 32); + assert_eq!(random_prefix_crockford_token("asdf", 32, true).unwrap().len(), 32); + } +} diff --git a/crates/lib/container_common/src/token/random_uuid.rs b/crates/lib/container_common/src/token/random_uuid.rs new file mode 100644 index 0000000000..806fa27145 --- /dev/null +++ b/crates/lib/container_common/src/token/random_uuid.rs @@ -0,0 +1,9 @@ +use uuid::Uuid; + +pub fn generate_random_uuid() -> String { + let uuid = Uuid::new_v4(); + let mut buffer = Uuid::encode_buffer(); + let hyphenated_uuid = uuid.to_hyphenated() + .encode_lower(&mut buffer); + hyphenated_uuid.to_string() +} diff --git a/crates/lib/database_queries/Cargo.toml b/crates/lib/database_queries/Cargo.toml new file mode 100644 index 0000000000..ab092d51d2 --- /dev/null +++ b/crates/lib/database_queries/Cargo.toml @@ -0,0 +1,44 @@ +[package] +name = "database_queries" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "database_queries" +path = "src/lib.rs" + +[dependencies] + +# Internal +config = { path = "../config" } +container_common = { path = "../container_common" } + +# External +anyhow = "1.0.40" +base64 = "0.13.0" +chrono = { version = "0.4.19", features = ["serde"] } +dotenv = "0.15.0" +log = "0.4.14" + +# Needed to safety check some queries we can't use escaping for +once_cell = "1.9.0" +regex = "1.5.4" + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" +toml = "0.5.8" + +# MySql Database +# "offline" supports typesafe builds without a live database +# "rustls" doesn't appear to work with IP addresses, per: +# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828 +#sqlx = { version = "0.4", features = [ "mysql", "runtime-tokio-native-tls", "offline", "chrono" ] } +sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] } + diff --git a/crates/lib/database_queries/src/builders.rs b/crates/lib/database_queries/src/builders.rs new file mode 100644 index 0000000000..5738718cef --- /dev/null +++ b/crates/lib/database_queries/src/builders.rs @@ -0,0 +1,5 @@ + +/// This is simply a semantic / literate way of saying +/// an Optional should be required. There should be a +/// better way to actually enforce this, though. +pub (crate) type RequiredOption = Option; diff --git a/crates/lib/database_queries/src/column_types/job_status.rs b/crates/lib/database_queries/src/column_types/job_status.rs new file mode 100644 index 0000000000..e3606e50ca --- /dev/null +++ b/crates/lib/database_queries/src/column_types/job_status.rs @@ -0,0 +1,44 @@ +//! These are columns where users can control the visibility of their data. + +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; + +/// To use this in a query, the query must have type annotations. +/// See: https://www.gitmemory.com/issue/launchbadge/sqlx/1241/847154375 +/// eg. job_status as `job_status: crate::column_types::job_status::JobStatus` +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "snake_case")] +#[sqlx(rename_all = "snake_case")] +pub enum JobStatus { + Pending, + Started, + CompleteSuccess, + CompleteFailure, + AttemptFailed, + Dead, +} + +impl JobStatus { + pub fn to_str(&self) -> &'static str { + match self { + Self::Pending => "pending", + Self::Started => "started", + Self::CompleteSuccess => "complete_success", + Self::CompleteFailure => "complete_failure", + Self::AttemptFailed => "attempt_failed", + Self::Dead => "dead", + } + } + + pub fn from_str(job_status: &str) -> AnyhowResult { + match job_status { + "pending" => Ok(Self::Pending), + "started" => Ok(Self::Started), + "complete_success" => Ok(Self::CompleteSuccess), + "complete_failure" => Ok(Self::CompleteFailure), + "attempt_failed" => Ok(Self::AttemptFailed), + "dead" => Ok(Self::Dead), + _ => Err(anyhow!("invalid job_status: {:?}", job_status)), + } + } +} diff --git a/crates/lib/database_queries/src/column_types/mod.rs b/crates/lib/database_queries/src/column_types/mod.rs new file mode 100644 index 0000000000..9b54f0c4a4 --- /dev/null +++ b/crates/lib/database_queries/src/column_types/mod.rs @@ -0,0 +1,5 @@ +pub mod job_status; +pub mod record_visibility; +pub mod tts_model_type; +pub mod twitch_event_category; +pub mod vocoder_type; diff --git a/crates/lib/database_queries/src/column_types/record_visibility.rs b/crates/lib/database_queries/src/column_types/record_visibility.rs new file mode 100644 index 0000000000..3f1ee15637 --- /dev/null +++ b/crates/lib/database_queries/src/column_types/record_visibility.rs @@ -0,0 +1,35 @@ +//! These are columns where users can control the visibility of their data. + +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; + +/// To use this in a query, the query must have type annotations. +/// See: https://www.gitmemory.com/issue/launchbadge/sqlx/1241/847154375 +/// eg. preferred_tts_result_visibility as `preferred_tts_result_visibility: crate::column_types::record_visibility::RecordVisibility` +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "lowercase")] +#[sqlx(rename_all = "lowercase")] +pub enum RecordVisibility { + Public, + Hidden, + Private, +} + +impl RecordVisibility { + pub fn to_str(&self) -> &'static str { + match self { + RecordVisibility::Public => "public", + RecordVisibility::Hidden => "hidden", + RecordVisibility::Private => "private", + } + } + + pub fn from_str(record_visibility: &str) -> AnyhowResult { + match record_visibility { + "public" => Ok(RecordVisibility::Public), + "hidden" => Ok(RecordVisibility::Hidden), + "private" => Ok(RecordVisibility::Private), + _ => Err(anyhow!("invalid value: {:?}", record_visibility)), + } + } +} diff --git a/crates/lib/database_queries/src/column_types/tts_model_type.rs b/crates/lib/database_queries/src/column_types/tts_model_type.rs new file mode 100644 index 0000000000..3492d744b3 --- /dev/null +++ b/crates/lib/database_queries/src/column_types/tts_model_type.rs @@ -0,0 +1,40 @@ +//! These denote the TTS synthesizer model types + +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "lowercase")] +#[sqlx(rename_all = "lowercase")] +pub enum TtsModelType { + #[serde(rename = "tacotron2")] + Tacotron2, + + #[serde(rename = "talknet")] + Talknet, + + // Currently unused enum values: + //#[serde(rename = "not-set")] + //NotSet, + //#[serde(rename = "glowtts")] + //GlowTts, + //#[serde(rename = "glowtts-vocodes")] + //GlowTtsVocodes, +} + +impl TtsModelType { + pub fn to_str(&self) -> &'static str { + match self { + TtsModelType::Tacotron2 => "tacotron2", + TtsModelType::Talknet => "talknet", + } + } + + pub fn from_str(vocoder_type: &str) -> AnyhowResult { + match vocoder_type { + "tacotron2" => Ok(TtsModelType::Tacotron2), + "talknet" => Ok(TtsModelType::Talknet), + _ => Err(anyhow!("invalid value: {:?}", vocoder_type)), + } + } +} diff --git a/crates/lib/database_queries/src/column_types/twitch_event_category.rs b/crates/lib/database_queries/src/column_types/twitch_event_category.rs new file mode 100644 index 0000000000..269b5d9ccf --- /dev/null +++ b/crates/lib/database_queries/src/column_types/twitch_event_category.rs @@ -0,0 +1,35 @@ +//! These are columns where users can control the visibility of their data. + +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; + +/// To use this in a query, the query must have type annotations. +/// See: https://www.gitmemory.com/issue/launchbadge/sqlx/1241/847154375 +/// eg. twitch_event_category as `twitch_event_category: crate::column_types::twitch_event_category::TwitchEventCategory` +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "snake_case")] +#[sqlx(rename_all = "snake_case")] +pub enum TwitchEventCategory { + Bits, + ChannelPoints, + ChatCommand, +} + +impl TwitchEventCategory { + pub fn to_str(&self) -> &'static str { + match self { + Self::Bits => "bits", + Self::ChannelPoints => "channel_points", + Self::ChatCommand => "chat_command", + } + } + + pub fn from_str(value: &str) -> AnyhowResult { + match value { + "bits" => Ok(Self::Bits), + "channel_points" => Ok(Self::ChannelPoints), + "chat_command" => Ok(Self::ChatCommand), + _ => Err(anyhow!("invalid value: {:?}", value)), + } + } +} diff --git a/crates/lib/database_queries/src/column_types/vocoder_type.rs b/crates/lib/database_queries/src/column_types/vocoder_type.rs new file mode 100644 index 0000000000..ec63f25a61 --- /dev/null +++ b/crates/lib/database_queries/src/column_types/vocoder_type.rs @@ -0,0 +1,33 @@ +//! These denote the pretrained vocoders + +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; + +#[derive(Clone, Copy, PartialEq, Eq, Debug, Serialize, Deserialize, sqlx::Type)] +#[serde(rename_all = "lowercase")] +#[sqlx(rename_all = "lowercase")] +pub enum VocoderType { + /// Current state of the art vocoder + #[serde(rename = "hifigan-superres")] + HifiGanSuperResolution, + + /// Older WaveGlow vocoder + WaveGlow, +} + +impl VocoderType { + pub fn to_str(&self) -> &'static str { + match self { + VocoderType::HifiGanSuperResolution => "hifigan-superres", + VocoderType::WaveGlow => "waveglow", + } + } + + pub fn from_str(vocoder_type: &str) -> AnyhowResult { + match vocoder_type { + "hifigan-superres" => Ok(VocoderType::HifiGanSuperResolution), + "waveglow" => Ok(VocoderType::WaveGlow), + _ => Err(anyhow!("invalid value: {:?}", vocoder_type)), + } + } +} diff --git a/crates/lib/database_queries/src/complex_models/event_match_predicate.rs b/crates/lib/database_queries/src/complex_models/event_match_predicate.rs new file mode 100644 index 0000000000..d989b4ca11 --- /dev/null +++ b/crates/lib/database_queries/src/complex_models/event_match_predicate.rs @@ -0,0 +1,62 @@ + +#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)] +#[serde(rename_all = "snake_case")] +pub enum EventMatchPredicate { + /// Default value + /// NB: This has an empty struct so it serializes as JSON and not a String! + NotSet {}, + + /// A predicate for an exactly matched cheermote name, including the bit value, + /// eg "Cheer1". Case insensitive. + BitsCheermoteNameExactMatch { + cheermote_name: String, + }, + + /// A predicate for a cheermote name without value, but with a minimum bits + /// threshold. eg "Cheer", minimum=150 bits. Case insensitive. + /// Twitch also calls the leading part "prefix" + BitsCheermotePrefixSpendThreshold { + cheermote_prefix: String, + minimum_bits_spent: u64, + }, + + /// Spend a minimum amount of bits. + BitsSpendThreshold { + minimum_bits_spent: u64, + }, + + ChannelPointsRewardNameExactMatch { + reward_name: String, + }, +} + +#[cfg(test)] +mod tests { + use crate::complex_models::event_match_predicate::EventMatchPredicate; + + #[test] + fn not_set() { + let rust_value = EventMatchPredicate::NotSet {}; + let json = "{\"not_set\":{}}"; + + let converted_to_json= serde_json::to_string(&rust_value).unwrap(); + assert_eq!(&converted_to_json, json); + + let converted_from_json : EventMatchPredicate = serde_json::from_str(json).unwrap(); + assert_eq!(&converted_from_json, &rust_value); + } + + #[test] + fn bits_cheermote_name_exact_match() { + let rust_value = EventMatchPredicate::BitsCheermoteNameExactMatch { + cheermote_name: "Cheer1".to_string() + }; + let json = "{\"bits_cheermote_name_exact_match\":{\"cheermote_name\":\"Cheer1\"}}"; + + let converted_to_json= serde_json::to_string(&rust_value).unwrap(); + assert_eq!(&converted_to_json, json); + + let converted_from_json : EventMatchPredicate = serde_json::from_str(json).unwrap(); + assert_eq!(&converted_from_json, &rust_value); + } +} diff --git a/crates/lib/database_queries/src/complex_models/event_responses.rs b/crates/lib/database_queries/src/complex_models/event_responses.rs new file mode 100644 index 0000000000..4c230e9545 --- /dev/null +++ b/crates/lib/database_queries/src/complex_models/event_responses.rs @@ -0,0 +1,76 @@ +use container_common::anyhow_result::AnyhowResult; +use std::collections::HashMap; + +/// How TTS commands should be prefixed: !voiceName, /voiceName, etc. +#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone, Copy)] +#[serde(rename_all = "snake_case")] +pub enum CommandPrefixType { + /// Prefix voice commands with "!" + Exclamation, + /// Prefix voice commands with "!" or "/" + ExclamationOrSlash, + /// Prefix voice commands with "/" + Slash, +} + +#[derive(Serialize, Deserialize, Debug, PartialEq, Eq, Clone)] +#[serde(rename_all = "snake_case")] +pub enum EventResponse { + /// Default value + /// NB: This has an empty struct so it serializes as JSON and not a String! + NotSet {}, + + /// Respond with a single TTS voice. + TtsSingleVoice { + tts_model_token: String, + }, + + /// Respond with a random TTS voice. + TtsRandomVoice { + tts_model_tokens: Vec, + }, + + /// Use TTS commands that FakeYou recommends. + TtsCommandPresets { + command_prefix_type: CommandPrefixType, + }, + + /// Use your own TTS commands. + /// The map is from "!command" (without the "/" or "!" prefix) to model token. + /// Command is lower case, alphanum only, and compared case insensitive. + TtsCommandCustom { + command_prefix_type: CommandPrefixType, + command_map: HashMap, + }, +} + +#[cfg(test)] +mod tests { + use crate::complex_models::event_responses::EventResponse; + + #[test] + fn not_set() { + let rust_value = EventResponse::NotSet {}; + let json = "{\"not_set\":{}}"; + + let converted_to_json= serde_json::to_string(&rust_value).unwrap(); + assert_eq!(&converted_to_json, json); + + let converted_from_json : EventResponse = serde_json::from_str(json).unwrap(); + assert_eq!(&converted_from_json, &rust_value); + } + + #[test] + fn tts_single_voice() { + let rust_value = EventResponse::TtsSingleVoice { + tts_model_token: "foo".to_string() + }; + let json = "{\"tts_single_voice\":{\"tts_model_token\":\"foo\"}}"; + + let converted_to_json= serde_json::to_string(&rust_value).unwrap(); + assert_eq!(&converted_to_json, json); + + let converted_from_json : EventResponse = serde_json::from_str(json).unwrap(); + assert_eq!(&converted_from_json, &rust_value); + } +} diff --git a/crates/lib/database_queries/src/complex_models/mod.rs b/crates/lib/database_queries/src/complex_models/mod.rs new file mode 100644 index 0000000000..ed032a6422 --- /dev/null +++ b/crates/lib/database_queries/src/complex_models/mod.rs @@ -0,0 +1,2 @@ +pub mod event_match_predicate; +pub mod event_responses; diff --git a/crates/lib/database_queries/src/helpers/boolean_converters.rs b/crates/lib/database_queries/src/helpers/boolean_converters.rs new file mode 100644 index 0000000000..c35678a450 --- /dev/null +++ b/crates/lib/database_queries/src/helpers/boolean_converters.rs @@ -0,0 +1,62 @@ + +/// MySQL stores non-nullable booleans as i8. +pub fn i8_to_bool(value: i8) -> bool { + if value == 0 { + false + } else { + true + } +} + +/// Bool conversion, but turn nulls to a default value. +pub fn nullable_i8_to_bool(value: Option, default_value: bool) -> bool { + value.map(|v| i8_to_bool(v)) + .unwrap_or(default_value) +} + +/// Bool conversion, but retain nulls. +pub fn nullable_i8_to_optional_bool(value: Option) -> Option { + value.map(|v| i8_to_bool(v)) +} + +#[cfg(test)] +mod tests { + use crate::helpers::boolean_converters::{ + i8_to_bool, + nullable_i8_to_bool, + nullable_i8_to_optional_bool + }; + + #[test] + fn test_i8_to_bool() { + assert_eq!(i8_to_bool(0), false); + assert_eq!(i8_to_bool(1), true); + assert_eq!(i8_to_bool(-1), true); + assert_eq!(i8_to_bool(120), true); + } + + #[test] + fn test_nullable_i8_to_bool() { + assert_eq!(nullable_i8_to_bool(None, false), false); + assert_eq!(nullable_i8_to_bool(None, true), true); + assert_eq!(nullable_i8_to_bool(Some(0), false), false); + assert_eq!(nullable_i8_to_bool(Some(0), true), false); + + assert_eq!(nullable_i8_to_bool(Some(1), false), true); + assert_eq!(nullable_i8_to_bool(Some(1), true), true); + + assert_eq!(nullable_i8_to_bool(Some(100), false), true); + assert_eq!(nullable_i8_to_bool(Some(100), true), true); + assert_eq!(nullable_i8_to_bool(Some(-100), false), true); + assert_eq!(nullable_i8_to_bool(Some(-100), true), true); + } + + #[test] + fn test_nullable_i8_to_optional_bool() { + assert_eq!(nullable_i8_to_optional_bool(None), None); + assert_eq!(nullable_i8_to_optional_bool(Some(0)), Some(false)); + assert_eq!(nullable_i8_to_optional_bool(Some(1)), Some(true)); + assert_eq!(nullable_i8_to_optional_bool(Some(100)), Some(true)); + assert_eq!(nullable_i8_to_optional_bool(Some(-100)), Some(true)); + } +} diff --git a/crates/lib/database_queries/src/helpers/mod.rs b/crates/lib/database_queries/src/helpers/mod.rs new file mode 100644 index 0000000000..62d38a8df4 --- /dev/null +++ b/crates/lib/database_queries/src/helpers/mod.rs @@ -0,0 +1 @@ +pub mod boolean_converters; \ No newline at end of file diff --git a/crates/lib/database_queries/src/lib.rs b/crates/lib/database_queries/src/lib.rs new file mode 100644 index 0000000000..db1243c4eb --- /dev/null +++ b/crates/lib/database_queries/src/lib.rs @@ -0,0 +1,10 @@ +#[macro_use] extern crate serde_derive; + +mod builders; + +pub mod column_types; +pub mod complex_models; +pub mod helpers; +pub mod mediators; +pub mod queries; +pub mod tokens; diff --git a/crates/lib/database_queries/src/mediators/badge_granter.rs b/crates/lib/database_queries/src/mediators/badge_granter.rs new file mode 100644 index 0000000000..d771744beb --- /dev/null +++ b/crates/lib/database_queries/src/mediators/badge_granter.rs @@ -0,0 +1,364 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::mediators::firehose_publisher::FirehosePublisher; +use log::{warn,info}; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlQueryResult; +use sqlx::{MySqlPool}; +use std::sync::Arc; + +#[derive(Debug, Clone, Copy)] +pub enum UserBadge { + // Granted for early vocodes users. + EarlyUser, + + // Granted for uploading models + TtsModelUploader1, + TtsModelUploader5, + TtsModelUploader10, + TtsModelUploader20, + TtsModelUploader50, + TtsModelUploader100, + TtsModelUploader150, + TtsModelUploader200, + TtsModelUploader250, + + // Granted for uploading templates + W2lTemplateUploader1, + W2lTemplateUploader10, + W2lTemplateUploader50, + W2lTemplateUploader100, + W2lTemplateUploader200, + W2lTemplateUploader500, + W2lTemplateUploader1000, + W2lTemplateUploader2000, + W2lTemplateUploader5000, + W2lTemplateUploader10000, + + //// Granted for using models + //TtsInferenceUser100, + //TtsInferenceUser500, + //TtsInferenceUser1000, + //TtsInferenceUser10000, + //TtsInferenceUser100000, + + //// Granted for using templates + //W2lInferenceUser100, + //W2lInferenceUser500, + //W2lInferenceUser1000, + //W2lInferenceUser10000, + //W2lInferenceUser100000, +} + +impl UserBadge { + pub fn to_db_value(&self) -> &'static str { + match self { + UserBadge::EarlyUser => "early_user", + UserBadge::TtsModelUploader1 => "tts_model_uploader_1", + UserBadge::TtsModelUploader5 => "tts_model_uploader_5", + UserBadge::TtsModelUploader10 => "tts_model_uploader_10", + UserBadge::TtsModelUploader20 => "tts_model_uploader_20", + UserBadge::TtsModelUploader50 => "tts_model_uploader_50", + UserBadge::TtsModelUploader100 => "tts_model_uploader_100", + UserBadge::TtsModelUploader150 => "tts_model_uploader_150", + UserBadge::TtsModelUploader200 => "tts_model_uploader_200", + UserBadge::TtsModelUploader250 => "tts_model_uploader_250", + UserBadge::W2lTemplateUploader1 => "w2l_template_uploader_1", + UserBadge::W2lTemplateUploader10 => "w2l_template_uploader_10", + UserBadge::W2lTemplateUploader50 => "w2l_template_uploader_50", + UserBadge::W2lTemplateUploader100 => "w2l_template_uploader_100", + UserBadge::W2lTemplateUploader200 => "w2l_template_uploader_200", + UserBadge::W2lTemplateUploader500 => "w2l_template_uploader_500", + UserBadge::W2lTemplateUploader1000 => "w2l_template_uploader_1000", + UserBadge::W2lTemplateUploader2000 => "w2l_template_uploader_2000", + UserBadge::W2lTemplateUploader5000 => "w2l_template_uploader_5000", + UserBadge::W2lTemplateUploader10000 => "w2l_template_uploader_10000", + //UserBadge::TtsInferenceUser100 => "tts_inference_100", + //UserBadge::TtsInferenceUser500 => "tts_inference_500", + //UserBadge::TtsInferenceUser1000 => "tts_inference_1000", + //UserBadge::TtsInferenceUser10000 =>"tts_inference_10000", + //UserBadge::TtsInferenceUser100000 => "tts_inference_100000", + //UserBadge::W2lInferenceUser100 => "w2l_inference_100", + //UserBadge::W2lInferenceUser500 => "w2l_inference_500", + //UserBadge::W2lInferenceUser1000 => "w2l_inference_1000", + //UserBadge::W2lInferenceUser10000 => "w2l_inference_10000", + //UserBadge::W2lInferenceUser100000 => "w2l_inference_100000", + } + } +} + +#[derive(Clone)] +pub struct BadgeGranter { + pub mysql_pool: MySqlPool, + pub firehose_publisher: FirehosePublisher, // NB: Type is Copy/Clone safe due to internal Arc. +} + +struct ExistenceRecord { + pub does_exist: i64, +} + +struct CountRecord { + pub count: i64, +} + +impl BadgeGranter { + + pub async fn grant_early_user_badge(&self, user_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + UserBadge::EarlyUser, + user_token, + ).await?; + Ok(()) + } + + /// This needs to be called *after* successful upload. + pub async fn maybe_grant_tts_model_uploads_badge(&self, user_token: &str) -> AnyhowResult<()> { + let count = self.count_tts_models_uploaded(user_token).await?; + + let mut maybe_badge = None; + + if count >= 250 { + maybe_badge = Some(UserBadge::TtsModelUploader250); + } else if count >= 200 { + maybe_badge = Some(UserBadge::TtsModelUploader200); + } else if count >= 150 { + maybe_badge = Some(UserBadge::TtsModelUploader150); + } else if count >= 100 { + maybe_badge = Some(UserBadge::TtsModelUploader100); + } else if count >= 50 { + maybe_badge = Some(UserBadge::TtsModelUploader50); + } else if count >= 20 { + maybe_badge = Some(UserBadge::TtsModelUploader20); + } else if count >= 10 { + maybe_badge = Some(UserBadge::TtsModelUploader10); + } else if count >= 5 { + maybe_badge = Some(UserBadge::TtsModelUploader5); + } else if count >= 1 { + maybe_badge = Some(UserBadge::TtsModelUploader1); + } + + let badge = match maybe_badge { + Some(badge) => badge, + None => return Ok(()), + }; + + if self.has_badge(user_token, badge).await? { + return Ok(()) + } + + let _record_id = self.insert( + badge, + user_token, + ).await?; + + self.firehose_publisher.publish_user_badge_granted(user_token, badge.to_db_value()) + .await?; + + Ok(()) + } + + + /// This needs to be called *after* successful upload. + pub async fn maybe_grant_w2l_template_uploads_badge(&self, user_token: &str) -> AnyhowResult<()> { + let count = self.count_w2l_templates_uploaded(user_token).await?; + + let mut maybe_badge = None; + + if count >= 10000 { + maybe_badge = Some(UserBadge::W2lTemplateUploader10000); + } else if count >= 5000 { + maybe_badge = Some(UserBadge::W2lTemplateUploader5000); + } else if count >= 2000 { + maybe_badge = Some(UserBadge::W2lTemplateUploader2000); + } else if count >= 1000 { + maybe_badge = Some(UserBadge::W2lTemplateUploader1000); + } else if count >= 500 { + maybe_badge = Some(UserBadge::W2lTemplateUploader500); + } else if count >= 200 { + maybe_badge = Some(UserBadge::W2lTemplateUploader200); + } else if count >= 100 { + maybe_badge = Some(UserBadge::W2lTemplateUploader100); + } else if count >= 50 { + maybe_badge = Some(UserBadge::W2lTemplateUploader50); + } else if count >= 10 { + maybe_badge = Some(UserBadge::W2lTemplateUploader10); + } else if count >= 1 { + maybe_badge = Some(UserBadge::W2lTemplateUploader1); + } + + let badge = match maybe_badge { + Some(badge) => badge, + None => return Ok(()), + }; + + if self.has_badge(user_token, badge).await? { + return Ok(()) + } + + let _record_id = self.insert( + badge, + user_token, + ).await?; + + self.firehose_publisher.publish_user_badge_granted(user_token, badge.to_db_value()) + .await?; + + Ok(()) + } + + + // ======================================================================= + + pub async fn has_badge(&self, user_token: &str, user_badge: UserBadge) -> AnyhowResult { + let maybe_result = sqlx::query_as!( + ExistenceRecord, + r#" +SELECT 1 as does_exist +FROM user_badges +WHERE + user_token = ? +AND + badge_slug = ? +LIMIT 1 + "#, + user_token, + user_badge.to_db_value() + ) + .fetch_one(&self.mysql_pool) + .await; + + let exists = match maybe_result { + Ok(_record) => true, + Err(err) => { + match err { + sqlx::Error::RowNotFound => false, + _ => { + warn!("query error: {:?}", err); + return Err(anyhow!("error querying: {:?}", err)); + } + } + } + }; + + Ok(exists) + } + + // ======================================================================= + + async fn count_tts_models_uploaded( + &self, + user_token: &str, + ) -> AnyhowResult { + // NB: This could get expensive! + let maybe_result = sqlx::query_as!( + CountRecord, + r#" +SELECT count(*) as count +FROM tts_models +WHERE + creator_user_token = ? +LIMIT 1 + "#, + user_token + ) + .fetch_one(&self.mysql_pool) + .await; + + self.handle_count_query(maybe_result) + } + + async fn count_w2l_templates_uploaded( + &self, + user_token: &str, + ) -> AnyhowResult { + // NB: This could get expensive! + // Especially at the scale we'll likely have W2L templates. + let maybe_result = sqlx::query_as!( + CountRecord, + r#" +SELECT count(*) as count +FROM w2l_templates +WHERE + creator_user_token = ? +LIMIT 1 + "#, + user_token + ) + .fetch_one(&self.mysql_pool) + .await; + + self.handle_count_query(maybe_result) + } + + async fn insert( + &self, + user_badge: UserBadge, + user_token: &str, + ) -> AnyhowResult { + let query_result = sqlx::query!( + r#" +INSERT INTO user_badges +SET + user_token = ?, + badge_slug = ? + "#, + user_token, + user_badge.to_db_value(), + ) + .execute(&self.mysql_pool) + .await; + + let record_id = Self::handle_results(query_result)?; + Ok(record_id) + } + + fn handle_results(query_result: Result) -> AnyhowResult { + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("Insert badge record DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(anyhow!("Error inserting record")); + } + }; + + Ok(record_id) + } + + fn handle_count_query(&self, query_result: Result) -> AnyhowResult { + let count = match query_result { + Ok(record) => record.count as u64, + Err(err) => { + match err { + sqlx::Error::RowNotFound => 0, + _ => { + warn!("query error: {:?}", err); + return Err(anyhow!("error querying: {:?}", err)); + } + } + } + }; + + Ok(count) + } +} diff --git a/crates/lib/database_queries/src/mediators/firehose_publisher.rs b/crates/lib/database_queries/src/mediators/firehose_publisher.rs new file mode 100644 index 0000000000..2499a2b144 --- /dev/null +++ b/crates/lib/database_queries/src/mediators/firehose_publisher.rs @@ -0,0 +1,241 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::tokens::Tokens; +use log::{warn,info}; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlQueryResult; +use sqlx::{MySqlPool}; +use std::sync::Arc; + +#[derive(Debug, Clone, Copy)] +enum FirehoseEvent { + UserSignUp, + + // We don't publish for all badges (eg. early user signup) + UserBadgeGranted, + + TtsModelUploadStarted, + TtsModelUploadCompleted, + TtsInferenceStarted, + TtsInferenceCompleted, + + W2lTemplateUploadStarted, + W2lTemplateUploadCompleted, + W2lInferenceStarted, + W2lInferenceCompleted, + + TwitterMention, + TwitterRetweet, + DiscordJoin, + DiscordMessage, + TwitchSubscribe, + TwitchFollow, +} + +impl FirehoseEvent { + pub fn to_db_value(&self) -> &'static str { + match self { + FirehoseEvent::UserSignUp => "user_sign_up", + FirehoseEvent::UserBadgeGranted=> "user_badge_granted", + FirehoseEvent::TtsModelUploadStarted => "tts_model_upload_started", + FirehoseEvent::TtsModelUploadCompleted => "tts_model_upload_completed", + FirehoseEvent::TtsInferenceStarted => "tts_inference_started", + FirehoseEvent::TtsInferenceCompleted => "tts_inference_completed", + FirehoseEvent::W2lTemplateUploadStarted => "w2l_template_upload_started", + FirehoseEvent::W2lTemplateUploadCompleted => "w2l_template_upload_completed", + FirehoseEvent::W2lInferenceStarted => "w2l_inference_started", + FirehoseEvent::W2lInferenceCompleted => "w2l_inference_completed", + FirehoseEvent::TwitterMention => "twitter_mention", + FirehoseEvent::TwitterRetweet => "twitter_retweet", + FirehoseEvent::DiscordJoin => "discord_join", + FirehoseEvent::DiscordMessage => "discord_message", + FirehoseEvent::TwitchSubscribe => "twitch_subscribe", + FirehoseEvent::TwitchFollow => "twitch_follow", + } + } +} + +#[derive(Clone)] +pub struct FirehosePublisher { + pub mysql_pool: MySqlPool, +} + +impl FirehosePublisher { + + pub async fn publish_user_sign_up(&self, user_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::UserSignUp, + Some(user_token), + Some(user_token), + Some(user_token) + ).await?; + Ok(()) + } + + pub async fn publish_user_badge_granted(&self, user_token: &str, badge_slug: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::UserBadgeGranted, + Some(user_token), + Some(badge_slug), + None, + ).await?; + Ok(()) + } + + pub async fn enqueue_tts_model_upload(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsModelUploadStarted, + Some(user_token), + None, + Some(job_token) + ).await?; + Ok(()) + } + + pub async fn publish_tts_model_upload_finished(&self, user_token: &str, model_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsModelUploadCompleted, + Some(user_token), + Some(model_token), + Some(model_token) + ).await?; + Ok(()) + } + + pub async fn enqueue_tts_inference( + &self, + maybe_user_token: Option<&str>, + job_token: &str, + model_token: &str + ) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsInferenceStarted, + maybe_user_token, + Some(model_token), + Some(job_token) + ).await?; + Ok(()) + } + + pub async fn tts_inference_finished( + &self, + maybe_user_token: Option<&str>, + model_token: &str, + result_token: &str + ) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsInferenceCompleted, + maybe_user_token, + Some(model_token), + Some(result_token) + ).await?; + Ok(()) + } + + pub async fn enqueue_w2l_template_upload(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::W2lTemplateUploadStarted, + Some(user_token), + None, + Some(job_token) + ).await?; + Ok(()) + } + + pub async fn enqueue_w2l_inference(&self, maybe_user_token: Option<&str>, job_token: &str, template_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::W2lInferenceStarted, + maybe_user_token, + Some(template_token), + Some(job_token) + ).await?; + Ok(()) + } + + pub async fn publish_w2l_template_upload_finished(&self, user_token: &str, template_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::W2lTemplateUploadCompleted, + Some(user_token), + Some(template_token), + Some(template_token) + ).await?; + Ok(()) + } + + pub async fn w2l_inference_finished(&self, maybe_user_token: Option<&str>, job_token: &str, result_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::W2lInferenceCompleted, + maybe_user_token, + Some(job_token), // TODO: This could be template_token + Some(result_token) + ).await?; + Ok(()) + } + + // ======================================================================= + + async fn insert( + &self, + event_type: FirehoseEvent, + user_token: Option<&str>, + entity_token: Option<&str>, + created_entity_token: Option<&str> + ) -> AnyhowResult { + let token = Tokens::new_firehose_event()?; + + let query_result = sqlx::query!( + r#" +INSERT INTO firehose_entries +SET + token = ?, + event_type = ?, + maybe_target_user_token = ?, + maybe_target_entity_token = ?, + maybe_created_entity_token = ? + "#, + token, + event_type.to_db_value(), + user_token, + entity_token, + created_entity_token + ) + .execute(&self.mysql_pool) + .await; + + let record_id = Self::handle_results(query_result)?; + Ok(record_id) + } + + fn handle_results(query_result: Result) -> AnyhowResult { + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("Insert record DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(anyhow!("Error inserting record")); + } + }; + + Ok(record_id) + } +} diff --git a/crates/lib/database_queries/src/mediators/mod.rs b/crates/lib/database_queries/src/mediators/mod.rs new file mode 100644 index 0000000000..5cc9617e88 --- /dev/null +++ b/crates/lib/database_queries/src/mediators/mod.rs @@ -0,0 +1,2 @@ +pub mod badge_granter; +pub mod firehose_publisher; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/api_tokens/create_api_token.rs b/crates/lib/database_queries/src/queries/api_tokens/create_api_token.rs new file mode 100644 index 0000000000..50ee9b4ad7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/api_tokens/create_api_token.rs @@ -0,0 +1,94 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::api_tokens::list_available_api_tokens_for_user::list_available_api_tokens_for_user; +use crate::tokens::Tokens; +use sqlx::MySqlPool; +use sqlx; + +/// Create a new API token for the user. +/// If the user has more than five tokens, delete the least recent. +pub async fn create_api_token_for_user( + user_token: &str, + uuid_idempotency_token: &str, + creator_ip_address: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult { + + // TODO: Database transaction. + + let internal_token = Tokens::new_api_token_internal_token()?; + let api_token = Tokens::new_api_token()?; + + let query = sqlx::query!( + r#" +INSERT INTO api_tokens +SET + internal_token = ?, + api_token = ?, + user_token = ?, + uuid_idempotency_token = ?, + ip_address_creation = ?, + ip_address_last_update = ? + "#, + internal_token, + api_token, + user_token, + uuid_idempotency_token, + creator_ip_address, + creator_ip_address + ); + + let query_result = query.execute(mysql_pool).await; + + let _record_id = match query_result { + Ok(res) => res.last_insert_id(), + Err(err) => return Err(anyhow!("Error creating API token: {:?}", err)), + }; + + // Delete older tokens. + let valid_api_tokens= list_available_api_tokens_for_user(user_token, mysql_pool).await?; + + let valid_api_tokens = valid_api_tokens.into_iter() + .map(|r| r.api_token) + .collect::>(); + + // TODO: This is a massive hack since SQLx doesn't support binding Vec + // This array's elements will be replaced with any valid tokens. + let mut api_tokens = [ + api_token.as_str(), + api_token.as_str(), + api_token.as_str(), + api_token.as_str(), + api_token.as_str(), + api_token.as_str(), // One extra copy of the latest token. + ]; + + for i in 0..5 { + if let Some(token) = valid_api_tokens.get(i).map(|s| s.as_str()) { + api_tokens[i] = token; + } + } + + let query = sqlx::query!( + r#" +UPDATE api_tokens +SET + deleted_at = CURRENT_TIMESTAMP, + ip_address_last_update = ? +WHERE + user_token = ? +AND api_token NOT IN (?, ?, ?, ?, ?) + "#, + creator_ip_address, + user_token, + &api_tokens[0], + &api_tokens[1], + &api_tokens[2], + &api_tokens[3], + &api_tokens[4], + ); + + let _query_result = query.execute(mysql_pool).await?; + + Ok(api_token) +} \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/api_tokens/delete_api_token.rs b/crates/lib/database_queries/src/queries/api_tokens/delete_api_token.rs new file mode 100644 index 0000000000..7f8f310ba0 --- /dev/null +++ b/crates/lib/database_queries/src/queries/api_tokens/delete_api_token.rs @@ -0,0 +1,34 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +pub async fn delete_api_token( + user_token: &str, + api_token: &str, + ip_address_update: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult { + + let query = sqlx::query!( + r#" +UPDATE api_tokens +SET + deleted_at = CURRENT_TIMESTAMP, + ip_address_last_update = ? +WHERE + user_token = ? + AND api_token = ? +LIMIT 1 + "#, + ip_address_update, + user_token, + api_token, + ); + + let result = query.execute(mysql_pool).await; + + match result { + Err(err) => Err(anyhow!("error with query: {:?}", err)), + Ok(_r) => Ok(true), + } +} diff --git a/crates/lib/database_queries/src/queries/api_tokens/edit_api_token.rs b/crates/lib/database_queries/src/queries/api_tokens/edit_api_token.rs new file mode 100644 index 0000000000..909075c105 --- /dev/null +++ b/crates/lib/database_queries/src/queries/api_tokens/edit_api_token.rs @@ -0,0 +1,38 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; +use std::collections::{HashSet, HashMap}; + +pub async fn edit_api_token( + user_token: &str, + api_token: &str, + maybe_short_description: Option<&str>, + ip_address_update: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult { + + let query = sqlx::query!( + r#" +UPDATE api_tokens +SET + maybe_short_description = ?, + ip_address_last_update = ? +WHERE + user_token = ? + AND api_token = ? +LIMIT 1 + "#, + maybe_short_description, + ip_address_update, + user_token, + api_token, + ); + + let result = query.execute(mysql_pool).await; + + match result { + Err(err) => Err(anyhow!("error with query: {:?}", err)), + Ok(_r) => Ok(true), + } +} diff --git a/crates/lib/database_queries/src/queries/api_tokens/list_available_api_tokens_for_user.rs b/crates/lib/database_queries/src/queries/api_tokens/list_available_api_tokens_for_user.rs new file mode 100644 index 0000000000..034d24d131 --- /dev/null +++ b/crates/lib/database_queries/src/queries/api_tokens/list_available_api_tokens_for_user.rs @@ -0,0 +1,69 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +/// table: api_tokens +#[derive(Debug)] +pub struct AvailableApiToken { + pub internal_token: String, + pub api_token: String, + pub maybe_short_description: Option, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} + +/// Query available (non-deleted) API tokens for a user. +/// A user can only have five active tokens at a time, so we only return the five most recent. +pub async fn list_available_api_tokens_for_user( + user_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + + let records : Vec = sqlx::query_as!( + AvailableApiTokenInternal, + r#" +SELECT + internal_token, + api_token, + maybe_short_description, + created_at, + updated_at +FROM api_tokens +WHERE + user_token = ? + AND deleted_at IS NULL +ORDER BY created_at DESC +LIMIT 5 + "#, + user_token, + ) + .fetch_all(pool) + .await?; + + let mut records = records.into_iter() + .map(|record : AvailableApiTokenInternal | { + AvailableApiToken { + internal_token: record.internal_token, + api_token: record.api_token, + maybe_short_description: record.maybe_short_description, + created_at: record.created_at, + updated_at: record.updated_at, + } + }) + .collect::>(); + + // Queried in DESC order, but sort returned results ordered ASC. + records.sort_by_key(|r| r.created_at); + + Ok(records) +} + +#[derive(Debug)] +struct AvailableApiTokenInternal { + pub internal_token: String, + pub api_token: String, + pub maybe_short_description: Option, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} diff --git a/crates/lib/database_queries/src/queries/api_tokens/mod.rs b/crates/lib/database_queries/src/queries/api_tokens/mod.rs new file mode 100644 index 0000000000..2ce84cd568 --- /dev/null +++ b/crates/lib/database_queries/src/queries/api_tokens/mod.rs @@ -0,0 +1,4 @@ +pub mod create_api_token; +pub mod delete_api_token; +pub mod edit_api_token; +pub mod list_available_api_tokens_for_user; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/health_check/health_check_query.rs b/crates/lib/database_queries/src/queries/health_check/health_check_query.rs new file mode 100644 index 0000000000..847f77e20d --- /dev/null +++ b/crates/lib/database_queries/src/queries/health_check/health_check_query.rs @@ -0,0 +1,30 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime, NaiveDateTime}; +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +/// This just queries for current time. If it fails, we know the connection pool is dead. +pub struct HealthCheckResult { + pub present_time: NaiveDateTime, +} + +/// Query the DB for time as a proxy for DB health +pub async fn health_check_db( + pool: &MySqlPool, +) -> AnyhowResult { + + let maybe_record = sqlx::query_as!( + HealthCheckResult, + r#" +SELECT NOW() as present_time + "#, + ) + .fetch_one(pool) + .await; + + match maybe_record { + Ok(record) => Ok(record), + Err(sqlx::Error::RowNotFound) => Err(anyhow!("query didn't return anything")), + Err(ref err) => Err(anyhow!("database query error: {:?}", err)), + } +} diff --git a/crates/lib/database_queries/src/queries/health_check/mod.rs b/crates/lib/database_queries/src/queries/health_check/mod.rs new file mode 100644 index 0000000000..3712d08210 --- /dev/null +++ b/crates/lib/database_queries/src/queries/health_check/mod.rs @@ -0,0 +1 @@ +pub mod health_check_query; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/ip_bans/list_ip_bans.rs b/crates/lib/database_queries/src/queries/ip_bans/list_ip_bans.rs new file mode 100644 index 0000000000..4bff642506 --- /dev/null +++ b/crates/lib/database_queries/src/queries/ip_bans/list_ip_bans.rs @@ -0,0 +1,66 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct IpBanRecordForList { + pub ip_address: String, + pub maybe_target_user_token: Option, + pub maybe_target_username: Option, + pub mod_user_token: String, + pub mod_username: String, + pub mod_display_name: String, + pub mod_notes: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +// NB: This is unconstrained because I don't anticipate a huge number of IP bans at launch +// We can reengineer with growth. +pub async fn list_ip_bans(mysql_pool: &MySqlPool) -> AnyhowResult> { + let maybe_results = sqlx::query_as!( + IpBanRecordForList, + r#" +SELECT + ip_bans.ip_address, + + ip_bans.maybe_target_user_token, + banned_users.username as maybe_target_username, + + ip_bans.mod_user_token, + mod_users.username as mod_username, + mod_users.display_name as mod_display_name, + + ip_bans.mod_notes, + ip_bans.created_at, + ip_bans.updated_at +FROM + ip_address_bans AS ip_bans +LEFT OUTER JOIN users as banned_users + ON ip_bans.maybe_target_user_token = banned_users.token +JOIN users as mod_users + ON ip_bans.mod_user_token = mod_users.token +WHERE + ip_bans.deleted_at IS NULL + "#, + ) + .fetch_all(mysql_pool) + .await; + + match maybe_results { + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + Ok(Vec::new()) + }, + _ => { + warn!("list ip bans db error: {:?}", err); + Err(anyhow!("error with query: {:?}", err)) + } + } + }, + Ok(results) => Ok(results), + } +} diff --git a/crates/lib/database_queries/src/queries/ip_bans/mod.rs b/crates/lib/database_queries/src/queries/ip_bans/mod.rs new file mode 100644 index 0000000000..2cebd7b6e9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/ip_bans/mod.rs @@ -0,0 +1 @@ +pub mod list_ip_bans; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/mod.rs b/crates/lib/database_queries/src/queries/mod.rs new file mode 100644 index 0000000000..eb9ea3bf70 --- /dev/null +++ b/crates/lib/database_queries/src/queries/mod.rs @@ -0,0 +1,9 @@ +pub mod api_tokens; +pub mod health_check; +pub mod ip_bans; +pub mod model_categories; +pub mod tts; +pub mod twitch; +pub mod users; +pub mod voice_clone_requests; +pub mod w2l; diff --git a/crates/lib/database_queries/src/queries/model_categories/get_category_by_token.rs b/crates/lib/database_queries/src/queries/model_categories/get_category_by_token.rs new file mode 100644 index 0000000000..5d5b71b0d0 --- /dev/null +++ b/crates/lib/database_queries/src/queries/model_categories/get_category_by_token.rs @@ -0,0 +1,158 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::{nullable_i8_to_optional_bool, i8_to_bool}; +use log::{info, warn, log}; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct ModelCategory { + // Non-included fields: + //pub id: i64, + //pub uuid_idempotency_token: String, + //pub version: i64, + + pub category_token: String, + pub model_type: String, // TODO: enum + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + + pub is_mod_approved: Option, + + pub maybe_mod_comments: Option, + pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} + +pub async fn get_category_by_token( + category_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + let maybe_category_record = sqlx::query_as!( + RawModelCategory, + r#" +SELECT + category.token as category_token, + category.model_type, + category.maybe_super_category_token, + category.can_directly_have_models, + category.can_have_subcategories, + category.can_only_mods_apply, + category.name, + category.maybe_dropdown_name, + category.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash AS creator_gravatar_hash, + category.creator_ip_address_creation, + category.creator_ip_address_last_update, + category.is_mod_approved, + category.maybe_mod_comments, + category.maybe_mod_user_token, + category.created_at, + category.updated_at, + category.deleted_at +FROM model_categories AS category +JOIN users +ON category.creator_user_token = users.token +WHERE + category.token = ? + "#, + category_token, + ) + .fetch_one(mysql_pool) + .await; + + let category : RawModelCategory = match maybe_category_record { + Ok(category) => category, + Err(err) => { + return match err { + sqlx::Error::RowNotFound => { + Ok(None) + }, + _ => { + warn!("User profile query error: {:?}", err); + Err(anyhow!("query error")) + } + } + } + }; + + let category = ModelCategory { + category_token: category.category_token, + model_type: category.model_type, + maybe_super_category_token: category.maybe_super_category_token, + can_directly_have_models: i8_to_bool(category.can_directly_have_models), + can_have_subcategories: i8_to_bool(category.can_have_subcategories), + can_only_mods_apply: i8_to_bool(category.can_only_mods_apply), + name: category.name.clone(), + maybe_dropdown_name: category.maybe_dropdown_name, + creator_user_token: category.creator_user_token, + creator_username: category.creator_username, + creator_display_name: category.creator_display_name, + creator_gravatar_hash: category.creator_gravatar_hash, + creator_ip_address_creation: category.creator_ip_address_creation, + creator_ip_address_last_update: category.creator_ip_address_last_update, + is_mod_approved: nullable_i8_to_optional_bool(category.is_mod_approved), + maybe_mod_comments: category.maybe_mod_comments, + maybe_mod_user_token: category.maybe_mod_user_token, + created_at: category.created_at, + updated_at: category.updated_at, + deleted_at: category.deleted_at, + }; + + Ok(Some(category)) +} + +#[derive(Serialize)] +pub struct RawModelCategory { + // Non-included fields: + //pub id: i64, + //pub uuid_idempotency_token: String, + //pub version: i64, + + pub category_token: String, + pub model_type: String, // TODO: enum + pub maybe_super_category_token: Option, + + pub can_directly_have_models: i8, + pub can_have_subcategories: i8, + pub can_only_mods_apply: i8, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + + pub is_mod_approved: Option, + + pub maybe_mod_comments: Option, + pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/model_categories/list_categories_query_builder.rs b/crates/lib/database_queries/src/queries/model_categories/list_categories_query_builder.rs new file mode 100644 index 0000000000..9c2235d402 --- /dev/null +++ b/crates/lib/database_queries/src/queries/model_categories/list_categories_query_builder.rs @@ -0,0 +1,385 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::{i8_to_bool, nullable_i8_to_optional_bool}; +use log::{warn, info}; +use sqlx::MySqlPool; + +/// List of categories +/// This is *NOT* paginated, as that defeats the purpose of having a category system. +#[derive(Serialize)] +pub struct CategoryList { + pub categories: Vec, +} + +#[derive(Serialize)] +pub struct Category { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub creator_user_token: Option, + pub creator_username: Option, + pub creator_display_name: Option, + pub creator_gravatar_hash: Option, + + // Moderator fields + pub is_mod_approved: Option, + pub maybe_mod_comments: Option, + //pub maybe_mod_user_token: Option, + //pub maybe_mod_username: Option, + //pub maybe_mod_display_name: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} + +/// Query Builder for listing categories +/// These are very difficult queries, so the builder helps for +/// testability, construction, and correctness. +pub struct ListCategoriesQueryBuilder { + // The public lists won't include unapproved, but mod views + // will want to see them. + show_unapproved: bool, + + // The public lists won't include deleted models, but mod views + // will want to see them. + show_deleted: bool, + + // Filter out approved records. + // Useful as a moderation filter. + // Doesn't make sense to use with `show_unapproved`, but don't want to refactor atm. + hide_approved: bool, + + // Filter out records with a non-deleted timestamp. + // Useful as a moderation filter. + // Doesn't make sense to use with `show_deleted`, but don't want to refactor atm. + hide_non_deleted: bool, + + // Useful for showing categories a user themselves has suggested + // We don't want users spamming creation. + scope_creator_user_token: Option, + + // Typically we won't show both types together. + // Except in a mod view or user view, perhaps. + scope_model_type: Option, // todo: enum +} + +impl ListCategoriesQueryBuilder { + pub fn new() -> Self { + Self { + show_unapproved: false, + show_deleted: false, + hide_approved: false, + hide_non_deleted: false, + scope_creator_user_token: None, + scope_model_type: None, + } + } + + pub fn show_unapproved(mut self, show_unapproved: bool) -> Self { + self.show_unapproved = show_unapproved; + self + } + + pub fn show_deleted(mut self, show_deleted: bool) -> Self { + self.show_deleted = show_deleted; + self + } + + pub fn hide_approved(mut self, hide_approved: bool) -> Self { + self.hide_approved = hide_approved; + self + } + + pub fn hide_non_deleted(mut self, hide_non_deleted: bool) -> Self { + self.hide_non_deleted = hide_non_deleted; + self + } + + pub fn scope_creator_user_token(mut self, scope_creator_user_token: Option<&str>) -> Self { + self.scope_creator_user_token = scope_creator_user_token.map(|u| u.to_string()); + self + } + + pub fn scope_model_type(mut self, scope_model_type: Option<&str>) -> Self { + self.scope_model_type = scope_model_type.map(|u| u.to_string()); + self + } + + /// Perform the query based on the set predicates. + pub async fn perform_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult { + + let internal_results = self.perform_internal_query(mysql_pool).await?; + + let categories = internal_results + .into_iter() + .map(|c| { + Category { + category_token: c.category_token, + model_type: c.model_type, + maybe_super_category_token: c.maybe_super_category_token, + can_directly_have_models: i8_to_bool(c.can_directly_have_models), + can_have_subcategories: i8_to_bool(c.can_have_subcategories), + can_only_mods_apply: i8_to_bool(c.can_only_mods_apply), + name: c.name, + maybe_dropdown_name: c.maybe_dropdown_name, + creator_user_token: c.creator_user_token, + creator_username: c.creator_username, + creator_display_name: c.creator_display_name, + creator_gravatar_hash: c.creator_gravatar_hash, + is_mod_approved: nullable_i8_to_optional_bool(c.is_mod_approved), + maybe_mod_comments: c.maybe_mod_comments, + created_at: c.created_at, + updated_at: c.updated_at, + deleted_at: c.deleted_at, + } + }) + .collect::>(); + + Ok(CategoryList { + categories, + }) + } + + /// Perform the query based on the set predicates. + async fn perform_internal_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + + let query = self.build_query_string(); + let mut query = sqlx::query_as::<_, RawInternalCategoryRecord>(&query); + + // NB: The following bindings must match the order of the query builder !! + + if let Some(user_token) = self.scope_creator_user_token.as_deref() { + query = query.bind(user_token); + } + + if let Some(model_type) = self.scope_model_type.as_deref() { + query = query.bind(model_type); + } + + let mut results = query.fetch_all(mysql_pool) + .await?; + + Ok(results) + } + + fn build_query_string(&self) -> String { + // TODO: I haven't figured out how to get field name disambiguation and type coercion working here. + // (1) tts_results.creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Query error: no column found for name: creator_set_visibility + // (2) creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Column 'creator_set_visibility' in field list is ambiguous + + // TODO/NB: Unfortunately SQLx can't statically typecheck this query + let mut query = r#" +SELECT + model_categories.token as category_token, + + model_categories.model_type, + + model_categories.maybe_super_category_token, + + model_categories.can_directly_have_models, + model_categories.can_have_subcategories, + model_categories.can_only_mods_apply, + + model_categories.name, + model_categories.maybe_dropdown_name, + + users.token as creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash AS creator_gravatar_hash, + + model_categories.is_mod_approved, + model_categories.maybe_mod_comments, + + model_categories.created_at, + model_categories.updated_at, + model_categories.deleted_at + +FROM model_categories +LEFT OUTER JOIN users + ON model_categories.creator_user_token = users.token + "#.to_string(); + + query.push_str(&self.build_predicates()); + query + } + + pub fn build_predicates(&self) -> String { + let mut first_predicate_added = false; + + let mut query = "".to_string(); + + if let Some(username) = self.scope_creator_user_token.as_deref() { + if !first_predicate_added { + query.push_str(" WHERE model_categories.creator_user_token = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND model_categories.creator_user_token = ?"); + } + } + + if let Some(model_type) = self.scope_model_type.as_deref() { + if !first_predicate_added { + query.push_str(" WHERE model_categories.model_type = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND model_categories.model_type = ?"); + } + } + + if !self.show_unapproved { + if !first_predicate_added { + query.push_str(" WHERE model_categories.is_mod_approved IS TRUE"); + first_predicate_added = true; + } else { + query.push_str(" AND model_categories.is_mod_approved IS TRUE"); + } + } + + if !self.show_deleted { + if !first_predicate_added { + query.push_str(" WHERE model_categories.deleted_at IS NULL"); + first_predicate_added = true; + } else { + query.push_str(" AND model_categories.deleted_at IS NULL"); + } + } + + if self.hide_approved { + if !first_predicate_added { + query.push_str(" WHERE (model_categories.is_mod_approved IS NULL OR model_categories.is_mod_approved IS FALSE) "); + first_predicate_added = true; + } else { + query.push_str(" AND (model_categories.is_mod_approved IS NULL OR model_categories.is_mod_approved IS FALSE) "); + } + } + + if self.hide_non_deleted { + if !first_predicate_added { + query.push_str(" WHERE model_categories.deleted_at IS NOT NULL"); + first_predicate_added = true; + } else { + query.push_str(" AND model_categories.deleted_at IS NOT NULL"); + } + } + + query + } +} + +#[derive(sqlx::FromRow)] +pub struct RawInternalCategoryRecord { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: i8, + pub can_have_subcategories: i8, + pub can_only_mods_apply: i8, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub creator_user_token: Option, + pub creator_username: Option, + pub creator_display_name: Option, + pub creator_gravatar_hash: Option, + + // Moderator fields + pub is_mod_approved: Option, + pub maybe_mod_comments: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} + +#[cfg(test)] +mod tests { + use crate::queries::model_categories::list_categories_query_builder::ListCategoriesQueryBuilder; + + #[test] + fn predicates_default_scoping() { + let query_builder = ListCategoriesQueryBuilder::new(); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.is_mod_approved IS TRUE \ + AND model_categories.deleted_at IS NULL"); + } + + #[test] + fn predicates_scoped_to_user() { + let query_builder = ListCategoriesQueryBuilder::new() + .scope_creator_user_token(Some("U:ASDF")); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.creator_user_token = ? \ + AND model_categories.is_mod_approved IS TRUE \ + AND model_categories.deleted_at IS NULL"); + } + + #[test] + fn predicates_scoped_to_model_type() { + let query_builder = ListCategoriesQueryBuilder::new() + .scope_model_type(Some("foo")); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.model_type = ? \ + AND model_categories.is_mod_approved IS TRUE \ + AND model_categories.deleted_at IS NULL"); + } + + #[test] + fn predicates_scoped_to_unapproved() { + let query_builder = ListCategoriesQueryBuilder::new() + .show_unapproved(false); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.is_mod_approved IS TRUE \ + AND model_categories.deleted_at IS NULL"); + + let query_builder = ListCategoriesQueryBuilder::new() + .show_unapproved(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.deleted_at IS NULL"); + } + + #[test] + fn predicates_scoped_to_deleted() { + let query_builder = ListCategoriesQueryBuilder::new() + .show_deleted(false); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.is_mod_approved IS TRUE \ + AND model_categories.deleted_at IS NULL"); + + let query_builder = ListCategoriesQueryBuilder::new() + .show_deleted(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE model_categories.is_mod_approved IS TRUE"); + } +} diff --git a/crates/lib/database_queries/src/queries/model_categories/mod.rs b/crates/lib/database_queries/src/queries/model_categories/mod.rs new file mode 100644 index 0000000000..f00a822919 --- /dev/null +++ b/crates/lib/database_queries/src/queries/model_categories/mod.rs @@ -0,0 +1,2 @@ +pub mod get_category_by_token; +pub mod list_categories_query_builder; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/tts/mod.rs b/crates/lib/database_queries/src/queries/tts/mod.rs new file mode 100644 index 0000000000..9b179450a1 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/mod.rs @@ -0,0 +1,5 @@ +pub mod stats; +pub mod tts_category_assignments; +pub mod tts_inference_jobs; +pub mod tts_models; +pub mod tts_results; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/tts/stats/calculate_tts_model_leaderboard.rs b/crates/lib/database_queries/src/queries/tts/stats/calculate_tts_model_leaderboard.rs new file mode 100644 index 0000000000..0f5b4a5e86 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/stats/calculate_tts_model_leaderboard.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct TtsLeaderboardRecordForList { + pub username: String, + pub display_name: String, + pub gravatar_hash: String, + pub creator_user_token: String, + pub uploaded_count: i64, +} + +#[derive(Serialize)] +struct TtsLeaderboardRecordForListRaw { + pub username: String, + pub display_name: String, + pub gravatar_hash: String, + pub creator_user_token: String, + pub uploaded_count: i64, +} + +// NB: This may need to become an offline query if it is expensive. +pub async fn calculate_tts_model_leaderboard( + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + // NB: We're not requiring "is_public_listing_approved IS TRUE" for TTS models + // This is the opposite for W2L templates (for now at least) + let maybe_results = sqlx::query_as!( + TtsLeaderboardRecordForListRaw, + r#" +SELECT + users.username, + users.display_name, + users.email_gravatar_hash as gravatar_hash, + creator_user_token, + count(*) as uploaded_count +FROM tts_models +JOIN users +ON + users.token = creator_user_token +WHERE + users.is_banned IS FALSE + AND tts_models.user_deleted_at IS NULL + AND tts_models.mod_deleted_at IS NULL +GROUP BY creator_user_token +ORDER BY uploaded_count desc +LIMIT 25; + "#, + ) + .fetch_all(mysql_pool) + .await; + + let results : Vec = match maybe_results { + Ok(results) => results, + Err(err) => { + warn!("Error: {:?}", err); + + match err { + sqlx::Error::RowNotFound => { + return Ok(Vec::new()); + }, + _ => { + warn!("tts leaderboard query error: {:?}", err); + return Err(anyhow!("tts leaderboard query error")); + } + } + } + }; + + Ok(results.into_iter() + .map(|result| { + TtsLeaderboardRecordForList { + username: result.username.to_string(), + display_name: result.display_name.to_string(), + gravatar_hash: result.gravatar_hash.to_string(), + creator_user_token: result.creator_user_token.to_string(), + uploaded_count: result.uploaded_count, + } + }) + .collect::>()) +} diff --git a/crates/lib/database_queries/src/queries/tts/stats/mod.rs b/crates/lib/database_queries/src/queries/tts/stats/mod.rs new file mode 100644 index 0000000000..c0e5981099 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/stats/mod.rs @@ -0,0 +1 @@ +pub mod calculate_tts_model_leaderboard; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/tts/tts_category_assignments/fetch_and_build_tts_model_category_map.rs b/crates/lib/database_queries/src/queries/tts/tts_category_assignments/fetch_and_build_tts_model_category_map.rs new file mode 100644 index 0000000000..e0c9066bb9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_category_assignments/fetch_and_build_tts_model_category_map.rs @@ -0,0 +1,79 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; +use std::collections::{HashSet, HashMap}; + +/// Map of model_token => vec +#[derive(Serialize, Default)] +pub struct TtsModelCategoryMap { + pub model_to_category_tokens: HashMap>, +} + +/// Fetch a map of every model to all of its categories. +pub async fn fetch_and_build_tts_model_category_map( + mysql_pool: &MySqlPool +) -> AnyhowResult { + let assignments = list_tts_model_category_assignments(mysql_pool).await?; + + let mut map : HashMap> = HashMap::new(); + + for assignment in assignments.iter() { + if !map.contains_key(&assignment.tts_model_token) { + map.insert(assignment.tts_model_token.clone(), HashSet::new()); + } + + map.get_mut(&assignment.tts_model_token).map(|hashset| { + hashset.insert(assignment.category_token.clone()) + }); + } + + Ok(TtsModelCategoryMap { + model_to_category_tokens: map, + }) +} + +#[derive(Serialize)] +pub struct CategoryAssignment { + pub tts_model_token: String, + pub category_token: String, +} + +async fn list_tts_model_category_assignments(mysql_pool: &MySqlPool) -> AnyhowResult> { + let maybe_results = sqlx::query_as!( + CategoryAssignment, + r#" +SELECT + assignments.model_token AS tts_model_token, + assignments.category_token +FROM + tts_category_assignments AS assignments +JOIN + tts_models AS tts + ON assignments.model_token = tts.token +WHERE + tts.is_locked_from_use IS FALSE + AND tts.user_deleted_at IS NULL + AND tts.mod_deleted_at IS NULL + AND assignments.deleted_at IS NULL + + "#, + ) + .fetch_all(mysql_pool) + .await; + + match maybe_results { + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + Ok(Vec::new()) + }, + _ => { + warn!("list tts model category assignments db error: {:?}", err); + Err(anyhow!("error with query: {:?}", err)) + } + } + }, + Ok(results) => Ok(results), + } +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_category_assignments/list_assigned_tts_categories_query_builder.rs b/crates/lib/database_queries/src/queries/tts/tts_category_assignments/list_assigned_tts_categories_query_builder.rs new file mode 100644 index 0000000000..eb3e225169 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_category_assignments/list_assigned_tts_categories_query_builder.rs @@ -0,0 +1,268 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::{i8_to_bool, nullable_i8_to_optional_bool}; +use log::{warn, info}; +use sqlx::MySqlPool; + +/// List of assigned categories for a TTS model +/// This is *NOT* paginated, as that defeats the purpose of having a category system. +#[derive(Serialize)] +pub struct AssignedCategoryList { + pub categories: Vec, +} + +#[derive(Serialize)] +pub struct AssignedCategory { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub category_creator_user_token: Option, + pub category_creator_username: Option, + pub category_creator_display_name: Option, + pub category_creator_gravatar_hash: Option, + + // Moderator fields + pub is_mod_approved: Option, + pub maybe_mod_comments: Option, + //pub maybe_mod_user_token: Option, + //pub maybe_mod_username: Option, + //pub maybe_mod_display_name: Option, + + pub category_created_at: DateTime, + pub category_updated_at: DateTime, + pub category_deleted_at: Option>, + + pub category_assignment_created_at: DateTime, + pub category_assignment_updated_at: DateTime, + pub category_assignment_deleted_at: Option>, +} + +/// Query Builder for listing categories +/// These are very difficult queries, so the builder helps for +/// testability, construction, and correctness. +pub struct ListAssignedTtsCategoriesQueryBuilder { + // REQUIRED. + scope_tts_model_token: String, + + // Sometimes "can_directly_have_models" categories change state. + // Mods will want to see invalid uses. + show_invalid_model_not_allowed_categories: bool, + + // The public lists won't include unapproved, but mod views + // will want to see them. + show_unapproved: bool, + + // The public lists won't include deleted models, but mod views + // will want to see them. + show_deleted: bool, +} + +impl ListAssignedTtsCategoriesQueryBuilder { + pub fn for_model_token(model_token: &str) -> Self { + Self { + scope_tts_model_token: model_token.to_string(), + show_invalid_model_not_allowed_categories: false, + show_unapproved: false, + show_deleted: false, + } + } + + pub fn show_invalid_model_not_allowed_categories(mut self, show_invalid_model_not_allowed_categories: bool) -> Self { + self.show_invalid_model_not_allowed_categories = show_invalid_model_not_allowed_categories; + self + } + + pub fn show_unapproved(mut self, show_unapproved: bool) -> Self { + self.show_unapproved = show_unapproved; + self + } + + pub fn show_deleted(mut self, show_deleted: bool) -> Self { + self.show_deleted = show_deleted; + self + } + + /// Perform the query based on the set predicates. + pub async fn perform_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult { + + let internal_results = self.perform_internal_query(mysql_pool).await?; + + let categories = internal_results + .into_iter() + .map(|c| { + AssignedCategory { + category_token: c.category_token, + model_type: c.model_type, + maybe_super_category_token: c.maybe_super_category_token, + can_directly_have_models: i8_to_bool(c.can_directly_have_models), + can_have_subcategories: i8_to_bool(c.can_have_subcategories), + can_only_mods_apply: i8_to_bool(c.can_only_mods_apply), + name: c.name, + maybe_dropdown_name: c.maybe_dropdown_name, + category_creator_user_token: c.category_creator_user_token, + category_creator_username: c.category_creator_username, + category_creator_display_name: c.category_creator_display_name, + category_creator_gravatar_hash: c.category_creator_gravatar_hash, + is_mod_approved: nullable_i8_to_optional_bool(c.is_mod_approved), + maybe_mod_comments: c.maybe_mod_comments, + category_created_at: c.category_created_at, + category_updated_at: c.category_updated_at, + category_deleted_at: c.category_deleted_at, + category_assignment_created_at: c.category_assignment_created_at, + category_assignment_updated_at: c.category_assignment_updated_at, + category_assignment_deleted_at: c.category_assignment_deleted_at, + } + }) + .collect::>(); + + Ok(AssignedCategoryList { + categories, + }) + } + + /// Perform the query based on the set predicates. + async fn perform_internal_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + + let query = self.build_query_string(); + let mut query = sqlx::query_as::<_, RawInternalCategoryRecord>(&query); + + // NB: The following bindings must match the order of the query builder !! + + // NB: Model token is always used for scoping. + query = query.bind(&self.scope_tts_model_token); + + let mut results = query.fetch_all(mysql_pool) + .await?; + + Ok(results) + } + + fn build_query_string(&self) -> String { + // TODO: I haven't figured out how to get field name disambiguation and type coercion working here. + // (1) tts_results.creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Query error: no column found for name: creator_set_visibility + // (2) creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Column 'creator_set_visibility' in field list is ambiguous + + // TODO/NB: Unfortunately SQLx can't statically typecheck this query + let mut query = r#" +SELECT + model_categories.token as category_token, + + model_categories.model_type, + + model_categories.maybe_super_category_token, + + model_categories.can_directly_have_models, + model_categories.can_have_subcategories, + model_categories.can_only_mods_apply, + + model_categories.name, + model_categories.maybe_dropdown_name, + + users.token as category_creator_user_token, + users.username as category_creator_username, + users.display_name as category_creator_display_name, + users.email_gravatar_hash as category_creator_gravatar_hash, + + model_categories.is_mod_approved, + model_categories.maybe_mod_comments, + + model_categories.created_at as category_created_at, + model_categories.updated_at as category_updated_at, + model_categories.deleted_at as category_deleted_at, + + tts_category_assignments.created_at as category_assignment_created_at, + tts_category_assignments.updated_at as category_assignment_updated_at, + tts_category_assignments.deleted_at as category_assignment_deleted_at + +FROM model_categories +JOIN tts_category_assignments + ON tts_category_assignments.category_token = model_categories.token +LEFT OUTER JOIN users + ON model_categories.creator_user_token = users.token + "#.to_string(); + + query.push_str(&self.build_predicates()); + query + } + + pub fn build_predicates(&self) -> String { + let mut query = "".to_string(); + + // Scoping to tts model is required. + query.push_str(" WHERE tts_category_assignments.model_token = ?"); + query.push_str(" AND model_categories.model_type = 'tts'"); + + // We also don't care to ever show soft-deleted assignments + query.push_str(" AND tts_category_assignments.deleted_at IS NULL"); + + if !self.show_invalid_model_not_allowed_categories { + query.push_str(" AND model_categories.can_directly_have_models IS TRUE"); + } + + if !self.show_unapproved { + query.push_str(" AND model_categories.is_mod_approved IS TRUE"); + } + + if !self.show_deleted { + query.push_str(" AND model_categories.deleted_at IS NULL"); + } + + query + } +} + +#[derive(sqlx::FromRow)] +pub struct RawInternalCategoryRecord { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: i8, + pub can_have_subcategories: i8, + pub can_only_mods_apply: i8, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub category_creator_user_token: Option, + pub category_creator_username: Option, + pub category_creator_display_name: Option, + pub category_creator_gravatar_hash: Option, + + // Moderator fields + pub is_mod_approved: Option, + pub maybe_mod_comments: Option, + + pub category_created_at: DateTime, + pub category_updated_at: DateTime, + pub category_deleted_at: Option>, + + pub category_assignment_created_at: DateTime, + pub category_assignment_updated_at: DateTime, + pub category_assignment_deleted_at: Option>, +} + +#[cfg(test)] +mod tests { +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_category_assignments/mod.rs b/crates/lib/database_queries/src/queries/tts/tts_category_assignments/mod.rs new file mode 100644 index 0000000000..dc821db4fb --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_category_assignments/mod.rs @@ -0,0 +1,2 @@ +pub mod fetch_and_build_tts_model_category_map; +pub mod list_assigned_tts_categories_query_builder; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/_keys.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/_keys.rs new file mode 100644 index 0000000000..0a6b60dd8f --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/_keys.rs @@ -0,0 +1,5 @@ +//! Primary and foreign keys for TTS-related models + +/// Since jobs are internal, we'll use IDs as primary keys. +#[derive(Clone, Copy, PartialOrd, PartialEq, Eq, Ord, Debug)] +pub struct TtsInferenceJobId(pub i64); diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/insert_tts_inference_job.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/insert_tts_inference_job.rs new file mode 100644 index 0000000000..ff95681774 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/insert_tts_inference_job.rs @@ -0,0 +1,182 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::token::random_uuid::generate_random_uuid; +use crate::builders::RequiredOption; +use sqlx::MySqlPool; + +pub struct TtsInferenceJobInsertBuilder { + // ========== Required ========== + // TODO: The builder should generate this + // This requires moving token code into the shared db crate. + job_token: RequiredOption, + uuid_idempotency_token: RequiredOption, + model_token: RequiredOption, + raw_inference_text: RequiredOption, + + // ========== Conditionally optional ========== + creator_ip_address: RequiredOption, + + // ========== Optional ========== + maybe_creator_user_token: Option, + creator_set_visibility: Option, + is_from_api: bool, + is_for_twitch: bool, + priority_level: u8, // Priority is 0 by default and optionally increases +} + +impl TtsInferenceJobInsertBuilder { + + /// Everything needs to be manually specified + pub fn new_for_fakeyou_request() -> Self { + Self { + job_token: None, + uuid_idempotency_token: None, + model_token: None, + raw_inference_text: None, + creator_ip_address: None, + maybe_creator_user_token: None, + creator_set_visibility: None, + is_from_api: false, + is_for_twitch: false, + priority_level: 0, + } + } + + /// We can default a lot of the fields to "empty"-ish values + pub fn new_for_internal_tts() -> Self { + let idempotency_token = generate_random_uuid(); + Self { + job_token: None, + model_token: None, + raw_inference_text: None, + // Defaults + uuid_idempotency_token: Some(idempotency_token), + maybe_creator_user_token: None, + creator_ip_address: Some("127.0.0.1".to_string()), + // hidden | public | ... + creator_set_visibility: Some("hidden".to_string()), + is_from_api: false, + is_for_twitch: false, + priority_level: 0, + } + } + + pub fn set_job_token(mut self, value: &str) -> Self { + self.job_token = Some(value.to_string()); + self + } + + pub fn set_model_token(mut self, value: &str) -> Self { + self.model_token = Some(value.to_string()); + self + } + + pub fn set_raw_inference_text(mut self, value: &str) -> Self { + self.raw_inference_text = Some(value.to_string()); + self + } + + pub fn set_uuid_idempotency_token(mut self, value: &str) -> Self { + self.uuid_idempotency_token = Some(value.to_string()); + self + } + + pub fn set_creator_ip_address(mut self, value: &str) -> Self { + self.creator_ip_address = Some(value.to_string()); + self + } + + pub fn set_maybe_creator_user_token(mut self, value: Option<&str>) -> Self { + self.maybe_creator_user_token = value.map(|s| s.to_string()); + self + } + + pub fn set_creator_set_visibility(mut self, value: &str) -> Self { + self.creator_set_visibility = Some(value.to_string()); + self + } + + pub fn set_is_from_api(mut self, value: bool) -> Self { + self.is_from_api = value; + self + } + + pub fn set_is_for_twitch(mut self, value: bool) -> Self { + self.is_for_twitch = value; + self + } + + pub fn set_priority_level(mut self, value: u8) -> Self { + self.priority_level = value; + self + } + + pub async fn insert(&mut self, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + // TODO: These should be custom error types with a custom macro to make this easy. + let job_token = self.job_token + .clone() + .ok_or(anyhow!("no job_token"))?; + + let uuid_idempotency_token = self.uuid_idempotency_token + .clone() + .ok_or(anyhow!("no uuid_idempotency_token"))?; + + let model_token = self.model_token + .clone() + .ok_or(anyhow!("no model_token"))?; + + let raw_inference_text = self.raw_inference_text + .clone() + .ok_or(anyhow!("no raw_inference_text"))?; + + let creator_ip_address = self.creator_ip_address + .clone() + .ok_or(anyhow!("no creator_ip_address"))?; + + let creator_set_visibility = self.creator_set_visibility + .clone() + .ok_or(anyhow!("no creator_set_visibility"))?; + + let query = sqlx::query!( + r#" +INSERT INTO tts_inference_jobs +SET + token = ?, + uuid_idempotency_token = ?, + model_token = ?, + raw_inference_text = ?, + maybe_creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + is_from_api = ?, + is_for_twitch = ?, + priority_level = ?, + status = "pending" + "#, + job_token, + uuid_idempotency_token, + model_token, + raw_inference_text, + self.maybe_creator_user_token.clone(), + creator_ip_address, + creator_set_visibility, + self.is_from_api, + self.is_for_twitch, + self.priority_level, + ); + + let query_result = query.execute(mysql_pool) + .await; + + let _record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("Twitch pubsub bits insert DB error: {:?}", err)); + } + }; + + Ok(()) + } +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/kill_tts_inference_jobs.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/kill_tts_inference_jobs.rs new file mode 100644 index 0000000000..5e5f4e6747 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/kill_tts_inference_jobs.rs @@ -0,0 +1,67 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use sqlx::mysql::MySqlArguments; +use sqlx::query::Query; +use sqlx::{MySqlPool, MySql}; + +#[derive(Copy, Clone)] +pub enum JobStatus { + /// Kill all "pending" jobs + AllPending, + /// Kill all "pending" and "attempt_failed" jobs + AllPendingAndFailed, + /// Kill "pending" jobs with priority_level = 0. + ZeroPriorityPending, +} + +pub async fn kill_tts_inference_jobs(job_status: JobStatus, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + + let query = match job_status { + JobStatus::AllPending => kill_all_pending(), + JobStatus::AllPendingAndFailed => kill_all_pending_and_failed(), + JobStatus::ZeroPriorityPending => kill_all_zero_priority_pending(), + }; + + let result = query.execute(mysql_pool).await; + + match result { + Err(err) => { + Err(anyhow!("error with query: {:?}", err)) + }, + Ok(_r) => Ok(()), + } +} + +fn kill_all_pending() -> Query<'static, MySql, MySqlArguments> { + sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET status='complete_failure' +WHERE + status IN ('pending'); + "# + ) +} + +fn kill_all_pending_and_failed() -> Query<'static, MySql, MySqlArguments> { + sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET status='complete_failure' +WHERE + status IN ('pending', 'attempt_failed'); + "# + ) +} + +fn kill_all_zero_priority_pending() -> Query<'static, MySql, MySqlArguments> { + sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET status='complete_failure' +WHERE + status IN ('pending') + AND priority_level = 0; + "# + ) +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs new file mode 100644 index 0000000000..c70f536f30 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs @@ -0,0 +1,190 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::job_status::JobStatus; +use crate::column_types::record_visibility::RecordVisibility; +use crate::helpers::boolean_converters::i8_to_bool; +use crate::queries::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use log::{warn, info}; +use sqlx::MySqlPool; + +/// table: tts_inference_jobs +#[derive(Debug)] +pub struct AvailableTtsInferenceJob { + pub id: TtsInferenceJobId, + pub inference_job_token: String, + pub uuid_idempotency_token: String, + + pub model_token: String, + pub raw_inference_text: String, + + pub creator_ip_address: String, + pub maybe_creator_user_token: Option, + pub creator_set_visibility: RecordVisibility, + + pub is_from_api: bool, + pub is_for_twitch: bool, + + pub status: JobStatus, + pub priority_level: u8, + pub attempt_count: i32, + pub failure_reason: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + +/// Query jobs that are ready to run +/// If sorting by priority, be careful not to starve lower priority jobs +/// (ie. if there's an issue with higher priorities.) +pub async fn list_available_tts_inference_jobs( + pool: &MySqlPool, + sort_by_priority: bool, + num_records: u32 +) -> AnyhowResult> { + + // NB: This query is awkwardly written twice because this is the only way the + // macro can statically type check the query, result types, and parameter bindings. + // + // The only difference here is the 'ORDER BY' clause ! + let job_records : Vec = if sort_by_priority { + sqlx::query_as!( + AvailableTtsInferenceJobRawInternal, + r#" +SELECT + id, + token AS inference_job_token, + uuid_idempotency_token, + + model_token, + raw_inference_text, + + creator_ip_address, + maybe_creator_user_token, + creator_set_visibility as `creator_set_visibility: crate::column_types::record_visibility::RecordVisibility`, + + is_from_api, + is_for_twitch, + + status as `status: crate::column_types::job_status::JobStatus`, + priority_level, + attempt_count, + failure_reason, + + created_at, + updated_at, + retry_at +FROM tts_inference_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + ORDER BY priority_level DESC, id ASC + LIMIT ? + "#, + num_records, + ).fetch_all(pool).await? + } else { + sqlx::query_as!( + AvailableTtsInferenceJobRawInternal, + r#" +SELECT + id, + token AS inference_job_token, + uuid_idempotency_token, + + model_token, + raw_inference_text, + + creator_ip_address, + maybe_creator_user_token, + creator_set_visibility as `creator_set_visibility: crate::column_types::record_visibility::RecordVisibility`, + + is_from_api, + is_for_twitch, + + status as `status: crate::column_types::job_status::JobStatus`, + priority_level, + attempt_count, + failure_reason, + + created_at, + updated_at, + retry_at +FROM tts_inference_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + LIMIT ? + "#, + num_records, + ).fetch_all(pool).await? + }; + + let job_records = job_records.into_iter() + .map(|record : AvailableTtsInferenceJobRawInternal| { + AvailableTtsInferenceJob { + id: TtsInferenceJobId(record.id), + inference_job_token: record.inference_job_token, + uuid_idempotency_token: record.uuid_idempotency_token, + model_token: record.model_token, + raw_inference_text: record.raw_inference_text, + creator_ip_address: record.creator_ip_address, + maybe_creator_user_token: record.maybe_creator_user_token, + creator_set_visibility: record.creator_set_visibility, + is_from_api: i8_to_bool(record.is_from_api), + is_for_twitch: i8_to_bool(record.is_for_twitch), + status: record.status, + priority_level: record.priority_level, + attempt_count: record.attempt_count, + failure_reason: record.failure_reason, + created_at: record.created_at, + updated_at: record.updated_at, + retry_at: record.retry_at, + } + }) + .collect::>(); + + Ok(job_records) +} + +#[derive(Debug)] +struct AvailableTtsInferenceJobRawInternal { + pub id: i64, + pub inference_job_token: String, + pub uuid_idempotency_token: String, + + pub model_token: String, + pub raw_inference_text: String, + + pub creator_ip_address: String, + pub maybe_creator_user_token: Option, + pub creator_set_visibility: RecordVisibility, + + pub is_from_api: i8, + pub is_for_twitch: i8, + + pub status: JobStatus, + pub priority_level: u8, + pub attempt_count: i32, + pub failure_reason: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_done.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_done.rs new file mode 100644 index 0000000000..abbd33a670 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_done.rs @@ -0,0 +1,39 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use sqlx::MySqlPool; +use sqlx; + +pub async fn mark_tts_inference_job_done( + pool: &MySqlPool, + job_id: TtsInferenceJobId, + success: bool, + maybe_result_token: Option<&str> +) -> AnyhowResult<()> { + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = ?, + on_success_result_token = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_result_token, + job_id.0 + ) + .execute(pool) + .await; + + match query_result { + Err(err) => { + Err(anyhow!("error with query: {:?}", err)) + }, + Ok(_r) => Ok(()), + } +} + diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_failure.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_failure.rs new file mode 100644 index 0000000000..a68476b592 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_failure.rs @@ -0,0 +1,47 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob; +use sqlx::MySqlPool; +use sqlx; + +// TODO: Rerun time shouldn't be 1-minute! + +/// Mark a single inference failure. The job may be re-run. +pub async fn mark_tts_inference_job_failure( + pool: &MySqlPool, + job: &AvailableTtsInferenceJob, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 1 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id.0, + ) + .execute(pool) + .await; + + match query_result { + Err(err) => { + Err(anyhow!("error with query: {:?}", err)) + }, + Ok(_r) => Ok(()), + } +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_pending_and_grab_lock.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_pending_and_grab_lock.rs new file mode 100644 index 0000000000..2c8ae40f9c --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_pending_and_grab_lock.rs @@ -0,0 +1,85 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use sqlx::MySqlPool; +use sqlx; + +// TODO(2022-02-09): It's unclear how stuck "pending" jobs don't become perma-dead +// I think this logic is flawed and the job system needs to pick up stuck "pending" jobs. + +pub struct TtsInferenceLockRecord { + id: i64, + status: String, +} + +pub async fn mark_tts_inference_job_pending_and_grab_lock( + pool: &MySqlPool, + job_id: TtsInferenceJobId +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + TtsInferenceLockRecord, + r#" +SELECT + id, + status +FROM tts_inference_jobs +WHERE id = ? +FOR UPDATE + "#, + job_id.0, + ) + .fetch_one(&mut transaction) + .await; + + let record : TtsInferenceLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job_id.0, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} + diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_permanently_dead.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_permanently_dead.rs new file mode 100644 index 0000000000..0b7c65e649 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mark_tts_inference_job_permanently_dead.rs @@ -0,0 +1,34 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::tts::tts_inference_jobs::_keys::TtsInferenceJobId; +use sqlx::MySqlPool; +use sqlx; + +/// Mark a TTS job as dead ("do not ever retry"). +pub async fn mark_tts_inference_job_permanently_dead( + pool: &MySqlPool, + job_id: TtsInferenceJobId, + failure_reason: &str, +) -> AnyhowResult<()> { + let query_result = sqlx::query!( + r#" +UPDATE tts_inference_jobs +SET + status = "dead", + failure_reason = ?, + retry_at = NULL +WHERE id = ? + "#, + failure_reason.to_string(), + job_id.0, + ) + .execute(pool) + .await; + + match query_result { + Err(err) => { + Err(anyhow!("error with query: {:?}", err)) + }, + Ok(_r) => Ok(()), + } +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mod.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mod.rs new file mode 100644 index 0000000000..b15fce6350 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/mod.rs @@ -0,0 +1,8 @@ +pub mod _keys; +pub mod insert_tts_inference_job; +pub mod kill_tts_inference_jobs; +pub mod list_available_tts_inference_jobs; +pub mod mark_tts_inference_job_done; +pub mod mark_tts_inference_job_failure; +pub mod mark_tts_inference_job_pending_and_grab_lock; +pub mod mark_tts_inference_job_permanently_dead; diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/delete_tts_model_various_scopes.rs b/crates/lib/database_queries/src/queries/tts/tts_models/delete_tts_model_various_scopes.rs new file mode 100644 index 0000000000..f01091e14b --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_models/delete_tts_model_various_scopes.rs @@ -0,0 +1,94 @@ +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +pub async fn delete_tts_model_as_user( + tts_model_token: &str, + creator_ip_address: &str, + mysql_pool: &MySqlPool, +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_models +SET + creator_ip_address_last_update = ?, + user_deleted_at = CURRENT_TIMESTAMP +WHERE + token = ? +LIMIT 1 + "#, + creator_ip_address, + tts_model_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn delete_tts_model_as_mod( + tts_model_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool, +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_models +SET + maybe_mod_user_token = ?, + mod_deleted_at = CURRENT_TIMESTAMP +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + tts_model_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_tts_model_as_user( + tts_model_token: &str, + creator_ip_address: &str, + mysql_pool: &MySqlPool, +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_models +SET + creator_ip_address_last_update = ?, + user_deleted_at = NULL +WHERE + token = ? +LIMIT 1 + "#, + creator_ip_address, + tts_model_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_tts_model_as_mod( + tts_model_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool, +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_models +SET + maybe_mod_user_token = ?, + mod_deleted_at = NULL +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + tts_model_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs new file mode 100644 index 0000000000..54b60c0033 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs @@ -0,0 +1,283 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::record_visibility::RecordVisibility; +use crate::column_types::vocoder_type::VocoderType; +use crate::helpers::boolean_converters::i8_to_bool; +use log::warn; +use sqlx::MySqlPool; + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// The moderator-only fields are good practice, though. + +// TODO/FIXME : This struct is returned publicly in some endpoints! +#[derive(Serialize)] +pub struct TtsModelRecordForResponse { + pub model_token: String, + pub tts_model_type: String, + pub maybe_default_pretrained_vocoder: Option, + pub text_preprocessing_algorithm: String, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + + pub is_front_page_featured: bool, + pub is_twitch_featured: bool, + + pub maybe_suggested_unique_bot_command: Option, + + pub creator_set_visibility: RecordVisibility, + + pub is_locked_from_use: bool, + pub is_locked_from_user_modification: bool, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub maybe_moderator_fields: Option, +} + +/// "Moderator-only fields" that we wouldn't want to expose to ordinary users. +/// It's the web endpoint controller's responsibility to clear these for non-mods. +#[derive(Serialize)] +pub struct TtsModelModeratorFields { + pub creator_is_banned: bool, + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// The moderator-only fields are good practice, though. + +pub async fn get_tts_model_by_token( + tts_model_token: &str, + can_see_deleted: bool, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + let maybe_record = if can_see_deleted { + select_including_deleted(tts_model_token, mysql_pool).await + } else { + select_without_deleted(tts_model_token, mysql_pool).await + }; + + let model : InternalTtsModelRecordRaw = match maybe_record { + Ok(model) => model, + Err(ref err) => { + match err { + sqlx::Error::RowNotFound => { + warn!("tts model not found: {:?}", &err); + return Ok(None); + }, + _ => { + warn!("tts model query error: {:?}", &err); + return Err(anyhow!("database error")); + } + } + } + }; + + let mut maybe_vocoder : Option = None; + if let Some(vocoder) = model.maybe_default_pretrained_vocoder.as_deref() { + maybe_vocoder = Some(VocoderType::from_str(vocoder)?); + } + + let model_for_response = TtsModelRecordForResponse { + model_token: model.model_token, + tts_model_type: model.tts_model_type, + maybe_default_pretrained_vocoder: maybe_vocoder, + text_preprocessing_algorithm: model.text_preprocessing_algorithm, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + title: model.title, + description_markdown: model.description_markdown, + description_rendered_html: model.description_rendered_html, + // NB: Fail open/public with creator_set_visibility since we're already looking at it + ietf_language_tag: model.ietf_language_tag, + ietf_primary_language_subtag: model.ietf_primary_language_subtag, + is_front_page_featured: i8_to_bool(model.is_front_page_featured), + is_twitch_featured: i8_to_bool(model.is_twitch_featured), + maybe_suggested_unique_bot_command: model.maybe_suggested_unique_bot_command, + creator_set_visibility: RecordVisibility::from_str(&model.creator_set_visibility) + .unwrap_or(RecordVisibility::Public), + is_locked_from_use: i8_to_bool(model.is_locked_from_use), + is_locked_from_user_modification: i8_to_bool(model.is_locked_from_user_modification), + created_at: model.created_at, + updated_at: model.updated_at, + maybe_moderator_fields: Some(TtsModelModeratorFields { + creator_is_banned: i8_to_bool(model.creator_is_banned), + creator_ip_address_creation: model.creator_ip_address_creation, + creator_ip_address_last_update: model.creator_ip_address_last_update, + user_deleted_at: model.user_deleted_at, + mod_deleted_at: model.mod_deleted_at, + }), + }; + + Ok(Some(model_for_response)) +} + +async fn select_including_deleted( + tts_model_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + InternalTtsModelRecordRaw, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.text_preprocessing_algorithm, + tts.maybe_default_pretrained_vocoder, + + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + users.is_banned as creator_is_banned, + + tts.title, + tts.description_markdown, + tts.description_rendered_html, + + tts.ietf_language_tag, + tts.ietf_primary_language_subtag, + + tts.is_front_page_featured, + tts.is_twitch_featured, + + tts.maybe_suggested_unique_bot_command, + + tts.creator_set_visibility, + + tts.is_locked_from_use, + tts.is_locked_from_user_modification, + + tts.created_at, + tts.updated_at, + + tts.creator_ip_address_creation, + tts.creator_ip_address_last_update, + tts.mod_deleted_at, + tts.user_deleted_at + +FROM tts_models as tts +JOIN users + ON users.token = tts.creator_user_token +WHERE tts.token = ? + "#, + tts_model_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} + +async fn select_without_deleted( + tts_model_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + InternalTtsModelRecordRaw, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.text_preprocessing_algorithm, + tts.maybe_default_pretrained_vocoder, + + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + users.is_banned as creator_is_banned, + + tts.title, + tts.description_markdown, + tts.description_rendered_html, + + tts.ietf_language_tag, + tts.ietf_primary_language_subtag, + + tts.is_front_page_featured, + tts.is_twitch_featured, + + tts.maybe_suggested_unique_bot_command, + + tts.creator_set_visibility, + + tts.is_locked_from_use, + tts.is_locked_from_user_modification, + + tts.created_at, + tts.updated_at, + + tts.creator_ip_address_creation, + tts.creator_ip_address_last_update, + tts.mod_deleted_at, + tts.user_deleted_at + +FROM tts_models as tts +JOIN users + ON users.token = tts.creator_user_token +WHERE + tts.token = ? + AND tts.user_deleted_at IS NULL + AND tts.mod_deleted_at IS NULL + "#, + tts_model_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} + +#[derive(Serialize)] +struct InternalTtsModelRecordRaw { + pub model_token: String, + pub tts_model_type: String, + pub maybe_default_pretrained_vocoder: Option, + pub text_preprocessing_algorithm: String, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + pub creator_is_banned: i8, + + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + + pub is_front_page_featured: i8, + pub is_twitch_featured: i8, + + pub maybe_suggested_unique_bot_command: Option, + + pub creator_set_visibility: String, + + pub is_locked_from_use: i8, + pub is_locked_from_user_modification: i8, + + pub created_at: DateTime, + pub updated_at: DateTime, + + // Moderator fields + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs new file mode 100644 index 0000000000..a5d131a7c4 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs @@ -0,0 +1,103 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use log::warn; +use sqlx::MySqlPool; +use sqlx; + +// TODO: Can probably just reuse another query. + +pub struct TtsModelForInferenceRecord { + pub model_token: String, + pub tts_model_type: String, + + pub maybe_default_pretrained_vocoder: Option, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + + pub title: String, + pub private_bucket_hash: String, + pub created_at: DateTime, + pub updated_at: DateTime, + pub mod_deleted_at: Option>, + pub user_deleted_at: Option>, +} + + +#[derive(Clone, Debug)] +pub enum TtsModelForInferenceError { + ModelNotFound, + ModelDeleted, + DatabaseError { reason: String }, +} + +impl std::fmt::Display for TtsModelForInferenceError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + TtsModelForInferenceError::ModelNotFound => write!(f, "ModelNotFound"), + TtsModelForInferenceError::ModelDeleted => write!(f, "ModelDeleted"), + TtsModelForInferenceError::DatabaseError { reason} => + write!(f, "Database error: {:?}", reason), + } + } +} + +impl std::error::Error for TtsModelForInferenceError {} + +pub async fn get_tts_model_for_inference( + pool: &MySqlPool, + model_token: &str +) -> Result +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_model = sqlx::query_as!( + TtsModelForInferenceRecord, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.maybe_default_pretrained_vocoder, + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + tts.title, + tts.private_bucket_hash, + tts.created_at, + tts.updated_at, + tts.user_deleted_at, + tts.mod_deleted_at +FROM tts_models as tts +JOIN users +ON users.token = tts.creator_user_token +WHERE tts.token = ? + "#, + &model_token + ) + .fetch_one(pool) + .await; // TODO: This will return error if it doesn't exist + + let model : TtsModelForInferenceRecord = match maybe_model { + Ok(model) => model, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(TtsModelForInferenceError::ModelNotFound); + }, + _ => { + warn!("tts model query error: {:?}", err); + return Err(TtsModelForInferenceError::DatabaseError { + reason: format!("Mysql error: {:?}", err) + }); + } + } + } + }; + + if model.mod_deleted_at.is_some() || model.user_deleted_at.is_some() { + return Err(TtsModelForInferenceError::ModelDeleted); + } + + Ok(model) +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs b/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs new file mode 100644 index 0000000000..c56ddd4427 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_models/list_tts_models.rs @@ -0,0 +1,269 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use log::{warn, info}; +use sqlx::MySqlPool; + +// FIXME: This is the old style of query scoping and shouldn't be copied. + +// TODO/FIXME : This struct is returned publicly in some endpoints! +#[derive(Serialize)] +pub struct TtsModelRecordForList { + pub model_token: String, + pub tts_model_type: String, + + pub title: String, + + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub is_locked_from_use: bool, + + pub is_front_page_featured: bool, + pub is_twitch_featured: bool, + + pub maybe_suggested_unique_bot_command: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +// FIXME: This is the old style of query scoping and shouldn't be copied. + +pub async fn list_tts_models( + mysql_pool: &MySqlPool, + scope_creator_username: Option<&str>, + require_mod_approved: bool +) -> AnyhowResult> { + + let maybe_models = match scope_creator_username { + Some(username) => { + list_tts_models_creator_scoped(mysql_pool, username, require_mod_approved) + .await + }, + None => { + list_tts_models_for_all_creators(mysql_pool, require_mod_approved) + .await + }, + }; + + let models : Vec = match maybe_models { + Ok(models) => models, + Err(err) => { + match err { + RowNotFound => { + return Ok(Vec::new()); + }, + _ => { + warn!("tts model list query error: {:?}", err); + return Err(anyhow!("tts model list query error")); + } + } + } + }; + + Ok(models.into_iter() + .map(|model| { + TtsModelRecordForList { + model_token: model.model_token, + tts_model_type: model.tts_model_type, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + title: model.title, + ietf_language_tag: model.ietf_language_tag, + ietf_primary_language_subtag: model.ietf_primary_language_subtag, + is_locked_from_use: i8_to_bool(model.is_locked_from_use), + is_front_page_featured: i8_to_bool(model.is_front_page_featured), + is_twitch_featured: i8_to_bool(model.is_twitch_featured), + maybe_suggested_unique_bot_command: model.maybe_suggested_unique_bot_command, + created_at: model.created_at, + updated_at: model.updated_at, + } + }) + .collect::>()) +} + +async fn list_tts_models_for_all_creators( + mysql_pool: &MySqlPool, + allow_mod_disabled: bool +) -> AnyhowResult> { + // TODO: There has to be a better way. + // Sqlx doesn't like anything except string literals. + let maybe_models = if !allow_mod_disabled { + info!("listing tts models for everyone; mod-approved only"); + sqlx::query_as!( + InternalRawTtsModelRecordForList, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + tts.title, + tts.ietf_language_tag, + tts.ietf_primary_language_subtag, + tts.is_locked_from_use, + tts.is_front_page_featured, + tts.is_twitch_featured, + tts.maybe_suggested_unique_bot_command, + tts.created_at, + tts.updated_at +FROM tts_models as tts +JOIN users + ON users.token = tts.creator_user_token +WHERE + tts.is_locked_from_use IS FALSE + AND tts.user_deleted_at IS NULL + AND tts.mod_deleted_at IS NULL + "#) + .fetch_all(mysql_pool) + .await? + } else { + info!("listing tts models for everyone; all"); + sqlx::query_as!( + InternalRawTtsModelRecordForList, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + tts.title, + tts.ietf_language_tag, + tts.ietf_primary_language_subtag, + tts.is_locked_from_use, + tts.is_front_page_featured, + tts.is_twitch_featured, + tts.maybe_suggested_unique_bot_command, + tts.created_at, + tts.updated_at +FROM tts_models as tts +JOIN users + ON users.token = tts.creator_user_token +WHERE + tts.user_deleted_at IS NULL + AND tts.mod_deleted_at IS NULL + "#) + .fetch_all(mysql_pool) + .await? + }; + + Ok(maybe_models) +} + +async fn list_tts_models_creator_scoped( + mysql_pool: &MySqlPool, + scope_creator_username: &str, + allow_mod_disabled: bool +) -> AnyhowResult> { + // TODO: There has to be a better way. + // Sqlx doesn't like anything except string literals. + let maybe_models = if !allow_mod_disabled { + info!("listing tts models for user; mod-approved only"); + sqlx::query_as!( + InternalRawTtsModelRecordForList, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + tts.title, + tts.ietf_language_tag, + tts.ietf_primary_language_subtag, + tts.is_locked_from_use, + tts.is_front_page_featured, + tts.is_twitch_featured, + tts.maybe_suggested_unique_bot_command, + tts.created_at, + tts.updated_at +FROM tts_models as tts +JOIN users +ON + users.token = tts.creator_user_token +WHERE + users.username = ? + AND tts.is_locked_from_use IS FALSE + AND tts.user_deleted_at IS NULL + AND tts.mod_deleted_at IS NULL + "#, + scope_creator_username) + .fetch_all(mysql_pool) + .await? + } else { + info!("listing tts models for user; all"); + sqlx::query_as!( + InternalRawTtsModelRecordForList, + r#" +SELECT + tts.token as model_token, + tts.tts_model_type, + tts.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + tts.title, + tts.ietf_language_tag, + tts.ietf_primary_language_subtag, + tts.is_locked_from_use, + tts.is_front_page_featured, + tts.is_twitch_featured, + tts.maybe_suggested_unique_bot_command, + tts.created_at, + tts.updated_at +FROM tts_models as tts +JOIN users +ON + users.token = tts.creator_user_token +WHERE + users.username = ? + AND tts.user_deleted_at IS NULL + AND tts.mod_deleted_at IS NULL + "#, + scope_creator_username) + .fetch_all(mysql_pool) + .await? + }; + + Ok(maybe_models) +} + +struct InternalRawTtsModelRecordForList { + pub model_token: String, + pub tts_model_type: String, + + pub title: String, + + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub is_locked_from_use: i8, // bool + + pub is_front_page_featured: i8, // bool + pub is_twitch_featured: i8, // bool + + pub maybe_suggested_unique_bot_command: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/mod.rs b/crates/lib/database_queries/src/queries/tts/tts_models/mod.rs new file mode 100644 index 0000000000..ecd8eec251 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_models/mod.rs @@ -0,0 +1,4 @@ +pub mod delete_tts_model_various_scopes; +pub mod get_tts_model; +pub mod get_tts_model_for_inference; +pub mod list_tts_models; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/tts/tts_results/delete_tts_result_various_scopes.rs b/crates/lib/database_queries/src/queries/tts/tts_results/delete_tts_result_various_scopes.rs new file mode 100644 index 0000000000..c5ca6f2f3f --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_results/delete_tts_result_various_scopes.rs @@ -0,0 +1,88 @@ +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +pub async fn delete_tts_inference_result_as_user( + inference_result_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_results +SET + user_deleted_at = CURRENT_TIMESTAMP +WHERE + token = ? +LIMIT 1 + "#, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn delete_tts_inference_result_as_mod( + inference_result_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_results +SET + mod_deleted_at = CURRENT_TIMESTAMP, + maybe_mod_user_token = ? +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_tts_inference_result_as_user( + inference_result_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_results +SET + user_deleted_at = NULL +WHERE + token = ? +LIMIT 1 + "#, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_tts_inference_result_as_mod( + inference_result_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE tts_results +SET + mod_deleted_at = NULL, + maybe_mod_user_token = ? +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_results/insert_tts_result.rs b/crates/lib/database_queries/src/queries/tts/tts_results/insert_tts_result.rs new file mode 100644 index 0000000000..00b75da4d1 --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_results/insert_tts_result.rs @@ -0,0 +1,165 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::vocoder_type::VocoderType; +use crate::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::AvailableTtsInferenceJob; +use crate::tokens::Tokens; +use log::warn; +use sqlx::MySqlPool; +use sqlx; +use std::path::Path; + +/// Used to give user-facing order to logged in user inference requests +pub struct SyntheticIdRecord { + pub next_id: i64, +} + +pub async fn insert_tts_result>( + pool: &MySqlPool, + job: &AvailableTtsInferenceJob, + text_hash: &str, + pretrained_vocoder_used: VocoderType, + bucket_audio_results_path: P, + bucket_spectrogram_results_path: P, + file_size_bytes: u64, + duration_millis: u64 +) -> AnyhowResult<(u64, String)> +{ + let inference_result_token = Tokens::new_tts_result()?; + + let bucket_audio_result_path = &bucket_audio_results_path + .as_ref() + .display() + .to_string(); + + let bucket_spectrogram_result_path = &bucket_spectrogram_results_path + .as_ref() + .display() + .to_string(); + + let normalized_inference_text = job.raw_inference_text.clone(); // TODO + + let maybe_creator_user_token = job.maybe_creator_user_token.clone(); + let mut maybe_creator_synthetic_id : Option = None; + + let mut transaction = pool.begin().await?; + + if let Some(creator_user_token) = maybe_creator_user_token.as_deref() { + let query_result = sqlx::query!( + r#" +INSERT INTO tts_result_synthetic_ids +SET + user_token = ?, + next_id = 1 +ON DUPLICATE KEY UPDATE + user_token = ?, + next_id = next_id + 1 + "#, + creator_user_token, + creator_user_token + ) + .execute(&mut transaction) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + //transaction.rollback().await?; + warn!("Transaction failure: {:?}", err); + } + } + + let query_result = sqlx::query_as!( + SyntheticIdRecord, + r#" +SELECT + next_id +FROM + tts_result_synthetic_ids +WHERE + user_token = ? +LIMIT 1 + "#, + creator_user_token, + ) + .fetch_one(&mut transaction) + .await; + + let record : SyntheticIdRecord = match query_result { + Ok(record) => record, + Err(err) => { + warn!("Transaction failure: {:?}", err); + transaction.rollback().await?; + return Err(anyhow!("Transaction failure: {:?}", err)); + } + }; + + let next_id = record.next_id as u64; + maybe_creator_synthetic_id = Some(next_id); + } + + let record_id = { + let query_result = sqlx::query!( + r#" +INSERT INTO tts_results +SET + token = ?, + + model_token = ?, + maybe_pretrained_vocoder_used = ?, + raw_inference_text = ?, + raw_inference_text_hash_sha2 = ?, + normalized_inference_text = ?, + + maybe_creator_user_token = ?, + maybe_creator_synthetic_id = ?, + + creator_ip_address = ?, + creator_set_visibility = ?, + + public_bucket_wav_audio_path = ?, + public_bucket_spectrogram_path = ?, + + file_size_bytes = ?, + duration_millis = ? + "#, + inference_result_token, + job.model_token.clone(), + pretrained_vocoder_used.to_str(), + job.raw_inference_text.clone(), + text_hash, + normalized_inference_text, + + maybe_creator_user_token, + maybe_creator_synthetic_id, + + job.creator_ip_address.clone(), + job.creator_set_visibility.to_str(), + + bucket_audio_result_path, + bucket_spectrogram_result_path, + + file_size_bytes, + duration_millis + ) + .execute(&mut transaction) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + //transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + record_id + }; + + transaction.commit().await?; + + Ok((record_id, inference_result_token.clone())) +} + diff --git a/crates/lib/database_queries/src/queries/tts/tts_results/list_tts_results_query_builder.rs b/crates/lib/database_queries/src/queries/tts/tts_results/list_tts_results_query_builder.rs new file mode 100644 index 0000000000..a60b5bc44e --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_results/list_tts_results_query_builder.rs @@ -0,0 +1,522 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use config::shared_constants::DEFAULT_MYSQL_QUERY_RESULT_PAGE_SIZE; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::record_visibility::RecordVisibility; +use log::{warn, info}; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct TtsInferenceListPage { + pub inference_records: Vec, + pub sort_ascending: bool, + + /// ID of the first record in `inference_records`. + pub first_id: Option, + + /// ID of the last record in `inference_records`. + pub last_id: Option, +} + +#[derive(Serialize)] +pub struct TtsInferenceRecordForList { + pub tts_result_token: String, + + pub tts_model_token: String, + pub tts_model_title: String, + pub raw_inference_text: String, + + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + + pub maybe_creator_result_id: Option, + + pub file_size_bytes: u32, + pub duration_millis: u32, + + pub visibility: RecordVisibility, + + //pub model_is_mod_approved: bool, // converted + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Query Builder for listing TTS results. +/// These are very difficult queries, so the builder helps for +/// testability, construction, and correctness. +pub struct ListTtsResultsQueryBuilder { + scope_creator_username: Option, + include_user_hidden: bool, + include_mod_disabled_results: bool, + sort_ascending: bool, + offset: Option, + limit: u16, + cursor_is_reversed: bool, +} + +impl ListTtsResultsQueryBuilder { + pub fn new() -> Self { + Self { + scope_creator_username: None, + include_user_hidden: false, + include_mod_disabled_results: false, + sort_ascending: false, + offset: None, + limit: DEFAULT_MYSQL_QUERY_RESULT_PAGE_SIZE, + cursor_is_reversed: false, + } + } + + pub fn scope_creator_username(mut self, scope_creator_username: Option<&str>) -> Self { + self.scope_creator_username = scope_creator_username.map(|u| u.to_string()); + self + } + + pub fn include_user_hidden(mut self, include_user_hidden: bool) -> Self { + self.include_user_hidden = include_user_hidden; + self + } + + pub fn include_mod_disabled_results(mut self, include_mod_disabled_results: bool) -> Self { + self.include_mod_disabled_results = include_mod_disabled_results; + self + } + + pub fn sort_ascending(mut self, sort_ascending: bool) -> Self { + self.sort_ascending = sort_ascending; + self + } + + pub fn offset(mut self, offset: Option) -> Self { + self.offset = offset; + self + } + + pub fn limit(mut self, limit: u16) -> Self { + self.limit = limit; + self + } + + pub fn cursor_is_reversed(mut self, cursor_is_reversed: bool) -> Self { + self.cursor_is_reversed = cursor_is_reversed; + self + } + + /// Perform the query based on the set predicates. + /// Wrap the results in a "Page" we can return to the frontend. + /// This will remove any IDs from the records, but the ids themselves also need scrubbing. + pub async fn perform_query_for_page( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult { + + let internal_results = self.perform_internal_query(mysql_pool).await?; + + let mut first_id = internal_results.first() + .map(|raw_result| raw_result.tts_result_id); + + let mut last_id = internal_results.last() + .map(|raw_result| raw_result.tts_result_id); + + let inference_results = internal_results + .iter() + .map(|r| { + TtsInferenceRecordForList { + tts_result_token: r.tts_result_token.clone(), + tts_model_token: r.tts_model_token.clone(), + tts_model_title: r.tts_model_title.clone(), + raw_inference_text: r.raw_inference_text.clone(), + maybe_creator_user_token: r.maybe_creator_user_token.clone(), + maybe_creator_username: r.maybe_creator_username.clone(), + maybe_creator_display_name: r.maybe_creator_display_name.clone(), + maybe_creator_result_id: r.maybe_creator_result_id.map(|v| v as u64).clone(), + file_size_bytes: if r.file_size_bytes > 0 { r.file_size_bytes as u32 } else { 0 }, + duration_millis: if r.duration_millis > 0 { r.duration_millis as u32 } else { 0 }, + visibility: RecordVisibility::from_str(&r.creator_set_visibility).unwrap_or(RecordVisibility::Public), + created_at: r.created_at, + updated_at: r.updated_at, + } + }) + .collect::>(); + + Ok(TtsInferenceListPage { + inference_records: inference_results, + sort_ascending: self.sort_ascending, + first_id, + last_id, + }) + } + + /// Perform the query based on the set predicates. + /// The records contain IDs, which we should not expose to the frontend. + pub async fn perform_internal_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + + let query = self.build_query_string(); + let mut query = sqlx::query_as::<_, RawInternalTtsRecord>(&query); + + // NB: The following bindings must match the order of the query builder !! + + if let Some(offset) = self.offset { + query = query.bind(offset); + } + + if let Some(username) = self.scope_creator_username.as_deref() { + query = query.bind(username); + } + + query = query.bind(self.limit); + + let mut results = query.fetch_all(mysql_pool) + .await?; + + if self.cursor_is_reversed { + results.reverse() + } + + Ok(results) + } + + pub fn build_query_string(&self) -> String { + // TODO: I haven't figured out how to get field name disambiguation and type coercion working here. + // (1) tts_results.creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Query error: no column found for name: creator_set_visibility + // (2) creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Column 'creator_set_visibility' in field list is ambiguous + + // TODO/NB: Unfortunately SQLx can't statically typecheck this query + let mut query = r#" +SELECT + tts_results.id as tts_result_id, + tts_results.token as tts_result_token, + + tts_results.model_token as tts_model_token, + tts_models.title as tts_model_title, + tts_results.raw_inference_text as raw_inference_text, + + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + + tts_results.maybe_creator_synthetic_id as maybe_creator_result_id, + + tts_results.file_size_bytes, + tts_results.duration_millis, + + tts_results.creator_set_visibility, + + tts_results.created_at, + tts_results.updated_at + +FROM tts_results +LEFT OUTER JOIN tts_models + ON tts_results.model_token = tts_models.token +LEFT OUTER JOIN users + ON tts_results.maybe_creator_user_token = users.token + "#.to_string(); + + query.push_str(&self.build_predicates()); + query + } + + pub fn build_predicates(&self) -> String { + // NB: Reverse cursors require us to invert the sort direction. + let mut sort_ascending = self.sort_ascending; + + let mut first_predicate_added = false; + + let mut query = "".to_string(); + + if let Some(offset) = self.offset { + if !first_predicate_added { + query.push_str(" WHERE"); + first_predicate_added = true; + } else { + query.push_str(" AND"); + } + + if sort_ascending { + if self.cursor_is_reversed { + // NB: We're searching backwards. + query.push_str(" tts_results.id < ?"); + sort_ascending = !sort_ascending; + } else { + query.push_str(" tts_results.id > ?"); + } + } else { + if self.cursor_is_reversed { + // NB: We're searching backwards. + query.push_str(" tts_results.id > ?"); + sort_ascending = !sort_ascending; + } else { + query.push_str(" tts_results.id < ?"); + } + } + } + + if let Some(username) = self.scope_creator_username.as_deref() { + if !first_predicate_added { + query.push_str(" WHERE users.username = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND users.username = ?"); + } + } + + if !self.include_user_hidden { + if !first_predicate_added { + query.push_str(" WHERE tts_results.creator_set_visibility = 'public'"); + first_predicate_added = true; + } else { + query.push_str(" AND tts_results.creator_set_visibility = 'public'"); + } + } + + if !self.include_mod_disabled_results { + if !first_predicate_added { + query.push_str(" WHERE tts_results.user_deleted_at IS NULL"); + query.push_str(" AND tts_results.mod_deleted_at IS NULL"); + first_predicate_added = true; + } else { + query.push_str(" AND tts_results.user_deleted_at IS NULL"); + query.push_str(" AND tts_results.mod_deleted_at IS NULL"); + } + } + + if sort_ascending { + query.push_str(" ORDER BY tts_results.id ASC"); + } else { + query.push_str(" ORDER BY tts_results.id DESC"); + } + + query.push_str(" LIMIT ?"); + + query + } +} + +#[derive(sqlx::FromRow)] +pub struct RawInternalTtsRecord { + pub tts_result_id: i64, + pub tts_result_token: String, + + pub tts_model_token: String, + pub tts_model_title: String, + pub raw_inference_text: String, + + pub maybe_creator_user_token : Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + + pub maybe_creator_result_id: Option, + + pub file_size_bytes : i64, + pub duration_millis : i64, + + pub creator_set_visibility: String, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[cfg(test)] +mod tests { + use crate::queries::tts::tts_results::list_tts_results_query_builder::ListTtsResultsQueryBuilder; + + #[test] + fn predicates_without_scoping() { + let query_builder = ListTtsResultsQueryBuilder::new(); + + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_scoped_to_user() { + let query_builder = ListTtsResultsQueryBuilder::new() + .scope_creator_username(Some("echelon")); + + assert_eq!(&query_builder.build_predicates(), + " WHERE users.username = ? \ + AND tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_user_hidden() { + let query_builder = ListTtsResultsQueryBuilder::new() + .include_user_hidden(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_deleted() { + let query_builder = ListTtsResultsQueryBuilder::new() + .include_mod_disabled_results(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.creator_set_visibility = 'public' \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_sort_ascending() { + let query_builder = ListTtsResultsQueryBuilder::new() + .sort_ascending(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id ASC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset() { + let query_builder = ListTtsResultsQueryBuilder::new() + .offset(Some(100)); + + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.id < ? \ + AND tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset_and_sort_ascending() { + let query_builder = ListTtsResultsQueryBuilder::new() + .sort_ascending(true) + .offset(Some(100)); + + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.id > ? \ + AND tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id ASC \ + LIMIT ?"); + } + + #[test] + fn predicates_limit() { + let query_builder = ListTtsResultsQueryBuilder::new() + .limit(15); + + // NB: Does not change the query itself! Just the downstream binding. + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_cursor_is_reversed_without_cursor() { + let query_builder = ListTtsResultsQueryBuilder::new() + .cursor_is_reversed(true); + + // NB: Without a cursor, nothing happens. + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset_cursor_is_reversed() { + let query_builder = ListTtsResultsQueryBuilder::new() + .offset(Some(100)) + .cursor_is_reversed(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.id > ? \ + AND tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id ASC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset_cursor_is_reversed_sort_ascending() { + let query_builder = ListTtsResultsQueryBuilder::new() + .offset(Some(100)) + .cursor_is_reversed(true) + .sort_ascending(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.id < ? \ + AND tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_limit_scope_user_offset_cursor_is_reversed_sort_ascending() { + let query_builder = ListTtsResultsQueryBuilder::new() + .limit(1000) + .scope_creator_username(Some("pikachu")) + .offset(Some(100)) + .cursor_is_reversed(true) + .sort_ascending(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.id < ? \ + AND users.username = ? \ + AND tts_results.creator_set_visibility = 'public' \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_limit_scope_user_show_hidden_offset_cursor_is_reversed_sort_ascending() { + let query_builder = ListTtsResultsQueryBuilder::new() + .limit(1000) + .include_user_hidden(true) + .scope_creator_username(Some("pikachu")) + .offset(Some(100)) + .cursor_is_reversed(true) + .sort_ascending(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE tts_results.id < ? \ + AND users.username = ? \ + AND tts_results.user_deleted_at IS NULL \ + AND tts_results.mod_deleted_at IS NULL \ + ORDER BY tts_results.id DESC \ + LIMIT ?"); + } +} diff --git a/crates/lib/database_queries/src/queries/tts/tts_results/mod.rs b/crates/lib/database_queries/src/queries/tts/tts_results/mod.rs new file mode 100644 index 0000000000..a4b88aa68b --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_results/mod.rs @@ -0,0 +1,4 @@ +pub mod delete_tts_result_various_scopes; +pub mod insert_tts_result; +pub mod list_tts_results_query_builder; +pub mod query_tts_result; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/tts/tts_results/query_tts_result.rs b/crates/lib/database_queries/src/queries/tts/tts_results/query_tts_result.rs new file mode 100644 index 0000000000..958145448b --- /dev/null +++ b/crates/lib/database_queries/src/queries/tts/tts_results/query_tts_result.rs @@ -0,0 +1,304 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::record_visibility::RecordVisibility; +use crate::column_types::vocoder_type::VocoderType; +use crate::helpers::boolean_converters::nullable_i8_to_bool; +use log::{info, warn, log}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct TtsResultRecordForResponse { + pub tts_result_token: String, + pub raw_inference_text: String, + + pub tts_model_token: String, + pub tts_model_title: Option, // TODO: Shouldn't be Option. + + pub maybe_pretrained_vocoder_used: Option, + + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + pub maybe_creator_gravatar_hash: Option, + + pub maybe_model_creator_user_token: Option, + pub maybe_model_creator_username: Option, + pub maybe_model_creator_display_name: Option, + pub maybe_model_creator_gravatar_hash: Option, + + pub public_bucket_wav_audio_path: String, + pub public_bucket_spectrogram_path: String, + + pub creator_set_visibility: RecordVisibility, + + pub file_size_bytes: u32, + pub duration_millis: u32, + + //pub model_is_mod_approved: bool, // converted + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub maybe_moderator_fields: Option, +} + +/// "Moderator-only fields" that we wouldn't want to expose to ordinary users. +/// It's the web endpoint controller's responsibility to clear these for non-mods. +#[derive(Serialize)] +pub struct TtsResultModeratorFields { + pub model_creator_is_banned: bool, + pub result_creator_is_banned_if_user: bool, + pub result_creator_ip_address: String, + pub result_creator_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +#[derive(Serialize)] +pub struct TtsResultRecordRaw { + pub tts_result_token: String, // from field `tts_results.token` + pub raw_inference_text: String, + + pub tts_model_token: String, + pub tts_model_title: Option, + + pub maybe_pretrained_vocoder_used: Option, + + pub maybe_creator_is_banned: Option, + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + pub maybe_creator_gravatar_hash: Option, + + pub maybe_model_creator_is_banned: Option, + pub maybe_model_creator_user_token: Option, + pub maybe_model_creator_username: Option, + pub maybe_model_creator_display_name: Option, + pub maybe_model_creator_gravatar_hash: Option, + + pub public_bucket_wav_audio_path: String, + pub public_bucket_spectrogram_path: String, + + pub creator_set_visibility: String, + + pub file_size_bytes: i32, + pub duration_millis: i32, + + //pub model_is_mod_approved: i8, // needs convert + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + + // Moderator fields + pub creator_ip_address: String, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +pub async fn select_tts_result_by_token( + tts_result_token: &str, + can_see_deleted: bool, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + let maybe_record = if can_see_deleted { + select_including_deleted(tts_result_token, mysql_pool).await + } else { + select_without_deleted(tts_result_token, mysql_pool).await + }; + + let ir : TtsResultRecordRaw = match maybe_record { + Ok(inference_result) => inference_result, + Err(ref err) => { + match err { + sqlx::Error::RowNotFound => { + warn!("tts result not found: {:?}", &err); + return Ok(None); + }, + _ => { + warn!("tts result query error: {:?}", &err); + return Err(anyhow!("database error")); + } + } + } + }; + + let mut pretrained_vocoder = None; + if let Some(vocoder) = ir.maybe_pretrained_vocoder_used.as_deref() { + pretrained_vocoder = Some(VocoderType::from_str(vocoder)?); + } + + let ir_for_response = TtsResultRecordForResponse { + tts_result_token: ir.tts_result_token, + + raw_inference_text: ir.raw_inference_text, + + tts_model_token: ir.tts_model_token, + tts_model_title: ir.tts_model_title, + + maybe_pretrained_vocoder_used: pretrained_vocoder, + + maybe_creator_user_token: ir.maybe_creator_user_token, + maybe_creator_username: ir.maybe_creator_username, + maybe_creator_display_name: ir.maybe_creator_display_name, + maybe_creator_gravatar_hash: ir.maybe_creator_gravatar_hash, + + maybe_model_creator_user_token: ir.maybe_model_creator_user_token, + maybe_model_creator_username: ir.maybe_model_creator_username, + maybe_model_creator_display_name: ir.maybe_model_creator_display_name, + maybe_model_creator_gravatar_hash: ir.maybe_model_creator_gravatar_hash, + + //model_is_mod_approved: if ir.model_is_mod_approved == 0 { false } else { true }, + + public_bucket_wav_audio_path: ir.public_bucket_wav_audio_path, + public_bucket_spectrogram_path: ir.public_bucket_spectrogram_path, + + // NB: Fail open/public since we're already looking at it + creator_set_visibility: RecordVisibility::from_str(&ir.creator_set_visibility) + .unwrap_or(RecordVisibility::Public), + + file_size_bytes: if ir.file_size_bytes > 0 { ir.file_size_bytes as u32 } else { 0 }, + duration_millis: if ir.duration_millis > 0 { ir.duration_millis as u32 } else { 0 }, + + created_at: ir.created_at, + updated_at: ir.updated_at, + + maybe_moderator_fields: Some(TtsResultModeratorFields { + model_creator_is_banned: + nullable_i8_to_bool(ir.maybe_model_creator_is_banned, false), + result_creator_is_banned_if_user: + nullable_i8_to_bool(ir.maybe_creator_is_banned, false), + result_creator_ip_address: ir.creator_ip_address, + result_creator_deleted_at: ir.user_deleted_at, + mod_deleted_at: ir.mod_deleted_at, + }), + }; + + Ok(Some(ir_for_response)) +} + +async fn select_including_deleted( + tts_result_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + TtsResultRecordRaw, + r#" +SELECT + tts_results.token as tts_result_token, + + tts_results.raw_inference_text, + + tts_results.model_token as tts_model_token, + tts_models.title as tts_model_title, + + tts_results.maybe_pretrained_vocoder_used, + + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + users.email_gravatar_hash as maybe_creator_gravatar_hash, + users.is_banned as maybe_creator_is_banned, + + model_users.token as maybe_model_creator_user_token, + model_users.username as maybe_model_creator_username, + model_users.display_name as maybe_model_creator_display_name, + model_users.email_gravatar_hash as maybe_model_creator_gravatar_hash, + model_users.is_banned as maybe_model_creator_is_banned, + + tts_results.public_bucket_wav_audio_path, + tts_results.public_bucket_spectrogram_path, + + tts_results.creator_set_visibility, + + tts_results.file_size_bytes, + tts_results.duration_millis, + tts_results.created_at, + tts_results.updated_at, + + tts_results.creator_ip_address, + tts_results.user_deleted_at, + tts_results.mod_deleted_at + +FROM tts_results +LEFT OUTER JOIN tts_models + ON tts_results.model_token = tts_models.token +LEFT OUTER JOIN users + ON tts_results.maybe_creator_user_token = users.token +LEFT OUTER JOIN users as model_users + ON tts_models.creator_user_token = model_users.token +WHERE + tts_results.token = ? + "#, + tts_result_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} + +async fn select_without_deleted( + tts_result_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + TtsResultRecordRaw, + r#" +SELECT + tts_results.token as tts_result_token, + + tts_results.raw_inference_text, + + tts_results.model_token as tts_model_token, + tts_models.title as tts_model_title, + + tts_results.maybe_pretrained_vocoder_used, + + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + users.email_gravatar_hash as maybe_creator_gravatar_hash, + users.is_banned as maybe_creator_is_banned, + + model_users.token as maybe_model_creator_user_token, + model_users.username as maybe_model_creator_username, + model_users.display_name as maybe_model_creator_display_name, + model_users.email_gravatar_hash as maybe_model_creator_gravatar_hash, + model_users.is_banned as maybe_model_creator_is_banned, + + tts_results.public_bucket_wav_audio_path, + tts_results.public_bucket_spectrogram_path, + + tts_results.creator_set_visibility, + + tts_results.file_size_bytes, + tts_results.duration_millis, + tts_results.created_at, + tts_results.updated_at, + + tts_results.creator_ip_address, + tts_results.user_deleted_at, + tts_results.mod_deleted_at + +FROM tts_results +LEFT OUTER JOIN tts_models + ON tts_results.model_token = tts_models.token +LEFT OUTER JOIN users + ON tts_results.maybe_creator_user_token = users.token +LEFT OUTER JOIN users as model_users + ON tts_models.creator_user_token = model_users.token +WHERE + tts_results.token = ? + AND tts_results.user_deleted_at IS NULL + AND tts_results.mod_deleted_at IS NULL + "#, + tts_result_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} diff --git a/crates/lib/database_queries/src/queries/twitch/mod.rs b/crates/lib/database_queries/src/queries/twitch/mod.rs new file mode 100644 index 0000000000..445b7129e7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/mod.rs @@ -0,0 +1,3 @@ +pub mod twitch_event_rules; +pub mod twitch_oauth; +pub mod twitch_pubsub; diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/delete_twitch_event_rule.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/delete_twitch_event_rule.rs new file mode 100644 index 0000000000..914856bb9d --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/delete_twitch_event_rule.rs @@ -0,0 +1,31 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +pub async fn delete_twitch_event_rule( + event_rule_token: &str, + ip_address_delete: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult { + + let query = sqlx::query!( + r#" +UPDATE twitch_event_rules +SET + deleted_at = CURRENT_TIMESTAMP, + ip_address_last_update = ? +WHERE + token = ? +LIMIT 1 + "#, + ip_address_delete, + event_rule_token, + ); + + let result = query.execute(mysql_pool).await; + + match result { + Err(err) => Err(anyhow!("error with query: {:?}", err)), + Ok(_r) => Ok(true), + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/get_twitch_event_rule_for_user.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/get_twitch_event_rule_for_user.rs new file mode 100644 index 0000000000..fb370cd058 --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/get_twitch_event_rule_for_user.rs @@ -0,0 +1,82 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::twitch_event_category::TwitchEventCategory; +use crate::helpers::boolean_converters::i8_to_bool; +use sqlx::MySqlPool; + +#[derive(Debug)] +pub struct TwitchEventRule { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: String, + pub event_response: String, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub deleted_at: Option>, +} + +/// Query record by token and user token +pub async fn get_twitch_event_rule_for_user( + twitch_event_rule_token: &str, + user_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + + let maybe_record = sqlx::query_as!( + TwitchEventRuleInternal, + r#" +SELECT + token, + event_category as `event_category: crate::column_types::twitch_event_category::TwitchEventCategory`, + event_match_predicate, + event_response, + user_specified_rule_order, + rule_is_disabled, + created_at, + updated_at, + deleted_at +FROM twitch_event_rules +WHERE + token = ? + AND user_token = ? + "#, + twitch_event_rule_token, + user_token, + ) + .fetch_one(pool) + .await; + + let record : TwitchEventRuleInternal = match maybe_record { + Ok(record) => record, + Err(sqlx::Error::RowNotFound) => return Ok(None), + Err(ref err) => return Err(anyhow!("database query error: {:?}", err)), + }; + + Ok(Some(TwitchEventRule { + token: record.token, + event_category: record.event_category, + event_match_predicate: record.event_match_predicate, + event_response: record.event_response, + user_specified_rule_order: record.user_specified_rule_order, + rule_is_disabled: i8_to_bool(record.rule_is_disabled), + created_at: record.created_at, + updated_at: record.updated_at, + deleted_at: record.deleted_at, + })) +} + +#[derive(Debug)] +struct TwitchEventRuleInternal { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: String, + pub event_response: String, + pub user_specified_rule_order: u32, + pub rule_is_disabled: i8, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub deleted_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/insert_twitch_event_rule_builder.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/insert_twitch_event_rule_builder.rs new file mode 100644 index 0000000000..d186c9a060 --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/insert_twitch_event_rule_builder.rs @@ -0,0 +1,61 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::twitch_event_category::TwitchEventCategory; +use crate::tokens::Tokens; +use sqlx::MySqlPool; + +pub struct InsertTwitchEventRuleBuilder { + pub uuid_idempotency_token: String, + pub user_token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: String, + pub event_response: String, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, + pub ip_address_creation: String, +} + +impl InsertTwitchEventRuleBuilder { + + /// Returns the newly generated token. + pub async fn insert(&self, mysql_pool: &MySqlPool) -> AnyhowResult { + + let token = Tokens::new_twitch_event_rule()?; + + let query = sqlx::query!( + r#" +INSERT INTO twitch_event_rules +SET + uuid_idempotency_token = ?, + token = ?, + user_token = ?, + event_category = ?, + event_match_predicate = ?, + event_response = ?, + user_specified_rule_order = ?, + rule_is_disabled = ?, + ip_address_creation = ?, + ip_address_last_update = ? + "#, + &self.uuid_idempotency_token, + &token, + &self.user_token, + &self.event_category.to_str(), + &self.event_match_predicate, + &self.event_response, + &self.user_specified_rule_order, + &self.rule_is_disabled, + &self.ip_address_creation, + &self.ip_address_creation, + ); + + let query_result = query.execute(mysql_pool).await; + + let _record_id = match query_result { + Ok(res) => res.last_insert_id(), + Err(err) => return Err(anyhow!("Error creating Twitch Event Rule: {:?}", err)), + }; + + Ok(token) + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/list_twitch_event_rules_for_user.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/list_twitch_event_rules_for_user.rs new file mode 100644 index 0000000000..6d7976dbde --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/list_twitch_event_rules_for_user.rs @@ -0,0 +1,79 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::twitch_event_category::TwitchEventCategory; +use crate::helpers::boolean_converters::i8_to_bool; +use sqlx::MySqlPool; + +#[derive(Debug, Serialize)] +pub struct TwitchEventRule { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: String, + pub event_response: String, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} + +/// Query non-deleted Twitch event rules for a user +pub async fn list_twitch_event_rules_for_user( + user_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + + let records : Vec = sqlx::query_as!( + TwitchEventRuleInternal, + r#" +SELECT + token, + event_category as `event_category: crate::column_types::twitch_event_category::TwitchEventCategory`, + event_match_predicate, + event_response, + user_specified_rule_order, + rule_is_disabled, + created_at, + updated_at +FROM twitch_event_rules +WHERE + user_token = ? + AND deleted_at IS NULL + "#, + user_token, + ) + .fetch_all(pool) + .await?; + + let mut records = records.into_iter() + .map(|record : TwitchEventRuleInternal | { + TwitchEventRule { + token: record.token, + event_category: record.event_category, + event_match_predicate: record.event_match_predicate, + event_response: record.event_response, + user_specified_rule_order: record.user_specified_rule_order, + rule_is_disabled: i8_to_bool(record.rule_is_disabled), + created_at: record.created_at, + updated_at: record.updated_at, + } + }) + .collect::>(); + + // Queried in DESC order, but sort returned results ordered ASC. + records.sort_by_key(|r| r.user_specified_rule_order); + + Ok(records) +} + +#[derive(Debug)] +struct TwitchEventRuleInternal { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: String, + pub event_response: String, + pub user_specified_rule_order: u32, + pub rule_is_disabled: i8, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/mod.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/mod.rs new file mode 100644 index 0000000000..b0db364aac --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/mod.rs @@ -0,0 +1,6 @@ +pub mod delete_twitch_event_rule; +pub mod get_twitch_event_rule_for_user; +pub mod insert_twitch_event_rule_builder; +pub mod list_twitch_event_rules_for_user; +pub mod reorder_twitch_event_rules; +pub mod update_twitch_event_rule_builder; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/reorder_twitch_event_rules.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/reorder_twitch_event_rules.rs new file mode 100644 index 0000000000..76a3f874b2 --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/reorder_twitch_event_rules.rs @@ -0,0 +1,132 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::error; +use once_cell::sync::Lazy; +use regex::Regex; +use sqlx::MySqlPool; +use std::collections::HashMap; + +/// Check that tokens are alphanumeric or can have colons. +static SAFE_TOKENS_REGEX : Lazy = Lazy::new(|| { + Regex::new("^[a-zA-Z0-9:]+$").unwrap() +}); + +/// Check that tokens are alphanumeric or can have colons. +#[must_use] +fn check_token_safety(token: &str) -> AnyhowResult<()> { + if SAFE_TOKENS_REGEX.is_match(token) { + Ok(()) + } else { + Err(anyhow!("Token contains invalid characters")) + } +} + +pub async fn reorder_twitch_event_rules( + rule_token_to_order_map: HashMap, + user_token: &str, + ip_address_update: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult { + + // THIS IS NOT SANITIZED, so we check for injection here. + for (rule_token, _position) in rule_token_to_order_map.iter() { + check_token_safety(rule_token.as_str())?; + } + + check_token_safety(user_token)?; + +// NB: Sadly upserts fail due to other required primary keys that we cannot supply here (eg. uuid). +// let query_inner = rule_token_to_order_map.iter() +// .map(|(token, position)| { +// format!("(\"{}\", {}, \"{}\")", token.as_str(), position, ip_address_update) +// }) +// .collect::>() +// .join(", "); +// +// // Upsert to accomplish multiple updates at once. +// // tl;dr https://stackoverflow.com/a/34866431 +// // ALSO THIS IS NOT SANITIZED. +// let query = format!(r#" +//INSERT INTO twitch_event_rules (token, user_specified_rule_order, ip_address_last_update) +//VALUES +// {} +//ON DUPLICATE KEY UPDATE +// token=VALUES(token), +// user_specified_rule_order=VALUES(user_specified_rule_order), +// ip_address_last_update=VALUES(ip_address_last_update) +// "#, query_inner); + + let rule_tokens = rule_token_to_order_map.iter() + .map(|(token, _position)| { + format!("'{}'", token) + }) + .collect::>() + .join(", "); + + let query_inner_positions = rule_token_to_order_map.iter() + .map(|(token, position)| { + format!("when token = '{}' then '{}'", token.as_str(), position) + }) + .collect::>() + .join("\n"); + + let query_inner_ip_addresses = rule_token_to_order_map.iter() + .map(|(token, position)| { + format!("when token = '{}' then '{}'", token.as_str(), ip_address_update) + }) + .collect::>() + .join("\n"); + + let query = format!(r#" +UPDATE twitch_event_rules +SET + user_specified_rule_order = (case + {} + end), + ip_address_last_update = (case + {} + end) + WHERE user_token = '{}' + AND deleted_at IS NULL + AND token IN ({}) + "#, query_inner_positions, query_inner_ip_addresses, user_token, rule_tokens); + + // NB: This, unfortunately, cannot be statically checked. + let mut query = sqlx::query(&query); + + let result = query.execute(mysql_pool).await; + + match result { + Err(err) => Err(anyhow!("error with query: {:?}", err)), + Ok(_r) => Ok(true), + } +} + +#[cfg(test)] +mod tests { + use crate::queries::twitch::twitch_event_rules::reorder_twitch_event_rules::check_token_safety; + + #[test] + fn safe_tokens() { + assert!(check_token_safety("foo").is_ok()); + assert!(check_token_safety("FOO").is_ok()); + assert!(check_token_safety("12345").is_ok()); + assert!(check_token_safety("abcdefABCDEF012345679").is_ok()); + assert!(check_token_safety("PREFIX:TOKEN1234").is_ok()); + } + + #[test] + fn unsafe_tokens() { + assert!(check_token_safety("").is_err()); + assert!(check_token_safety(",").is_err()); + assert!(check_token_safety("ABCD,").is_err()); + assert!(check_token_safety("foo;").is_err()); + assert!(check_token_safety("foo;foo").is_err()); + assert!(check_token_safety("foo;foo").is_err()); + assert!(check_token_safety("\"foo").is_err()); + assert!(check_token_safety("foo\"").is_err()); + assert!(check_token_safety("foo'").is_err()); + assert!(check_token_safety("foo'bar").is_err()); + assert!(check_token_safety("foo`foo").is_err()); + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/update_twitch_event_rule_builder.rs b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/update_twitch_event_rule_builder.rs new file mode 100644 index 0000000000..08b07b384f --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_event_rules/update_twitch_event_rule_builder.rs @@ -0,0 +1,47 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::tokens::Tokens; +use sqlx::MySqlPool; + +// TODO: This doesn't handle sparse updates. + +pub struct UpdateTwitchEventRuleBuilder { + pub token: String, + pub event_match_predicate: String, + pub event_response: String, + pub rule_is_disabled: bool, + pub ip_address_update: String, +} + +impl UpdateTwitchEventRuleBuilder { + + pub async fn update(&self, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + + let query = sqlx::query!( + r#" +UPDATE twitch_event_rules +SET + event_match_predicate = ?, + event_response = ?, + rule_is_disabled = ?, + ip_address_last_update = ? +WHERE + token = ? + AND deleted_at IS NULL +LIMIT 1 + "#, + &self.event_match_predicate, + &self.event_response, + &self.rule_is_disabled, + &self.ip_address_update, + &self.token, + ); + + let query_result = query.execute(mysql_pool).await; + + match query_result { + Err(err) => return Err(anyhow!("Error updating Twitch Event Rule: {:?}", err)), + Ok(_r) => Ok(()), + } + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_oauth/find.rs b/crates/lib/database_queries/src/queries/twitch/twitch_oauth/find.rs new file mode 100644 index 0000000000..09274990a9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_oauth/find.rs @@ -0,0 +1,290 @@ +//! Return the *SINGLE MOST RECENT* Twitch OAuth token per the search params. + +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use log::error; +use log::warn; +use sqlx::MySqlPool; + +#[derive(Serialize, Clone)] +pub struct TwitchOauthTokenRecord { + /// Our internal token metadata / bookkeeping + pub internal_token: String, + pub oauth_refresh_grouping_token: String, + pub refresh_count: u32, + pub ip_address_creation: Option, + + /// NB: Vocodes/FakeYou/Storyteller user + pub maybe_user_token: Option, + pub maybe_user_display_name: Option, + pub maybe_user_gravatar_hash: Option, + + /// Twitch user + pub twitch_user_id: String, + pub twitch_username: String, + pub twitch_username_lowercase: String, + + /// Twitch token details + pub access_token: String, + pub maybe_refresh_token: Option, + pub token_type: Option, + pub expires_in_seconds: Option, + pub has_bits_read: bool, + pub has_channel_read_redemptions: bool, + pub has_channel_read_subscriptions: bool, + pub has_chat_edit: bool, + pub has_chat_read: bool, + + /// Potentially when the token is expected to expire. + /// Do not eagerly refresh. Lazily renew. + pub expires_at: Option>, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +/// Return the *SINGLE MOST RECENT* Twitch OAuth token per the search params. +pub struct TwitchOauthTokenFinder { + /// Storyteller/FakeYou username + scope_user_token: Option, + + scope_twitch_user_id: Option, + scope_twitch_username_lowercase: Option, + + /// Allow return of expired tokens? Expired tokens can be renewed. + allow_expired_tokens: bool, +} + +impl TwitchOauthTokenFinder { + pub fn new() -> Self { + Self { + scope_user_token: None, + scope_twitch_user_id: None, + scope_twitch_username_lowercase: None, + allow_expired_tokens: false, + } + } + + pub fn scope_user_token(mut self, user_token: Option<&str>) -> Self { + self.scope_user_token = user_token.map(|t| t.to_string()); + self + } + + pub fn scope_twitch_user_id(mut self, user_id: Option) -> Self { + self.scope_twitch_user_id = user_id; + self + } + + /// This will automatically look up against a lowercased username + pub fn scope_twitch_username(mut self, twitch_username: Option<&str>) -> Self { + self.scope_twitch_username_lowercase = twitch_username + .map(|t| t.to_string().to_lowercase()); + self + } + + pub fn allow_expired_tokens(mut self, allow_expired: bool) -> Self { + self.allow_expired_tokens = allow_expired; + self + } + + pub async fn perform_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + Ok(self.perform_query_internal(mysql_pool).await? + .map(|record| { + TwitchOauthTokenRecord { + internal_token: record.internal_token, + oauth_refresh_grouping_token: record.oauth_refresh_grouping_token, + maybe_user_token: record.maybe_user_token, + maybe_user_display_name: record.maybe_user_display_name, + maybe_user_gravatar_hash: record.maybe_user_gravatar_hash, + twitch_user_id: record.twitch_user_id, + twitch_username: record.twitch_username, + twitch_username_lowercase: record.twitch_username_lowercase, + access_token: record.access_token, + maybe_refresh_token: record.maybe_refresh_token, + token_type: record.token_type, + expires_in_seconds: record.expires_in_seconds, + refresh_count: record.refresh_count, + has_bits_read: i8_to_bool(record.has_bits_read), + has_channel_read_redemptions: i8_to_bool(record.has_channel_read_redemptions), + has_channel_read_subscriptions: i8_to_bool(record.has_channel_read_subscriptions), + has_chat_edit: i8_to_bool(record.has_chat_edit), + has_chat_read: i8_to_bool(record.has_chat_read), + ip_address_creation: record.ip_address_creation, + expires_at: record.expires_at, + user_deleted_at: record.user_deleted_at, + mod_deleted_at: record.mod_deleted_at, + } + })) + } + + pub async fn perform_query_internal( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + + let query = self.build_query_string(); + let mut query = sqlx::query_as::<_, TwitchOauthTokenRecordInternal>(&query); + //let mut query = sqlx::query(&query); + + // NB: The following bindings must match the order of the query builder !! + + if let Some(user_token) = self.scope_user_token.as_deref() { + query = query.bind(user_token); + } + + if let Some(user_id) = self.scope_twitch_user_id { + query = query.bind(user_id); + } + + if let Some(twitch_username) = self.scope_twitch_username_lowercase.as_deref() { + query = query.bind(twitch_username); + } + + let mut result = query.fetch_optional(mysql_pool).await; + + match result { + Ok(Some(record)) => { + Ok(Some(record)) + } + Ok(None) => { + Ok(None) + }, + Err(err) => { + error!("twitch oauth token query error: {:?}", err); + Err(anyhow!("twitch oauth token query error: {:?}", err)) + } + } + } + + pub fn build_query_string(&self) -> String { + // TODO/NB: Unfortunately SQLx can't statically typecheck this query + let mut query = r#" +SELECT + twitch_oauth_tokens.internal_token, + twitch_oauth_tokens.oauth_refresh_grouping_token, + users.username as maybe_username, + users.display_name as maybe_user_display_name, + users.email_gravatar_hash as maybe_user_gravatar_hash, + twitch_oauth_tokens.maybe_user_token, + + twitch_oauth_tokens.twitch_user_id, + twitch_oauth_tokens.twitch_username, + twitch_oauth_tokens.twitch_username_lowercase, + + twitch_oauth_tokens.access_token, + twitch_oauth_tokens.maybe_refresh_token, + twitch_oauth_tokens.token_type, + twitch_oauth_tokens.expires_in_seconds, + twitch_oauth_tokens.refresh_count, + twitch_oauth_tokens.has_bits_read, + twitch_oauth_tokens.has_channel_read_redemptions, + twitch_oauth_tokens.has_channel_read_subscriptions, + twitch_oauth_tokens.has_chat_edit, + twitch_oauth_tokens.has_chat_read, + twitch_oauth_tokens.ip_address_creation, + twitch_oauth_tokens.created_at, + twitch_oauth_tokens.updated_at, + twitch_oauth_tokens.expires_at, + twitch_oauth_tokens.user_deleted_at, + twitch_oauth_tokens.mod_deleted_at + +FROM twitch_oauth_tokens +LEFT OUTER JOIN users + ON twitch_oauth_tokens.maybe_user_token = users.token + "#.to_string(); + + query.push_str(&self.build_predicates()); + + query + } + + pub fn build_predicates(&self) -> String { + let mut query = "".to_string(); + let mut first_predicate_added = false; + + if let Some(_user_token) = self.scope_user_token.as_deref() { + if !first_predicate_added { + query.push_str(" WHERE users.token = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND users.token = ?"); + } + } + + if let Some(_user_id) = self.scope_twitch_user_id { + if !first_predicate_added { + query.push_str(" WHERE twitch_oauth_tokens.twitch_user_id = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND twitch_oauth_tokens.twitch_user_id = ?"); + } + } + + if let Some(_username) = self.scope_twitch_username_lowercase.as_deref() { + if !first_predicate_added { + query.push_str(" WHERE twitch_oauth_tokens.twitch_username_lowercase = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND twitch_oauth_tokens.twitch_username_lowercase = ?"); + } + } + + if !self.allow_expired_tokens { + if !first_predicate_added { + query.push_str(" WHERE ( twitch_oauth_tokens.expires_at IS NULL OR twitch_oauth_tokens.expires_at > NOW() ) "); + } else { + query.push_str(" AND ( twitch_oauth_tokens.expires_at IS NULL OR twitch_oauth_tokens.expires_at > NOW() ) "); + } + } + + // NB: Return the most recent. + // This might have stupid race condition/transaction weirdness. + // Sorry future me. + query.push_str(" ORDER BY twitch_oauth_tokens.id DESC"); + + // NB: Only a single record. + query.push_str(" LIMIT 1"); + + query + } +} + +#[derive(sqlx::FromRow)] +pub struct TwitchOauthTokenRecordInternal { + /// Our internal token metadata / bookkeeping + pub internal_token: String, + pub oauth_refresh_grouping_token: String, + pub refresh_count: u32, + pub ip_address_creation: Option, + + /// NB: Vocodes/FakeYou/Storyteller user + pub maybe_user_token: Option, + pub maybe_user_display_name: Option, + pub maybe_user_gravatar_hash: Option, + + /// Twitch user + pub twitch_user_id: String, + pub twitch_username: String, + pub twitch_username_lowercase: String, + + /// Token details + pub access_token: String, + pub maybe_refresh_token: Option, + pub token_type: Option, + pub expires_in_seconds: Option, + pub has_bits_read: i8, + pub has_channel_read_subscriptions: i8, + pub has_channel_read_redemptions: i8, + pub has_chat_edit: i8, + pub has_chat_read: i8, + + /// Potentially when the token is expected to expire. + /// Do not eagerly refresh. Lazily renew. + pub expires_at: Option>, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_oauth/insert.rs b/crates/lib/database_queries/src/queries/twitch/twitch_oauth/insert.rs new file mode 100644 index 0000000000..3786075c5d --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_oauth/insert.rs @@ -0,0 +1,242 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use log::info; +use log::warn; +use sqlx::MySqlPool; +use crate::tokens::Tokens; + +pub struct TwitchOauthTokenInsertBuilder { + // ===== Required Fields ===== + + /// Old APIs return a u32, but this should be a string. + /// Twitch is migrating IDs to strings. + twitch_user_id: String, + + /// Twitch username / channel name. + twitch_username: String, + + /// The secret we use + access_token: String, + + /// Internal mechanism we use to group tokens across refresh events + oauth_grouping_token: String, + + // ===== Optional / Default Fields ===== + + /// Other half of the access token, for renewal + maybe_refresh_token: Option, + + /// Storyteller/FakeYou username. + maybe_user_token: Option, + + /// Abuse tracking. + maybe_ip_address_creation: Option, + + /// Probably 'bearer' + token_type: Option, + + /// TTL seconds from issue + expires_in_seconds: Option, + + /// Number of times we've refreshed the token + refresh_count: u32, + + // ===== OAuth Scopes ===== + + has_bits_read: bool, + has_channel_read_redemptions: bool, + has_channel_read_subscriptions: bool, + has_chat_edit: bool, + has_chat_read: bool, +} + +impl TwitchOauthTokenInsertBuilder { + pub fn new( + twitch_user_id: &str, + twitch_username: &str, + access_token: &str, + oauth_grouping_token: &str + ) -> Self { + Self { + twitch_user_id: twitch_user_id.to_string(), + twitch_username: twitch_username.to_string(), + access_token: access_token.to_string(), + oauth_grouping_token: oauth_grouping_token.to_string(), + maybe_refresh_token: None, + maybe_user_token: None, + maybe_ip_address_creation: None, + token_type: None, + expires_in_seconds: None, + has_bits_read: false, + has_channel_read_subscriptions: false, + has_channel_read_redemptions: false, + has_chat_edit: false, + has_chat_read: false, + refresh_count: 0, + } + } + + pub fn set_refresh_token(mut self, refresh_token: Option<&str>) -> Self { + self.maybe_refresh_token = refresh_token.map(|t| t.to_string()); + self + } + + /// NB: This is the FakeYou/Storyteller user account, not Twitch! + pub fn set_user_token(mut self, maybe_user_token: Option<&str>) -> Self { + self.maybe_user_token = maybe_user_token.map(|t| t.to_string()); + self + } + + pub fn set_ip_address_creation(mut self, ip_address_creation: Option<&str>) -> Self { + self.maybe_ip_address_creation = ip_address_creation.map(|t| t.to_string()); + self + } + + pub fn set_token_type(mut self, token_type: Option<&str>) -> Self { + self.token_type = token_type.map(|t| t.to_string()); + self + } + + pub fn set_expires_in_seconds(mut self, expires_in_seconds: Option) -> Self { + self.expires_in_seconds = expires_in_seconds; + self + } + + pub fn set_refresh_count(mut self, refresh_count: u32) -> Self { + self.refresh_count = refresh_count; + self + } + + pub fn has_bits_read(mut self, value: bool) -> Self { + self.has_bits_read = value; + self + } + + pub fn has_channel_read_redemptions(mut self, value: bool) -> Self { + self.has_channel_read_redemptions = value; + self + } + + pub fn has_channel_read_subscriptions(mut self, value: bool) -> Self { + self.has_channel_read_subscriptions = value; + self + } + + pub fn has_chat_edit(mut self, value: bool) -> Self { + self.has_chat_edit = value; + self + } + + pub fn has_chat_read(mut self, value: bool) -> Self { + self.has_chat_read = value; + self + } + + pub async fn insert(&mut self, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + // An internal token that we may use in the future. + let internal_token = Tokens::new_twitch_oauth_internal_token()?; + + // NB: We have to duplicate these since the string literals must not + // include concatenation. Boo. + let query = if let Some(expires_in_seconds) = self.expires_in_seconds { + sqlx::query!( + r#" +INSERT INTO twitch_oauth_tokens +SET + internal_token = ?, + oauth_refresh_grouping_token = ?, + maybe_user_token = ?, + twitch_user_id = ?, + twitch_username = ?, + twitch_username_lowercase = ?, + access_token = ?, + maybe_refresh_token = ?, + refresh_count = ?, + token_type = ?, + expires_in_seconds = ?, + has_bits_read = ?, + has_channel_read_redemptions= ?, + has_channel_read_subscriptions = ?, + has_chat_edit = ?, + has_chat_read = ?, + ip_address_creation = ?, + expires_at = DATE_ADD(NOW(), INTERVAL ? SECOND) + "#, + internal_token.clone(), + self.oauth_grouping_token.clone(), + self.maybe_user_token.clone(), + self.twitch_user_id.clone(), + self.twitch_username.clone(), + self.twitch_username.clone().to_lowercase(), + self.access_token.clone(), + self.maybe_refresh_token.clone(), + self.refresh_count, + self.token_type.clone(), + expires_in_seconds, + self.has_bits_read.clone(), + self.has_channel_read_redemptions.clone(), + self.has_channel_read_subscriptions.clone(), + self.has_chat_edit.clone(), + self.has_chat_read.clone(), + self.maybe_ip_address_creation.clone(), + expires_in_seconds, + ) + } else { + sqlx::query!( + r#" +INSERT INTO twitch_oauth_tokens +SET + internal_token = ?, + oauth_refresh_grouping_token = ?, + maybe_user_token = ?, + twitch_user_id = ?, + twitch_username = ?, + twitch_username_lowercase = ?, + access_token = ?, + maybe_refresh_token = ?, + refresh_count = ?, + token_type = ?, + expires_in_seconds = NULL, + has_bits_read = ?, + has_channel_read_redemptions= ?, + has_channel_read_subscriptions = ?, + has_chat_edit = ?, + has_chat_read = ?, + ip_address_creation = ?, + expires_at = NULL + "#, + internal_token.clone(), + self.oauth_grouping_token.clone(), + self.maybe_user_token.clone(), + self.twitch_user_id.clone(), + self.twitch_username.clone(), + self.twitch_username.clone().to_lowercase(), + self.access_token.clone(), + self.maybe_refresh_token.clone(), + self.refresh_count, + self.token_type.clone(), + self.has_bits_read.clone(), + self.has_channel_read_redemptions.clone(), + self.has_channel_read_subscriptions.clone(), + self.has_chat_edit.clone(), + self.has_chat_read.clone(), + self.maybe_ip_address_creation.clone(), + ) + }; + + let query_result = query.execute(mysql_pool) + .await; + + let _record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("Twitch token insert DB error: {:?}", err)); + } + }; + + Ok(()) + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_oauth/mod.rs b/crates/lib/database_queries/src/queries/twitch/twitch_oauth/mod.rs new file mode 100644 index 0000000000..38216ccb54 --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_oauth/mod.rs @@ -0,0 +1,2 @@ +pub mod find; +pub mod insert; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/insert_bits.rs b/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/insert_bits.rs new file mode 100644 index 0000000000..889d25dd4e --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/insert_bits.rs @@ -0,0 +1,144 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use log::info; +use log::warn; +use sqlx::MySqlPool; + +pub struct TwitchPubsubBitsInsertBuilder { + // ===== Required Fields ===== + + sender_twitch_user_id: Option, + sender_twitch_username: Option, + sender_twitch_username_lowercase: Option, + + destination_channel_id: Option, + destination_channel_name: Option, + + // ===== Optional / Default Fields ===== + + is_anonymous: bool, + + bits_used: u64, + total_bits_used: u64, + + chat_message: String, +} + +impl TwitchPubsubBitsInsertBuilder { + pub fn new() -> Self { + Self { + sender_twitch_user_id: None, + sender_twitch_username: None, + sender_twitch_username_lowercase: None, + destination_channel_id: None, + destination_channel_name: None, + is_anonymous: false, + bits_used: 0, + total_bits_used: 0, + chat_message: "".to_string(), + } + } + + pub fn set_sender_twitch_user_id(mut self, value: &str) -> Self { + self.sender_twitch_user_id = Some(value.to_string()); + self + } + + pub fn set_sender_twitch_username(mut self, value: &str) -> Self { + self.sender_twitch_username = Some(value.to_string()); + self.sender_twitch_username_lowercase = Some(value.to_lowercase()); + self + } + + pub fn set_destination_channel_id(mut self, value: &str) -> Self { + self.destination_channel_id = Some(value.to_string()); + self + } + + pub fn set_destination_channel_name(mut self, value: &str) -> Self { + self.destination_channel_name = Some(value.to_string()); + self + } + + pub fn set_is_anonymous(mut self, value: bool) -> Self { + self.is_anonymous = value; + self + } + + pub fn set_bits_used(mut self, value: u64) -> Self { + self.bits_used = value; + self + } + + pub fn set_total_bits_used(mut self, value: u64) -> Self { + self.total_bits_used = value; + self + } + + pub fn set_chat_message(mut self, value: &str) -> Self { + self.chat_message = value.to_string(); + self + } + + pub async fn insert(&mut self, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + let sender_twitch_user_id = self.sender_twitch_user_id + .clone() + .ok_or(anyhow!("no sender_twitch_user_id"))?; + + let sender_twitch_username = self.sender_twitch_username + .clone() + .ok_or(anyhow!("no sender_twitch_username"))?; + + let sender_twitch_username_lowercase = self.sender_twitch_username_lowercase + .clone() + .ok_or(anyhow!("no sender_twitch_username_lowercase"))?; + + let destination_channel_id = self.destination_channel_id + .clone() + .ok_or(anyhow!("no destination_channel_id"))?; + + let destination_channel_name = self.destination_channel_name + .clone() + .ok_or(anyhow!("no destination_channel_name"))?; + + let query = sqlx::query!( + r#" +INSERT INTO twitch_bits_events +SET + sender_twitch_user_id = ?, + sender_twitch_username = ?, + sender_twitch_username_lowercase = ?, + destination_channel_id = ?, + destination_channel_name = ?, + is_anonymous= ?, + bits_used = ?, + total_bits_used = ?, + chat_message = ? + "#, + sender_twitch_user_id, + sender_twitch_username, + sender_twitch_username_lowercase, + destination_channel_id, + destination_channel_name, + self.is_anonymous, + self.bits_used, + self.total_bits_used, + self.chat_message.clone(), + ); + + let query_result = query.execute(mysql_pool) + .await; + + let _record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("Twitch pubsub bits insert DB error: {:?}", err)); + } + }; + + Ok(()) + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/insert_channel_points.rs b/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/insert_channel_points.rs new file mode 100644 index 0000000000..001055e6ba --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/insert_channel_points.rs @@ -0,0 +1,193 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use log::info; +use log::warn; +use sqlx::MySqlPool; + +pub struct TwitchPubsubChannelPointsInsertBuilder { + // ===== Required Fields ===== + + sender_twitch_user_id: Option, + sender_twitch_username: Option, + sender_twitch_username_lowercase: Option, + + destination_channel_id: Option, + destination_channel_name: Option, + + // ===== Optional / Default Fields ===== + + title: String, + prompt: String, + + user_text_input: String, + + redemption_id: String, + reward_id: String, + + reward_cost: u64, + + is_sub_only: bool, + max_per_stream: u64, + max_per_user_per_stream: u64, +} + +impl TwitchPubsubChannelPointsInsertBuilder { + pub fn new() -> Self { + Self { + sender_twitch_user_id: None, + sender_twitch_username: None, + sender_twitch_username_lowercase: None, + destination_channel_id: None, + destination_channel_name: None, + title: "".to_string(), + prompt: "".to_string(), + user_text_input: "".to_string(), + redemption_id: "".to_string(), + reward_id: "".to_string(), + reward_cost: 0, + is_sub_only: false, + max_per_stream: 0, + max_per_user_per_stream: 0, + } + } + + pub fn set_sender_twitch_user_id(mut self, value: &str) -> Self { + self.sender_twitch_user_id = Some(value.to_string()); + self + } + + pub fn set_sender_twitch_username(mut self, value: &str) -> Self { + self.sender_twitch_username = Some(value.to_string()); + self.sender_twitch_username_lowercase = Some(value.to_lowercase()); + self + } + + pub fn set_destination_channel_id(mut self, value: &str) -> Self { + self.destination_channel_id = Some(value.to_string()); + self + } + + pub fn set_destination_channel_name(mut self, value: &str) -> Self { + self.destination_channel_name = Some(value.to_string()); + self + } + + pub fn set_title(mut self, value: &str) -> Self { + self.title = value.to_string(); + self + } + + pub fn set_prompt(mut self, value: &str) -> Self { + self.prompt = value.to_string(); + self + } + + pub fn set_user_text_input(mut self, value: Option<&str>) -> Self { + self.user_text_input = value + .map(|s| s.to_string()) + .unwrap_or("".to_string()); + self + } + + pub fn set_redemption_id(mut self, value: &str) -> Self { + self.redemption_id = value.to_string(); + self + } + + pub fn set_reward_id(mut self, value: &str) -> Self { + self.reward_id = value.to_string(); + self + } + + pub fn set_is_sub_only(mut self, value: bool) -> Self { + self.is_sub_only = value; + self + } + + pub fn set_reward_cost(mut self, value: u64) -> Self { + self.reward_cost = value; + self + } + + pub fn set_max_per_stream(mut self, value: u64) -> Self { + self.max_per_stream = value; + self + } + + pub fn set_max_per_user_per_stream(mut self, value: u64) -> Self { + self.max_per_user_per_stream = value; + self + } + + pub async fn insert(&mut self, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + let sender_twitch_user_id = self.sender_twitch_user_id + .clone() + .ok_or(anyhow!("no sender_twitch_user_id"))?; + + let sender_twitch_username = self.sender_twitch_username + .clone() + .ok_or(anyhow!("no sender_twitch_username"))?; + + let sender_twitch_username_lowercase = self.sender_twitch_username_lowercase + .clone() + .ok_or(anyhow!("no sender_twitch_username_lowercase"))?; + + let destination_channel_id = self.destination_channel_id + .clone() + .ok_or(anyhow!("no destination_channel_id"))?; + + let destination_channel_name = self.destination_channel_name + .clone() + .ok_or(anyhow!("no destination_channel_name"))?; + + let query = sqlx::query!( + r#" +INSERT INTO twitch_channel_point_events +SET + sender_twitch_user_id = ?, + sender_twitch_username = ?, + sender_twitch_username_lowercase = ?, + destination_channel_id = ?, + destination_channel_name = ?, + title = ?, + prompt = ?, + user_text_input = ?, + redemption_id = ?, + reward_id = ?, + reward_cost = ?, + is_sub_only = ?, + max_per_stream = ?, + max_per_user_per_stream = ? + "#, + sender_twitch_user_id, + sender_twitch_username, + sender_twitch_username_lowercase, + destination_channel_id, + destination_channel_name, + self.title.clone(), + self.prompt.clone(), + self.user_text_input.clone(), + self.redemption_id.clone(), + self.reward_id.clone(), + self.reward_cost, + self.is_sub_only, + self.max_per_stream, + self.max_per_user_per_stream, + ); + + let query_result = query.execute(mysql_pool) + .await; + + let _record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("Twitch pubsub bits insert DB error: {:?}", err)); + } + }; + + Ok(()) + } +} diff --git a/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/mod.rs b/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/mod.rs new file mode 100644 index 0000000000..1f8758da1b --- /dev/null +++ b/crates/lib/database_queries/src/queries/twitch/twitch_pubsub/mod.rs @@ -0,0 +1,2 @@ +pub mod insert_bits; +pub mod insert_channel_points; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/users/get_user_profile_by_username.rs b/crates/lib/database_queries/src/queries/users/get_user_profile_by_username.rs new file mode 100644 index 0000000000..f7c739d4a4 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/get_user_profile_by_username.rs @@ -0,0 +1,120 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use log::{info, warn, log}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +// TODO: This duplicates the get_profile_handler. + +#[derive(Serialize)] +pub struct UserProfileResult { + pub user_token: String, + pub username: String, + pub display_name: String, + pub email_gravatar_hash: String, + pub profile_markdown: String, + pub profile_rendered_html: String, + pub user_role_slug: String, + pub is_banned: bool, + pub disable_gravatar: bool, + pub discord_username: Option, + pub twitch_username: Option, + pub twitter_username: Option, + pub patreon_username: Option, + pub github_username: Option, + pub created_at: DateTime, +} + +#[derive(Serialize)] +struct RawUserProfileRecord { + pub user_token: String, + pub username: String, + pub email_gravatar_hash: String, + pub display_name: String, + pub profile_markdown: String, + pub profile_rendered_html: String, + pub user_role_slug: String, + pub is_banned: i8, + pub disable_gravatar: i8, + pub discord_username: Option, + pub twitch_username: Option, + pub twitter_username: Option, + pub patreon_username: Option, + pub github_username: Option, + pub created_at: DateTime, +} + +pub async fn get_user_profile_by_username( + username: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + let maybe_profile_record = sqlx::query_as!( + RawUserProfileRecord, + r#" +SELECT + users.token as user_token, + username, + display_name, + email_gravatar_hash, + profile_markdown, + profile_rendered_html, + user_role_slug, + is_banned, + disable_gravatar, + discord_username, + twitch_username, + twitter_username, + patreon_username, + github_username, + created_at +FROM users +WHERE + users.username = ? + AND users.user_deleted_at IS NULL + AND users.mod_deleted_at IS NULL + "#, + username, + ) + .fetch_one(mysql_pool) + .await; + + let profile_record : RawUserProfileRecord = match maybe_profile_record { + Ok(profile_record) => profile_record, + Err(err) => { + return match err { + sqlx::Error::RowNotFound => { + Ok(None) + }, + _ => { + warn!("User profile query error: {:?}", err); + Err(anyhow!("query error")) + } + } + } + }; + + let profile_for_response = UserProfileResult { + user_token: profile_record.user_token.clone(), + username: profile_record.username.clone(), + display_name: profile_record.display_name.clone(), + email_gravatar_hash: profile_record.email_gravatar_hash.clone(), + profile_markdown: profile_record.profile_markdown.clone(), + profile_rendered_html: profile_record.profile_rendered_html.clone(), + user_role_slug: profile_record.user_role_slug.clone(), + is_banned: i8_to_bool(profile_record.is_banned), + disable_gravatar: i8_to_bool(profile_record.disable_gravatar), + discord_username: profile_record.discord_username.clone(), + twitch_username: profile_record.twitch_username.clone(), + twitter_username: profile_record.twitter_username.clone(), + patreon_username: profile_record.patreon_username.clone(), + github_username: profile_record.github_username.clone(), + created_at: profile_record.created_at.clone(), + }; + + Ok(Some(profile_for_response)) +} diff --git a/crates/lib/database_queries/src/queries/users/list_users_query_builder.rs b/crates/lib/database_queries/src/queries/users/list_users_query_builder.rs new file mode 100644 index 0000000000..6c40f62bc0 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/list_users_query_builder.rs @@ -0,0 +1,158 @@ +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use sqlx::MySqlPool; + +/// NB: This is not to be shared externally. +/// Only for trusted mods and staff. +#[derive(Serialize)] +pub struct ListUsersPage { + pub users: Vec, + pub sorted_ascending: bool, + pub sorted_by_key: String, + + /// ID of the first record in `users` + pub first_id: Option, + + /// ID of the last record in `users` + pub last_id: Option, +} + +#[derive(Serialize)] +pub struct UserForList { + pub user_id: i64, + pub user_token: String, + + pub username: String, + pub display_name: String, + pub gravatar_hash: String, + + pub is_banned: bool, + pub user_role_slug: String, + + pub ip_address_creation: String, + pub ip_address_last_login: String, + pub ip_address_last_update: String, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +pub struct ListUsersQueryBuilder { + sort_ascending: bool, + per_page: i64, +} + +impl ListUsersQueryBuilder { + pub fn new() -> Self { + Self { + sort_ascending: false, + per_page: 50, + } + } + + pub fn sort_ascending(mut self, sort_ascending: bool) -> Self { + self.sort_ascending = sort_ascending; + self + } + + pub async fn query_for_page( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult { + + let internal_results = self.internal_query(mysql_pool).await?; + + let mut first_id = internal_results.first() + .map(|raw_result| raw_result.user_id); + + let mut last_id = internal_results.last() + .map(|raw_result| raw_result.user_id); + + let users = internal_results + .into_iter() + .map(|r| { + UserForList { + user_id: r.user_id, + user_token: r.user_token, + username: r.username, + display_name: r.display_name, + gravatar_hash: r.gravatar_hash, + is_banned: i8_to_bool(r.is_banned), + user_role_slug: r.user_role_slug, + ip_address_creation: r.ip_address_creation, + ip_address_last_login: r.ip_address_last_login, + ip_address_last_update: r.ip_address_last_update, + created_at: r.created_at, + updated_at: r.updated_at, + } + }) + .collect::>(); + + Ok(ListUsersPage { + users, + sorted_ascending: false, + sorted_by_key: "".to_string(), + first_id, + last_id, + }) + } + + async fn internal_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + + let query = self.build_query_string(); + let mut query = sqlx::query_as::<_, UserForListRaw>(&query); + + + let mut results = query.fetch_all(mysql_pool) + .await?; + + Ok(results) + } + + pub fn build_query_string(&self) -> String { + let mut query_string = r#" +SELECT + id as user_id, + token as user_token, + username, + display_name, + display_name, + email_gravatar_hash as gravatar_hash, + is_banned, + user_role_slug, + ip_address_creation, + ip_address_last_login, + ip_address_last_update, + created_at, + updated_at +FROM users +"#; + + query_string.to_string() + } +} + +#[derive(sqlx::FromRow)] +struct UserForListRaw { + pub user_id: i64, + pub user_token: String, + + pub username: String, + pub display_name: String, + pub gravatar_hash: String, + + pub is_banned: i8, + pub user_role_slug: String, + + pub ip_address_creation: String, + pub ip_address_last_login: String, + pub ip_address_last_update: String, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + diff --git a/crates/lib/database_queries/src/queries/users/mod.rs b/crates/lib/database_queries/src/queries/users/mod.rs new file mode 100644 index 0000000000..a4b95c51c7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/mod.rs @@ -0,0 +1,5 @@ +pub mod get_user_profile_by_username; +pub mod list_users_query_builder; +pub mod user_badges; +pub mod user_roles; +pub mod user_sessions; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/users/user_badges/list_user_badges.rs b/crates/lib/database_queries/src/queries/users/user_badges/list_user_badges.rs new file mode 100644 index 0000000000..80769eb9a0 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/user_badges/list_user_badges.rs @@ -0,0 +1,76 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use log::{warn, info}; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct UserBadgeForList { + pub slug: String, + pub title: String, + pub description: String, + pub image_url: String, + pub granted_at: DateTime, +} + +struct RawDbUserBadgeForList { + slug: String, + title: String, + description: String, + image_url: String, + user_created_at : DateTime, +} + +pub async fn list_user_badges( + mysql_pool: &MySqlPool, + user_token: &str, +) -> AnyhowResult> { + info!("listing user badges"); + let maybe_user_badges = sqlx::query_as!( + RawDbUserBadgeForList, + r#" +SELECT + badges.slug, + badges.title, + badges.description, + badges.image_url, + user_badges.created_at as user_created_at + +FROM badges +JOIN user_badges +ON + badges.slug = user_badges.badge_slug +WHERE + user_badges.user_token = ? + "#, + user_token + ) + .fetch_all(mysql_pool) + .await; + + let user_badges : Vec = match maybe_user_badges { + Ok(badges) => badges, + Err(err) => { + warn!("Error: {:?}", err); + match err { + sqlx::Error::RowNotFound => Vec::new(), + _ => { + warn!("user badges query error: {:?}", err); + return Err(anyhow!("error querying user badges")); + } + } + } + }; + + Ok(user_badges.into_iter() + .map(|badge| { + UserBadgeForList { + slug: badge.slug, + title: badge.title, + description: badge.description, + image_url: badge.image_url, + granted_at: badge.user_created_at, + } + }) + .collect::>()) +} diff --git a/crates/lib/database_queries/src/queries/users/user_badges/mod.rs b/crates/lib/database_queries/src/queries/users/user_badges/mod.rs new file mode 100644 index 0000000000..cd1a70c4b8 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/user_badges/mod.rs @@ -0,0 +1 @@ +pub mod list_user_badges; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/users/user_roles/list_user_roles.rs b/crates/lib/database_queries/src/queries/users/user_roles/list_user_roles.rs new file mode 100644 index 0000000000..652c5402dc --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/user_roles/list_user_roles.rs @@ -0,0 +1,154 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use log::{warn, info}; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct UserRoleForList { + pub slug: String, + pub name: String, + + pub can_use_tts: bool, + pub can_use_w2l: bool, + pub can_delete_own_tts_results: bool, + pub can_delete_own_w2l_results: bool, + pub can_delete_own_account: bool, + + pub can_upload_tts_models: bool, + pub can_upload_w2l_templates: bool, + pub can_delete_own_tts_models: bool, + pub can_delete_own_w2l_templates: bool, + + pub can_approve_w2l_templates: bool, + pub can_edit_other_users_profiles: bool, + pub can_edit_other_users_tts_models: bool, + pub can_edit_other_users_w2l_templates: bool, + pub can_delete_other_users_tts_models: bool, + pub can_delete_other_users_tts_results: bool, + pub can_delete_other_users_w2l_templates: bool, + pub can_delete_other_users_w2l_results: bool, + pub can_ban_users: bool, + pub can_delete_users: bool, + + pub created_at : DateTime, + pub updated_at : DateTime, +} + +struct RawDbUserRoleForList { + slug: String, + name: String, + + can_use_tts: i8, + can_use_w2l: i8, + can_delete_own_tts_results: i8, + can_delete_own_w2l_results: i8, + can_delete_own_account: i8, + + can_upload_tts_models: i8, + can_upload_w2l_templates: i8, + can_delete_own_tts_models: i8, + can_delete_own_w2l_templates: i8, + + can_approve_w2l_templates: i8, + can_edit_other_users_profiles: i8, + can_edit_other_users_tts_models: i8, + can_edit_other_users_w2l_templates: i8, + can_delete_other_users_tts_models: i8, + can_delete_other_users_tts_results: i8, + can_delete_other_users_w2l_templates: i8, + can_delete_other_users_w2l_results: i8, + can_ban_users: i8, + can_delete_users: i8, + + created_at : DateTime, + updated_at : DateTime, +} + +pub async fn list_user_roles( + mysql_pool: &MySqlPool, +) -> AnyhowResult> { + info!("listing user roles"); + let maybe_user_roles = sqlx::query_as!( + RawDbUserRoleForList, + r#" +SELECT + slug, + name, + + can_use_tts, + can_use_w2l, + can_delete_own_tts_results, + can_delete_own_w2l_results, + can_delete_own_account, + + can_upload_tts_models, + can_upload_w2l_templates, + can_delete_own_tts_models, + can_delete_own_w2l_templates, + + can_approve_w2l_templates, + can_edit_other_users_profiles, + can_edit_other_users_tts_models, + can_edit_other_users_w2l_templates, + can_delete_other_users_tts_models, + can_delete_other_users_tts_results, + can_delete_other_users_w2l_templates, + can_delete_other_users_w2l_results, + can_ban_users, + can_delete_users, + + created_at, + updated_at + +FROM user_roles + "#, + ) + .fetch_all(mysql_pool) + .await; + + let user_roles : Vec = match maybe_user_roles { + Ok(roles) => roles, + Err(err) => { + warn!("Error: {:?}", err); + match err { + sqlx::Error::RowNotFound => Vec::new(), + _ => { + warn!("user role query error: {:?}", err); + return Err(anyhow!("error querying user roles")); + } + } + } + }; + + Ok(user_roles.into_iter() + .map(|ur| { + UserRoleForList { + slug: ur.slug, + name: ur.name, + can_use_tts: i8_to_bool(ur.can_use_tts), + can_use_w2l: i8_to_bool(ur.can_use_w2l), + can_delete_own_tts_results: i8_to_bool(ur.can_delete_own_tts_results), + can_delete_own_w2l_results: i8_to_bool(ur.can_delete_own_w2l_results), + can_delete_own_account: i8_to_bool(ur.can_delete_own_account), + can_upload_tts_models: i8_to_bool(ur.can_upload_tts_models), + can_upload_w2l_templates: i8_to_bool(ur.can_upload_w2l_templates), + can_delete_own_tts_models: i8_to_bool(ur.can_delete_own_tts_models), + can_delete_own_w2l_templates: i8_to_bool(ur.can_delete_own_w2l_templates), + can_approve_w2l_templates: i8_to_bool(ur.can_approve_w2l_templates), + can_edit_other_users_profiles: i8_to_bool(ur.can_edit_other_users_profiles), + can_edit_other_users_tts_models: i8_to_bool(ur.can_edit_other_users_tts_models), + can_edit_other_users_w2l_templates: i8_to_bool(ur.can_edit_other_users_w2l_templates), + can_delete_other_users_tts_models: i8_to_bool(ur.can_delete_other_users_tts_models), + can_delete_other_users_tts_results: i8_to_bool(ur.can_delete_other_users_tts_results), + can_delete_other_users_w2l_templates: i8_to_bool(ur.can_delete_other_users_w2l_templates), + can_delete_other_users_w2l_results: i8_to_bool(ur.can_delete_other_users_w2l_results), + can_ban_users: i8_to_bool(ur.can_ban_users), + can_delete_users: i8_to_bool(ur.can_delete_users), + created_at: ur.created_at, + updated_at: ur.updated_at, + } + }) + .collect::>()) +} diff --git a/crates/lib/database_queries/src/queries/users/user_roles/mod.rs b/crates/lib/database_queries/src/queries/users/user_roles/mod.rs new file mode 100644 index 0000000000..e227b96ba9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/user_roles/mod.rs @@ -0,0 +1 @@ +pub mod list_user_roles; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/users/user_sessions/create_session.rs b/crates/lib/database_queries/src/queries/users/user_sessions/create_session.rs new file mode 100644 index 0000000000..3748195e67 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/user_sessions/create_session.rs @@ -0,0 +1,41 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::tokens::Tokens; +use log::{info, warn}; +use sqlx::MySqlPool; + +pub async fn create_session_for_user(user_token: &str, ip_address: &str, mysql_pool: &MySqlPool) + -> AnyhowResult +{ + let session_token = Tokens::new_session()?; + + let query_result = sqlx::query!( + r#" +INSERT INTO user_sessions ( + token, + user_token, + ip_address_creation, + expires_at +) +VALUES ( ?, ?, ?, NOW() + interval 1 year ) + "#, + session_token.to_string(), + user_token.to_string(), + ip_address.to_string(), + ) + .execute(mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("session creation DB error: {:?}", err)); + } + }; + + info!("Created session id: {}", record_id); + + Ok(session_token) +} diff --git a/crates/lib/database_queries/src/queries/users/user_sessions/mod.rs b/crates/lib/database_queries/src/queries/users/user_sessions/mod.rs new file mode 100644 index 0000000000..6a3f6dcb74 --- /dev/null +++ b/crates/lib/database_queries/src/queries/users/user_sessions/mod.rs @@ -0,0 +1 @@ +pub mod create_session; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/voice_clone_requests/get_voice_clone_request.rs b/crates/lib/database_queries/src/queries/voice_clone_requests/get_voice_clone_request.rs new file mode 100644 index 0000000000..8da3ac7f51 --- /dev/null +++ b/crates/lib/database_queries/src/queries/voice_clone_requests/get_voice_clone_request.rs @@ -0,0 +1,236 @@ +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::twitch_event_category::TwitchEventCategory; +use crate::helpers::boolean_converters::i8_to_bool; +use sqlx::MySqlPool; + +/// Lookup by voice clone request token +pub async fn get_voice_clone_request_by_token( + voice_clone_request_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + internal_voice_clone_request_by_token(voice_clone_request_token, pool).await +} + +/// Lookup by user +pub async fn get_voice_clone_request_by_user_token( + user_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + internal_voice_clone_request_by_user_token(user_token, pool).await +} + +/// Used for both *BY TOKEN* and for *BY USER* lookups. +/// The user field should be optional for one and not for the other, but code reuse is more +/// important and this is a quick / hasty feature. +#[derive(Debug)] +pub struct VoiceCloneRequest { + pub token: String, + + // Contact + pub maybe_user_token: Option, + pub email_address: String, + pub discord_username: String, + + // Visibility + pub is_for_private_use: bool, + pub is_for_public_use: bool, + + // Use + pub is_for_studio: bool, + pub is_for_twitch_tts: bool, + pub is_for_api_use: bool, + pub is_for_music: bool, + pub is_for_games: bool, + pub is_for_other: bool, + pub optional_notes_on_use: Option, + + // Subject/Ownership + pub is_own_voice: bool, + pub is_third_party_voice: bool, + pub optional_notes_on_subject: Option, + + // Equipment + pub has_clean_audio_recordings: bool, + pub has_good_microphone: bool, + + // Comments + pub optional_questions: Option, + pub optional_extra_comments: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} + +async fn internal_voice_clone_request_by_token( + voice_clone_request_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + let maybe_record = sqlx::query_as!( + VoiceCloneRequestInternal, + r#" +SELECT + token, + + maybe_user_token, + email_address, + discord_username, + + is_for_private_use, + is_for_public_use, + + is_for_studio, + is_for_twitch_tts, + is_for_api_use, + is_for_music, + is_for_games, + is_for_other, + optional_notes_on_use, + + is_own_voice, + is_third_party_voice, + optional_notes_on_subject, + + has_clean_audio_recordings, + has_good_microphone, + + optional_questions, + optional_extra_comments, + + created_at, + updated_at +FROM voice_clone_requests +WHERE + token = ? + "#, + voice_clone_request_token, + ) + .fetch_one(pool) + .await; + + handle_result(maybe_record) +} + +async fn internal_voice_clone_request_by_user_token( + user_token: &str, + pool: &MySqlPool, +) -> AnyhowResult> { + let maybe_record = sqlx::query_as!( + VoiceCloneRequestInternal, + r#" +SELECT + token, + + maybe_user_token, + email_address, + discord_username, + + is_for_private_use, + is_for_public_use, + + is_for_studio, + is_for_twitch_tts, + is_for_api_use, + is_for_music, + is_for_games, + is_for_other, + optional_notes_on_use, + + is_own_voice, + is_third_party_voice, + optional_notes_on_subject, + + has_clean_audio_recordings, + has_good_microphone, + + optional_questions, + optional_extra_comments, + + created_at, + updated_at +FROM voice_clone_requests +WHERE + maybe_user_token = ? + "#, + user_token, + ) + .fetch_one(pool) + .await; + + handle_result(maybe_record) +} + +fn handle_result( + maybe_record: Result, +) -> AnyhowResult> { + let record : VoiceCloneRequestInternal = match maybe_record { + Ok(record) => record, + Err(sqlx::Error::RowNotFound) => return Ok(None), + Err(ref err) => return Err(anyhow!("database query error: {:?}", err)), + }; + + Ok(Some(VoiceCloneRequest { + token: record.token, + maybe_user_token: record.maybe_user_token, + email_address: record.email_address, + discord_username: record.discord_username, + is_for_private_use: i8_to_bool(record.is_for_private_use), + is_for_public_use: i8_to_bool(record.is_for_public_use), + is_for_studio: i8_to_bool(record.is_for_studio), + is_for_twitch_tts: i8_to_bool(record.is_for_twitch_tts), + is_for_api_use: i8_to_bool(record.is_for_api_use), + is_for_music: i8_to_bool(record.is_for_music), + is_for_games: i8_to_bool(record.is_for_games), + is_for_other: i8_to_bool(record.is_for_other), + optional_notes_on_use: record.optional_notes_on_use, + is_own_voice: i8_to_bool(record.is_own_voice), + is_third_party_voice: i8_to_bool(record.is_third_party_voice), + optional_notes_on_subject: record.optional_notes_on_subject, + has_clean_audio_recordings: i8_to_bool(record.has_clean_audio_recordings), + has_good_microphone: i8_to_bool(record.has_good_microphone), + optional_questions: record.optional_questions, + optional_extra_comments: record.optional_extra_comments, + created_at: record.created_at, + updated_at: record.updated_at, + })) +} + +#[derive(Debug)] +struct VoiceCloneRequestInternal { + pub token: String, + + // Contact + pub maybe_user_token: Option, + pub email_address: String, + pub discord_username: String, + + // Visibility + pub is_for_private_use: i8, + pub is_for_public_use: i8, + + // Use + pub is_for_studio: i8, + pub is_for_twitch_tts: i8, + pub is_for_api_use: i8, + pub is_for_music: i8, + pub is_for_games: i8, + pub is_for_other: i8, + pub optional_notes_on_use: Option, + + // Subject/Ownership + pub is_own_voice: i8, + pub is_third_party_voice: i8, + pub optional_notes_on_subject: Option, + + // Equipment + pub has_clean_audio_recordings: i8, + pub has_good_microphone: i8, + + // Comments + pub optional_questions: Option, + pub optional_extra_comments: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} diff --git a/crates/lib/database_queries/src/queries/voice_clone_requests/insert_voice_clone_requests_builder.rs b/crates/lib/database_queries/src/queries/voice_clone_requests/insert_voice_clone_requests_builder.rs new file mode 100644 index 0000000000..009d87c05a --- /dev/null +++ b/crates/lib/database_queries/src/queries/voice_clone_requests/insert_voice_clone_requests_builder.rs @@ -0,0 +1,125 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::twitch_event_category::TwitchEventCategory; +use crate::tokens::Tokens; +use sqlx::MySqlPool; + +pub struct InsertVoiceCloneRequestBuilder { + pub uuid_idempotency_token: String, + + // Contact + pub maybe_user_token: Option, + pub email_address: String, + pub discord_username: String, + + // Visibility + pub is_for_private_use: bool, + pub is_for_public_use: bool, + + // Use + pub is_for_studio: bool, + pub is_for_twitch_tts: bool, + pub is_for_api_use: bool, + pub is_for_music: bool, + pub is_for_games: bool, + pub is_for_other: bool, + pub optional_notes_on_use: Option, + + // Subject/Ownership + pub is_own_voice: bool, + pub is_third_party_voice: bool, + pub optional_notes_on_subject: Option, + + // Equipment + pub has_clean_audio_recordings: bool, + pub has_good_microphone: bool, + + // Comments + pub optional_questions: Option, + pub optional_extra_comments: Option, + + pub ip_address_creation: String, +} + +impl InsertVoiceCloneRequestBuilder { + + /// Returns the newly generated token. + pub async fn insert(&self, mysql_pool: &MySqlPool) -> AnyhowResult { + + let token = Tokens::new_voice_clone_request_token()?; + + let query = sqlx::query!( + r#" +INSERT INTO voice_clone_requests +SET + token = ?, + uuid_idempotency_token = ?, + + maybe_user_token = ?, + email_address = ?, + discord_username = ?, + + is_for_private_use = ?, + is_for_public_use = ?, + + is_for_studio = ?, + is_for_twitch_tts = ?, + is_for_api_use = ?, + is_for_music = ?, + is_for_games = ?, + is_for_other = ?, + optional_notes_on_use = ?, + + is_own_voice = ?, + is_third_party_voice = ?, + optional_notes_on_subject = ?, + + has_clean_audio_recordings = ?, + has_good_microphone = ?, + + optional_questions = ?, + optional_extra_comments = ?, + + ip_address_creation = ? + "#, + &token, + &self.uuid_idempotency_token, + + &self.maybe_user_token, + &self.email_address, + &self.discord_username, + + &self.is_for_private_use, + &self.is_for_public_use, + + &self.is_for_studio, + &self.is_for_twitch_tts, + &self.is_for_api_use, + &self.is_for_music, + &self.is_for_games, + &self.is_for_other, + &self.optional_notes_on_use, + + &self.is_own_voice, + &self.is_third_party_voice, + &self.optional_notes_on_subject, + + &self.has_clean_audio_recordings, + &self.has_good_microphone, + + &self.optional_questions, + &self.optional_extra_comments, + + &self.ip_address_creation, + ); + + let query_result = query.execute(mysql_pool).await; + + let _record_id = match query_result { + Ok(res) => res.last_insert_id(), + Err(err) => return Err(anyhow!("Error creating Voice Clone Request: {:?}", err)), + }; + + Ok(token) + } +} diff --git a/crates/lib/database_queries/src/queries/voice_clone_requests/mod.rs b/crates/lib/database_queries/src/queries/voice_clone_requests/mod.rs new file mode 100644 index 0000000000..cd31493d36 --- /dev/null +++ b/crates/lib/database_queries/src/queries/voice_clone_requests/mod.rs @@ -0,0 +1,2 @@ +pub mod get_voice_clone_request; +pub mod insert_voice_clone_requests_builder; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/w2l/mod.rs b/crates/lib/database_queries/src/queries/w2l/mod.rs new file mode 100644 index 0000000000..9cd08d628b --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/mod.rs @@ -0,0 +1,4 @@ +pub mod stats; +pub mod w2l_inference_jobs; +pub mod w2l_results; +pub mod w2l_templates; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/w2l/stats/calculate_w2l_template_leaderboard.rs b/crates/lib/database_queries/src/queries/w2l/stats/calculate_w2l_template_leaderboard.rs new file mode 100644 index 0000000000..52ed9d5979 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/stats/calculate_w2l_template_leaderboard.rs @@ -0,0 +1,85 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use log::warn; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct W2lLeaderboardRecordForList { + pub username: String, + pub display_name: String, + pub gravatar_hash: String, + pub creator_user_token: String, + pub uploaded_count: i64, +} + +#[derive(Serialize)] +struct W2lLeaderboardRecordForListRaw { + pub username: String, + pub display_name: String, + pub gravatar_hash: String, + pub creator_user_token: String, + pub uploaded_count: i64, +} + +// NB: This may need to become an offline query if it is expensive. +pub async fn calculate_w2l_template_leaderboard( + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + // NB: We're requiring "is_public_listing_approved IS TRUE" for W2L templates! + // This is the opposite for TTS models (for now at least) + let maybe_results = sqlx::query_as!( + W2lLeaderboardRecordForListRaw, + r#" +SELECT + users.username, + users.display_name, + users.email_gravatar_hash as gravatar_hash, + creator_user_token, + count(*) as uploaded_count +FROM w2l_templates +JOIN users +ON + users.token = creator_user_token +WHERE + is_public_listing_approved IS TRUE + AND users.is_banned IS FALSE + AND w2l_templates.user_deleted_at IS NULL + AND w2l_templates.mod_deleted_at IS NULL +GROUP BY creator_user_token +ORDER BY uploaded_count desc +LIMIT 25; + "#, + ) + .fetch_all(mysql_pool) + .await; + + let results : Vec = match maybe_results { + Ok(results) => results, + Err(err) => { + warn!("Error: {:?}", err); + + match err { + sqlx::Error::RowNotFound => { + return Ok(Vec::new()); + }, + _ => { + warn!("w2l leaderboard query error: {:?}", err); + return Err(anyhow!("w2l leaderboard query error")); + } + } + } + }; + + Ok(results.into_iter() + .map(|result| { + W2lLeaderboardRecordForList { + username: result.username.to_string(), + display_name: result.display_name.to_string(), + gravatar_hash: result.gravatar_hash.to_string(), + creator_user_token: result.creator_user_token.to_string(), + uploaded_count: result.uploaded_count, + } + }) + .collect::>()) +} diff --git a/crates/lib/database_queries/src/queries/w2l/stats/mod.rs b/crates/lib/database_queries/src/queries/w2l/stats/mod.rs new file mode 100644 index 0000000000..0a326d2cd9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/stats/mod.rs @@ -0,0 +1 @@ +pub mod calculate_w2l_template_leaderboard; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_inference_jobs/mod.rs b/crates/lib/database_queries/src/queries/w2l/w2l_inference_jobs/mod.rs new file mode 100644 index 0000000000..4e398c43ff --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_inference_jobs/mod.rs @@ -0,0 +1 @@ +// TODO: Move the "job queries" here and break them down into several files. \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_results/delete_w2l_result_various_scopes.rs b/crates/lib/database_queries/src/queries/w2l/w2l_results/delete_w2l_result_various_scopes.rs new file mode 100644 index 0000000000..adaced9ac7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_results/delete_w2l_result_various_scopes.rs @@ -0,0 +1,88 @@ +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +pub async fn delete_w2l_inference_result_as_user( + inference_result_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_results +SET + user_deleted_at = CURRENT_TIMESTAMP +WHERE + token = ? +LIMIT 1 + "#, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn delete_w2l_inference_result_as_mod( + inference_result_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_results +SET + mod_deleted_at = CURRENT_TIMESTAMP, + maybe_mod_user_token = ? +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_w2l_inference_result_as_user( + inference_result_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_results +SET + user_deleted_at = NULL +WHERE + token = ? +LIMIT 1 + "#, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_w2l_inference_result_as_mod( + inference_result_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_results +SET + mod_deleted_at = NULL, + maybe_mod_user_token = ? +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + inference_result_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_results/list_w2l_inference_results_query_builder.rs b/crates/lib/database_queries/src/queries/w2l/w2l_results/list_w2l_inference_results_query_builder.rs new file mode 100644 index 0000000000..6c4a6f1760 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_results/list_w2l_inference_results_query_builder.rs @@ -0,0 +1,660 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use config::shared_constants::DEFAULT_MYSQL_QUERY_RESULT_PAGE_SIZE; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::record_visibility::RecordVisibility; +use log::{warn, info}; +use sqlx::MySqlPool; + +#[derive(Serialize)] +pub struct W2lInferenceListPage { + pub inference_records: Vec, + pub sort_ascending: bool, + + /// ID of the first record in `inference_records`. + pub first_id: Option, + + /// ID of the last record in `inference_records`. + pub last_id: Option, +} + +#[derive(Serialize)] +pub struct W2lInferenceRecordForList { + pub w2l_result_token: String, + pub maybe_w2l_template_token: Option, + pub maybe_tts_inference_result_token: Option, + + pub template_type: Option, + pub template_title: Option, + + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + + pub maybe_creator_result_id: Option, + + pub file_size_bytes: u32, + pub frame_width: u32, + pub frame_height: u32, + pub duration_millis: u32, + + pub visibility: RecordVisibility, + + //pub template_is_mod_approved: bool, // converted + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Query Builder for listing W2L results. +/// These are very difficult queries, so the builder helps for +/// testability, construction, and correctness. +pub struct ListW2lResultsQueryBuilder { + scope_creator_username: Option, + include_user_hidden: bool, + include_mod_deleted_results: bool, + include_user_deleted_results: bool, + include_templates_not_approved_for_public_listing: bool, + sort_ascending: bool, + offset: Option, + limit: u16, + cursor_is_reversed: bool, +} + +impl ListW2lResultsQueryBuilder { + pub fn new() -> Self { + Self { + scope_creator_username: None, + include_user_hidden: false, + include_mod_deleted_results: false, + include_user_deleted_results: false, + include_templates_not_approved_for_public_listing: false, + sort_ascending: false, + offset: None, + limit: DEFAULT_MYSQL_QUERY_RESULT_PAGE_SIZE, + cursor_is_reversed: false, + } + } + + pub fn scope_creator_username(mut self, scope_creator_username: Option<&str>) -> Self { + self.scope_creator_username = scope_creator_username.map(|u| u.to_string()); + self + } + + pub fn include_user_hidden(mut self, include_user_hidden: bool) -> Self { + self.include_user_hidden = include_user_hidden; + self + } + + pub fn include_mod_deleted_results(mut self, include_mod_deleted_results: bool) -> Self { + self.include_mod_deleted_results = include_mod_deleted_results; + self + } + + pub fn include_user_deleted_results(mut self, include_user_deleted_results: bool) -> Self { + self.include_user_deleted_results = include_user_deleted_results; + self + } + + pub fn include_templates_not_approved_for_public_listing(mut self, include_templates_not_approved_for_public_listing: bool) -> Self { + self.include_templates_not_approved_for_public_listing = include_templates_not_approved_for_public_listing; + self + } + + pub fn sort_ascending(mut self, sort_ascending: bool) -> Self { + self.sort_ascending = sort_ascending; + self + } + + pub fn offset(mut self, offset: Option) -> Self { + self.offset = offset; + self + } + + pub fn limit(mut self, limit: u16) -> Self { + self.limit = limit; + self + } + + pub fn cursor_is_reversed(mut self, cursor_is_reversed: bool) -> Self { + self.cursor_is_reversed = cursor_is_reversed; + self + } + + /// Perform the query based on the set predicates. + /// Wrap the results in a "Page" we can return to the frontend. + /// This will remove any IDs from the records, but the ids themselves also need scrubbing. + pub async fn perform_query_for_page( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult { + + let internal_results = self.perform_internal_query(mysql_pool).await?; + + let mut first_id = internal_results.first() + .map(|raw_result| raw_result.w2l_result_id); + + let mut last_id = internal_results.last() + .map(|raw_result| raw_result.w2l_result_id); + + let inference_results = internal_results + .iter() + .map(|r| { + W2lInferenceRecordForList { + w2l_result_token: r.w2l_result_token.clone(), + maybe_w2l_template_token: r.maybe_w2l_template_token.clone(), + maybe_tts_inference_result_token: r.maybe_tts_inference_result_token.clone(), + template_type: r.template_type.clone(), + template_title: r.template_title.clone(), + maybe_creator_user_token: r.maybe_creator_user_token.clone(), + maybe_creator_username: r.maybe_creator_username.clone(), + maybe_creator_display_name: r.maybe_creator_display_name.clone(), + maybe_creator_result_id: r.maybe_creator_result_id.map(|v| v as u64).clone(), + file_size_bytes: if r.file_size_bytes > 0 { r.file_size_bytes as u32 } else { 0 }, + frame_width: if r.frame_width > 0 { r.frame_width as u32 } else { 0 }, + frame_height: if r.frame_height > 0 { r.frame_height as u32 } else { 0 }, + duration_millis: if r.duration_millis > 0 { r.duration_millis as u32 } else { 0 }, + visibility: RecordVisibility::from_str(&r.creator_set_visibility).unwrap_or(RecordVisibility::Public), + created_at: r.created_at, + updated_at: r.updated_at, + } + }) + .collect::>(); + + Ok(W2lInferenceListPage { + inference_records: inference_results, + sort_ascending: self.sort_ascending, + first_id, + last_id, + }) + } + + /// Perform the query based on the set predicates. + /// The records contain IDs, which we should not expose to the frontend. + pub async fn perform_internal_query( + &self, + mysql_pool: &MySqlPool + ) -> AnyhowResult> { + + let query = self.build_query_string(); + let mut query = sqlx::query_as::<_, RawInternalTtsRecord>(&query); + + // NB: The following bindings must match the order of the query builder !! + + if let Some(offset) = self.offset { + query = query.bind(offset); + } + + if let Some(username) = self.scope_creator_username.as_deref() { + query = query.bind(username); + } + + query = query.bind(self.limit); + + let mut results = query.fetch_all(mysql_pool) + .await?; + + if self.cursor_is_reversed { + results.reverse() + } + + Ok(results) + } + + pub fn build_query_string(&self) -> String { + // TODO: I haven't figured out how to get field name disambiguation and type coercion working here. + // (1) w2l_results.creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Query error: no column found for name: creator_set_visibility + // (2) creator_set_visibility `creator_set_visibility: crate::database::enums::record_visibility::RecordVisibility`, + // Column 'creator_set_visibility' in field list is ambiguous + + // TODO/NB: Unfortunately SQLx can't statically typecheck this query + let mut query = r#" +SELECT + w2l_results.id as w2l_result_id, + w2l_results.token as w2l_result_token, + w2l_results.maybe_tts_inference_result_token, + + w2l_templates.token as maybe_w2l_template_token, + w2l_templates.template_type, + w2l_templates.title as template_title, + + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + + w2l_results.maybe_creator_synthetic_id as maybe_creator_result_id, + + w2l_results.file_size_bytes, + w2l_results.frame_width, + w2l_results.frame_height, + w2l_results.duration_millis, + + w2l_results.creator_set_visibility, + + w2l_results.created_at, + w2l_results.updated_at + +FROM w2l_results +LEFT OUTER JOIN w2l_templates + ON w2l_results.maybe_w2l_template_token = w2l_templates.token +LEFT OUTER JOIN users + ON w2l_results.maybe_creator_user_token = users.token + "#.to_string(); + + query.push_str(&self.build_predicates()); + query + } + + pub fn build_predicates(&self) -> String { + // NB: Reverse cursors require us to invert the sort direction. + let mut sort_ascending = self.sort_ascending; + + let mut first_predicate_added = false; + + let mut query = "".to_string(); + + if let Some(offset) = self.offset { + if !first_predicate_added { + query.push_str(" WHERE"); + first_predicate_added = true; + } else { + query.push_str(" AND"); + } + + if sort_ascending { + if self.cursor_is_reversed { + // NB: We're searching backwards. + query.push_str(" w2l_results.id < ?"); + sort_ascending = !sort_ascending; + } else { + query.push_str(" w2l_results.id > ?"); + } + } else { + if self.cursor_is_reversed { + // NB: We're searching backwards. + query.push_str(" w2l_results.id > ?"); + sort_ascending = !sort_ascending; + } else { + query.push_str(" w2l_results.id < ?"); + } + } + } + + if let Some(username) = self.scope_creator_username.as_deref() { + if !first_predicate_added { + query.push_str(" WHERE users.username = ?"); + first_predicate_added = true; + } else { + query.push_str(" AND users.username = ?"); + } + } + + if !self.include_templates_not_approved_for_public_listing { + if !first_predicate_added { + query.push_str(" WHERE w2l_templates.is_public_listing_approved IS TRUE"); + first_predicate_added = true; + } else { + query.push_str(" AND w2l_templates.is_public_listing_approved IS TRUE"); + } + } + + if !self.include_user_hidden { + if !first_predicate_added { + query.push_str(" WHERE w2l_results.creator_set_visibility = 'public'"); + first_predicate_added = true; + } else { + query.push_str(" AND w2l_results.creator_set_visibility = 'public'"); + } + } + + if !self.include_mod_deleted_results { + if !first_predicate_added { + query.push_str(" WHERE w2l_results.mod_deleted_at IS NULL"); + first_predicate_added = true; + } else { + query.push_str(" AND w2l_results.mod_deleted_at IS NULL"); + } + } + + if !self.include_user_deleted_results { + if !first_predicate_added { + query.push_str(" WHERE w2l_results.user_deleted_at IS NULL"); + first_predicate_added = true; + } else { + query.push_str(" AND w2l_results.user_deleted_at IS NULL"); + } + } + + if sort_ascending { + query.push_str(" ORDER BY w2l_results.id ASC"); + } else { + query.push_str(" ORDER BY w2l_results.id DESC"); + } + + query.push_str(" LIMIT ?"); + + query + } +} + +#[derive(sqlx::FromRow)] +pub struct RawInternalTtsRecord { + pub w2l_result_id: i64, + pub w2l_result_token: String, + + pub maybe_w2l_template_token: Option, + pub maybe_tts_inference_result_token: Option, + + pub template_type: Option, + pub template_title: Option, // from field `w2l_templates.title` + + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + + pub maybe_creator_result_id: Option, + + pub file_size_bytes: i32, + pub frame_width: i32, + pub frame_height: i32, + pub duration_millis: i32, + + pub creator_set_visibility: String, + + //pub template_is_mod_approved: i8, // needs convert + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[cfg(test)] +mod tests { + use crate::queries::w2l::w2l_results::list_w2l_inference_results_query_builder::ListW2lResultsQueryBuilder; + + #[test] + fn predicates_without_scoping() { + let query_builder = ListW2lResultsQueryBuilder::new(); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_scoped_to_user() { + let query_builder = ListW2lResultsQueryBuilder::new() + .scope_creator_username(Some("echelon")); + + assert_eq!(&query_builder.build_predicates(), + " WHERE users.username = ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_user_hidden() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_user_hidden(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_mod_deleted() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_mod_deleted_results(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_user_deleted() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_user_deleted_results(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_mod_deleted_and_user_deleted() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_mod_deleted_results(true) + .include_user_deleted_results(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_unapproved() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_templates_not_approved_for_public_listing(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_unapproved_and_mod_deleted() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_templates_not_approved_for_public_listing(true) + .include_mod_deleted_results(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_unapproved_and_user_deleted() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_templates_not_approved_for_public_listing(true) + .include_user_deleted_results(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_including_mod_deleted_user_deleted_and_unapproved() { + let query_builder = ListW2lResultsQueryBuilder::new() + .include_mod_deleted_results(true) + .include_user_deleted_results(true) + .include_templates_not_approved_for_public_listing(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.creator_set_visibility = 'public' \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_sort_ascending() { + let query_builder = ListW2lResultsQueryBuilder::new() + .sort_ascending(true); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id ASC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset() { + let query_builder = ListW2lResultsQueryBuilder::new() + .offset(Some(100)); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.id < ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset_and_sort_ascending() { + let query_builder = ListW2lResultsQueryBuilder::new() + .sort_ascending(true) + .offset(Some(100)); + + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.id > ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id ASC \ + LIMIT ?"); + } + + #[test] + fn predicates_limit() { + let query_builder = ListW2lResultsQueryBuilder::new() + .limit(15); + + // NB: Does not change the query itself! Just the downstream binding. + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_cursor_is_reversed_without_cursor() { + let query_builder = ListW2lResultsQueryBuilder::new() + .cursor_is_reversed(true); + + // NB: Without a cursor, nothing happens. + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset_cursor_is_reversed() { + let query_builder = ListW2lResultsQueryBuilder::new() + .offset(Some(100)) + .cursor_is_reversed(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.id > ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id ASC \ + LIMIT ?"); + } + + #[test] + fn predicates_offset_cursor_is_reversed_sort_ascending() { + let query_builder = ListW2lResultsQueryBuilder::new() + .offset(Some(100)) + .cursor_is_reversed(true) + .sort_ascending(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.id < ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_limit_scope_user_offset_cursor_is_reversed_sort_ascending() { + let query_builder = ListW2lResultsQueryBuilder::new() + .limit(1000) + .scope_creator_username(Some("pikachu")) + .offset(Some(100)) + .cursor_is_reversed(true) + .sort_ascending(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.id < ? \ + AND users.username = ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.creator_set_visibility = 'public' \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } + + #[test] + fn predicates_limit_scope_user_show_hidden_offset_cursor_is_reversed_sort_ascending() { + let query_builder = ListW2lResultsQueryBuilder::new() + .limit(1000) + .include_user_hidden(true) + .scope_creator_username(Some("pikachu")) + .offset(Some(100)) + .cursor_is_reversed(true) + .sort_ascending(true); + + // NB: This will change the sort order and greater/less than direction! + assert_eq!(&query_builder.build_predicates(), + " WHERE w2l_results.id < ? \ + AND users.username = ? \ + AND w2l_templates.is_public_listing_approved IS TRUE \ + AND w2l_results.mod_deleted_at IS NULL \ + AND w2l_results.user_deleted_at IS NULL \ + ORDER BY w2l_results.id DESC \ + LIMIT ?"); + } +} diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_results/mod.rs b/crates/lib/database_queries/src/queries/w2l/w2l_results/mod.rs new file mode 100644 index 0000000000..8189a79f96 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_results/mod.rs @@ -0,0 +1,3 @@ +pub mod delete_w2l_result_various_scopes; +pub mod list_w2l_inference_results_query_builder; +pub mod query_w2l_result; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_results/query_w2l_result.rs b/crates/lib/database_queries/src/queries/w2l/w2l_results/query_w2l_result.rs new file mode 100644 index 0000000000..2d8339d6ce --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_results/query_w2l_result.rs @@ -0,0 +1,294 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::record_visibility::RecordVisibility; +use crate::helpers::boolean_converters::nullable_i8_to_bool; +use log::{info, warn, log}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; + +#[derive(Serialize)] +pub struct W2lResultRecordForResponse { + pub w2l_result_token: String, + pub maybe_w2l_template_token: Option, + pub maybe_tts_inference_result_token: Option, + + pub public_bucket_video_path: String, + + pub template_type: Option, + pub template_title: Option, + + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + pub maybe_creator_gravatar_hash: Option, + + pub maybe_template_creator_user_token: Option, + pub maybe_template_creator_username: Option, + pub maybe_template_creator_display_name: Option, + pub maybe_template_creator_gravatar_hash: Option, + + pub creator_set_visibility: RecordVisibility, + + pub file_size_bytes: u32, + pub frame_width: u32, + pub frame_height: u32, + pub duration_millis: u32, + + //pub template_is_mod_approved: bool, // converted + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub maybe_moderator_fields: Option, +} + +/// "Moderator-only fields" that we wouldn't want to expose to ordinary users. +/// It's the web endpoint controller's responsibility to clear these for non-mods. +#[derive(Serialize)] +pub struct W2lResultModeratorFields { + pub template_creator_is_banned: bool, + pub result_creator_is_banned_if_user: bool, + pub result_creator_ip_address: String, + pub result_creator_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +#[derive(Serialize)] +pub struct W2lResultRecordRaw { + pub w2l_result_token: String, // from field `w2l_results.token` + + pub maybe_w2l_template_token: Option, + pub maybe_tts_inference_result_token: Option, + + pub public_bucket_video_path: String, + + pub template_type: Option, + pub template_title: Option, // from field `w2l_templates.title` + + pub maybe_creator_is_banned: Option, + pub maybe_creator_user_token: Option, + pub maybe_creator_username: Option, + pub maybe_creator_display_name: Option, + pub maybe_creator_gravatar_hash: Option, + + pub maybe_template_creator_is_banned: Option, + pub maybe_template_creator_user_token: Option, + pub maybe_template_creator_username: Option, + pub maybe_template_creator_display_name: Option, + pub maybe_template_creator_gravatar_hash: Option, + + pub creator_set_visibility: String, + + pub file_size_bytes: i32, + pub frame_width: i32, + pub frame_height: i32, + pub duration_millis: i32, + + //pub template_is_mod_approved: i8, // needs convert + //pub maybe_mod_user_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub creator_ip_address: String, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +pub async fn select_w2l_result_by_token( + w2l_result_token: &str, + can_see_deleted: bool, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + let maybe_record = if can_see_deleted { + select_including_deleted(w2l_result_token, mysql_pool).await + } else { + select_without_deleted(w2l_result_token, mysql_pool).await + }; + + let ir : W2lResultRecordRaw = match maybe_record { + Ok(inference_result) => inference_result, + Err(ref err) => { + match err { + sqlx::Error::RowNotFound => { + warn!("w2l result not found: {:?}", &err); + return Ok(None); + }, + _ => { + warn!("w2l result query error: {:?}", &err); + return Err(anyhow!("database error")); + } + } + } + }; + + let ir_for_response = W2lResultRecordForResponse { + w2l_result_token: ir.w2l_result_token, + maybe_w2l_template_token: ir.maybe_w2l_template_token, + maybe_tts_inference_result_token: ir.maybe_tts_inference_result_token, + + public_bucket_video_path: ir.public_bucket_video_path, + + template_type: ir.template_type, + template_title: ir.template_title, + + maybe_creator_user_token: ir.maybe_creator_user_token, + maybe_creator_username: ir.maybe_creator_username, + maybe_creator_display_name: ir.maybe_creator_display_name, + maybe_creator_gravatar_hash: ir.maybe_creator_gravatar_hash, + + maybe_template_creator_user_token: ir.maybe_template_creator_user_token, + maybe_template_creator_username: ir.maybe_template_creator_username, + maybe_template_creator_display_name: ir.maybe_template_creator_display_name, + maybe_template_creator_gravatar_hash: ir.maybe_template_creator_gravatar_hash, + + // NB: Fail open/public since we're already looking at it + creator_set_visibility: RecordVisibility::from_str(&ir.creator_set_visibility) + .unwrap_or(RecordVisibility::Public), + + //template_is_mod_approved: if ir.template_is_mod_approved == 0 { false } else { true }, + + file_size_bytes: if ir.file_size_bytes > 0 { ir.file_size_bytes as u32 } else { 0 }, + frame_width: if ir.frame_width > 0 { ir.frame_width as u32 } else { 0 }, + frame_height: if ir.frame_height > 0 { ir.frame_height as u32 } else { 0 }, + duration_millis: if ir.duration_millis > 0 { ir.duration_millis as u32 } else { 0 }, + + created_at: ir.created_at, + updated_at: ir.updated_at, + + maybe_moderator_fields: Some(W2lResultModeratorFields { + template_creator_is_banned: + nullable_i8_to_bool(ir.maybe_template_creator_is_banned, false), + result_creator_is_banned_if_user: + nullable_i8_to_bool(ir.maybe_creator_is_banned, false), + result_creator_ip_address: ir.creator_ip_address, + result_creator_deleted_at: ir.user_deleted_at, + mod_deleted_at: ir.mod_deleted_at, + }), + }; + + Ok(Some(ir_for_response)) +} + +async fn select_including_deleted( + w2l_result_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + W2lResultRecordRaw, + r#" +SELECT + w2l_results.token as w2l_result_token, + w2l_results.maybe_tts_inference_result_token, + + w2l_results.public_bucket_video_path, + + w2l_templates.token as maybe_w2l_template_token, + w2l_templates.template_type, + w2l_templates.title as template_title, + + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + users.email_gravatar_hash as maybe_creator_gravatar_hash, + users.is_banned as maybe_creator_is_banned, + + template_users.token as maybe_template_creator_user_token, + template_users.username as maybe_template_creator_username, + template_users.display_name as maybe_template_creator_display_name, + template_users.email_gravatar_hash as maybe_template_creator_gravatar_hash, + template_users.is_banned as maybe_template_creator_is_banned, + + w2l_results.creator_set_visibility, + + w2l_results.file_size_bytes, + w2l_results.frame_width, + w2l_results.frame_height, + w2l_results.duration_millis, + w2l_results.created_at, + w2l_results.updated_at, + + w2l_results.creator_ip_address, + w2l_results.user_deleted_at, + w2l_results.mod_deleted_at + +FROM w2l_results +LEFT OUTER JOIN w2l_templates + ON w2l_results.maybe_w2l_template_token = w2l_templates.token +LEFT OUTER JOIN users + ON w2l_results.maybe_creator_user_token = users.token +LEFT OUTER JOIN users as template_users + ON w2l_templates.creator_user_token = template_users.token +WHERE + w2l_results.token = ? + "#, + w2l_result_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} + +async fn select_without_deleted( + w2l_result_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + W2lResultRecordRaw, + r#" +SELECT + w2l_results.token as w2l_result_token, + w2l_results.maybe_tts_inference_result_token, + + w2l_results.public_bucket_video_path, + + w2l_templates.token as maybe_w2l_template_token, + w2l_templates.template_type, + w2l_templates.title as template_title, + + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + users.email_gravatar_hash as maybe_creator_gravatar_hash, + users.is_banned as maybe_creator_is_banned, + + template_users.token as maybe_template_creator_user_token, + template_users.username as maybe_template_creator_username, + template_users.display_name as maybe_template_creator_display_name, + template_users.email_gravatar_hash as maybe_template_creator_gravatar_hash, + template_users.is_banned as maybe_template_creator_is_banned, + + w2l_results.creator_set_visibility, + + w2l_results.file_size_bytes, + w2l_results.frame_width, + w2l_results.frame_height, + w2l_results.duration_millis, + w2l_results.created_at, + w2l_results.updated_at, + + w2l_results.creator_ip_address, + w2l_results.user_deleted_at, + w2l_results.mod_deleted_at + +FROM w2l_results +LEFT OUTER JOIN w2l_templates + ON w2l_results.maybe_w2l_template_token = w2l_templates.token +LEFT OUTER JOIN users + ON w2l_results.maybe_creator_user_token = users.token +LEFT OUTER JOIN users as template_users + ON w2l_templates.creator_user_token = template_users.token +WHERE + w2l_results.token = ? + AND w2l_results.user_deleted_at IS NULL + AND w2l_results.mod_deleted_at IS NULL + "#, + w2l_result_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_templates/delete_w2l_template_various_scopes.rs b/crates/lib/database_queries/src/queries/w2l/w2l_templates/delete_w2l_template_various_scopes.rs new file mode 100644 index 0000000000..308684434d --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_templates/delete_w2l_template_various_scopes.rs @@ -0,0 +1,94 @@ +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; + +pub async fn delete_w2l_template_as_user( + template_token: &str, + creator_ip_address: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_templates +SET + creator_ip_address_last_update = ?, + user_deleted_at = CURRENT_TIMESTAMP +WHERE + token = ? +LIMIT 1 + "#, + creator_ip_address, + template_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn delete_w2l_template_as_mod( + template_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_templates +SET + maybe_mod_user_token = ?, + mod_deleted_at = CURRENT_TIMESTAMP +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + template_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_w2l_template_as_user( + template_token: &str, + creator_ip_address: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_templates +SET + creator_ip_address_last_update = ?, + user_deleted_at = NULL +WHERE + token = ? +LIMIT 1 + "#, + creator_ip_address, + template_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} + +pub async fn undelete_w2l_template_as_mod( + template_token: &str, + mod_user_token: &str, + mysql_pool: &MySqlPool +) -> AnyhowResult<()> { + let _r = sqlx::query!( + r#" +UPDATE w2l_templates +SET + maybe_mod_user_token = ?, + mod_deleted_at = NULL +WHERE + token = ? +LIMIT 1 + "#, + mod_user_token, + template_token, + ) + .execute(mysql_pool) + .await?; + Ok(()) +} diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_templates/get_w2l_template.rs b/crates/lib/database_queries/src/queries/w2l/w2l_templates/get_w2l_template.rs new file mode 100644 index 0000000000..3c9df4abb3 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_templates/get_w2l_template.rs @@ -0,0 +1,233 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::record_visibility::RecordVisibility; +use crate::helpers::boolean_converters::{nullable_i8_to_optional_bool, i8_to_bool}; +use log::{info, warn, log}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// The moderator-only fields are good practice, though. + +#[derive(Serialize)] +pub struct W2lTemplateRecordForResponse { + pub template_token: String, + pub template_type: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + pub frame_width: u32, + pub frame_height: u32, + pub duration_millis: u32, + pub maybe_image_object_name: Option, + pub maybe_video_object_name: Option, + pub creator_set_visibility: RecordVisibility, + pub is_public_listing_approved: Option, + pub is_locked_from_use: bool, + pub is_locked_from_user_modification: bool, + pub created_at: DateTime, + pub updated_at: DateTime, + pub maybe_moderator_fields: Option, +} + +/// "Moderator-only fields" that we wouldn't want to expose to ordinary users. +/// It's the web endpoint controller's responsibility to clear these for non-mods. +#[derive(Serialize)] +pub struct W2lTemplateModeratorFields { + pub creator_is_banned: bool, + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +#[derive(Serialize)] +pub struct W2lTemplateRecordRaw { + pub template_token: String, + pub template_type: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + pub creator_is_banned: i8, + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + pub frame_width: i32, + pub frame_height: i32, + pub duration_millis: i32, + pub maybe_public_bucket_preview_image_object_name: Option, + pub maybe_public_bucket_preview_video_object_name: Option, + pub creator_set_visibility: String, + pub is_public_listing_approved: Option, + pub is_locked_from_use: i8, + pub is_locked_from_user_modification: i8, + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + pub created_at: DateTime, + pub updated_at: DateTime, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +pub async fn select_w2l_template_by_token( + w2l_template_token: &str, + can_see_deleted: bool, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + let maybe_record = if can_see_deleted { + select_including_deleted(w2l_template_token, mysql_pool).await + } else { + select_without_deleted(w2l_template_token, mysql_pool).await + }; + + let template : W2lTemplateRecordRaw = match maybe_record { + Ok(template) => template, + Err(ref err) => { + match err { + sqlx::Error::RowNotFound => { + warn!("w2l template not found: {:?}", &err); + return Ok(None); + }, + _ => { + warn!("w2l template query error: {:?}", &err); + return Err(anyhow!("database error")); + } + } + } + }; + + let template_for_response = W2lTemplateRecordForResponse { + template_token: template.template_token, + template_type: template.template_type, + creator_user_token: template.creator_user_token, + creator_username: template.creator_username, + creator_display_name: template.creator_display_name, + creator_gravatar_hash: template.creator_gravatar_hash, + title: template.title, + description_markdown: template.description_markdown, + description_rendered_html: template.description_rendered_html, + frame_width: if template.frame_width > 0 { template.frame_width as u32 } else { 0 }, + frame_height: if template.frame_height > 0 { template.frame_height as u32 } else { 0 }, + duration_millis: if template.duration_millis > 0 { template.duration_millis as u32 } else { 0 }, + maybe_image_object_name: template.maybe_public_bucket_preview_image_object_name, + maybe_video_object_name: template.maybe_public_bucket_preview_video_object_name, + // NB: Fail open/public with creator_set_visibility since we're already looking at it + creator_set_visibility: RecordVisibility::from_str(&template.creator_set_visibility) + .unwrap_or(RecordVisibility::Public), + is_public_listing_approved: nullable_i8_to_optional_bool(template.is_public_listing_approved), + is_locked_from_use: i8_to_bool(template.is_locked_from_use), + is_locked_from_user_modification: i8_to_bool(template.is_locked_from_user_modification), + created_at: template.created_at, + updated_at: template.updated_at, + maybe_moderator_fields: Some(W2lTemplateModeratorFields { + creator_is_banned: i8_to_bool(template.creator_is_banned), + creator_ip_address_creation: template.creator_ip_address_creation, + creator_ip_address_last_update: template.creator_ip_address_last_update, + user_deleted_at: template.user_deleted_at, + mod_deleted_at: template.mod_deleted_at, + }), + }; + + Ok(Some(template_for_response)) +} + +async fn select_including_deleted( + w2l_template_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + W2lTemplateRecordRaw, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.is_banned as creator_is_banned, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + w2l.title, + w2l.description_markdown, + w2l.description_rendered_html, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.creator_set_visibility, + w2l.is_public_listing_approved, + w2l.is_locked_from_use, + w2l.is_locked_from_user_modification, + w2l.creator_ip_address_creation, + w2l.creator_ip_address_last_update, + w2l.created_at, + w2l.updated_at, + w2l.user_deleted_at, + w2l.mod_deleted_at +FROM w2l_templates as w2l +JOIN users + ON users.token = w2l.creator_user_token +WHERE w2l.token = ? + "#, + w2l_template_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} + +async fn select_without_deleted( + w2l_template_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + W2lTemplateRecordRaw, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.is_banned as creator_is_banned, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + w2l.title, + w2l.description_markdown, + w2l.description_rendered_html, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.creator_set_visibility, + w2l.is_public_listing_approved, + w2l.is_locked_from_use, + w2l.is_locked_from_user_modification, + w2l.creator_ip_address_creation, + w2l.creator_ip_address_last_update, + w2l.created_at, + w2l.updated_at, + w2l.user_deleted_at, + w2l.mod_deleted_at +FROM w2l_templates as w2l +JOIN users + ON users.token = w2l.creator_user_token +WHERE + w2l.token = ? + AND w2l.user_deleted_at IS NULL + AND w2l.mod_deleted_at IS NULL + "#, + w2l_template_token + ) + .fetch_one(mysql_pool) + .await // TODO: This will return error if it doesn't exist +} diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_templates/list_w2l_templates.rs b/crates/lib/database_queries/src/queries/w2l/w2l_templates/list_w2l_templates.rs new file mode 100644 index 0000000000..0ce308eba7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_templates/list_w2l_templates.rs @@ -0,0 +1,249 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::nullable_i8_to_optional_bool; +use log::{warn, info}; +use sqlx::MySqlPool; + +// FIXME: This is the old style of query scoping and shouldn't be copied. + +#[derive(Serialize)] +pub struct W2lTemplateRecordForList { + pub template_token: String, + pub template_type: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub title: String, + pub frame_width: u32, + pub frame_height: u32, + pub duration_millis: u32, + pub maybe_image_object_name: Option, + pub maybe_video_object_name: Option, + pub is_public_listing_approved: Option, // converted + pub created_at: DateTime, + pub updated_at: DateTime, +} + +struct RawW2lTemplateRecordForList { + pub template_token: String, + pub template_type: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub title: String, + pub frame_width: i32, + pub frame_height: i32, + pub duration_millis: i32, + pub maybe_public_bucket_preview_image_object_name: Option, + pub maybe_public_bucket_preview_video_object_name: Option, + pub is_public_listing_approved: Option, // NB: needs conversion + pub created_at: DateTime, + pub updated_at: DateTime, +} + +pub async fn list_w2l_templates( + mysql_pool: &MySqlPool, + scope_creator_username: Option<&str>, + require_mod_approved: bool +) -> AnyhowResult> { + + let maybe_templates = match scope_creator_username { + Some(username) => { + list_w2l_templates_creator_scoped(mysql_pool, username, require_mod_approved) + .await + }, + None => { + list_w2l_templates_for_all_creators(mysql_pool, require_mod_approved) + .await + }, + }; + + let templates : Vec = match maybe_templates { + Ok(templates) => { + info!("Template length: {}", templates.len()); + templates + }, + Err(err) => { + warn!("Error: {:?}", err); + + match err { + RowNotFound => { + return Ok(Vec::new()); + }, + _ => { + warn!("w2l template list query error: {:?}", err); + return Err(anyhow!("w2l template list query error")); + } + } + } + }; + + Ok(templates.into_iter() + .map(|template| { + W2lTemplateRecordForList { + template_token: template.template_token, + template_type: template.template_type, + creator_user_token: template.creator_user_token, + creator_username: template.creator_username, + creator_display_name: template.creator_display_name, + title: template.title, + frame_width: if template.frame_width > 0 { template.frame_width as u32 } else { 0 }, + frame_height: if template.frame_height > 0 { template.frame_height as u32 } else { 0 }, + duration_millis: if template.duration_millis > 0 { template.duration_millis as u32 } else { 0 }, + maybe_image_object_name: template.maybe_public_bucket_preview_image_object_name, + maybe_video_object_name: template.maybe_public_bucket_preview_video_object_name, + is_public_listing_approved: nullable_i8_to_optional_bool(template.is_public_listing_approved), + created_at: template.created_at, + updated_at: template.updated_at, + } + }) + .collect::>()) +} + +async fn list_w2l_templates_for_all_creators( + mysql_pool: &MySqlPool, + require_mod_approved: bool +) -> AnyhowResult> { + // TODO: There has to be a better way. + // Sqlx doesn't like anything except string literals. + let maybe_templates = if require_mod_approved { + info!("listing w2l templates for everyone; mod-approved only"); + sqlx::query_as!( + RawW2lTemplateRecordForList, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + w2l.title, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.is_public_listing_approved, + w2l.created_at, + w2l.updated_at +FROM w2l_templates as w2l +JOIN users + ON users.token = w2l.creator_user_token +WHERE + w2l.is_public_listing_approved IS TRUE + AND w2l.user_deleted_at IS NULL + AND w2l.mod_deleted_at IS NULL + "#) + .fetch_all(mysql_pool) + .await? + } else { + info!("listing w2l templates for everyone; all"); + sqlx::query_as!( + RawW2lTemplateRecordForList, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + w2l.title, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.is_public_listing_approved, + w2l.created_at, + w2l.updated_at +FROM w2l_templates as w2l +JOIN users + ON users.token = w2l.creator_user_token +WHERE + w2l.user_deleted_at IS NULL + AND w2l.mod_deleted_at IS NULL + "#) + .fetch_all(mysql_pool) + .await? + }; + + Ok(maybe_templates) +} + +async fn list_w2l_templates_creator_scoped( + mysql_pool: &MySqlPool, + scope_creator_username: &str, + require_mod_approved: bool +) -> AnyhowResult> { + // TODO: There has to be a better way. + // Sqlx doesn't like anything except string literals. + let maybe_templates = if require_mod_approved { + info!("listing w2l templates for user; mod-approved only"); + sqlx::query_as!( + RawW2lTemplateRecordForList, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + w2l.title, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.is_public_listing_approved, + w2l.created_at, + w2l.updated_at +FROM w2l_templates as w2l +JOIN users +ON + users.token = w2l.creator_user_token +WHERE + users.username = ? + AND w2l.is_public_listing_approved IS TRUE + AND w2l.user_deleted_at IS NULL + AND w2l.mod_deleted_at IS NULL + "#, + scope_creator_username) + .fetch_all(mysql_pool) + .await? + } else { + info!("listing w2l templates for user; all"); + sqlx::query_as!( + RawW2lTemplateRecordForList, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + w2l.title, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.is_public_listing_approved, + w2l.created_at, + w2l.updated_at +FROM w2l_templates as w2l +JOIN users +ON + users.token = w2l.creator_user_token +WHERE + users.username = ? + AND w2l.user_deleted_at IS NULL + AND w2l.mod_deleted_at IS NULL + "#, + scope_creator_username) + .fetch_all(mysql_pool) + .await? + }; + + Ok(maybe_templates) +} diff --git a/crates/lib/database_queries/src/queries/w2l/w2l_templates/mod.rs b/crates/lib/database_queries/src/queries/w2l/w2l_templates/mod.rs new file mode 100644 index 0000000000..d6ae263916 --- /dev/null +++ b/crates/lib/database_queries/src/queries/w2l/w2l_templates/mod.rs @@ -0,0 +1,3 @@ +pub mod delete_w2l_template_various_scopes; +pub mod get_w2l_template; +pub mod list_w2l_templates; \ No newline at end of file diff --git a/crates/lib/database_queries/src/tokens.rs b/crates/lib/database_queries/src/tokens.rs new file mode 100644 index 0000000000..44df715317 --- /dev/null +++ b/crates/lib/database_queries/src/tokens.rs @@ -0,0 +1,87 @@ +use container_common::anyhow_result::AnyhowResult; +use container_common::token::random_prefix_crockford_token::random_prefix_crockford_token; + +/// Central place to define all the token types. +// TODO: Make tokens strongly typed. +pub struct Tokens {} + +impl Tokens { + + // ========== Tokens with high visibility (make these nice) ========== + + pub fn new_tts_model() -> AnyhowResult { + random_prefix_crockford_token("TM:", 15, false) + } + + pub fn new_tts_result() -> AnyhowResult { + random_prefix_crockford_token("TR:", 32, false) + } + + pub fn new_w2l_template() -> AnyhowResult { + random_prefix_crockford_token("WT:", 32, false) + } + + pub fn new_w2l_result() -> AnyhowResult { + random_prefix_crockford_token("WR:", 32, false) + } + + pub fn new_api_token() -> AnyhowResult { + random_prefix_crockford_token("API:", 32, true) + } + + // ========== Tokens not typically visible to users ========== + + pub fn new_user() -> AnyhowResult { + random_prefix_crockford_token("U:", 15, true) + } + + pub fn new_session() -> AnyhowResult { + random_prefix_crockford_token("SESSION:", 32, false) + } + + pub fn new_tts_model_upload_job() -> AnyhowResult { + random_prefix_crockford_token("JTUP:", 32, false) + } + + pub fn new_tts_inference_job() -> AnyhowResult { + random_prefix_crockford_token("JTINF:", 32, false) + } + + pub fn new_w2l_template_upload_job() -> AnyhowResult { + random_prefix_crockford_token("JWUP:", 32, false) + } + + pub fn new_w2l_inference_job() -> AnyhowResult { + random_prefix_crockford_token("JWINF:", 32, false) + } + + pub fn new_firehose_event() -> AnyhowResult { + random_prefix_crockford_token("EV:", 32, false) + } + + pub fn new_voice_clone_request_token() -> AnyhowResult { + random_prefix_crockford_token("VCR:", 32, false) + } + + pub fn new_category() -> AnyhowResult { + random_prefix_crockford_token("CAT:", 15, false) + } + + pub fn new_api_token_internal_token() -> AnyhowResult { + random_prefix_crockford_token("INT_API:", 32, false) + } + + // ========== [Twitch] Tokens not typically visible to users ========== + + pub fn new_twitch_oauth_internal_token() -> AnyhowResult { + random_prefix_crockford_token("TOI:", 32, false) + } + + pub fn new_twitch_oauth_grouping_token() -> AnyhowResult { + random_prefix_crockford_token("OG:", 32, false) + } + + pub fn new_twitch_event_rule() -> AnyhowResult { + random_prefix_crockford_token("TER:", 32, false) + } +} diff --git a/crates/lib/http_server_common/Cargo.toml b/crates/lib/http_server_common/Cargo.toml new file mode 100644 index 0000000000..b351f38410 --- /dev/null +++ b/crates/lib/http_server_common/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "http_server_common" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "http_server_common" +path = "src/lib.rs" + +[dependencies] + +# Internal +container_common = { path = "../container_common" } + +# External +anyhow = "1.0.40" +base64 = "0.13.0" +chrono = { version = "0.4.19", features = ["serde"] } +dotenv = "0.15.0" +log = "0.4.14" + +# Web +actix-cors = "0.6.0-beta.3" +actix-http = "3.0.0-beta.11" +actix-web = "4.0.0-beta.10" + +# Language +language-tags = "0.3.2" # Parsing single tags richly +#accept-language = "2.0.0" # Parsing many tags poorly + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" +toml = "0.5.8" diff --git a/crates/lib/http_server_common/src/cors.rs b/crates/lib/http_server_common/src/cors.rs new file mode 100644 index 0000000000..f072f8857f --- /dev/null +++ b/crates/lib/http_server_common/src/cors.rs @@ -0,0 +1,75 @@ +use actix_cors::Cors; +use actix_http::http; + +// TODO: Take an environment. +/// Return cors config for FakeYou / Vocodes / OBS / local development +pub fn build_common_cors_config() -> Cors { + Cors::default() + // Local Development (Localhost) + .allowed_origin("http://localhost:12345") + .allowed_origin("http://localhost:3000") + .allowed_origin("http://localhost:54321") + .allowed_origin("http://localhost:5555") + .allowed_origin("http://localhost:7000") + .allowed_origin("http://localhost:7001") + .allowed_origin("http://localhost:7002") + .allowed_origin("http://localhost:7003") + .allowed_origin("http://localhost:7004") + .allowed_origin("http://localhost:7005") + .allowed_origin("http://localhost:7006") + .allowed_origin("http://localhost:7007") + .allowed_origin("http://localhost:7008") + .allowed_origin("http://localhost:7009") + .allowed_origin("http://localhost:8000") + .allowed_origin("http://localhost:8080") + // Local Development (JungleHorse) + .allowed_origin("http://api.jungle.horse") + .allowed_origin("http://jungle.horse") + .allowed_origin("http://jungle.horse:12345") + .allowed_origin("http://jungle.horse:7000") + .allowed_origin("http://obs.jungle.horse") + .allowed_origin("http://ws.jungle.horse") + .allowed_origin("https://api.jungle.horse") + .allowed_origin("https://jungle.horse") + .allowed_origin("https://obs.jungle.horse") + .allowed_origin("https://ws.jungle.horse") + // FakeYou + .allowed_origin("http://api.fakeyou.com") + .allowed_origin("http://fakeyou.com") + .allowed_origin("https://api.fakeyou.com") + .allowed_origin("https://fakeyou.com") + // Storyteller + .allowed_origin("http://api.storyteller.io") + .allowed_origin("http://create.storyteller.io") + .allowed_origin("http://obs.storyteller.io") + .allowed_origin("http://storyteller.io") + .allowed_origin("http://ws.storyteller.io") + .allowed_origin("https://api.storyteller.io") + .allowed_origin("https://create.storyteller.io") + .allowed_origin("https://obs.storyteller.io") + .allowed_origin("https://storyteller.io") + .allowed_origin("https://ws.storyteller.io") + // Trumped (legacy) + .allowed_origin("http://trumped.com") + .allowed_origin("https://trumped.com") + // Vocodes (legacy) + .allowed_origin("http://api.vo.codes") + .allowed_origin("http://mumble.stream") + .allowed_origin("http://vo.codes") + .allowed_origin("http://vocodes.com") + .allowed_origin("https://api.vo.codes") + .allowed_origin("https://mumble.stream") + .allowed_origin("https://vo.codes") + .allowed_origin("https://vocodes.com") + // Remaining setup + .allowed_methods(vec!["GET", "POST", "OPTIONS", "DELETE"]) + .supports_credentials() + .allowed_headers(vec![ + http::header::ACCEPT, + http::header::ACCESS_CONTROL_ALLOW_ORIGIN, // Tabulator Ajax + http::header::CONTENT_TYPE, + http::header::ACCESS_CONTROL_ALLOW_CREDENTIALS, // https://stackoverflow.com/a/46412839 + http::header::HeaderName::from_static("x-requested-with") // Tabulator Ajax sends + ]) + .max_age(3600) +} \ No newline at end of file diff --git a/crates/lib/http_server_common/src/endpoints/default_route_404.rs b/crates/lib/http_server_common/src/endpoints/default_route_404.rs new file mode 100644 index 0000000000..117ed03336 --- /dev/null +++ b/crates/lib/http_server_common/src/endpoints/default_route_404.rs @@ -0,0 +1,9 @@ +use actix_web::{HttpResponse, Responder}; +use log::warn; + +pub async fn default_route_404() -> impl Responder { + warn!("404 not found"); + HttpResponse::NotFound() + .content_type("text/html; charset=utf-8") + .body("

not found

") +} diff --git a/crates/lib/http_server_common/src/endpoints/mod.rs b/crates/lib/http_server_common/src/endpoints/mod.rs new file mode 100644 index 0000000000..9ae4ad6cf5 --- /dev/null +++ b/crates/lib/http_server_common/src/endpoints/mod.rs @@ -0,0 +1,2 @@ +pub mod default_route_404; +pub mod root_index; diff --git a/crates/lib/http_server_common/src/endpoints/root_index.rs b/crates/lib/http_server_common/src/endpoints/root_index.rs new file mode 100644 index 0000000000..ba745827f5 --- /dev/null +++ b/crates/lib/http_server_common/src/endpoints/root_index.rs @@ -0,0 +1,13 @@ +use actix_http::http::StatusCode; +use actix_web::HttpResponse; + + + +use log::debug; + +pub async fn get_root_index() -> HttpResponse { + debug!("GET /"); // NB: Google load balancer hits this a lot, and it spams. + HttpResponse::build(StatusCode::OK) + .content_type("text/html; charset=utf-8") + .body("

hello!

Are you looking for an API? Join our Discord!

Maybe you want to work with us? We can pay! Get in touch!

") +} diff --git a/crates/lib/http_server_common/src/error/common_server_error.rs b/crates/lib/http_server_common/src/error/common_server_error.rs new file mode 100644 index 0000000000..875d6a57de --- /dev/null +++ b/crates/lib/http_server_common/src/error/common_server_error.rs @@ -0,0 +1,49 @@ +use actix_web::HttpResponse; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use std::fmt; +use crate::error::simple_json_error_response::simple_json_error_response; + +/// Common server errors that should handle 90% of cases. +#[derive(Debug, Serialize)] +pub enum CommonServerError { + BadInput, + BadInputWithReason(String), + NotAuthorized, + NotFound, + ServerError, + ServerErrorWithReason(String), +} + +impl ResponseError for CommonServerError { + fn status_code(&self) -> StatusCode { + match *self { + Self::BadInput => StatusCode::BAD_REQUEST, + Self::BadInputWithReason(_) => StatusCode::BAD_REQUEST, + Self::NotAuthorized => StatusCode::UNAUTHORIZED, + Self::NotFound => StatusCode::NOT_FOUND, + Self::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + Self::ServerErrorWithReason(_) => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + Self::BadInput => "bad input", + Self::BadInputWithReason(_) => "bad input", + Self::NotAuthorized => "not authorized", + Self::NotFound => "not found", + Self::ServerError => "server error", + Self::ServerErrorWithReason(_) => "server error", + }; + + simple_json_error_response(error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for CommonServerError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} diff --git a/crates/lib/http_server_common/src/error/mod.rs b/crates/lib/http_server_common/src/error/mod.rs new file mode 100644 index 0000000000..9bdbd64051 --- /dev/null +++ b/crates/lib/http_server_common/src/error/mod.rs @@ -0,0 +1,2 @@ +pub mod common_server_error; +pub mod simple_json_error_response; \ No newline at end of file diff --git a/crates/lib/http_server_common/src/error/simple_json_error_response.rs b/crates/lib/http_server_common/src/error/simple_json_error_response.rs new file mode 100644 index 0000000000..b769e241ef --- /dev/null +++ b/crates/lib/http_server_common/src/error/simple_json_error_response.rs @@ -0,0 +1,55 @@ +use actix_web::http::{StatusCode, header}; +use actix_web::{HttpResponse, HttpResponseBuilder}; + +/// Convert a string and status code to a JSON error response. +pub fn simple_json_error_response( + error_reason: &str, + status_code: StatusCode, +) -> HttpResponse { + + let response = SimpleGenericJsonError { + success: false, + error_reason : error_reason.to_string(), + }; + + let body = match serde_json::to_string(&response) { + Ok(json) => json, + Err(_) => "{}".to_string(), + }; + + HttpResponseBuilder::new(status_code) + .content_type("application/json") + .body(body) +} + +#[derive(Serialize)] +struct SimpleGenericJsonError { + pub success: bool, + pub error_reason: String, +} + +#[cfg(test)] +mod tests { + use crate::error::simple_json_error_response::simple_json_error_response; + use actix_web::http::StatusCode; + use actix_http::body::AnyBody; + + #[test] + pub fn serialization() { + let response = simple_json_error_response( + "foo", + StatusCode::TOO_MANY_REQUESTS); + + assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS); + + let body = response.into_body(); + match body { + AnyBody::Bytes(bytes) => { + let bytes = bytes.to_vec(); + let body = String::from_utf8(bytes).unwrap(); + assert_eq!(body, "{\"success\":false,\"error_reason\":\"foo\"}"); + } + _ => panic!("this failed"), + } + } +} diff --git a/crates/lib/http_server_common/src/lib.rs b/crates/lib/http_server_common/src/lib.rs new file mode 100644 index 0000000000..3867702b0b --- /dev/null +++ b/crates/lib/http_server_common/src/lib.rs @@ -0,0 +1,17 @@ +//! +//! Common HTTP Server components +//! + +#[macro_use] extern crate serde_derive; + +// This is a little messy since we're supporting two versions of the Actix crate ecosystem. +// Long story short: one service requires an old Tokio runtime, and another needs the newer one. +// +// I don't think this will be too hard to maintain. AFAICT, the only major differences are which +// package "HttpResponseBuilder" is located in. + +pub mod cors; +pub mod endpoints; +pub mod error; +pub mod request; +pub mod response; diff --git a/crates/lib/http_server_common/src/request/get_request_header_optional.rs b/crates/lib/http_server_common/src/request/get_request_header_optional.rs new file mode 100644 index 0000000000..a3892150e3 --- /dev/null +++ b/crates/lib/http_server_common/src/request/get_request_header_optional.rs @@ -0,0 +1,15 @@ +use actix_web::HttpRequest; + +/// Pull out the named request header, if present. +pub fn get_request_header_optional(request: &HttpRequest, header_name: &str) -> Option { + // TODO: Clean up with transpose() once stable + let result = request.headers() + .get(header_name) + .map(|h| h.to_str()); + match result { + Some(Ok(header)) => Some(header.to_string()), + Some(Err(_)) => None, + None => None, + _ => None, + } +} \ No newline at end of file diff --git a/crates/lib/http_server_common/src/request/get_request_host.rs b/crates/lib/http_server_common/src/request/get_request_host.rs new file mode 100644 index 0000000000..159d69facc --- /dev/null +++ b/crates/lib/http_server_common/src/request/get_request_host.rs @@ -0,0 +1,8 @@ +use actix_web::HttpRequest; +use actix_web::http::header::HOST; +use crate::request::get_request_header_optional::get_request_header_optional; + +/// Get the request host +pub fn get_request_host(request: &HttpRequest) -> Option { + get_request_header_optional(request, HOST.as_str()) +} diff --git a/crates/lib/http_server_common/src/request/get_request_ip.rs b/crates/lib/http_server_common/src/request/get_request_ip.rs new file mode 100644 index 0000000000..77eb336a6d --- /dev/null +++ b/crates/lib/http_server_common/src/request/get_request_ip.rs @@ -0,0 +1,100 @@ +use actix_web::HttpRequest; +use actix_web::dev::ServiceRequest; +use log::{info}; +use actix_web::http::{HeaderMap, HeaderName}; +use std::str::FromStr; + +pub fn get_request_ip(request: &HttpRequest) -> String { + let headers = request.headers(); + let maybe_x_forwarded = get_ip_from_header(headers, "x-forwarded-for"); + let maybe_forwarded = get_ip_from_header(headers, "forwarded"); + + info!("(1) x-forwarded-for: {:?}, forwarded: {:?}", maybe_x_forwarded, maybe_forwarded); + + let maybe_ip = maybe_x_forwarded.or(maybe_forwarded); + + match maybe_ip { + Some(ip_address) => ip_address, + None => { + // If we're running without the upstream Rust proxy, we can grab 'x-forarded-for', which is + // populated by the DigitalOcean load balancer. + let ip_address_and_port = request.connection_info() + .remote_addr() + .unwrap_or("") + .to_string(); + let ip_address = ip_address_and_port.split(":") + .collect::>() + .get(0) + .copied() + .unwrap_or("") + .to_string(); + info!("Forwarded IP address (1): {}", &ip_address); + ip_address + }, + } +} + +// TODO: De-duplicate +pub fn get_service_request_ip(request: &ServiceRequest) -> String { + let headers = request.headers(); + let maybe_x_forwarded = get_ip_from_header(headers, "x-forwarded-for"); + let maybe_forwarded = get_ip_from_header(headers, "forwarded"); + + info!("(2) x-forwarded-for: {:?}, forwarded: {:?}", maybe_x_forwarded, maybe_forwarded); + + let maybe_ip = maybe_x_forwarded.or(maybe_forwarded); + + match maybe_ip { + Some(ip_address) => ip_address, + None => { + // If we're running without the upstream Rust proxy, we can grab 'x-forarded-for', which is + // populated by the DigitalOcean load balancer. + let ip_address_and_port = request.connection_info() + .remote_addr() + .unwrap_or("") + .to_string(); + let ip_address = ip_address_and_port.split(":") + .collect::>() + .get(0) + .copied() + .unwrap_or("") + .to_string(); + info!("(2) Forwarded IP address: {}", &ip_address); + ip_address + }, + } +} + +fn get_ip_from_header(headers: &HeaderMap, header_name: &str) -> Option { + if let Ok(header_name) = HeaderName::from_str(header_name) { + headers.get(&header_name) + .and_then(|value| value.to_str().ok()) + .and_then(|value| first_ip_from_list(value)) + } else { + None + } +} + +// NB: GKE packs in multiple IP addresses into the header: +// Example of header: x-forwarded-for: Some("136.55.189.34, 34.117.9.171") +fn first_ip_from_list(ip_list: &str) -> Option { + ip_list.split(",") + .into_iter() + .map(|ip| ip.trim()) + .find(|ip| !ip.is_empty()) + .map(|ip| ip.to_string()) +} + +#[cfg(test)] +mod tests { + use crate::request::get_request_ip::first_ip_from_list; + + #[test] + fn test_first_ip() { + assert_eq!(first_ip_from_list("1.2.3.4"), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list("1.2.3.4, 34.117.9.171"), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list(", 1.2.3.4"), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list("1.2.3.4, "), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list(""), None); + } +} \ No newline at end of file diff --git a/crates/lib/http_server_common/src/request/mod.rs b/crates/lib/http_server_common/src/request/mod.rs new file mode 100644 index 0000000000..c4a13b3a69 --- /dev/null +++ b/crates/lib/http_server_common/src/request/mod.rs @@ -0,0 +1,4 @@ +pub mod get_request_header_optional; +pub mod get_request_host; +pub mod get_request_ip; +pub mod parse_accept_language; \ No newline at end of file diff --git a/crates/lib/http_server_common/src/request/parse_accept_language.rs b/crates/lib/http_server_common/src/request/parse_accept_language.rs new file mode 100644 index 0000000000..982fb05e85 --- /dev/null +++ b/crates/lib/http_server_common/src/request/parse_accept_language.rs @@ -0,0 +1,74 @@ +use language_tags::{LanguageTag, ParseError}; +use log::warn; + +/// Parse out accept languages +/// Does not error so that the endpoint won't degrade +pub fn parse_accept_language(accept_languages_header: &str) -> Vec { + let unparsed_tags = accept_languages_header.split(",") + .into_iter() + .map(|tag| tag.trim().to_string()) + .collect::>(); + + let mut parsed_tags = Vec::new(); + + for tag in unparsed_tags.iter() { + match LanguageTag::parse(tag.as_ref()) { + Ok(t) => { + parsed_tags.push(t); + } + Err(e) => { + warn!("Error parsing language tag: {:?}", e); + } + } + } + + parsed_tags +} + +#[cfg(test)] +mod tests { + use crate::request::parse_accept_language::parse_accept_language; + + #[test] + fn test_empty() { + assert_eq!(parse_accept_language(""), Vec::new()); + } + + #[test] + fn test_garbage() { + assert_eq!(parse_accept_language("HJHKDJSHJKHSF"), Vec::new()); + assert_eq!(parse_accept_language("\n\n"), Vec::new()); + } + + #[test] + fn test_single_language() { + let list = parse_accept_language("en-US"); + let lang = list.get(0).unwrap(); + assert_eq!(lang.primary_language(), "en"); + + let list = parse_accept_language("en-GB"); + let lang = list.get(0).unwrap(); + assert_eq!(lang.primary_language(), "en"); + + let list = parse_accept_language("ja-JP"); + let lang = list.get(0).unwrap(); + assert_eq!(lang.primary_language(), "ja"); + } + + #[test] + fn test_multiple() { + let list = parse_accept_language("en-US, es-419"); + let first = list.get(0).unwrap(); + let second = list.get(1).unwrap(); + assert_eq!(first.primary_language(), "en"); + assert_eq!(second.primary_language(), "es"); + + let list = parse_accept_language("en, es, en-GB"); + let first = list.get(0).unwrap(); + let second = list.get(1).unwrap(); + let third= list.get(2).unwrap(); + assert_eq!(first.primary_language(), "en"); + assert_eq!(second.primary_language(), "es"); + assert_eq!(third.primary_language(), "en"); + } +} diff --git a/crates/lib/http_server_common/src/response/mod.rs b/crates/lib/http_server_common/src/response/mod.rs new file mode 100644 index 0000000000..6bcb009c18 --- /dev/null +++ b/crates/lib/http_server_common/src/response/mod.rs @@ -0,0 +1,4 @@ +pub mod response_error_helpers; +pub mod response_success_helpers; +pub mod serialize_as_json_error; +pub mod to_json_success_response; diff --git a/crates/lib/http_server_common/src/response/response_error_helpers.rs b/crates/lib/http_server_common/src/response/response_error_helpers.rs new file mode 100644 index 0000000000..20e0010716 --- /dev/null +++ b/crates/lib/http_server_common/src/response/response_error_helpers.rs @@ -0,0 +1,30 @@ +use actix_web::HttpResponseBuilder; +use actix_web::http::StatusCode; +use actix_web::HttpResponse; + + +#[derive(Serialize)] +pub struct SimpleGenericJsonError { + pub success: bool, + pub error_reason: String, +} + +pub fn to_simple_json_error( + error_reason: &str, + status_code: StatusCode, +) -> HttpResponse { + + let response = SimpleGenericJsonError { + success: false, + error_reason : error_reason.to_string(), + }; + + let body = match serde_json::to_string(&response) { + Ok(json) => json, + Err(_) => "{}".to_string(), + }; + + HttpResponseBuilder::new(status_code) + .content_type("application/json") + .body(body) +} \ No newline at end of file diff --git a/crates/lib/http_server_common/src/response/response_success_helpers.rs b/crates/lib/http_server_common/src/response/response_success_helpers.rs new file mode 100644 index 0000000000..dae959059f --- /dev/null +++ b/crates/lib/http_server_common/src/response/response_success_helpers.rs @@ -0,0 +1,19 @@ +use actix_web::HttpResponse; + +#[derive(Serialize)] +pub struct SimpleGenericJsonSuccess { + pub success: bool, +} + +pub fn simple_json_success() -> HttpResponse { + let response = SimpleGenericJsonSuccess { + success: true, + }; + + let body = serde_json::to_string(&response) + .unwrap_or("{ \"success\": true }".to_string()); + + HttpResponse::Ok() + .content_type("application/json") + .body(body) +} diff --git a/crates/lib/http_server_common/src/response/serialize_as_json_error.rs b/crates/lib/http_server_common/src/response/serialize_as_json_error.rs new file mode 100644 index 0000000000..9dec81d745 --- /dev/null +++ b/crates/lib/http_server_common/src/response/serialize_as_json_error.rs @@ -0,0 +1,20 @@ +use actix_web::HttpResponse; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use serde::Serialize; + +/// Turn error responses into JSON HTTP responses +pub fn serialize_as_json_error( + error_payload: &T, +) -> HttpResponse + where T: ?Sized + Serialize + ResponseError, +{ + let body = match serde_json::to_string(error_payload) { + Ok(json) => json, + Err(_) => "{}".to_string(), + }; + + HttpResponseBuilder::new(error_payload.status_code()) + .content_type("application/json") + .body(body) +} diff --git a/crates/lib/http_server_common/src/response/to_json_success_response.rs b/crates/lib/http_server_common/src/response/to_json_success_response.rs new file mode 100644 index 0000000000..6d981a9d1b --- /dev/null +++ b/crates/lib/http_server_common/src/response/to_json_success_response.rs @@ -0,0 +1,11 @@ +use serde::Serialize; +use actix_web::HttpResponse; +use container_common::anyhow_result::AnyhowResult; + +pub fn to_json_success_response(response: &impl Serialize) -> AnyhowResult { + let body = serde_json::to_string(&response)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/lib/redis_common/Cargo.toml b/crates/lib/redis_common/Cargo.toml new file mode 100644 index 0000000000..79c7d5f2eb --- /dev/null +++ b/crates/lib/redis_common/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "redis_common" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "redis_common" +path = "src/lib.rs" + +[dependencies] + +# Internal +container_common = { path = "../container_common" } + +# External +anyhow = "1.0.40" + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" + +# TODO: Get rid of in favor of std::time +time = "0.2.26" diff --git a/crates/lib/redis_common/src/lib.rs b/crates/lib/redis_common/src/lib.rs new file mode 100644 index 0000000000..ffd322f242 --- /dev/null +++ b/crates/lib/redis_common/src/lib.rs @@ -0,0 +1,5 @@ +#[macro_use] extern crate serde_derive; + +pub mod payloads; +pub mod redis_keys; +pub mod shared_constants; diff --git a/crates/lib/redis_common/src/payloads/lease_payload.rs b/crates/lib/redis_common/src/payloads/lease_payload.rs new file mode 100644 index 0000000000..6d1b415e65 --- /dev/null +++ b/crates/lib/redis_common/src/payloads/lease_payload.rs @@ -0,0 +1,81 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::thread::thread_id::ThreadId; +use container_common::token::random_crockford_token::random_crockford_token; + +// TODO: Rename, document + +// TODO: This should be JSON for future compat. +/// Lease values stored in Redis are in the form "{server_id}:{thread_id}" +#[derive(Debug, Eq, PartialEq, Clone)] +pub struct LeasePayload { + pub server_id: String, + pub thread_id: String, +} + +impl LeasePayload { + pub fn with_random_thread_id(server_id: &str) -> Self { + let thread_id = random_crockford_token(10); + Self::from_string_id(server_id, &thread_id) + } + + pub fn from_thread_id(server_id: &str, thread_id: &ThreadId) -> Self { + Self::from_string_id(server_id, thread_id.get_id()) + } + + pub fn from_string_id(server_id: &str, thread_id: &str) -> Self { + Self { + server_id: server_id.to_string(), + thread_id: thread_id.to_string(), + } + } + + pub fn serialize(&self) -> String { + format!("{}:{}", &self.server_id, &self.thread_id) + } + + pub fn deserialize(payload: &str) -> AnyhowResult { + let pieces = payload.split(":").collect::>(); + if pieces.len() != 2 { + return Err(anyhow!("Invalid payload: {}", payload)); + } + + let pair = (pieces.get(0), pieces.get(1)); + + if let (Some(k), Some(v)) = pair { + return Ok(Self::from_string_id(k, v)) + } else { + return Err(anyhow!("Invalid payload: {}", payload)); + } + } +} + +#[cfg(test)] +mod tests { + use container_common::thread::thread_id::ThreadId; + use crate::payloads::lease_payload::LeasePayload; + + #[test] + fn equals() { + let a = LeasePayload::from_string_id("foo", "bar"); + let b = LeasePayload::from_string_id("foo", "bar"); + assert_eq!(a, b); + let t = ThreadId::with_id("bar"); + let b = LeasePayload::from_thread_id("foo", &t); + assert_eq!(a, b); + } + + #[test] + fn not_equals() { + let a = LeasePayload::from_string_id("foo", "bar"); + let b = LeasePayload::from_string_id("asdf", "qwerty"); + assert_ne!(a, b); + let b = LeasePayload::from_string_id("foo", "ASDF"); + assert_ne!(a, b); + let b = LeasePayload::from_string_id("ASDF", "bar"); + assert_ne!(a, b); + let t = ThreadId::with_id("ASDF"); + let b = LeasePayload::from_thread_id("foo", &t); + assert_ne!(a, b); + } +} diff --git a/crates/lib/redis_common/src/payloads/mod.rs b/crates/lib/redis_common/src/payloads/mod.rs new file mode 100644 index 0000000000..6ed256ff9b --- /dev/null +++ b/crates/lib/redis_common/src/payloads/mod.rs @@ -0,0 +1,2 @@ +pub mod lease_payload; +pub mod obs_active_payload; \ No newline at end of file diff --git a/crates/lib/redis_common/src/payloads/obs_active_payload.rs b/crates/lib/redis_common/src/payloads/obs_active_payload.rs new file mode 100644 index 0000000000..cca1154553 --- /dev/null +++ b/crates/lib/redis_common/src/payloads/obs_active_payload.rs @@ -0,0 +1,24 @@ +use container_common::anyhow_result::AnyhowResult; + +// TODO: Rename, document + +#[derive(Serialize, Deserialize, Debug)] +pub struct ObsActivePayload { + pub twitch_user_id: String, +} + +impl ObsActivePayload { + pub fn new(twitch_user_id: &str) -> Self { + Self { + twitch_user_id: twitch_user_id.to_string(), + } + } + + pub fn from_json_str(json: &str) -> AnyhowResult { + Ok(serde_json::from_str(json)?) + } + + pub fn serialize(&self) -> AnyhowResult { + Ok(serde_json::to_string(&self)?) + } +} \ No newline at end of file diff --git a/crates/lib/redis_common/src/redis_keys.rs b/crates/lib/redis_common/src/redis_keys.rs new file mode 100644 index 0000000000..e38f826b30 --- /dev/null +++ b/crates/lib/redis_common/src/redis_keys.rs @@ -0,0 +1,81 @@ +/// Centralize all of the Redis key naming and dereferencing. +pub struct RedisKeys {} + +impl RedisKeys { + + // =============== FAKEYOU / TTS / W2L =============== + + /// This is a counter incremented with TTS inference. + /// We can use a job to read MySql and fix counts if they get out of sync. + /// These should be long-lived keys. + pub fn tts_model_usage_count(model_token: &str) -> String { + format!("ttsUseCount:{}", model_token) + } + + /// This is a counter incremented with TTS inference. + /// We can use a job to read MySql and fix counts if they get out of sync. + /// These should be long-lived keys. + pub fn w2l_template_usage_count(template_token: &str) -> String { + format!("w2lUseCount:{}", template_token) + } + + /// We write extra in-progress status information to keys. + /// These keys should have a TTL. + pub fn tts_download_extra_status_info(inference_job_token: &str) -> String { + format!("ttsDownloadExtraStatus:{}", inference_job_token) + } + + /// We write extra in-progress status information to keys. + /// These keys should have a TTL. + pub fn tts_inference_extra_status_info(inference_job_token: &str) -> String { + format!("ttsInferenceExtraStatus:{}", inference_job_token) + } + + /// We write extra in-progress status information to keys. + /// These keys should have a TTL. + pub fn w2l_download_extra_status_info(inference_job_token: &str) -> String { + format!("w2lDownloadExtraStatus:{}", inference_job_token) + } + + /// We write extra in-progress status information to keys. + /// These keys should have a TTL. + pub fn w2l_inference_extra_status_info(inference_job_token: &str) -> String { + format!("w2lInferenceExtraStatus:{}", inference_job_token) + } + + // =============== OBS / TWITCH =============== + + /// This is a PubSub topic. + /// We publish when new viewers arrive and as a "keep-alive ping". + pub fn obs_active_sessions_topic() -> &'static str { + "obsActiveSessionsTopic" + } + + /// This is a key that denotes a user is actively on the OBS gateway. + /// These keys have a short expiry and OBS gateway must continually bump them. + /// If a key goes away, the thread monitoring that user should exit. + pub fn obs_active_session_keepalive(twitch_user_id: &str) -> String { + format!("obsActiveSessionKeepalive:{}", twitch_user_id) + } + + /// This is a key that denotes a PubSub lease subscriber. + /// Only one thread should be PubSub subscribed to any given user at a time. + /// These keys should have a TTL so that they naturally expire without maintenance. + pub fn twitch_pubsub_lease(twitch_user_id: &str) -> String { + format!("twitchPubsubLease:{}", twitch_user_id) + } + + /// This is a PubSub topic. + /// These are "unenriched" Twitch events from the PubSub (and eventually IRC) + /// subscribers. Downstream listeners will enrich these for user-facing functionality. + pub fn twitch_events_topic(twitch_user_id: &str) -> String { + format!("twitchEventsTopic:{}", twitch_user_id) + } + + /// This is a lazy first approximation of the worker -> OBS messaging. + /// We push new TTS job tokens into a Redis list (with long expiry), and OBS + /// will dequeue the events in order. + pub fn twitch_tts_job_queue(twitch_user_id: &str) -> String { + format!("twitchTtsJobs:{}", twitch_user_id) + } +} diff --git a/crates/lib/redis_common/src/shared_constants.rs b/crates/lib/redis_common/src/shared_constants.rs new file mode 100644 index 0000000000..19dd683298 --- /dev/null +++ b/crates/lib/redis_common/src/shared_constants.rs @@ -0,0 +1,25 @@ +use time::Duration; + +// TODO: Think about these timings more. + +/// Our Twitch PubSub subscriber threads must renew leases in +/// order to be granted an exclusive "lock" on a single user. +/// If the underlying Redis key goes away, another thread can +/// begin leasing. +pub const LEASE_TIMEOUT_SECONDS : usize = 60; + +/// The amount of time to wait before renewing our lease. +pub const LEASE_RENEW_PERIOD : Duration = Duration::seconds(30); + +/// The amount of time to wait before checking or lease is valid. +pub const LEASE_CHECK_PERIOD : Duration = Duration::seconds(15); + +/// The OBS websocket polls a key to inform the backend that the +/// user is still active. +pub const OBS_ACTIVE_TTL_SECONDS : usize = 120; + +/// How often to check the OBS session activity. +pub const OBS_ACTIVE_CHECK_PERIOD : Duration = Duration::seconds(30); + +/// When to expire the streamer-specific TTS job queue notifications +pub const STREAMER_TTS_JOB_QUEUE_TTL_SECONDS : usize = 600; diff --git a/crates/lib/tts_common/Cargo.toml b/crates/lib/tts_common/Cargo.toml new file mode 100644 index 0000000000..c4578df6be --- /dev/null +++ b/crates/lib/tts_common/Cargo.toml @@ -0,0 +1,27 @@ +[package] +name = "tts_common" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "tts_common" +path = "src/lib.rs" + +[dependencies] + +# Internal +container_common = { path = "../container_common" } +user_input_common = { path = "../user_input_common" } + +# External +anyhow = "1.0.40" +log = "0.4.14" +once_cell = "1.9.0" +regex = "1.5.4" +unicode-segmentation = "1.7.1" + diff --git a/crates/lib/tts_common/src/clean_symbols.rs b/crates/lib/tts_common/src/clean_symbols.rs new file mode 100644 index 0000000000..b4f6cdd0ff --- /dev/null +++ b/crates/lib/tts_common/src/clean_symbols.rs @@ -0,0 +1,449 @@ +use once_cell::sync::Lazy; +use std::collections::HashMap; +use unicode_segmentation::UnicodeSegmentation; +use user_input_common::latin_alphabet::LATIN_TO_ASCII_CHARACTER_MAP; + +// Used to insert tables into hashmap +fn deref_to_owned(item: (&String, &String)) -> (String, String) { + (item.0.to_string(), item.1.to_string()) +} + +fn to_owned(item: &(&str, &str)) -> (String, String) { + (item.0.to_string(), item.1.to_string()) +} + +// https://unicodelookup.com/#quo +// http://www.geocities.ws/click2speak/unicode/chars_es.html +static REPLACEMENTS : Lazy> = Lazy::new(|| { + let mut map = HashMap::new(); + + // Latin characters such as àáâãäå + map.extend(LATIN_TO_ASCII_CHARACTER_MAP.iter().map(&deref_to_owned)); + + // Spacing character replacements + map.extend([ + ("\u{00A0}", " "), // Non-break space (aka  ) \xa0 + ("\u{2003}", " "), // Em Space + ("\u{2004}", " "), // Three-Per-Em Space + ("\u{2005}", " "), // Four-Per-Em Space + ("\u{2006}", " "), // Six-Per-Em Space + ("\u{2007}", " "), // Figure Space + ("\u{2008}", " "), // Punctuation Space + ("\u{2009}", " "), // Thin Space + ("\u{200A}", " "), // Hair Space + ("\u{200B}", " "), // Zero Width Space + ("\u{200C}", " "), // Zero Width Non-Joiner + ("\u{200D}", " "), // Zero Width Joiner + ("\u{2028}", " "), // Line Separator + ("\u{2029}", " "), // Paragraph Separator + ("\u{205F}", " "), // Medium Mathematical Space (MMSP) + ("\u{2588}", " "), // Full Block + ("\u{2800}", " "), // Braille Pattern Blank + ("\u{3000}", " "), // Ideographic Space + ("\u{3164}", " "), // Hangul Filler + ("\u{FEFF}", " "), // Zero Width No-Break Space + ].iter().map(&to_owned)); + + // Punctuation that should turn into spaces + map.extend([ + ("\u{b7}", " "), // · Middle Dot + ("\u{2022}", " "), // • Bullet + ].iter().map(&to_owned)); + + // Quotes (single) + map.extend([ + ("\u{0060}", "'"), // Grave Accent + ("\u{00B4}", "'"), // Acute Accent + ("\u{2018}", "'"), // Left Single Quotation Mark + ("\u{2019}", "'"), // Right Single Quotation Mark + ("\u{201A}", "'"), // Single Low-9 Quotation Mark + ("\u{201B}", "'"), // Single High-Revered-9 Quotation Mark + ].iter().map(&to_owned)); + + // Quotes (double) + map.extend([ + ("\u{201C}", "\""), // Left Double Quotation Mark + ("\u{201D}", "\""), // Right Double Quotation Mark + ("\u{201E}", "\""), // Double Low-9 Quotation Mark + ("\u{201F}", "\""), // Double High-Reversed-9 Quotation Mark + ("\u{301D}", "\""), // Reversed Double Prime Quotation Mark + ("\u{301E}", "\""), // Double Prime Quotation Mark + ("\u{301F}", "\""), // Low Double Prime Quotation Mark + ("\u{FF02}", "\""), // Fullwidth Quotation Mark + ].iter().map(&to_owned)); + + // Dashes + map.extend([ + ("\u{00AD}", "-"), // ­ Soft Hyphen + ("\u{2010}", "-"), // ‐ Hyphen + ("\u{2011}", "-"), // ‑ Non-Breaking Hyphen + ("\u{2013}", "-"), // – En Dash + ("\u{2014}", "-"), // — Em Dash + ("\u{2015}", "-"), // ― Horizontal Bar + ("\u{2E3A}", "-"), // ⸺ Two-Em Dash + ("\u{2E3B}", "-"), // ⸻ Three-Em Dash + ("\u{FE58}", "-"), // ﹘ Small Em Dash + ("\u{FE63}", "-"), // ﹣ Small Hyphen-Minus + ("\u{FF0D}", "-"), // － Fullwidth Hyphen-Minus + ].iter().map(&to_owned)); + + // Close enough to existing allowed punctuation + map.extend([ + ("\u{00A1}", "!"), // ¡ Inverted Exclamation Mark + ("\u{00BF}", "?"), // ¿ Inverted Question Mark + ("\u{2024}", "."), // ․ One Dot Leader + ("\u{2025}", ".."), // ‥ Two Dot Leader + ("\u{2026}", "..."), // … Horizontal Ellipsis + ("\u{203C}", "!!"), // ‼ Double Exclamation Mark Emoji + ("\u{203D}", "!?"), // ‽ Interrobang + ("\u{2588}", " "), // █ Full Block + ("\u{3001}", ","), // 、 Ideographic Comma + ("\u{3002}", "."), // 。 Ideographic Full Stop + ("\u{FF01}", "!"), // ！ Fullwidth Exclamation Mark + ].iter().map(&to_owned)); + + // Spanish special characters + map.extend([ + ("\u{AA}", "a"), // Feminine ordinal + ("\u{BA}", "o"), // Masculine ordinal + ].iter().map(&to_owned)); + + // Symbols we can insert as words + map.extend([ + ("\u{00A9}", " copyright "), // © Copyright Sign Emoji + ("\u{00B0}", " degrees "), // ° Degree Sign + ("\u{03C0}", " pie "), // greek small letter pi (TODO: Incorrect handling) + ("\u{2122}", " trademark "), // ™ Trade Mark Sign Emoji + ].iter().map(&to_owned)); + + // Emoji we can insert as words + // https://unicode-table.com/en/blocks/emoticons/ + map.extend([ + ("\u{1F33D}", " corn on the cobb "), // 🌽 Ear of Maize Emoji + ("\u{1F436}", " dog face "), // 🐶 Dog Face Emoji + ("\u{1F44C}", " okay "), // 👌 Ok Hand Sign Emoji + ("\u{1F4A6}", " splashing sweat "), // 💦 Splashing Sweat Symbol Emoji + ("\u{1F4A9}", " poop "), // 💩 Pile of Poo Emoji + ("\u{1F4AF}", " hundred points "), // 💯 Hundred Points Symbol Emoji + ("\u{1F525}", " fire "), // 🔥 Fire Emoji + ("\u{1F530}", " japanese symbol for beginner "), // 🔰 Japanese Symbol for Beginner Emoji + ("\u{1F602}", " face with tears of joy "), // 😂 Face with Tears of Joy Emoji + ("\u{1F60D}", " heart eyes "), // 😍 Smiling Face with Heart-Shaped Eyes Emoji + ("\u{1F62D}", " loudly crying face "), // 😭 Loudly Crying Face Emoji + ("\u{1F633}", " flushed face "), // 😳 Flushed Face Emoji + ("\u{1F923}", " rolling on the floor laughing "), // 🤣 Rolling On The Floor Laughing Emoji + ].iter().map(&to_owned)); + + map +}); + +/// Tacotron chokes on Unicode, latin characters, special punctuation, emoji, etc. +/// We attempt to turn as much as possible into ASCII and a reduced set of symbols. +pub fn clean_symbols(input_text: &str) -> String { + let segmented= UnicodeSegmentation::graphemes(input_text, true) + .map(|segment| { + if let Some(replace) = REPLACEMENTS.get(segment) { + return replace.as_str(); + } + return segment; + }) + .collect::>(); + + let before_clean : String = segmented.join(""); + + before_clean.chars() + .filter(|segment| { + match segment { + '\x00'..='\x7f' => true, // Full ASCII range + //'\x01'...'\x08' | '\u{10FFFE}'...'\u{10FFFF}' => true, + _ => false, + } + }) + .collect::() +} + +#[cfg(test)] +mod tests { + use crate::clean_symbols::clean_symbols; + + fn assert_converted(original: &str, expected: &str) { + assert_eq!(clean_symbols(original), expected.to_string()); + } + + #[test] + fn neutral_tests() { + assert_eq!(clean_symbols(""), "".to_string()); // Empty check + assert_eq!(clean_symbols("this should be the same."), "this should be the same.".to_string()); + assert_eq!(clean_symbols("one\ntwo\r\nthree "), "one\ntwo\r\nthree ".to_string()); + } + + #[test] + fn assert_ascii_retained() { + assert_converted("This, sentence. It\nhas\nnewlines.", + "This, sentence. It\nhas\nnewlines."); + assert_converted("12:34", "12:34"); + assert_converted("Punctuation!?.", "Punctuation!?."); + } + + #[test] + fn assert_non_handled_emoji_removed() { + assert_converted("👹", ""); + assert_converted("👹👍✅", ""); + assert_converted("✨what✨", "what"); + } + + #[test] + fn assert_non_handled_languages_removed() { + assert_converted("これはテスト", ""); + assert_converted("これはテストdesu", "desu"); + assert_converted("你好", ""); + assert_converted("hello你好hello", "hellohello"); + } + + #[test] + fn filters_smart_quotes() { + assert_eq!(clean_symbols("That’s ok"), "That's ok".to_string()); + assert_eq!(clean_symbols("That’s it’s new home"), "That's it's new home".to_string()); + assert_eq!(clean_symbols("‘foo’"), "'foo'".to_string()); + assert_eq!(clean_symbols("“as ‘shown’ here.”"), "\"as 'shown' here.\"".to_string()); + } + + #[test] + fn filters_dashes() { + assert_eq!(clean_symbols("en – dash"), "en - dash".to_string()); + assert_eq!(clean_symbols("em — dash"), "em - dash".to_string()); + assert_eq!(clean_symbols("three em ⸻ dash"), "three em - dash".to_string()); + } + + #[test] + fn filters_ellipsis() { + assert_eq!(clean_symbols("test…"), "test...".to_string()); + } + + #[test] + fn symbol_expansion() { + assert_eq!(clean_symbols("Pokémon™ is a popular video game series"), + "Pokemon trademark is a popular video game series".to_string()); // NB: Extra space + } + + #[test] + fn filters_latin_characters() { + assert_eq!(clean_symbols("pokémon"), "pokemon".to_string()); + assert_eq!(clean_symbols("POKÉMON"), "POKEMON".to_string()); + assert_eq!(clean_symbols("Æther"), "AEther".to_string()); + assert_eq!(clean_symbols("æther"), "aether".to_string()); + // Almost exhaustive + assert_eq!(clean_symbols("ÀÁÂÃÄÅ"), "AAAAAA".to_string()); + assert_eq!(clean_symbols("Æ"), "AE".to_string()); + assert_eq!(clean_symbols("Ç"), "C".to_string()); + assert_eq!(clean_symbols("ÈÉÊË"), "EEEE".to_string()); + assert_eq!(clean_symbols("ÌÍÎÏ"), "IIII".to_string()); + assert_eq!(clean_symbols("ÒÓÔÕÖØ"), "OOOOOO".to_string()); + assert_eq!(clean_symbols("ÙÚÛÜ"), "UUUU".to_string()); + assert_eq!(clean_symbols("Ý"), "Y".to_string()); + assert_eq!(clean_symbols("àáâãäå"), "aaaaaa".to_string()); + assert_eq!(clean_symbols("æ"), "ae".to_string()); + assert_eq!(clean_symbols("ç"), "c".to_string()); + assert_eq!(clean_symbols("èéêë"), "eeee".to_string()); + assert_eq!(clean_symbols("ìíîï"), "iiii".to_string()); + } + + #[test] + fn filters_spanish_characters() { + assert_eq!(clean_symbols("¡"), "!".to_string()); + assert_eq!(clean_symbols("¿"), "?".to_string()); + } + + #[test] + fn actual_database_failures() { + assert_eq!(clean_symbols( + "Sabías que?,tu papá es el tercer planeta del sistema solar"), + "Sabias que?,tu papa es el tercer planeta del sistema solar".to_string()); + assert_eq!(clean_symbols("señoras"), "senoras".to_string()); + + // Most frequent according to partial histogram + assert_eq!(clean_symbols("á"), "a".to_string()); + assert_eq!(clean_symbols("í"), "i".to_string()); + assert_eq!(clean_symbols("ó"), "o".to_string()); + assert_eq!(clean_symbols("é"), "e".to_string()); + assert_eq!(clean_symbols("ñ"), "n".to_string()); + assert_eq!(clean_symbols("ú"), "u".to_string()); + assert_eq!(clean_symbols("ı"), "i".to_string()); // Dotless i + assert_eq!(clean_symbols("ü"), "u".to_string()); + assert_eq!(clean_symbols("¿"), "?".to_string()); + assert_eq!(clean_symbols("…"), "...".to_string()); + assert_eq!(clean_symbols("¡"), "!".to_string()); + assert_eq!(clean_symbols("ş"), "s".to_string()); + assert_eq!(clean_symbols("ç"), "c".to_string()); + assert_eq!(clean_symbols("”"), "\"".to_string()); + assert_eq!(clean_symbols("“"), "\"".to_string()); + assert_eq!(clean_symbols("ö"), "o".to_string()); + assert_eq!(clean_symbols("ğ"), "g".to_string()); + assert_eq!(clean_symbols("\u{00a0}"), " ".to_string()); + assert_eq!(clean_symbols("ã"), "a".to_string()); + assert_eq!(clean_symbols("à"), "a".to_string()); + assert_eq!(clean_symbols("Á"), "A".to_string()); + assert_eq!(clean_symbols("İ"), "I".to_string()); + assert_eq!(clean_symbols(" "), " ".to_string()); + assert_eq!(clean_symbols(" "), " ".to_string()); + assert_eq!(clean_symbols("—"), "-".to_string()); + assert_eq!(clean_symbols("Ñ"), "N".to_string()); + assert_eq!(clean_symbols("´"), "'".to_string()); + assert_eq!(clean_symbols("ê"), "e".to_string()); + assert_eq!(clean_symbols("ə"), "e".to_string()); + assert_eq!(clean_symbols("è"), "e".to_string()); + assert_eq!(clean_symbols("ī"), "i".to_string()); + assert_eq!(clean_symbols("ᴺ"), "n".to_string()); + assert_eq!(clean_symbols(" "), " ".to_string()); + assert_eq!(clean_symbols("ä"), "a".to_string()); + assert_eq!(clean_symbols("ǐ"), "i".to_string()); + assert_eq!(clean_symbols("█"), " ".to_string()); + assert_eq!(clean_symbols("ō"), "o".to_string()); + assert_eq!(clean_symbols("‘"), "'".to_string()); + assert_eq!(clean_symbols("Í"), "I".to_string()); + assert_eq!(clean_symbols("É"), "E".to_string()); + assert_eq!(clean_symbols("ā"), "a".to_string()); + assert_eq!(clean_symbols("ù"), "u".to_string()); + assert_eq!(clean_symbols("、"), ",".to_string()); + assert_eq!(clean_symbols("ǎ"), "a".to_string()); + assert_eq!(clean_symbols("‼"), "!!".to_string()); + + // TODO: assert_eq!(clean_symbols("ß"), "B".to_string()); + // TODO: assert_eq!(clean_symbols("°"), "degrees".to_string()); + } + + #[test] + pub fn most_frequent_failures_100k_usages() { + // The leftmost number is the number of occurrences in our histogram of TTS failures + // The order of the lines is reversed. + assert_converted("ñ", "n"); // b'\xf1' 105898 + assert_converted("ó", "o"); // b'\xf3' 127099 + assert_converted("í", "i"); // b'\xed' 132716 + assert_converted("é", "e"); // b'\xe9' 140278 + assert_converted("á", "a"); // b'\xe1' 184138 + assert_converted("’", "'"); // b'\\u2019' 739079 + } + + #[test] + pub fn most_frequent_failures_20k_usages() { + assert_converted("ǐ", "i"); // b'\\u01d0' 20236 + assert_converted("ī", "i"); // b'\\u012b' 20257 + assert_converted("è", "e"); // b'\xe8' 21121 + assert_converted("°", " degrees "); // b'\xb0' 22808 + assert_converted("“", "\""); // b'\\u201c' 27343 + assert_converted("ü", "u"); // b'\xfc' 27955 + assert_converted("”", "\""); // b'\\u201d' 28171 + assert_converted("à", "a"); // b'\xe0' 34414 + assert_converted("ı", "i"); // b'\\u0131' 44032 + assert_converted("ú", "u"); // b'\xfa' 44043 + assert_converted("…", "..."); // b'\\u2026' 49348 + } + + #[test] + pub fn most_frequent_failures_10k_usages() { + assert_converted("ş", "s"); // b'\\u015f' 10006 + assert_converted(" ", " "); // b'\\u2005' 10118 + assert_converted("—", "-"); // b'\\u2014' 10555 + assert_converted("ã", "a"); // b'\xe3' 10586 + assert_converted("ǎ", "a"); // b'\\u01ce' 10843 + assert_converted("ə", "e"); // b'\\u0259' 11508 + assert_converted("ö", "o"); // b'\xf6' 11908 + assert_converted("¿", "?"); // b'\xbf' 13501 + assert_converted("ç", "c"); // b'\xe7' 13784 + assert_converted(" ", " "); // b'\xa0' 14184 + assert_converted("ō", "o"); // b'\\u014d' 14329 + assert_converted("ā", "a"); // b'\\u0101' 15927 + assert_converted("¡", "!"); // b'\xa1' 17933 + } + + #[test] + pub fn most_frequent_failures_5k_usages() { + assert_converted("​", " "); // b'\\u200b' 6019 + assert_converted("ğ", "g"); // b'\\u011f' 6103 + assert_converted("–", "-"); // b'\\u2013' 6172 + assert_converted("ì", "i"); // b'\xec' 6209 + assert_converted("ē", "e"); // b'\\u0113' 6312 + assert_converted("ù", "u"); // b'\xf9' 6907 + assert_converted("‘", "'"); // b'\\u2018' 7151 + assert_converted("ǔ", "u"); // b'\\u01d4' 7422 + assert_converted("´", "'"); // b'\xb4' 8210 + assert_converted("ê", "e"); // b'\xea' 8882 + assert_converted("ä", "a"); // b'\xe4' 9202 + } + + #[test] + pub fn most_frequent_failures_2p5k_usages() { + assert_converted("～", "~"); // b'\\uff5e' 2507 + assert_converted("â", "a"); // b'\xe2' 2607 + assert_converted("α", "a"); // b'\\u03b1' 2625 + assert_converted("å", "a"); // b'\xe5' 2753 + assert_converted("🐶", " dog face "); // b'\\U0001f436' 2782 + assert_converted("™", " trademark "); // b'\\u2122' 2869 + assert_converted("É", "E"); // b'\xc9' 3040 + assert_converted("æ", "ae"); // b'\xe6' 3142 + assert_converted("¨", "\""); // b'\xa8' 3236 + assert_converted("ò", "o"); // b'\xf2' 3250 + assert_converted("đ", "d"); // b'\\u0111' 3420 + assert_converted("，", ","); // b'\\uff0c' 3487 + assert_converted("ô", "o"); // b'\xf4' 3568 + assert_converted("Á", "A"); // b'\xc1' 3779 + assert_converted("ß", "B"); // b'\xdf' 3779 + assert_converted("Ñ", "N"); // b'\xd1' 3963 + assert_converted("ǒ", "o"); // b'\\u01d2' 4613 + assert_converted("İ", "I"); // b'\\u0130' 4808 + } + + #[test] + pub fn most_frequent_failures_2k_usages() { + assert_converted("、", ","); // b'\\u3001' 2001 + assert_converted("ć", "c"); // b'\\u0107' 2017 + assert_converted("ę", "e"); // b'\\u0119' 2236 + assert_converted("。", "."); // b'\\u3002' 2288 + assert_converted(" ", " "); // b'\\u205f' 2347 + assert_converted("ᴺ", "n"); // b'\\u1d3a' 2383 + assert_converted("ě", "e"); // b'\\u011b' 2441 + assert_converted("ᴾ", "p"); // b'\\u1d3e' 2479 + assert_converted("ł", "s"); // b'\\u0142' 2480 + } + + #[test] + pub fn most_frequent_failures_1p5k_usages() { + assert_converted("·", " "); // b'\xb7' 1514 + assert_converted("š", "s"); // b'\\u0161' 1514 + assert_converted("ο", "o"); // b'\\u03bf' 1515 + assert_converted("ū", "u"); // b'\\u016b' 1523 + assert_converted("Ö", "O"); // b'\xd6' 1543 + assert_converted("ι", "l"); // b'\\u03b9' 1559 + assert_converted("ε", "e"); // b'\\u03b5' 1575 + assert_converted("ă", "a"); // b'\\u0103' 1576 + //assert_converted("😂", ""); // b'\\U0001f602' 1741 + assert_converted("Ó", "O"); // b'\xd3' 1774 + assert_converted("ư", "u"); // b'\\u01b0' 1794 + assert_converted("•", " "); // b'\\u2022' 1948 + } + + #[test] + pub fn most_frequent_failures_1k_usages() { + assert_converted("　", " "); // b'\\u3000' 1005 + //assert_converted("🤣", " laugh "); // b'\\U0001f923' 1010 + assert_converted("ø", "o"); // b'\xf8' 1017 + assert_converted("！", "!"); // b'\\uff01' 1019 + //assert_converted("¥", " yen "); // b'\xa5' 1033 + //assert_converted("😭", "cry "); // b'\\U0001f62d' 1061 + assert_converted("Ü", "U"); // b'\xdc' 1062 + assert_converted("č", "c"); // b'\\u010d' 1092 + assert_converted("ν", "v"); // b'\\u03bd' 1096 + assert_converted("ż", "z"); // b'\\u017c' 1099 + assert_converted("⠀", " "); // b'\\u2800' 1166 + assert_converted("ą", "a"); // b'\\u0105' 1169 + //assert_converted("£", " pounds "); // b'\xa3' 1186 + assert_converted("ë", "e"); // b'\xeb' 1189 + assert_converted("Ç", "C"); // b'\xc7' 1213 + //assert_converted("τ", "t"); // b'\\u03c4' 1217 + assert_converted("ẹ", "e"); // b'\\u1eb9' 1228 + assert_converted("î", "i"); // b'\xee' 1349 + assert_converted("ś", "s"); // b'\\u015b' 1402 + assert_converted("ạ", "a"); // b'\\u1ea1' 1413 + assert_converted("Í", "I"); // b'\xcd' 1434 + } +} diff --git a/crates/lib/tts_common/src/lib.rs b/crates/lib/tts_common/src/lib.rs new file mode 100644 index 0000000000..82df584c77 --- /dev/null +++ b/crates/lib/tts_common/src/lib.rs @@ -0,0 +1 @@ +pub mod clean_symbols; \ No newline at end of file diff --git a/crates/lib/twitch_common/Cargo.toml b/crates/lib/twitch_common/Cargo.toml new file mode 100644 index 0000000000..11fc427157 --- /dev/null +++ b/crates/lib/twitch_common/Cargo.toml @@ -0,0 +1,33 @@ +[package] +name = "twitch_common" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "twitch_common" +path = "src/lib.rs" + +[dependencies] + +# Internal +container_common = { path = "../container_common" } + +# External +anyhow = "1.0.40" +once_cell = "1.9.0" +regex = "1.5.4" +easyenv = "0.0.6" + +# Twitch +# NB "unsupported" pubsub features may disappear from the Twitch API at any time. It's an unofficial API. +# twitch_api2 = { version = "0.6.0-rc.2", features = ["helix", "reqwest_client", "pubsub", "chrono", "unsupported"] } +twitch_oauth2 = "0.6.0-rc.2" + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" diff --git a/crates/lib/twitch_common/src/cheers.rs b/crates/lib/twitch_common/src/cheers.rs new file mode 100644 index 0000000000..4d55ede5a0 --- /dev/null +++ b/crates/lib/twitch_common/src/cheers.rs @@ -0,0 +1,106 @@ +use once_cell::sync::Lazy; +use std::collections::HashSet; +use regex::Regex; + +// NB: This can be used as a nice, succinct macro with once_cell: +//macro_rules! regex { +// ($re:literal $(,)?) => {{ +// static RE: once_cell::sync::OnceCell = once_cell::sync::OnceCell::new(); +// RE.get_or_init(|| regex::Regex::new($re).unwrap()) +// }}; +//} + +/// Remove cheermotes from text, eg "this is great Cheer1" -> "this is great" +pub fn remove_cheers(text: &str) -> String { + let result = TWITCH_CHEER_REGEX.replace_all(text, ""); + let result = MULTI_SPACE_REGEX.replace_all(&result, " ").to_string(); + result.trim().to_string() +} + +/// These cheer emotes always combine with a number +/// This list seems pretty comprehensive +/// https://github.com/nossebro/TwitchPubSubMirror/blob/master/TwitchPubSubMirror_StreamlabsSystem.py + +static TWITCH_CHEERS : Lazy> = Lazy::new(|| { + let cheers = vec![ + "Cheer", + "DoodleCheer", + "BibleThump", + "cheerwhal", + "Corgo", + "uni", + "ShowLove", + "Party", + "SeemsGood", + "Pride", + "Kappa", + "FrankerZ", + "HeyGuys", + "DansGame", + "EleGiggle", + "TriHard", + "Kreygasm", + "4Head", + "SwiftRage", + "NotLikeThis", + "FailFish", + "VoHiYo", + "PJSalt", + "MrDestructoid", + "bday", + "RIPCheer", + "Shamrock", + "BitBoss", + "Streamlabs", + "Muxy", + "HolidayCheer" + ]; + let mut s = HashSet::new(); + for cheer in cheers.into_iter() { + s.insert(cheer.to_string()); + } + s +}); + +static MULTI_SPACE_REGEX : Lazy = Lazy::new(|| { + Regex::new("(\\s{2,})").unwrap() +}); + +static TWITCH_CHEER_REGEX : Lazy = Lazy::new(|| { + let mut regex_pieces = Vec::new(); + + for cheer in TWITCH_CHEERS.iter() { + regex_pieces.push(format!("({}\\d+)", cheer)); + } + + let regex = regex_pieces.join("|"); + + // NB: (?i) is case insensitive search + // https://docs.rs/regex/1.5.4/regex/#structs + let regex = format!("(?i)({})", regex); + + Regex::new(®ex).unwrap() +}); + +#[cfg(test)] +mod tests { + use crate::cheers::remove_cheers; + + #[test] + fn test_remove_cheers() { + let r = remove_cheers("testing Kappa1 SeemsGood100 removal"); + assert_eq!(&r, "testing removal"); + + let r = remove_cheers("Kappa1 SeemsGood100 no more Cheer1 cheers"); + assert_eq!(&r, "no more cheers"); + } + + #[test] + fn test_remove_cheers_case_insensitive() { + let r = remove_cheers("testing cheer1"); + assert_eq!(&r, "testing"); + + let r = remove_cheers("TESTING CHEER1"); + assert_eq!(&r, "TESTING"); + } +} diff --git a/crates/lib/twitch_common/src/lib.rs b/crates/lib/twitch_common/src/lib.rs new file mode 100644 index 0000000000..6ac11f6e4e --- /dev/null +++ b/crates/lib/twitch_common/src/lib.rs @@ -0,0 +1,4 @@ +pub mod cheers; +pub mod oauth_token_builder; +pub mod twitch_secrets; +pub mod twitch_user_id; \ No newline at end of file diff --git a/crates/lib/twitch_common/src/oauth_token_builder.rs b/crates/lib/twitch_common/src/oauth_token_builder.rs new file mode 100644 index 0000000000..49588c7ed2 --- /dev/null +++ b/crates/lib/twitch_common/src/oauth_token_builder.rs @@ -0,0 +1,66 @@ +use twitch_oauth2::tokens::UserTokenBuilder; +use twitch_oauth2::{Scope, ClientSecret, ClientId, CsrfToken}; +use twitch_oauth2::url::Url; + +/// Get the token builder with the scopes our app needs +/// This is used in several places, so centralize creation here. +pub fn get_oauth_token_builder( + client_id: &str, + client_secret: &str, + redirect_url: &Url, + force_verify: bool +) -> UserTokenBuilder { + + let client_id = ClientId::new(client_id); + let client_secret = ClientSecret::new(client_secret); + + /* + - AnalyticsReadExtensions, scope: "analytics:read:extensions", doc: "View analytics data for the Twitch Extensions owned by the authenticated account."; + - AnalyticsReadGames, scope: "analytics:read:games", doc: "View analytics data for the games owned by the authenticated account."; + [YES] BitsRead, scope: "bits:read", doc: "View Bits information for a channel."; + - ChannelEditCommercial, scope: "channel:edit:commercial", doc: "Run commercials on a channel."; + - ChannelManageBroadcast, scope: "channel:manage:broadcast", doc: "Manage a channel’s broadcast configuration, including updating channel configuration and managing stream markers and stream tags."; + - ChannelManageExtensions, scope: "channel:manage:extensions", doc: "Manage a channel’s Extension configuration, including activating Extensions."; + (?) ChannelManagePolls, scope: "channel:manage:polls", doc: "Manage a channel’s polls."; + (?) ChannelManagePredictions, scope: "channel:manage:predictions", doc: "Manage of channel’s Channel Points Predictions"; + - ChannelManageRedemptions, scope: "channel:manage:redemptions", doc: "Manage Channel Points custom rewards and their redemptions on a channel."; + - ChannelManageSchedule, scope: "channel:manage:schedule", doc: "Manage a channel’s stream schedule."; + - ChannelManageVideos, scope: "channel:manage:videos", doc: "Manage a channel’s videos, including deleting videos."; + - ChannelModerate, scope: "channel:moderate", doc: "Perform moderation actions in a channel. The user requesting the scope must be a moderator in the channel."; + - ChannelReadEditors, scope: "channel:read:editors", doc: "View a list of users with the editor role for a channel."; + - ChannelReadGoals, scope: "channel:read:goals", doc: "View Creator Goals for a channel."; + - ChannelReadHypeTrain, scope: "channel:read:hype_train", doc: "View Hype Train information for a channel."; + (?) ChannelReadPolls, scope: "channel:read:polls", doc: "View a channel’s polls."; + (?) ChannelReadPredictions, scope: "channel:read:predictions", doc: "View a channel’s Channel Points Predictions."; + + ChannelReadRedemptions, scope: "channel:read:redemptions", doc: "View Channel Points custom rewards and their redemptions on a channel."; + - ChannelReadStreamKey, scope: "channel:read:stream_key", doc: "View an authorized user’s stream key."; + + ChannelReadSubscriptions, scope: "channel:read:subscriptions", doc: "View a list of all subscribers to a channel and check if a user is subscribed to a channel."; + - ChannelSubscriptions, scope: "channel_subscriptions", doc: "\\[DEPRECATED\\] Read all subscribers to your channel."; + [YES] ChatEdit, scope: "chat:edit", doc: "Send live stream chat and rooms messages."; + [YES] ChatRead, scope: "chat:read", doc: "View live stream chat and rooms messages."; + (?) ClipsEdit, scope: "clips:edit", doc: "Manage Clips for a channel."; + (?) ModerationRead, scope: "moderation:read", doc: "View a channel’s moderation data including Moderators, Bans, Timeouts, and Automod settings."; + (?) ModeratorManageAutoMod, scope: "moderator:manage:automod", doc: "Manage messages held for review by AutoMod in channels where you are a moderator."; + - UserEdit, scope: "user:edit", doc: "Manage a user object."; + - UserEditBroadcast, scope: "user:edit:broadcast", doc: "Edit your channel's broadcast configuration, including extension configuration. (This scope implies user:read:broadcast capability.)"; + - UserEditFollows, scope: "user:edit:follows", doc: "Edit a user’s follows."; + - UserManageBlockedUsers, scope: "user:manage:blocked_users", doc: "Manage the block list of a user."; + - UserReadBlockedUsers, scope: "user:read:blocked_users", doc: "View the block list of a user."; + - UserReadBroadcast, scope: "user:read:broadcast", doc: "View a user’s broadcasting configuration, including Extension configurations."; + - UserReadEmail, scope: "user:read:email", doc: "Read an authorized user’s email address."; + - UserReadFollows, scope: "user:read:follows", doc: "View the list of channels a user follows."; + (?) UserReadSubscriptions, scope: "user:read:subscriptions", doc: "View if an authorized user is subscribed to specific channels."; + - WhispersEdit, scope: "whispers:edit", doc: "Send whisper messages."; + - WhispersRead, scope: "whispers:read", doc: "View your whisper messages."; + */ + + UserTokenBuilder::new(client_id, client_secret, redirect_url.clone()) + .set_scopes(vec![ + Scope::BitsRead, // Using bits to cheer (cost money, larger streams) + Scope::ChannelReadRedemptions, // Channel points redeemed (free points, smaller streams) + Scope::ChannelReadSubscriptions, // List subscribing (paying) users + Scope::ChatEdit, // IRC send + Scope::ChatRead, // IRC read + ]) + .force_verify(force_verify) +} diff --git a/crates/lib/twitch_common/src/twitch_secrets.rs b/crates/lib/twitch_common/src/twitch_secrets.rs new file mode 100644 index 0000000000..dad0d7eb20 --- /dev/null +++ b/crates/lib/twitch_common/src/twitch_secrets.rs @@ -0,0 +1,29 @@ +use container_common::anyhow_result::AnyhowResult; + +/// Standardized way to read Twitch secrets. +/// These are used in several different apps. +#[derive(Clone, serde::Deserialize)] +pub struct TwitchSecrets { + pub app_client_id: String, + pub app_client_secret: String, +} + +impl TwitchSecrets { + pub fn new( + app_client_id: &str, + app_client_secret: &str, + ) -> Self { + Self { + app_client_id: app_client_id.to_string(), + app_client_secret: app_client_secret.to_string(), + } + } + + pub fn from_env() -> AnyhowResult { + Ok(Self::new( + &easyenv::get_env_string_required("TWITCH_APP_CLIENT_ID")?, + &easyenv::get_env_string_required("TWITCH_APP_CLIENT_SECRET")?, + )) + } +} + diff --git a/crates/lib/twitch_common/src/twitch_user_id.rs b/crates/lib/twitch_common/src/twitch_user_id.rs new file mode 100644 index 0000000000..ae1b83aa41 --- /dev/null +++ b/crates/lib/twitch_common/src/twitch_user_id.rs @@ -0,0 +1,35 @@ +use container_common::anyhow_result::AnyhowResult; + +/// Twitch says to encode these as strings, but the Twitch Rust libraries +/// still deal in integers (uh oh). This centralizes type juggling. +#[derive(Clone)] +pub struct TwitchUserId { + string_user_id: String, + numeric_user_id: u32, +} + +impl TwitchUserId { + + pub fn from_str(twitch_user_id: &str) -> AnyhowResult { + let numeric_id = twitch_user_id.parse::()?; + Ok(Self { + string_user_id: twitch_user_id.to_string(), + numeric_user_id: numeric_id, + }) + } + + pub fn from_numeric(twitch_user_id: u32) -> Self { + Self { + string_user_id: twitch_user_id.to_string(), + numeric_user_id: twitch_user_id, + } + } + + pub fn get_numeric(&self) -> u32 { + self.numeric_user_id + } + + pub fn get_str(&self) -> &str { + &self.string_user_id + } +} diff --git a/crates/lib/user_input_common/Cargo.toml b/crates/lib/user_input_common/Cargo.toml new file mode 100644 index 0000000000..2c1d369bf4 --- /dev/null +++ b/crates/lib/user_input_common/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "user_input_common" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "user_input_common" +path = "src/lib.rs" + +[dependencies] + +# Internal +container_common = { path = "../container_common" } + +# External +anyhow = "1.0.40" +log = "0.4.14" +once_cell = "1.9.0" +regex = "1.5.4" +unicode-segmentation = "1.7.1" diff --git a/crates/lib/user_input_common/src/check_for_slurs.rs b/crates/lib/user_input_common/src/check_for_slurs.rs new file mode 100644 index 0000000000..36f0e87455 --- /dev/null +++ b/crates/lib/user_input_common/src/check_for_slurs.rs @@ -0,0 +1,67 @@ +use crate::BANNED_SLURS; +use once_cell::sync::Lazy; +use std::collections::HashSet; +use crate::latin_alphabet::latin_to_ascii; + +static BANNED_SLURS_SET : Lazy> = Lazy::new(|| { + BANNED_SLURS.lines() + .map(|line| line.trim()) + .filter(|line| !(line.starts_with("#") || line.is_empty())) + .map(|line| line.to_string()) + .collect::>() +}); + +pub fn contains_slurs(unparsed_text: &str) -> bool { + let simplified = latin_to_ascii(unparsed_text).to_lowercase(); + for wordlike in simplified.split_ascii_whitespace() { + if BANNED_SLURS_SET.contains(wordlike) { + return true; + } + } + false +} + +#[cfg(test)] +mod tests { + use crate::check_for_slurs::contains_slurs; + + #[test] + fn valid_text_passes() { + assert_eq!(contains_slurs(""), false); + assert_eq!(contains_slurs("this is a test."), false); + assert_eq!(contains_slurs("this\nis\na\ntest\n\n"), false); + assert_eq!(contains_slurs(" this is a test"), false); + } + + #[test] + fn text_with_slurs_fails() { + assert_eq!(contains_slurs("a sentence containing fag is banned"), true); + assert_eq!(contains_slurs("a\nsentence\ncontaining fags\nis banned"), true); + } + + #[test] + fn text_with_mixed_case_slurs_fails() { + assert_eq!(contains_slurs("FAG"), true); + assert_eq!(contains_slurs("FaG"), true); + assert_eq!(contains_slurs("fAg"), true); + assert_eq!(contains_slurs("A SENTENCE CONTAINING FAG IS BANNED"), true); + assert_eq!(contains_slurs("a\nsentence\ncontaining FAGS\nis banned"), true); + } + + #[test] + fn text_with_latin_obfuscated_slurs_fails() { + assert_eq!(contains_slurs("FÀG"), true); + assert_eq!(contains_slurs("FÁG"), true); + assert_eq!(contains_slurs("FÂG"), true); + assert_eq!(contains_slurs("FÃG"), true); + assert_eq!(contains_slurs("FÄG"), true); + assert_eq!(contains_slurs("FÅG"), true); + + assert_eq!(contains_slurs("fàg"), true); + assert_eq!(contains_slurs("fág"), true); + assert_eq!(contains_slurs("fâg"), true); + assert_eq!(contains_slurs("fãg"), true); + assert_eq!(contains_slurs("fäg"), true); + assert_eq!(contains_slurs("fåg"), true); + } +} diff --git a/crates/lib/user_input_common/src/latin_alphabet.rs b/crates/lib/user_input_common/src/latin_alphabet.rs new file mode 100644 index 0000000000..efcd7d3172 --- /dev/null +++ b/crates/lib/user_input_common/src/latin_alphabet.rs @@ -0,0 +1,202 @@ +use once_cell::sync::Lazy; +use std::collections::HashMap; +use unicode_segmentation::UnicodeSegmentation; + +fn to_owned(item: &(&str, &str)) -> (String, String) { + (item.0.to_string(), item.1.to_string()) +} + +// http://www.geocities.ws/click2speak/unicode/chars_es.html +pub static LATIN_TO_ASCII_CHARACTER_MAP : Lazy> = Lazy::new(|| { + // NB: Note that certain latin characters map to multiple ascii characters, eg. "AE". + let table = [ + ("\u{AA}", "a"), // Feminine ordinal + ("\u{BA}", "o"), // Masculine ordinal + ("\u{C0}", "A"), // Latin capital letter A with grave + ("\u{C1}", "A"), // Latin capital letter A with acute + ("\u{C2}", "A"), // Latin capital letter A with circumflex + ("\u{C3}", "A"), // Latin capital letter A with tilde + ("\u{C4}", "A"), // Latin capital letter A diaeresis + ("\u{C5}", "A"), // Latin capital letter A with ring above + ("\u{C6}", "AE"), // Latin capital letter AE + ("\u{C7}", "C"), // Latin capital letter C with cedilla + ("\u{C8}", "E"), // Latin capital letter E with grave + ("\u{C9}", "E"), // Latin capital letter E with acute + ("\u{CA}", "E"), // Latin capital letter E with circumflex + ("\u{CB}", "E"), // Latin capital letter E with diaeresis + ("\u{CC}", "I"), // Latin capital letter I with grave + ("\u{CD}", "I"), // Latin capital letter I with acute + ("\u{CE}", "I"), // Latin capital letter I with circumflex + ("\u{CF}", "I"), // Latin capital letter I with diaeresis + ("\u{D0}", "D"), // Latin capital letter ETH (~D) TODO: How to handle this? + ("\u{D1}", "N"), // Latin capital letter N with tilde + ("\u{D2}", "O"), // Latin capital letter O with grave + ("\u{D3}", "O"), // Latin capital letter O with acute + ("\u{D4}", "O"), // Latin capital letter O with circumflex + ("\u{D5}", "O"), // Latin capital letter O with tilde + ("\u{D6}", "O"), // Latin capital letter O with diaeresis + ("\u{D8}", "O"), // Latin capital letter O with stroke (NB: Skips D7 multiplication sign) + ("\u{D9}", "U"), // Latin capital letter U with grave + ("\u{DA}", "U"), // Latin capital letter U with acute + ("\u{DB}", "U"), // Latin capital letter U with circumflex + ("\u{DC}", "U"), // Latin capital letter U with diaeresis + ("\u{DD}", "Y"), // Latin capital letter Y with acute + ("\u{DE}", "P"), // Latin capital letter THORN (~P) TODO: How to handle this? + ("\u{DF}", "b"), // Latin small letter sharp s (~B) TODO: How to handle this? + ("\u{E0}", "a"), // Latin small letter a with grave + ("\u{E1}", "a"), // Latin small letter a with acute + ("\u{E2}", "a"), // Latin small letter a with circumflex + ("\u{E3}", "a"), // Latin small letter a with tilde + ("\u{E4}", "a"), // Latin small letter a with diaeresis + ("\u{E5}", "a"), // Latin small letter a with ring above + ("\u{E6}", "ae"), // Latin small letter ae + ("\u{E7}", "c"), // Latin small letter c with cedilla + ("\u{E8}", "e"), // Latin small letter e with grave + ("\u{E9}", "e"), // Latin small letter e with acute + ("\u{EA}", "e"), // Latin small letter e with circumflex + ("\u{EB}", "e"), // Latin small letter e with diaeresis + ("\u{EC}", "i"), // Latin small letter i with grave + ("\u{ED}", "i"), // Latin small letter i with acute + ("\u{EE}", "i"), // Latin small letter i with circumflex + ("\u{EF}", "i"), // Latin small letter i with diaeresis + ("\u{F0}", "d"), // Latin capital letter eth (~d) TODO: How to handle this? + ("\u{F1}", "n"), // Lower n tilde + ("\u{F2}", "o"), // Latin small letter o with grave + ("\u{F3}", "o"), // Latin small letter o with acute + ("\u{F4}", "o"), // Latin small letter o with circumflex + ("\u{F5}", "o"), // Latin small letter o with tilde + ("\u{F6}", "o"), // Latin small letter o with diaeresis + ("\u{F8}", "o"), // Latin small letter o with stroke (NB: Skips F7 division sign) + ("\u{F9}", "u"), // Latin small letter u with grave + ("\u{FA}", "u"), // Latin small letter u with acute + ("\u{FB}", "u"), // Latin small letter u with circumflex + ("\u{FC}", "u"), // Latin small letter u with diaeresis + ("\u{FD}", "y"), // Latin small letter y with acute + ("\u{FE}", "p"), // Latin small letter thorn (~p) TODO: How to handle this? + ("\u{FF}", "y"), // Latin small letter y with diaeresis + + // TODO: There are a lot more to go... + ]; + + let mut map = HashMap::new(); + + map.extend(table.iter().map(|item| (item.0.to_string(), item.1.to_string()))); + + // Latin Extended-A + // https://unicode-table.com/en/blocks/latin-extended-a/ + map.extend([ + ("\u{0100}", "A"), // Ā Latin Capital Letter a with Macron + ("\u{0101}", "a"), // ā Latin Small Letter a with Macron + ("\u{0102}", "A"), // Ă Latin Capital Letter a with Breve + ("\u{0103}", "a"), // ă Latin Small Letter a with Breve + ("\u{0104}", "A"), // Ą Latin Capital Letter a with Ogonek + ("\u{0105}", "a"), // ą Latin Small Letter a with Ogonek + ("\u{0106}", "C"), // Ć Latin Capital Letter C with Acute + ("\u{0107}", "c"), // ć Latin Small Letter C with Acute + ("\u{0108}", "C"), // Ĉ Latin Capital Letter C with Circumflex + ("\u{0109}", "c"), // ĉ Latin Small Letter C with Circumflex + ("\u{010A}", "C"), // Ċ Latin Capital Letter C with Dot Above + ("\u{010B}", "c"), // ċ Latin Small Letter C with Dot Above + ("\u{010C}", "C"), // Č Latin Capital Letter C with Caron + ("\u{010D}", "c"), // č Latin Small Letter C with Caron + ("\u{010E}", "D"), // Ď Latin Capital Letter D with Caron + ("\u{010F}", "d"), // ď Latin Small Letter D with Caron + // TODO: Not done... + ("\u{0113}", "e"), // ē Latin Small Letter E with Macron + ("\u{015B}", "s"), // ś Latin Small Letter S with Acute + ("\u{0161}", "s"), // š Latin Small Letter S with Caron + ("\u{017C}", "z"), // ż Latin Small Letter Z with Dot Above + // TODO: Incomplete + ].iter().map(&to_owned)); + + // Latin Extended-B + // https://unicode-table.com/en/blocks/latin-extended-b/ + map.extend([ + ("\u{01B0}", "u"), // ư Latin Small Letter U with Horn + ("\u{01B1}", "U"), // Ʊ Latin Capital Letter Upsilon (TODO: Not correct handling) + ("\u{01B3}", "Y"), // Ƴ Latin Capital Letter Y with Hook + ("\u{01B4}", "y"), // ƴ Latin Small Letter Y with Hook + ("\u{01B5}", "Z"), // Ƶ Latin Capital Letter Z with Stroke + // TODO: Incomplete... + ("\u{01D0}", "i"), // ǐ Latin Small Letter I with Caron + ("\u{01D1}", "O"), // Ǒ Latin Capital Letter O with Caron + ("\u{01D2}", "o"), // ǒ Latin Small Letter O with Caron + ("\u{01D3}", "U"), // Ǔ Latin Capital Letter U with Caron + ("\u{01D4}", "u"), // ǔ Latin Small Letter U with Caron + ].iter().map(&to_owned)); + + // Greek and Coptic + // https://unicode-table.com/en/blocks/greek-coptic/ + map.extend([ + ("\u{03BD}", "v"), // ν Greek Small Letter Nu (TODO: Incorrect handling, but very frequent.) + ("\u{03BF}", "o"), // ο Greek Small Letter Omicron (TODO: Incorrect handling, but very frequent.) + // TODO: Incomplete + ].iter().map(&to_owned)); + + // Latin Extended Additional + // https://unicode-table.com/en/blocks/latin-extended-additional/ + map.extend([ + ("\u{1EA1}", "a"), // ạ Latin Small Letter a with Dot Below + ("\u{1EB9}", "e"), // ẹ Latin Small Letter E with Dot Below + // TODO: Incomplete + ].iter().map(&to_owned)); + + // Misc characters that frequently occur + map.insert("\u{00F3}".to_string(), "o".to_string()); // Latin Small Letter O with Acute + map.insert("\u{0131}".to_string(), "i".to_string()); // Latin Small Letter Dotless I + map.insert("\u{00FC}".to_string(), "u".to_string()); // Latin Small Letter U with Diaeresis + map.insert("\u{015F}".to_string(), "s".to_string()); // Latin Small Letter S with Cedilla + map.insert("\u{00F6}".to_string(), "o".to_string()); // Latin Small Letter O with Diaeresis + map.insert("\u{011F}".to_string(), "g".to_string()); // Latin Small Letter G with Breve + map.insert("\u{0130}".to_string(), "I".to_string()); // Latin Capital Letter I with Dot Above + map.insert("\u{0259}".to_string(), "e".to_string()); // Latin Small Letter Schwa (TODO: Inaccurate?) + map.insert("\u{012B}".to_string(), "i".to_string()); // Latin Small Letter I with Macron + map.insert("\u{1D3A}".to_string(), "n".to_string()); // Modifier Letter Capital N + map.insert("\u{01D0}".to_string(), "i".to_string()); // Latin Small Letter I with Caron + map.insert("\u{014D}".to_string(), "o".to_string()); // Latin Small Letter O with Macron + map.insert("\u{0101}".to_string(), "a".to_string()); // Latin Small Letter A with Macron + map.insert("\u{00F9}".to_string(), "u".to_string()); // Latin Small Letter U with Grave + map.insert("\u{01Ce}".to_string(), "a".to_string()); // Latin Small Letter A with Caron + map +}); + +pub fn latin_to_ascii(input_text: &str) -> String { + let segmented = UnicodeSegmentation::graphemes(input_text, true) + .map(|segment| { + LATIN_TO_ASCII_CHARACTER_MAP.get(segment) + .map(|replace| replace.as_str()) + .unwrap_or(segment) + }) + .collect::>(); + + segmented.join("") +} + +#[cfg(test)] +mod tests { + use crate::latin_alphabet::latin_to_ascii; + + #[test] + fn test_latin_to_ascii() { + assert_eq!(latin_to_ascii("pokémon"), "pokemon".to_string()); + assert_eq!(latin_to_ascii("POKÉMON"), "POKEMON".to_string()); + assert_eq!(latin_to_ascii("Æther"), "AEther".to_string()); + assert_eq!(latin_to_ascii("æther"), "aether".to_string()); + // Almost exhaustive + assert_eq!(latin_to_ascii("ÀÁÂÃÄÅ"), "AAAAAA".to_string()); + assert_eq!(latin_to_ascii("Æ"), "AE".to_string()); + assert_eq!(latin_to_ascii("Ç"), "C".to_string()); + assert_eq!(latin_to_ascii("ÈÉÊË"), "EEEE".to_string()); + assert_eq!(latin_to_ascii("ÌÍÎÏ"), "IIII".to_string()); + assert_eq!(latin_to_ascii("ÒÓÔÕÖØ"), "OOOOOO".to_string()); + assert_eq!(latin_to_ascii("ÙÚÛÜ"), "UUUU".to_string()); + assert_eq!(latin_to_ascii("Ý"), "Y".to_string()); + assert_eq!(latin_to_ascii("àáâãäå"), "aaaaaa".to_string()); + assert_eq!(latin_to_ascii("æ"), "ae".to_string()); + assert_eq!(latin_to_ascii("ç"), "c".to_string()); + assert_eq!(latin_to_ascii("èéêë"), "eeee".to_string()); + assert_eq!(latin_to_ascii("ìíîï"), "iiii".to_string()); + assert_eq!(latin_to_ascii("òóôõöø"), "oooooo".to_string()); + assert_eq!(latin_to_ascii("ùúûü"), "uuuu".to_string()); + } +} diff --git a/crates/lib/user_input_common/src/lib.rs b/crates/lib/user_input_common/src/lib.rs new file mode 100644 index 0000000000..4439295579 --- /dev/null +++ b/crates/lib/user_input_common/src/lib.rs @@ -0,0 +1,13 @@ +#![forbid(dead_code)] +#![forbid(non_snake_case)] +#![forbid(private_in_public)] +#![forbid(unused_imports)] +#![forbid(unused_must_use)] // Important +#![forbid(unused_mut)] +#![forbid(unused_variables)] +#![forbid(warnings)] + +pub (crate) const BANNED_SLURS : &'static str = include_str!("../../../../db/banned_slurs.txt"); + +pub mod check_for_slurs; +pub mod latin_alphabet; diff --git a/crates/service/storyteller_web/Cargo.toml b/crates/service/storyteller_web/Cargo.toml new file mode 100644 index 0000000000..5f5c91b8fa --- /dev/null +++ b/crates/service/storyteller_web/Cargo.toml @@ -0,0 +1,151 @@ +[package] +name = "storyteller-web" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +# Main server monolith +[[bin]] +name = "storyteller-web" +path = "src/main.rs" + +## OBS websocket server +#[[bin]] +#name = "obs-gateway-server" +#path = "src/obs_gateway_server.rs" + +# Jobs for managing downloads and inference +[[bin]] +name = "tts-download-job" +path = "src/tts_download_job.rs" + +[[bin]] +name = "tts-inference-job" +path = "src/tts_inference_job.rs" + +[[bin]] +name = "w2l-download-job" +path = "src/w2l_download_job.rs" + +[[bin]] +name = "w2l-inference-job" +path = "src/w2l_inference_job.rs" + +[dependencies] + +# Internal +config = { path = "../../lib/config" } +container_common = { path = "../../lib/container_common" } +database_queries = { path = "../../lib/database_queries" } +http_server_common = { path = "../../lib/http_server_common" } +redis_common = { path = "../../lib/redis_common" } +tts_common = { path = "../../lib/tts_common" } +twitch_common = { path = "../../lib/twitch_common" } +user_input_common = { path = "../../lib/user_input_common" } + +# Vendored code +newrelic-telemetry = { path = "../../vendor/newrelic-telemetry-sdk-rust" } + +# External +anyhow = "1.0.40" +base64 = "0.13.0" +chrono = { version = "0.4.19", features = ["serde"] } +derive_more = "0.99.14" +dotenv = "0.15.0" +easyenv = "0.0.6" +hex = "0.4.3" +hostname = "0.3.1" +lazy_static = "1.4.0" +lexical-sort = "0.3.1" +lfu = "0.2.5" +log = "0.4.14" +rand = "0.8.3" +regex = "1.5.4" +time = "0.2.26" +uuid = { version = "0.8.2", features = ["v4"] } + +# TODO/TEMPORARY: +# Twitch +# NB: "unsupported" pubsub features may disappear from the Twitch API at any time. It's an unofficial API. +twitch_api2 = { version = "0.6.0-rc.2", features = ["helix", "reqwest_client", "pubsub", "chrono", "unsupported"] } +twitch_oauth2 = "0.6.0-rc.2" + +# Language tag parsing +# TODO: Move to language support crate +language-tags = "0.3.2" # Parsing single tags richly + +# HTTP +# TODO/NB: This is just for Twitch, but we might need it later. +reqwest = { version = "0.11.4" } +native-tls = "0.2.8" + +# Runtime +# NB: We switched to Tokio from `async-std` for broader ecosystem support. +tokio = { version = "1.6.0", features = ["macros"] } + +# Futures +futures = { version = "0.3.17", features = ["thread-pool"] } # feature for ThreadPool +futures-util = "0.3.9" +futures_old_for_limiter = { package = "futures", version = "0.1.29" } # Older version of futures for `limiter` crate + +# MySql Database +# "offline" supports typesafe builds without a live database +# "rustls" doesn't appear to work with IP addresses, per: +# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828 +sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] } + +# Redis Database +limitation = "0.1.1" # Redis rate limiter +r2d2_redis = "0.14.0" + +# In-memory caching +lru_time_cache = "0.11.11" + +# Files +data-encoding = "2.3.2" +ring = "0.17.0-alpha.10" +rust-s3 = { version = "0.28", features = ["never-encode-slash"] } +infer = "0.5.0" + +subprocess = "0.2.6" +tempdir = "0.3.7" + +actix = "0.12.0" +actix-cors = "0.6.0-beta.3" +actix-files= "0.6.0-beta.8" +actix-http = "3.0.0-beta.11" +actix-multipart = "0.4.0-beta.7" +actix-rt = "2.4.0" +actix-service = "2.0.1" +actix-utils = "3.0.0" +actix-web = "4.0.0-beta.10" +actix-web-actors = "4.0.0-beta.7" +jwt = "0.13.0" + +# NB: Hyper is just for HTTP client use in jobs. I'd use Actix client, but I need tokio runtime +# NB: Furthermore, the Tokio runtime is old, so I have to use an older Hyper. +hyper = { version = "0.14.14" } + +# Crypto +bcrypt = "0.8" # password hashes +hmac = "0.10.0" # cookies +magic-crypt = "3.1.8" # encrypting sort IDs for the frontend so we don't leak info +md-5 = "0.9.1" # gravatar +sha2 = "0.9.5" # file hashes, inference text hashes, etc. + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" +toml = "0.5.8" + +# Text processing +pulldown-cmark = "0.8.0" + +[build-dependencies] +# NB: Only using diesel for migrations support +diesel = { version = "1.4.6", features = ["mysql"] } diff --git a/crates/service/storyteller_web/src/common_env.rs b/crates/service/storyteller_web/src/common_env.rs new file mode 100644 index 0000000000..a2e272762a --- /dev/null +++ b/crates/service/storyteller_web/src/common_env.rs @@ -0,0 +1,56 @@ +use config::shared_constants::DEFAULT_REDIS_DATABASE_0_CONNECTION_STRING; +use config::shared_constants::DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING; +use container_common::anyhow_result::AnyhowResult; + +// TODO: Move more shared configs here. + +/// Environment variables shared between one or more jobs or the main server. +/// These do not have to have the same value for each instance, but they behave +/// similarly across all usages. +pub struct CommonEnv { + + /// The Redis connection string. + /// This has FakeYou voice stats, rate limiters, etc. + pub redis_0_connection_string: String, + + /// The Redis connection string. + /// This has OBS/Twitch related functionality + pub redis_1_connection_string: String, + + /// The amount of time to wait between job batches (not individual jobs). + /// This prevents the outer loop of querying batches from flooding the DB. + /// (In theory, there's work within jobs that prevents rapidly pegging the DB.) + pub job_batch_wait_millis: u64, + + /// The maximum number of retries a job will get. + /// After a job exhausts these attempts, it will become "dead". + pub job_max_attempts: u8, + + /// Number of job records to query at once. + pub job_batch_size: u32, + + /// Time to wait between no-op logger log lines. + pub no_op_logger_millis: u64, + + /// For debugging. + /// If nonzero, sleep before exiting the job for this amount of millis. + /// This allows introspection of temp directories before they are unlinked. + pub debug_job_end_sleep_millis: u64, +} + +impl CommonEnv { + + pub fn read_from_env() -> AnyhowResult { + Ok(Self { + redis_0_connection_string: easyenv::get_env_string_or_default("REDIS_0_URL", + DEFAULT_REDIS_DATABASE_0_CONNECTION_STRING), + redis_1_connection_string: easyenv::get_env_string_or_default("REDIS_1_URL", + DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING), + job_batch_wait_millis: easyenv::get_env_num("JOB_BATCH_WAIT_MILLIS", 100)?, + job_max_attempts: easyenv::get_env_num("JOB_MAX_ATTEMPTS", 3)?, + job_batch_size: easyenv::get_env_num("JOB_BATCH_SIZE", 10)?, + no_op_logger_millis: easyenv::get_env_num("NO_OP_LOGGER_MILLIS", 15_000)?, + debug_job_end_sleep_millis: easyenv::get_env_num("DEBUG_JOB_END_SLEEP_MILLIS", 0)?, + }) + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_clients/mod.rs b/crates/service/storyteller_web/src/http_clients/mod.rs new file mode 100644 index 0000000000..8f7ae3bad7 --- /dev/null +++ b/crates/service/storyteller_web/src/http_clients/mod.rs @@ -0,0 +1 @@ +pub mod tts_inference_sidecar_client; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs b/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs new file mode 100644 index 0000000000..4e990e78b3 --- /dev/null +++ b/crates/service/storyteller_web/src/http_clients/tts_inference_sidecar_client.rs @@ -0,0 +1,143 @@ +use actix_web::http::header; +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use database_queries::column_types::vocoder_type::VocoderType; +use hyper::client::Client; +use hyper::{Body, Request}; +use log::info; +use std::path::Path; + +pub struct TtsInferenceSidecarClient { + hostname: String, +} + +#[derive(Serialize)] +struct InferenceRequest { + // Vocoder information + pub vocoder_type: VocoderType, + pub waveglow_vocoder_checkpoint_path : String, + pub hifigan_vocoder_checkpoint_path: String, + pub hifigan_superres_vocoder_checkpoint_path: String, + + // Synthesizer information + pub synthesizer_checkpoint_path : String, + + // Text + pub inference_text : String, + + // Output information + pub output_audio_filename : String, + pub output_spectrogram_filename : String, + pub output_metadata_filename : String, + + /// To instruct the sidecar to unload the model from memory + pub maybe_clear_synthesizer_checkpoint_path: Option, +} + +impl TtsInferenceSidecarClient { + pub fn new(hostname: &str) -> Self { + //let client = Client::builder() + // .header("User-Agent", "actix/tts_inference_job") + // .finish(); + Self { + hostname: hostname.to_string(), + // client, + } + } + + pub async fn request_inference>( + &self, + raw_text: &str, + synthesizer_checkpoint_path: P, + vocoder_type: VocoderType, + hifigan_vocoder_checkpoint_path: P, + hifigan_superres_vocoder_checkpoint_path: P, + waveglow_vocoder_checkpoint_path: P, + output_audio_filename: P, + output_spectrogram_filename: P, + output_metadata_filename: P, + maybe_unload_model_path: Option, + ) -> AnyhowResult<()> { + + let waveglow_vocoder_checkpoint_path = waveglow_vocoder_checkpoint_path + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad waveglow vocoder path"))?; + + let hifigan_vocoder_checkpoint_path = hifigan_vocoder_checkpoint_path + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad hifigan vocoder path"))?; + + let hifigan_superres_vocoder_checkpoint_path = hifigan_superres_vocoder_checkpoint_path + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad hifigan super resolution vocoder path"))?; + + let synthesizer_checkpoint_path = synthesizer_checkpoint_path + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad synthesizer path"))?; + + let output_audio_filename = output_audio_filename + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad output audio path"))?; + + let output_spectrogram_filename = output_spectrogram_filename + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad output spectrogram path"))?; + + let output_metadata_filename = output_metadata_filename + .as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("bad output metadata path"))?; + + let request = InferenceRequest { + inference_text: raw_text.to_string(), + vocoder_type, + waveglow_vocoder_checkpoint_path, + hifigan_vocoder_checkpoint_path, + hifigan_superres_vocoder_checkpoint_path, + synthesizer_checkpoint_path, + output_audio_filename, + output_spectrogram_filename, + output_metadata_filename, + maybe_clear_synthesizer_checkpoint_path: maybe_unload_model_path, + }; + + let url = format!("http://{}/infer", self.hostname); + info!("Requesting {}", url); + + //let maybe_response = self.client.get(&url) + // .header(header::CONTENT_TYPE, "application/json") + // .send_json(&request) + // .await; + + let request = serde_json::to_string(&request)?; + + let req = Request::builder() + .method(hyper::Method::POST) + .uri(url) + .header("content-type", "application/json") + .body(Body::from(request))?; + + let client = Client::new(); + + let maybe_response = client.request(req).await?; + + //match maybe_response { + // Err(e) => Err(anyhow!("Error talking to sidecar: {:?}", e)), + // Ok(_) => Ok(()), + //} + Ok(()) + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/create_api_token.rs b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/create_api_token.rs new file mode 100644 index 0000000000..7d8169f506 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/create_api_token.rs @@ -0,0 +1,125 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::api_tokens::create_api_token::create_api_token_for_user; +use database_queries::tokens::Tokens; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct CreateApiTokenRequest { + pub idempotency_token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct CreateApiTokenResponse { + pub success: bool, + pub api_token: String, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum CreateApiTokenError { + BadInput(String), + NotAuthorized, + ServerError, +} + +impl ResponseError for CreateApiTokenError { + fn status_code(&self) -> StatusCode { + match *self { + CreateApiTokenError::BadInput(_) => StatusCode::BAD_REQUEST, + CreateApiTokenError::NotAuthorized => StatusCode::UNAUTHORIZED, + CreateApiTokenError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for CreateApiTokenError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn create_api_token_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CreateApiTokenError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(CreateApiTokenError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot create API tokens"); + return Err(CreateApiTokenError::NotAuthorized); + } + + let creator_ip_address = get_request_ip(&http_request); + + let api_token = create_api_token_for_user( + &user_session.user_token, + &request.idempotency_token, + &creator_ip_address, + &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("API token creation errror: {:?}", e); + CreateApiTokenError::ServerError + })?; + + let response = CreateApiTokenResponse { + success: true, + api_token + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CreateApiTokenError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/delete_api_token.rs b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/delete_api_token.rs new file mode 100644 index 0000000000..c5a303d36f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/delete_api_token.rs @@ -0,0 +1,142 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::queries::api_tokens::list_available_api_tokens_for_user::list_available_api_tokens_for_user; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use database_queries::queries::api_tokens::delete_api_token::delete_api_token; + +// =============== Request =============== + +// NB: `api_token` is only in PathInfo because we're deleting it. + +#[derive(Deserialize)] +pub struct DeleteApiTokenPathInfo { + api_token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct DeleteApiTokenResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum DeleteApiTokenError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for DeleteApiTokenError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteApiTokenError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteApiTokenError::NotFound => StatusCode::NOT_FOUND, + DeleteApiTokenError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteApiTokenError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DeleteApiTokenError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn delete_api_token_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteApiTokenError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteApiTokenError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot edit API tokens"); + return Err(DeleteApiTokenError::NotAuthorized); + } + + let tokens = list_available_api_tokens_for_user(&user_session.user_token, + &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Error querying tokens: {:?}", e); + return DeleteApiTokenError::ServerError; + })?; + + let valid_token = tokens.iter() + .find(|t| t.api_token.eq(&path.api_token)) + .is_some(); + + if !valid_token { + warn!("Invalid API Token"); + return Err(DeleteApiTokenError::NotFound); + } + + let creator_ip_address = get_request_ip(&http_request); + + let _r = delete_api_token( + &user_session.user_token, + &path.api_token, + &creator_ip_address, + &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return DeleteApiTokenError::ServerError; + }); + + let response = DeleteApiTokenResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| DeleteApiTokenError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/edit_api_token.rs b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/edit_api_token.rs new file mode 100644 index 0000000000..213cb0a584 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/edit_api_token.rs @@ -0,0 +1,145 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use database_queries::queries::api_tokens::list_available_api_tokens_for_user::list_available_api_tokens_for_user; +use database_queries::queries::api_tokens::edit_api_token::edit_api_token; +use crate::http_server::web_utils::ip_address::get_request_ip; + +// =============== Request =============== + +// NB: Ordinarily, `api_token` would be PathInfo, but these are secret values. + +// These are not sparse updates! +#[derive(Deserialize)] +pub struct EditApiTokenRequest { + pub api_token: String, + pub maybe_short_description: Option, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct EditApiTokenResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum EditApiTokenError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for EditApiTokenError { + fn status_code(&self) -> StatusCode { + match *self { + EditApiTokenError::BadInput(_) => StatusCode::BAD_REQUEST, + EditApiTokenError::NotFound => StatusCode::NOT_FOUND, + EditApiTokenError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditApiTokenError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditApiTokenError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn edit_api_token_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditApiTokenError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditApiTokenError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot edit API tokens"); + return Err(EditApiTokenError::NotAuthorized); + } + + let tokens = list_available_api_tokens_for_user(&user_session.user_token, + &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Error querying tokens: {:?}", e); + return EditApiTokenError::ServerError; + })?; + + let valid_token = tokens.iter() + .find(|t| t.api_token.eq(&request.api_token)) + .is_some(); + + if !valid_token { + warn!("Invalid API Token"); + return Err(EditApiTokenError::NotFound); + } + + let creator_ip_address = get_request_ip(&http_request); + + let _r = edit_api_token( + &user_session.user_token, + &request.api_token, + request.maybe_short_description.as_deref(), + &creator_ip_address, + &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return EditApiTokenError::ServerError; + }); + + let response = EditApiTokenResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| EditApiTokenError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/list_api_tokens.rs b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/list_api_tokens.rs new file mode 100644 index 0000000000..b9bb38519f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/list_api_tokens.rs @@ -0,0 +1,130 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::api_tokens::list_available_api_tokens_for_user::list_available_api_tokens_for_user; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use lexical_sort::natural_lexical_cmp; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct ListApiTokensResponse { + pub success: bool, + pub api_tokens: Vec, +} + +/// Public-facing and optimized (non-irrelevant fields) +#[derive(Serialize)] +pub struct ApiToken { + pub api_token: String, + pub maybe_short_description: Option, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum ListApiTokensError { + NotAuthorized, + ServerError, +} + +impl ResponseError for ListApiTokensError { + fn status_code(&self) -> StatusCode { + match *self { + ListApiTokensError::NotAuthorized => StatusCode::UNAUTHORIZED, + ListApiTokensError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListApiTokensError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn list_api_tokens_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListApiTokensError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListApiTokensError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot use API tokens"); + return Err(ListApiTokensError::NotAuthorized); + } + + let api_tokens = list_available_api_tokens_for_user(&user_session.user_token, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("API token query error: {:?}", e); + ListApiTokensError::ServerError + })?; + + let mut api_tokens = api_tokens.into_iter() + .map(|r| { + ApiToken { + api_token: r.api_token, + maybe_short_description: r.maybe_short_description, + created_at: r.created_at, + updated_at: r.updated_at, + } + }) + .collect::>(); + + let response = ListApiTokensResponse { + success: true, + api_tokens, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListApiTokensError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/mod.rs new file mode 100644 index 0000000000..e3c4f3427d --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/api_tokens/mod.rs @@ -0,0 +1,4 @@ +pub mod create_api_token; +pub mod delete_api_token; +pub mod edit_api_token; +pub mod list_api_tokens; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/categories/assign_tts_category.rs b/crates/service/storyteller_web/src/http_server/endpoints/categories/assign_tts_category.rs new file mode 100644 index 0000000000..015ba9fa10 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/categories/assign_tts_category.rs @@ -0,0 +1,268 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use anyhow::anyhow; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::model_categories::get_category_by_token::get_category_by_token; +use database_queries::queries::tts::tts_models::get_tts_model::get_tts_model_by_token; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct AssignTtsCategoryRequest { + category_token: String, + tts_model_token: String, + + /// Whether to assign or delete the category. + assign: bool, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct AssignTtsCategoryResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum AssignTtsCategoryError { + CategoryNotFound, + CategoryNotApplicable, + ModelNotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for AssignTtsCategoryError { + fn status_code(&self) -> StatusCode { + match *self { + AssignTtsCategoryError::CategoryNotFound => StatusCode::NOT_FOUND, + AssignTtsCategoryError::CategoryNotApplicable => StatusCode::BAD_REQUEST, + AssignTtsCategoryError::ModelNotFound => StatusCode::NOT_FOUND, + AssignTtsCategoryError::NotAuthorized => StatusCode::UNAUTHORIZED, + AssignTtsCategoryError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for AssignTtsCategoryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn assign_tts_category_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + AssignTtsCategoryError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(AssignTtsCategoryError::NotAuthorized); + } + }; + + let category_lookup_result + = get_category_by_token(&request.category_token, &server_state.mysql_pool).await; + + let category = match category_lookup_result { + Ok(Some(result)) => { + result + }, + Ok(None) => { + warn!("could not find category"); + return Err(AssignTtsCategoryError::CategoryNotFound); + }, + Err(err) => { + warn!("error looking up category: {:?}", err); + return Err(AssignTtsCategoryError::ServerError); + }, + }; + + // NB: First permission check. + // TODO: We don't have proper category permissions yet, so this is a proxy. + let is_category_mod= user_session.can_delete_other_users_tts_models; + + // If category is exclusively for mods + if !is_category_mod { + if category.can_only_mods_apply { + warn!("user is not allowed to assign this category: {}", user_session.user_token); + return Err(AssignTtsCategoryError::NotAuthorized); + } + + if category.deleted_at.is_some() || + !category.is_mod_approved.unwrap_or(false) { + warn!("user is not allowed to see this category: {}", user_session.user_token); + return Err(AssignTtsCategoryError::ModelNotFound); + } + } + + if request.assign && !category.can_directly_have_models { + warn!("category cannot have models: {}", category.category_token); + return Err(AssignTtsCategoryError::CategoryNotApplicable); + } + + // NB: Second set of permission checks. + // Only mods should see deleted models (both user_* and mod_* deleted). + let is_mod_that_can_see_deleted = user_session.can_delete_other_users_tts_models; + + let model_lookup_result = get_tts_model_by_token( + &request.tts_model_token, + is_mod_that_can_see_deleted, + &server_state.mysql_pool).await; + + let model_record = match model_lookup_result { + Ok(Some(result)) => { + result + }, + Ok(None) => { + warn!("could not find model"); + return Err(AssignTtsCategoryError::ModelNotFound); + }, + Err(err) => { + warn!("error looking up model: {:?}", err); + return Err(AssignTtsCategoryError::ServerError); + }, + }; + + // NB: Third set of permission checks + let is_author = model_record.creator_user_token == user_session.user_token; + let is_mod = user_session.can_edit_other_users_tts_models ; + + if !is_author && !is_mod { + warn!("user is not allowed to add categories to model: {}", user_session.user_token); + return Err(AssignTtsCategoryError::NotAuthorized); + } + + if !is_mod { + if model_record.is_locked_from_user_modification || model_record.is_locked_from_use { + return Err(AssignTtsCategoryError::NotAuthorized); + } + } + + let ip_address = get_request_ip(&http_request); + + let query_builder = if request.assign { + // NB: deleted_at = NULL + sqlx::query!(r#" + INSERT INTO tts_category_assignments + SET + model_token = ?, + category_token = ?, + category_addition_user_token = ?, + ip_address_creation = ?, + ip_address_last_update = ?, + deleted_at = NULL + ON DUPLICATE KEY UPDATE + model_token = ?, + category_token = ?, + category_addition_user_token = ?, + ip_address_last_update = ?, + deleted_at = NULL + "#, + // Insert + &request.tts_model_token, + &request.category_token, + &user_session.user_token, + &ip_address, + &ip_address, + // Update + &request.tts_model_token, + &request.category_token, + &user_session.user_token, + &ip_address, + ) + + } else { + // NB: deleted_at = CURRENT_TIMESTAMP + sqlx::query!(r#" + INSERT INTO tts_category_assignments + SET + model_token = ?, + category_token = ?, + category_removal_user_token = ?, + ip_address_creation = ?, + ip_address_last_update = ?, + deleted_at = CURRENT_TIMESTAMP + ON DUPLICATE KEY UPDATE + model_token = ?, + category_token = ?, + category_removal_user_token = ?, + ip_address_last_update = ?, + deleted_at = CURRENT_TIMESTAMP + "#, + // Insert + &request.tts_model_token, + &request.category_token, + &user_session.user_token, + &ip_address, + &ip_address, + // Update + &request.tts_model_token, + &request.category_token, + &user_session.user_token, + &ip_address, + ) + }; + + // NB: We're soft deleting so we don't delete the associations. + let query_result = query_builder.execute(&server_state.mysql_pool).await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Assign category edit DB error: {:?}", err); + return Err(AssignTtsCategoryError::ServerError); + } + }; + + let response = AssignTtsCategoryResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| AssignTtsCategoryError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/categories/create_category.rs b/crates/service/storyteller_web/src/http_server/endpoints/categories/create_category.rs new file mode 100644 index 0000000000..bb2f9f8d54 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/categories/create_category.rs @@ -0,0 +1,217 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::tokens::Tokens; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +const DEFAULT_CAN_DIRECTLY_HAVE_MODELS : bool = true; +const DEFAULT_CAN_HAVE_SUBCATEGORIES : bool = false; +const DEFAULT_CAN_ONLY_MODS_APPLY : bool = false; + +// =============== Request =============== + +#[derive(Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum ModelType { + /// 'tts' in database + Tts, + /// 'w2l' in database + W2l, +} + +#[derive(Deserialize)] +pub struct CreateCategoryRequest { + // Fields for everyone + pub name: Option, + pub model_type: Option, + pub idempotency_token: Option, + + // Fields for moderators only + pub can_directly_have_models: Option, + pub can_have_subcategories: Option, + pub can_only_mods_apply: Option, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct CreateCategoryResponse { + pub success: bool, + pub token: Option, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum CreateCategoryError { + BadInput(String), + NotAuthorized, + ServerError, +} + +impl ResponseError for CreateCategoryError { + fn status_code(&self) -> StatusCode { + match *self { + CreateCategoryError::BadInput(_) => StatusCode::BAD_REQUEST, + CreateCategoryError::NotAuthorized => StatusCode::UNAUTHORIZED, + CreateCategoryError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for CreateCategoryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn create_category_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CreateCategoryError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(CreateCategoryError::NotAuthorized); + } + }; + + // NB: First permission check. + // TODO: We don't have a permission for categories, so we use this as a proxy. + let is_mod = user_session.can_ban_users; + + let model_type = match request.model_type { + None => { + return Err(CreateCategoryError::BadInput("no model type".to_string())); + } + Some(ModelType::Tts) => "tts", + Some(ModelType::W2l) => "w2l", + }; + + let idempotency_token = request.idempotency_token + .clone() + .ok_or(CreateCategoryError::BadInput("no idempotency token provided".to_string()))?; + + let name = request.name + .clone() + .ok_or(CreateCategoryError::BadInput("no name provided".to_string()))?; + + let category_token = Tokens::new_category() + .map_err(|e| { + warn!("Bad crockford token: {:?}", e); + CreateCategoryError::ServerError + })?; + + let creator_ip_address = get_request_ip(&http_request); + + let mut is_mod_approved = false; + let mut maybe_mod_user_token = None; + let mut can_directly_have_models = DEFAULT_CAN_DIRECTLY_HAVE_MODELS; + let mut can_have_subcategories = DEFAULT_CAN_HAVE_SUBCATEGORIES; + let mut can_only_mods_apply = DEFAULT_CAN_ONLY_MODS_APPLY; + + if is_mod { + // Moderator fields and adjustments + is_mod_approved = true; + maybe_mod_user_token = Some(user_session.user_token.clone()); + can_directly_have_models = request.can_directly_have_models + .unwrap_or(DEFAULT_CAN_DIRECTLY_HAVE_MODELS); + can_have_subcategories = request.can_have_subcategories + .unwrap_or(DEFAULT_CAN_HAVE_SUBCATEGORIES); + can_only_mods_apply = request.can_only_mods_apply + .unwrap_or(DEFAULT_CAN_ONLY_MODS_APPLY); + } + + let query_result = sqlx::query!( + r#" +INSERT INTO model_categories +SET + token = ?, + uuid_idempotency_token = ?, + model_type = ?, + name = ?, + + creator_user_token = ?, + creator_ip_address_creation = ?, + creator_ip_address_last_update = ?, + + is_mod_approved = ?, + maybe_mod_user_token = ?, + can_directly_have_models = ?, + can_have_subcategories = ?, + can_only_mods_apply = ? + "#, + + category_token, + idempotency_token, + model_type, + name, + &user_session.user_token, + &creator_ip_address, + &creator_ip_address, + is_mod_approved, + maybe_mod_user_token, + can_directly_have_models, + can_have_subcategories, + can_only_mods_apply + ) + .execute(&server_state.mysql_pool) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Create category edit DB error: {:?}", err); + return Err(CreateCategoryError::ServerError); + } + }; + + let response = CreateCategoryResponse { + success: true, + token: Some(category_token.to_string()) + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CreateCategoryError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/categories/get_category.rs b/crates/service/storyteller_web/src/http_server/endpoints/categories/get_category.rs new file mode 100644 index 0000000000..00fea85139 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/categories/get_category.rs @@ -0,0 +1,175 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::model_categories::get_category_by_token::get_category_by_token; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetCategoryPathInfo { + /// TTS model token + token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct GetCategoryResponse { + pub success: bool, + pub category: DisplayCategory, +} + +/// Public-facing and optimized (non-irrelevant fields) category list +/// Used for the main TTS dropdown as well as the TTS edit/tagging UI +#[derive(Serialize)] +pub struct DisplayCategory { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub is_mod_approved: Option, + + /// Absent for non-mods + pub maybe_mod_comments: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum GetCategoryError { + NotFound, + ServerError, +} + +impl ResponseError for GetCategoryError { + fn status_code(&self) -> StatusCode { + match *self { + GetCategoryError::NotFound => StatusCode::NOT_FOUND, + GetCategoryError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetCategoryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn get_category_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetCategoryError::ServerError + })?; + + // TODO: We don't have any permissions for categories. This is a proxy. + let is_mod = maybe_user_session + .map(|session| session.can_ban_users) + .unwrap_or(false); + + let category_lookup_result + = get_category_by_token(&path.token, &server_state.mysql_pool).await; + + let category = match category_lookup_result { + Ok(Some(result)) => { + result + }, + Ok(None) => { + warn!("could not find category"); + return Err(GetCategoryError::NotFound); + }, + Err(err) => { + warn!("error looking up category: {:?}", err); + return Err(GetCategoryError::ServerError); + }, + }; + + let category_for_response = DisplayCategory { + category_token: category.category_token.clone(), + model_type: category.model_type.clone(), + maybe_super_category_token: category.maybe_super_category_token.clone(), + can_directly_have_models: category.can_directly_have_models, + can_have_subcategories: category.can_have_subcategories, + can_only_mods_apply: category.can_only_mods_apply, + name: category.name.clone(), + maybe_dropdown_name: category.maybe_dropdown_name.clone(), + creator_user_token: category.creator_user_token.clone(), + creator_username: category.creator_username.clone(), + creator_display_name: category.creator_display_name.clone(), + creator_gravatar_hash: category.creator_gravatar_hash.clone(), + is_mod_approved: category.is_mod_approved.clone(), + created_at: category.created_at.clone(), + updated_at: category.updated_at.clone(), + deleted_at: category.deleted_at.clone(), + // Moderator fields + // Clear out fields for non-mods + maybe_mod_comments: if is_mod { category.maybe_mod_comments.clone() } else { None }, + }; + + let response = GetCategoryResponse { + success: true, + category: category_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetCategoryError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/categories/list_tts_categories.rs b/crates/service/storyteller_web/src/http_server/endpoints/categories/list_tts_categories.rs new file mode 100644 index 0000000000..5ccf74d83f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/categories/list_tts_categories.rs @@ -0,0 +1,178 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::model_categories::list_categories_query_builder::ListCategoriesQueryBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use lexical_sort::natural_lexical_cmp; +use log::{info, warn, log, error}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct ListTtsCategoriesResponse { + pub success: bool, + pub categories: Vec, +} + +/// Public-facing and optimized (non-irrelevant fields) category list +/// Used for the main TTS dropdown as well as the TTS edit/tagging UI +#[derive(Serialize, Clone)] +pub struct DisplayCategory { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + /// Instead of making the frontend deal with "maybe_dropdown_name", always use this: + pub name_for_dropdown: String, + + // Moderator fields + // It's okay to leak this since we do for assigned categories and for mods to see + // which assigned categories might be invalid. + pub is_mod_approved: Option, + + //pub creator_user_token: Option, + //pub creator_username: Option, + //pub creator_display_name: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum ListTtsCategoriesError { + NotAuthorized, + ServerError, +} + +impl ResponseError for ListTtsCategoriesError { + fn status_code(&self) -> StatusCode { + match *self { + ListTtsCategoriesError::NotAuthorized => StatusCode::UNAUTHORIZED, + ListTtsCategoriesError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTtsCategoriesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn list_tts_categories_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let maybe_categories = server_state.caches.category_list.copy_without_bump_if_unexpired() + .map_err(|e| { + error!("Error consulting cache: {:?}", e); + ListTtsCategoriesError::ServerError + })?; + + let categories = match maybe_categories { + Some(categories) => { + info!("Serving TTS categories from cache"); + categories + }, + None => { + let query_builder = ListCategoriesQueryBuilder::new() + .show_deleted(false) + .show_unapproved(false) + .scope_model_type(Some("tts")); + + let query_result = + query_builder.perform_query(&server_state.mysql_pool).await; + + let results = match query_result { + Ok(results) => results, + Err(err) => { + warn!("DB error: {:?}", err); + return Err(ListTtsCategoriesError::ServerError); + } + }; + + warn!("Number of categories: {:?}", results.categories.len()); + + let mut categories = results.categories.iter() + .map(|c| { + DisplayCategory { + category_token: c.category_token.clone(), + model_type: c.model_type.clone(), + maybe_super_category_token: c.maybe_super_category_token.clone(), + can_directly_have_models: c.can_directly_have_models, + can_have_subcategories: c.can_have_subcategories, + can_only_mods_apply: c.can_only_mods_apply, + is_mod_approved: c.is_mod_approved, + name: c.name.clone(), + name_for_dropdown: c.maybe_dropdown_name.clone().unwrap_or(c.name.clone()), + created_at: c.created_at.clone(), + updated_at: c.updated_at.clone(), + deleted_at: c.deleted_at.clone(), + } + }) + .collect::>(); + + // NB: This might produce weird sorting resorts relative to the "name" field, + // but the typical way this should be consumed is via dropdowns. + categories.sort_by(|c1, c2| + natural_lexical_cmp(&c1.name_for_dropdown, &c2.name_for_dropdown)); + + server_state.caches.category_list.store_copy(&categories) + .map_err(|e| { + error!("Error storing cache: {:?}", e); + ListTtsCategoriesError::ServerError + })?; + + categories + }, + }; + + // TODO: Cache results in Redis w/ TTL. + let response = ListTtsCategoriesResponse { + success: true, + categories, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTtsCategoriesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/categories/list_tts_model_assigned_categories.rs b/crates/service/storyteller_web/src/http_server/endpoints/categories/list_tts_model_assigned_categories.rs new file mode 100644 index 0000000000..3daff48f93 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/categories/list_tts_model_assigned_categories.rs @@ -0,0 +1,166 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::queries::tts::tts_category_assignments::list_assigned_tts_categories_query_builder::ListAssignedTtsCategoriesQueryBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct ListTtsModelAssignedCategoriesPathInfo { + /// TTS model token + token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct ListTtsModelAssignedCategoriesResponse { + pub success: bool, + pub categories: Vec, +} + +/// Public-facing and optimized (non-irrelevant fields) category list +/// Used for the main TTS dropdown as well as the TTS edit/tagging UI +#[derive(Serialize)] +pub struct DisplayCategory { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + pub maybe_dropdown_name: Option, + + // NB: It's okay for non-mods to see this. + pub is_mod_approved: Option, + + pub category_created_at: DateTime, + pub category_updated_at: DateTime, + pub category_deleted_at: Option>, + + //pub creator_user_token: Option, + //pub creator_username: Option, + //pub creator_display_name: Option, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum ListTtsModelAssignedCategoriesError { + ServerError, +} + +impl ResponseError for ListTtsModelAssignedCategoriesError { + fn status_code(&self) -> StatusCode { + match *self { + ListTtsModelAssignedCategoriesError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTtsModelAssignedCategoriesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn list_tts_model_assigned_categories_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListTtsModelAssignedCategoriesError::ServerError + })?; + + // TODO: We don't have any permissions for categories. This is a proxy. + let is_mod = maybe_user_session + .map(|session| session.can_ban_users) + .unwrap_or(false); + + let query_builder = ListAssignedTtsCategoriesQueryBuilder::for_model_token(&path.token) + .show_invalid_model_not_allowed_categories(is_mod) + .show_deleted(is_mod) + .show_unapproved(is_mod); + + let query_result = + query_builder.perform_query(&server_state.mysql_pool).await; + + let results = match query_result { + Ok(results) => results, + Err(err) => { + warn!("DB error: {:?}", err); + return Err(ListTtsModelAssignedCategoriesError::ServerError); + } + }; + + let mut categories = results.categories.iter() + .map(|c| { + DisplayCategory { + category_token: c.category_token.clone(), + model_type: c.model_type.clone(), + maybe_super_category_token: c.maybe_super_category_token.clone(), + can_directly_have_models: c.can_directly_have_models, + can_have_subcategories: c.can_have_subcategories, + can_only_mods_apply: c.can_only_mods_apply, + name: c.name.clone(), + maybe_dropdown_name:c.maybe_dropdown_name.clone(), + is_mod_approved: c.is_mod_approved.clone(), + category_created_at: c.category_created_at.clone(), + category_updated_at: c.category_updated_at.clone(), + category_deleted_at: c.category_deleted_at.clone(), + } + }) + .collect::>(); + + // TODO: Sort by dropdown name too. + categories.sort_by(|c1, c2| c1.name.cmp(&c2.name)); + + let response = ListTtsModelAssignedCategoriesResponse { + success: true, + categories, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTtsModelAssignedCategoriesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/categories/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/categories/mod.rs new file mode 100644 index 0000000000..a8e6e001ae --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/categories/mod.rs @@ -0,0 +1,5 @@ +pub mod assign_tts_category; +pub mod create_category; +pub mod get_category; +pub mod list_tts_categories; +pub mod list_tts_model_assigned_categories; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/events/list_events.rs b/crates/service/storyteller_web/src/http_server/endpoints/events/list_events.rs new file mode 100644 index 0000000000..250a09e651 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/events/list_events.rs @@ -0,0 +1,116 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct EventRecord { + pub event_token: String, + pub event_type: String, + pub maybe_target_user_token: Option, + pub maybe_target_username: Option, + pub maybe_target_display_name: Option, + pub maybe_target_user_gravatar_hash: Option, + pub maybe_target_entity_token: Option, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct ListEventsSuccessResponse { + pub success: bool, + pub events: Vec, +} + +#[derive(Debug, Display)] +pub enum ListEventsError { + ServerError, +} + +impl ResponseError for ListEventsError { + fn status_code(&self) -> StatusCode { + match *self { + ListEventsError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListEventsError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn list_events_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_events = sqlx::query_as!( + EventRecord, + r#" +SELECT + events.token as event_token, + events.event_type, + events.maybe_target_user_token, + users.username as maybe_target_username, + users.display_name as maybe_target_display_name, + users.email_gravatar_hash as maybe_target_user_gravatar_hash, + events.maybe_target_entity_token, + events.created_at, + events.updated_at +FROM firehose_entries as events +LEFT OUTER JOIN users +ON events.maybe_target_user_token = users.token +ORDER BY events.id DESC +LIMIT 25 + "#, + ) + .fetch_all(&server_state.mysql_pool) + .await; + + let events : Vec = match maybe_events { + Ok(events) => events, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(ListEventsError::ServerError); + }, + _ => { + warn!("tts models list query error: {:?}", err); + return Err(ListEventsError::ServerError); + } + } + } + }; + + let response = ListEventsSuccessResponse { + success: true, + events, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListEventsError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/events/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/events/mod.rs new file mode 100644 index 0000000000..409b5152d2 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/events/mod.rs @@ -0,0 +1 @@ +pub mod list_events; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/default_redirect.rs b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/default_redirect.rs new file mode 100644 index 0000000000..95ac8fbc41 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/default_redirect.rs @@ -0,0 +1,18 @@ + +pub const DEFAULT_INVESTOR_REDIRECT : &'static str = "https://fakeyou.com"; + +pub fn redirect_is_allowed(redirect_url: &str) -> bool { + match redirect_url { + "http://jungle.horse" => true, + "http://jungle.horse/" => true, + "https://create.storyteller.io" => true, + "https://create.storyteller.io/" => true, + "https://fakeyou.com" => true, + "https://fakeyou.com/clone" => true, + "https://jungle.horse" => true, + "https://jungle.horse/" => true, + "https://storyteller.io" => true, + "https://storyteller.io/" => true, + _ => false, + } +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/demo_cookie.rs b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/demo_cookie.rs new file mode 100644 index 0000000000..540ffe58a0 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/demo_cookie.rs @@ -0,0 +1,10 @@ +//! A temporary cookie for investor demo + +use actix_web::HttpRequest; + +pub const STORYTELLER_DEMO_COOKIE_NAME : &'static str = "storyteller_demo"; + +pub fn request_has_demo_cookie(http_request: &HttpRequest) -> bool { + // NB: Payload of the cookie does not matter. + http_request.cookie(STORYTELLER_DEMO_COOKIE_NAME).is_some() +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/disable_demo_mode_handler.rs b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/disable_demo_mode_handler.rs new file mode 100644 index 0000000000..ddbd070863 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/disable_demo_mode_handler.rs @@ -0,0 +1,69 @@ +use actix_http::header; +use actix_http::http::StatusCode; +use actix_web::cookie::Cookie; +use actix_web::web::Query; +use actix_web::{HttpResponse, HttpRequest, Responder, get, web, HttpMessage}; +use crate::http_server::endpoints::investor_demo::default_redirect::{DEFAULT_INVESTOR_REDIRECT, redirect_is_allowed}; +use crate::http_server::endpoints::investor_demo::demo_cookie::STORYTELLER_DEMO_COOKIE_NAME; +use crate::server_state::ServerState; +use hyper::header::LOCATION; +use log::info; +use std::ops::Deref; +use std::sync::Arc; +use time::OffsetDateTime; +use http_server_common::request::get_request_host::get_request_host; + +#[derive(Deserialize)] +pub struct QueryFields { + redirect_to: Option, +} + +pub async fn disable_demo_mode_handler( + http_request: HttpRequest, + query: Query, + server_state: web::Data> +) -> impl Responder +{ + let unsafe_redirect = query.redirect_to + .as_deref() + .map(|r| r.to_string()) + .unwrap_or(DEFAULT_INVESTOR_REDIRECT.to_string()); + + let redirect_allowed = redirect_is_allowed(&unsafe_redirect); + + let safe_redirect_url = if redirect_allowed { + unsafe_redirect + } else { + DEFAULT_INVESTOR_REDIRECT.to_string() + }; + + let maybe_host = get_request_host(&http_request); + let maybe_cookie_domain = match maybe_host.as_deref() { + Some("jungle.horse") | Some("api.jungle.horse") => Some(".jungle.horse"), + Some("fakeyou.com") | Some("api.fakeyou.com") => Some(".fakeyou.com"), + Some("storyteller.io") | Some("api.storyteller.io") => Some(".storyteller.io"), + _ => None, + }; + + // Kill the cookie + let mut cookie_builder = Cookie::build(STORYTELLER_DEMO_COOKIE_NAME, "") + .secure(server_state.env_config.cookie_secure) // HTTPS-only + .expires(OffsetDateTime::unix_epoch()) + .http_only(false) // This is meant to be exposed to Javascript! + // NB: Since we're setting this from direct HTTP access rather than XHR, the browser + // implicitly sets the path used to access the cookie, ie. '/demo_mode', which is not what + // we want! + .path("/") + .permanent(); + + if let Some(cookie_domain) = maybe_cookie_domain { + cookie_builder = cookie_builder.domain(cookie_domain); + } + + let cookie = cookie_builder.finish(); + + HttpResponse::build(StatusCode::FOUND) + .append_header((header::LOCATION, safe_redirect_url.to_string())) + .del_cookie(&cookie) + .finish() +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/enable_demo_mode_handler.rs b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/enable_demo_mode_handler.rs new file mode 100644 index 0000000000..2f75d9d206 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/enable_demo_mode_handler.rs @@ -0,0 +1,68 @@ +use actix_http::header; +use actix_http::http::StatusCode; +use actix_web::cookie::Cookie; +use actix_web::web::Query; +use actix_web::{HttpResponse, HttpRequest, Responder, get, web, HttpMessage}; +use crate::http_server::endpoints::investor_demo::default_redirect::{DEFAULT_INVESTOR_REDIRECT, redirect_is_allowed}; +use crate::http_server::endpoints::investor_demo::demo_cookie::STORYTELLER_DEMO_COOKIE_NAME; +use crate::server_state::ServerState; +use hyper::header::LOCATION; +use log::info; +use std::ops::Deref; +use std::sync::Arc; +use time::OffsetDateTime; +use http_server_common::request::get_request_host::get_request_host; + +#[derive(Deserialize)] +pub struct QueryFields { + redirect_to: Option, +} + +pub async fn enable_demo_mode_handler( + http_request: HttpRequest, + query: Query, + server_state: web::Data> +) -> impl Responder +{ + let unsafe_redirect = query.redirect_to + .as_deref() + .map(|r| r.to_string()) + .unwrap_or(DEFAULT_INVESTOR_REDIRECT.to_string()); + + let redirect_allowed = redirect_is_allowed(&unsafe_redirect); + + let safe_redirect_url = if redirect_allowed { + unsafe_redirect + } else { + DEFAULT_INVESTOR_REDIRECT.to_string() + }; + + let maybe_host = get_request_host(&http_request); + let maybe_cookie_domain = match maybe_host.as_deref() { + Some("jungle.horse") | Some("api.jungle.horse") => Some(".jungle.horse"), + Some("fakeyou.com") | Some("api.fakeyou.com") => Some(".fakeyou.com"), + Some("storyteller.io") | Some("api.storyteller.io") => Some(".storyteller.io"), + _ => None, + }; + + // Add the cookie. + let mut cookie_builder = Cookie::build(STORYTELLER_DEMO_COOKIE_NAME, "true") + .secure(server_state.env_config.cookie_secure) // HTTPS-only + .http_only(false) // This is meant to be exposed to Javascript! + // NB: Since we're setting this from direct HTTP access rather than XHR, the browser + // implicitly sets the path used to access the cookie, ie. '/demo_mode', which is not what + // we want! + .path("/") + .permanent(); + + if let Some(cookie_domain) = maybe_cookie_domain { + cookie_builder = cookie_builder.domain(cookie_domain); + } + + let cookie = cookie_builder.finish(); + + HttpResponse::build(StatusCode::FOUND) + .append_header((header::LOCATION, safe_redirect_url.to_string())) + .cookie(cookie) + .finish() +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/mod.rs new file mode 100644 index 0000000000..92c427c8cf --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/investor_demo/mod.rs @@ -0,0 +1,4 @@ +pub mod demo_cookie; +pub mod disable_demo_mode_handler; +pub mod enable_demo_mode_handler; +pub mod default_redirect; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/leaderboard/get_leaderboard.rs b/crates/service/storyteller_web/src/http_server/endpoints/leaderboard/get_leaderboard.rs new file mode 100644 index 0000000000..de65564b56 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/leaderboard/get_leaderboard.rs @@ -0,0 +1,112 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use database_queries::queries::tts::stats::calculate_tts_model_leaderboard::TtsLeaderboardRecordForList; +use database_queries::queries::tts::stats::calculate_tts_model_leaderboard::calculate_tts_model_leaderboard; +use database_queries::queries::w2l::stats::calculate_w2l_template_leaderboard::W2lLeaderboardRecordForList; +use database_queries::queries::w2l::stats::calculate_w2l_template_leaderboard::calculate_w2l_template_leaderboard; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct LeaderboardResponse { + success: bool, + tts_leaderboard: Vec, + w2l_leaderboard: Vec, +} + +#[derive(Serialize, Debug)] +pub struct LeaderboardErrorResponse { + pub success: bool, + pub error_type: LeaderboardErrorType, + pub error_message: String, +} + +#[derive(Copy, Clone, Debug, Serialize)] +pub enum LeaderboardErrorType { + ServerError, +} + +impl LeaderboardErrorResponse { + fn server_error() -> Self { + Self { + success: false, + error_type: LeaderboardErrorType::ServerError, + error_message: "server error".to_string() + } + } +} + +impl ResponseError for LeaderboardErrorResponse { + fn status_code(&self) -> StatusCode { + match self.error_type { + LeaderboardErrorType::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for LeaderboardErrorResponse { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.error_type) + } +} + +pub async fn leaderboard_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_tts_results = + calculate_tts_model_leaderboard(&server_state.mysql_pool); + + let maybe_w2l_results = + calculate_w2l_template_leaderboard(&server_state.mysql_pool); + + let tts_results = match maybe_tts_results.await { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(LeaderboardErrorResponse::server_error()); + } + }; + + let w2l_results = match maybe_w2l_results.await { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(LeaderboardErrorResponse::server_error()); + } + }; + + let response = LeaderboardResponse { + success: true, + tts_leaderboard: tts_results, + w2l_leaderboard: w2l_results, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| LeaderboardErrorResponse::server_error())?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/leaderboard/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/leaderboard/mod.rs new file mode 100644 index 0000000000..4fad116a64 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/leaderboard/mod.rs @@ -0,0 +1 @@ +pub mod get_leaderboard; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/alpha_cookie.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/alpha_cookie.rs new file mode 100644 index 0000000000..da5bdb88b7 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/alpha_cookie.rs @@ -0,0 +1,13 @@ +use actix_web::cookie::Cookie; +use crate::server_state::ServerState; + +pub const ALPHA_COOKIE_NAME : &'static str = "enable-alpha"; + +pub fn alpha_cookie(server_state: &ServerState) -> Cookie { + Cookie::build(ALPHA_COOKIE_NAME, "true") + .domain(&server_state.env_config.cookie_domain) + .secure(server_state.env_config.cookie_secure) // HTTPS-only + .http_only(false) // This is meant to be exposed to Javascript! + .permanent() + .finish() +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/default_route_404.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/default_route_404.rs new file mode 100644 index 0000000000..6bc7403c15 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/default_route_404.rs @@ -0,0 +1,10 @@ +use actix_http::http::StatusCode; +use actix_web::{HttpResponse, HttpRequest, Responder, get}; +use log::warn; + +pub async fn default_route_404() -> impl Responder { + warn!("404 not found"); + HttpResponse::NotFound() + .content_type("text/html; charset=utf-8") + .body("

not found

") +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/detect_locale_handler.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/detect_locale_handler.rs new file mode 100644 index 0000000000..ce2515f879 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/detect_locale_handler.rs @@ -0,0 +1,104 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use http_server_common::request::get_request_header_optional::get_request_header_optional; +use http_server_common::request::parse_accept_language::parse_accept_language; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Success Response =============== + +#[derive(Serialize, Default)] +pub struct DetectLocaleResponse { + pub success: bool, + /// Full BCP47 language tags + pub full_language_tags: Vec, + /// Parsed out languages + pub language_codes: Vec, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum DetectLocaleError { + ServerError, +} + +impl ResponseError for DetectLocaleError { + fn status_code(&self) -> StatusCode { + match *self { + DetectLocaleError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DetectLocaleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub const FORCE_LOCALE_COOKIE_HEADER_NAME : &'static str = "force-locale"; + +pub async fn detect_locale_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let mut maybe_accept_language = + get_request_header_optional(&http_request, "accept-language"); + + if let Some(cookie) = http_request.cookie(FORCE_LOCALE_COOKIE_HEADER_NAME) { + warn!("Overriding default accept language with custom value (from cookie)"); + maybe_accept_language = Some(cookie.value().to_string()); + + } else if let Some(header) = get_request_header_optional(&http_request, FORCE_LOCALE_COOKIE_HEADER_NAME) { + warn!("Overriding default accept language with custom value (from header)"); + maybe_accept_language = Some(header); + } + + let accept_language = maybe_accept_language.unwrap_or("en".to_string()); + + let mut response = DetectLocaleResponse::default(); + response.success = true; + + let language_tags = parse_accept_language(&accept_language); + + for language_tag in language_tags.iter() { + response.full_language_tags.push(language_tag.to_string()); + response.language_codes.push(language_tag.primary_language().to_string()); + } + + let body = serde_json::to_string(&response) + .map_err(|e| DetectLocaleError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/enable_alpha_easy_handler.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/enable_alpha_easy_handler.rs new file mode 100644 index 0000000000..c93aea0d49 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/enable_alpha_easy_handler.rs @@ -0,0 +1,28 @@ +use actix_web::cookie::Cookie; +use actix_http::http::StatusCode; +use actix_web::web::Query; +use actix_web::{HttpResponse, HttpRequest, Responder, get, web, HttpMessage}; +use crate::http_server::endpoints::misc::alpha_cookie::alpha_cookie; +use crate::server_state::ServerState; +use hyper::header::LOCATION; +use log::info; +use std::ops::Deref; +use std::sync::Arc; + +#[get("/enable")] +pub async fn enable_alpha_easy_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> impl Responder +{ + info!("GET /enable"); + + let cookie = alpha_cookie(&server_state); + + let response = HttpResponse::Found() + .cookie(cookie) + .header(LOCATION, server_state.env_config.website_homepage_redirect.deref()) + .finish(); + + response +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/enable_alpha_handler.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/enable_alpha_handler.rs new file mode 100644 index 0000000000..b07d763759 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/enable_alpha_handler.rs @@ -0,0 +1,55 @@ +use actix_web::cookie::Cookie; +use actix_http::http::StatusCode; +use actix_web::web::Query; +use actix_web::{HttpResponse, HttpRequest, Responder, get, web, HttpMessage}; +use crate::http_server::endpoints::misc::alpha_cookie::{ALPHA_COOKIE_NAME, alpha_cookie}; +use crate::server_state::ServerState; +use log::info; +use std::sync::Arc; + +const CONTENT_TYPE : &'static str = "text/html; charset=utf-8"; + +const ENABLE_LINK : &'static str = "I want Vocodes 2.0"; +const DISABLE_LINK : &'static str = "I don't want Vocodes 2.0"; +const STATUS_LINK : &'static str = "status"; + + +#[derive(Deserialize)] +pub struct QueryFields { + enable: Option, +} + +#[get("/alpha")] +pub async fn enable_alpha_handler(http_request: HttpRequest, + query: Query, + server_state: web::Data>) -> impl Responder +{ + info!("GET /alpha"); + + let cookie = alpha_cookie(&server_state); + + match query.enable { + None => { + let cookie_exists = http_request.cookie(ALPHA_COOKIE_NAME).is_some(); + + HttpResponse::build(StatusCode::OK) + .content_type(CONTENT_TYPE) + .body(format!("

Vocodes 2.0

Current `{}` cookie state = {}

{} | {} | {}", + ALPHA_COOKIE_NAME, cookie_exists, ENABLE_LINK, DISABLE_LINK, STATUS_LINK)) + } + Some(true) => { + HttpResponse::build(StatusCode::OK) + .content_type(CONTENT_TYPE) + .cookie(cookie) + .body(format!("

Vocodes 2.0

setting `{}` cookie

{} | {} | {}", + ALPHA_COOKIE_NAME, ENABLE_LINK, DISABLE_LINK, STATUS_LINK)) + } + Some(false) => { + HttpResponse::build(StatusCode::OK) + .content_type(CONTENT_TYPE) + .del_cookie(&cookie) + .body(format!("

Vocodes 2.0

unsetting `{}` cookie

{} | {} | {}", + ALPHA_COOKIE_NAME, ENABLE_LINK, DISABLE_LINK, STATUS_LINK)) + } + } +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/mod.rs new file mode 100644 index 0000000000..417922d9aa --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/mod.rs @@ -0,0 +1,6 @@ +pub mod alpha_cookie; +pub mod default_route_404; +pub mod detect_locale_handler; +pub mod enable_alpha_easy_handler; +pub mod enable_alpha_handler; +pub mod root_index; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/misc/root_index.rs b/crates/service/storyteller_web/src/http_server/endpoints/misc/root_index.rs new file mode 100644 index 0000000000..960d858d00 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/misc/root_index.rs @@ -0,0 +1,22 @@ +use actix_http::http::StatusCode; +use actix_web::Error; +use actix_web::Responder; +use actix_web::Result; +use actix_web::web::Data; +use actix_web::{HttpResponse, HttpRequest, get}; +use crate::server_state::ServerState; +use log::debug; +use std::sync::Arc; + +pub async fn get_root_index(server_state: Data>) -> HttpResponse { + debug!("GET /"); // NB: Google load balancer hits this a lot, and it spams. + let body = format!(" \ +

hello!

\ +

Are you looking for an API? Join our Discord!

\ +

Maybe you want to work with us? We can pay! Get in touch!

\ +

storyteller-web instance: {}

\ + ", server_state.hostname); + HttpResponse::build(StatusCode::OK) + .content_type("text/html; charset=utf-8") + .body(body) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/mod.rs new file mode 100644 index 0000000000..e7dc1acb4f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/mod.rs @@ -0,0 +1,13 @@ +pub mod api_tokens; +pub mod categories; +pub mod events; +pub mod investor_demo; +pub mod leaderboard; +pub mod misc; +pub mod moderation; +pub mod service; +pub mod tts; +pub mod twitch; +pub mod users; +pub mod voice_clone_requests; +pub mod w2l; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/approval/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/approval/mod.rs new file mode 100644 index 0000000000..9f34ef78bf --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/approval/mod.rs @@ -0,0 +1 @@ +pub mod pending_w2l_templates; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/approval/pending_w2l_templates.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/approval/pending_w2l_templates.rs new file mode 100644 index 0000000000..5752c04ac5 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/approval/pending_w2l_templates.rs @@ -0,0 +1,182 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log}; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct GetPendingW2lTemplatesResponse { + pub success: bool, + pub templates: Vec, +} + +#[derive(Serialize)] +pub struct PendingW2lTemplate { + pub template_token: String, + pub title: String, + pub template_type: String, + pub duration_millis: i32, + pub frame_width: i32, + pub frame_height: i32, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + pub created_at: DateTime, +} + +#[derive(Debug, Serialize)] +pub enum GetPendingW2lTemplatesError { + ServerError, + Unauthorized, +} + +impl ResponseError for GetPendingW2lTemplatesError { + fn status_code(&self) -> StatusCode { + match *self { + GetPendingW2lTemplatesError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetPendingW2lTemplatesError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for GetPendingW2lTemplatesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_pending_w2l_templates_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetPendingW2lTemplatesError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetPendingW2lTemplatesError::Unauthorized); + } + }; + + if !user_session.can_approve_w2l_templates { + warn!("user is not allowed to approve w2l templates: {}", user_session.user_token); + return Err(GetPendingW2lTemplatesError::Unauthorized); + } + + // NB: Lookup failure is Err(RowNotFound). + let maybe_results = sqlx::query_as!( + PendingW2lTemplatesRaw, + r#" +SELECT + w2l_templates.token as template_token, + w2l_templates.title, + w2l_templates.template_type, + w2l_templates.duration_millis, + w2l_templates.frame_width, + w2l_templates.frame_height, + w2l_templates.creator_user_token, + users.username AS creator_username, + users.display_name AS creator_display_name, + users.email_gravatar_hash AS creator_gravatar_hash, + w2l_templates.created_at +FROM + w2l_templates +JOIN + users +ON + users.token = w2l_templates.creator_user_token +WHERE + w2l_templates.is_public_listing_approved IS NULL + AND w2l_templates.user_deleted_at IS NULL + AND w2l_templates.mod_deleted_at IS NULL + AND w2l_templates.is_locked_from_use IS FALSE + "#, + ) + .fetch_all(&server_state.mysql_pool) + .await; + + let results : Vec = match maybe_results { + Ok(results) => results, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + Vec::new() + }, + _ => { + warn!("list pending w2l db error: {:?}", err); + return Err(GetPendingW2lTemplatesError::ServerError) + } + } + } + }; + + let results = results.iter().map(|r| { + PendingW2lTemplate { + template_token: r.template_token.clone(), + title: r.title.clone(), + template_type: r.template_type.clone(), + duration_millis: r.duration_millis, + frame_width: r.frame_width, + frame_height: r.frame_height, + creator_user_token: r.creator_user_token.clone(), + creator_username: r.creator_username.clone(), + creator_display_name: r.creator_display_name.clone(), + creator_gravatar_hash: r.creator_gravatar_hash.clone(), + created_at: r.created_at.clone(), + } + }).collect::>(); + + let response = GetPendingW2lTemplatesResponse { + success: true, + templates: results, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetPendingW2lTemplatesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +#[derive(Serialize)] +pub struct PendingW2lTemplatesRaw { + pub template_token: String, + pub title: String, + pub template_type: String, + pub duration_millis: i32, + pub frame_width: i32, + pub frame_height: i32, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + pub created_at: DateTime, +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/delete_category.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/delete_category.rs new file mode 100644 index 0000000000..b8f8021d0f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/delete_category.rs @@ -0,0 +1,150 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; +use std::fmt; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; + +// =============== Request =============== + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct DeleteCategoryPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct DeleteCategoryRequest { + set_delete: bool, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct DeleteCategoryResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum DeleteCategoryError { + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for DeleteCategoryError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteCategoryError::NotFound => StatusCode::NOT_FOUND, + DeleteCategoryError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteCategoryError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for DeleteCategoryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn delete_category_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteCategoryError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteCategoryError::NotAuthorized); + } + }; + + // TODO: We don't have a permission for categories, so we use this as a proxy. + if !user_session.can_ban_users { + warn!("no permission to delete categories"); + return Err(DeleteCategoryError::NotAuthorized); + } + + let query_builder = if request.set_delete { + sqlx::query!(r#" + UPDATE model_categories + SET + deleted_at = CURRENT_TIMESTAMP + WHERE + token = ? + LIMIT 1 + "#, + path.token) + + } else { + sqlx::query!(r#" + UPDATE model_categories + SET + deleted_at = NULL + WHERE + token = ? + LIMIT 1 + "#, + path.token) + }; + + // NB: We're soft deleting so we don't delete the associations. + let query_result = query_builder.execute(&server_state.mysql_pool).await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Delete/undelete category edit DB error: {:?}", err); + return Err(DeleteCategoryError::ServerError); + } + }; + + let response = DeleteCategoryResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| DeleteCategoryError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/edit_category.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/edit_category.rs new file mode 100644 index 0000000000..5c19e5dee7 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/edit_category.rs @@ -0,0 +1,206 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::model_categories::get_category_by_token::get_category_by_token; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct EditCategoryPathInfo { + token: String, +} + +// NB: ONLY MODERATORS CAN EDIT CATEGORIES. +// These are not sparse updates! +#[derive(Deserialize)] +pub struct EditCategoryRequest { + pub name: String, + + // If absent, null the fields. + pub maybe_dropdown_name: Option, + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + // Moderation fields + pub is_mod_approved: bool, + + // If absent, null the field. + pub maybe_mod_comments: Option, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct EditCategoryResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum EditCategoryError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for EditCategoryError { + fn status_code(&self) -> StatusCode { + match *self { + EditCategoryError::BadInput(_) => StatusCode::BAD_REQUEST, + EditCategoryError::NotFound => StatusCode::NOT_FOUND, + EditCategoryError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditCategoryError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditCategoryError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn edit_category_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditCategoryError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditCategoryError::NotAuthorized); + } + }; + + // TODO: We don't have a permission for categories, so we use this as a proxy. + if !user_session.can_ban_users { + warn!("no permission to edit categories"); + return Err(EditCategoryError::NotAuthorized); + } + + // Category tree integrity + if let Some(parent_category_token) = request.maybe_super_category_token.as_deref() { + if parent_category_token == &path.token { + return Err(EditCategoryError::BadInput( + "category cannot have itself as a parent".to_string())); + } + + let parent_category_lookup + = get_category_by_token(parent_category_token, &server_state.mysql_pool).await; + + match parent_category_lookup { + Ok(Some(parent_category)) => { + if !parent_category.can_have_subcategories { + return Err(EditCategoryError::BadInput( + "parent category cannot have children".to_string())); + } + }, + Ok(None) => return Err(EditCategoryError::NotFound), + Err(err) => { + warn!("Category lookup DB error: {:?}", err); + return Err(EditCategoryError::ServerError) + }, + } + } + + let query_result = + // We need to store the IP address details. + sqlx::query!( + r#" +UPDATE model_categories +SET + name = ?, + maybe_dropdown_name = ?, + + can_directly_have_models = ?, + can_have_subcategories = ?, + can_only_mods_apply = ?, + + maybe_super_category_token = ?, + + is_mod_approved = ?, + maybe_mod_user_token = ?, + maybe_mod_comments = ?, + + version = version + 1 + +WHERE token = ? +LIMIT 1 + "#, + &request.name, + &request.maybe_dropdown_name, + &request.can_directly_have_models, + &request.can_have_subcategories, + &request.can_only_mods_apply, + &request.maybe_super_category_token, + &request.is_mod_approved, + &user_session.user_token, + &request.maybe_mod_comments, + &path.token, + ) + .execute(&server_state.mysql_pool) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Edit category DB error: {:?}", err); + return Err(EditCategoryError::ServerError); + } + }; + + let response = EditCategoryResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| EditCategoryError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/list_tts_categories_for_moderation.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/list_tts_categories_for_moderation.rs new file mode 100644 index 0000000000..e471a96184 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/list_tts_categories_for_moderation.rs @@ -0,0 +1,212 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::queries::model_categories::list_categories_query_builder::ListCategoriesQueryBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize, Copy, Clone, Debug)] +#[serde(rename_all = "snake_case")] +pub enum ViewState { + /// 'exclude'; keep the objects out + Exclude, + /// 'include'; include the objects + Include, + /// 'only'; limit to just the objects + Only, +} + +// Only show approved&non-deleted , only show un-approved, only show deleted +#[derive(Deserialize)] +pub struct QueryParams { + deleted: Option, + unapproved: Option, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct ListTtsCategoriesForModerationResponse { + pub success: bool, + pub categories: Vec, +} + +#[derive(Serialize)] +pub struct CategoryForModeration { + pub category_token: String, + + pub model_type: String, // TODO: ENUM + + pub maybe_super_category_token: Option, + + pub can_directly_have_models: bool, + pub can_have_subcategories: bool, + pub can_only_mods_apply: bool, + + pub name: String, + pub maybe_dropdown_name: Option, + + pub creator_user_token: Option, + pub creator_username: Option, + pub creator_display_name: Option, + pub creator_gravatar_hash: Option, + + // Moderator fields + pub is_mod_approved: Option, + pub maybe_mod_comments: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, + pub deleted_at: Option>, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum ListTtsCategoriesForModerationError { + NotAuthorized, + ServerError, +} + +impl ResponseError for ListTtsCategoriesForModerationError { + fn status_code(&self) -> StatusCode { + match *self { + ListTtsCategoriesForModerationError::NotAuthorized => StatusCode::UNAUTHORIZED, + ListTtsCategoriesForModerationError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTtsCategoriesForModerationError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn list_tts_categories_for_moderation_handler( + http_request: HttpRequest, + query: web::Query, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListTtsCategoriesForModerationError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListTtsCategoriesForModerationError::NotAuthorized); + } + }; + + // TODO: We don't have a permission for categories, so we use this as a proxy. + if !user_session.can_ban_users { + warn!("no permission to edit categories"); + return Err(ListTtsCategoriesForModerationError::NotAuthorized); + } + + // The objects we'll show or hide + let delete_state = query.deleted.unwrap_or(ViewState::Exclude); + let unapproved_state = query.unapproved.unwrap_or(ViewState::Include); + + let mut query_builder = ListCategoriesQueryBuilder::new() + .scope_model_type(Some("tts")); + + query_builder = match delete_state { + ViewState::Exclude => query_builder.show_deleted(false), + ViewState::Include => query_builder.show_deleted(true), + ViewState::Only => query_builder.show_deleted(true).hide_non_deleted(true), + }; + + query_builder = match unapproved_state { + ViewState::Exclude => query_builder.show_unapproved(false), + ViewState::Include => query_builder.show_unapproved(true), + ViewState::Only => query_builder.show_unapproved(true).hide_approved(true), + }; + + let query_result = + query_builder.perform_query(&server_state.mysql_pool).await; + + let results = match query_result { + Ok(results) => results, + Err(err) => { + warn!("DB error: {:?}", err); + return Err(ListTtsCategoriesForModerationError::ServerError); + } + }; + + warn!("Number of categories: {:?}", results.categories.len()); + + let mut categories = results.categories.iter() + .map(|c| { + CategoryForModeration { + category_token: c.category_token.clone(), + model_type: c.model_type.clone(), + maybe_super_category_token: c.maybe_super_category_token.clone(), + can_directly_have_models: c.can_directly_have_models, + can_have_subcategories: c.can_have_subcategories, + can_only_mods_apply: c.can_only_mods_apply, + name: c.name.clone(), + maybe_dropdown_name:c.maybe_dropdown_name.clone(), + creator_user_token: c.creator_user_token.clone(), + creator_username: c.creator_username.clone(), + creator_display_name: c.creator_display_name.clone(), + creator_gravatar_hash: c.creator_gravatar_hash.clone(), + is_mod_approved: c.is_mod_approved.clone(), + maybe_mod_comments: c.maybe_mod_comments.clone(), + created_at: c.created_at.clone(), + updated_at: c.updated_at.clone(), + deleted_at: c.deleted_at.clone() + } + }) + .collect::>(); + + // TODO: Sort by dropdown name too. + categories.sort_by(|c1, c2| c1.name.cmp(&c2.name)); + + let response = ListTtsCategoriesForModerationResponse { + success: true, + categories, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTtsCategoriesForModerationError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/mod.rs new file mode 100644 index 0000000000..0f457bdcc0 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/categories/mod.rs @@ -0,0 +1,3 @@ +pub mod delete_category; +pub mod edit_category; +pub mod list_tts_categories_for_moderation; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/add_ip_ban.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/add_ip_ban.rs new file mode 100644 index 0000000000..3be5a80fe9 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/add_ip_ban.rs @@ -0,0 +1,135 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::ip_addresses::validate_moderator_provided_ip_address; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct AddIpBanRequest { + ip_address: String, + mod_notes: String, + maybe_target_user_token: Option, +} + +#[derive(Debug, Display)] +pub enum AddIpBanError { + BadInput(String), + ServerError, + Unauthorized, +} + +impl ResponseError for AddIpBanError { + fn status_code(&self) -> StatusCode { + match *self { + AddIpBanError::BadInput(_) => StatusCode::BAD_REQUEST, + AddIpBanError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + AddIpBanError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + AddIpBanError::BadInput(reason) => reason.to_string(), + AddIpBanError::ServerError => "server error".to_string(), + AddIpBanError::Unauthorized => "unauthorized".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn add_ip_ban_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + AddIpBanError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(AddIpBanError::Unauthorized); + } + }; + + if !user_session.can_ban_users { + warn!("user is not allowed to add bans: {}", user_session.user_token); + return Err(AddIpBanError::Unauthorized); + } + + let ip_address = request.ip_address.trim(); + + if let Err(e) = validate_moderator_provided_ip_address(&ip_address) { + warn!("Bad ip address: {}", e); + return Err(AddIpBanError::BadInput(e.to_string())); + } + + info!("Creating ban..."); + + let query_result = sqlx::query!( + r#" +INSERT INTO + ip_address_bans +SET + ip_address = ?, + maybe_target_user_token = ?, + mod_user_token = ?, + mod_notes = ? +ON DUPLICATE KEY UPDATE + expires_at = NULL, + deleted_at = NULL, + ip_address = ?, + maybe_target_user_token = ?, + mod_user_token = ?, + mod_notes = ? + "#, + // Insert + &request.ip_address, + &request.maybe_target_user_token, + &user_session.user_token, + &request.mod_notes, + // Update + &ip_address, + &request.maybe_target_user_token, + &user_session.user_token, + &request.mod_notes, + ) + .execute(&server_state.mysql_pool) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Add IP ban DB error: {:?}", err); + return Err(AddIpBanError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/delete_ip_ban.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/delete_ip_ban.rs new file mode 100644 index 0000000000..eba7577ba4 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/delete_ip_ban.rs @@ -0,0 +1,140 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct DeleteIpBanPathInfo { + ip_address: String, +} + +#[derive(Deserialize)] +pub struct DeleteIpBanRequest { + delete: bool, +} + +#[derive(Debug, Display)] +pub enum DeleteIpBanError { + BadInput(String), + ServerError, + Unauthorized, +} + +impl ResponseError for DeleteIpBanError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteIpBanError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteIpBanError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + DeleteIpBanError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + DeleteIpBanError::BadInput(reason) => reason.to_string(), + DeleteIpBanError::ServerError => "server error".to_string(), + DeleteIpBanError::Unauthorized => "unauthorized".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn delete_ip_ban_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteIpBanError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteIpBanError::Unauthorized); + } + }; + + if !user_session.can_ban_users { + warn!("user is not allowed to delete bans: {}", user_session.user_token); + return Err(DeleteIpBanError::Unauthorized); + } + + let query_result = if request.delete { + info!("Deleting IP ban: {}", &path.ip_address); + + sqlx::query!( + r#" +UPDATE + ip_address_bans +SET + mod_user_token = ?, + deleted_at = CURRENT_TIMESTAMP +WHERE + ip_address = ? +LIMIT 1 + "#, + &user_session.user_token, + &path.ip_address, + ) + .execute(&server_state.mysql_pool) + .await + } else { + info!("Restoring IP ban: {}", &path.ip_address); + + sqlx::query!( + r#" +UPDATE + ip_address_bans +SET + mod_user_token = ?, + deleted_at = NULL +WHERE + ip_address = ? +LIMIT 1 + "#, + &user_session.user_token, + &path.ip_address, + ) + .execute(&server_state.mysql_pool) + .await + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("(un)delete IP ban DB error: {:?}", err); + return Err(DeleteIpBanError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/get_ip_ban.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/get_ip_ban.rs new file mode 100644 index 0000000000..ea1336c298 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/get_ip_ban.rs @@ -0,0 +1,170 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetIpBanPathInfo { + ip_address: String, +} + +#[derive(Serialize)] +pub struct GetIpBanResponse { + pub success: bool, + pub ip_address_ban: IpBanRecord, +} + +#[derive(Serialize)] +pub struct IpBanRecord { + pub ip_address: String, + pub maybe_target_user_token: Option, + pub maybe_target_username: Option, + pub mod_user_token: String, + pub mod_username: String, + pub mod_display_name: String, + pub mod_notes: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Debug, Display)] +pub enum GetIpBanError { + BadInput(String), + ServerError, + NotFound, + Unauthorized, +} + +impl ResponseError for GetIpBanError { + fn status_code(&self) -> StatusCode { + match *self { + GetIpBanError::BadInput(_) => StatusCode::BAD_REQUEST, + GetIpBanError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetIpBanError::NotFound => StatusCode::NOT_FOUND, + GetIpBanError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetIpBanError::BadInput(reason) => reason.to_string(), + GetIpBanError::ServerError => "server error".to_string(), + GetIpBanError::NotFound => "not found".to_string(), + GetIpBanError::Unauthorized => "unauthorized".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn get_ip_ban_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetIpBanError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetIpBanError::Unauthorized); + } + }; + + if !user_session.can_ban_users { + warn!("user is not allowed to view bans: {}", user_session.user_token); + return Err(GetIpBanError::Unauthorized); + } + + let ip_address = path.ip_address.trim(); + + // NB: Lookup failure is Err(RowNotFound). + let maybe_result = sqlx::query_as!( + IpBanRecord, + r#" +SELECT + ip_bans.ip_address, + + ip_bans.maybe_target_user_token, + banned_users.username as maybe_target_username, + + ip_bans.mod_user_token, + mod_users.username as mod_username, + mod_users.display_name as mod_display_name, + + ip_bans.mod_notes, + ip_bans.created_at, + ip_bans.updated_at +FROM + ip_address_bans AS ip_bans +LEFT OUTER JOIN users as banned_users + ON ip_bans.maybe_target_user_token = banned_users.token +JOIN users as mod_users + ON ip_bans.mod_user_token = mod_users.token +WHERE + ip_bans.ip_address = ? +LIMIT 1 + "#, + ip_address + ) + .fetch_one(&server_state.mysql_pool) + .await; + + let result : IpBanRecord = match maybe_result { + Ok(result) => { + result + }, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(GetIpBanError::NotFound); + }, + _ => { + warn!("get ip ban db error: {:?}", err); + return Err(GetIpBanError::ServerError); + } + } + } + }; + + let response = GetIpBanResponse { + success: true, + ip_address_ban: result, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetIpBanError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/list_ip_bans.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/list_ip_bans.rs new file mode 100644 index 0000000000..9c7a00a993 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/list_ip_bans.rs @@ -0,0 +1,157 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct ListIpBansResponse { + pub success: bool, + pub ip_address_bans: Vec, +} + +#[derive(Serialize)] +pub struct IpBanRecordForList { + pub ip_address: String, + pub maybe_target_user_token: Option, + pub maybe_target_username: Option, + pub mod_user_token: String, + pub mod_username: String, + pub mod_display_name: String, + pub mod_notes: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Debug, Display)] +pub enum ListIpBansError { + BadInput(String), + ServerError, + Unauthorized, +} + +impl ResponseError for ListIpBansError { + fn status_code(&self) -> StatusCode { + match *self { + ListIpBansError::BadInput(_) => StatusCode::BAD_REQUEST, + ListIpBansError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + ListIpBansError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListIpBansError::BadInput(reason) => reason.to_string(), + ListIpBansError::ServerError => "server error".to_string(), + ListIpBansError::Unauthorized => "unauthorized".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn list_ip_bans_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListIpBansError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListIpBansError::Unauthorized); + } + }; + + if !user_session.can_ban_users { + warn!("user is not allowed to see bans: {}", user_session.user_token); + return Err(ListIpBansError::Unauthorized); + } + + // NB: Lookup failure is Err(RowNotFound). + let maybe_results = sqlx::query_as!( + IpBanRecordForList, + r#" +SELECT + ip_bans.ip_address, + + ip_bans.maybe_target_user_token, + banned_users.username as maybe_target_username, + + ip_bans.mod_user_token, + mod_users.username as mod_username, + mod_users.display_name as mod_display_name, + + ip_bans.mod_notes, + ip_bans.created_at, + ip_bans.updated_at +FROM + ip_address_bans AS ip_bans +LEFT OUTER JOIN users as banned_users + ON ip_bans.maybe_target_user_token = banned_users.token +JOIN users as mod_users + ON ip_bans.mod_user_token = mod_users.token +WHERE + ip_bans.deleted_at IS NULL + "#, + ) + .fetch_all(&server_state.mysql_pool) + .await; + + let results : Vec = match maybe_results { + Ok(results) => { + info!("Results length: {}", results.len()); + results + }, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + Vec::new() + }, + _ => { + warn!("list ip bans db error: {:?}", err); + return Err(ListIpBansError::ServerError); + } + } + } + }; + + let response = ListIpBansResponse { + success: true, + ip_address_bans: results, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListIpBansError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/mod.rs new file mode 100644 index 0000000000..cc016c5c2b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/ip_bans/mod.rs @@ -0,0 +1,4 @@ +pub mod add_ip_ban; +pub mod delete_ip_ban; +pub mod get_ip_ban; +pub mod list_ip_bans; diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/get_tts_inference_queue_count.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/get_tts_inference_queue_count.rs new file mode 100644 index 0000000000..081f061513 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/get_tts_inference_queue_count.rs @@ -0,0 +1,194 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct GetTtsInferenceQueueCountResponse { + pub success: bool, + // Seconds since job at head of queue was enqueued (proxy for queue wait time) + pub seconds_since_first: i64, + + // All "pending" jobs + pub pending_count: i64, + + // All "pending" jobs with priority > 0 + pub pending_priority_nonzero_count: i64, + + // All "pending" jobs with priority > 1 + pub pending_priority_gt_one_count: i64, + + // Failed, but not permanently dead + pub attempt_failed_count: i64, +} + +#[derive(Debug, Serialize)] +pub enum GetTtsInferenceQueueCountError { + ServerError, + Unauthorized, +} + +impl ResponseError for GetTtsInferenceQueueCountError { + fn status_code(&self) -> StatusCode { + match *self { + GetTtsInferenceQueueCountError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetTtsInferenceQueueCountError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for GetTtsInferenceQueueCountError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_tts_inference_queue_count_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetTtsInferenceQueueCountError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetTtsInferenceQueueCountError::Unauthorized); + } + }; + + // TODO: Not a good fit for this permission. + if !user_session.can_ban_users { + warn!("user is not allowed to view bans: {}", user_session.user_token); + return Err(GetTtsInferenceQueueCountError::Unauthorized); + } + + // NB(old?): Lookup failure is Err(RowNotFound). + // NB(2022-03-05): The "seconds_since_first" result might return null if no pending, so IFNULL + // means we won't fail the full query. + let maybe_result = sqlx::query_as!( + PendingCountResult, + r#" +SELECT + IFNULL(( + SELECT + NOW() - t1.created_at AS seconds_since_first + FROM tts_inference_jobs AS t1 + WHERE t1.status = "pending" + ORDER BY t1.id ASC + LIMIT 1 + ), 0) as seconds_since_first, + sub2.pending_count, + sub3.pending_priority_nonzero_count, + sub4.pending_priority_gt_one_count, + sub5.attempt_failed_count +FROM + ( + SELECT + count(t2.id) as pending_count + FROM tts_inference_jobs AS t2 + WHERE t2.status = "pending" + ) as sub2, + ( + SELECT + count(t3.id) as pending_priority_nonzero_count + FROM tts_inference_jobs AS t3 + WHERE t3.status = "pending" + AND t3.priority_level > 0 + ) as sub3, + ( + SELECT + count(t4.id) as pending_priority_gt_one_count + FROM tts_inference_jobs AS t4 + WHERE t4.status = "pending" + AND t4.priority_level > 1 + ) as sub4, + ( + SELECT + count(t5.id) as attempt_failed_count + FROM tts_inference_jobs AS t5 + WHERE t5.status = "attempt_failed" + ) as sub5 + "#, + ) + .fetch_one(&server_state.mysql_pool) + .await; + + let result : PendingCountResult = match maybe_result { + Ok(result) => result, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + // NB: Not Found for null results means nothing is pending in the queue + PendingCountResult { + seconds_since_first: 0, + pending_count: 0, + pending_priority_nonzero_count: 0, + pending_priority_gt_one_count: 0, + attempt_failed_count: 0, + } + }, + _ => { + warn!("get tts pending count error: {:?}", err); + return Err(GetTtsInferenceQueueCountError::ServerError) + } + } + }, + }; + + println!("Foo: {:?}", &result); + + let response = GetTtsInferenceQueueCountResponse { + success: true, + seconds_since_first: result.seconds_since_first, + pending_count: result.pending_count, + pending_priority_nonzero_count: result.pending_priority_nonzero_count, + pending_priority_gt_one_count: result.pending_priority_gt_one_count, + attempt_failed_count: result.attempt_failed_count, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetTtsInferenceQueueCountError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +#[derive(Serialize, Debug)] +pub struct PendingCountResult { + pub seconds_since_first: i64, + pub pending_count: i64, + pub pending_priority_nonzero_count: i64, + pub pending_priority_gt_one_count: i64, + pub attempt_failed_count: i64, +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/get_w2l_inference_queue_count.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/get_w2l_inference_queue_count.rs new file mode 100644 index 0000000000..3188055b6f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/get_w2l_inference_queue_count.rs @@ -0,0 +1,141 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct GetW2lInferenceQueueCountResponse { + pub success: bool, + pub pending_count: i64, + pub seconds_since_first: i64, +} + +#[derive(Debug, Serialize)] +pub enum GetW2lInferenceQueueCountError { + ServerError, + Unauthorized, +} + +impl ResponseError for GetW2lInferenceQueueCountError { + fn status_code(&self) -> StatusCode { + match *self { + GetW2lInferenceQueueCountError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetW2lInferenceQueueCountError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for GetW2lInferenceQueueCountError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_w2l_inference_queue_count_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetW2lInferenceQueueCountError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetW2lInferenceQueueCountError::Unauthorized); + } + }; + + // TODO: Not a good fit for this permission. + if !user_session.can_ban_users { + warn!("user is not allowed to view bans: {}", user_session.user_token); + return Err(GetW2lInferenceQueueCountError::Unauthorized); + } + + // NB: Lookup failure is Err(RowNotFound). + let maybe_result = sqlx::query_as!( + PendingCountResult, + r#" +SELECT + NOW() - t2.created_at AS seconds_since_first, + ( + SELECT + count(t1.id) as pending_count + FROM w2l_inference_jobs AS t1 + WHERE t1.status = "pending" + ) as pending_count +FROM w2l_inference_jobs AS t2 +WHERE t2.status = "pending" +ORDER BY t2.id ASC +LIMIT 1 + "#, + ) + .fetch_one(&server_state.mysql_pool) + .await; + + let result : PendingCountResult = match maybe_result { + Ok(result) => result, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + // NB: Not Found for null results means nothing is pending in the queue + PendingCountResult { + pending_count: None, + seconds_since_first: 0, + } + }, + _ => { + warn!("get w2l pending count error: {:?}", err); + return Err(GetW2lInferenceQueueCountError::ServerError) + } + } + }, + }; + + let response = GetW2lInferenceQueueCountResponse { + success: true, + pending_count: result.pending_count.unwrap_or(0), + seconds_since_first: result.seconds_since_first, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetW2lInferenceQueueCountError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +#[derive(Serialize)] +pub struct PendingCountResult { + pub pending_count: Option, + pub seconds_since_first: i64, +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/kill_tts_inference_jobs.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/kill_tts_inference_jobs.rs new file mode 100644 index 0000000000..ec4df966b2 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/kill_tts_inference_jobs.rs @@ -0,0 +1,130 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::server_state::ServerState; +use database_queries::queries::tts::tts_inference_jobs::kill_tts_inference_jobs::{kill_tts_inference_jobs, JobStatus}; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use regex::Regex; +use sqlx::MySqlPool; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Copy, Clone, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum KillAction { + /// Kill all "pending" jobs + AllPending, + /// Kill all "pending" and "attempt_failed" jobs + AllPendingAndFailed, + /// Kill "pending" jobs with priority_level = 0. + ZeroPriorityPending, +} + +// NB: ONLY MODERATORS CAN EDIT CATEGORIES. +// These are not sparse updates! +#[derive(Deserialize)] +pub struct KillTtsInferenceJobsRequest { + kill_action: KillAction, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct KillTtsInferenceJobsResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum KillTtsInferenceJobsError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for KillTtsInferenceJobsError { + fn status_code(&self) -> StatusCode { + match *self { + KillTtsInferenceJobsError::BadInput(_) => StatusCode::BAD_REQUEST, + KillTtsInferenceJobsError::NotFound => StatusCode::NOT_FOUND, + KillTtsInferenceJobsError::NotAuthorized => StatusCode::UNAUTHORIZED, + KillTtsInferenceJobsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for KillTtsInferenceJobsError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn kill_tts_inference_jobs_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + KillTtsInferenceJobsError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(KillTtsInferenceJobsError::NotAuthorized); + } + }; + + // TODO: We don't have a permission for this, so use this as a proxy permission + if !user_session.can_ban_users { + warn!("no permission to edit categories"); + return Err(KillTtsInferenceJobsError::NotAuthorized); + } + + let job_status = match request.kill_action { + KillAction::AllPending => JobStatus::AllPending, + KillAction::AllPendingAndFailed => JobStatus::AllPendingAndFailed, + KillAction::ZeroPriorityPending => JobStatus::ZeroPriorityPending, + }; + + kill_tts_inference_jobs(job_status, &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return KillTtsInferenceJobsError::ServerError; + })?; + + let response = KillTtsInferenceJobsResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| KillTtsInferenceJobsError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/mod.rs new file mode 100644 index 0000000000..0b81519ead --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/jobs/mod.rs @@ -0,0 +1,3 @@ +pub mod get_tts_inference_queue_count; +pub mod get_w2l_inference_queue_count; +pub mod kill_tts_inference_jobs; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/mod.rs new file mode 100644 index 0000000000..95f5febebb --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/mod.rs @@ -0,0 +1,8 @@ +pub mod approval; +pub mod categories; +pub mod ip_bans; +pub mod jobs; +pub mod stats; +pub mod user_bans; +pub mod user_roles; +pub mod users; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs new file mode 100644 index 0000000000..38eb438994 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/get_voice_count_stats.rs @@ -0,0 +1,137 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct GetVoiceCountStatsResponse { + pub success: bool, + pub all_voices_count_including_deleted: i64, + pub public_voices_count: i64, +} + +#[derive(Debug, Serialize)] +pub enum GetVoiceCountStatsError { + ServerError, + Unauthorized, +} + +impl ResponseError for GetVoiceCountStatsError { + fn status_code(&self) -> StatusCode { + match *self { + GetVoiceCountStatsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetVoiceCountStatsError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for GetVoiceCountStatsError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_voice_count_stats_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetVoiceCountStatsError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetVoiceCountStatsError::Unauthorized); + } + }; + + // TODO: Not a good fit for this permission. + if !user_session.can_edit_other_users_tts_models { + warn!("user is not allowed to edit user tts: {}", user_session.user_token); + return Err(GetVoiceCountStatsError::Unauthorized); + } + + // NB: Lookup failure is Err(RowNotFound). + let maybe_result = sqlx::query_as!( + VoiceStats, + r#" +SELECT public_count, all_count FROM ( + (SELECT count(*) as public_count + FROM tts_models + WHERE user_deleted_at IS NULL + AND mod_deleted_at IS NULL) as pc, + (SELECT count(*) as all_count + FROM tts_models) as ac +) + "#, + ) + .fetch_one(&server_state.mysql_pool) + .await; + + let result : VoiceStats = match maybe_result { + Ok(result) => result, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + // NB: Not Found for null results means nothing is pending in the queue + VoiceStats { + all_count: -1, + public_count: -1, + } + }, + _ => { + warn!("get tts pending count error: {:?}", err); + return Err(GetVoiceCountStatsError::ServerError) + } + } + }, + }; + + let response = GetVoiceCountStatsResponse { + success: true, + all_voices_count_including_deleted: result.all_count, + public_voices_count: result.public_count, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetVoiceCountStatsError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +#[derive(Serialize)] +pub struct VoiceStats { + pub public_count: i64, + pub all_count: i64, +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs new file mode 100644 index 0000000000..a2441d1780 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/stats/mod.rs @@ -0,0 +1 @@ +pub mod get_voice_count_stats; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/ban_user.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/ban_user.rs new file mode 100644 index 0000000000..9683d99b9b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/ban_user.rs @@ -0,0 +1,163 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use database_queries::queries::users::get_user_profile_by_username::get_user_profile_by_username; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct BanUserRequest { + pub username: String, + pub mod_notes: String, + pub is_banned: bool, +} + +#[derive(Serialize)] +pub struct BanUserSuccessResponse { + pub success: bool, +} + +#[derive(Serialize, Debug)] +pub struct BanUserErrorResponse { + pub success: bool, + pub error_type: BanUserErrorType, + pub error_message: String, +} + +#[derive(Copy, Clone, Debug, Serialize)] +pub enum BanUserErrorType { + ServerError, + Unauthorized, + UserNotFound, +} + +impl BanUserErrorResponse { + fn unauthorized() -> Self { + Self { + success: false, + error_type: BanUserErrorType::Unauthorized, + error_message: "unauthorized".to_string() + } + } + fn server_error() -> Self { + Self { + success: false, + error_type: BanUserErrorType::ServerError, + error_message: "server error".to_string() + } + } + fn not_found() -> Self { + Self { + success: false, + error_type: BanUserErrorType::UserNotFound, + error_message: "user not found".to_string() + } + } +} + +impl ResponseError for BanUserErrorResponse { + fn status_code(&self) -> StatusCode { + match self.error_type { + BanUserErrorType::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + BanUserErrorType::Unauthorized => StatusCode::UNAUTHORIZED, + BanUserErrorType::UserNotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for BanUserErrorResponse { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.error_type) + } +} + +pub async fn ban_user_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + BanUserErrorResponse::server_error() + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(BanUserErrorResponse::unauthorized()); + } + }; + + if !user_session.can_ban_users { + warn!("user is not allowed to add bans: {}", user_session.user_token); + return Err(BanUserErrorResponse::unauthorized()); + } + + let username_lower = request.username.to_lowercase(); + + let user_lookup_result = get_user_profile_by_username( + &username_lower, &server_state.mysql_pool).await; + + let user_profile = match user_lookup_result { + Ok(Some(result)) => result, + Ok(None) => return Err(BanUserErrorResponse::not_found()), + Err(err) => { + warn!("lookup error: {:?}", err); + return Err(BanUserErrorResponse::server_error()); + } + }; + + let query_result = sqlx::query!( + r#" +UPDATE users +SET + is_banned = ?, + maybe_mod_comments = ?, + maybe_mod_user_token = ?, + version = version + 1 + +WHERE users.token = ? +LIMIT 1 + "#, + &request.is_banned, + &request.mod_notes, + &user_session.user_token, + &user_profile.user_token + ) + .execute(&server_state.mysql_pool) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Add IP ban DB error: {:?}", err); + return Err(BanUserErrorResponse::server_error()); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/list_banned_users.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/list_banned_users.rs new file mode 100644 index 0000000000..310c64a9a9 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/list_banned_users.rs @@ -0,0 +1,121 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct ListBannedUsersRequest { + // TODO: Implement the endpoint +} + +#[derive(Serialize)] +pub struct ListBannedUsersResponse { + // TODO: Implement the endpoint +} + +#[derive(Serialize, Debug)] +pub struct ListBannedUsersErrorResponse { + pub success: bool, + pub error_type: ListBannedUsersErrorType, + pub error_message: String, +} + +#[derive(Copy, Clone, Debug, Serialize)] +pub enum ListBannedUsersErrorType { + BadInput, + ServerError, + Unauthorized, + UserNotFound, +} + +impl ListBannedUsersErrorResponse { + fn unauthorized() -> Self { + Self { + success: false, + error_type: ListBannedUsersErrorType::Unauthorized, + error_message: "unauthorized".to_string() + } + } + fn server_error() -> Self { + Self { + success: false, + error_type: ListBannedUsersErrorType::ServerError, + error_message: "server error".to_string() + } + } + fn bad_request(error_message: &str) -> Self { + Self { + success: false, + error_type: ListBannedUsersErrorType::BadInput, + error_message: error_message.to_string() + } + } +} + +impl ResponseError for ListBannedUsersErrorResponse { + fn status_code(&self) -> StatusCode { + match self.error_type { + ListBannedUsersErrorType::BadInput => StatusCode::BAD_REQUEST, + ListBannedUsersErrorType::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + ListBannedUsersErrorType::Unauthorized => StatusCode::UNAUTHORIZED, + ListBannedUsersErrorType::UserNotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for ListBannedUsersErrorResponse { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.error_type) + } +} + +pub async fn list_banned_users_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListBannedUsersErrorResponse::server_error() + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListBannedUsersErrorResponse::unauthorized()); + } + }; + + if !user_session.can_ban_users { + warn!("user is not allowed to add bans: {}", user_session.user_token); + return Err(ListBannedUsersErrorResponse::unauthorized()); + } + + // TODO: Implement the endpoint + unimplemented!() +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/mod.rs new file mode 100644 index 0000000000..eec5556ade --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_bans/mod.rs @@ -0,0 +1,2 @@ +pub mod ban_user; +pub mod list_banned_users; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/list_roles.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/list_roles.rs new file mode 100644 index 0000000000..63bf3f01a7 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/list_roles.rs @@ -0,0 +1,116 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::users::user_roles::list_user_roles::{UserRoleForList, list_user_roles}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct ListUserRolesResponse { + pub success: bool, + pub user_roles: Vec, +} + +#[derive(Debug)] +pub enum ListUserRolesError { + BadInput(String), + ServerError, + Unauthorized, +} + +impl ResponseError for ListUserRolesError { + fn status_code(&self) -> StatusCode { + match *self { + ListUserRolesError::BadInput(_) => StatusCode::BAD_REQUEST, + ListUserRolesError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + ListUserRolesError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListUserRolesError::BadInput(reason) => reason.to_string(), + ListUserRolesError::ServerError => "server error".to_string(), + ListUserRolesError::Unauthorized => "unauthorized".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListUserRolesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_user_roles_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListUserRolesError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListUserRolesError::Unauthorized); + } + }; + + // TODO: Add new permission for this. + if !user_session.can_ban_users { + warn!("user is not allowed to view user roles: {}", user_session.user_token); + return Err(ListUserRolesError::Unauthorized); + } + + let maybe_user_roles = + list_user_roles(&server_state.mysql_pool).await; + + let user_roles = match maybe_user_roles { + Ok(results) => results, + Err(e) => { + warn!("Error querying user roles: {:?}", e); + return Err(ListUserRolesError::ServerError); + } + }; + + let response = ListUserRolesResponse { + success: true, + user_roles, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListUserRolesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/list_staff.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/list_staff.rs new file mode 100644 index 0000000000..37216b3ec2 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/list_staff.rs @@ -0,0 +1,146 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + + +#[derive(Serialize)] +pub struct ListStaffResponse { + pub success: bool, + pub staff: Vec, +} + +#[derive(Serialize)] +pub struct StaffRecordForList { + pub user_token: String, + pub username: String, + pub display_name: String, + pub user_role_slug: String, + pub user_role_name: String, +} + +#[derive(Debug, Display)] +pub enum ListStaffError { + BadInput(String), + ServerError, + Unauthorized, +} + +impl ResponseError for ListStaffError { + fn status_code(&self) -> StatusCode { + match *self { + ListStaffError::BadInput(_) => StatusCode::BAD_REQUEST, + ListStaffError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + ListStaffError::Unauthorized => StatusCode::UNAUTHORIZED, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListStaffError::BadInput(reason) => reason.to_string(), + ListStaffError::ServerError => "server error".to_string(), + ListStaffError::Unauthorized => "unauthorized".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn list_staff_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListStaffError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListStaffError::Unauthorized); + } + }; + + // TODO: This is not the correct permission. + if !user_session.can_ban_users { + warn!("user is not allowed to delete bans: {}", user_session.user_token); + return Err(ListStaffError::Unauthorized); + } + + // NB: Lookup failure is Err(RowNotFound). + let maybe_results = sqlx::query_as!( + StaffRecordForList, + r#" +SELECT + users.token as user_token, + users.username, + users.display_name, + user_roles.slug as user_role_slug, + user_roles.name as user_role_name +FROM + users +JOIN user_roles + ON users.user_role_slug = user_roles.slug +WHERE + user_roles.slug != 'user' + "#, + ) + .fetch_all(&server_state.mysql_pool) + .await; + + let results : Vec = match maybe_results { + Ok(results) => { + info!("Results length: {}", results.len()); + results + }, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + Vec::new() + }, + _ => { + warn!("list staff db error: {:?}", err); + return Err(ListStaffError::ServerError); + } + } + } + }; + + let response = ListStaffResponse { + success: true, + staff: results, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListStaffError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/mod.rs new file mode 100644 index 0000000000..8e10abf233 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/mod.rs @@ -0,0 +1,3 @@ +pub mod list_staff; +pub mod list_roles; +pub mod set_user_role; diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/set_user_role.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/set_user_role.rs new file mode 100644 index 0000000000..acd6c34dfa --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/user_roles/set_user_role.rs @@ -0,0 +1,155 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::users::get_user_profile_by_username::get_user_profile_by_username; +use database_queries::queries::users::user_roles::list_user_roles::list_user_roles; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct SetUserRolePathInfo { + username: String, +} + +#[derive(Deserialize)] +pub struct SetUserRoleRequest { + user_role_slug: String, +} + +#[derive(Debug)] +pub enum SetUserRoleError { + BadInput(String), + NotAuthorized, + NotFound, + ServerError, +} + +impl ResponseError for SetUserRoleError { + fn status_code(&self) -> StatusCode { + match *self { + SetUserRoleError::BadInput(_) => StatusCode::BAD_REQUEST, + SetUserRoleError::NotAuthorized => StatusCode::UNAUTHORIZED, + SetUserRoleError::NotFound => StatusCode::NOT_FOUND, + SetUserRoleError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + SetUserRoleError::BadInput(reason) => reason.to_string(), + SetUserRoleError::NotAuthorized => "unauthorized".to_string(), + SetUserRoleError::NotFound => "not found".to_string(), + SetUserRoleError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for SetUserRoleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn set_user_role_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + SetUserRoleError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(SetUserRoleError::NotAuthorized); + } + }; + + // TODO: This is not the correct permission + if !user_session.can_ban_users { + warn!("user is not allowed to change user roles: {}", user_session.user_token); + return Err(SetUserRoleError::NotAuthorized); + } + + // TODO: This is lazy and inefficient + let user_roles = list_user_roles(&server_state.mysql_pool) + .await + .map_err(|err| { + warn!("error listing roles: {:?}", err); + return SetUserRoleError::ServerError; + })?; + + let role_exists = user_roles.into_iter() + .find(|user_role| user_role.slug == request.user_role_slug) + .is_some(); + + if !role_exists { + return Err(SetUserRoleError::BadInput("invalid user role".to_string())); + } + + let user_lookup_result = + get_user_profile_by_username(&path.username, &server_state.mysql_pool).await; + + let target_user = match user_lookup_result { + Ok(Some(result)) => result, + Ok(None) => return Err(SetUserRoleError::NotFound), + Err(err) => { + warn!("lookup error: {:?}", err); + return Err(SetUserRoleError::ServerError); + } + }; + + let query_result = sqlx::query!( + r#" +UPDATE users +SET + user_role_slug = ? +WHERE + token = ? +LIMIT 1 + "#, + &request.user_role_slug, + &target_user.user_token, + ) + .execute(&server_state.mysql_pool) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("unable to update user role: {:?}", err); + return Err(SetUserRoleError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/users/list_users.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/users/list_users.rs new file mode 100644 index 0000000000..07e1c22a68 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/users/list_users.rs @@ -0,0 +1,146 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use database_queries::queries::users::list_users_query_builder::{UserForList, ListUsersQueryBuilder}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct ListUsersRequest { + // TODO: Implement the endpoint +} + +#[derive(Serialize)] +pub struct ListUsersResponse { + success: bool, + users: Vec +} + +#[derive(Serialize, Debug)] +pub struct ListUsersErrorResponse { + pub success: bool, + pub error_type: ListUsersErrorType, + pub error_message: String, +} + +#[derive(Copy, Clone, Debug, Serialize)] +pub enum ListUsersErrorType { + BadInput, + ServerError, + Unauthorized, + UserNotFound, +} + +impl ListUsersErrorResponse { + fn unauthorized() -> Self { + Self { + success: false, + error_type: ListUsersErrorType::Unauthorized, + error_message: "unauthorized".to_string() + } + } + fn server_error() -> Self { + Self { + success: false, + error_type: ListUsersErrorType::ServerError, + error_message: "server error".to_string() + } + } + fn bad_request(error_message: &str) -> Self { + Self { + success: false, + error_type: ListUsersErrorType::BadInput, + error_message: error_message.to_string() + } + } +} + +impl ResponseError for ListUsersErrorResponse { + fn status_code(&self) -> StatusCode { + match self.error_type { + ListUsersErrorType::BadInput => StatusCode::BAD_REQUEST, + ListUsersErrorType::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + ListUsersErrorType::Unauthorized => StatusCode::UNAUTHORIZED, + ListUsersErrorType::UserNotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for ListUsersErrorResponse { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.error_type) + } +} + +pub async fn list_users_handler( + http_request: HttpRequest, + request: web::Query, + server_state: web::Data> +) -> Result { + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListUsersErrorResponse::server_error() + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListUsersErrorResponse::unauthorized()); + } + }; + + // TODO: Not the correct permission + if !user_session.can_ban_users { + warn!("user is not allowed to add bans: {}", user_session.user_token); + return Err(ListUsersErrorResponse::unauthorized()); + } + + let query_results = ListUsersQueryBuilder::new() + .sort_ascending(true) + .query_for_page(&server_state.mysql_pool) + .await; + + let results_page = match query_results { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(ListUsersErrorResponse::server_error()); + } + }; + + let response = ListUsersResponse { + success: true, + users: results_page.users, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListUsersErrorResponse::server_error())?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/moderation/users/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/moderation/users/mod.rs new file mode 100644 index 0000000000..6589194aa5 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/moderation/users/mod.rs @@ -0,0 +1 @@ +pub mod list_users; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/service/health_check_handler.rs b/crates/service/storyteller_web/src/http_server/endpoints/service/health_check_handler.rs new file mode 100644 index 0000000000..a6e8353236 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/service/health_check_handler.rs @@ -0,0 +1,97 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc, NaiveDateTime}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::complex_models::event_responses::EventResponse; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use lexical_sort::natural_lexical_cmp; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct HealthCheckResponse { + pub success: bool, + pub is_healthy: bool, + pub last_db_time: Option, + pub healthy_check_consecutive_count: Option, + pub unhealthy_check_consecutive_count: Option, +} + + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum HealthCheckError { + ServerError, +} + +impl ResponseError for HealthCheckError { + fn status_code(&self) -> StatusCode { + match *self { + HealthCheckError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for HealthCheckError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn get_health_check_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let health_check_status = server_state.health_check_status.get_health_check_status() + .map_err(|e| { + error!("Error serving health check status: {:?}", e); + HealthCheckError::ServerError + })?; + + let is_healthy = health_check_status.is_healthy; + + let response = HealthCheckResponse { + success: true, + is_healthy: health_check_status.is_healthy, + last_db_time: health_check_status.last_db_time, + healthy_check_consecutive_count: health_check_status.healthy_check_consecutive_count, + unhealthy_check_consecutive_count: health_check_status.unhealthy_check_consecutive_count, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| HealthCheckError::ServerError)?; + + if is_healthy { + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) + } else { + Ok(HttpResponse::InternalServerError() + .content_type("application/json") + .body(body)) + } +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/service/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/service/mod.rs new file mode 100644 index 0000000000..f6eeac1ec2 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/service/mod.rs @@ -0,0 +1 @@ +pub mod health_check_handler; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/delete_tts_model.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/delete_tts_model.rs new file mode 100644 index 0000000000..398392ad0b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/delete_tts_model.rs @@ -0,0 +1,194 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::delete_role_disambiguation::DeleteRole; +use crate::util::delete_role_disambiguation::delete_role_disambiguation; +use database_queries::queries::tts::tts_models::delete_tts_model_various_scopes::delete_tts_model_as_mod; +use database_queries::queries::tts::tts_models::delete_tts_model_various_scopes::delete_tts_model_as_user; +use database_queries::queries::tts::tts_models::delete_tts_model_various_scopes::undelete_tts_model_as_mod; +use database_queries::queries::tts::tts_models::delete_tts_model_various_scopes::undelete_tts_model_as_user; +use database_queries::queries::tts::tts_models::get_tts_model::get_tts_model_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct DeleteTtsModelPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct DeleteTtsModelRequest { + set_delete: bool, + /// NB: this is only to disambiguate when a user is both a mod and an author. + as_mod: Option, +} + +#[derive(Debug)] +pub enum DeleteTtsModelError { + BadInput(String), + NotAuthorized, + NotFound, + ServerError, +} + +impl ResponseError for DeleteTtsModelError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteTtsModelError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteTtsModelError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteTtsModelError::NotFound => StatusCode::NOT_FOUND, + DeleteTtsModelError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + DeleteTtsModelError::BadInput(reason) => reason.to_string(), + DeleteTtsModelError::NotAuthorized => "unauthorized".to_string(), + DeleteTtsModelError::NotFound => "not found".to_string(), + DeleteTtsModelError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DeleteTtsModelError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn delete_tts_model_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data> +) -> Result { + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteTtsModelError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteTtsModelError::NotAuthorized); + } + }; + + // NB: First permission check. + // Only mods should see deleted models (both user_* and mod_* deleted). + let is_mod = user_session.can_delete_other_users_tts_models; + + let model_query_result = get_tts_model_by_token( + &path.token, + is_mod, + &server_state.mysql_pool, + ).await; + + let tts_model = match model_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(DeleteTtsModelError::ServerError); + } + Ok(None) => return Err(DeleteTtsModelError::NotFound), + Ok(Some(model)) => model, + }; + + // NB: Second set of permission checks + let is_author = &tts_model.creator_user_token == &user_session.user_token; + + if !is_author && !is_mod { + warn!("user is not allowed to delete models: {}", user_session.user_token); + return Err(DeleteTtsModelError::NotAuthorized); + } + + if !is_mod { + if tts_model.is_locked_from_user_modification || tts_model.is_locked_from_use { + warn!("user is not allowed to delete models (locked): {}", user_session.user_token); + return Err(DeleteTtsModelError::NotAuthorized); + } + } + + let ip_address = get_request_ip(&http_request); + + let delete_role = delete_role_disambiguation(is_mod, is_author, request.as_mod); + + let query_result = if request.set_delete { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to delete model: {}", user_session.user_token); + return Err(DeleteTtsModelError::NotAuthorized); + } + DeleteRole::AsUser => { + delete_tts_model_as_user( + &path.token, + &ip_address, + &server_state.mysql_pool, + ).await + } + DeleteRole::AsMod => { + delete_tts_model_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + } else { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to undelete model: {}", user_session.user_token); + return Err(DeleteTtsModelError::NotAuthorized); + } + DeleteRole::AsUser => { + // NB: Technically only mods can see their own inference_results here + undelete_tts_model_as_user( + &path.token, + &ip_address, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + undelete_tts_model_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Delete tts model DB error: {:?}", err); + return Err(DeleteTtsModelError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/delete_tts_result.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/delete_tts_result.rs new file mode 100644 index 0000000000..2f4c5b45a0 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/delete_tts_result.rs @@ -0,0 +1,191 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::delete_role_disambiguation::DeleteRole; +use crate::util::delete_role_disambiguation::delete_role_disambiguation; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::tts::tts_results::delete_tts_result_various_scopes::delete_tts_inference_result_as_mod; +use database_queries::queries::tts::tts_results::delete_tts_result_various_scopes::delete_tts_inference_result_as_user; +use database_queries::queries::tts::tts_results::delete_tts_result_various_scopes::undelete_tts_inference_result_as_mod; +use database_queries::queries::tts::tts_results::delete_tts_result_various_scopes::undelete_tts_inference_result_as_user; +use database_queries::queries::tts::tts_results::query_tts_result::select_tts_result_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct DeleteTtsInferenceResultPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct DeleteTtsInferenceResultRequest { + set_delete: bool, + /// NB: this is only to disambiguate when a user is both a mod and an author. + as_mod: Option, +} + +#[derive(Debug)] +pub enum DeleteTtsInferenceResultError { + BadInput(String), + NotAuthorized, + NotFound, + ServerError, +} + +impl ResponseError for DeleteTtsInferenceResultError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteTtsInferenceResultError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteTtsInferenceResultError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteTtsInferenceResultError::NotFound => StatusCode::NOT_FOUND, + DeleteTtsInferenceResultError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + DeleteTtsInferenceResultError::BadInput(reason) => reason.to_string(), + DeleteTtsInferenceResultError::NotAuthorized => "unauthorized".to_string(), + DeleteTtsInferenceResultError::NotFound => "not found".to_string(), + DeleteTtsInferenceResultError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DeleteTtsInferenceResultError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn delete_tts_inference_result_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data> +) -> Result { + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteTtsInferenceResultError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteTtsInferenceResultError::NotAuthorized); + } + }; + + // NB: First permission check. + // Only mods should see deleted models (both user_* and mod_* deleted). + let is_mod_that_can_see_deleted = user_session.can_delete_other_users_tts_results; + + let inference_result_query_result = select_tts_result_by_token( + &path.token, + is_mod_that_can_see_deleted, + &server_state.mysql_pool, + ).await; + + let tts_inference_result = match inference_result_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(DeleteTtsInferenceResultError::ServerError); + } + Ok(None) => return Err(DeleteTtsInferenceResultError::NotFound), + Ok(Some(inference_result)) => inference_result, + }; + + // NB: Second set of permission checks + let is_author = tts_inference_result.maybe_creator_user_token + .as_deref() + .map(|creator_token| creator_token == &user_session.user_token) + .unwrap_or(false); + + let is_mod = user_session.can_delete_other_users_tts_results; + + if !is_author && !is_mod { + warn!("user is not allowed to delete inference results: {}", user_session.user_token); + return Err(DeleteTtsInferenceResultError::NotAuthorized); + } + + let delete_role = delete_role_disambiguation(is_mod, is_author, request.as_mod); + + let query_result = if request.set_delete { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to delete inference results: {}", user_session.user_token); + return Err(DeleteTtsInferenceResultError::NotAuthorized); + } + DeleteRole::AsUser => { + delete_tts_inference_result_as_user( + &path.token, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + delete_tts_inference_result_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + } else { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to undelete inference results: {}", user_session.user_token); + return Err(DeleteTtsInferenceResultError::NotAuthorized); + } + DeleteRole::AsUser => { + // NB: Technically only mods can see their own inference_results here + undelete_tts_inference_result_as_user( + &path.token, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + undelete_tts_inference_result_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update tts mod approval status DB error: {:?}", err); + return Err(DeleteTtsInferenceResultError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs new file mode 100644 index 0000000000..28d2c215fa --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs @@ -0,0 +1,400 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use container_common::i18n::supported_languages_for_models::{is_valid_language_for_models, get_canonicalized_language_tag_for_model}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::column_types::vocoder_type::VocoderType; +use database_queries::queries::tts::tts_models::get_tts_model::get_tts_model_by_token; +use language_tags::{LanguageTag, ParseError}; +use log::{info, warn, log, error}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use user_input_common::check_for_slurs::contains_slurs; + +const DEFAULT_IETF_LANGUAGE_TAG : &'static str = "en-US"; +const DEFAULT_IETF_PRIMARY_LANGUAGE_SUBTAG : &'static str = "en"; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct EditTtsModelPathInfo { + model_token: String, +} + +#[derive(Deserialize)] +pub struct EditTtsModelRequest { + // ========== Author + Moderator options ========== + + pub title: Option, + pub description_markdown: Option, + pub creator_set_visibility: Option, + pub maybe_default_pretrained_vocoder: Option, + + // NB: We calculate 'ietf_primary_language_subtag' from this value. + pub ietf_language_tag: Option, + + //pub updatable_slug: Option, + //pub tts_model_type: Option, + //pub text_preprocessing_algorithm: Option, + //pub vocoder_token: Option, + + // ========== Moderator options (protection) ========== + + pub is_public_listing_approved: Option, + pub is_locked_from_user_modification: Option, + pub is_locked_from_use: Option, + pub maybe_mod_comments: Option, + + // ========== Moderator options (front page, Discord, Twitch, etc.) ========== + + pub is_front_page_featured: Option, + pub is_twitch_featured: Option, + + // NB: We take "empty string" to mean removal. + pub maybe_suggested_unique_bot_command: Option, +} + +#[derive(Debug)] +pub enum EditTtsModelError { + BadInput(String), + NotAuthorized, + ModelNotFound, + ServerError, +} + +impl ResponseError for EditTtsModelError { + fn status_code(&self) -> StatusCode { + match *self { + EditTtsModelError::BadInput(_) => StatusCode::BAD_REQUEST, + EditTtsModelError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditTtsModelError::ModelNotFound => StatusCode::NOT_FOUND, + EditTtsModelError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EditTtsModelError::BadInput(reason) => reason.to_string(), + EditTtsModelError::NotAuthorized=> "unauthorized".to_string(), + EditTtsModelError::ModelNotFound => "not found".to_string(), + EditTtsModelError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditTtsModelError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn edit_tts_model_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditTtsModelError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditTtsModelError::NotAuthorized); + } + }; + + // NB: First permission check. + // Only mods should see deleted models (both user_* and mod_* deleted). + let is_mod_that_can_see_deleted = user_session.can_delete_other_users_tts_models; + + let model_lookup_result = get_tts_model_by_token( + &path.model_token, + is_mod_that_can_see_deleted, + &server_state.mysql_pool).await; + + let model_record = match model_lookup_result { + Ok(Some(result)) => { + info!("Found model: {}", result.model_token); + result + }, + Ok(None) => { + warn!("could not find model"); + return Err(EditTtsModelError::ModelNotFound); + }, + Err(err) => { + warn!("error looking up model: {:?}", err); + return Err(EditTtsModelError::ModelNotFound); + }, + }; + + // NB: Second set of permission checks + let is_author = model_record.creator_user_token == user_session.user_token; + let is_mod = user_session.can_edit_other_users_tts_models ; + + if !is_author && !is_mod { + warn!("user is not allowed to edit model: {}", user_session.user_token); + return Err(EditTtsModelError::NotAuthorized); + } + + if !is_mod { + if model_record.is_locked_from_user_modification || model_record.is_locked_from_use { + return Err(EditTtsModelError::NotAuthorized); + } + } + + // ============================================= + // Author + Mod fields. + // These fields must be present on all requests. + // ============================================= + + let mut title = None; + let mut description_markdown = None; + let mut description_html = None; + let mut ietf_language_tag = None; + let mut ietf_primary_language_subtag = None; + + let mut creator_set_visibility = RecordVisibility::Public; + let mut maybe_default_pretrained_vocoder = + model_record.maybe_default_pretrained_vocoder + .clone(); + + if let Some(payload) = request.title.as_deref() { + if contains_slurs(payload) { + return Err(EditTtsModelError::BadInput("title contains slurs".to_string())); + } + + title = Some(payload.to_string()); + } + + if let Some(markdown) = request.description_markdown.as_deref() { + if contains_slurs(markdown) { + return Err(EditTtsModelError::BadInput("description contains slurs".to_string())); + } + + let markdown = markdown.trim().to_string(); + let html = markdown_to_html(&markdown); + + description_markdown = Some(markdown); + description_html = Some(html); + } + + if let Some(tag) = request.ietf_language_tag.as_deref() { + // eg. en, en-US, es-419, ja-JP, etc. + let maybe_full_canonical_tag = get_canonicalized_language_tag_for_model(tag); + + // eg. en, es, ja, etc. + let maybe_primary_language_subtag = maybe_full_canonical_tag + .map(|t| LanguageTag::parse(t) + .map(|language_tag| language_tag.primary_language().to_string()) + ) + .transpose() + .map_err(|e| { + error!("Error parsing language tag '{}': {:?}", tag, e); + EditTtsModelError::BadInput("bad locale string".to_string()) + })?; + + if let Some(full_tag) = maybe_full_canonical_tag { + if let Some(primary_subtag) = maybe_primary_language_subtag.as_deref() { + ietf_language_tag = Some(full_tag.to_string()); + ietf_primary_language_subtag = Some(primary_subtag.to_string()); + } + } + } + + let mut ietf_language_tag = ietf_language_tag + .unwrap_or(DEFAULT_IETF_LANGUAGE_TAG.to_string()); + + let mut ietf_primary_language_subtag = ietf_primary_language_subtag + .unwrap_or(DEFAULT_IETF_PRIMARY_LANGUAGE_SUBTAG.to_string()); + + if let Some(visibility) = request.creator_set_visibility.as_deref() { + creator_set_visibility = RecordVisibility::from_str(visibility) + .map_err(|_| EditTtsModelError::BadInput("bad record visibility".to_string()))?; + } + + if let Some(vocoder) = request.maybe_default_pretrained_vocoder { + maybe_default_pretrained_vocoder = Some(vocoder); + } + + let ip_address = get_request_ip(&http_request); + + let query_result = if is_author { + // We need to store the IP address details. + sqlx::query!( + r#" +UPDATE tts_models +SET + maybe_default_pretrained_vocoder = ?, + title = ?, + description_markdown = ?, + description_rendered_html = ?, + ietf_language_tag = ?, + ietf_primary_language_subtag = ?, + creator_set_visibility = ?, + creator_ip_address_last_update = ?, + version = version + 1 +WHERE token = ? +LIMIT 1 + "#, + maybe_default_pretrained_vocoder.map(|v| v.to_str()), + &title, + &description_markdown, + &description_html, + &ietf_language_tag, + &ietf_primary_language_subtag, + &creator_set_visibility.to_str(), + &ip_address, + &model_record.model_token, + ) + .execute(&server_state.mysql_pool) + .await + } else { + // We need to store the moderator details. + sqlx::query!( + r#" +UPDATE tts_models +SET + maybe_default_pretrained_vocoder = ?, + title = ?, + description_markdown = ?, + description_rendered_html = ?, + ietf_language_tag = ?, + ietf_primary_language_subtag = ?, + creator_set_visibility = ?, + maybe_mod_user_token = ?, + version = version + 1 +WHERE token = ? +LIMIT 1 + "#, + maybe_default_pretrained_vocoder.map(|v| v.to_str()), + &title, + &description_markdown, + &description_html, + &ietf_language_tag, + &ietf_primary_language_subtag, + &creator_set_visibility.to_str(), + &user_session.user_token, + &model_record.model_token, + ) + .execute(&server_state.mysql_pool) + .await + }; + + // ============================================= + // Mod-only fields. + // ============================================= + + // TODO: This is lazy and suboptimal af to UPDATE again. + // The reason we're doing this is because `sqlx` only does static type checking of queries + // with string literals. It does not support dynamic query building, thus the PREDICATES + // MUST BE HELD CONSTANT (at least in type signature). :( + if is_mod { + update_mod_details( + &request, + &user_session.user_token, + &model_record.model_token, + &server_state.mysql_pool + ).await?; + } + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update W2L model edit DB error: {:?}", err); + return Err(EditTtsModelError::ServerError); + } + }; + + Ok(simple_json_success()) +} + +async fn update_mod_details( + request: &Json, + mod_user_token: &str, + model_token: &str, + mysql_pool: &MySqlPool +) -> Result<(), EditTtsModelError> { + + let is_public_listing_approved= request.is_public_listing_approved.unwrap_or(false); + let is_locked_from_user_modification = request.is_locked_from_user_modification.unwrap_or(false); + let is_locked_from_use = request.is_locked_from_use.unwrap_or(false); + + let is_front_page_featured = request.is_front_page_featured.unwrap_or(false); + let is_twitch_featured = request.is_twitch_featured.unwrap_or(false); + + // Commands must be non-empty, alphanumeric, and lowercase. + let mut maybe_suggested_unique_bot_command = None; + if let Some(command) = request.maybe_suggested_unique_bot_command.as_deref() { + // Clear empty commands to null. + let command = command.trim().to_lowercase(); + + let is_alphanumeric = command.chars().all(|c| c.is_ascii_alphanumeric()); + + if !command.is_empty() && is_alphanumeric { + maybe_suggested_unique_bot_command = Some(command); + } + } + + let query_result = sqlx::query!( + r#" +UPDATE tts_models +SET + is_public_listing_approved = ?, + is_locked_from_user_modification = ?, + is_locked_from_use = ?, + maybe_mod_comments = ?, + maybe_mod_user_token = ?, + is_front_page_featured = ?, + is_twitch_featured = ?, + maybe_suggested_unique_bot_command = ?, + version = version + 1 +WHERE token = ? +LIMIT 1 + "#, + is_public_listing_approved, + is_locked_from_user_modification, + is_locked_from_use, + request.maybe_mod_comments, + mod_user_token, + is_front_page_featured, + is_twitch_featured, + maybe_suggested_unique_bot_command, + model_token + ) + .execute(mysql_pool) + .await; + + match query_result { + Ok(_) => Ok(()), + Err(err) => { + warn!("Update TTS model (mod details) DB error: {:?}", err); + Err(EditTtsModelError::ServerError) + } + } +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_result.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_result.rs new file mode 100644 index 0000000000..15aff691b5 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_result.rs @@ -0,0 +1,189 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::queries::tts::tts_results::query_tts_result::select_tts_result_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct EditTtsResultPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct EditTtsResultRequest { + // ========== Author + Moderator options ========== + pub creator_set_visibility: Option, +} + +#[derive(Debug)] +pub enum EditTtsResultError { + BadInput(String), + NotAuthorized, + ResultNotFound, + ServerError, +} + +impl ResponseError for EditTtsResultError { + fn status_code(&self) -> StatusCode { + match *self { + EditTtsResultError::BadInput(_) => StatusCode::BAD_REQUEST, + EditTtsResultError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditTtsResultError::ResultNotFound => StatusCode::NOT_FOUND, + EditTtsResultError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EditTtsResultError::BadInput(reason) => reason.to_string(), + EditTtsResultError::NotAuthorized=> "unauthorized".to_string(), + EditTtsResultError::ResultNotFound => "not found".to_string(), + EditTtsResultError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditTtsResultError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn edit_tts_inference_result_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditTtsResultError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditTtsResultError::NotAuthorized); + } + }; + + // NB: Moderators can see deleted results. + // Original creators cannot see them (unless they're moderators!) + let show_deleted_results = user_session.can_delete_other_users_tts_results; + + // Moderators get to see all the fields. + let is_moderator = user_session.can_delete_other_users_tts_results + || user_session.can_edit_other_users_tts_models; // TODO: Not an exact permission fit + + let inference_result_query_result = select_tts_result_by_token( + &path.token, + show_deleted_results, + &server_state.mysql_pool + ).await; + + let mut inference_result = match inference_result_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(EditTtsResultError::ServerError); + } + Ok(None) => return Err(EditTtsResultError::ResultNotFound), + Ok(Some(inference_result)) => inference_result, + }; + + // NB: Second set of permission checks + let mut is_author = false; + if let Some(creator_user_token) = inference_result.maybe_creator_user_token.as_deref() { + is_author = creator_user_token == &user_session.user_token; + } + + if !is_author && !is_moderator { + warn!("user is not allowed to edit result: {}", user_session.user_token); + return Err(EditTtsResultError::NotAuthorized); + } + + // Author + Mod fields. + // These fields must be present on all requests. + let mut creator_set_visibility = RecordVisibility::Public; + + if let Some(visibility) = request.creator_set_visibility.as_deref() { + creator_set_visibility = RecordVisibility::from_str(visibility) + .map_err(|_| EditTtsResultError::BadInput("bad record visibility".to_string()))?; + } + + let ip_address = get_request_ip(&http_request); + + let query_result = if is_author { + // TODO: Don't update the original IP address. Create a new field. + // We need to store the IP address details. + sqlx::query!( + r#" +UPDATE tts_results +SET + creator_set_visibility = ?, + creator_ip_address = ? +WHERE token = ? +LIMIT 1 + "#, + &creator_set_visibility.to_str(), + &ip_address, + &inference_result.tts_result_token, + ) + .execute(&server_state.mysql_pool) + .await + } else { + // We need to store the moderator details. + sqlx::query!( + r#" +UPDATE tts_results +SET + creator_set_visibility = ?, + maybe_mod_user_token = ? +WHERE token = ? +LIMIT 1 + "#, + &creator_set_visibility.to_str(), + &user_session.user_token, + &inference_result.tts_result_token, + ) + .execute(&server_state.mysql_pool) + .await + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update TTS result DB error: {:?}", err); + return Err(EditTtsResultError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts.rs new file mode 100644 index 0000000000..c68791dc24 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts.rs @@ -0,0 +1,286 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::tokens::Tokens; +use log::{info, warn, log}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use user_input_common::check_for_slurs::contains_slurs; +use crate::http_server::endpoints::investor_demo::demo_cookie::request_has_demo_cookie; + +// TODO: Temporary for investor demo +const STORYTELLER_DEMO_COOKIE_NAME : &'static str = "storyteller_demo"; + +#[derive(Deserialize)] +pub struct InferTtsRequest { + uuid_idempotency_token: String, + tts_model_token: String, + inference_text: String, + creator_set_visibility: Option, + is_storyteller_demo: Option, +} + +#[derive(Serialize)] +pub struct InferTtsSuccessResponse { + pub success: bool, + pub inference_job_token: String, +} + +#[derive(Debug)] +pub enum InferTtsError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for InferTtsError { + fn status_code(&self) -> StatusCode { + match *self { + InferTtsError::BadInput(_) => StatusCode::BAD_REQUEST, + InferTtsError::NotAuthorized => StatusCode::UNAUTHORIZED, + InferTtsError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + InferTtsError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + InferTtsError::BadInput(reason) => reason.to_string(), + InferTtsError::NotAuthorized => "unauthorized".to_string(), + InferTtsError::ServerError => "server error".to_string(), + InferTtsError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for InferTtsError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn infer_tts_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + InferTtsError::ServerError + })?; + + // ==================== PRIORITY ==================== // + + // Give logged in users execution priority. + let mut priority_level = if maybe_user_session.is_some() { 1 } else { 0 }; + + // TODO/TEMP: Give investors even more priority + let mut is_investor = false; + + // TODO/TEMP: The storyteller.io website's AJAX calls will set this. + // This is just for the YCombinator demo. + match request.is_storyteller_demo { + Some(true) => { + is_investor = true; + }, + _ => {}, + }; + + // TODO/TEMP: The storyteller.io website will redirect and establish this cookie. + // This is just for the YCombinator demo. + if request_has_demo_cookie(&http_request) { + is_investor = true; + } + + if is_investor { + priority_level = 2; + } + + // ==================== RATE LIMIT ==================== // + + let mut rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.is_banned { + return Err(InferTtsError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + // TODO/TEMP + if is_investor { + rate_limiter = &server_state.redis_rate_limiters.logged_in; + } + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(InferTtsError::RateLimited); + } + + // ==================== CHECK AND PERFORM TTS ==================== // + + let mut maybe_user_token : Option = maybe_user_session + .as_ref() + .map(|user_session| user_session.user_token.to_string()); + + let inference_text = &request.inference_text.trim().to_string(); + + if let Err(reason) = validate_inference_text(&inference_text) { + return Err(InferTtsError::BadInput(reason)); + } + + if contains_slurs(&inference_text) { + return Err(InferTtsError::BadInput("text contains slurs".to_string())); + } + + // TODO(bt): CHECK DATABASE! + let model_token = request.tts_model_token.to_string(); + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + InferTtsError::ServerError + })?; + + let redis_count_key = RedisKeys::tts_model_usage_count(&model_token); + + redis.incr(&redis_count_key, 1) + .map_err(|e| { + warn!("redis error: {:?}", e); + InferTtsError::ServerError + })?; + + let ip_address = get_request_ip(&http_request); + + let maybe_user_preferred_visibility : Option = maybe_user_session + .as_ref() + .map(|user_session| user_session.preferred_tts_result_visibility); + + let set_visibility = request.creator_set_visibility + .or(maybe_user_preferred_visibility) + .unwrap_or(RecordVisibility::Public); + + // This token is returned to the client. + let job_token = Tokens::new_tts_inference_job() + .map_err(|e| { + warn!("Error creating token"); + InferTtsError::ServerError + })?; + + info!("Creating w2l inference job record..."); + + let query_result = sqlx::query!( + r#" +INSERT INTO tts_inference_jobs +SET + token = ?, + uuid_idempotency_token = ?, + + model_token = ?, + raw_inference_text = ?, + maybe_creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + priority_level = ?, + status = "pending" + "#, + &job_token, + request.uuid_idempotency_token.clone(), + model_token, + inference_text, + maybe_user_token, + ip_address, + set_visibility.to_str(), + priority_level, + ) + .execute(&server_state.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("New w2l inference job creation DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(InferTtsError::ServerError); + } + }; + + info!("new w2l inference job id: {}", record_id); + + server_state.firehose_publisher.enqueue_tts_inference(maybe_user_token.as_deref(), &job_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + InferTtsError::ServerError + })?; + + let response = InferTtsSuccessResponse { + success: true, + inference_job_token: job_token.to_string(), + }; + + let body = serde_json::to_string(&response) + .map_err(|e| InferTtsError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +pub fn validate_inference_text(text: &str) -> Result<(), String> { + if text.len() < 3 { + return Err("text is too short".to_string()); + } + + if text.len() > 1024 { + return Err("text is too long".to_string()); + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs new file mode 100644 index 0000000000..d4dddfe6c5 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/enqueue_upload_tts_model.rs @@ -0,0 +1,232 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::tokens::Tokens; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub enum TtsModelType { + /// tacotron2 + Tacotron2, + /// glowtts + GlowTts, + /// glowtts-vocodes + GlowTts_Vocodes, + /// talknet + Talknet, +} + +#[deprecated(note = "Use `RecordVisibility` instead!")] +#[derive(Deserialize)] +pub enum CreatorSetVisibility { + /// public + Public, + /// hidden + Hidden, + /// private + Private, +} + +#[derive(Deserialize)] +pub struct UploadTtsModelRequest { + idempotency_token: String, + title: String, + download_url: String, + tts_model_type: Option, + creator_set_visibility: Option, +} + +#[derive(Serialize)] +pub struct UploadTtsModelSuccessResponse { + pub success: bool, + /// This is how frontend clients can request the job execution status. + pub job_token: String, +} + +#[derive(Debug)] +pub enum UploadTtsModelError { + BadInput(String), + MustBeLoggedIn, + ServerError, + RateLimited, +} + +impl ResponseError for UploadTtsModelError { + fn status_code(&self) -> StatusCode { + match *self { + UploadTtsModelError::BadInput(_) => StatusCode::BAD_REQUEST, + UploadTtsModelError::MustBeLoggedIn => StatusCode::UNAUTHORIZED, + UploadTtsModelError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + UploadTtsModelError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + UploadTtsModelError::BadInput(reason) => reason.to_string(), + UploadTtsModelError::MustBeLoggedIn => "user must be logged in".to_string(), + UploadTtsModelError::ServerError => "server error".to_string(), + UploadTtsModelError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for UploadTtsModelError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn upload_tts_model_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) { + return Err(UploadTtsModelError::RateLimited); + } + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + UploadTtsModelError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(UploadTtsModelError::MustBeLoggedIn); + } + }; + + if let Err(reason) = validate_idempotency_token(&request.idempotency_token) { + return Err(UploadTtsModelError::BadInput(reason)); + } + + if let Err(reason) = validate_model_title(&request.title) { + return Err(UploadTtsModelError::BadInput(reason)); + } + + let ip_address = get_request_ip(&http_request); + + let uuid = request.idempotency_token.to_string(); + let title = request.title.to_string(); + let download_url = request.download_url.to_string(); + + let tts_model_type = "tacotron2".to_string(); + let creator_set_visibility = "public".to_string(); + + // This token is returned to the client. + let job_token = Tokens::new_tts_model_upload_job() + .map_err(|e| { + warn!("Error creating token"); + UploadTtsModelError::ServerError + })?; + + let query_result = sqlx::query!( + r#" +INSERT INTO tts_model_upload_jobs +SET + token = ?, + uuid_idempotency_token = ?, + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + title = ?, + tts_model_type = ?, + download_url = ?, + status = "pending" + "#, + job_token.to_string(), + uuid.to_string(), + user_session.user_token.to_string(), + ip_address.to_string(), + creator_set_visibility.to_string(), + title.to_string(), + tts_model_type, + download_url, + ) + .execute(&server_state.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("New tts model upload creation DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(UploadTtsModelError::ServerError); + } + }; + + info!("new model upload job id: {}", record_id); + + server_state.firehose_publisher.enqueue_tts_model_upload(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + UploadTtsModelError::ServerError + })?; + + let response = UploadTtsModelSuccessResponse { + success: true, + job_token: job_token.to_string(), + }; + + let body = serde_json::to_string(&response) + .map_err(|e| UploadTtsModelError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +fn validate_idempotency_token(token: &str) -> Result<(), String> { + if token.len() != 36 { + return Err("idempotency token should be 36 characters".to_string()); + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs new file mode 100644 index 0000000000..640076528a --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_inference_job_status.rs @@ -0,0 +1,219 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log, error}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2::{Pool, PooledConnection}; +use r2d2_redis::redis::{Commands, RedisError, RedisResult}; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::borrow::BorrowMut; +use std::fmt; +use std::ops::Deref; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetTtsInferenceStatusPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct TtsInferenceJobStatusForResponse { + pub job_token: String, + + /// Primary status from the database (a state machine). + pub status: String, + + /// Extra, temporary status from Redis. + /// This can denote inference progress, and the Python code can write to it. + pub maybe_extra_status_description: Option, + + pub attempt_count: u8, + + pub maybe_result_token: Option, + pub maybe_public_bucket_wav_audio_path: Option, + + pub model_token: String, + pub tts_model_type: String, + pub title: String, // Name of the TTS model + + pub raw_inference_text: String, // User text + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct GetTtsInferenceStatusSuccessResponse { + pub success: bool, + pub state: TtsInferenceJobStatusForResponse, +} + +#[derive(Debug)] +pub enum GetTtsInferenceStatusError { + ServerError, +} + +#[derive(Serialize)] +pub struct TtsInferenceJobStatusRecord { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + + pub maybe_result_token: Option, + pub maybe_public_bucket_wav_audio_path: Option, + + pub model_token: String, + pub tts_model_type: String, + pub title: String, // Name of the TTS model + + pub raw_inference_text: String, // User text + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl ResponseError for GetTtsInferenceStatusError { + fn status_code(&self) -> StatusCode { + match *self { + GetTtsInferenceStatusError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetTtsInferenceStatusError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetTtsInferenceStatusError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + + +pub async fn get_tts_inference_job_status_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_status = sqlx::query_as!( + TtsInferenceJobStatusRecord, + r#" +SELECT + jobs.token as job_token, + + jobs.status, + jobs.attempt_count, + jobs.on_success_result_token as maybe_result_token, + results.public_bucket_wav_audio_path as maybe_public_bucket_wav_audio_path, + + jobs.model_token, + tts.tts_model_type, + tts.title, + + jobs.raw_inference_text, + + jobs.created_at, + jobs.updated_at + +FROM tts_inference_jobs as jobs +JOIN tts_models as tts + ON tts.token = jobs.model_token +LEFT OUTER JOIN tts_results as results + ON jobs.on_success_result_token = results.token + +WHERE jobs.token = ? + "#, + &path.token + ) + .fetch_one(&server_state.mysql_pool) + .await; // TODO: This will return error if it doesn't exist + + let record : TtsInferenceJobStatusRecord = match maybe_status { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + warn!("tts job record not found"); + return Err(GetTtsInferenceStatusError::ServerError); + }, + _ => { + error!("tts job query error: {:?}", err); + return Err(GetTtsInferenceStatusError::ServerError); + } + } + } + }; + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + error!("redis error: {:?}", e); + GetTtsInferenceStatusError::ServerError + })?; + + let extra_status_key = RedisKeys::tts_inference_extra_status_info(&path.token); + let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { + Ok(Some(status)) => { + Some(status) + }, + Ok(None) => None, + Err(e) => { + error!("redis error: {:?}", e); + None // Fail open + }, + }; + + let record_for_response = TtsInferenceJobStatusForResponse { + job_token: record.job_token, + status: record.status, + maybe_extra_status_description, + attempt_count: record.attempt_count as u8, + maybe_result_token: record.maybe_result_token, + maybe_public_bucket_wav_audio_path: record.maybe_public_bucket_wav_audio_path, + model_token: record.model_token, + tts_model_type: record.tts_model_type, + title: record.title, + raw_inference_text: record.raw_inference_text, + created_at: record.created_at, + updated_at: record.updated_at, + }; + + let response = GetTtsInferenceStatusSuccessResponse { + success: true, + state: record_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| { + error!("error returning response: {:?}", e); + GetTtsInferenceStatusError::ServerError + })?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs new file mode 100644 index 0000000000..954e5e93cc --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs @@ -0,0 +1,138 @@ +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use anyhow::Error; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::tts::tts_models::get_tts_model::{TtsModelRecordForResponse, get_tts_model_by_token}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetTtsModelPathInfo { + token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct GetTtsModelSuccessResponse { + pub success: bool, + pub model: TtsModelRecordForResponse, +} + +// =============== Error Response =============== + +#[derive(Debug)] +pub enum GetTtsModelError { + ServerError, + NotFound, +} + +impl ResponseError for GetTtsModelError { + fn status_code(&self) -> StatusCode { + match *self { + GetTtsModelError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + GetTtsModelError::NotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetTtsModelError::ServerError => "server error".to_string(), + GetTtsModelError::NotFound => "not found".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetTtsModelError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn get_tts_model_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetTtsModelError::ServerError + })?; + + let mut show_deleted_models = false; + let mut is_moderator = false; + + if let Some(user_session) = maybe_user_session { + // NB: Moderators can see deleted models + // Original creators cannot see them (unless they're moderators!) + show_deleted_models = user_session.can_delete_other_users_tts_models; + // Moderators get to see all the fields. + is_moderator = user_session.can_delete_other_users_tts_results + || user_session.can_edit_other_users_tts_models; + } + + let model_query_result = get_tts_model_by_token( + &path.token, + show_deleted_models, + &server_state.mysql_pool + ).await; + + let mut model = match model_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(GetTtsModelError::ServerError); + } + Ok(None) => return Err(GetTtsModelError::NotFound), + Ok(Some(model)) => model, + }; + + if let Some(moderator_fields) = model.maybe_moderator_fields.as_ref() { + // NB: The moderator fields will always be present before removal + // We don't want non-mods seeing stuff made by banned users. + if moderator_fields.creator_is_banned && !is_moderator { + return Err(GetTtsModelError::NotFound); + } + } + + if !is_moderator { + model.maybe_moderator_fields = None; + } + + let response = GetTtsModelSuccessResponse { + success: true, + model, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetTtsModelError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model_use_count.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model_use_count.rs new file mode 100644 index 0000000000..a1936be3aa --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model_use_count.rs @@ -0,0 +1,91 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use r2d2_redis::redis::{Commands, RedisError}; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetTtsModelUseCountPathInfo { + model_token: String, +} + +#[derive(Serialize)] +pub struct GetTtsModelUseCountSuccessResponse { + pub success: bool, + pub count: Option, +} + +#[derive(Debug, Display)] +pub enum GetTtsModelUseCountError { + ServerError, + NotFoundError, +} + +impl ResponseError for GetTtsModelUseCountError { + fn status_code(&self) -> StatusCode { + match *self { + GetTtsModelUseCountError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + GetTtsModelUseCountError::NotFoundError => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetTtsModelUseCountError::ServerError => "server error".to_string(), + GetTtsModelUseCountError::NotFoundError => "not found".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn get_tts_model_use_count_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + GetTtsModelUseCountError::ServerError + })?; + + let redis_count_key = RedisKeys::tts_model_usage_count(&path.model_token); + + let count : Option = redis.get(&redis_count_key) + .map_err(|e| { + warn!("redis error: {:?}", e); + GetTtsModelUseCountError::ServerError + })?; + + let response = GetTtsModelUseCountSuccessResponse { + success: true, + count, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetTtsModelUseCountError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_result.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_result.rs new file mode 100644 index 0000000000..3da2574f26 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_result.rs @@ -0,0 +1,136 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::session_checker::SessionRecord; +use crate::server_state::ServerState; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::tts::tts_results::query_tts_result::TtsResultRecordForResponse; +use database_queries::queries::tts::tts_results::query_tts_result::select_tts_result_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetTtsResultPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct GetTtsResultSuccessResponse { + pub success: bool, + pub result: TtsResultRecordForResponse, +} + +#[derive(Debug)] +pub enum GetTtsResultError { + ServerError, + NotFound, +} + +impl ResponseError for GetTtsResultError { + fn status_code(&self) -> StatusCode { + match *self { + GetTtsResultError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + GetTtsResultError::NotFound=> StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetTtsResultError::ServerError => "server error".to_string(), + GetTtsResultError::NotFound => "not found".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetTtsResultError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_tts_inference_result_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetTtsResultError::ServerError + })?; + + let mut show_deleted_results = false; + let mut is_moderator = false; + + if let Some(user_session) = maybe_user_session { + // NB: Moderators can see deleted results. + // Original creators cannot see them (unless they're moderators!) + show_deleted_results = user_session.can_delete_other_users_tts_results; + // Moderators get to see all the fields. + is_moderator = user_session.can_delete_other_users_tts_results + || user_session.can_edit_other_users_tts_models; + } + + let inference_result_query_result = select_tts_result_by_token( + &path.token, + show_deleted_results, + &server_state.mysql_pool + ).await; + + let mut inference_result = match inference_result_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(GetTtsResultError::ServerError); + } + Ok(None) => return Err(GetTtsResultError::NotFound), + Ok(Some(inference_result)) => inference_result, + }; + + if let Some(moderator_fields) = inference_result.maybe_moderator_fields.as_ref() { + // NB: The moderator fields will always be present before removal + // We don't want non-mods seeing stuff made by banned users. + if (moderator_fields.model_creator_is_banned || moderator_fields.result_creator_is_banned_if_user) + && !is_moderator{ + return Err(GetTtsResultError::NotFound); + } + } + + if !is_moderator { + inference_result.maybe_moderator_fields = None; + } + + let response = GetTtsResultSuccessResponse { + success: true, + result: inference_result, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetTtsResultError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_upload_model_job_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_upload_model_job_status.rs new file mode 100644 index 0000000000..17d24a204a --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_upload_model_job_status.rs @@ -0,0 +1,172 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetTtsUploadModelStatusPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct TtsUploadModelJobStatusForResponse { + pub job_token: String, + + /// Primary status from the database (a state machine). + pub status: String, + + // TODO: Not yet used. + /// Extra, temporary status from Redis. + /// This can denote inference progress, and the Python code can write to it. + pub maybe_extra_status_description: Option, + + pub attempt_count: u8, + pub maybe_model_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct GetTtsUploadModelStatusSuccessResponse { + pub success: bool, + pub state: TtsUploadModelJobStatusForResponse, +} + +#[derive(Debug, Display)] +pub enum GetTtsUploadModelStatusError { + ServerError, +} + +#[derive(Serialize)] +pub struct TtsUploadModelJobStatusRecord { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + pub maybe_model_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl ResponseError for GetTtsUploadModelStatusError { + fn status_code(&self) -> StatusCode { + match *self { + GetTtsUploadModelStatusError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetTtsUploadModelStatusError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn get_tts_upload_model_job_status_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_status = sqlx::query_as!( + TtsUploadModelJobStatusRecord, + r#" +SELECT + jobs.token as job_token, + + jobs.status, + jobs.attempt_count, + jobs.on_success_result_token as maybe_model_token, + + jobs.created_at, + jobs.updated_at + +FROM tts_model_upload_jobs as jobs + +WHERE jobs.token = ? + "#, + &path.token + ) + .fetch_one(&server_state.mysql_pool) + .await; // TODO: This will return error if it doesn't exist + + let record : TtsUploadModelJobStatusRecord = match maybe_status { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(GetTtsUploadModelStatusError::ServerError); + }, + _ => { + warn!("tts template query error: {:?}", err); + return Err(GetTtsUploadModelStatusError::ServerError); + } + } + } + }; + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + GetTtsUploadModelStatusError::ServerError + })?; + + let extra_status_key = RedisKeys::tts_download_extra_status_info(&path.token); + let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { + Ok(Some(status)) => { + Some(status) + }, + Ok(None) => None, + Err(e) => { + warn!("redis error: {:?}", e); + None // Fail open + }, + }; + + let model_for_response = TtsUploadModelJobStatusForResponse { + job_token: record.job_token.clone(), + status: record.status.clone(), + maybe_extra_status_description, + attempt_count: record.attempt_count as u8, + maybe_model_token: record.maybe_model_token.clone(), + created_at: record.created_at.clone(), + updated_at: record.updated_at.clone(), + }; + + let response = GetTtsUploadModelStatusSuccessResponse { + success: true, + state: model_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetTtsUploadModelStatusError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs new file mode 100644 index 0000000000..f094766144 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs @@ -0,0 +1,168 @@ +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::tts::tts_category_assignments::fetch_and_build_tts_model_category_map::fetch_and_build_tts_model_category_map; +use database_queries::queries::tts::tts_models::list_tts_models::list_tts_models; +use lexical_sort::natural_lexical_cmp; +use log::{info, warn, log, error}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::collections::HashSet; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize, Clone)] +pub struct TtsModelRecordForResponse { + pub model_token: String, + pub tts_model_type: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + pub title: String, + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + pub is_front_page_featured: bool, + pub is_twitch_featured: bool, + pub maybe_suggested_unique_bot_command: Option, + + /// Category assignments + /// From non-deleted, mod-approved categories only + pub category_tokens: HashSet, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct ListTtsModelsSuccessResponse { + pub success: bool, + pub models: Vec, +} + +#[derive(Debug)] +pub enum ListTtsModelsError { + ServerError, +} + +impl ResponseError for ListTtsModelsError { + fn status_code(&self) -> StatusCode { + match *self { + ListTtsModelsError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListTtsModelsError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTtsModelsError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_tts_models_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let maybe_models = server_state.caches.voice_list.copy_without_bump_if_unexpired() + .map_err(|e| { + error!("Error consulting cache: {:?}", e); + ListTtsModelsError::ServerError + })?; + + let models = match maybe_models { + Some(models) => { + info!("Serving TTS models from cache"); + models + }, + None => { + info!("Populating TTS models from database"); + let models = get_all_models(&server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error querying database: {:?}", e); + ListTtsModelsError::ServerError + })?; + + server_state.caches.voice_list.store_copy(&models) + .map_err(|e| { + error!("Error storing cache: {:?}", e); + ListTtsModelsError::ServerError + })?; + models + }, + }; + + let response = ListTtsModelsSuccessResponse { + success: true, + models, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTtsModelsError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +async fn get_all_models(mysql_pool: &MySqlPool) -> AnyhowResult> { + let mut models = list_tts_models( + mysql_pool, + None, + false + ).await?; + + let model_categories_map + = fetch_and_build_tts_model_category_map(mysql_pool).await?; + + // Make the list nice for human readers. + models.sort_by(|a, b| + natural_lexical_cmp(&a.title, &b.title)); + + let models_for_response = models.into_iter() + .map(|model| { + let model_token = model.model_token.clone(); + TtsModelRecordForResponse { + model_token: model.model_token, + tts_model_type: model.tts_model_type, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + title: model.title, + ietf_language_tag: model.ietf_language_tag, + ietf_primary_language_subtag: model.ietf_primary_language_subtag, + is_front_page_featured: model.is_front_page_featured, + is_twitch_featured: model.is_twitch_featured, + maybe_suggested_unique_bot_command: model.maybe_suggested_unique_bot_command, + category_tokens: model_categories_map.model_to_category_tokens.get(&model_token) + .map(|hash| hash.clone()) + .unwrap_or(HashSet::new()), + created_at: model.created_at, + updated_at: model.updated_at, + } + }) + .collect::>(); + + Ok(models_for_response) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/mod.rs new file mode 100644 index 0000000000..5d7e38d8c1 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/mod.rs @@ -0,0 +1,12 @@ +pub mod delete_tts_model; +pub mod delete_tts_result; +pub mod edit_tts_model; +pub mod edit_tts_result; +pub mod enqueue_infer_tts; +pub mod enqueue_upload_tts_model; +pub mod get_tts_inference_job_status; +pub mod get_tts_model; +pub mod get_tts_model_use_count; +pub mod get_tts_result; +pub mod get_tts_upload_model_job_status; +pub mod list_tts_models; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/create_event_rule.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/create_event_rule.rs new file mode 100644 index 0000000000..509e49a1a9 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/create_event_rule.rs @@ -0,0 +1,157 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::column_types::twitch_event_category::TwitchEventCategory; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_event_rules::insert_twitch_event_rule_builder::InsertTwitchEventRuleBuilder; +use database_queries::tokens::Tokens; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct CreateTwitchEventRuleRequest { + pub idempotency_token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: Option, + pub event_response: Option, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct CreateTwitchEventRuleResponse { + pub success: bool, + pub twitch_event_rule_token: String, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum CreateTwitchEventRuleError { + BadInput(String), + NotAuthorized, + ServerError, +} + +impl ResponseError for CreateTwitchEventRuleError { + fn status_code(&self) -> StatusCode { + match *self { + CreateTwitchEventRuleError::BadInput(_) => StatusCode::BAD_REQUEST, + CreateTwitchEventRuleError::NotAuthorized => StatusCode::UNAUTHORIZED, + CreateTwitchEventRuleError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for CreateTwitchEventRuleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn create_twitch_event_rule_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CreateTwitchEventRuleError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(CreateTwitchEventRuleError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot use"); + return Err(CreateTwitchEventRuleError::NotAuthorized); + } + + let creator_ip_address = get_request_ip(&http_request); + + let event_match_predicate = request.event_match_predicate + .clone() + .unwrap_or(EventMatchPredicate::NotSet {}); + + let mut event_match_predicate = serde_json::to_string(&event_match_predicate) + .map_err(|e| { + return CreateTwitchEventRuleError::BadInput( + "improper EventMatchPredicate".to_string()); + })?; + + let event_response = request.event_response + .clone() + .unwrap_or(EventResponse::NotSet {}); + + let mut event_response = serde_json::to_string(&event_response) + .map_err(|e| { + return CreateTwitchEventRuleError::BadInput( + "improper EventResponse".to_string()); + })?; + + let insert_builder = InsertTwitchEventRuleBuilder { + uuid_idempotency_token: request.idempotency_token.clone(), + user_token: user_session.user_token, + event_category: request.event_category, + event_match_predicate, + event_response, + user_specified_rule_order: request.user_specified_rule_order, + rule_is_disabled: request.rule_is_disabled, + ip_address_creation: creator_ip_address, + }; + + let twitch_event_rule_token = insert_builder.insert(&server_state.mysql_pool) + .await + .map_err(|e| { + warn!("creation error: {:?}", e); + CreateTwitchEventRuleError::ServerError + })?; + + let response = CreateTwitchEventRuleResponse { + success: true, + twitch_event_rule_token + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CreateTwitchEventRuleError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/delete_event_rule.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/delete_event_rule.rs new file mode 100644 index 0000000000..070c4fddc6 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/delete_event_rule.rs @@ -0,0 +1,134 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::queries::twitch::twitch_event_rules::delete_twitch_event_rule::delete_twitch_event_rule; +use database_queries::queries::twitch::twitch_event_rules::get_twitch_event_rule_for_user::get_twitch_event_rule_for_user; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct DeleteTwitchEventRulePathInfo { + token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct DeleteTwitchEventRuleResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum DeleteTwitchEventRuleError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for DeleteTwitchEventRuleError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteTwitchEventRuleError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteTwitchEventRuleError::NotFound => StatusCode::NOT_FOUND, + DeleteTwitchEventRuleError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteTwitchEventRuleError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DeleteTwitchEventRuleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn delete_twitch_event_rule_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteTwitchEventRuleError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteTwitchEventRuleError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot delete"); + return Err(DeleteTwitchEventRuleError::NotAuthorized); + } + + let twitch_event_rule = get_twitch_event_rule_for_user( + &path.token, &user_session.user_token, &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return DeleteTwitchEventRuleError::ServerError; + })?; + + if twitch_event_rule.is_none() { + info!("event rule not found"); + return Err(DeleteTwitchEventRuleError::NotFound); + } + + let creator_ip_address = get_request_ip(&http_request); + + let _r = delete_twitch_event_rule( + &path.token, + &creator_ip_address, + &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return DeleteTwitchEventRuleError::ServerError; + }); + + let response = DeleteTwitchEventRuleResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| DeleteTwitchEventRuleError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/edit_event_rule.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/edit_event_rule.rs new file mode 100644 index 0000000000..cbfe0b9128 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/edit_event_rule.rs @@ -0,0 +1,178 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_event_rules::get_twitch_event_rule_for_user::get_twitch_event_rule_for_user; +use database_queries::queries::twitch::twitch_event_rules::update_twitch_event_rule_builder::UpdateTwitchEventRuleBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use crate::http_server::endpoints::twitch::event_rules::validations::validate_event_response::validate_event_response; +use crate::http_server::endpoints::twitch::event_rules::validations::validate_event_match_predicate::validate_event_match_predicate; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct EditTwitchEventRulePathInfo { + token: String, +} + +// These are not sparse updates! +#[derive(Deserialize)] +pub struct EditTwitchEventRuleRequest { + pub event_match_predicate: Option, + pub event_response: Option, + pub rule_is_disabled: bool, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct EditTwitchEventRuleResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum EditTwitchEventRuleError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for EditTwitchEventRuleError { + fn status_code(&self) -> StatusCode { + match *self { + EditTwitchEventRuleError::BadInput(_) => StatusCode::BAD_REQUEST, + EditTwitchEventRuleError::NotFound => StatusCode::NOT_FOUND, + EditTwitchEventRuleError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditTwitchEventRuleError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditTwitchEventRuleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn edit_twitch_event_rule_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditTwitchEventRuleError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditTwitchEventRuleError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot edit"); + return Err(EditTwitchEventRuleError::NotAuthorized); + } + + let twitch_event_rule = get_twitch_event_rule_for_user( + &path.token, &user_session.user_token, &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return EditTwitchEventRuleError::ServerError; + })?; + + if twitch_event_rule.is_none() { + info!("event rule not found"); + return Err(EditTwitchEventRuleError::NotFound); + } + + let creator_ip_address = get_request_ip(&http_request); + + let event_match_predicate = request.event_match_predicate + .clone() + .unwrap_or(EventMatchPredicate::NotSet {}); + + validate_event_match_predicate(&event_match_predicate) + .map_err(|reason| EditTwitchEventRuleError::BadInput(reason))?; + + let mut event_match_predicate = serde_json::to_string(&event_match_predicate) + .map_err(|e| { + return EditTwitchEventRuleError::BadInput( + "improper EventMatchPredicate".to_string()); + })?; + + let event_response = request.event_response + .clone() + .unwrap_or(EventResponse::NotSet {}); + + validate_event_response(&event_response) + .map_err(|reason| EditTwitchEventRuleError::BadInput(reason))?; + + let mut event_response = serde_json::to_string(&event_response) + .map_err(|e| { + return EditTwitchEventRuleError::BadInput( + "improper EventResponse".to_string()); + })?; + + let update_builder = UpdateTwitchEventRuleBuilder { + token: path.token.clone(), + event_match_predicate, + event_response, + rule_is_disabled: request.rule_is_disabled, + ip_address_update: creator_ip_address, + }; + + let _r = update_builder.update(&server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return EditTwitchEventRuleError::ServerError; + }); + + let response = EditTwitchEventRuleResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| EditTwitchEventRuleError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/get_event_rule.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/get_event_rule.rs new file mode 100644 index 0000000000..e1664019b2 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/get_event_rule.rs @@ -0,0 +1,165 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::column_types::twitch_event_category::TwitchEventCategory; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_event_rules::list_twitch_event_rules_for_user::{TwitchEventRule, list_twitch_event_rules_for_user}; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use lexical_sort::natural_lexical_cmp; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use database_queries::queries::twitch::twitch_event_rules::get_twitch_event_rule_for_user::get_twitch_event_rule_for_user; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct GetTwitchEventRulePathInfo { + token: String, +} + +// =============== Success Response =============== + +#[derive(Debug, Serialize)] +pub struct HydratedTwitchEventRule { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: EventMatchPredicate, + pub event_response: EventResponse, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} + +#[derive(Serialize)] +pub struct GetTwitchEventRuleResponse { + pub success: bool, + pub twitch_event_rule: HydratedTwitchEventRule, +} + + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum GetTwitchEventRuleError { + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for GetTwitchEventRuleError { + fn status_code(&self) -> StatusCode { + match *self { + GetTwitchEventRuleError::NotFound => StatusCode::NOT_FOUND, + GetTwitchEventRuleError::NotAuthorized => StatusCode::UNAUTHORIZED, + GetTwitchEventRuleError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetTwitchEventRuleError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn get_twitch_event_rule_for_user_handler( + path: Path, + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetTwitchEventRuleError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(GetTwitchEventRuleError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot use"); + return Err(GetTwitchEventRuleError::NotAuthorized); + } + + let maybe_twitch_event_rule = get_twitch_event_rule_for_user( + &path.token, &user_session.user_token, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("query error: {:?}", e); + GetTwitchEventRuleError::ServerError + })?; + + let rule = match maybe_twitch_event_rule { + Some(rule) => rule, + None => { + return Err(GetTwitchEventRuleError::NotFound); + } + }; + + let event_match_predicate = serde_json::from_str(&rule.event_match_predicate) + .unwrap_or_else(|e| { + error!("Issue with deserializing: {}", e); + EventMatchPredicate::NotSet {} + }); + + let event_response = serde_json::from_str(&rule.event_response) + .unwrap_or_else(|e| { + error!("Issue with deserializing: {}", e); + EventResponse::NotSet {} + }); + + let twitch_event_rule = HydratedTwitchEventRule { + token: rule.token, + event_category: rule.event_category, + event_match_predicate, + event_response, + user_specified_rule_order: rule.user_specified_rule_order, + rule_is_disabled: rule.rule_is_disabled, + created_at: rule.created_at, + updated_at: rule.updated_at, + }; + + let response = GetTwitchEventRuleResponse { + success: true, + twitch_event_rule, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetTwitchEventRuleError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/list_event_rules_for_user.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/list_event_rules_for_user.rs new file mode 100644 index 0000000000..015d580b59 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/list_event_rules_for_user.rs @@ -0,0 +1,152 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::column_types::twitch_event_category::TwitchEventCategory; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_event_rules::list_twitch_event_rules_for_user::{TwitchEventRule, list_twitch_event_rules_for_user}; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use lexical_sort::natural_lexical_cmp; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Success Response =============== + +#[derive(Debug, Serialize)] +pub struct HydratedTwitchEventRule { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: EventMatchPredicate, + pub event_response: EventResponse, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, +} + +#[derive(Serialize)] +pub struct ListTwitchEventRulesResponse { + pub success: bool, + pub twitch_event_rules: Vec, +} + + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum ListTwitchEventRulesError { + NotAuthorized, + ServerError, +} + +impl ResponseError for ListTwitchEventRulesError { + fn status_code(&self) -> StatusCode { + match *self { + ListTwitchEventRulesError::NotAuthorized => StatusCode::UNAUTHORIZED, + ListTwitchEventRulesError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTwitchEventRulesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn list_twitch_event_rules_for_user_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListTwitchEventRulesError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ListTwitchEventRulesError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot use"); + return Err(ListTwitchEventRulesError::NotAuthorized); + } + + let twitch_event_rules = list_twitch_event_rules_for_user( + &user_session.user_token, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("query error: {:?}", e); + ListTwitchEventRulesError::ServerError + })?; + + let twitch_event_rules = twitch_event_rules.into_iter() + .map(|rule| { + + let event_match_predicate = serde_json::from_str(&rule.event_match_predicate) + .unwrap_or_else(|e| { + error!("Issue with deserializing: {}", e); + EventMatchPredicate::NotSet {} + }); + + let event_response = serde_json::from_str(&rule.event_response) + .unwrap_or_else(|e| { + error!("Issue with deserializing: {}", e); + EventResponse::NotSet {} + }); + + HydratedTwitchEventRule { + token: rule.token, + event_category: rule.event_category, + event_match_predicate, + event_response, + user_specified_rule_order: rule.user_specified_rule_order, + rule_is_disabled: rule.rule_is_disabled, + created_at: rule.created_at, + updated_at: rule.updated_at, + } + }) + .collect(); + + let response = ListTwitchEventRulesResponse { + success: true, + twitch_event_rules, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTwitchEventRulesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/mod.rs new file mode 100644 index 0000000000..56d071b288 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/mod.rs @@ -0,0 +1,7 @@ +pub mod create_event_rule; +pub mod delete_event_rule; +pub mod edit_event_rule; +pub mod get_event_rule; +pub mod list_event_rules_for_user; +pub mod reorder_twitch_event_rules; +pub mod validations; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/reorder_twitch_event_rules.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/reorder_twitch_event_rules.rs new file mode 100644 index 0000000000..75e1937573 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/reorder_twitch_event_rules.rs @@ -0,0 +1,168 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::endpoints::twitch::event_rules::validations::validate_event_match_predicate::validate_event_match_predicate; +use crate::http_server::endpoints::twitch::event_rules::validations::validate_event_response::validate_event_response; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_event_rules::get_twitch_event_rule_for_user::get_twitch_event_rule_for_user; +use database_queries::queries::twitch::twitch_event_rules::list_twitch_event_rules_for_user::list_twitch_event_rules_for_user; +use database_queries::queries::twitch::twitch_event_rules::reorder_twitch_event_rules::reorder_twitch_event_rules; +use database_queries::queries::twitch::twitch_event_rules::update_twitch_event_rule_builder::UpdateTwitchEventRuleBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log, error}; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::collections::{HashSet, HashMap}; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct RuleTokenPositionPair { + pub rule_token: String, + pub position: u32, +} + +#[derive(Deserialize)] +pub struct ReorderTwitchEventRulesRequest { + pub rule_token_position_pairs: Vec +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct ReorderTwitchEventRulesResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum ReorderTwitchEventRulesError { + BadInput(String), + NotFound, + NotAuthorized, + ServerError, +} + +impl ResponseError for ReorderTwitchEventRulesError { + fn status_code(&self) -> StatusCode { + match *self { + ReorderTwitchEventRulesError::BadInput(_) => StatusCode::BAD_REQUEST, + ReorderTwitchEventRulesError::NotFound => StatusCode::NOT_FOUND, + ReorderTwitchEventRulesError::NotAuthorized => StatusCode::UNAUTHORIZED, + ReorderTwitchEventRulesError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ReorderTwitchEventRulesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn reorder_twitch_event_rules_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ReorderTwitchEventRulesError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(ReorderTwitchEventRulesError::NotAuthorized); + } + }; + + if user_session.is_banned { + warn!("banned users cannot edit"); + return Err(ReorderTwitchEventRulesError::NotAuthorized); + } + + let rules = list_twitch_event_rules_for_user( + &user_session.user_token, + &server_state.mysql_pool) + .await + .map_err(|e| { + error!("Error with query: {:?}", e); + return ReorderTwitchEventRulesError::ServerError; + })?; + + // Here we check that the input set is entirely valid. + // The update query is pretty dangerous (though also uses input checks), so this + // should doubly guarantee nothing bad slips through. + // Also it serves as a permission check. + + let valid_rule_tokens = rules.iter() + .map(|rule| rule.token.to_string()) + .collect::>(); + + let check_tokens = request.rule_token_position_pairs.iter() + .map(|pair| pair.rule_token.to_string()) + .collect::>(); + + for check_token in check_tokens.iter() { + if !valid_rule_tokens.contains(check_token.as_str()) { + return Err(ReorderTwitchEventRulesError::BadInput("invalid token supplied".to_string())); + } + } + + let rule_token_to_order_map = request.rule_token_position_pairs.iter() + .map(|pair| (pair.rule_token.to_string(), pair.position)) + .collect::>(); + + let ip_address = get_request_ip(&http_request); + + reorder_twitch_event_rules( + rule_token_to_order_map, + &user_session.user_token, + &ip_address, + &server_state.mysql_pool) + .await + .map_err(|err| { + error!("Error with query: {:?}", err); + return ReorderTwitchEventRulesError::ServerError; + })?; + + + let response = ReorderTwitchEventRulesResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ReorderTwitchEventRulesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/mod.rs new file mode 100644 index 0000000000..9334525946 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/mod.rs @@ -0,0 +1,2 @@ +pub mod validate_event_match_predicate; +pub mod validate_event_response; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/validate_event_match_predicate.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/validate_event_match_predicate.rs new file mode 100644 index 0000000000..443cc02b2b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/validate_event_match_predicate.rs @@ -0,0 +1,37 @@ +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; + +// TODO: Check that the tokens are valid models. + +/// Try to prevent saving garbage to the database. +pub fn validate_event_match_predicate(event_match_predicate: &EventMatchPredicate) -> Result<(), String> { + match event_match_predicate { + EventMatchPredicate::NotSet {} => { + // Implicitly valid + }, + EventMatchPredicate::BitsCheermoteNameExactMatch { cheermote_name } => { + if cheermote_name.trim().is_empty() { + return Err("cheermote_name is empty".to_string()); + } + } + EventMatchPredicate::BitsCheermotePrefixSpendThreshold { cheermote_prefix, minimum_bits_spent } => { + if cheermote_prefix.trim().is_empty() { + return Err("cheermote_prefix is empty".to_string()); + } + if *minimum_bits_spent < 1 { + return Err("minimum_bits_spent should be a positive integer".to_string()); + } + } + EventMatchPredicate::BitsSpendThreshold { minimum_bits_spent } => { + if *minimum_bits_spent < 1 { + return Err("minimum_bits_spent should be a positive integer".to_string()); + } + } + EventMatchPredicate::ChannelPointsRewardNameExactMatch { reward_name } => { + if reward_name.trim().is_empty() { + return Err("reward_name is empty".to_string()); + } + } + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/validate_event_response.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/validate_event_response.rs new file mode 100644 index 0000000000..3356246df4 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/event_rules/validations/validate_event_response.rs @@ -0,0 +1,39 @@ +use database_queries::complex_models::event_responses::EventResponse; + +// TODO: Check that the tokens are valid models. + +/// Try to prevent saving garbage to the database. +pub fn validate_event_response(event_response: &EventResponse) -> Result<(), String> { + match event_response { + EventResponse::NotSet {} => { + // Implicitly valid + }, + EventResponse::TtsSingleVoice { tts_model_token } => { + if tts_model_token.is_empty() { + return Err("tts_model_token is empty".to_string()); + } + }, + EventResponse::TtsRandomVoice { tts_model_tokens } => { + for model_token in tts_model_tokens { + if model_token.is_empty() { + return Err("one or more tts_model_tokens is invalid".to_string()); + } + } + }, + EventResponse::TtsCommandPresets { .. } => { + // Implicitly valid + } + EventResponse::TtsCommandCustom { command_map, .. } => { + for (command, model_token) in command_map.iter() { + if command.trim().is_empty() { + return Err("one or more commands is invalid".to_string()); + } + if model_token.is_empty() { + return Err("one or more model tokens is invalid".to_string()); + } + } + }, + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/mod.rs new file mode 100644 index 0000000000..8ad19b93b8 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/mod.rs @@ -0,0 +1,2 @@ +pub mod event_rules; +pub mod oauth; diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/check_oauth_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/check_oauth_status.rs new file mode 100644 index 0000000000..60d1c24476 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/check_oauth_status.rs @@ -0,0 +1,113 @@ +use actix_http::{StatusCode, header}; +use actix_web::{HttpResponse, HttpRequest, web, ResponseError}; +use crate::server_state::ServerState; +use database_queries::queries::twitch::twitch_oauth::insert::TwitchOauthTokenInsertBuilder; +use http_server_common::request::get_request_ip::get_request_ip; +use http_server_common::response::response_error_helpers::to_simple_json_error; +use log::error; +use log::info; +use log::warn; +use std::fmt; +use std::sync::Arc; +use twitch_common::oauth_token_builder::get_oauth_token_builder; +use twitch_oauth2::tokens::BearerTokenType::UserToken; +use twitch_oauth2::{CsrfToken, TwitchToken}; +use database_queries::tokens::Tokens; +use database_queries::queries::twitch::twitch_oauth::find::{TwitchOauthTokenFinder, TwitchOauthTokenRecord}; + +#[derive(Serialize)] +pub struct CheckOauthResponse { + pub success: bool, + + /// This is false if the user isn't logged in, doesn't have oauth, or oauth is invalid + pub oauth_token_found: bool, + + /// Shouldn't be harmful to send this back since user must be logged in. + pub maybe_twitch_username: Option, + pub maybe_twitch_username_lowercase: Option, +} + +#[derive(Debug)] +pub enum CheckOauthStatusError { + ServerError, +} + +impl ResponseError for CheckOauthStatusError { + fn status_code(&self) -> StatusCode { + match *self { + Self::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + Self::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for CheckOauthStatusError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn check_oauth_status_handler( + http_request: HttpRequest, + server_state: web::Data> +) + -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CheckOauthStatusError::ServerError + })?; + + let mut oauth_token_found = false; + let mut maybe_twitch_username = None; + let mut maybe_twitch_username_lowercase = None; + + match maybe_user_session { + None => {}, // Not logged in - not found + Some(session) => { + let finder = TwitchOauthTokenFinder::new() + .allow_expired_tokens(true) + .scope_user_token(Some(&session.user_token)); + + let maybe_record = finder.perform_query(&server_state.mysql_pool) + .await + .map_err(|e| { + warn!("lookup error: {:?}", e); + CheckOauthStatusError::ServerError + })?; + + // Existence is sufficient. + if let Some(record) = maybe_record { + oauth_token_found = true; + maybe_twitch_username = Some(record.twitch_username.to_string()); + maybe_twitch_username_lowercase = Some(record.twitch_username_lowercase.to_string()); + } + }, + } + + let response = CheckOauthResponse { + success: true, + oauth_token_found, + maybe_twitch_username, + maybe_twitch_username_lowercase, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CheckOauthStatusError::ServerError)?; + + return Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)); +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/mod.rs new file mode 100644 index 0000000000..35e370f133 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/mod.rs @@ -0,0 +1,4 @@ +pub mod check_oauth_status; +pub mod oauth_begin_json; +pub mod oauth_begin_redirect; +pub mod oauth_end; diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_begin_json.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_begin_json.rs new file mode 100644 index 0000000000..5c1ef6cd82 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_begin_json.rs @@ -0,0 +1,77 @@ +use actix_http::{StatusCode, header}; +use actix_web::{HttpRequest, web, HttpResponse, ResponseError, HttpResponseBuilder}; +use crate::server_state::ServerState; +use http_server_common::response::response_error_helpers::to_simple_json_error; +use http_server_common::response::response_success_helpers::simple_json_success; +use http_server_common::response::to_json_success_response::to_json_success_response; +use log::info; +use log::warn; +use std::fmt; +use std::sync::Arc; +use twitch_api2::twitch_oauth2::ClientId; +use twitch_common::oauth_token_builder::get_oauth_token_builder; +use twitch_oauth2::tokens::UserTokenBuilder; +use twitch_oauth2::{Scope, ClientSecret}; + +#[derive(Serialize)] +pub struct OauthBeginEnrollResult { + redirect_url: String, +} + +#[derive(Debug)] +pub enum OauthBeginEnrollError { + ServerError, +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for OauthBeginEnrollError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "OauthBeginEnrollError") + } +} + +impl ResponseError for OauthBeginEnrollError { + fn status_code(&self) -> StatusCode { + match *self { + Self::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + Self::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn oauth_begin_enroll_json( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let redirect_url = + twitch_oauth2::url::Url::parse(&server_state.twitch_oauth.redirect_landing_url) + .map_err(|e| { + warn!("Error parsing url: {:?}", e); + OauthBeginEnrollError::ServerError + })?; + + let mut builder = get_oauth_token_builder( + &server_state.twitch_oauth.secrets.client_id, + &server_state.twitch_oauth.secrets.client_secret, + &redirect_url, + true); + + let (url, _csrf_token) = builder.generate_url(); + + let response = OauthBeginEnrollResult { + redirect_url: url.to_string(), + }; + + let response = to_json_success_response(&response) + .map_err(|_| OauthBeginEnrollError::ServerError)?; + + Ok(response) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_begin_redirect.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_begin_redirect.rs new file mode 100644 index 0000000000..a2e39bea79 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_begin_redirect.rs @@ -0,0 +1,72 @@ +use actix_http::{StatusCode, header}; +use actix_web::{HttpRequest, web, HttpResponse, ResponseError, HttpResponseBuilder}; +use crate::server_state::ServerState; +use http_server_common::response::response_error_helpers::to_simple_json_error; +use http_server_common::response::response_success_helpers::simple_json_success; +use http_server_common::response::to_json_success_response::to_json_success_response; +use log::info; +use log::warn; +use std::fmt; +use std::sync::Arc; +use twitch_api2::twitch_oauth2::ClientId; +use twitch_common::oauth_token_builder::get_oauth_token_builder; +use twitch_oauth2::tokens::UserTokenBuilder; +use twitch_oauth2::{Scope, ClientSecret}; + +#[derive(Serialize)] +pub struct OauthBeginEnrollResult { + redirect_url: String, +} + +#[derive(Debug)] +pub enum OauthBeginEnrollError { + ServerError, +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for OauthBeginEnrollError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "OauthBeginEnrollError") + } +} + +impl ResponseError for OauthBeginEnrollError { + fn status_code(&self) -> StatusCode { + match *self { + Self::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + Self::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn oauth_begin_enroll_redirect( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let redirect_url = + twitch_oauth2::url::Url::parse(&server_state.twitch_oauth.redirect_landing_url) + .map_err(|e| { + warn!("Error parsing url: {:?}", e); + OauthBeginEnrollError::ServerError + })?; + + let mut builder = get_oauth_token_builder( + &server_state.twitch_oauth.secrets.client_id, + &server_state.twitch_oauth.secrets.client_secret, + &redirect_url, + true); + + let (url, _csrf_token) = builder.generate_url(); + + Ok(HttpResponse::Found() + .append_header((header::LOCATION, url.to_string())) + .finish()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_end.rs b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_end.rs new file mode 100644 index 0000000000..1bab1a8e9b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/twitch/oauth/oauth_end.rs @@ -0,0 +1,188 @@ +use actix_http::{StatusCode, header}; +use actix_web::{HttpResponse, HttpRequest, web, ResponseError}; +use crate::server_state::ServerState; +use database_queries::queries::twitch::twitch_oauth::insert::TwitchOauthTokenInsertBuilder; +use http_server_common::request::get_request_ip::get_request_ip; +use http_server_common::response::response_error_helpers::to_simple_json_error; +use log::error; +use log::info; +use log::warn; +use std::fmt; +use std::sync::Arc; +use twitch_common::oauth_token_builder::get_oauth_token_builder; +use twitch_oauth2::tokens::BearerTokenType::UserToken; +use twitch_oauth2::{CsrfToken, TwitchToken}; +use database_queries::tokens::Tokens; + +#[derive(Deserialize)] +pub struct QueryParams { + /// OAuth authorization code + pub code: Option, + /// Opaque value used to avoid CSRF attacks. + pub state: Option, + /// List of permission scopes + pub scope: Option, + + /// Only set in the event of an error + pub error: Option, + /// Only set in the event of an error + pub error_description: Option, +} + +#[derive(Debug)] +pub enum OauthEndEnrollFromRedirectError { + ServerError, + TwitchOauthError { reason: String }, +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for OauthEndEnrollFromRedirectError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "OauthEndEnrollFromRedirectError") + } +} + +impl ResponseError for OauthEndEnrollFromRedirectError { + fn status_code(&self) -> StatusCode { + match *self { + Self::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + OauthEndEnrollFromRedirectError::TwitchOauthError { .. } => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + Self::ServerError => "server error".to_string(), + OauthEndEnrollFromRedirectError::TwitchOauthError { .. } => "twitch error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn oauth_end_enroll_from_redirect( + http_request: HttpRequest, + query: web::Query, + server_state: web::Data> +) + -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + OauthEndEnrollFromRedirectError::ServerError + })?; + + if let Some(error) = query.error.as_deref() { + return Err(OauthEndEnrollFromRedirectError::TwitchOauthError { reason: error.to_string() }) + } + if let Some(error) = query.error_description.as_deref() { + return Err(OauthEndEnrollFromRedirectError::TwitchOauthError { reason: error.to_string() }) + } + + let (code, state) = match (query.code.as_deref(), query.state.as_deref()) { + (Some(code), Some(state)) => (code.to_string(), state.to_string()), + _ => { + return Err(OauthEndEnrollFromRedirectError::TwitchOauthError { + reason: "params not set".to_string() + }); + } + }; + + // TODO/FIXME: This is a major security issue (part 1). + // These need to be persisted in the database and associated with the user. + let csrf_token = CsrfToken::new(&state); + + let redirect_url = + twitch_oauth2::url::Url::parse(&server_state.twitch_oauth.redirect_landing_url) + .map_err(|e| { + warn!("Error parsing url: {:?}", e); + OauthEndEnrollFromRedirectError::ServerError + })?; + + let mut builder = get_oauth_token_builder( + &server_state.twitch_oauth.secrets.client_id, + &server_state.twitch_oauth.secrets.client_secret, + &redirect_url, + true); + + // TODO/FIXME: This is a major security issue (part 2). + builder.set_csrf(csrf_token); + + let http_client = &reqwest::Client::builder() + .redirect(reqwest::redirect::Policy::none()) + .build() + .map_err(|e| { + error!("Problem creating HTTP client: {:?}", e); + OauthEndEnrollFromRedirectError::ServerError + })?; + + let user_token = builder + .get_user_token( + http_client, + &state, + &code, + ) + .await + .map_err(|e| { + warn!("Error fetching user token: {:?}", e); + OauthEndEnrollFromRedirectError::TwitchOauthError { + reason: "error fetching user token".to_string() + } + })?; + + let user_id = user_token.user_id.to_string(); + let auth_token = user_token.access_token.secret().to_string(); + let refresh_token: Option = user_token.refresh_token + .as_ref() + .map(|token| token.secret().to_string()); + + let twitch_username = user_token.login.to_string(); + let never_expiring = user_token.never_expiring; + let expires_in = user_token.expires_in(); + + info!("User id: {:?}", user_id); + info!("Twitch username: {:?}", twitch_username); + + let ip_address = get_request_ip(&http_request); + + let expires_seconds = expires_in.as_secs() as u32; // NB: Silent overflow + + let oauth_grouping_token = Tokens::new_twitch_oauth_grouping_token() + .map_err(|e| { + error!("token creation error: {:?}", e); + OauthEndEnrollFromRedirectError::ServerError + })?; + + let mut insert_builder = + TwitchOauthTokenInsertBuilder::new(&user_id, &twitch_username, &auth_token, &oauth_grouping_token) + .set_expires_in_seconds(Some(expires_seconds)) + .set_refresh_token(refresh_token.as_deref()) + .set_ip_address_creation(Some(&ip_address)) + .has_bits_read(true) + .has_channel_read_redemptions(true) + .has_channel_read_subscriptions(true) + .has_chat_edit(true) + .has_chat_read(true); + + if let Some(user_session) = maybe_user_session { + insert_builder = insert_builder.set_user_token(Some(&user_session.user_token)); + } + + let _result = insert_builder.insert(&server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Error saving to db: {:?}", e); + OauthEndEnrollFromRedirectError::ServerError + })?; + + Ok(HttpResponse::Found() + .append_header(( + header::LOCATION, + server_state.twitch_oauth.redirect_landing_finished_url.as_str() + )) + .finish()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/create_account.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/create_account.rs new file mode 100644 index 0000000000..d61fb24fd9 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/create_account.rs @@ -0,0 +1,284 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::endpoints::users::create_account::CreateAccountErrorType::{BadInput, ServerError, UsernameTaken, EmailTaken}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use crate::validations::username_reservations::is_reserved_username; +use database_queries::tokens::Tokens; +use database_queries::queries::users::user_sessions::create_session::create_session_for_user; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::collections::HashMap; +use std::fmt::Formatter; +use std::fmt; +use std::sync::Arc; +use user_input_common::check_for_slurs::contains_slurs; + +const NEW_USER_ROLE: &'static str = "user"; + +#[derive(Deserialize)] +pub struct CreateAccountRequest { + pub username: String, + pub password: String, + pub password_confirmation: String, + pub email_address: String, +} + +#[derive(Serialize)] +pub struct CreateAccountSuccessResponse { + pub success: bool, +} + +#[derive(Serialize, Debug)] +pub struct CreateAccountErrorResponse { + pub success: bool, + pub error_type: CreateAccountErrorType, + pub error_fields: HashMap, +} + +#[derive(Copy, Clone, Debug, Serialize)] +pub enum CreateAccountErrorType { + BadInput, + EmailTaken, + ServerError, + UsernameReserved, + UsernameTaken, +} + +impl CreateAccountErrorResponse { + fn server_error() -> Self { + Self { + success: false, + error_type: CreateAccountErrorType::ServerError, + error_fields: HashMap::new(), + } + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for CreateAccountErrorResponse { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.error_type) + } +} + +impl ResponseError for CreateAccountErrorResponse { + fn status_code(&self) -> StatusCode { + match self.error_type { + CreateAccountErrorType::BadInput => StatusCode::BAD_REQUEST, + CreateAccountErrorType::EmailTaken => StatusCode::BAD_REQUEST, + CreateAccountErrorType::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + CreateAccountErrorType::UsernameReserved => StatusCode::BAD_REQUEST, + CreateAccountErrorType::UsernameTaken => StatusCode::BAD_REQUEST, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +pub async fn create_account_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let mut error_fields = HashMap::new(); + + if let Err(reason) = validate_username(&request.username) { + error_fields.insert("username".to_string(), reason); + } + + if let Err(reason) = validate_passwords(&request.password, &request.password_confirmation) { + error_fields.insert("password".to_string(), reason); + } + + if contains_slurs(&request.username) { + error_fields.insert("username".to_string(), "username contains slurs".to_string()); + } + + if !request.email_address.contains("@") { + error_fields.insert("email_address".to_string(), "invalid email address".to_string()); + } + + if is_reserved_username(&request.username) { + error_fields.insert("username".to_string(), "username is reserved".to_string()); + + return Err(CreateAccountErrorResponse { + success: false, + error_type: CreateAccountErrorType::UsernameReserved, + error_fields + }); + } + + if !error_fields.is_empty() { + return Err(CreateAccountErrorResponse { + success: false, + error_type: CreateAccountErrorType::BadInput, + error_fields + }); + } + + let user_token = Tokens::new_user() + .map_err(|e| { + warn!("Bad crockford token: {:?}", e); + CreateAccountErrorResponse::server_error() + })?; + + let password_hash = match bcrypt::hash(&request.password, bcrypt::DEFAULT_COST) { + Ok(hash) => hash, + Err(err) => { + warn!("Bcrypt error: {:?}", err); + return Err(CreateAccountErrorResponse::server_error()); + } + }; + + let username = request.username.trim().to_lowercase(); + let display_name = request.username.trim().to_string(); + + let email_address = request.email_address.trim().to_lowercase(); + + let email_gravatar_hash = email_to_gravatar(&email_address); + + let profile_markdown = ""; + let profile_rendered_html = ""; + + let ip_address = get_request_ip(&http_request); + + let query_result = sqlx::query!( + r#" +INSERT INTO users ( + token, + username, + display_name, + email_address, + email_gravatar_hash, + profile_markdown, + profile_rendered_html, + user_role_slug, + password_hash, + ip_address_creation, + ip_address_last_login, + ip_address_last_update +) +VALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? ) + "#, + user_token.to_string(), + username, + display_name, + email_address, + email_gravatar_hash, + profile_markdown, + profile_rendered_html, + NEW_USER_ROLE, + password_hash, + ip_address.to_string(), + ip_address.to_string(), + ip_address.to_string(), + ) + .execute(&server_state.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("New user creation DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + let mut error_fields = HashMap::new(); + let mut maybe_error_type = None; + + match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + maybe_error_type = Some(CreateAccountErrorType::UsernameTaken); + error_fields.insert("username".to_string(), "username is taken".to_string()); + } else if err.message().contains("email_address") { + maybe_error_type = Some(CreateAccountErrorType::EmailTaken); + error_fields.insert("email_address".to_string(), "email is taken".to_string()); + } + } + _ => {}, + } + + if let Some(error_type) = maybe_error_type { + return Err(CreateAccountErrorResponse { + success: false, + error_type, + error_fields + }) + } + }, + _ => {}, + } + return Err(CreateAccountErrorResponse::server_error()); + } + }; + + info!("new user id: {}", record_id); + + //// TODO: Remove me in the future. Too lazy to check the date, plus this will + //// vary with adoption rate. + //server_state.badge_granter.grant_early_user_badge(&user_token) + // .await + // .map_err(|e| { + // warn!("error creating badge: {:?}", e); + // CreateAccountErrorResponse::server_error() + // })?; + + let create_session_result = + create_session_for_user(&user_token, &ip_address, &server_state.mysql_pool).await; + + let session_token = match create_session_result { + Ok(token) => token, + Err(e) => { + warn!("create account session creation error : {:?}", e); + return Err(CreateAccountErrorResponse::server_error()); + } + }; + + info!("new user session created"); + + server_state.firehose_publisher.publish_user_sign_up(&user_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + CreateAccountErrorResponse::server_error() + })?; + + let session_cookie = match server_state.cookie_manager.create_cookie(&session_token) { + Ok(cookie) => cookie, + Err(_) => return Err(CreateAccountErrorResponse::server_error()), + }; + + let response = CreateAccountSuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CreateAccountErrorResponse::server_error())?; + + Ok(HttpResponse::Ok() + .cookie(session_cookie) + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/edit_profile.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/edit_profile.rs new file mode 100644 index 0000000000..83682fdc82 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/edit_profile.rs @@ -0,0 +1,349 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use crate::validations::cashapp_username::{validate_cashapp_username, normalize_cashapp_username_for_storage}; +use crate::validations::discord_username::validate_discord_username; +use crate::validations::github_username::validate_github_username; +use crate::validations::passwords::validate_passwords; +use crate::validations::twitch_username::validate_twitch_username; +use crate::validations::twitter_username::{validate_twitter_username, normalize_twitter_username_for_storage}; +use crate::validations::username::validate_username; +use crate::validations::username_reservations::is_reserved_username; +use crate::validations::website_url::validate_website_url; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::queries::users::get_user_profile_by_username::get_user_profile_by_username; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use user_input_common::check_for_slurs::contains_slurs; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct EditProfilePathInfo { + username: String, +} + +#[derive(Deserialize)] +pub struct EditProfileRequest { + pub display_name: Option, + + pub profile_markdown: Option, + + pub discord_username: Option, + pub twitter_username: Option, + pub twitch_username: Option, + pub patreon_username: Option, + pub github_username: Option, + pub cashapp_username: Option, + pub website_url: Option, + + pub preferred_tts_result_visibility: Option, + pub preferred_w2l_result_visibility: Option, +} + +#[derive(Serialize)] +pub struct EditProfileSuccessResponse { + pub success: bool, +} + +#[derive(Debug)] +pub enum EditProfileError { + BadInput(String), + NotAuthorized, + UserNotFound, + ServerError, +} + +impl ResponseError for EditProfileError { + fn status_code(&self) -> StatusCode { + match *self { + EditProfileError::BadInput(_) => StatusCode::BAD_REQUEST, + EditProfileError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditProfileError::UserNotFound => StatusCode::NOT_FOUND, + EditProfileError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EditProfileError::BadInput(reason) => reason.to_string(), + EditProfileError::NotAuthorized=> "unauthorized".to_string(), + EditProfileError::UserNotFound => "user not found".to_string(), + EditProfileError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditProfileError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn edit_profile_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditProfileError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditProfileError::NotAuthorized); + } + }; + + if user_session.is_banned { + // Banned users can't edit anything + return Err(EditProfileError::NotAuthorized); + } + + let user_lookup_result = + get_user_profile_by_username(&path.username, &server_state.mysql_pool) + .await; + + let user_record = match user_lookup_result { + Ok(Some(result)) => result, + Ok(None) => return Err(EditProfileError::UserNotFound), + Err(err) => { + warn!("lookup error: {:?}", err); + return Err(EditProfileError::ServerError); + } + }; + + let mut editor_is_original_user = false; + let mut editor_is_moderator = false; + + if path.username == user_session.username { + editor_is_original_user = true; + } + + if user_session.can_edit_other_users_profiles { + editor_is_moderator = true; + } + + if !editor_is_original_user && !editor_is_moderator { + return Err(EditProfileError::NotAuthorized); + } + + // Fields to set + let mut twitter_username = None; + let mut twitch_username = None; + let mut discord_username = None; + let mut cashapp_username = None; + let mut github_username = None; + let mut website_url = None; + let mut profile_markdown = None; + let mut profile_html = None; + + if let Some(twitter) = request.twitter_username.as_deref() { + let trimmed = twitter.trim(); + if trimmed.is_empty() { + twitter_username = None; + } else { + if let Err(reason) = validate_twitter_username(trimmed) { + return Err(EditProfileError::BadInput(reason)); + } + let normalized = normalize_twitter_username_for_storage(trimmed); + twitter_username = Some(normalized); + } + } + + if let Some(twitch) = request.twitch_username.as_deref() { + let trimmed = twitch.trim(); + if trimmed.is_empty() { + twitch_username = None; + } else { + if let Err(reason) = validate_twitch_username(trimmed) { + return Err(EditProfileError::BadInput(reason)); + } + twitch_username = Some(trimmed); + } + } + + if let Some(discord) = request.discord_username.as_deref() { + let trimmed = discord.trim(); + if trimmed.is_empty() { + discord_username = None; + } else { + if let Err(reason) = validate_discord_username(trimmed) { + return Err(EditProfileError::BadInput(reason)); + } + discord_username = Some(trimmed); + } + } + + if let Some(github) = request.github_username.as_deref() { + let trimmed = github.trim(); + if trimmed.is_empty() { + github_username = None; + } else { + if let Err(reason) = validate_github_username(trimmed) { + return Err(EditProfileError::BadInput(reason)); + } + github_username = Some(trimmed); + } + } + + if let Some(cashapp) = request.cashapp_username.as_deref() { + let trimmed = cashapp.trim(); + if trimmed.is_empty() { + cashapp_username = None; + } else { + if let Err(reason) = validate_cashapp_username(trimmed) { + return Err(EditProfileError::BadInput(reason)); + } + let normalized = normalize_cashapp_username_for_storage(trimmed); + cashapp_username = Some(normalized); + } + } + + if let Some(website) = request.website_url.as_deref() { + let trimmed = website.trim(); + if trimmed.is_empty() { + website_url = None; + } else { + if let Err(reason) = validate_website_url(trimmed) { + return Err(EditProfileError::BadInput(reason)); + } + website_url = Some(trimmed); + } + } + + if let Some(markdown) = request.profile_markdown.as_deref() { + if contains_slurs(markdown) { + return Err(EditProfileError::BadInput("profile contains slurs".to_string())); + } + + let markdown = markdown.trim().to_string(); + let html = markdown_to_html(&markdown); + + profile_markdown = Some(markdown); + profile_html = Some(html); + } + + let ip_address = get_request_ip(&http_request); + + let preferred_tts_result_visibility = request.preferred_tts_result_visibility + .unwrap_or(RecordVisibility::Hidden); + + let preferred_w2l_result_visibility = request.preferred_w2l_result_visibility + .unwrap_or(RecordVisibility::Hidden); + + let query_result = if editor_is_original_user { + // We need to store the IP address details. + sqlx::query!( + r#" +UPDATE users +SET + profile_markdown = ?, + profile_rendered_html = ?, + preferred_tts_result_visibility = ?, + preferred_w2l_result_visibility = ?, + discord_username = ?, + twitter_username = ?, + twitch_username = ?, + github_username = ?, + cashapp_username = ?, + website_url = ?, + ip_address_last_update = ?, + version = version + 1 + +WHERE users.token = ? +LIMIT 1 + "#, + profile_markdown, + profile_html, + preferred_tts_result_visibility.to_str(), + preferred_w2l_result_visibility.to_str(), + discord_username, + twitter_username, + twitch_username, + github_username, + cashapp_username, + website_url, + ip_address.clone(), + user_record.user_token.clone(), + ) + .execute(&server_state.mysql_pool) + .await + } else { + // We need to store the moderator details. + // Also, mods shouldn't change user preferences. + sqlx::query!( + r#" +UPDATE users +SET + profile_markdown = ?, + profile_rendered_html = ?, + discord_username = ?, + twitter_username = ?, + twitch_username = ?, + github_username = ?, + cashapp_username = ?, + website_url = ?, + version = version + 1 + +WHERE users.token = ? +LIMIT 1 + "#, + profile_markdown, + profile_html, + discord_username, + twitter_username, + twitch_username, + github_username, + cashapp_username, + website_url, + user_record.user_token.clone(), + ) + .execute(&server_state.mysql_pool) + .await + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Profile edit DB error: {:?}", err); + return Err(EditProfileError::ServerError); + } + }; + + let response = EditProfileSuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| EditProfileError::BadInput("".to_string()))?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/get_profile.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/get_profile.rs new file mode 100644 index 0000000000..b838e2ed10 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/get_profile.rs @@ -0,0 +1,261 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::helpers::boolean_converters::i8_to_bool; +use database_queries::queries::users::user_badges::list_user_badges::UserBadgeForList; +use database_queries::queries::users::user_badges::list_user_badges::list_user_badges; +use log::{info, warn, log}; +use md5::{Md5, Digest}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// TODO: This is duplicated in query_user_profile +// TODO: This handler has embedded queries. + +#[derive(Serialize)] +pub struct UserProfileRecordForResponse { + pub user_token: String, + pub username: String, + pub display_name: String, + pub email_gravatar_hash: String, + pub profile_markdown: String, + pub profile_rendered_html: String, + pub user_role_slug: String, + pub disable_gravatar: bool, + pub preferred_tts_result_visibility: RecordVisibility, + pub preferred_w2l_result_visibility: RecordVisibility, + pub discord_username: Option, + pub twitch_username: Option, + pub twitter_username: Option, + pub patreon_username: Option, + pub github_username: Option, + pub cashapp_username: Option, + pub website_url: Option, + pub badges: Vec, + pub created_at: DateTime, + pub maybe_moderator_fields: Option, +} + +#[derive(Serialize)] +pub struct UserProfileModeratorFields { + pub is_banned: bool, + pub maybe_mod_comments: Option, + pub maybe_mod_user_token: Option, +} + +#[derive(Serialize)] +pub struct ProfileSuccessResponse { + pub success: bool, + pub user: Option, +} + +#[derive(Serialize)] +pub struct RawUserProfileRecord { + pub user_token: String, + pub username: String, + pub email_gravatar_hash: String, + pub display_name: String, + pub profile_markdown: String, + pub profile_rendered_html: String, + pub user_role_slug: String, + pub disable_gravatar: i8, + pub preferred_tts_result_visibility: RecordVisibility, + pub preferred_w2l_result_visibility: RecordVisibility, + pub discord_username: Option, + pub twitch_username: Option, + pub twitter_username: Option, + pub patreon_username: Option, + pub github_username: Option, + pub cashapp_username: Option, + pub website_url: Option, + pub is_banned: i8, + pub maybe_mod_comments: Option, + pub maybe_mod_user_token: Option, + pub created_at: DateTime, +} + +#[derive(Debug)] +pub enum ProfileError { + ServerError, + NotFound, +} + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetProfilePathInfo { + username: String, +} + +impl ResponseError for ProfileError { + fn status_code(&self) -> StatusCode { + match *self { + ProfileError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + ProfileError::NotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ProfileError::ServerError => "server error".to_string(), + ProfileError::NotFound => "not found error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ProfileError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_profile_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ProfileError::ServerError + })?; + + let mut is_mod = false; + + if let Some(user_session) = &maybe_user_session { + is_mod = user_session.can_ban_users; + } + + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_profile_record = sqlx::query_as!( + RawUserProfileRecord, + r#" +SELECT + users.token as user_token, + username, + display_name, + email_gravatar_hash, + profile_markdown, + profile_rendered_html, + user_role_slug, + disable_gravatar, + preferred_tts_result_visibility as `preferred_tts_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`, + preferred_w2l_result_visibility as `preferred_w2l_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`, + discord_username, + twitch_username, + twitter_username, + patreon_username, + github_username, + cashapp_username, + website_url, + is_banned, + maybe_mod_comments, + maybe_mod_user_token, + created_at +FROM users +WHERE + users.username = ? + AND users.user_deleted_at IS NULL + AND users.mod_deleted_at IS NULL + "#, + &path.username, + ) + .fetch_one(&server_state.mysql_pool) + .await; // TODO: This will return error if it doesn't exist + + let profile_record : RawUserProfileRecord = match maybe_profile_record { + Ok(profile_record) => profile_record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + warn!("Invalid user"); + return Err(ProfileError::NotFound); + }, + _ => { + warn!("User profile query error: {:?}", err); + return Err(ProfileError::ServerError); + } + } + } + }; + + let is_banned = i8_to_bool(profile_record.is_banned); + + if is_banned && !is_mod { + // Can't see banned users. + return Err(ProfileError::NotFound); + } + + let badges = list_user_badges(&server_state.mysql_pool, &profile_record.user_token) + .await + .unwrap_or_else(|err| { + warn!("Error querying badges: {:?}", err); + return Vec::new(); // NB: Fine if this fails. Not sure why it would. + }); + + let maybe_mod_fields = if is_mod { + Some(UserProfileModeratorFields { + is_banned, + maybe_mod_comments: profile_record.maybe_mod_comments.clone(), + maybe_mod_user_token: profile_record.maybe_mod_user_token.clone(), + }) + } else { + None + }; + + let profile_for_response = UserProfileRecordForResponse { + user_token: profile_record.user_token, + username: profile_record.username, + display_name: profile_record.display_name, + email_gravatar_hash: profile_record.email_gravatar_hash, + profile_markdown: profile_record.profile_markdown, + profile_rendered_html: profile_record.profile_rendered_html, + user_role_slug: profile_record.user_role_slug, + disable_gravatar: i8_to_bool(profile_record.disable_gravatar), + preferred_tts_result_visibility: profile_record.preferred_tts_result_visibility, + preferred_w2l_result_visibility: profile_record.preferred_w2l_result_visibility, + discord_username: profile_record.discord_username, + twitch_username: profile_record.twitch_username, + twitter_username: profile_record.twitter_username, + patreon_username: profile_record.patreon_username, + github_username: profile_record.github_username, + cashapp_username: profile_record.cashapp_username, + website_url: profile_record.website_url, + created_at: profile_record.created_at, + maybe_moderator_fields: maybe_mod_fields, + badges, + }; + + let response = ProfileSuccessResponse { + success: true, + user: Some(profile_for_response), + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ProfileError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_tts_inference_results.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_tts_inference_results.rs new file mode 100644 index 0000000000..8197eceaaa --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_tts_inference_results.rs @@ -0,0 +1,179 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Query}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::tts::tts_results::list_tts_results_query_builder::{ListTtsResultsQueryBuilder, TtsInferenceRecordForList}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct ListTtsInferenceResultsForUserPathInfo { + pub username: String, +} + +#[derive(Deserialize)] +pub struct ListTtsInferenceResultsForUserQuery { + pub sort_ascending: Option, + pub limit: Option, + pub cursor: Option, + pub cursor_is_reversed: Option, +} + +#[derive(Serialize)] +pub struct ListTtsInferenceResultsForUserSuccessResponse { + pub success: bool, + pub results: Vec, + pub cursor_next: Option, + pub cursor_previous: Option, +} + +#[derive(Debug)] +pub enum ListTtsInferenceResultsForUserError { + ServerError, +} + +impl ResponseError for ListTtsInferenceResultsForUserError { + fn status_code(&self) -> StatusCode { + match *self { + ListTtsInferenceResultsForUserError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListTtsInferenceResultsForUserError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTtsInferenceResultsForUserError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_user_tts_inference_results_handler( + http_request: HttpRequest, + path: Path, + query: Query, + server_state: web::Data> +) -> Result +{ + info!("Fetching inference results for user: {}", &path.username); + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListTtsInferenceResultsForUserError::ServerError + })?; + + // Permissions & ACLs + let mut viewer_is_user = false; + let mut is_mod_that_can_see_deleted = false; + + match maybe_user_session { + None => {}, + Some(session) => { + viewer_is_user = path.username == session.username; + is_mod_that_can_see_deleted = session.can_delete_other_users_tts_results; + }, + }; + + // TODO: Enforce real maximums and defaults + let limit = query.limit.unwrap_or(25); + //let limit = std::cmp::max(limit, 100); + + let sort_ascending = query.sort_ascending.unwrap_or(false); + let cursor_is_reversed = query.cursor_is_reversed.unwrap_or(false); + + let cursor = if let Some(cursor) = query.cursor.as_deref() { + let cursor = server_state.sort_key_crypto.decrypt_id(cursor) + .map_err(|e| { + warn!("crypto error: {:?}", e); + return ListTtsInferenceResultsForUserError::ServerError + })?; + Some(cursor) + } else { + None + }; + + let include_user_hidden = viewer_is_user || is_mod_that_can_see_deleted; + + let mut query_builder = ListTtsResultsQueryBuilder::new() + .sort_ascending(sort_ascending) + .scope_creator_username(Some(path.username.as_ref())) + .include_user_hidden(include_user_hidden) + .include_mod_disabled_results(is_mod_that_can_see_deleted) + .limit(limit) + .cursor_is_reversed(cursor_is_reversed) + .offset(cursor); + + let query_results = query_builder.perform_query_for_page(&server_state.mysql_pool) + .await; + + let results_page = match query_results { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(ListTtsInferenceResultsForUserError::ServerError); + } + }; + + let cursor_next = if let Some(id) = results_page.last_id { + let cursor = server_state.sort_key_crypto.encrypt_id(id as u64) + .map_err(|e| { + warn!("crypto error: {:?}", e); + return ListTtsInferenceResultsForUserError::ServerError + })?; + Some(cursor) + } else { + None + }; + + let cursor_previous = if let Some(id) = results_page.first_id { + let cursor = server_state.sort_key_crypto.encrypt_id(id as u64) + .map_err(|e| { + warn!("crypto error: {:?}", e); + return ListTtsInferenceResultsForUserError::ServerError + })?; + Some(cursor) + } else { + None + }; + + let response = ListTtsInferenceResultsForUserSuccessResponse { + success: true, + results: results_page.inference_records, + cursor_next, + cursor_previous, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTtsInferenceResultsForUserError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_tts_models.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_tts_models.rs new file mode 100644 index 0000000000..ace2c6a14a --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_tts_models.rs @@ -0,0 +1,96 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::tts::tts_models::list_tts_models::{TtsModelRecordForList, list_tts_models}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetProfilePathInfo { + username: String, +} + +#[derive(Serialize)] +pub struct ListTtsModelsForUserSuccessResponse { + pub success: bool, + pub models: Vec, +} + +#[derive(Debug)] +pub enum ListTtsModelsForUserError { + ServerError, +} + +impl ResponseError for ListTtsModelsForUserError { + fn status_code(&self) -> StatusCode { + match *self { + ListTtsModelsForUserError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListTtsModelsForUserError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListTtsModelsForUserError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_user_tts_models_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data> +) -> Result +{ + info!("Fetching templates for user: {}", &path.username); + + let query_results = list_tts_models( + &server_state.mysql_pool, + Some(path.username.as_ref()), + false + ).await; + + let models = match query_results { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(ListTtsModelsForUserError::ServerError); + } + }; + + let response = ListTtsModelsForUserSuccessResponse { + success: true, + models, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListTtsModelsForUserError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_w2l_inference_results.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_w2l_inference_results.rs new file mode 100644 index 0000000000..5f8f585cdd --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_w2l_inference_results.rs @@ -0,0 +1,183 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Query}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::w2l::w2l_results::list_w2l_inference_results_query_builder::ListW2lResultsQueryBuilder; +use database_queries::queries::w2l::w2l_results::list_w2l_inference_results_query_builder::W2lInferenceRecordForList; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct ListW2lInferenceResultsForUserPathInfo { + username: String, +} + +#[derive(Deserialize)] +pub struct ListW2lInferenceResultsForUserQuery { + pub sort_ascending: Option, + pub limit: Option, + pub cursor: Option, + pub cursor_is_reversed: Option, +} + +#[derive(Serialize)] +pub struct ListW2lInferenceResultsForUserSuccessResponse { + pub success: bool, + pub results: Vec, + pub cursor_next: Option, + pub cursor_previous: Option, +} + +#[derive(Debug)] +pub enum ListW2lInferenceResultsForUserError { + ServerError, +} + +impl ResponseError for ListW2lInferenceResultsForUserError { + fn status_code(&self) -> StatusCode { + match *self { + ListW2lInferenceResultsForUserError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListW2lInferenceResultsForUserError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListW2lInferenceResultsForUserError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_user_w2l_inference_results_handler( + http_request: HttpRequest, + path: Path, + query: Query, + server_state: web::Data> +) -> Result +{ + info!("Fetching inference results for user: {}", &path.username); + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListW2lInferenceResultsForUserError::ServerError + })?; + + let mut is_author = false; + let mut is_mod_that_can_see_deleted = false; + let mut is_mod_that_can_approve_w2l_templates = false; + + match maybe_user_session { + None => {}, + Some(session) => { + is_author = session.username == path.username; + is_mod_that_can_see_deleted = session.can_delete_other_users_w2l_results; + is_mod_that_can_approve_w2l_templates = session.can_approve_w2l_templates; + }, + }; + + // TODO: Enforce real maximums and defaults + let limit = query.limit.unwrap_or(25); + //let limit = std::cmp::max(limit, 100); + + let sort_ascending = query.sort_ascending.unwrap_or(false); + let cursor_is_reversed = query.cursor_is_reversed.unwrap_or(false); + + let cursor = if let Some(cursor) = query.cursor.as_deref() { + let cursor = server_state.sort_key_crypto.decrypt_id(cursor) + .map_err(|e| { + warn!("crypto error: {:?}", e); + return ListW2lInferenceResultsForUserError::ServerError + })?; + Some(cursor) + } else { + None + }; + + let include_user_hidden = is_author || is_mod_that_can_see_deleted; + let include_unapproved = is_author || is_mod_that_can_approve_w2l_templates; + + let mut query_builder = ListW2lResultsQueryBuilder::new() + .sort_ascending(sort_ascending) + .scope_creator_username(Some(path.username.as_ref())) + .include_user_hidden(include_user_hidden) + .include_user_deleted_results(is_mod_that_can_see_deleted) + .include_mod_deleted_results(is_mod_that_can_see_deleted) + .include_templates_not_approved_for_public_listing(include_unapproved) + .limit(limit) + .cursor_is_reversed(cursor_is_reversed) + .offset(cursor); + + let query_results = query_builder.perform_query_for_page(&server_state.mysql_pool) + .await; + + let results_page = match query_results { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(ListW2lInferenceResultsForUserError::ServerError); + } + }; + + let cursor_next = if let Some(id) = results_page.last_id { + let cursor = server_state.sort_key_crypto.encrypt_id(id as u64) + .map_err(|e| { + warn!("crypto error: {:?}", e); + return ListW2lInferenceResultsForUserError::ServerError + })?; + Some(cursor) + } else { + None + }; + + let cursor_previous = if let Some(id) = results_page.first_id { + let cursor = server_state.sort_key_crypto.encrypt_id(id as u64) + .map_err(|e| { + warn!("crypto error: {:?}", e); + return ListW2lInferenceResultsForUserError::ServerError + })?; + Some(cursor) + } else { + None + }; + + let response = ListW2lInferenceResultsForUserSuccessResponse { + success: true, + results: results_page.inference_records, + cursor_next, + cursor_previous, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListW2lInferenceResultsForUserError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_w2l_templates.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_w2l_templates.rs new file mode 100644 index 0000000000..8c6f9f1f5d --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_w2l_templates.rs @@ -0,0 +1,120 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::w2l::w2l_templates::list_w2l_templates::{W2lTemplateRecordForList, list_w2l_templates}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetProfilePathInfo { + username: String, +} + +#[derive(Serialize)] +pub struct ListW2lTemplatesForUserSuccessResponse { + pub success: bool, + pub templates: Vec, +} + +#[derive(Debug)] +pub enum ListW2lTemplatesForUserError { + ServerError, +} + +impl ResponseError for ListW2lTemplatesForUserError { + fn status_code(&self) -> StatusCode { + match *self { + ListW2lTemplatesForUserError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListW2lTemplatesForUserError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListW2lTemplatesForUserError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_user_w2l_templates_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data> +) -> Result +{ + info!("Fetching templates for user: {}", &path.username); + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListW2lTemplatesForUserError::ServerError + })?; + + let mut viewer_is_original_user = false; + let mut viewer_is_moderator = false; + + match maybe_user_session { + None => {}, + Some(session) => { + viewer_is_moderator = session.can_delete_other_users_w2l_templates; + viewer_is_original_user = &session.username == &path.username; + }, + }; + + // The original user can see all their uploads, as can mods. + // Once they've been approved, everyone can see them. + let require_mod_approved = !(viewer_is_original_user || viewer_is_moderator); + + let query_results = list_w2l_templates( + &server_state.mysql_pool, + Some(path.username.as_ref()), + require_mod_approved + ).await; + + let templates = match query_results { + Ok(results) => results, + Err(e) => { + warn!("Query error: {:?}", e); + return Err(ListW2lTemplatesForUserError::ServerError); + } + }; + + let response = ListW2lTemplatesForUserSuccessResponse { + success: true, + templates, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListW2lTemplatesForUserError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/login.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/login.rs new file mode 100644 index 0000000000..c1b9745be0 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/login.rs @@ -0,0 +1,216 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::AnyhowResult; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::serialize_as_json_error::serialize_as_json_error; +use crate::server_state::ServerState; +use database_queries::helpers::boolean_converters::i8_to_bool; +use database_queries::queries::users::user_sessions::create_session::create_session_for_user; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt::Formatter; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct LoginRequest { + pub username_or_email: String, + pub password: String, +} + +#[derive(Serialize)] +pub struct LoginSuccessResponse { + pub success: bool, +} + +#[derive(Serialize, Debug)] +pub struct LoginErrorResponse { + pub success: bool, + pub error_type: LoginErrorType, + pub error_message: String, +} + +#[derive(Copy, Clone, Debug, Serialize)] +pub enum LoginErrorType { + InvalidCredentials, + ServerError, +} + +impl LoginErrorResponse { + fn invalid_credentials() -> Self { + Self { + success: false, + error_type: LoginErrorType::InvalidCredentials, + error_message: "invalid credentials".to_string() + } + } + fn server_error() -> Self { + Self { + success: false, + error_type: LoginErrorType::ServerError, + error_message: "server error".to_string() + } + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for LoginErrorResponse { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self.error_type) + } +} + +impl ResponseError for LoginErrorResponse { + fn status_code(&self) -> StatusCode { + match self.error_type { + LoginErrorType::InvalidCredentials => StatusCode::UNAUTHORIZED, + LoginErrorType::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +pub async fn login_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let check_username_or_email = request.username_or_email.to_lowercase(); + + let maybe_user = if check_username_or_email.contains("@") { + lookup_by_email(&check_username_or_email, &server_state.mysql_pool).await + } else { + lookup_by_username(&check_username_or_email, &server_state.mysql_pool).await + }; + + let user = match maybe_user { + Ok(user) => user, + Err(e) => { + // TODO: This isn't necessarily user error. I need to fix the above code to not lose error + // semantics. I also need to prevent user lookup attacks. + warn!("Login lookup error: {:?}", e); + return Err(LoginErrorResponse::invalid_credentials()); + } + }; + + let is_banned = i8_to_bool(user.is_banned); + if is_banned { + // We don't allow banned users back in. + return Err(LoginErrorResponse::invalid_credentials()); + } + + info!("login user found"); + + let actual_hash = match String::from_utf8(user.password_hash.clone()) { + Ok(hash) => hash, + Err(e) => { + warn!("Login hash hydration error: {:?}", e); + return Err(LoginErrorResponse::server_error()); + } + }; + + match bcrypt::verify(&request.password, &actual_hash) { + Err(e) => { + warn!("Login hash comparison error: {:?}", e); + return Err(LoginErrorResponse::server_error()); + } + Ok(is_valid) => { + if !is_valid { + info!("invalid credentials"); + return Err(LoginErrorResponse::invalid_credentials()); + } + // Good to go...! + }, + }; + + let ip_address = get_request_ip(&http_request); + + let create_session_result = + create_session_for_user(&user.token, &ip_address, &server_state.mysql_pool).await; + + let session_token = match create_session_result { + Ok(token) => token, + Err(e) => { + warn!("login create session error : {:?}", e); + return Err(LoginErrorResponse::server_error()); + } + }; + + info!("login session created"); + + let session_cookie = match server_state.cookie_manager.create_cookie(&session_token) { + Ok(cookie) => cookie, + Err(_) => return Err(LoginErrorResponse::server_error()), + }; + + let response = LoginSuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| LoginErrorResponse::server_error())?; + + Ok(HttpResponse::Ok() + .cookie(session_cookie) + .content_type("application/json") + .body(body)) +} + +#[derive(Debug)] +pub struct UserRecordForLogin { + token: String, + username: String, + email_address: String, + password_hash: Vec, + is_banned: i8, +} + +async fn lookup_by_username(username: &str, pool: &MySqlPool) -> AnyhowResult +{ + // NB: Lookup failure is Err(RowNotFound). + let record = sqlx::query_as!( + UserRecordForLogin, + r#" +SELECT token, username, email_address, password_hash, is_banned +FROM users +WHERE username = ? +LIMIT 1 + "#, + username.to_string(), + ) + .fetch_one(pool) + .await?; + + Ok(record) +} + +async fn lookup_by_email(email: &str, pool: &MySqlPool) -> AnyhowResult { + // NB: Lookup failure is Err(RowNotFound). + let record = sqlx::query_as!( + UserRecordForLogin, + r#" +SELECT token, username, email_address, password_hash, is_banned +FROM users +WHERE email_address = ? +LIMIT 1 + "#, + email.to_string(), + ) + .fetch_one(pool) + .await?; + + Ok(record) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/logout.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/logout.rs new file mode 100644 index 0000000000..f892574399 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/logout.rs @@ -0,0 +1,96 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct LogoutSuccessResponse { + pub success: bool, +} + +#[derive(Debug, Display)] +pub enum LogoutError { + ServerError, +} + +impl ResponseError for LogoutError { + fn status_code(&self) -> StatusCode { + match *self { + LogoutError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + LogoutError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn logout_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let delete_cookie = match http_request.cookie("session") { + Some(cookie) => { + match server_state.cookie_manager.decode_session_token(&cookie) { + Err(e) => { + warn!("Session cookie decode error: {:?}", e); + }, + Ok(session_token) => { + let _r = delete_session(&session_token, &server_state.mysql_pool).await; + } + } + + cookie // delete this cookie + }, + None => { + server_state.cookie_manager.delete_cookie() + } + }; + + let response = LogoutSuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| LogoutError::ServerError)?; + + Ok(HttpResponse::Ok() + .del_cookie(&delete_cookie) + .content_type("application/json") + .body(body)) +} + +async fn delete_session(session_token: &str, mysql_pool: &MySqlPool) -> AnyhowResult<()> { + let query_result = sqlx::query!( + r#" +UPDATE user_sessions +SET deleted_at = CURRENT_TIMESTAMP() +WHERE + token = ? + AND deleted_at IS NULL + "#, + session_token.to_string(), + ) + .execute(mysql_pool) + .await; + + Ok(()) +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs new file mode 100644 index 0000000000..8d5324ab12 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs @@ -0,0 +1,10 @@ +pub mod create_account; +pub mod edit_profile; +pub mod get_profile; +pub mod list_user_tts_inference_results; +pub mod list_user_tts_models; +pub mod list_user_w2l_inference_results; +pub mod list_user_w2l_templates; +pub mod login; +pub mod logout; +pub mod session_info; diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/session_info.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/session_info.rs new file mode 100644 index 0000000000..b13a984319 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/session_info.rs @@ -0,0 +1,150 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct UserInfo { + pub user_token: String, + pub username: String, + pub display_name: String, + pub email_gravatar_hash: String, + + // Usage permissions: + pub can_use_tts: bool, + pub can_use_w2l: bool, + pub can_delete_own_tts_results: bool, + pub can_delete_own_w2l_results: bool, + pub can_delete_own_account: bool, + + // Contribution permissions: + pub can_upload_tts_models: bool, + pub can_upload_w2l_templates: bool, + pub can_delete_own_tts_models: bool, + pub can_delete_own_w2l_templates: bool, + + // Moderation permissions: + pub can_approve_w2l_templates: bool, + pub can_edit_other_users_profiles: bool, + pub can_edit_other_users_tts_models: bool, + pub can_edit_other_users_w2l_templates: bool, + pub can_delete_other_users_tts_models: bool, + pub can_delete_other_users_tts_results: bool, + pub can_delete_other_users_w2l_templates: bool, + pub can_delete_other_users_w2l_results: bool, + pub can_ban_users: bool, + pub can_delete_users: bool, +} + +#[derive(Serialize)] +pub struct SessionInfoSuccessResponse { + pub success: bool, + pub logged_in: bool, + pub user: Option, +} + +#[derive(Debug, Display)] +pub enum SessionInfoError { + ServerError, +} + +impl ResponseError for SessionInfoError { + fn status_code(&self) -> StatusCode { + match *self { + SessionInfoError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + SessionInfoError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn session_info_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + SessionInfoError::ServerError + })?; + + let mut logged_in = false; + let mut user_info = None; + + match maybe_user_session { + None => {} + Some(session_data) => { + if !session_data.is_banned { + // NB: Banned users can't be logged in + logged_in = true; + user_info = Some(UserInfo { + user_token: session_data.user_token.clone(), + username: session_data.username.to_string(), + display_name: session_data.display_name.to_string(), + email_gravatar_hash: session_data.email_gravatar_hash.to_string(), + + // Usage permissions: + can_use_tts: session_data.can_use_tts, + can_use_w2l: session_data.can_use_w2l, + can_delete_own_tts_results: session_data.can_delete_own_tts_results, + can_delete_own_w2l_results: session_data.can_delete_own_w2l_results, + can_delete_own_account: session_data.can_delete_own_account, + + // Contribution permissions: + can_upload_tts_models: session_data.can_upload_tts_models, + can_upload_w2l_templates: session_data.can_upload_w2l_templates, + can_delete_own_tts_models: session_data.can_delete_own_tts_models, + can_delete_own_w2l_templates: session_data.can_delete_own_w2l_templates, + + // Moderation permissions: + can_approve_w2l_templates: session_data.can_approve_w2l_templates, + can_edit_other_users_profiles: session_data.can_edit_other_users_profiles, + can_edit_other_users_tts_models: session_data.can_edit_other_users_tts_models, + can_edit_other_users_w2l_templates: session_data.can_edit_other_users_w2l_templates, + can_delete_other_users_tts_models: session_data.can_delete_other_users_tts_models, + can_delete_other_users_tts_results: session_data.can_delete_other_users_tts_results, + can_delete_other_users_w2l_templates: session_data.can_delete_other_users_w2l_templates, + can_delete_other_users_w2l_results: session_data.can_delete_other_users_w2l_results, + can_ban_users: session_data.can_ban_users, + can_delete_users: session_data.can_delete_users, + }); + } + } + } + + let response = SessionInfoSuccessResponse { + success: true, + logged_in, + user: user_info, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| SessionInfoError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/check_if_voice_clone_request_submitted.rs b/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/check_if_voice_clone_request_submitted.rs new file mode 100644 index 0000000000..418cb41748 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/check_if_voice_clone_request_submitted.rs @@ -0,0 +1,130 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::queries::voice_clone_requests::get_voice_clone_request::{get_voice_clone_request_by_user_token, get_voice_clone_request_by_token}; +use database_queries::queries::voice_clone_requests::insert_voice_clone_requests_builder::InsertVoiceCloneRequestBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct CheckIfVoiceRequestSubmittedRequest { + /// Though we look up logged in users by user token, we can also look up non-logged-in users + /// with a token lookup. We can store this in frontend state or a cookie. + pub maybe_request_token: Option, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct CheckIfVoiceRequestSubmittedResponse { + pub success: bool, + pub has_submitted: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum CheckIfVoiceRequestSubmittedError { + BadInput(String), + NotAuthorized, + ServerError, +} + +impl ResponseError for CheckIfVoiceRequestSubmittedError { + fn status_code(&self) -> StatusCode { + match *self { + CheckIfVoiceRequestSubmittedError::BadInput(_) => StatusCode::BAD_REQUEST, + CheckIfVoiceRequestSubmittedError::NotAuthorized => StatusCode::UNAUTHORIZED, + CheckIfVoiceRequestSubmittedError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for CheckIfVoiceRequestSubmittedError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn check_if_voice_clone_request_submitted_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CheckIfVoiceRequestSubmittedError::ServerError + })?; + + let mut submitted_by_token = false; + let mut submitted_by_user = false; + + // Non-logged in accounts (via cookie or frontend mechanism) + if let Some(token) = request.maybe_request_token.as_deref() { + let request = get_voice_clone_request_by_token( + token, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Database error: {:?}", e); + CheckIfVoiceRequestSubmittedError::ServerError + })?; + + submitted_by_token = request.is_some(); + } + + // Logged in users + if let Some(user) = maybe_user_session.as_ref() { + let request = get_voice_clone_request_by_user_token( + &user.user_token, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Database error: {:?}", e); + CheckIfVoiceRequestSubmittedError::ServerError + })?; + + submitted_by_user = request.is_some(); + } + + let has_submitted = submitted_by_token || submitted_by_user; + + let response = CheckIfVoiceRequestSubmittedResponse { + success: true, + has_submitted, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CheckIfVoiceRequestSubmittedError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/create_voice_clone_request.rs b/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/create_voice_clone_request.rs new file mode 100644 index 0000000000..024b3e1fc0 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/create_voice_clone_request.rs @@ -0,0 +1,158 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use database_queries::queries::voice_clone_requests::insert_voice_clone_requests_builder::InsertVoiceCloneRequestBuilder; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +#[derive(Deserialize)] +pub struct CreateVoiceRequestRequest { + pub idempotency_token: String, + + // Contact + pub email_address: String, + pub discord_username: String, + + // TODO: Make enum. + // Visibility + pub is_for_private_use: bool, + pub is_for_public_use: bool, + + // Use + pub is_for_studio: bool, + pub is_for_twitch_tts: bool, + pub is_for_api_use: bool, + pub is_for_music: bool, + pub is_for_games: bool, + pub is_for_other: bool, + pub optional_notes_on_use: Option, + + // Subject/Ownership + pub is_own_voice: bool, + pub is_third_party_voice: bool, + pub optional_notes_on_subject: Option, + + // Equipment + pub has_clean_audio_recordings: bool, + pub has_good_microphone: bool, + + // Comments + pub optional_questions: Option, + pub optional_extra_comments: Option, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct CreateVoiceRequestResponse { + pub success: bool, +} + +// =============== Error Response =============== + +#[derive(Debug, Serialize)] +pub enum CreateVoiceRequestError { + BadInput(String), + NotAuthorized, + ServerError, +} + +impl ResponseError for CreateVoiceRequestError { + fn status_code(&self) -> StatusCode { + match *self { + CreateVoiceRequestError::BadInput(_) => StatusCode::BAD_REQUEST, + CreateVoiceRequestError::NotAuthorized => StatusCode::UNAUTHORIZED, + CreateVoiceRequestError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using DeriveMore since Clion doesn't understand it. +impl fmt::Display for CreateVoiceRequestError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +// =============== Handler =============== + +pub async fn create_voice_clone_request_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CreateVoiceRequestError::ServerError + })?; + + let creator_ip_address = get_request_ip(&http_request); + + let builder = InsertVoiceCloneRequestBuilder { + uuid_idempotency_token: request.idempotency_token.clone(), + maybe_user_token: maybe_user_session.map(|user| user.user_token.clone()), + email_address: request.email_address.clone(), + discord_username: request.discord_username.clone(), + is_for_private_use: request.is_for_private_use, + is_for_public_use: request.is_for_public_use, + is_for_studio: request.is_for_studio, + is_for_twitch_tts: request.is_for_twitch_tts, + is_for_api_use: request.is_for_api_use, + is_for_music: request.is_for_music, + is_for_games: request.is_for_games, + is_for_other: request.is_for_other, + optional_notes_on_use: request.optional_notes_on_use.clone(), + is_own_voice: request.is_own_voice, + is_third_party_voice: request.is_third_party_voice, + optional_notes_on_subject: request.optional_notes_on_subject.clone(), + has_clean_audio_recordings: request.has_clean_audio_recordings, + has_good_microphone: request.has_good_microphone, + optional_questions: request.optional_questions.clone(), + optional_extra_comments: request.optional_extra_comments.clone(), + ip_address_creation: creator_ip_address, + }; + + builder.insert(&server_state.mysql_pool) + .await + .map_err(|e| { + warn!("creation error: {:?}", e); + CreateVoiceRequestError::ServerError + })?; + + let response = CreateVoiceRequestResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| CreateVoiceRequestError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/mod.rs new file mode 100644 index 0000000000..8cea07f4c1 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/voice_clone_requests/mod.rs @@ -0,0 +1,2 @@ +pub mod create_voice_clone_request; +pub mod check_if_voice_clone_request_submitted; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/delete_w2l_result.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/delete_w2l_result.rs new file mode 100644 index 0000000000..4302111506 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/delete_w2l_result.rs @@ -0,0 +1,192 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::delete_role_disambiguation::DeleteRole; +use crate::util::delete_role_disambiguation::delete_role_disambiguation; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::w2l::w2l_results::delete_w2l_result_various_scopes::delete_w2l_inference_result_as_mod; +use database_queries::queries::w2l::w2l_results::delete_w2l_result_various_scopes::delete_w2l_inference_result_as_user; +use database_queries::queries::w2l::w2l_results::delete_w2l_result_various_scopes::undelete_w2l_inference_result_as_mod; +use database_queries::queries::w2l::w2l_results::delete_w2l_result_various_scopes::undelete_w2l_inference_result_as_user; +use database_queries::queries::w2l::w2l_results::query_w2l_result::select_w2l_result_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct DeleteW2lInferenceResultPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct DeleteW2lInferenceResultRequest { + set_delete: bool, + /// NB: this is only to disambiguate when a user is both a mod and an author. + as_mod: Option, +} + +#[derive(Debug)] +pub enum DeleteW2lInferenceResultError { + BadInput(String), + NotAuthorized, + NotFound, + ServerError, +} + +impl ResponseError for DeleteW2lInferenceResultError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteW2lInferenceResultError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteW2lInferenceResultError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteW2lInferenceResultError::NotFound => StatusCode::NOT_FOUND, + DeleteW2lInferenceResultError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + DeleteW2lInferenceResultError::BadInput(reason) => reason.to_string(), + DeleteW2lInferenceResultError::NotAuthorized => "unauthorized".to_string(), + DeleteW2lInferenceResultError::NotFound => "not found".to_string(), + DeleteW2lInferenceResultError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DeleteW2lInferenceResultError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn delete_w2l_inference_result_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data> +) -> Result { + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteW2lInferenceResultError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteW2lInferenceResultError::NotAuthorized); + } + }; + + // NB: First permission check. + // Only mods should see deleted models (both user_* and mod_* deleted). + let is_mod_that_can_see_deleted = user_session.can_delete_other_users_w2l_results; + + let inference_result_query_result = select_w2l_result_by_token( + &path.token, + is_mod_that_can_see_deleted, + &server_state.mysql_pool, + ).await; + + let w2l_inference_result = match inference_result_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(DeleteW2lInferenceResultError::ServerError); + } + Ok(None) => return Err(DeleteW2lInferenceResultError::NotFound), + Ok(Some(inference_result)) => inference_result, + }; + + // NB: Second set of permission checks + let is_author = w2l_inference_result.maybe_creator_user_token + .as_deref() + .map(|creator_token| creator_token == &user_session.user_token) + .unwrap_or(false); + + let is_mod = user_session.can_delete_other_users_w2l_results; + + if !is_author && !is_mod { + warn!("user is not allowed to delete inference results: {}", user_session.user_token); + return Err(DeleteW2lInferenceResultError::NotAuthorized); + } + + let delete_role = delete_role_disambiguation(is_mod, is_author, request.as_mod); + + let query_result = if request.set_delete { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to delete inference results: {}", user_session.user_token); + return Err(DeleteW2lInferenceResultError::NotAuthorized); + } + DeleteRole::AsUser => { + delete_w2l_inference_result_as_user( + &path.token, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + delete_w2l_inference_result_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + } else { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to delete inference results: {}", user_session.user_token); + return Err(DeleteW2lInferenceResultError::NotAuthorized); + } + DeleteRole::AsUser => { + // NB: Technically only mods can see their own inference_results here + undelete_w2l_inference_result_as_user( + &path.token, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + undelete_w2l_inference_result_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update w2l mod approval status DB error: {:?}", err); + return Err(DeleteW2lInferenceResultError::ServerError); + } + }; + + Ok(simple_json_success()) +} + diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/delete_w2l_template.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/delete_w2l_template.rs new file mode 100644 index 0000000000..a23d27f0b0 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/delete_w2l_template.rs @@ -0,0 +1,199 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use database_queries::queries::w2l::w2l_templates::delete_w2l_template_various_scopes::delete_w2l_template_as_mod; +use database_queries::queries::w2l::w2l_templates::delete_w2l_template_various_scopes::delete_w2l_template_as_user; +use database_queries::queries::w2l::w2l_templates::delete_w2l_template_various_scopes::undelete_w2l_template_as_mod; +use database_queries::queries::w2l::w2l_templates::delete_w2l_template_various_scopes::undelete_w2l_template_as_user; +use database_queries::queries::w2l::w2l_templates::get_w2l_template::select_w2l_template_by_token; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::delete_role_disambiguation::DeleteRole; +use crate::util::delete_role_disambiguation::delete_role_disambiguation; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use log::{info, warn, log}; +use std::fmt; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct DeleteW2lTemplatePathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct DeleteW2lTemplateRequest { + set_delete: bool, + /// NB: this is only to disambiguate when a user is both a mod and an author. + as_mod: Option, +} + +#[derive(Debug)] +pub enum DeleteW2lTemplateError { + BadInput(String), + NotAuthorized, + NotFound, + ServerError, +} + +impl ResponseError for DeleteW2lTemplateError { + fn status_code(&self) -> StatusCode { + match *self { + DeleteW2lTemplateError::BadInput(_) => StatusCode::BAD_REQUEST, + DeleteW2lTemplateError::NotAuthorized => StatusCode::UNAUTHORIZED, + DeleteW2lTemplateError::NotFound => StatusCode::NOT_FOUND, + DeleteW2lTemplateError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + DeleteW2lTemplateError::BadInput(reason) => reason.to_string(), + DeleteW2lTemplateError::NotAuthorized => "unauthorized".to_string(), + DeleteW2lTemplateError::NotFound => "not found".to_string(), + DeleteW2lTemplateError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for DeleteW2lTemplateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn delete_w2l_template_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data> +) -> Result { + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + DeleteW2lTemplateError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(DeleteW2lTemplateError::NotAuthorized); + } + }; + + // NB: First permission check. + // Only mods should see deleted models (both user_* and mod_* deleted). + let is_mod_that_can_see_deleted = user_session.can_delete_other_users_w2l_templates; + + let template_query_result = select_w2l_template_by_token( + &path.token, + is_mod_that_can_see_deleted, + &server_state.mysql_pool, + ).await; + + let w2l_template = match template_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(DeleteW2lTemplateError::ServerError); + } + Ok(None) => return Err(DeleteW2lTemplateError::NotFound), + Ok(Some(template)) => template, + }; + + // NB: Second set of permission checks + let is_author = w2l_template.creator_user_token == user_session.user_token; + let is_mod = user_session.can_delete_other_users_w2l_templates; + + if !is_author && !is_mod { + warn!("user is not allowed to delete templates: {}", user_session.user_token); + return Err(DeleteW2lTemplateError::NotAuthorized); + } + + if !is_mod { + if w2l_template.is_locked_from_user_modification || w2l_template.is_locked_from_use { + warn!("user is not allowed to delete templates (locked): {}", user_session.user_token); + return Err(DeleteW2lTemplateError::NotAuthorized); + } + } + + let ip_address = get_request_ip(&http_request); + + let delete_role = delete_role_disambiguation(is_mod, is_author, request.as_mod); + + let query_result = if request.set_delete { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to delete template: {}", user_session.user_token); + return Err(DeleteW2lTemplateError::NotAuthorized); + } + DeleteRole::AsUser => { + delete_w2l_template_as_user( + &path.token, + &ip_address, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + delete_w2l_template_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + } else { + match delete_role { + DeleteRole::ErrorDoNotDelete => { + warn!("user is not allowed to delete template: {}", user_session.user_token); + return Err(DeleteW2lTemplateError::NotAuthorized); + } + DeleteRole::AsUser => { + // NB: Technically only mods can see their own templates here + undelete_w2l_template_as_user( + &path.token, + &ip_address, + &server_state.mysql_pool + ).await + } + DeleteRole::AsMod => { + undelete_w2l_template_as_mod( + &path.token, + &user_session.user_token, + &server_state.mysql_pool + ).await + } + } + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update w2l mod approval status DB error: {:?}", err); + return Err(DeleteW2lTemplateError::ServerError); + } + }; + + Ok(simple_json_success()) +} + diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/edit_w2l_result.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/edit_w2l_result.rs new file mode 100644 index 0000000000..98c1301f5b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/edit_w2l_result.rs @@ -0,0 +1,189 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::queries::w2l::w2l_results::query_w2l_result::select_w2l_result_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct EditW2lResultPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct EditW2lResultRequest { + // ========== Author + Moderator options ========== + pub creator_set_visibility: Option, +} + +#[derive(Debug)] +pub enum EditW2lResultError { + BadInput(String), + NotAuthorized, + ResultNotFound, + ServerError, +} + +impl ResponseError for EditW2lResultError { + fn status_code(&self) -> StatusCode { + match *self { + EditW2lResultError::BadInput(_) => StatusCode::BAD_REQUEST, + EditW2lResultError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditW2lResultError::ResultNotFound => StatusCode::NOT_FOUND, + EditW2lResultError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EditW2lResultError::BadInput(reason) => reason.to_string(), + EditW2lResultError::NotAuthorized=> "unauthorized".to_string(), + EditW2lResultError::ResultNotFound => "not found".to_string(), + EditW2lResultError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditW2lResultError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn edit_w2l_inference_result_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditW2lResultError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditW2lResultError::NotAuthorized); + } + }; + + // NB: Moderators can see deleted results. + // Original creators cannot see them (unless they're moderators!) + let show_deleted_results = user_session.can_delete_other_users_w2l_results; + + // Moderators get to see all the fields. + let is_moderator = user_session.can_delete_other_users_w2l_results + || user_session.can_edit_other_users_w2l_templates; // TODO: Not an exact permission fit + + let inference_result_query_result = select_w2l_result_by_token( + &path.token, + show_deleted_results, + &server_state.mysql_pool + ).await; + + let mut inference_result = match inference_result_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(EditW2lResultError::ServerError); + } + Ok(None) => return Err(EditW2lResultError::ResultNotFound), + Ok(Some(inference_result)) => inference_result, + }; + + // NB: Second set of permission checks + let mut is_author = false; + if let Some(creator_user_token) = inference_result.maybe_creator_user_token.as_deref() { + is_author = creator_user_token == &user_session.user_token; + } + + if !is_author && !is_moderator { + warn!("user is not allowed to edit result: {}", user_session.user_token); + return Err(EditW2lResultError::NotAuthorized); + } + + // Author + Mod fields. + // These fields must be present on all requests. + let mut creator_set_visibility = RecordVisibility::Public; + + if let Some(visibility) = request.creator_set_visibility.as_deref() { + creator_set_visibility = RecordVisibility::from_str(visibility) + .map_err(|_| EditW2lResultError::BadInput("bad record visibility".to_string()))?; + } + + let ip_address = get_request_ip(&http_request); + + let query_result = if is_author { + // TODO: Don't update the original IP address. Create a new field. + // We need to store the IP address details. + sqlx::query!( + r#" +UPDATE w2l_results +SET + creator_set_visibility = ?, + creator_ip_address = ? +WHERE token = ? +LIMIT 1 + "#, + &creator_set_visibility.to_str(), + &ip_address, + &inference_result.w2l_result_token, + ) + .execute(&server_state.mysql_pool) + .await + } else { + // We need to store the moderator details. + sqlx::query!( + r#" +UPDATE w2l_results +SET + creator_set_visibility = ?, + maybe_mod_user_token = ? +WHERE token = ? +LIMIT 1 + "#, + &creator_set_visibility.to_str(), + &user_session.user_token, + &inference_result.w2l_result_token, + ) + .execute(&server_state.mysql_pool) + .await + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update W2L result DB error: {:?}", err); + return Err(EditW2lResultError::ServerError); + } + }; + + Ok(simple_json_success()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/edit_w2l_template.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/edit_w2l_template.rs new file mode 100644 index 0000000000..20a4f93214 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/edit_w2l_template.rs @@ -0,0 +1,296 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::{Path, Json}; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::http_server::web_utils::response_success_helpers::simple_json_success; +use crate::server_state::ServerState; +use crate::util::email_to_gravatar::email_to_gravatar; +use crate::util::markdown_to_html::markdown_to_html; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::queries::w2l::w2l_templates::get_w2l_template::select_w2l_template_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use user_input_common::check_for_slurs::contains_slurs; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct EditW2lTemplatePathInfo { + template_token: String, +} + +#[derive(Deserialize)] +pub struct EditW2lTemplateRequest { + // ========== Author + Moderator options ========== + pub title: Option, + pub description_markdown: Option, + pub creator_set_visibility: Option, + //pub updatable_slug: Option, + + // ========== Moderator options ========== + + pub is_public_listing_approved: Option, + pub is_locked_from_user_modification: Option, + pub is_locked_from_use: Option, + pub maybe_mod_comments: Option, +} + +#[derive(Debug)] +pub enum EditW2lTemplateError { + BadInput(String), + NotAuthorized, + TemplateNotFound, + ServerError, +} + +impl ResponseError for EditW2lTemplateError { + fn status_code(&self) -> StatusCode { + match *self { + EditW2lTemplateError::BadInput(_) => StatusCode::BAD_REQUEST, + EditW2lTemplateError::NotAuthorized => StatusCode::UNAUTHORIZED, + EditW2lTemplateError::TemplateNotFound => StatusCode::NOT_FOUND, + EditW2lTemplateError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EditW2lTemplateError::BadInput(reason) => reason.to_string(), + EditW2lTemplateError::NotAuthorized=> "unauthorized".to_string(), + EditW2lTemplateError::TemplateNotFound => "not found".to_string(), + EditW2lTemplateError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EditW2lTemplateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn edit_w2l_template_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EditW2lTemplateError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EditW2lTemplateError::NotAuthorized); + } + }; + + // NB: First permission check. + // Only mods should see deleted templates (both user_* and mod_* deleted). + let is_mod_that_can_see_deleted = user_session.can_delete_other_users_w2l_templates; + + let template_lookup_result = select_w2l_template_by_token( + &path.template_token, + is_mod_that_can_see_deleted, + &server_state.mysql_pool).await; + + let template_record = match template_lookup_result { + Ok(Some(result)) => { + info!("Found template: {}", result.template_token); + result + }, + Ok(None) => { + warn!("could not find template"); + return Err(EditW2lTemplateError::TemplateNotFound); + }, + Err(err) => { + warn!("error looking up template: {:?}", err); + return Err(EditW2lTemplateError::TemplateNotFound); + }, + }; + + // NB: Second set of permission checks + let is_author = template_record.creator_user_token == user_session.user_token; + let is_mod = user_session.can_edit_other_users_w2l_templates ; + + if !is_author && !is_mod { + return Err(EditW2lTemplateError::NotAuthorized); + } + + if !is_mod { + if template_record.is_locked_from_user_modification || template_record.is_locked_from_use { + return Err(EditW2lTemplateError::NotAuthorized); + } + } + + // Author + Mod fields. + // These fields must be present on all requests. + let mut title = None; + let mut description_markdown = None; + let mut description_html = None; + let mut creator_set_visibility = RecordVisibility::Public; + + if let Some(payload) = request.title.as_deref() { + if contains_slurs(payload) { + return Err(EditW2lTemplateError::BadInput("title contains slurs".to_string())); + } + + title = Some(payload.to_string()); + } + + if let Some(markdown) = request.description_markdown.as_deref() { + if contains_slurs(markdown) { + return Err(EditW2lTemplateError::BadInput("description contains slurs".to_string())); + } + + let markdown = markdown.trim().to_string(); + let html = markdown_to_html(&markdown); + + description_markdown = Some(markdown); + description_html = Some(html); + } + + if let Some(visibility) = request.creator_set_visibility.as_deref() { + creator_set_visibility = RecordVisibility::from_str(visibility) + .map_err(|_| EditW2lTemplateError::BadInput("bad record visibility".to_string()))?; + } + + let ip_address = get_request_ip(&http_request); + + let query_result = if is_author { + // We need to store the IP address details. + sqlx::query!( + r#" +UPDATE w2l_templates +SET + title = ?, + description_markdown = ?, + description_rendered_html = ?, + creator_set_visibility = ?, + creator_ip_address_last_update = ?, + version = version + 1 +WHERE token = ? +LIMIT 1 + "#, + &title, + &description_markdown, + &description_html, + &creator_set_visibility.to_str(), + &ip_address, + &template_record.template_token, + ) + .execute(&server_state.mysql_pool) + .await + } else { + // We need to store the moderator details. + sqlx::query!( + r#" +UPDATE w2l_templates +SET + title = ?, + description_markdown = ?, + description_rendered_html = ?, + creator_set_visibility = ?, + maybe_mod_user_token = ?, + version = version + 1 +WHERE token = ? +LIMIT 1 + "#, + &title, + &description_markdown, + &description_html, + &creator_set_visibility.to_str(), + &user_session.user_token, + &template_record.template_token, + ) + .execute(&server_state.mysql_pool) + .await + }; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update W2L template edit DB error: {:?}", err); + return Err(EditW2lTemplateError::ServerError); + } + }; + + // TODO: This is lazy and suboptimal af to UPDATE again. + // The reason we're doing this is because `sqlx` only does static type checking of queries + // with string literals. It does not support dynamic query building, thus the PREDICATES + // MUST BE HELD CONSTANT (at least in type signature). :( + if is_mod { + update_mod_details( + &request, + &user_session.user_token, + &template_record.template_token, + &server_state.mysql_pool + ).await?; + } + + Ok(simple_json_success()) +} + +async fn update_mod_details( + request: &Json, + mod_user_token: &str, + template_token: &str, + mysql_pool: &MySqlPool +) -> Result<(), EditW2lTemplateError> { + + let is_public_listing_approved= request.is_public_listing_approved.unwrap_or(false); + let is_locked_from_user_modification = request.is_locked_from_user_modification.unwrap_or(false); + let is_locked_from_use = request.is_locked_from_use.unwrap_or(false); + + let query_result = sqlx::query!( + r#" +UPDATE w2l_templates +SET + is_public_listing_approved = ?, + is_locked_from_user_modification = ?, + is_locked_from_use = ?, + maybe_mod_comments = ?, + maybe_mod_user_token = ?, + version = version + 1 +WHERE token = ? +LIMIT 1 + "#, + is_public_listing_approved, + is_locked_from_user_modification, + is_locked_from_use, + request.maybe_mod_comments, + mod_user_token, + template_token + ) + .execute(mysql_pool) + .await; + + match query_result { + Ok(_) => Ok(()), + Err(err) => { + warn!("Update W2L template (mod details) DB error: {:?}", err); + Err(EditW2lTemplateError::ServerError) + } + } +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_infer_w2l.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_infer_w2l.rs new file mode 100644 index 0000000000..4ec1a9413b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_infer_w2l.rs @@ -0,0 +1,211 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::column_types::record_visibility::RecordVisibility; +use log::{info, warn, log}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct InferW2lRequest { + w2l_template_token: Option, + tts_inference_result_token: Option, + creator_set_visibility: Option, +} + +#[derive(Serialize)] +pub struct InferW2lSuccessResponse { + pub success: bool, +} + +#[derive(Debug)] +pub enum InferW2lError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for InferW2lError { + fn status_code(&self) -> StatusCode { + match *self { + InferW2lError::BadInput(_) => StatusCode::BAD_REQUEST, + InferW2lError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + InferW2lError::NotAuthorized => StatusCode::UNAUTHORIZED, + InferW2lError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + InferW2lError::BadInput(reason) => reason.to_string(), + InferW2lError::ServerError => "server error".to_string(), + InferW2lError::NotAuthorized => "not authorized".to_string(), + InferW2lError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for InferW2lError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn infer_w2l_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + + if true { + unimplemented!("this isn't finished"); + } + + let maybe_session = server_state + .session_checker + .maybe_get_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + InferW2lError::ServerError + })?; + + // TODO + //if let Err(_err) = server_state.redis_rate_limiter.rate_limit_request(&http_request) { + // return Err(InferW2lError::RateLimited); + //} + + let mut maybe_user_token : Option = maybe_session + .as_ref() + .map(|user_session| user_session.user_token.to_string()); + + info!("Enqueue infer w2l by user token: {:?}", maybe_user_token); + + let w2l_template_token = match &request.w2l_template_token { + None => { + // TODO: Allow image uploads. + return Err(InferW2lError::BadInput("w2l token is required".to_string())); + }, + Some(t) => { + // TODO: CHECK DATABASE! + t.to_string() + }, + }; + + let tts_inference_result_token = match &request.tts_inference_result_token { + None => { + // TODO: Allow audio uploads. + return Err(InferW2lError::BadInput("tts token is required".to_string())); + }, + Some(t) => { + // TODO: CHECK DATABASE! + t.to_string() + }, + }; + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + InferW2lError::ServerError + })?; + + let redis_count_key = RedisKeys::w2l_template_usage_count(&w2l_template_token); + + redis.incr(&redis_count_key, 1) + .map_err(|e| { + warn!("redis error: {:?}", e); + InferW2lError::ServerError + })?; + + let ip_address = get_request_ip(&http_request); + let creator_set_visibility = "public".to_string(); + + if true { + unimplemented!("this isn't finished"); + } + + let query_result = sqlx::query!( + r#" +INSERT INTO w2l_inference_jobs +SET + maybe_w2l_template_token = ?, + maybe_tts_inference_result_token = ?, + maybe_public_audio_bucket_location = NULL, + maybe_public_image_bucket_location = NULL, + maybe_creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + status = "pending" + "#, + w2l_template_token, + tts_inference_result_token, + maybe_user_token, + ip_address, + creator_set_visibility + ) + .execute(&server_state.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("New w2l inference job creation DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(InferW2lError::ServerError); + } + }; + + info!("new w2l inference job id: {}", record_id); + + let response = InferW2lSuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| InferW2lError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs new file mode 100644 index 0000000000..97257db6e9 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_infer_w2l_with_uploads.rs @@ -0,0 +1,413 @@ +use actix_http::http::header; +use actix_multipart::{Multipart, Field}; +use actix_web::HttpResponseBuilder; +use actix_web::body::ResponseBody::Body; +use actix_web::http::StatusCode; +use actix_web::http::header::ContentDisposition; +use actix_web::web::{Data, Json, BytesMut}; +use actix_web::{middleware, web, App, Error, HttpResponse, HttpServer, HttpRequest, Either, ResponseError}; +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::token::random_uuid::generate_random_uuid; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::read_multipart_field_bytes::checked_read_multipart_bytes; +use crate::http_server::web_utils::read_multipart_field_bytes::read_multipart_field_as_boolean; +use crate::http_server::web_utils::read_multipart_field_bytes::read_multipart_field_as_text; +use crate::http_server::web_utils::read_multipart_field_bytes::read_multipart_field_bytes; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::buckets::bucket_paths::hash_to_bucket_path; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::tokens::Tokens; +use futures::{StreamExt, TryStreamExt}; +use log::{warn, info}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use sqlx::MySqlPool; +use sqlx::error::Error::Database; +use std::fmt; +use std::io::Write; +use std::sync::Arc; + +const BUCKET_AUDIO_FILE_NAME : &'static str = "input_audio_file"; +const BUCKET_IMAGE_FILE_NAME: &'static str = "input_image_file"; +const BUCKET_VIDEO_FILE_NAME : &'static str = "input_video_file"; + +const MIN_BYTES : usize = 10; +const MAX_BYTES : usize = 1024 * 1024 * 20; + +/// Just to query for existence +#[derive(Serialize)] +pub struct W2lTemplateExistenceRecord { + pub template_token: String, +} + +#[derive(Serialize)] +pub struct InferW2lWithUploadSuccessResponse { + pub success: bool, + /// This is how frontend clients can request the job execution status. + pub inference_job_token: String, +} + +#[derive(Debug)] +pub enum InferW2lWithUploadError { + BadInput(String), + NotAuthorized, + EmptyFileUploaded, + ServerError, + RateLimited, +} + +impl ResponseError for InferW2lWithUploadError { + fn status_code(&self) -> StatusCode { + match *self { + InferW2lWithUploadError::BadInput(_) => StatusCode::BAD_REQUEST, + InferW2lWithUploadError::NotAuthorized => StatusCode::UNAUTHORIZED, + InferW2lWithUploadError::EmptyFileUploaded => StatusCode::BAD_REQUEST, + InferW2lWithUploadError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + InferW2lWithUploadError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + InferW2lWithUploadError::BadInput(reason) => reason.to_string(), + InferW2lWithUploadError::NotAuthorized => "unauthorized".to_string(), + InferW2lWithUploadError::EmptyFileUploaded => "empty file uploaded".to_string(), + InferW2lWithUploadError::ServerError => "server error".to_string(), + InferW2lWithUploadError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for InferW2lWithUploadError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +/// This handles audio uploads w/ W2L templates. +pub async fn enqueue_infer_w2l_with_uploads( + http_request: HttpRequest, + server_state: web::Data>, + mut payload: Multipart +) -> Result { + + let ip_address = get_request_ip(&http_request); + + // ==================== READ SESSION ==================== // + + let maybe_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref user) => { + if user.is_banned { + return Err(InferW2lWithUploadError::NotAuthorized); + } + &server_state.redis_rate_limiters.logged_in + }, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(InferW2lWithUploadError::RateLimited); + } + + // ==================== SESSION DETAILS ==================== // + + let mut maybe_user_token : Option = maybe_session + .as_ref() + .map(|user_session| user_session.user_token.to_string()); + + let maybe_user_preferred_visibility : Option = maybe_session + .as_ref() + .map(|user_session| user_session.preferred_tts_result_visibility); + + let set_visibility = maybe_user_preferred_visibility + .unwrap_or(RecordVisibility::Public); + + info!("Enqueue infer w2l by user token: {:?}", maybe_user_token); + + // ==================== READ MULTIPART REQUEST ==================== // + + info!("Reading multipart request..."); + + let mut maybe_uuid_idempotency_token: Option = None; + let mut maybe_template_token: Option = None; + let mut maybe_audio_file_name : Option = None; + let mut audio_bytes = BytesMut::with_capacity(0); + + while let Ok(Some(mut field)) = payload.try_next().await { + let mut field_name = "".to_string(); + let mut filename = "".to_string(); + + if let Some(content_disposition) = field.content_disposition() { + field_name = content_disposition.get_name() + .map(|s| s.to_string()) + .unwrap_or("".to_string()); + filename = content_disposition.get_filename() + .map(|s| s.to_string()) + .unwrap_or("".to_string()); + } + + match field_name.as_ref() { + "uuid_idempotency_token" => { + // Form text field. + maybe_uuid_idempotency_token = read_multipart_field_as_text(&mut field).await + .map_err(|e| { + warn!("Error reading idempotency token: {:}", e); + InferW2lWithUploadError::ServerError + })?; + }, + "template_token" => { + // Form text field. + maybe_template_token = read_multipart_field_as_text(&mut field).await + .map_err(|e| { + warn!("Error reading template token: {:}", e); + InferW2lWithUploadError::ServerError + })?; + }, + "audio" => { + // Form binary data. + maybe_audio_file_name = Some(filename.to_string()); + + let maybe_bytes = checked_read_multipart_bytes(&mut field).await + .map_err(|e| { + warn!("Error reading audio upload: {:}", e); + InferW2lWithUploadError::ServerError + })?; + + audio_bytes = match maybe_bytes { + Some(bytes) => bytes, + None => { + warn!("Empty file uploaded"); + return Err(InferW2lWithUploadError::EmptyFileUploaded); // Nothing was uploaded! + }, + }; + }, + _ => continue, + } + + info!("Saved file: {}.", &filename); + } + + // ==================== CHECK REQUEST ==================== // + + let template_token = match &maybe_template_token { + Some(ref token) => token.to_string(), + None => { + return Err(InferW2lWithUploadError::BadInput("No template selected".to_string())); + } + }; + + let uuid_idempotency_token = match maybe_uuid_idempotency_token { + Some(token) => token, + None => { + return Err(InferW2lWithUploadError::BadInput("No uuid idempotency token".to_string())); + } + }; + + let exists = check_template_exists(&template_token, &server_state.mysql_pool).await + .map_err(|e| { + warn!("error checking tmpl existence : {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + if !exists { + return Err(InferW2lWithUploadError::BadInput("Template does not exist".to_string())); + } + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + let redis_count_key = RedisKeys::w2l_template_usage_count(&template_token); + + redis.incr(&redis_count_key, 1) + .map_err(|e| { + warn!("redis error: {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + // ==================== ANALYZE AND UPLOAD AUDIO FILE ==================== // + + let mut audio_type = "application/octet-stream".to_string(); + + if let Some(maybe_type) = infer::get(audio_bytes.as_ref()) { + audio_type = maybe_type.mime_type().to_string(); + } + + let upload_uuid = generate_random_uuid(); + + let audio_upload_bucket_hash = upload_uuid.clone(); + + let audio_upload_bucket_path = hash_to_bucket_path( + &upload_uuid, + Some(&server_state.audio_uploads_bucket_root) + ).map_err(|e| { + warn!("Hash bucket path error: {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + info!("Uploading audio to bucket..."); + server_state.private_bucket_client.upload_file_with_content_type( + &audio_upload_bucket_path, + audio_bytes.as_ref(), + &audio_type) + .await + .map_err(|e| { + warn!("Upload audio bytes to bucket error: {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + // ==================== SAVE JOB RECORD ==================== // + + // This token is returned to the client. + let job_token = Tokens::new_w2l_inference_job() + .map_err(|e| { + warn!("Error creating token"); + InferW2lWithUploadError::ServerError + })?; + + info!("Creating w2l inference job record..."); + + let query_result = sqlx::query!( + r#" +INSERT INTO w2l_inference_jobs +SET + token = ?, + uuid_idempotency_token = ?, + + maybe_w2l_template_token = ?, + maybe_public_audio_bucket_hash = ?, + maybe_public_audio_bucket_location = ?, + + maybe_original_audio_filename = ?, + maybe_audio_mime_type = ?, + + maybe_creator_user_token = ?, + creator_ip_address = ?, + disable_end_bump = false, + creator_set_visibility = ?, + status = "pending" + "#, + job_token.to_string(), + uuid_idempotency_token.to_string(), + maybe_template_token.clone(), + Some(audio_upload_bucket_hash.clone()), + Some(audio_upload_bucket_path.clone()), + maybe_audio_file_name.clone(), + Some(audio_type.clone()), + maybe_user_token.clone(), + ip_address.to_string(), + set_visibility.to_str(), + ) + .execute(&server_state.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("New w2l template upload creation DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(InferW2lWithUploadError::ServerError); + } + }; + + server_state.firehose_publisher.enqueue_w2l_inference(maybe_user_token.as_deref(), &job_token, &template_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + InferW2lWithUploadError::ServerError + })?; + + + let response = InferW2lWithUploadSuccessResponse { + success: true, + inference_job_token: job_token.to_string(), + }; + + let body = serde_json::to_string(&response) + .map_err(|e| InferW2lWithUploadError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +async fn check_template_exists(template_token: &str, mysql_pool: &MySqlPool) -> AnyhowResult +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_template = sqlx::query_as!( + W2lTemplateExistenceRecord, + r#" +SELECT + token as template_token +FROM w2l_templates +WHERE + token = ? + AND user_deleted_at IS NULL + AND mod_deleted_at IS NULL + "#, + &template_token + ) + .fetch_one(mysql_pool) + .await; // TODO: This will return error if it doesn't exist + + let record_exists = match maybe_template { + Ok(record) => { + true + }, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + false + }, + _ => { + warn!("Infer w2l query error: {:?}", err); + return Err(anyhow!("infer w2l query error: {:?}", err)); + } + } + } + }; + + Ok(record_exists) +} + diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs new file mode 100644 index 0000000000..0087eafedf --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/enqueue_upload_w2l_template.rs @@ -0,0 +1,236 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::tokens::Tokens; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; +use user_input_common::check_for_slurs::contains_slurs; + +#[derive(Deserialize)] +pub enum W2lTemplateType { + /// unknown + Unknown, + /// video + Video, + /// image + Image, +} + +#[deprecated(note = "Use `RecordVisibility` instead!")] +#[derive(Deserialize)] +pub enum CreatorSetVisibility { + /// public + Public, + /// hidden + Hidden, + /// private + Private, +} + +#[derive(Deserialize)] +pub struct UploadW2lTemplateRequest { + idempotency_token: String, + title: String, + download_url: String, + template_type: Option, + creator_set_visibility: Option, +} + +#[derive(Serialize)] +pub struct UploadW2lTemplateSuccessResponse { + pub success: bool, + /// This is how frontend clients can request the job execution status. + pub job_token: String, +} + +#[derive(Debug)] +pub enum UploadW2lTemplateError { + BadInput(String), + MustBeLoggedIn, + ServerError, + RateLimited, +} + +impl ResponseError for UploadW2lTemplateError { + fn status_code(&self) -> StatusCode { + match *self { + UploadW2lTemplateError::BadInput(_) => StatusCode::BAD_REQUEST, + UploadW2lTemplateError::MustBeLoggedIn => StatusCode::UNAUTHORIZED, + UploadW2lTemplateError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + UploadW2lTemplateError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + UploadW2lTemplateError::BadInput(reason) => reason.to_string(), + UploadW2lTemplateError::MustBeLoggedIn => "user must be logged in".to_string(), + UploadW2lTemplateError::ServerError => "server error".to_string(), + UploadW2lTemplateError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for UploadW2lTemplateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn upload_w2l_template_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) { + return Err(UploadW2lTemplateError::RateLimited); + } + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + UploadW2lTemplateError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(UploadW2lTemplateError::MustBeLoggedIn); + } + }; + + if let Err(reason) = validate_idempotency_token(&request.idempotency_token) { + return Err(UploadW2lTemplateError::BadInput(reason)); + } + + if let Err(reason) = validate_model_title(&request.title) { + return Err(UploadW2lTemplateError::BadInput(reason)); + } + + if contains_slurs(&request.title) { + return Err(UploadW2lTemplateError::BadInput("title contains slurs".to_string())); + } + + let ip_address = get_request_ip(&http_request); + + let uuid = request.idempotency_token.to_string(); + let title = request.title.to_string(); + let download_url = request.download_url.to_string(); + + let template_type = "unknown".to_string(); + let creator_set_visibility = "public".to_string(); + + // This token is returned to the client. + let job_token = Tokens::new_w2l_template_upload_job() + .map_err(|e| { + warn!("Error creating token"); + UploadW2lTemplateError::ServerError + })?; + + + let query_result = sqlx::query!( + r#" +INSERT INTO w2l_template_upload_jobs +SET + token = ?, + uuid_idempotency_token = ?, + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + title = ?, + template_type = ?, + download_url = ?, + status = "pending" + "#, + job_token.to_string(), + uuid.to_string(), + user_session.user_token.to_string(), + ip_address.to_string(), + creator_set_visibility.to_string(), + title.to_string(), + template_type, + download_url, + ) + .execute(&server_state.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + warn!("New w2l template upload creation DB error: {:?}", err); + + // NB: SQLSTATE[23000]: Integrity constraint violation + // NB: MySQL Error Code 1062: Duplicate key insertion (this is harder to access) + match err { + Database(err) => { + let maybe_code = err.code().map(|c| c.into_owned()); + /*match maybe_code.as_deref() { + Some("23000") => { + if err.message().contains("username") { + return Err(UsernameTaken); + } else if err.message().contains("email_address") { + return Err(EmailTaken); + } + } + _ => {}, + }*/ + }, + _ => {}, + } + return Err(UploadW2lTemplateError::ServerError); + } + }; + + info!("new w2l template upload job id: {}", record_id); + + server_state.firehose_publisher.enqueue_w2l_template_upload(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + UploadW2lTemplateError::ServerError + })?; + + let response = UploadW2lTemplateSuccessResponse { + success: true, + job_token: job_token.to_string(), + }; + + let body = serde_json::to_string(&response) + .map_err(|e| UploadW2lTemplateError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} + +fn validate_idempotency_token(token: &str) -> Result<(), String> { + if token.len() != 36 { + return Err("idempotency token should be 36 characters".to_string()); + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_inference_job_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_inference_job_status.rs new file mode 100644 index 0000000000..475d8f28b1 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_inference_job_status.rs @@ -0,0 +1,202 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use log::{info, warn, log}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetW2lInferenceStatusPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct W2lInferenceJobStatusForResponse { + pub job_token: String, + + /// Primary status from the database (a state machine). + pub status: String, + + /// Extra, temporary status from Redis. + /// This can denote inference progress, and the Python code can write to it. + pub maybe_extra_status_description: Option, + + pub attempt_count: u8, + pub maybe_result_token: Option, + pub maybe_public_bucket_video_path: Option, + + pub maybe_w2l_template_token: Option, + pub w2l_template_type: String, + pub title: String, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct GetW2lInferenceStatusSuccessResponse { + pub success: bool, + pub state: W2lInferenceJobStatusForResponse, +} + +#[derive(Debug)] +pub enum GetW2lInferenceStatusError { + ServerError, +} + +#[derive(Serialize)] +pub struct W2lInferenceJobStatusRecord { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + pub maybe_result_token: Option, + pub maybe_public_bucket_video_path: Option, + + pub maybe_w2l_template_token: Option, + pub w2l_template_type: String, + + pub title: String, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl ResponseError for GetW2lInferenceStatusError { + fn status_code(&self) -> StatusCode { + match *self { + GetW2lInferenceStatusError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetW2lInferenceStatusError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetW2lInferenceStatusError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_w2l_inference_job_status_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_status = sqlx::query_as!( + W2lInferenceJobStatusRecord, + r#" +SELECT + jobs.token as job_token, + + jobs.status, + jobs.attempt_count, + jobs.on_success_result_token as maybe_result_token, + results.public_bucket_video_path as maybe_public_bucket_video_path, + + jobs.maybe_w2l_template_token, + w2l.template_type as w2l_template_type, + w2l.title, + + jobs.created_at, + jobs.updated_at + +FROM w2l_inference_jobs as jobs +JOIN w2l_templates as w2l + ON w2l.token = jobs.maybe_w2l_template_token +LEFT OUTER JOIN w2l_results as results + ON jobs.on_success_result_token = results.token + +WHERE jobs.token = ? + "#, + &path.token + ) + .fetch_one(&server_state.mysql_pool) + .await; // TODO: This will return error if it doesn't exist + + let record : W2lInferenceJobStatusRecord = match maybe_status { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(GetW2lInferenceStatusError::ServerError); + }, + _ => { + warn!("w2l template query error: {:?}", err); + return Err(GetW2lInferenceStatusError::ServerError); + } + } + } + }; + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + GetW2lInferenceStatusError::ServerError + })?; + + let extra_status_key = RedisKeys::w2l_inference_extra_status_info(&path.token); + let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { + Ok(Some(status)) => { + Some(status) + }, + Ok(None) => None, + Err(e) => { + warn!("redis error: {:?}", e); + None // Fail open + }, + }; + + let template_for_response = W2lInferenceJobStatusForResponse { + job_token: record.job_token.clone(), + status: record.status.clone(), + maybe_extra_status_description, + attempt_count: record.attempt_count as u8, + maybe_result_token: record.maybe_result_token.clone(), + maybe_public_bucket_video_path: record.maybe_public_bucket_video_path.clone(), + maybe_w2l_template_token: record.maybe_w2l_template_token.clone(), + w2l_template_type: record.w2l_template_type.clone(), + title: record.title.clone(), + created_at: record.created_at.clone(), + updated_at: record.updated_at.clone(), + }; + + let response = GetW2lInferenceStatusSuccessResponse { + success: true, + state: template_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetW2lInferenceStatusError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_result.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_result.rs new file mode 100644 index 0000000000..0204780f6a --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_result.rs @@ -0,0 +1,131 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::w2l::w2l_results::query_w2l_result::W2lResultRecordForResponse; +use database_queries::queries::w2l::w2l_results::query_w2l_result::select_w2l_result_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetW2lResultPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct GetW2lResultSuccessResponse { + pub success: bool, + pub result: W2lResultRecordForResponse, +} + +#[derive(Debug)] +pub enum GetW2lResultError { + ServerError, + NotFound, +} + +impl ResponseError for GetW2lResultError { + fn status_code(&self) -> StatusCode { + match *self { + GetW2lResultError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + GetW2lResultError::NotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetW2lResultError::ServerError => "server error".to_string(), + GetW2lResultError::NotFound => "not found".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetW2lResultError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_w2l_inference_result_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data> +) -> Result { + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetW2lResultError::ServerError + })?; + + let mut show_deleted_results = false; + let mut is_moderator = false; + + if let Some(user_session) = maybe_user_session { + // NB: Moderators can see deleted results. + // Original creators cannot see them (unless they're moderators!) + show_deleted_results = user_session.can_delete_other_users_w2l_results; + // Moderators get to see all the fields. + is_moderator = user_session.can_delete_other_users_w2l_results + || user_session.can_edit_other_users_w2l_templates; + } + + let inference_result_query_result = select_w2l_result_by_token( + &path.token, + show_deleted_results, + &server_state.mysql_pool + ).await; + + let mut inference_result = match inference_result_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(GetW2lResultError::ServerError); + } + Ok(None) => return Err(GetW2lResultError::NotFound), + Ok(Some(inference_result)) => inference_result, + }; + + if let Some(moderator_fields) = inference_result.maybe_moderator_fields.as_ref() { + // NB: The moderator fields will always be present before removal + // We don't want non-mods seeing stuff made by banned users. + if (moderator_fields.template_creator_is_banned || moderator_fields.result_creator_is_banned_if_user) + && !is_moderator{ + return Err(GetW2lResultError::NotFound); + } + } + + if !is_moderator { + inference_result.maybe_moderator_fields = None; + } + + let response = GetW2lResultSuccessResponse { + success: true, + result: inference_result, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetW2lResultError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_template.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_template.rs new file mode 100644 index 0000000000..ba3f883c4f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_template.rs @@ -0,0 +1,133 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::w2l::w2l_templates::get_w2l_template::W2lTemplateRecordForResponse; +use database_queries::queries::w2l::w2l_templates::get_w2l_template::select_w2l_template_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetW2lTemplatePathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct GetW2lTemplateSuccessResponse { + pub success: bool, + pub template: W2lTemplateRecordForResponse, +} + +#[derive(Debug)] +pub enum GetW2lTemplateError { + ServerError, + NotFound, +} + +impl ResponseError for GetW2lTemplateError { + fn status_code(&self) -> StatusCode { + match *self { + GetW2lTemplateError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + GetW2lTemplateError::NotFound => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetW2lTemplateError::ServerError => "server error".to_string(), + GetW2lTemplateError::NotFound=> "not found".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetW2lTemplateError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_w2l_template_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetW2lTemplateError::ServerError + })?; + + let mut show_deleted_templates = false; + let mut is_moderator = false; + + if let Some(user_session) = maybe_user_session { + // NB: Moderators can see deleted templates. + // Original creators cannot see them (unless they're moderators!) + show_deleted_templates = user_session.can_delete_other_users_w2l_templates; + // Moderators get to see all the fields. + is_moderator = user_session.can_delete_other_users_w2l_results + || user_session.can_edit_other_users_w2l_templates; + } + + let template_query_result = select_w2l_template_by_token( + &path.token, + show_deleted_templates, + &server_state.mysql_pool + ).await; + + let mut template = match template_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(GetW2lTemplateError::ServerError); + } + Ok(None) => return Err(GetW2lTemplateError::NotFound), + Ok(Some(template)) => template, + }; + + if let Some(moderator_fields) = template.maybe_moderator_fields.as_ref() { + // NB: The moderator fields will always be present before removal + // We don't want non-mods seeing stuff made by banned users. + if moderator_fields.creator_is_banned && !is_moderator { + return Err(GetW2lTemplateError::NotFound); + } + } + + if !is_moderator { + template.maybe_moderator_fields = None; + } + + let response = GetW2lTemplateSuccessResponse { + success: true, + template, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetW2lTemplateError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_template_use_count.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_template_use_count.rs new file mode 100644 index 0000000000..e754d93702 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_template_use_count.rs @@ -0,0 +1,92 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use r2d2_redis::redis::{Commands, RedisError}; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetW2lTemplateUseCountPathInfo { + template_token: String, +} + +#[derive(Serialize)] +pub struct GetW2lTemplateUseCountSuccessResponse { + pub success: bool, + pub count: Option, +} + +#[derive(Debug, Display)] +pub enum GetW2lTemplateUseCountError { + ServerError, + NotFoundError, +} + +impl ResponseError for GetW2lTemplateUseCountError { + fn status_code(&self) -> StatusCode { + match *self { + GetW2lTemplateUseCountError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + GetW2lTemplateUseCountError::NotFoundError => StatusCode::NOT_FOUND, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetW2lTemplateUseCountError::ServerError => "server error".to_string(), + GetW2lTemplateUseCountError::NotFoundError => "not found".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn get_w2l_template_use_count_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + GetW2lTemplateUseCountError::ServerError + })?; + + let redis_count_key = RedisKeys::w2l_template_usage_count(&path.template_token); + + let count : Option = redis.get(&redis_count_key) + .map_err(|e| { + warn!("redis error: {:?}", e); + GetW2lTemplateUseCountError::ServerError + })?; + + let response = GetW2lTemplateUseCountSuccessResponse { + success: true, + count, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetW2lTemplateUseCountError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs new file mode 100644 index 0000000000..c279af0aa5 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/get_w2l_upload_template_job_status.rs @@ -0,0 +1,179 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::cookie::Cookie; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use derive_more::{Display, Error}; +use log::{info, warn, log}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetW2lUploadTemplateStatusPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct W2lUploadTemplateJobStatusForResponse { + pub job_token: String, + + /// Primary status from the database (a state machine). + pub status: String, + + /// Extra, temporary status from Redis. + /// This can denote inference progress, and the Python code can write to it. + pub maybe_extra_status_description: Option, + + /// If the job failed or permanently died, we might have a reason. + pub maybe_failure_reason: Option, + + pub attempt_count: u8, + pub maybe_template_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct GetW2lUploadTemplateStatusSuccessResponse { + pub success: bool, + pub state: W2lUploadTemplateJobStatusForResponse, +} + +#[derive(Debug, Display)] +pub enum GetW2lUploadTemplateStatusError { + ServerError, +} + +#[derive(Serialize)] +pub struct W2lUploadTemplateJobStatusRecord { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + pub maybe_template_token: Option, + + pub maybe_failure_reason: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl ResponseError for GetW2lUploadTemplateStatusError { + fn status_code(&self) -> StatusCode { + match *self { + GetW2lUploadTemplateStatusError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetW2lUploadTemplateStatusError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +pub async fn get_w2l_upload_template_job_status_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_status = sqlx::query_as!( + W2lUploadTemplateJobStatusRecord, + r#" +SELECT + jobs.token as job_token, + + jobs.status, + jobs.attempt_count, + jobs.on_success_result_token as maybe_template_token, + + jobs.failure_reason as maybe_failure_reason, + + jobs.created_at, + jobs.updated_at + +FROM w2l_template_upload_jobs as jobs + +WHERE jobs.token = ? + "#, + &path.token + ) + .fetch_one(&server_state.mysql_pool) + .await; // TODO: This will return error if it doesn't exist + + let record : W2lUploadTemplateJobStatusRecord = match maybe_status { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(GetW2lUploadTemplateStatusError::ServerError); + }, + _ => { + warn!("w2l template query error: {:?}", err); + return Err(GetW2lUploadTemplateStatusError::ServerError); + } + } + } + }; + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + GetW2lUploadTemplateStatusError::ServerError + })?; + + let extra_status_key = RedisKeys::w2l_download_extra_status_info(&path.token); + let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { + Ok(Some(status)) => { + Some(status) + }, + Ok(None) => None, + Err(e) => { + warn!("redis error: {:?}", e); + None // Fail open + }, + }; + + let template_for_response = W2lUploadTemplateJobStatusForResponse { + job_token: record.job_token.clone(), + status: record.status.clone(), + maybe_extra_status_description, + maybe_failure_reason: record.maybe_failure_reason.clone(), + attempt_count: record.attempt_count as u8, + maybe_template_token: record.maybe_template_token.clone(), + created_at: record.created_at.clone(), + updated_at: record.updated_at.clone(), + }; + + let response = GetW2lUploadTemplateStatusSuccessResponse { + success: true, + state: template_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetW2lUploadTemplateStatusError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/list_w2l_templates.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/list_w2l_templates.rs new file mode 100644 index 0000000000..4c7a0e0236 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/list_w2l_templates.rs @@ -0,0 +1,89 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::w2l::w2l_templates::list_w2l_templates::{W2lTemplateRecordForList, list_w2l_templates}; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct ListW2lTemplatesSuccessResponse { + pub success: bool, + pub templates: Vec, +} + +#[derive(Debug)] +pub enum ListW2lTemplatesError { + ServerError, +} + +impl ResponseError for ListW2lTemplatesError { + fn status_code(&self) -> StatusCode { + match *self { + ListW2lTemplatesError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListW2lTemplatesError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListW2lTemplatesError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_w2l_templates_handler( + http_request: HttpRequest, + server_state: web::Data>) -> Result +{ + let no_scope_creator = None; + + // TODO(bt): Ideally we show users their own W2L templates here before they're mod approved. + + let query_results = list_w2l_templates( + &server_state.mysql_pool, + no_scope_creator, + true, + ).await; + + let templates = match query_results { + Ok(results) => results, + Err(e) => { + warn!("w2l template list query error: {:?}", e); + return Err(ListW2lTemplatesError::ServerError); + } + }; + + let response = ListW2lTemplatesSuccessResponse { + success: true, + templates, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| ListW2lTemplatesError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/mod.rs new file mode 100644 index 0000000000..0f55ac2244 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/mod.rs @@ -0,0 +1,14 @@ +pub mod delete_w2l_result; +pub mod delete_w2l_template; +pub mod edit_w2l_result; +pub mod edit_w2l_template; +pub mod enqueue_infer_w2l; +pub mod enqueue_infer_w2l_with_uploads; +pub mod enqueue_upload_w2l_template; +pub mod get_w2l_inference_job_status; +pub mod get_w2l_result; +pub mod get_w2l_template; +pub mod get_w2l_template_use_count; +pub mod get_w2l_upload_template_job_status; +pub mod list_w2l_templates; +pub mod set_w2l_template_mod_approval; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/w2l/set_w2l_template_mod_approval.rs b/crates/service/storyteller_web/src/http_server/endpoints/w2l/set_w2l_template_mod_approval.rs new file mode 100644 index 0000000000..b6aab394a4 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/w2l/set_w2l_template_mod_approval.rs @@ -0,0 +1,157 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use database_queries::queries::w2l::w2l_templates::get_w2l_template::select_w2l_template_by_token; +use log::{info, warn, log}; +use regex::Regex; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct SetW2lTemplateModApprovalPathInfo { + token: String, +} + +#[derive(Deserialize)] +pub struct SetW2lTemplateModApprovalRequest { + is_approved: bool, +} + +#[derive(Serialize)] +pub struct SetW2lTemplateModApprovalSuccessResponse { + pub success: bool, +} + +#[derive(Debug)] +pub enum SetW2lTemplateModApprovalError { + BadInput(String), + NotAuthorized, + NotFound, + ServerError, +} + +impl ResponseError for SetW2lTemplateModApprovalError { + fn status_code(&self) -> StatusCode { + match *self { + SetW2lTemplateModApprovalError::BadInput(_) => StatusCode::BAD_REQUEST, + SetW2lTemplateModApprovalError::NotAuthorized => StatusCode::UNAUTHORIZED, + SetW2lTemplateModApprovalError::NotFound => StatusCode::NOT_FOUND, + SetW2lTemplateModApprovalError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + SetW2lTemplateModApprovalError::BadInput(reason) => reason.to_string(), + SetW2lTemplateModApprovalError::NotAuthorized => "unauthorized".to_string(), + SetW2lTemplateModApprovalError::NotFound => "not found".to_string(), + SetW2lTemplateModApprovalError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for SetW2lTemplateModApprovalError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn set_w2l_template_mod_approval_handler( + http_request: HttpRequest, + path: Path, + request: web::Json, + server_state: web::Data>) -> Result +{ + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + SetW2lTemplateModApprovalError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(SetW2lTemplateModApprovalError::NotAuthorized); + } + }; + + if !user_session.can_approve_w2l_templates { + warn!("user is not allowed to approve templates: {}", user_session.user_token); + return Err(SetW2lTemplateModApprovalError::NotAuthorized); + } + + let template_query_result = select_w2l_template_by_token( + &path.token, + true, // Only mods can perform this action + &server_state.mysql_pool, + ).await; + + let w2l_template = match template_query_result { + Err(e) => { + warn!("query error: {:?}", e); + return Err(SetW2lTemplateModApprovalError::ServerError); + } + Ok(None) => return Err(SetW2lTemplateModApprovalError::NotFound), + Ok(Some(template)) => template, + }; + + let ip_address = get_request_ip(&http_request); + + let query_result = sqlx::query!( + r#" +UPDATE w2l_templates +SET + is_public_listing_approved = ?, + maybe_mod_user_token = ? +WHERE + token = ? +LIMIT 1 + "#, + + &request.is_approved, + &user_session.user_token, + &path.token, + ) + .execute(&server_state.mysql_pool) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + warn!("Update w2l mod approval status DB error: {:?}", err); + return Err(SetW2lTemplateModApprovalError::ServerError); + } + }; + let response = SetW2lTemplateModApprovalSuccessResponse { + success: true, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| SetW2lTemplateModApprovalError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/middleware/ip_filter_middleware.rs b/crates/service/storyteller_web/src/http_server/middleware/ip_filter_middleware.rs new file mode 100644 index 0000000000..68f0a5fcad --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/middleware/ip_filter_middleware.rs @@ -0,0 +1,110 @@ +// There are two steps in middleware processing. +// 1. Middleware initialization, middleware factory gets called with +// next service in chain as parameter. +// 2. Middleware's call method gets called with normal request. + +use actix_http::StatusCode; +use actix_http::http::{header, HeaderMap, HeaderValue}; +use actix_http::{error, body::Body, Response}; +use actix_web::dev::{Service, Transform}; +use actix_web::dev::{ServiceRequest, ServiceResponse}; +use actix_web::web::{BytesMut, Buf, BufMut}; +use actix_web::{Error, HttpResponse}; +use actix_web::{ResponseError, HttpMessage, HttpRequest, HttpResponseBuilder}; +use crate::http_server::web_utils::ip_address::get_service_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::threads::ip_banlist_set::IpBanlistSet; +use futures_util::future::{err, ok, Either, Ready}; +use std::io::Write; +use std::task::{Context, Poll}; + +#[derive(Debug)] +pub struct BannedError {} + +impl std::error::Error for BannedError {} + +impl std::fmt::Display for BannedError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "BannedError") + } +} + +impl ResponseError for BannedError { + fn status_code(&self) -> StatusCode { + StatusCode::UNAUTHORIZED + } + + fn error_response(&self) -> HttpResponse { + // NB: I'm setting a string error because I mistakenly got caught by this in local dev + // and couldn't figure out the issue for a bit. At least I can grep for this string. + // However, I need to balance this requirement with not cluing in those that are banned. + to_simple_json_error( + "ERR64: storyteller-web database error", + self.status_code()) + } +} + +#[derive(Clone)] // NB: Clone is safe because IpBanlist is clone-safe (internal Arc) +pub struct IpFilter { + ip_banlist: IpBanlistSet, +} + +impl IpFilter { + pub fn new(ip_banlist: IpBanlistSet) -> Self { + Self { + ip_banlist, + } + } +} + +impl Transform for IpFilter + where + S: Service, + S::Future: 'static, +{ + type Response = ServiceResponse; + type Error = Error; + type InitError = (); + type Transform = IpFilterMiddleware; + type Future = Ready>; + + fn new_transform(&self, service: S) -> Self::Future { + // NB: IpBanlist is clone-safe due to internal Arc. + ok(IpFilterMiddleware { service, ip_banlist: self.ip_banlist.clone() }) + } +} + +pub struct IpFilterMiddleware { + service: S, + ip_banlist: IpBanlistSet, +} + +impl Service for IpFilterMiddleware + where + S: Service, + S::Future: 'static, +{ + type Response = ServiceResponse; + type Error = Error; + type Future = Either>>; + + // alternatively(?), actix_service::forward_ready!(service); + fn poll_ready(&self, cx: &mut Context) -> Poll> { + self.service.poll_ready(cx) + } + + fn call(&self, req: ServiceRequest) -> Self::Future { + let ip_address = get_service_request_ip(&req); + + // NB: Fail open. + let is_banned = self.ip_banlist + .is_banned(&ip_address) + .unwrap_or(false); + + if is_banned { + return Either::Right(err(Error::from(BannedError {}))) + } + + Either::Left(self.service.call(req)) + } +} diff --git a/crates/service/storyteller_web/src/http_server/middleware/mod.rs b/crates/service/storyteller_web/src/http_server/middleware/mod.rs new file mode 100644 index 0000000000..a2369086a3 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/middleware/mod.rs @@ -0,0 +1 @@ +pub mod ip_filter_middleware; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/mod.rs b/crates/service/storyteller_web/src/http_server/mod.rs new file mode 100644 index 0000000000..43daf6085b --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/mod.rs @@ -0,0 +1,3 @@ +pub mod endpoints; +pub mod middleware; +pub mod web_utils; diff --git a/crates/service/storyteller_web/src/http_server/web_utils/cookie_manager.rs b/crates/service/storyteller_web/src/http_server/web_utils/cookie_manager.rs new file mode 100644 index 0000000000..c2de649aff --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/cookie_manager.rs @@ -0,0 +1,96 @@ +use actix_web::cookie::{Cookie, SameSite}; +use actix_web::{HttpRequest, HttpMessage}; +use anyhow::anyhow; +use crate::AnyhowResult; +use hmac::Hmac; +use hmac::NewMac; +use jwt::SignWithKey; +use jwt::VerifyWithKey; +use log::LevelFilter::Off; +use log::{info, warn}; +use sha2::Sha256; +use std::collections::BTreeMap; +use std::ops::Sub; +use time::OffsetDateTime; + +const COOKIE_VERSION : u32 = 1; + +const SESSION_COOKIE_NAME : &'static str = "session"; + +#[derive(Clone)] +pub struct CookieManager { + cookie_domain: String, + hmac_secret: String, +} + +impl CookieManager { + pub fn new(cookie_domain: &str, hmac_secret: &str) -> Self { + Self { + cookie_domain: cookie_domain.to_string(), + hmac_secret: hmac_secret.to_string(), + } + } + + pub fn create_cookie(&self, session_token: &str) -> AnyhowResult { + let key: Hmac = Hmac::new_varkey(self.hmac_secret.as_bytes()) + .map_err(|e| anyhow!("invalid hmac: {:?}", e))?; + + let cookie_version = COOKIE_VERSION.to_string(); + + let mut claims = BTreeMap::new(); + claims.insert("session_token", session_token); + claims.insert("cookie_version", &cookie_version); + + let jwt_string = claims.sign_with_key(&key)?; + + let make_secure = !self.cookie_domain.to_lowercase().contains("jungle.horse") + && !self.cookie_domain.to_lowercase().contains("localhost"); + + Ok(Cookie::build(SESSION_COOKIE_NAME, jwt_string) + //.domain(&self.cookie_domain) + .secure(make_secure) // HTTPS-only + //.path("/") + //.http_only(true) // Not exposed to Javascript + //.expires(OffsetDateTime::now_utc() + time::Duration::days(365)) + .permanent() + //.same_site(SameSite::Lax) + .finish()) + } + + pub fn delete_cookie(&self) -> Cookie { + Cookie::build(SESSION_COOKIE_NAME, "DELETED") + .expires(OffsetDateTime::unix_epoch()) + .finish() + } + + pub fn decode_session_token(&self, session_cookie: &Cookie) -> AnyhowResult { + let key: Hmac = Hmac::new_varkey(self.hmac_secret.as_bytes()) + .map_err(|e| anyhow!("invalid hmac: {:?}", e))?; + + let cookie_contents = session_cookie.value().to_string(); + + let claims: BTreeMap = cookie_contents.verify_with_key(&key)?; + + let session_token = claims["session_token"].clone(); + + Ok(session_token) + } + + pub fn decode_session_token_from_request(&self, request: &HttpRequest) + -> AnyhowResult> + { + let cookie = match request.cookie(SESSION_COOKIE_NAME) { + None => return Ok(None), + Some(cookie) => cookie, + }; + + match self.decode_session_token(&cookie) { + Err(e) => { + warn!("Session cookie decode error: {:?}", e); + Err(anyhow!("Could not decode session cookie: {:?}", e)) + }, + Ok(payload) => Ok(Some(payload)), + } + } +} + diff --git a/crates/service/storyteller_web/src/http_server/web_utils/get_host_header.rs b/crates/service/storyteller_web/src/http_server/web_utils/get_host_header.rs new file mode 100644 index 0000000000..57b4ec37e8 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/get_host_header.rs @@ -0,0 +1,18 @@ +use actix_http::http::HeaderMap; +use actix_http::http::HeaderName; +use actix_web::HttpRequest; +use actix_web::http::HeaderValue; +use container_common::anyhow_result::AnyhowResult; + +// NB: This is extracted because CLion is having incredible difficulty with these types. +pub fn get_host_header(http_request: &HttpRequest) -> AnyhowResult> { + let host_header_name = HeaderName::from_static("host"); + let header_map : &HeaderMap = http_request.headers(); + let maybe_hostname = header_map.get(host_header_name); + + let hostname = maybe_hostname.map(|s| s.to_str()) + .map(|inner| inner.map(|s| s.to_string())) + .transpose()?; + + Ok(hostname) +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/web_utils/ip_address.rs b/crates/service/storyteller_web/src/http_server/web_utils/ip_address.rs new file mode 100644 index 0000000000..c1b097410e --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/ip_address.rs @@ -0,0 +1,100 @@ +use actix_web::HttpRequest; +use actix_web::dev::ServiceRequest; +use log::{info}; +use actix_web::http::{HeaderMap, HeaderName}; +use std::str::FromStr; + +pub fn get_request_ip(request: &HttpRequest) -> String { + let headers = request.headers(); + let maybe_x_forwarded = get_ip_from_header(headers, "x-forwarded-for"); + let maybe_forwarded = get_ip_from_header(headers, "forwarded"); + + info!("(1) x-forwarded-for: {:?}, forwarded: {:?}", maybe_x_forwarded, maybe_forwarded); + + let maybe_ip = maybe_x_forwarded.or(maybe_forwarded); + + match maybe_ip { + Some(ip_address) => ip_address, + None => { + // If we're running without the upstream Rust proxy, we can grab 'x-forarded-for', which is + // populated by the DigitalOcean load balancer. + let ip_address_and_port = request.connection_info() + .remote_addr() + .unwrap_or("") + .to_string(); + let ip_address = ip_address_and_port.split(":") + .collect::>() + .get(0) + .copied() + .unwrap_or("") + .to_string(); + info!("Forwarded IP address (1): {}", &ip_address); + ip_address + }, + } +} + +// TODO: De-duplicate +pub fn get_service_request_ip(request: &ServiceRequest) -> String { + let headers = request.headers(); + let maybe_x_forwarded = get_ip_from_header(headers, "x-forwarded-for"); + let maybe_forwarded = get_ip_from_header(headers, "forwarded"); + + info!("(2) x-forwarded-for: {:?}, forwarded: {:?}", maybe_x_forwarded, maybe_forwarded); + + let maybe_ip = maybe_x_forwarded.or(maybe_forwarded); + + match maybe_ip { + Some(ip_address) => ip_address, + None => { + // If we're running without the upstream Rust proxy, we can grab 'x-forarded-for', which is + // populated by the DigitalOcean load balancer. + let ip_address_and_port = request.connection_info() + .remote_addr() + .unwrap_or("") + .to_string(); + let ip_address = ip_address_and_port.split(":") + .collect::>() + .get(0) + .copied() + .unwrap_or("") + .to_string(); + info!("(2) Forwarded IP address: {}", &ip_address); + ip_address + }, + } +} + +fn get_ip_from_header(headers: &HeaderMap, header_name: &str) -> Option { + if let Ok(header_name) = HeaderName::from_str(header_name) { + headers.get(&header_name) + .and_then(|value| value.to_str().ok()) + .and_then(|value| first_ip_from_list(value)) + } else { + None + } +} + +// NB: GKE packs in multiple IP addresses into the header: +// Example of header: x-forwarded-for: Some("136.55.189.34, 34.117.9.171") +fn first_ip_from_list(ip_list: &str) -> Option { + ip_list.split(",") + .into_iter() + .map(|ip| ip.trim()) + .find(|ip| !ip.is_empty()) + .map(|ip| ip.to_string()) +} + +#[cfg(test)] +mod tests { + use crate::http_server::web_utils::ip_address::first_ip_from_list; + + #[test] + fn test_first_ip() { + assert_eq!(first_ip_from_list("1.2.3.4"), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list("1.2.3.4, 34.117.9.171"), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list(", 1.2.3.4"), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list("1.2.3.4, "), Some("1.2.3.4".to_string())); + assert_eq!(first_ip_from_list(""), None); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/web_utils/mod.rs b/crates/service/storyteller_web/src/http_server/web_utils/mod.rs new file mode 100644 index 0000000000..8037e1b8c8 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/mod.rs @@ -0,0 +1,9 @@ +pub mod cookie_manager; +pub mod get_host_header; +pub mod ip_address; +pub mod read_multipart_field_bytes; +pub mod redis_rate_limiter; +pub mod response_error_helpers; +pub mod response_success_helpers; +pub mod serialize_as_json_error; +pub mod session_checker; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs b/crates/service/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs new file mode 100644 index 0000000000..57fa2fd81f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs @@ -0,0 +1,66 @@ +//use bytes::BytesMut; +use actix_multipart::Field; +use actix_web::web::BytesMut; +use anyhow::anyhow; +use futures::{StreamExt, TryStreamExt}; +use log::warn; + +const MIN_BYTES : usize = 10; +const MAX_BYTES : usize = 1024 * 1024 * 20; + +/// Read any field from a multipart form as bytes, but do some error checking. +pub async fn checked_read_multipart_bytes(field: &mut Field) -> anyhow::Result> { + let bytes = read_multipart_field_bytes(field).await?; + + if bytes.is_empty() { + // This field is empty + return Ok(None); + } + + if bytes.len() < MIN_BYTES { + warn!("Bytes of upload too few: {}", bytes.len()); + return Err(anyhow!("Too few bytes")); + } + + if bytes.len() > MAX_BYTES { + warn!("Bytes of upload too much: {}", bytes.len()); + return Err(anyhow!("Too many bytes")); + } + + Ok(Some(bytes)) +} + +/// Read any field from a multipart form as bytes. +pub async fn read_multipart_field_bytes(field: &mut Field) -> anyhow::Result { + let mut bytes = BytesMut::new(); + while let Some(chunk) = field.next().await { + let chunk_data = chunk.map_err(|err| anyhow!("Error: {:?}", err))?; + bytes.extend_from_slice(&chunk_data); + } + Ok(bytes) +} + +/// This reads inputs, text boxes, etc. because form-multipart is hard. +pub async fn read_multipart_field_as_text(field: &mut Field) -> anyhow::Result> { + let bytes = read_multipart_field_bytes(field).await?; + let value = String::from_utf8(bytes.to_vec())?; + + if value.is_empty() || value == "none" { + return Ok(None); + } + + Ok(Some(value)) +} + +/// Read a html form boolean (eg. input type=checkbox); assumes default is false. +pub async fn read_multipart_field_as_boolean(field: &mut Field) -> anyhow::Result { + let result = read_multipart_field_as_text(field).await?; + Ok(match result.as_ref().map(String::as_str) { + Some("true") => true, + Some("TRUE") => true, + Some("false") => false, + Some("FALSE") => false, + None => false, + _ => true, + }) +} diff --git a/crates/service/storyteller_web/src/http_server/web_utils/redis_rate_limiter.rs b/crates/service/storyteller_web/src/http_server/web_utils/redis_rate_limiter.rs new file mode 100644 index 0000000000..1a7004ec42 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/redis_rate_limiter.rs @@ -0,0 +1,73 @@ +//! NB: This doesn't live in the `util.redis` package because it depends on `http_server`. +//! It's only used for the server anyway, so living here makes sense. + +use actix_http::http::{HeaderName, HeaderMap}; +use actix_web::HttpRequest; +use crate::http_server::web_utils::ip_address::get_request_ip; +use futures_old_for_limiter::Future; +use limitation::{Limiter, Error}; +use log::info; + +/// If this HTTP header is set, the rate limiter can be bypassed (eg. for debugging) +const RATE_LIMIT_BYPASS_HEADER : &'static str = "limitless"; + +#[derive(Clone)] // NB: Limiter is `Clone` +pub struct RedisRateLimiter { + limiter: Limiter, + limit_key_prefix: String, + limiter_enabled: bool, + rate_limit_bypass_header: HeaderName, +} + +pub enum RateLimiterError { + RateLimitExceededError, + ClientError, // NB: Fail open for this! +} + +impl RedisRateLimiter { + pub fn new(limiter: Limiter, limit_key_prefix: &str, limiter_enabled: bool) -> Self { + let rate_limit_bypass_header = HeaderName::from_static(RATE_LIMIT_BYPASS_HEADER); + RedisRateLimiter { + limiter, + limit_key_prefix: limit_key_prefix.to_string(), + limiter_enabled, + rate_limit_bypass_header, + } + } + + pub fn rate_limit_request(&self, request: &HttpRequest) -> Result<(), RateLimiterError> { + if !self.limiter_enabled { + return Ok(()) + } + + let headers = request.headers(); + if headers.contains_key(&self.rate_limit_bypass_header) { + info!("Bypassing rate limiter with special bypass/debug header."); + return Ok(()); + } + + let ip_address = get_request_ip(&request); + + let rate_limit_key = format!("rate_limit:{}:{}", &self.limit_key_prefix, ip_address); + self.rate_limit_key(&rate_limit_key) + } + + pub fn rate_limit_key(&self, rate_limit_key: &str) -> Result<(), RateLimiterError> { + if !self.limiter_enabled { + return Ok(()) + } + + // NB: This library uses old-school futures (pre-async/await) + let permit = self.limiter.count(rate_limit_key); + + // TODO/FIXME: This could block. + match permit.wait() { + Ok(_) => Ok(()), + Err(err) => match err { + Error::Client(_) => Ok(()), // NB: Fail open for failure to connect to Redis + Error::Time(_) => Ok(()), // NB: Fail open for key parsing + Error::LimitExceeded(_) => Err(RateLimiterError::RateLimitExceededError), + }, + } + } +} diff --git a/crates/service/storyteller_web/src/http_server/web_utils/response_error_helpers.rs b/crates/service/storyteller_web/src/http_server/web_utils/response_error_helpers.rs new file mode 100644 index 0000000000..e950731d73 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/response_error_helpers.rs @@ -0,0 +1,30 @@ +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::http::StatusCode; +use actix_web::{HttpResponse, ResponseError}; + +#[derive(Serialize)] +pub struct SimpleGenericJsonError { + pub success: bool, + pub error_reason: String, +} + +pub fn to_simple_json_error( + error_reason: &str, + status_code: StatusCode, +) -> HttpResponse { + + let response = SimpleGenericJsonError { + success: false, + error_reason : error_reason.to_string(), + }; + + let body = match serde_json::to_string(&response) { + Ok(json) => json, + Err(_) => "{}".to_string(), + }; + + HttpResponseBuilder::new(status_code) + .set_header(header::CONTENT_TYPE, "application/json") + .body(body) +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/web_utils/response_success_helpers.rs b/crates/service/storyteller_web/src/http_server/web_utils/response_success_helpers.rs new file mode 100644 index 0000000000..dae959059f --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/response_success_helpers.rs @@ -0,0 +1,19 @@ +use actix_web::HttpResponse; + +#[derive(Serialize)] +pub struct SimpleGenericJsonSuccess { + pub success: bool, +} + +pub fn simple_json_success() -> HttpResponse { + let response = SimpleGenericJsonSuccess { + success: true, + }; + + let body = serde_json::to_string(&response) + .unwrap_or("{ \"success\": true }".to_string()); + + HttpResponse::Ok() + .content_type("application/json") + .body(body) +} diff --git a/crates/service/storyteller_web/src/http_server/web_utils/serialize_as_json_error.rs b/crates/service/storyteller_web/src/http_server/web_utils/serialize_as_json_error.rs new file mode 100644 index 0000000000..d288c9205e --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/serialize_as_json_error.rs @@ -0,0 +1,21 @@ +use actix_http::http::{StatusCode, header}; +use actix_web::HttpResponse; +use actix_web::HttpResponseBuilder; +use actix_web::error::ResponseError; +use serde::Serialize; + +/// Turn error responses into JSON HTTP responses +pub fn serialize_as_json_error( + error_payload: &T, +) -> HttpResponse + where T: ?Sized + Serialize + ResponseError, +{ + let body = match serde_json::to_string(error_payload) { + Ok(json) => json, + Err(_) => "{}".to_string(), + }; + + HttpResponseBuilder::new(error_payload.status_code()) + .set_header(header::CONTENT_TYPE, "application/json") + .body(body) +} diff --git a/crates/service/storyteller_web/src/http_server/web_utils/session_checker.rs b/crates/service/storyteller_web/src/http_server/web_utils/session_checker.rs new file mode 100644 index 0000000000..19f64d8de1 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/web_utils/session_checker.rs @@ -0,0 +1,304 @@ +use actix_web::HttpRequest; +use anyhow::anyhow; +use crate::AnyhowResult; +use crate::http_server::web_utils::cookie_manager::CookieManager; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::helpers::boolean_converters::{nullable_i8_to_optional_bool, i8_to_bool}; +use log::{info, warn}; +use sqlx::MySqlPool; +use sqlx::error::Error::RowNotFound; + +#[derive(Clone)] +pub struct SessionChecker { + cookie_manager: CookieManager, +} + +pub struct SessionRecord { + pub session_token: String, + pub user_token: String, +} + +pub struct SessionUserRecord { + pub user_token: String, + pub username: String, + pub display_name: String, + + pub email_address: String, + pub email_confirmed: bool, + pub email_gravatar_hash: String, + + // ===== PREFERENCES ===== // + + pub disable_gravatar: bool, + pub auto_play_audio_preference: Option, + pub preferred_tts_result_visibility: RecordVisibility, + pub preferred_w2l_result_visibility: RecordVisibility, + pub auto_play_video_preference: Option, + + // ===== ROLE ===== // + + pub user_role_slug: String, + pub is_banned: bool, + + // ===== PERMISSIONS FLAGS ===== // + + // Usage + pub can_use_tts: bool, + pub can_use_w2l: bool, + pub can_delete_own_tts_results: bool, + pub can_delete_own_w2l_results: bool, + pub can_delete_own_account: bool, + + // Contribution + pub can_upload_tts_models: bool, + pub can_upload_w2l_templates: bool, + pub can_delete_own_tts_models: bool, + pub can_delete_own_w2l_templates: bool, + + // Moderation + pub can_approve_w2l_templates: bool, + pub can_edit_other_users_profiles: bool, + pub can_edit_other_users_tts_models: bool, + pub can_edit_other_users_w2l_templates: bool, + pub can_delete_other_users_tts_models: bool, + pub can_delete_other_users_tts_results: bool, + pub can_delete_other_users_w2l_templates: bool, + pub can_delete_other_users_w2l_results: bool, + pub can_ban_users: bool, + pub can_delete_users: bool, +} + +pub struct SessionUserRawDbRecord { + pub user_token: String, + pub username: String, + pub display_name: String, + + pub email_address: String, + pub email_confirmed: i8, + pub email_gravatar_hash: String, + + pub disable_gravatar: i8, + pub auto_play_audio_preference: Option, + pub auto_play_video_preference: Option, + pub preferred_tts_result_visibility: RecordVisibility, + pub preferred_w2l_result_visibility: RecordVisibility, + + pub user_role_slug: String, + pub is_banned: i8, + + // NB: These are `Option` due to the JOIN not being compile-time assured. + // Usage + pub can_use_tts: Option, + pub can_use_w2l: Option, + pub can_delete_own_tts_results: Option, + pub can_delete_own_w2l_results: Option, + pub can_delete_own_account: Option, + + // Contribution + pub can_upload_tts_models: Option, + pub can_upload_w2l_templates: Option, + pub can_delete_own_tts_models: Option, + pub can_delete_own_w2l_templates: Option, + + // Moderation + pub can_approve_w2l_templates: Option, + pub can_edit_other_users_profiles: Option, + pub can_edit_other_users_tts_models: Option, + pub can_edit_other_users_w2l_templates: Option, + pub can_delete_other_users_tts_models: Option, + pub can_delete_other_users_tts_results: Option, + pub can_delete_other_users_w2l_templates: Option, + pub can_delete_other_users_w2l_results: Option, + pub can_ban_users: Option, + pub can_delete_users: Option, +} + +impl SessionChecker { + + pub fn new(cookie_manager: &CookieManager) -> Self { + Self { + cookie_manager: cookie_manager.clone(), + } + } + + pub async fn maybe_get_session(&self, request: &HttpRequest, pool: &MySqlPool) + -> AnyhowResult> + { + let session_token = match self.cookie_manager.decode_session_token_from_request(request)? { + None => return Ok(None), + Some(session_token) => session_token, + }; + + // NB: Lookup failure is Err(RowNotFound). + let maybe_session_record = sqlx::query_as!( + SessionRecord, + r#" +SELECT + token as session_token, + user_token +FROM user_sessions +WHERE token = ? +AND deleted_at IS NULL + "#, + session_token.to_string(), + ) + .fetch_one(pool) + .await; // TODO: This will return error if it doesn't exist + + match maybe_session_record { + Ok(session_record) => { + Ok(Some(session_record)) + }, + Err(err) => { + match err { + RowNotFound => { + warn!("Valid cookie; invalid session: {}", session_token); + Ok(None) + }, + _ => { + warn!("Session query error: {:?}", err); + Err(anyhow!("session query error: {:?}", err)) + } + } + } + } + } + + pub async fn maybe_get_user_session(&self, request: &HttpRequest, pool: &MySqlPool) + -> AnyhowResult> + { + + let session_token = match self.cookie_manager.decode_session_token_from_request(request)? { + None => return Ok(None), + Some(session_token) => session_token, + }; + + // NB: Lookup failure is Err(RowNotFound). + let maybe_user_record = sqlx::query_as!( + SessionUserRawDbRecord, + r#" +SELECT + users.token as user_token, + users.username, + users.display_name, + + users.email_address, + users.email_confirmed, + users.email_gravatar_hash, + + users.disable_gravatar, + users.auto_play_audio_preference, + users.auto_play_video_preference, + users.preferred_tts_result_visibility as `preferred_tts_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`, + users.preferred_w2l_result_visibility as `preferred_w2l_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`, + + users.user_role_slug, + users.is_banned, + + user_roles.can_use_tts, + user_roles.can_use_w2l, + user_roles.can_delete_own_tts_results, + user_roles.can_delete_own_w2l_results, + user_roles.can_delete_own_account, + + user_roles.can_upload_tts_models, + user_roles.can_upload_w2l_templates, + user_roles.can_delete_own_tts_models, + user_roles.can_delete_own_w2l_templates, + + user_roles.can_approve_w2l_templates, + user_roles.can_edit_other_users_profiles, + user_roles.can_edit_other_users_tts_models, + user_roles.can_edit_other_users_w2l_templates, + user_roles.can_delete_other_users_tts_models, + user_roles.can_delete_other_users_tts_results, + user_roles.can_delete_other_users_w2l_templates, + user_roles.can_delete_other_users_w2l_results, + user_roles.can_ban_users, + user_roles.can_delete_users + +FROM users +LEFT OUTER JOIN user_sessions + ON users.token = user_sessions.user_token +LEFT OUTER JOIN user_roles + ON users.user_role_slug = user_roles.slug +WHERE user_sessions.token = ? + AND user_sessions.deleted_at IS NULL + AND users.user_deleted_at IS NULL + AND users.mod_deleted_at IS NULL + "#, + session_token.to_string(), + ) + .fetch_one(pool) + .await; // TODO: This will return error if it doesn't exist + + match maybe_user_record { + Ok(raw_user_record) => { + let result_user_record = SessionUserRecord { + user_token: raw_user_record.user_token, + username: raw_user_record.username, + display_name: raw_user_record.display_name, + email_address: raw_user_record.email_address, + email_confirmed: i8_to_bool(raw_user_record.email_confirmed), + email_gravatar_hash: raw_user_record.email_gravatar_hash, + // Preference + disable_gravatar: i8_to_bool(raw_user_record.disable_gravatar), + auto_play_audio_preference: nullable_i8_to_optional_bool(raw_user_record.auto_play_audio_preference), + auto_play_video_preference: nullable_i8_to_optional_bool(raw_user_record.auto_play_video_preference), + user_role_slug: raw_user_record.user_role_slug, + preferred_tts_result_visibility: raw_user_record.preferred_tts_result_visibility, + preferred_w2l_result_visibility: raw_user_record.preferred_w2l_result_visibility, + + is_banned: i8_to_bool(raw_user_record.is_banned), + + // Usage + can_use_tts: convert_optional_db_bool_default_false(raw_user_record.can_use_tts), + can_use_w2l: convert_optional_db_bool_default_false(raw_user_record.can_use_w2l), + can_delete_own_tts_results: convert_optional_db_bool_default_false(raw_user_record.can_delete_own_tts_results), + can_delete_own_w2l_results: convert_optional_db_bool_default_false(raw_user_record.can_delete_own_w2l_results), + can_delete_own_account: convert_optional_db_bool_default_false(raw_user_record.can_delete_own_account), + // Contribution + can_upload_tts_models: convert_optional_db_bool_default_false(raw_user_record.can_upload_tts_models), + can_upload_w2l_templates: convert_optional_db_bool_default_false(raw_user_record.can_upload_w2l_templates), + can_delete_own_tts_models: convert_optional_db_bool_default_false(raw_user_record.can_delete_own_tts_models), + can_delete_own_w2l_templates: convert_optional_db_bool_default_false(raw_user_record.can_delete_own_w2l_templates), + // Moderation + can_approve_w2l_templates: convert_optional_db_bool_default_false(raw_user_record.can_approve_w2l_templates), + can_edit_other_users_profiles: convert_optional_db_bool_default_false(raw_user_record.can_edit_other_users_profiles), + can_edit_other_users_tts_models: convert_optional_db_bool_default_false(raw_user_record.can_edit_other_users_tts_models), + can_edit_other_users_w2l_templates: convert_optional_db_bool_default_false(raw_user_record.can_edit_other_users_w2l_templates), + can_delete_other_users_tts_models: convert_optional_db_bool_default_false(raw_user_record.can_delete_other_users_tts_models), + can_delete_other_users_tts_results: convert_optional_db_bool_default_false(raw_user_record.can_delete_other_users_tts_results), + can_delete_other_users_w2l_templates: convert_optional_db_bool_default_false(raw_user_record.can_delete_other_users_w2l_templates), + can_delete_other_users_w2l_results: convert_optional_db_bool_default_false(raw_user_record.can_delete_other_users_w2l_results), + can_ban_users: convert_optional_db_bool_default_false(raw_user_record.can_ban_users), + can_delete_users: convert_optional_db_bool_default_false(raw_user_record.can_delete_users) , + }; + + Ok(Some(result_user_record)) + }, + Err(err) => { + match err { + RowNotFound => { + warn!("Valid cookie; invalid session: {}", session_token); + Ok(None) + }, + _ => { + warn!("Session query error: {:?}", err); + Err(anyhow!("session query error: {:?}", err)) + } + } + } + } + } +} + +fn convert_optional_db_bool_default_false(value: Option) -> bool { + match value { + None => false, + Some(v) => match v { + 0 => false, + _ => true, + } + } +} diff --git a/crates/service/storyteller_web/src/job_queries/mod.rs b/crates/service/storyteller_web/src/job_queries/mod.rs new file mode 100644 index 0000000000..c25ac9cc81 --- /dev/null +++ b/crates/service/storyteller_web/src/job_queries/mod.rs @@ -0,0 +1,6 @@ +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + +pub mod tts_download_job_queries; +pub mod w2l_download_job_queries; +pub mod w2l_inference_job_queries; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs b/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs new file mode 100644 index 0000000000..3108000ded --- /dev/null +++ b/crates/service/storyteller_web/src/job_queries/tts_download_job_queries.rs @@ -0,0 +1,259 @@ +//! NB: This seems required for sqlx to generate the cached queries. +//! Sqlx's prepare needs a *single* binary to work against, so we need to +//! include these in the main binary to generate all the queries. + +use anyhow::anyhow; +use chrono::Utc; +use container_common::anyhow_result::AnyhowResult; +use database_queries::tokens::Tokens; +use sqlx::MySqlPool; +use std::path::Path; + +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + +/// table: tts_model_upload_jobs +#[derive(Debug)] +pub struct TtsUploadJobRecord { + pub id: i64, + pub token: String, + pub uuid_idempotency_token: String, + pub on_success_result_token: Option, + pub creator_user_token: String, + pub creator_ip_address: String, + pub creator_set_visibility: String, // TODO + pub title: String, + pub tts_model_type: String, // TODO + pub download_url: Option, + pub status: String, // TODO + pub attempt_count: i32, + pub failure_reason: Option, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + +pub async fn query_tts_upload_job_records(pool: &MySqlPool, num_records: u32) + -> AnyhowResult> +{ + let job_records = sqlx::query_as!( + TtsUploadJobRecord, + r#" +SELECT * +FROM tts_model_upload_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + LIMIT ? + "#, + num_records, + ) + .fetch_all(pool) + .await?; + + Ok(job_records) +} + +pub struct TtsUploadLockRecord { + id: i64, + status: String, +} + +pub async fn grab_job_lock_and_mark_pending( + pool: &MySqlPool, + job: &TtsUploadJobRecord +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + TtsUploadLockRecord, + r#" +SELECT + id, + status +FROM tts_model_upload_jobs +WHERE id = ? +FOR UPDATE + "#, + job.id, + ) + .fetch_one(&mut transaction) + .await; + + let record : TtsUploadLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE tts_model_upload_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job.id, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} + +pub async fn mark_tts_upload_job_failure( + pool: &MySqlPool, + job: &TtsUploadJobRecord, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE tts_model_upload_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn mark_tts_upload_job_done( + pool: &MySqlPool, + job: &TtsUploadJobRecord, + success: bool, + maybe_model_token: Option<&str>, +) -> AnyhowResult<()> +{ + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE tts_model_upload_jobs +SET + status = ?, + on_success_result_token = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_model_token, + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn insert_tts_model>( + pool: &MySqlPool, + job: &TtsUploadJobRecord, + private_bucket_hash: &str, + private_bucket_object_name: P, + file_size_bytes: u64 +) -> AnyhowResult<(u64, String)> { + + let model_token = Tokens::new_tts_model()?; + + let private_bucket_object_name = &private_bucket_object_name + .as_ref() + .display() + .to_string(); + + let query_result = sqlx::query!( + r#" +INSERT INTO tts_models +SET + token = ?, + tts_model_type = "tacotron2", + title = ?, + description_markdown = '', + description_rendered_html = '', + creator_user_token = ?, + creator_ip_address_creation = ?, + creator_ip_address_last_update = ?, + original_download_url = ?, + private_bucket_hash = ?, + private_bucket_object_name = ?, + file_size_bytes = ? + "#, + &model_token, + job.title.to_string(), + job.creator_user_token.clone(), + job.creator_ip_address.clone(), + job.creator_ip_address.clone(), + job.download_url.clone(), + private_bucket_hash.to_string(), + private_bucket_object_name.to_string(), + file_size_bytes + ) + .execute(pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + Ok((record_id, model_token)) +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs b/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs new file mode 100644 index 0000000000..fd8c8347ad --- /dev/null +++ b/crates/service/storyteller_web/src/job_queries/w2l_download_job_queries.rs @@ -0,0 +1,305 @@ +//! NB: This seems required for sqlx to generate the cached queries. +//! Sqlx's prepare needs a *single* binary to work against, so we need to +//! include these in the main binary to generate all the queries. + +use anyhow::anyhow; +use chrono::Utc; +use container_common::anyhow_result::AnyhowResult; +use database_queries::tokens::Tokens; +use sqlx::MySqlPool; + +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + +/// table: w2l_template_upload_jobs +#[derive(Debug)] +pub struct W2lTemplateUploadJobRecord { + pub id: i64, + pub token: String, + pub uuid_idempotency_token: String, + pub on_success_result_token: Option, + pub creator_user_token: String, + pub creator_ip_address: String, + pub creator_set_visibility: String, // TODO + pub title: String, + pub template_type: String, // TODO + pub download_url: Option, + pub status: String, // TODO + pub attempt_count: i32, + pub failure_reason: Option, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + +pub async fn query_w2l_template_upload_job_records(pool: &MySqlPool, num_records: u32) + -> AnyhowResult> +{ + let job_records = sqlx::query_as!( + W2lTemplateUploadJobRecord, + r#" +SELECT * +FROM w2l_template_upload_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + LIMIT ? + "#, + num_records, + ) + .fetch_all(pool) + .await?; + + Ok(job_records) +} + +pub struct W2lDownloadLockRecord { + id: i64, + status: String, +} + +pub async fn grab_job_lock_and_mark_pending( + pool: &MySqlPool, + job: &W2lTemplateUploadJobRecord +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + W2lDownloadLockRecord, + r#" +SELECT + id, + status +FROM w2l_template_upload_jobs +WHERE id = ? +FOR UPDATE + "#, + job.id, + ) + .fetch_one(&mut transaction) + .await; + + let record : W2lDownloadLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE w2l_template_upload_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job.id, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} + +pub async fn mark_w2l_template_upload_job_failure( + pool: &MySqlPool, + job: &W2lTemplateUploadJobRecord, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE w2l_template_upload_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn mark_w2l_template_upload_job_permanently_dead( + pool: &MySqlPool, + job: &W2lTemplateUploadJobRecord, + failure_reason: &str, +) -> AnyhowResult<()> { + let query_result = sqlx::query!( + r#" +UPDATE w2l_template_upload_jobs +SET + status = "dead", + failure_reason = ?, + retry_at = NULL +WHERE id = ? + "#, + failure_reason.to_string(), + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn mark_w2l_template_upload_job_done( + pool: &MySqlPool, + job: &W2lTemplateUploadJobRecord, + success: bool, + maybe_model_token: Option<&str> +) -> AnyhowResult<()> +{ + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE w2l_template_upload_jobs +SET + status = ?, + on_success_result_token = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_model_token, + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn insert_w2l_template( + pool: &MySqlPool, + template_type: &str, // TODO: ENUM! + job: &W2lTemplateUploadJobRecord, + private_bucket_hash: &str, + private_bucket_object_name: &str, + private_bucket_cached_faces_object_name: &str, + maybe_image_preview_object_name: Option<&str>, + maybe_video_preview_object_name: Option<&str>, + file_size_bytes: u64, + maybe_mime_type: Option<&str>, + frame_width: u32, + frame_height: u32, + frame_count: u64, + fps: f32, + duration_millis: u64 +) -> AnyhowResult<(u64, String)> { + + let model_token = Tokens::new_w2l_template()?; + + let query_result = sqlx::query!( + r#" +INSERT INTO w2l_templates +SET + token = ?, + template_type = ?, + title = ?, + description_markdown = '', + description_rendered_html = '', + creator_user_token = ?, + creator_ip_address_creation = ?, + creator_ip_address_last_update = ?, + original_download_url = ?, + private_bucket_hash = ?, + private_bucket_object_name = ?, + private_bucket_cached_faces_object_name = ?, + maybe_public_bucket_preview_image_object_name = ?, + maybe_public_bucket_preview_video_object_name = ?, + file_size_bytes = ?, + mime_type = ?, + frame_width = ?, + frame_height = ?, + frame_count = ?, + fps = ?, + duration_millis = ? + "#, + model_token, + template_type, + job.title.to_string(), + job.creator_user_token.clone(), + job.creator_ip_address.clone(), + job.creator_ip_address.clone(), + job.download_url.clone(), + private_bucket_hash.to_string(), + private_bucket_object_name.to_string(), + private_bucket_cached_faces_object_name.to_string(), + maybe_image_preview_object_name, + maybe_video_preview_object_name, + file_size_bytes, + maybe_mime_type.unwrap_or(""), + frame_width, + frame_height, + frame_count, + fps, + duration_millis + ) + .execute(pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + Ok((record_id, model_token.clone())) +} diff --git a/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs b/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs new file mode 100644 index 0000000000..0db4a6b258 --- /dev/null +++ b/crates/service/storyteller_web/src/job_queries/w2l_inference_job_queries.rs @@ -0,0 +1,455 @@ +//! NB: This seems required for sqlx to generate the cached queries. +//! Sqlx's prepare needs a *single* binary to work against, so we need to +//! include these in the main binary to generate all the queries. + +use anyhow::anyhow; +use chrono::{Utc, DateTime}; +use container_common::anyhow_result::AnyhowResult; +use database_queries::column_types::record_visibility::RecordVisibility; +use database_queries::tokens::Tokens; +use log::{warn, info}; +use sqlx::MySqlPool; +use std::path::Path; + +// TODO(2022-02-10): Move these queries to the 'database_queries' crate +// I moved tts_inference_job queries out already. + +/// table: w2l_inference_jobs +#[derive(Debug)] +pub struct W2lInferenceJobRecord { + pub id: i64, + pub inference_job_token: String, + pub uuid_idempotency_token: String, + + // ===== FACE TEMPLATE OPTIONS ===== + pub maybe_w2l_template_token: Option, + pub maybe_public_image_bucket_location: Option, + + // ===== AUDIO SOURCE OPTIONS ===== + pub maybe_tts_inference_result_token: Option, + pub maybe_public_audio_bucket_hash: Option, + pub maybe_public_audio_bucket_location: Option, + + pub maybe_original_audio_filename: Option, + pub maybe_original_audio_download_url: Option, + pub maybe_audio_mime_type: Option, + + pub creator_ip_address: String, + pub maybe_creator_user_token: Option, + + pub creator_set_visibility: RecordVisibility, + pub disable_end_bump: i8, // bool + pub disable_watermark: i8, // bool + + //pub maybe_subject_token: Option, + //pub maybe_actor_subject_token: Option, + pub status: String, // TODO + pub attempt_count: i32, + pub failure_reason: Option, + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + +pub async fn query_w2l_inference_job_records(pool: &MySqlPool, num_records: u32) + -> AnyhowResult> +{ + let job_records = sqlx::query_as!( + W2lInferenceJobRecord, + r#" +SELECT + id, + token AS inference_job_token, + uuid_idempotency_token, + + maybe_w2l_template_token, + maybe_public_image_bucket_location, + maybe_tts_inference_result_token, + maybe_public_audio_bucket_hash, + maybe_public_audio_bucket_location, + + maybe_original_audio_filename, + maybe_original_audio_download_url, + maybe_audio_mime_type, + + creator_ip_address, + maybe_creator_user_token, + + creator_set_visibility as `creator_set_visibility: database_queries::column_types::record_visibility::RecordVisibility`, + disable_end_bump, + disable_watermark, + + status, + attempt_count, + failure_reason, + created_at, + updated_at, + retry_at +FROM w2l_inference_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + LIMIT ? + "#, + num_records, + ) + .fetch_all(pool) + .await?; + + Ok(job_records) +} + +pub struct W2lInferenceLockRecord { + id: i64, + status: String, +} + +pub async fn grab_job_lock_and_mark_pending( + pool: &MySqlPool, + job: &W2lInferenceJobRecord +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + W2lInferenceLockRecord, + r#" +SELECT + id, + status +FROM w2l_inference_jobs +WHERE id = ? +FOR UPDATE + "#, + job.id, + ) + .fetch_one(&mut transaction) + .await; + + let record : W2lInferenceLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE w2l_inference_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job.id, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} + + +pub async fn mark_w2l_inference_job_failure( + pool: &MySqlPool, + job: &W2lInferenceJobRecord, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE w2l_inference_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub async fn mark_w2l_inference_job_done( + pool: &MySqlPool, + job: &W2lInferenceJobRecord, + success: bool, + maybe_result_token: Option<&str> +) -> AnyhowResult<()> { + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE w2l_inference_jobs +SET + status = ?, + on_success_result_token = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_result_token, + job.id, + ) + .execute(pool) + .await?; + + Ok(()) +} + +pub struct SyntheticIdRecord { + pub next_id: i64, +} + +pub async fn insert_w2l_result>( + pool: &MySqlPool, + job: &W2lInferenceJobRecord, + bucket_video_results_path: P, + file_size_bytes: u64, + maybe_mime_type: Option<&str>, + frame_width: u32, + frame_height: u32, + duration_millis: u64 +) -> AnyhowResult<(u64, String)> +{ + let inference_result_token = Tokens::new_w2l_result()?; + + let bucket_video_result_path = &bucket_video_results_path + .as_ref() + .display() + .to_string(); + + let maybe_creator_user_token = job.maybe_creator_user_token.clone(); + let mut maybe_creator_synthetic_id : Option = None; + + let mut transaction = pool.begin().await?; + + if let Some(creator_user_token) = maybe_creator_user_token.as_deref() { + let query_result = sqlx::query!( + r#" +INSERT INTO w2l_result_synthetic_ids +SET + user_token = ?, + next_id = 1 +ON DUPLICATE KEY UPDATE + user_token = ?, + next_id = next_id + 1 + "#, + creator_user_token, + creator_user_token + ) + .execute(&mut transaction) + .await; + + match query_result { + Ok(_) => {}, + Err(err) => { + //transaction.rollback().await?; + warn!("Transaction failure: {:?}", err); + } + } + + let query_result = sqlx::query_as!( + SyntheticIdRecord, + r#" +SELECT + next_id +FROM + w2l_result_synthetic_ids +WHERE + user_token = ? +LIMIT 1 + "#, + creator_user_token, + ) + .fetch_one(&mut transaction) + .await; + + let record : SyntheticIdRecord = match query_result { + Ok(record) => record, + Err(err) => { + warn!("Transaction failure: {:?}", err); + //transaction.rollback().await?; + return Err(anyhow!("Transaction failure: {:?}", err)); + } + }; + + let next_id = record.next_id as u64; + maybe_creator_synthetic_id = Some(next_id); + } + + let record_id = { + let query_result = sqlx::query!( + r#" +INSERT INTO w2l_results +SET + token = ?, + maybe_w2l_template_token = ?, + maybe_creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + + maybe_creator_synthetic_id = ?, + + public_bucket_video_path = ?, + + file_size_bytes = ?, + mime_type = ?, + frame_width = ?, + frame_height = ?, + duration_millis = ? + "#, + inference_result_token, + job.maybe_w2l_template_token.clone(), + job.maybe_creator_user_token.clone(), + job.creator_ip_address.clone(), + job.creator_set_visibility.to_str(), + + maybe_creator_synthetic_id, + + bucket_video_result_path, + + file_size_bytes, + maybe_mime_type.unwrap_or(""), + frame_width, + frame_height, + duration_millis + ) + .execute(&mut transaction) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + record_id + }; + + transaction.commit().await?; + + Ok((record_id, inference_result_token.clone())) +} + +pub struct W2lTemplateRecord2 { + pub template_token: String, + pub template_type: String, + pub creator_user_token: String, + pub creator_username: String, + pub private_bucket_hash: String, + pub creator_display_name: String, + pub title: String, + pub frame_width: i32, + pub frame_height: i32, + pub duration_millis: i32, + pub maybe_public_bucket_preview_image_object_name: Option, + pub maybe_public_bucket_preview_video_object_name: Option, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +pub async fn get_w2l_template_by_token(pool: &MySqlPool, template_token: &str) + -> AnyhowResult> +{ + + // NB: Lookup failure is Err(RowNotFound). + // NB: Since this is publicly exposed, we don't query sensitive data. + let maybe_template = sqlx::query_as!( + W2lTemplateRecord2, + r#" +SELECT + w2l.token as template_token, + w2l.template_type, + w2l.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + w2l.title, + w2l.frame_width, + w2l.frame_height, + w2l.duration_millis, + w2l.private_bucket_hash, + w2l.maybe_public_bucket_preview_image_object_name, + w2l.maybe_public_bucket_preview_video_object_name, + w2l.created_at, + w2l.updated_at +FROM w2l_templates as w2l +JOIN users +ON users.token = w2l.creator_user_token +WHERE w2l.token = ? +AND w2l.user_deleted_at IS NULL +AND w2l.mod_deleted_at IS NULL + "#, + &template_token + ) + .fetch_one(pool) + .await; // TODO: This will return error if it doesn't exist + + match maybe_template { + Ok(template) => Ok(Some(template)), + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + Ok(None) + }, + _ => { + warn!("w2l template query error: {:?}", err); + Err(anyhow!("Mysql error: {:?}", err)) + } + } + } + } +} diff --git a/crates/service/storyteller_web/src/main.rs b/crates/service/storyteller_web/src/main.rs new file mode 100644 index 0000000000..7b7c0fd2e2 --- /dev/null +++ b/crates/service/storyteller_web/src/main.rs @@ -0,0 +1,419 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate lazy_static; +#[macro_use] extern crate magic_crypt; +#[macro_use] extern crate serde_derive; + +pub const RESERVED_USERNAMES : &'static str = include_str!("../../../../db/reserved_usernames.txt"); +pub const RESERVED_SUBSTRINGS : &'static str = include_str!("../../../../db/reserved_usernames_including.txt"); + +pub mod common_env; +pub mod http_clients; +pub mod http_server; +pub mod routes; +pub mod server_state; +pub mod threads; +pub mod util; +pub mod validations; + +// NB: This is included so sqlx can generate all the queries. +mod job_queries; + +// NB: Also so sqlx codegens everything. +// Not sure if this is strictly necessary. +mod shared_queries { + use database_queries::queries::twitch::twitch_oauth::find; + use database_queries::queries::twitch::twitch_oauth::insert; +} + +use actix_cors::Cors; +use actix_http::http; +use actix_web::middleware::{Logger, DefaultHeaders}; +use actix_web::{HttpServer, web, HttpResponse, App}; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use crate::common_env::CommonEnv; +use crate::http_server::endpoints::events::list_events::list_events_handler; +use crate::http_server::endpoints::leaderboard::get_leaderboard::leaderboard_handler; +use crate::http_server::endpoints::misc::default_route_404::default_route_404; +use crate::http_server::endpoints::misc::enable_alpha_easy_handler::enable_alpha_easy_handler; +use crate::http_server::endpoints::misc::enable_alpha_handler::enable_alpha_handler; +use crate::http_server::endpoints::misc::root_index::get_root_index; +use crate::http_server::endpoints::moderation::approval::pending_w2l_templates::get_pending_w2l_templates_handler; +use crate::http_server::endpoints::moderation::ip_bans::add_ip_ban::add_ip_ban_handler; +use crate::http_server::endpoints::moderation::ip_bans::delete_ip_ban::delete_ip_ban_handler; +use crate::http_server::endpoints::moderation::ip_bans::get_ip_ban::get_ip_ban_handler; +use crate::http_server::endpoints::moderation::ip_bans::list_ip_bans::list_ip_bans_handler; +use crate::http_server::endpoints::moderation::jobs::get_tts_inference_queue_count::get_tts_inference_queue_count_handler; +use crate::http_server::endpoints::moderation::jobs::get_w2l_inference_queue_count::get_w2l_inference_queue_count_handler; +use crate::http_server::endpoints::moderation::stats::get_voice_count_stats::get_voice_count_stats_handler; +use crate::http_server::endpoints::moderation::user_bans::ban_user::ban_user_handler; +use crate::http_server::endpoints::moderation::user_bans::list_banned_users::list_banned_users_handler; +use crate::http_server::endpoints::moderation::user_roles::list_roles::list_user_roles_handler; +use crate::http_server::endpoints::moderation::user_roles::list_staff::list_staff_handler; +use crate::http_server::endpoints::moderation::user_roles::set_user_role::set_user_role_handler; +use crate::http_server::endpoints::moderation::users::list_users::list_users_handler; +use crate::http_server::endpoints::tts::delete_tts_model::delete_tts_model_handler; +use crate::http_server::endpoints::tts::delete_tts_result::delete_tts_inference_result_handler; +use crate::http_server::endpoints::tts::edit_tts_model::edit_tts_model_handler; +use crate::http_server::endpoints::tts::edit_tts_result::edit_tts_inference_result_handler; +use crate::http_server::endpoints::tts::enqueue_infer_tts::infer_tts_handler; +use crate::http_server::endpoints::tts::enqueue_upload_tts_model::upload_tts_model_handler; +use crate::http_server::endpoints::tts::get_tts_inference_job_status::get_tts_inference_job_status_handler; +use crate::http_server::endpoints::tts::get_tts_model::get_tts_model_handler; +use crate::http_server::endpoints::tts::get_tts_model_use_count::get_tts_model_use_count_handler; +use crate::http_server::endpoints::tts::get_tts_result::get_tts_inference_result_handler; +use crate::http_server::endpoints::tts::get_tts_upload_model_job_status::get_tts_upload_model_job_status_handler; +use crate::http_server::endpoints::tts::list_tts_models::list_tts_models_handler; +use crate::http_server::endpoints::users::create_account::create_account_handler; +use crate::http_server::endpoints::users::edit_profile::edit_profile_handler; +use crate::http_server::endpoints::users::get_profile::get_profile_handler; +use crate::http_server::endpoints::users::list_user_tts_inference_results::list_user_tts_inference_results_handler; +use crate::http_server::endpoints::users::list_user_tts_models::list_user_tts_models_handler; +use crate::http_server::endpoints::users::list_user_w2l_inference_results::list_user_w2l_inference_results_handler; +use crate::http_server::endpoints::users::list_user_w2l_templates::list_user_w2l_templates_handler; +use crate::http_server::endpoints::users::login::login_handler; +use crate::http_server::endpoints::users::logout::logout_handler; +use crate::http_server::endpoints::users::session_info::session_info_handler; +use crate::http_server::endpoints::w2l::delete_w2l_result::delete_w2l_inference_result_handler; +use crate::http_server::endpoints::w2l::delete_w2l_template::delete_w2l_template_handler; +use crate::http_server::endpoints::w2l::edit_w2l_result::edit_w2l_inference_result_handler; +use crate::http_server::endpoints::w2l::edit_w2l_template::edit_w2l_template_handler; +use crate::http_server::endpoints::w2l::enqueue_infer_w2l::infer_w2l_handler; +use crate::http_server::endpoints::w2l::enqueue_infer_w2l_with_uploads::enqueue_infer_w2l_with_uploads; +use crate::http_server::endpoints::w2l::enqueue_upload_w2l_template::upload_w2l_template_handler; +use crate::http_server::endpoints::w2l::get_w2l_inference_job_status::get_w2l_inference_job_status_handler; +use crate::http_server::endpoints::w2l::get_w2l_result::get_w2l_inference_result_handler; +use crate::http_server::endpoints::w2l::get_w2l_template::get_w2l_template_handler; +use crate::http_server::endpoints::w2l::get_w2l_template_use_count::get_w2l_template_use_count_handler; +use crate::http_server::endpoints::w2l::get_w2l_upload_template_job_status::get_w2l_upload_template_job_status_handler; +use crate::http_server::endpoints::w2l::list_w2l_templates::list_w2l_templates_handler; +use crate::http_server::endpoints::w2l::set_w2l_template_mod_approval::set_w2l_template_mod_approval_handler; +use crate::http_server::middleware::ip_filter_middleware::IpFilter; +use crate::http_server::web_utils::cookie_manager::CookieManager; +use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter; +use crate::http_server::web_utils::session_checker::SessionChecker; +use crate::routes::add_routes; +use crate::server_state::{ServerState, EnvConfig, TwitchOauthSecrets, TwitchOauth, RedisRateLimiters, InMemoryCaches}; +use crate::threads::ip_banlist_set::IpBanlistSet; +use crate::threads::poll_ip_banlist_thread::poll_ip_bans; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::caching::single_item_ttl_cache::SingleItemTtlCache; +use crate::util::encrypted_sort_id::SortKeyCrypto; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use futures::Future; +use http_server_common::cors::build_common_cors_config; +use limitation::Limiter; +use log::{info}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::sync::Arc; +use std::time::Duration; +use tokio::runtime::Runtime; +use twitch_common::twitch_secrets::TwitchSecrets; +use crate::threads::db_health_checker_thread::db_health_checker_thread::db_health_checker_thread; +use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus; + +// TODO TODO TODO TODO +// TODO TODO TODO TODO +// TODO TODO TODO TODO +// https://github.com/TensorSpeech/TensorFlowTTS (MAYBE USE THIS) +// TODO TODO TODO TODO +// TODO TODO TODO TODO +// TODO TODO TODO TODO + +// TODO TODO TODO -- ON signup, add an "early adopter" badge. And a tool for making it easy to add badges. +// TODO - badge for uploading template, badge for uploading model, etc. + +// TODO TODO -- also this: https://material-ui.com + + +const DEFAULT_BIND_ADDRESS : &'static str = "0.0.0.0:12345"; + +// Buckets (shared config) +const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; +const ENV_SECRET_KEY : &'static str = "SECRET_KEY"; +const ENV_REGION_NAME : &'static str = "REGION_NAME"; + +// Buckets (private data) +const ENV_PRIVATE_BUCKET_NAME : &'static str = "W2L_PRIVATE_DOWNLOAD_BUCKET_NAME"; +// Buckets (public data) +const ENV_PUBLIC_BUCKET_NAME : &'static str = "W2L_PUBLIC_DOWNLOAD_BUCKET_NAME"; + +// Various bucket roots +const ENV_AUDIO_UPLOADS_BUCKET_ROOT : &'static str = "AUDIO_UPLOADS_BUCKET_ROOT"; + +#[actix_web::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + // NB: Do not check this secrets-containing dotenv file into VCS. + // This file should only contain *development* secrets, never production. + let _ = dotenv::from_filename(".env-secrets").ok(); + + let common_env = CommonEnv::read_from_env()?; + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("storyteller-web-unknown".to_string()); + + info!("Hostname: {}", &server_hostname); + + info!("Connecting to database..."); + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + let pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + let firehose_publisher = FirehosePublisher { + mysql_pool: pool.clone(), // NB: Pool is clone/sync/send-safe + }; + + let badge_granter = BadgeGranter { + mysql_pool: pool.clone(), // NB: Pool is clone/sync/send-safe + firehose_publisher: firehose_publisher.clone(), // NB: Also safe + }; + + let redis_manager = RedisConnectionManager::new( + common_env.redis_0_connection_string.clone())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + info!("Setting up Redis rate limiters..."); + + // Old env vars: + // + // "LIMITER_ENABLED" + // "LIMITER_MAX_REQUESTS" + // "LIMITER_WINDOW_SECONDS" + // + let logged_out_redis_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_OUT_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_OUT_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_OUT_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "logged_out", limiter_enabled) + }; + + let logged_in_redis_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_LOGGED_IN_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_LOGGED_IN_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_LOGGED_IN_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "logged_in", limiter_enabled) + }; + + let model_upload_rate_limiter = { + let limiter_enabled = easyenv::get_env_bool_or_default("LIMITER_MODEL_UPLOAD_ENABLED", true); + let limiter_max_requests = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_MAX_REQUESTS", 3)?; + let limiter_window_seconds = easyenv::get_env_num("LIMITER_MODEL_UPLOAD_WINDOW_SECONDS", 10)?; + + let limiter = Limiter::build(&common_env.redis_0_connection_string) + .limit(limiter_max_requests) + .period(Duration::from_secs(limiter_window_seconds)) + .finish()?; + + RedisRateLimiter::new(limiter, "model_upload", limiter_enabled) + }; + + info!("Reading env vars and setting up utils..."); + + let bind_address = easyenv::get_env_string_or_default("BIND_ADDRESS", DEFAULT_BIND_ADDRESS); + let num_workers = easyenv::get_env_num("NUM_WORKERS", 8)?; + let hmac_secret = easyenv::get_env_string_or_default("COOKIE_SECRET", "notsecret"); + let cookie_domain = easyenv::get_env_string_or_default("COOKIE_DOMAIN", ".vo.codes"); + let cookie_secure = easyenv::get_env_bool_or_default("COOKIE_SECURE", true); + let cookie_http_only = easyenv::get_env_bool_or_default("COOKIE_HTTP_ONLY", true); + let website_homepage_redirect = + easyenv::get_env_string_or_default("WEBSITE_HOMEPAGE_REDIRECT", "https://vo.codes/"); + + let cookie_manager = CookieManager::new(&cookie_domain, &hmac_secret); + let session_checker = SessionChecker::new(&cookie_manager); + + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + + // Private and Public Buckets + let private_bucket_name = easyenv::get_env_string_required(ENV_PRIVATE_BUCKET_NAME)?; + let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?; + + // Bucket roots + let audio_uploads_bucket_root= easyenv::get_env_string_required(ENV_AUDIO_UPLOADS_BUCKET_ROOT)?; + + let private_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &private_bucket_name, + None, + )?; + + let public_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &public_bucket_name, + None, + )?; + + // In-Memory Cache + let voice_list_cache_ttl = easyenv::get_env_duration_seconds_or_default("VOICE_LIST_CACHE_TTL_SECONDS", Duration::from_secs(60)); + let category_list_cache_ttl = easyenv::get_env_duration_seconds_or_default("CATEGORY_LIST_CACHE_TTL_SECONDS", Duration::from_secs(60)); + let voice_list_cache = SingleItemTtlCache::create_with_duration(voice_list_cache_ttl); + let category_list_cache = SingleItemTtlCache::create_with_duration(category_list_cache_ttl); + + // NB: This secret really isn't too important. + // We can even rotate it without too much impact to users. + let sort_key_crypto_secret = + easyenv::get_env_string_or_default("SORT_KEY_SECRET", "webscale"); + let sort_key_crypto = SortKeyCrypto::new(&sort_key_crypto_secret); + + let twitch_oauth_redirect_landing_url = easyenv::get_env_string_or_default( + "TWITCH_OAUTH_REDIRECT_LANDING_URL", + "https://api.jungle.horse/twitch/oauth/enroll_redirect_landing"); + + let twitch_oauth_redirect_landing_finished_url = easyenv::get_env_string_or_default( + "TWITCH_OAUTH_REDIRECT_LANDING_FINISHED_URL", + "https://jungle.horse/"); + + let twitch_secrets = TwitchSecrets::from_env()?; + + // Background jobs. + + let health_check_status = HealthCheckStatus::new(); + let health_check_status2 = health_check_status.clone(); + let ip_banlist = IpBanlistSet::new(); + let ip_banlist2 = ip_banlist.clone(); + let mysql_pool3 = pool.clone(); + let mysql_pool4 = pool.clone(); + + let tokio_runtime = Runtime::new()?; + + info!("Spawning DB health checker thread."); + + tokio_runtime.spawn(async { + db_health_checker_thread(health_check_status2, mysql_pool3).await; + }); + + info!("Spawning IP ban polling thread."); + + tokio_runtime.spawn(async { + poll_ip_bans(ip_banlist2, mysql_pool4).await; + }); + + let server_state = ServerState { + env_config: EnvConfig { + num_workers, + bind_address, + cookie_domain, + cookie_secure, + cookie_http_only, + website_homepage_redirect, + }, + hostname: server_hostname, + health_check_status, + mysql_pool: pool, + redis_pool, + redis_rate_limiters: RedisRateLimiters { + logged_out: logged_out_redis_rate_limiter, + logged_in: logged_in_redis_rate_limiter, + model_upload: model_upload_rate_limiter, + }, + firehose_publisher, + badge_granter, + cookie_manager, + session_checker, + private_bucket_client, + public_bucket_client, + audio_uploads_bucket_root, + sort_key_crypto, + ip_banlist, + caches: InMemoryCaches { + voice_list: voice_list_cache, + category_list: category_list_cache, + }, + twitch_oauth: TwitchOauth { + secrets: TwitchOauthSecrets { + client_id: twitch_secrets.app_client_id, + client_secret: twitch_secrets.app_client_secret, + }, + redirect_landing_url: twitch_oauth_redirect_landing_url, + redirect_landing_finished_url: twitch_oauth_redirect_landing_finished_url, + } + }; + + serve(server_state) + .await?; + Ok(()) +} + +pub async fn serve(server_state: ServerState) -> AnyhowResult<()> +{ + let bind_address = server_state.env_config.bind_address.clone(); + let num_workers = server_state.env_config.num_workers.clone(); + let hostname = server_state.hostname.clone(); + + let server_state_arc = web::Data::new(Arc::new(server_state)); + + info!("Starting HTTP service."); + + let log_format = "[%{HOSTNAME}e] %a \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T"; + + HttpServer::new(move || { + // NB: Safe to clone due to internal arc + let ip_banlist = server_state_arc.ip_banlist.clone(); + + let app = App::new() + .app_data(server_state_arc.clone()) + .wrap(build_common_cors_config()) + .wrap(DefaultHeaders::new() + .header("X-Backend-Hostname", &hostname) + .header("X-Build-Sha", "")) + .wrap(IpFilter::new(ip_banlist)) + .wrap(Logger::new(&log_format) + .exclude("/liveness") + .exclude("/readiness")); + + add_routes(app) + }) + .bind(bind_address)? + .workers(num_workers) + .run() + .await?; + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/routes.rs b/crates/service/storyteller_web/src/routes.rs new file mode 100644 index 0000000000..432503075d --- /dev/null +++ b/crates/service/storyteller_web/src/routes.rs @@ -0,0 +1,709 @@ +use actix_http::body::MessageBody; +use actix_service::ServiceFactory; +use actix_web::dev::{ServiceRequest, ServiceResponse}; +use actix_web::error::Error; +use actix_web::{App, web, HttpResponse}; +use crate::http_server::endpoints::api_tokens::create_api_token::create_api_token_handler; +use crate::http_server::endpoints::api_tokens::delete_api_token::delete_api_token_handler; +use crate::http_server::endpoints::api_tokens::edit_api_token::edit_api_token_handler; +use crate::http_server::endpoints::api_tokens::list_api_tokens::list_api_tokens_handler; +use crate::http_server::endpoints::categories::assign_tts_category::assign_tts_category_handler; +use crate::http_server::endpoints::categories::create_category::create_category_handler; +use crate::http_server::endpoints::categories::get_category::get_category_handler; +use crate::http_server::endpoints::categories::list_tts_categories::list_tts_categories_handler; +use crate::http_server::endpoints::categories::list_tts_model_assigned_categories::list_tts_model_assigned_categories_handler; +use crate::http_server::endpoints::events::list_events::list_events_handler; +use crate::http_server::endpoints::investor_demo::disable_demo_mode_handler::disable_demo_mode_handler; +use crate::http_server::endpoints::investor_demo::enable_demo_mode_handler::enable_demo_mode_handler; +use crate::http_server::endpoints::leaderboard::get_leaderboard::leaderboard_handler; +use crate::http_server::endpoints::misc::default_route_404::default_route_404; +use crate::http_server::endpoints::misc::detect_locale_handler::detect_locale_handler; +use crate::http_server::endpoints::misc::enable_alpha_easy_handler::enable_alpha_easy_handler; +use crate::http_server::endpoints::misc::enable_alpha_handler::enable_alpha_handler; +use crate::http_server::endpoints::misc::root_index::get_root_index; +use crate::http_server::endpoints::moderation::approval::pending_w2l_templates::get_pending_w2l_templates_handler; +use crate::http_server::endpoints::moderation::categories::delete_category::delete_category_handler; +use crate::http_server::endpoints::moderation::categories::edit_category::edit_category_handler; +use crate::http_server::endpoints::moderation::categories::list_tts_categories_for_moderation::list_tts_categories_for_moderation_handler; +use crate::http_server::endpoints::moderation::ip_bans::add_ip_ban::add_ip_ban_handler; +use crate::http_server::endpoints::moderation::ip_bans::delete_ip_ban::delete_ip_ban_handler; +use crate::http_server::endpoints::moderation::ip_bans::get_ip_ban::get_ip_ban_handler; +use crate::http_server::endpoints::moderation::ip_bans::list_ip_bans::list_ip_bans_handler; +use crate::http_server::endpoints::moderation::jobs::get_tts_inference_queue_count::get_tts_inference_queue_count_handler; +use crate::http_server::endpoints::moderation::jobs::get_w2l_inference_queue_count::get_w2l_inference_queue_count_handler; +use crate::http_server::endpoints::moderation::jobs::kill_tts_inference_jobs::kill_tts_inference_jobs_handler; +use crate::http_server::endpoints::moderation::stats::get_voice_count_stats::get_voice_count_stats_handler; +use crate::http_server::endpoints::moderation::user_bans::ban_user::ban_user_handler; +use crate::http_server::endpoints::moderation::user_bans::list_banned_users::list_banned_users_handler; +use crate::http_server::endpoints::moderation::user_roles::list_roles::list_user_roles_handler; +use crate::http_server::endpoints::moderation::user_roles::list_staff::list_staff_handler; +use crate::http_server::endpoints::moderation::user_roles::set_user_role::set_user_role_handler; +use crate::http_server::endpoints::moderation::users::list_users::list_users_handler; +use crate::http_server::endpoints::service::health_check_handler::get_health_check_handler; +use crate::http_server::endpoints::tts::delete_tts_model::delete_tts_model_handler; +use crate::http_server::endpoints::tts::delete_tts_result::delete_tts_inference_result_handler; +use crate::http_server::endpoints::tts::edit_tts_model::edit_tts_model_handler; +use crate::http_server::endpoints::tts::edit_tts_result::edit_tts_inference_result_handler; +use crate::http_server::endpoints::tts::enqueue_infer_tts::infer_tts_handler; +use crate::http_server::endpoints::tts::enqueue_upload_tts_model::upload_tts_model_handler; +use crate::http_server::endpoints::tts::get_tts_inference_job_status::get_tts_inference_job_status_handler; +use crate::http_server::endpoints::tts::get_tts_model::get_tts_model_handler; +use crate::http_server::endpoints::tts::get_tts_model_use_count::get_tts_model_use_count_handler; +use crate::http_server::endpoints::tts::get_tts_result::get_tts_inference_result_handler; +use crate::http_server::endpoints::tts::get_tts_upload_model_job_status::get_tts_upload_model_job_status_handler; +use crate::http_server::endpoints::tts::list_tts_models::list_tts_models_handler; +use crate::http_server::endpoints::twitch::event_rules::create_event_rule::create_twitch_event_rule_handler; +use crate::http_server::endpoints::twitch::event_rules::delete_event_rule::delete_twitch_event_rule_handler; +use crate::http_server::endpoints::twitch::event_rules::edit_event_rule::edit_twitch_event_rule_handler; +use crate::http_server::endpoints::twitch::event_rules::get_event_rule::get_twitch_event_rule_for_user_handler; +use crate::http_server::endpoints::twitch::event_rules::list_event_rules_for_user::list_twitch_event_rules_for_user_handler; +use crate::http_server::endpoints::twitch::event_rules::reorder_twitch_event_rules::reorder_twitch_event_rules_handler; +use crate::http_server::endpoints::twitch::oauth::check_oauth_status::check_oauth_status_handler; +use crate::http_server::endpoints::twitch::oauth::oauth_begin_json::oauth_begin_enroll_json; +use crate::http_server::endpoints::twitch::oauth::oauth_begin_redirect::oauth_begin_enroll_redirect; +use crate::http_server::endpoints::twitch::oauth::oauth_end::oauth_end_enroll_from_redirect; +use crate::http_server::endpoints::users::create_account::create_account_handler; +use crate::http_server::endpoints::users::edit_profile::edit_profile_handler; +use crate::http_server::endpoints::users::get_profile::get_profile_handler; +use crate::http_server::endpoints::users::list_user_tts_inference_results::list_user_tts_inference_results_handler; +use crate::http_server::endpoints::users::list_user_tts_models::list_user_tts_models_handler; +use crate::http_server::endpoints::users::list_user_w2l_inference_results::list_user_w2l_inference_results_handler; +use crate::http_server::endpoints::users::list_user_w2l_templates::list_user_w2l_templates_handler; +use crate::http_server::endpoints::users::login::login_handler; +use crate::http_server::endpoints::users::logout::logout_handler; +use crate::http_server::endpoints::users::session_info::session_info_handler; +use crate::http_server::endpoints::voice_clone_requests::check_if_voice_clone_request_submitted::check_if_voice_clone_request_submitted_handler; +use crate::http_server::endpoints::voice_clone_requests::create_voice_clone_request::create_voice_clone_request_handler; +use crate::http_server::endpoints::w2l::delete_w2l_result::delete_w2l_inference_result_handler; +use crate::http_server::endpoints::w2l::delete_w2l_template::delete_w2l_template_handler; +use crate::http_server::endpoints::w2l::edit_w2l_result::edit_w2l_inference_result_handler; +use crate::http_server::endpoints::w2l::edit_w2l_template::edit_w2l_template_handler; +use crate::http_server::endpoints::w2l::enqueue_infer_w2l::infer_w2l_handler; +use crate::http_server::endpoints::w2l::enqueue_infer_w2l_with_uploads::enqueue_infer_w2l_with_uploads; +use crate::http_server::endpoints::w2l::enqueue_upload_w2l_template::upload_w2l_template_handler; +use crate::http_server::endpoints::w2l::get_w2l_inference_job_status::get_w2l_inference_job_status_handler; +use crate::http_server::endpoints::w2l::get_w2l_result::get_w2l_inference_result_handler; +use crate::http_server::endpoints::w2l::get_w2l_template::get_w2l_template_handler; +use crate::http_server::endpoints::w2l::get_w2l_template_use_count::get_w2l_template_use_count_handler; +use crate::http_server::endpoints::w2l::get_w2l_upload_template_job_status::get_w2l_upload_template_job_status_handler; +use crate::http_server::endpoints::w2l::list_w2l_templates::list_w2l_templates_handler; +use crate::http_server::endpoints::w2l::set_w2l_template_mod_approval::set_w2l_template_mod_approval_handler; + +pub fn add_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + let mut app = add_moderator_routes(app); /* /moderation */ + app = add_tts_routes(app); /* /tts */ + app = add_w2l_routes(app); /* /w2l */ + app = add_category_routes(app); /* /category */ + app = add_user_profile_routes(app); /* /user */ + app = add_api_token_routes(app); /* /api_tokens */ + app = add_voice_clone_request_routes(app); /* /voice_clone_requests */ + app = add_twitch_routes(app); /* /twitch */ // TODO: MAYBE TEMPORARY + app = add_investor_demo_routes(app); /* /demo_mode */ // TODO: DEFINITELY TEMPORARY + + // ==================== SERVICE ==================== + app.service( + web::resource("/_status") + .route(web::get().to(get_health_check_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + // ==================== ACCOUNT CREATION / SESSION MANAGEMENT ==================== + .service( + web::resource("/create_account") + .route(web::post().to(create_account_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/login") + .route(web::post().to(login_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/logout") + .route(web::post().to(logout_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/session") + .route(web::get().to(session_info_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + // ==================== MISC ==================== + .service( + web::resource("/events") + .route(web::get().to(list_events_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/detect_locale") + .route(web::get().to(detect_locale_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/leaderboard") + .route(web::get().to(leaderboard_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/") + .route(web::get().to(get_root_index)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(enable_alpha_handler) + .service(enable_alpha_easy_handler) + .default_service( web::route().to(default_route_404)) +} + +// ==================== MODERATOR ROUTES ==================== + +fn add_moderator_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/moderation") + .service( + web::resource("/staff") + .route(web::get().to(list_staff_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::scope("/ip_bans") + .service( + web::resource("/list") + .route(web::get().to(list_ip_bans_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/add") + .route(web::post().to(add_ip_ban_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{ip_address}") + .route(web::get().to(get_ip_ban_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{ip_address}/delete") + .route(web::post().to(delete_ip_ban_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/user") + .service( + web::resource("/list") + .route(web::get().to(list_users_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/user_bans") + .service( + web::resource("/list") + .route(web::get().to(list_banned_users_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/manage_ban") + .route(web::post().to(ban_user_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/roles") + .service( + web::resource("/list") + .route(web::get().to(list_user_roles_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{username}/edit") + .route(web::post().to(set_user_role_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/jobs") + .service( + web::resource("/tts_inference_queue_stats") + .route(web::get().to(get_tts_inference_queue_count_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/kill_tts_inference_jobs") + .route(web::post().to(kill_tts_inference_jobs_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/w2l_inference_queue_stats") + .route(web::get().to(get_w2l_inference_queue_count_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/pending") + .service( + web::resource("/w2l_templates") + .route(web::get().to(get_pending_w2l_templates_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/w2l_inference_queue_stats") + .route(web::get().to(get_w2l_inference_queue_count_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/stats") + .service( + web::resource("/tts_voices") + .route(web::get().to(get_voice_count_stats_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/categories") + .service( + web::resource("/{token}/edit") + .route(web::post().to(edit_category_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{token}/delete") + .route(web::post().to(delete_category_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/tts/list") + .route(web::get().to(list_tts_categories_for_moderation_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + ) +} + +// ==================== TTS ROUTES ==================== + +fn add_tts_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/tts") + .service( + web::resource("/upload") + .route(web::post().to(upload_tts_model_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/inference") + .route(web::post().to(infer_tts_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/list") + .route(web::get().to(list_tts_models_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/model/{token}") + .route(web::get().to(get_tts_model_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/model/{token}/delete") + .route(web::post().to(delete_tts_model_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/model/{model_token}/count") + .route(web::get().to(get_tts_model_use_count_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/model/{model_token}/edit") + .route(web::post().to(edit_tts_model_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/result/{token}") + .route(web::get().to(get_tts_inference_result_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/result/{token}/edit") + .route(web::post().to(edit_tts_inference_result_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/result/{token}/delete") + .route(web::post().to(delete_tts_inference_result_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/job/{token}") + .route(web::get().to(get_tts_inference_job_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/upload_model_job/{token}") + .route(web::get().to(get_tts_upload_model_job_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + +// ==================== WAV2LIP ROUTES ==================== + +fn add_w2l_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/w2l") + .service( + web::resource("/upload") + .route(web::post().to(upload_w2l_template_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/inference") + .route(web::post().to(enqueue_infer_w2l_with_uploads)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/list") + .route(web::get().to(list_w2l_templates_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/template/{token}") + .route(web::get().to(get_w2l_template_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/template/{template_token}/count") + .route(web::get().to(get_w2l_template_use_count_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/template/{template_token}/edit") + .route(web::post().to(edit_w2l_template_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/template/{token}/moderate") + .route(web::post().to(set_w2l_template_mod_approval_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/template/{token}/delete") + .route(web::post().to(delete_w2l_template_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/result/{token}") + .route(web::get().to(get_w2l_inference_result_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/result/{token}/edit") + .route(web::post().to(edit_w2l_inference_result_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/result/{token}/delete") + .route(web::post().to(delete_w2l_inference_result_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/job/{token}") + .route(web::get().to(get_w2l_inference_job_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/upload_template_job/{token}") + .route(web::get().to(get_w2l_upload_template_job_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + +// ==================== CATEGORY ROUTES ==================== + +fn add_category_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/category") + .service( + web::resource("/create") + .route(web::post().to(create_category_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/view/{token}") + .route(web::get().to(get_category_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::scope("/list") + .service( + web::resource("/tts") + .route(web::get().to(list_tts_categories_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/assign") + .service( + web::resource("/tts") + .route(web::post().to(assign_tts_category_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service( + web::scope("/assignments") + .service( + web::resource("/tts/{token}") + .route(web::get().to(list_tts_model_assigned_categories_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + ) +} + +// ==================== USER PROFILE ROUTES ==================== + +fn add_user_profile_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/user") + .service( + web::resource("/{username}/profile") + .route(web::get().to(get_profile_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{username}/edit_profile") + .route(web::post().to(edit_profile_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{username}/tts_models") + .route(web::get().to(list_user_tts_models_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{username}/tts_results") + .route(web::get().to(list_user_tts_inference_results_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{username}/w2l_templates") + .route(web::get().to(list_user_w2l_templates_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/{username}/w2l_results") + .route(web::get().to(list_user_w2l_inference_results_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + +// ==================== API TOKEN ROUTES ==================== + +fn add_api_token_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service(web::scope("/api_tokens") + .service(web::resource("/create") + .route(web::post().to(create_api_token_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/edit") + .route(web::post().to(edit_api_token_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/{api_token}/delete") + .route(web::post().to(delete_api_token_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/list") + .route(web::get().to(list_api_tokens_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + +// ==================== API TOKEN ROUTES ==================== + +fn add_voice_clone_request_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service(web::scope("/voice_clone_requests") + .service(web::resource("/create") + .route(web::post().to(create_voice_clone_request_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/check") + .route(web::post().to(check_if_voice_clone_request_submitted_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + +// ==================== INVESTOR DEMO MODE ==================== + +fn add_investor_demo_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service(web::scope("/demo_mode") + .service(web::resource("/enable") + .route(web::get().to(enable_demo_mode_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/disable") + .route(web::get().to(disable_demo_mode_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + +// ==================== TWITCH ROUTES ==================== + +// TODO: Maybe move these into an "oauth-gateway" type http service. +// We'll want the domain to have reputation and not confuse people. +// It'll also be nice to accrue all oauth things here. +fn add_twitch_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + // TODO: Move oauth endpoints under /twitch/oauth (requires updating Twitch configs too.) + app.service(web::scope("/twitch") + .service(web::scope("/oauth") + .service(web::resource("/enroll_json") + .route(web::get().to(oauth_begin_enroll_json)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/enroll_redirect_begin") + .route(web::get().to(oauth_begin_enroll_redirect)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/enroll_redirect_landing") + .route(web::get().to(oauth_end_enroll_from_redirect)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/check") + .route(web::get().to(check_oauth_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + .service(web::scope("/event_rule") + .service(web::resource("/create") + .route(web::post().to(create_twitch_event_rule_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/reorder") + .route(web::post().to(reorder_twitch_event_rules_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/list") + .route(web::get().to(list_twitch_event_rules_for_user_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/{token}/info") + .route(web::get().to(get_twitch_event_rule_for_user_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/{token}/update") + .route(web::post().to(edit_twitch_event_rule_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/{token}/delete") + .route(web::delete().to(delete_twitch_event_rule_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) + ) +} diff --git a/crates/service/storyteller_web/src/script_execution/ffmpeg_generate_preview_image_command.rs b/crates/service/storyteller_web/src/script_execution/ffmpeg_generate_preview_image_command.rs new file mode 100644 index 0000000000..310e1b2061 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/ffmpeg_generate_preview_image_command.rs @@ -0,0 +1,105 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; + +/// This command is used to generate ffmpeg previews for videos. +#[derive(Clone)] +pub struct FfmpegGeneratePreviewImageCommand { +} + +impl FfmpegGeneratePreviewImageCommand { + pub fn new() -> Self { + Self { + } + } + + pub fn execute(&self, + input_image_filename: &str, + output_image_filename: &str, + output_width: u32, + output_height: u32, + spawn_process: bool) -> AnyhowResult<()> + { + let mut command = String::new(); + + let scale = format!("scale={}:{}", output_width, output_height); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str("ffmpeg"); + command.push_str(" -i "); + command.push_str(&input_image_filename); + command.push_str(" -vf "); + command.push_str(&scale); + + command.push_str(" "); // NB: no arg flag for output filename + command.push_str(&output_image_filename); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/ffmpeg_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/ffmpeg_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/ffmpeg_generate_preview_video_command.rs b/crates/service/storyteller_web/src/script_execution/ffmpeg_generate_preview_video_command.rs new file mode 100644 index 0000000000..d160c6e077 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/ffmpeg_generate_preview_video_command.rs @@ -0,0 +1,113 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; + +/// This command is used to generate ffmpeg previews for videos. +#[derive(Clone)] +pub struct FfmpegGeneratePreviewVideoCommand { +} + +impl FfmpegGeneratePreviewVideoCommand { + pub fn new() -> Self { + Self { + } + } + + pub fn execute(&self, + input_video_filename: &str, + output_video_filename: &str, + output_width: u32, + output_height: u32, + looping: bool, + spawn_process: bool) -> AnyhowResult<()> + { + let mut command = String::new(); + + let dimensions = format!("{}x{}", output_width, output_height); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str("ffmpeg"); + command.push_str(" -i "); + command.push_str(&input_video_filename); + command.push_str(" -lossless 0 "); + command.push_str(" -ss 00:00:00 "); + command.push_str(" -t 00:00:03 "); + command.push_str(" -s "); + command.push_str(&dimensions); + + if looping { + command.push_str(" -loop 0 "); // NB: -loop 1 disables looping + } + + command.push_str(" "); // NB: no arg flag for output filename + command.push_str(&output_video_filename); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/ffmpeg_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/ffmpeg_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/google_drive_download_command.rs b/crates/service/storyteller_web/src/script_execution/google_drive_download_command.rs new file mode 100644 index 0000000000..b147b93205 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/google_drive_download_command.rs @@ -0,0 +1,66 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::token::random_crockford_token::random_crockford_token; +use log::{info, warn}; +use std::io::{BufReader, Read}; +use std::path::{PathBuf, Path}; +use std::process::Command; +use tempdir::TempDir; + +/// This is a Python script that uses the `gdown` package to download from Google Drive. +/// We're using this because it's a hack that gets around OAuth gateways. All the Rust +/// crates require OAuth permissions. Ugh. +/// +/// Filename: `download_internet_file.py` +/// Arguments: +/// --url (google drive, web, or youtube url) +/// --output_file (local download filename) +pub struct GoogleDriveDownloadCommand { + command: String, +} + +impl GoogleDriveDownloadCommand { + pub fn new(command: &str) -> Self { + Self { + command: command.to_string(), + } + } + + /// Download file from Google Drive into the `TempDir`. + /// Return the local filesystem filename. + pub async fn download_file(&self, + download_url: &str, + temp_dir: &TempDir) -> AnyhowResult + { + let temp_dir_path = temp_dir.path() + .to_str() + .unwrap_or("/tmp") + .to_string(); + + let temp_filename = random_crockford_token(10); + let temp_filename = format!("{}/{}.bin", temp_dir_path, temp_filename); + + info!("Downloading {} to: {}", download_url, temp_filename); + + let command = format!("{} --url {} --output_file {}", + &self.command, + &download_url, + &temp_filename); + + info!("Running command: {}", command); + + let result = Command::new("sh") + .arg("-c") + .arg(command) + .output()?; + + info!("Downloader Result: {:?}", result); + + if !result.status.success() { + let reason = String::from_utf8(result.stderr).unwrap_or("UNKNOWN".to_string()); + return Err(anyhow!("Failure to execute command: {:?}", reason)) + } + + Ok(temp_filename) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/imagemagick_generate_preview_image_command.rs b/crates/service/storyteller_web/src/script_execution/imagemagick_generate_preview_image_command.rs new file mode 100644 index 0000000000..cb6a4b7eb8 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/imagemagick_generate_preview_image_command.rs @@ -0,0 +1,109 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; + +/// This command is used to generate ffmpeg previews for videos. +#[derive(Clone)] +pub struct ImagemagickGeneratePreviewImageCommand { +} + +impl ImagemagickGeneratePreviewImageCommand { + pub fn new() -> Self { + Self { + } + } + + pub fn execute(&self, + input_image_filename: &str, + output_image_filename: &str, + output_width: u32, + output_height: u32, + spawn_process: bool) -> AnyhowResult<()> + { + let mut command = String::new(); + + // NB: `\!` ignores aspect ratio preservation + let dimensions = format!("{}x{}\\!", output_width, output_height); + + // magick wizard.png -quality 50 -define webp:lossless=true wizard.webp + command.push_str("echo 'test'"); + command.push_str(" && "); + //command.push_str("magick"); + command.push_str("convert "); + command.push_str(&input_image_filename); + command.push_str(" -resize "); + command.push_str(&dimensions); + command.push_str(" -quality 50 "); + command.push_str(" -define webp:loseless=false "); + + command.push_str(" "); // NB: no arg flag for output filename + command.push_str(&output_image_filename); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/ffmpeg_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/ffmpeg_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/mod.rs b/crates/service/storyteller_web/src/script_execution/mod.rs new file mode 100644 index 0000000000..a4c05b7424 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/mod.rs @@ -0,0 +1,9 @@ +pub mod ffmpeg_generate_preview_image_command; +pub mod ffmpeg_generate_preview_video_command; +pub mod google_drive_download_command; +pub mod imagemagick_generate_preview_image_command; +pub mod tacotron_inference_command; +pub mod tacotron_model_check_command; +pub mod talknet_model_check_command; +pub mod wav2lip_inference_command; +pub mod wav2lip_process_upload_command; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/script_execution/tacotron_inference_command.rs b/crates/service/storyteller_web/src/script_execution/tacotron_inference_command.rs new file mode 100644 index 0000000000..7249c1cabd --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/tacotron_inference_command.rs @@ -0,0 +1,124 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; +use std::path::Path; + +/// This command is used to run inference. +#[derive(Clone)] +pub struct TacotronInferenceCommand { + tacotron_directory: String, + inference_script_name: String, +} + +impl TacotronInferenceCommand { + pub fn new( + tacotron_directory: &str, + inference_script_name: &str, + ) -> Self { + Self { + tacotron_directory: tacotron_directory.to_string(), + inference_script_name: inference_script_name.to_string(), + } + } + + pub fn execute>( + &self, + synthesizer_checkpoint_path: P, + vocoder_checkpoint_path: P, + input_text_filename: P, + output_audio_filename: P, + output_spectrogram_filename: P, + output_metadata_filename: P, + spawn_process: bool + ) -> AnyhowResult<()> { + let mut command = String::new(); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str(&format!("cd {}", self.tacotron_directory)); + command.push_str(" && "); + command.push_str("source python/bin/activate"); + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.inference_script_name); + command.push_str(" --synthesizer_checkpoint_path "); + command.push_str(&synthesizer_checkpoint_path.as_ref().display().to_string()); + command.push_str(" --vocoder_checkpoint_path "); + command.push_str(&vocoder_checkpoint_path.as_ref().display().to_string()); + command.push_str(" --input_text_filename "); + command.push_str(&input_text_filename.as_ref().display().to_string()); + command.push_str(" --output_audio_filename "); + command.push_str(&output_audio_filename.as_ref().display().to_string()); + command.push_str(" --output_spectrogram_filename "); + command.push_str(&output_spectrogram_filename.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/tacotron_upload_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/tacotron_upload_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/tacotron_model_check_command.rs b/crates/service/storyteller_web/src/script_execution/tacotron_model_check_command.rs new file mode 100644 index 0000000000..e637d20290 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/tacotron_model_check_command.rs @@ -0,0 +1,119 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; +use std::path::Path; + +/// This command is used to check tacotron for being a real model +#[derive(Clone)] +pub struct TacotronModelCheckCommand { + /// Where the Tacotron code lives + tacotron_root_code_directory: String, + + /// eg. `source python/bin/activate` + virtual_env_activation_command: String, + + tacotron_model_check_script_name: String, +} + +impl TacotronModelCheckCommand { + pub fn new( + tacotron_root_code_directory: &str, + virtual_env_activation_command: &str, + tacotron_model_check_script_name: &str, + ) -> Self { + Self { + tacotron_root_code_directory: tacotron_root_code_directory.to_string(), + virtual_env_activation_command: virtual_env_activation_command.to_string(), + tacotron_model_check_script_name: tacotron_model_check_script_name.to_string(), + } + } + + pub fn execute>( + &self, + synthesizer_checkpoint_path: P, + output_metadata_filename: P, + spawn_process: bool + ) -> AnyhowResult<()> { + let mut command = String::new(); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str(&format!("cd {}", self.tacotron_root_code_directory)); + command.push_str(" && "); + command.push_str(&self.virtual_env_activation_command); + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.tacotron_model_check_script_name); + command.push_str(" --synthesizer_checkpoint_path "); + command.push_str(&synthesizer_checkpoint_path.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/tacotron_upload_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/tacotron_upload_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/talknet_model_check_command.rs b/crates/service/storyteller_web/src/script_execution/talknet_model_check_command.rs new file mode 100644 index 0000000000..dc7d2968bc --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/talknet_model_check_command.rs @@ -0,0 +1,112 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; +use std::path::Path; + +/// This command is used to check talknet for being a real model +#[derive(Clone)] +pub struct TalknetModelCheckCommand { + talknet_directory: String, + check_model_script_name: String, +} + +impl TalknetModelCheckCommand { + pub fn new( + talknet_directory: &str, + check_model_script_name: &str, + ) -> Self { + Self { + talknet_directory: talknet_directory.to_string(), + check_model_script_name: check_model_script_name.to_string(), + } + } + + pub fn execute>( + &self, + synthesizer_checkpoint_path: P, + output_metadata_filename: P, + spawn_process: bool + ) -> AnyhowResult<()> { + let mut command = String::new(); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str(&format!("cd {}", self.talknet_directory)); + command.push_str(" && "); + command.push_str("source python/bin/activate"); + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.check_model_script_name); + command.push_str(" --synthesizer_checkpoint_path "); + command.push_str(&synthesizer_checkpoint_path.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/talknet_upload_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/talknet_upload_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/wav2lip_inference_command.rs b/crates/service/storyteller_web/src/script_execution/wav2lip_inference_command.rs new file mode 100644 index 0000000000..bfe83a6be9 --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/wav2lip_inference_command.rs @@ -0,0 +1,136 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; +use std::path::Path; + +/// This command is used to run inference. +/// It uses preprocessed face files so that it's much faster. +#[derive(Clone)] +pub struct Wav2LipInferenceCommand { + w2l_directory: String, + script_name: String, + checkpoint_path: String, +} + +impl Wav2LipInferenceCommand { + pub fn new( + w2l_directory: &str, + script_name: &str, + checkpoint_path: &str, + ) -> Self { + Self { + w2l_directory: w2l_directory.to_string(), + script_name: script_name.to_string(), + checkpoint_path: checkpoint_path.to_string(), + } + } + + pub fn execute>( + &self, + checkpoint_path: P, + audio_filename: P, + end_bump_video_filename: P, + media_template_filename: P, + cached_faces_filename: P, + output_metadata_filename: P, + output_video_filename: P, + is_image: bool, + spawn_process: bool + ) -> AnyhowResult<()> { + let mut command = String::new(); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str(&format!("cd {}", self.w2l_directory)); + command.push_str(" && "); + command.push_str("source python/bin/activate"); + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.script_name); + command.push_str(" --checkpoint_path "); + command.push_str(&checkpoint_path.as_ref().display().to_string()); + command.push_str(" --image_or_video_filename "); + command.push_str(&media_template_filename.as_ref().display().to_string()); + command.push_str(" --audio_filename "); + command.push_str(&audio_filename.as_ref().display().to_string()); + command.push_str(" --end_bump_file "); + command.push_str(&end_bump_video_filename.as_ref().display().to_string()); + command.push_str(" --cached_faces_filename "); + command.push_str(&cached_faces_filename.as_ref().display().to_string()); + command.push_str(" --output_video_filename "); + command.push_str(&output_video_filename.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + if is_image { + command.push_str(" --is_image "); + } + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + if spawn_process { + // NB: This forks and returns immediately. + //let _child_pid = command_builder.spawn()?; + + let stdout_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/wav2lip_upload_stdout.txt")?; + + let stderr_file = OpenOptions::new() + .read(true) + .write(true) + .create(true) + .truncate(true) + .open("/tmp/wav2lip_upload_stderr.txt")?; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + stdout: Redirection::File(stdout_file), + stderr: Redirection::File(stderr_file), + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + p.detach(); + + } else { + // NB: This is a blocking call. + /*let output = command_builder.output()?; + + info!("Output status: {}", output.status); + info!("Stdout: {:?}", String::from_utf8(output.stdout)); + error!("Stderr: {:?}", String::from_utf8(output.stderr)); + + if !output.status.success() { + bail!("Bad error code: {:?}", output.status); + }*/ + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + } + + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/script_execution/wav2lip_process_upload_command.rs b/crates/service/storyteller_web/src/script_execution/wav2lip_process_upload_command.rs new file mode 100644 index 0000000000..e60047fdea --- /dev/null +++ b/crates/service/storyteller_web/src/script_execution/wav2lip_process_upload_command.rs @@ -0,0 +1,135 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection, ExitStatus}; +use std::fs::OpenOptions; +use std::error::Error; +use std::fmt; +use std::fmt::Formatter; + +/// The python script uses this exit code when face detection fails. +const FACE_DETECT_FAILURE_CODE : u32 = 5; + +/// This command is used to preprocess the face detection frames from user-submitted video. +/// This should only ever need to run once. The frames can then be uploaded to Buckets and saved. +#[derive(Clone)] +pub struct Wav2LipPreprocessClient { + w2l_directory: String, + script_name: String, + checkpoint_path: String, +} + +/// Command failure +#[derive(Debug, Copy, Clone)] +pub enum Wav2LipPreprocessError { + /// Permanent failure + NoFacesDetected, + /// Unknown. Retry allowed. + UnknownError, +} + +impl Error for Wav2LipPreprocessError {} + +impl fmt::Display for Wav2LipPreprocessError { + fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result { + match self { + Wav2LipPreprocessError::NoFacesDetected => write!(f, "NoFacesDetected"), + Wav2LipPreprocessError::UnknownError => write!(f, "UnknownError"), + } + } +} + +impl Wav2LipPreprocessClient { + pub fn new( + w2l_directory: &str, + script_name: &str, + checkpoint_path: &str, + ) -> Self { + Self { + w2l_directory: w2l_directory.to_string(), + script_name: script_name.to_string(), + checkpoint_path: checkpoint_path.to_string(), + } + } + + pub fn execute(&self, + image_or_video_filename: &str, + output_cached_faces_filename: &str, + output_metadata_filename: &str, + is_image: bool, + spawn_process: bool) -> Result<(), Wav2LipPreprocessError> + { + let mut command = String::new(); + + command.push_str("echo 'test'"); + command.push_str(" && "); + command.push_str(&format!("cd {}", self.w2l_directory)); + command.push_str(" && "); + command.push_str("source python/bin/activate"); + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.script_name); + command.push_str(" --checkpoint_path "); + command.push_str(&self.checkpoint_path); + command.push_str(" --image_or_video_filename "); + command.push_str(image_or_video_filename); + command.push_str(" --output_cached_faces_filename "); + command.push_str(output_cached_faces_filename); + command.push_str(" --output_metadata_filename "); + command.push_str(output_metadata_filename); + + if is_image { + command.push_str(" --is_image "); + } + + info!("Command: {:?}", command); + + // NB: Got rid of the previous (unused) child process spawning + + let exit_status = match self.do_execute(&command) { + Err(e) => { + warn!("Unknown execution error: {:?}", e); + return Err(Wav2LipPreprocessError::UnknownError); + }, + Ok(exit_status) => exit_status, + }; + + info!("Exit status: {:?}", exit_status); + + if !exit_status.success() { + if let ExitStatus::Exited(code) = exit_status { + if code == FACE_DETECT_FAILURE_CODE { + // We want to permanently fail retries. + warn!("Failure to detect faces error code returned. This is a permanent failure."); + return Err(Wav2LipPreprocessError::NoFacesDetected); + } + } + warn!("Unknown failure reason."); + return Err(Wav2LipPreprocessError::UnknownError); + } + + Ok(()) + } + + fn do_execute(&self, command: &str) -> AnyhowResult { + let command_parts = [ + "bash", + "-c", + command + ]; + + let mut p = Popen::create(&command_parts, PopenConfig { + //stdout: Redirection::Pipe, + //stderr: Redirection::Pipe, + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + + Ok(exit_status) + } +} diff --git a/crates/service/storyteller_web/src/server_state.rs b/crates/service/storyteller_web/src/server_state.rs new file mode 100644 index 0000000000..3d62a80690 --- /dev/null +++ b/crates/service/storyteller_web/src/server_state.rs @@ -0,0 +1,101 @@ +use crate::http_server::endpoints::tts::list_tts_models::TtsModelRecordForResponse; +use crate::http_server::web_utils::cookie_manager::CookieManager; +use crate::http_server::web_utils::redis_rate_limiter::RedisRateLimiter; +use crate::http_server::web_utils::session_checker::SessionChecker; +use crate::threads::ip_banlist_set::IpBanlistSet; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::caching::single_item_ttl_cache::SingleItemTtlCache; +use crate::util::encrypted_sort_id::SortKeyCrypto; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use r2d2_redis::{r2d2, RedisConnectionManager}; +use sqlx::MySqlPool; +use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus; +use crate::http_server::endpoints::categories::list_tts_categories::DisplayCategory; + +/// State that is injected into every endpoint. +#[derive(Clone)] +pub struct ServerState { + /// Configuration from ENV vars. + pub env_config: EnvConfig, + + pub hostname: String, + + pub health_check_status: HealthCheckStatus, + + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub redis_rate_limiters: RedisRateLimiters, + + pub cookie_manager: CookieManager, + + pub session_checker: SessionChecker, + + pub firehose_publisher: FirehosePublisher, + pub badge_granter: BadgeGranter, + + pub private_bucket_client: BucketClient, + pub public_bucket_client: BucketClient, + + /// Where to store audio uploads for w2l + pub audio_uploads_bucket_root: String, + + pub sort_key_crypto: SortKeyCrypto, + + pub ip_banlist: IpBanlistSet, + + pub caches: InMemoryCaches, + + pub twitch_oauth: TwitchOauth, +} + +#[derive(Clone)] +pub struct EnvConfig { + // Number of thread workers. + pub num_workers: usize, + pub bind_address: String, + pub cookie_domain: String, + pub cookie_secure: bool, + pub cookie_http_only: bool, + pub website_homepage_redirect: String, +} + +/// Necessary to run the OAuth flow. +#[derive(Clone)] +pub struct TwitchOauth { + pub secrets: TwitchOauthSecrets, + pub redirect_landing_url: String, + pub redirect_landing_finished_url: String, +} + +/// Necessary to run the OAuth flow. +#[derive(Clone)] +pub struct TwitchOauthSecrets { + pub client_id: String, + pub client_secret: String, +} + +/// Different rate limiters for different users +#[derive(Clone)] +pub struct RedisRateLimiters { + /// Logged out users have stricter limits + pub logged_out: RedisRateLimiter, + + /// Logged in users have a little more leeway + pub logged_in: RedisRateLimiter, + + /// A rate limiter for TTS and W2L uploads + pub model_upload: RedisRateLimiter, +} + +/// In-memory caches +#[derive(Clone)] +pub struct InMemoryCaches { + /// In-memory caches with TTL-based eviction. Contains a list of all voices. + pub voice_list: SingleItemTtlCache>, + + /// In-memory caches with TTL-based eviction. Contains a list of all TTS categories. + pub category_list: SingleItemTtlCache>, +} diff --git a/crates/service/storyteller_web/src/threads/db_health_checker_thread/db_health_check_status.rs b/crates/service/storyteller_web/src/threads/db_health_checker_thread/db_health_check_status.rs new file mode 100644 index 0000000000..3c63842e11 --- /dev/null +++ b/crates/service/storyteller_web/src/threads/db_health_checker_thread/db_health_check_status.rs @@ -0,0 +1,78 @@ +use anyhow::anyhow; +use crate::AnyhowResult; +use std::collections::HashSet; +use std::collections::hash_map::RandomState; +use std::sync::{RwLock, PoisonError, RwLockWriteGuard, RwLockReadGuard, Arc}; +use chrono::{Utc, DateTime, NaiveDateTime}; + +#[derive(Clone)] +pub struct HealthCheckStatus { + internal_status: Arc>, +} + +#[derive(Clone)] +pub struct HealthCheckStatusData { + pub last_db_time: Option, + pub is_healthy: bool, + /// Number of consecutive tries the DB has been good. + /// Resets at the first sign of bad. + pub healthy_check_consecutive_count: Option, + /// Number of consecutive tries the DB has been bad. + /// Resets at the first sign of good. + pub unhealthy_check_consecutive_count: Option, +} + +impl HealthCheckStatus { + + pub fn new() -> Self { + Self { + internal_status: Arc::new(RwLock::new(HealthCheckStatusData { + last_db_time: None, + is_healthy: false, // Start false + healthy_check_consecutive_count: None, + unhealthy_check_consecutive_count: None, + })), + } + } + + pub fn record_ping_success(&self, db_time: NaiveDateTime) -> AnyhowResult<()> { + match self.internal_status.write() { + Err(_) => Err(anyhow!("Can't write lock")), + Ok(mut lock) => { + (*lock).last_db_time = Some(db_time); + (*lock).is_healthy = true; + (*lock).unhealthy_check_consecutive_count = None; + + let healthy_count = (*lock).healthy_check_consecutive_count.unwrap_or(0); + (*lock).healthy_check_consecutive_count = Some(healthy_count.saturating_add(1)); + + Ok(()) + }, + } + } + + pub fn record_ping_failure(&self) -> AnyhowResult<()> { + match self.internal_status.write() { + Err(_) => Err(anyhow!("Can't write lock")), + Ok(mut lock) => { + (*lock).is_healthy = false; + (*lock).healthy_check_consecutive_count = None; + + let unhealthy_count = (*lock).unhealthy_check_consecutive_count.unwrap_or(0); + (*lock).unhealthy_check_consecutive_count = Some(unhealthy_count.saturating_add(1)); + + Ok(()) + }, + } + } + + /// Read the status + pub fn get_health_check_status(&self) -> AnyhowResult { + match self.internal_status.read() { + Err(_) => Err(anyhow!("Can't read lock")), + Ok(lock) => { + Ok(lock.clone()) + }, + } + } +} diff --git a/crates/service/storyteller_web/src/threads/db_health_checker_thread/db_health_checker_thread.rs b/crates/service/storyteller_web/src/threads/db_health_checker_thread/db_health_checker_thread.rs new file mode 100644 index 0000000000..7698d3a4c5 --- /dev/null +++ b/crates/service/storyteller_web/src/threads/db_health_checker_thread/db_health_checker_thread.rs @@ -0,0 +1,39 @@ +use anyhow::Error; +use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus; +use database_queries::queries::health_check::health_check_query::{health_check_db, HealthCheckResult}; +use log::debug; +use log::error; +use log::info; +use log::warn; +use sqlx::MySqlPool; +use std::thread; +use std::time::Duration; + +pub async fn db_health_checker_thread( + health_check_status: HealthCheckStatus, + mysql_pool: MySqlPool, +) { + loop { + debug!("Checking DB health..."); + + match health_check_db(&mysql_pool).await { + Ok(result) => { + match health_check_status.record_ping_success(result.present_time) { + Err(e) => error!("Problem updating application health checks!"), + Ok(_) => {}, + } + } + Err(e) => { + error!("Problem health checking database: {:?}", e); + match health_check_status.record_ping_failure() { + Err(e) => error!("Problem updating application health checks!"), + Ok(_) => {}, + } + } + } + + thread::sleep(Duration::from_millis(3_000)); + } + + warn!("Should never happen: Health Checker Exits"); +} diff --git a/crates/service/storyteller_web/src/threads/db_health_checker_thread/mod.rs b/crates/service/storyteller_web/src/threads/db_health_checker_thread/mod.rs new file mode 100644 index 0000000000..18698ccfad --- /dev/null +++ b/crates/service/storyteller_web/src/threads/db_health_checker_thread/mod.rs @@ -0,0 +1,2 @@ +pub mod db_health_check_status; +pub mod db_health_checker_thread; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/threads/ip_banlist_set.rs b/crates/service/storyteller_web/src/threads/ip_banlist_set.rs new file mode 100644 index 0000000000..f0ec2611bf --- /dev/null +++ b/crates/service/storyteller_web/src/threads/ip_banlist_set.rs @@ -0,0 +1,49 @@ +use anyhow::anyhow; +use crate::AnyhowResult; +use std::collections::HashSet; +use std::collections::hash_map::RandomState; +use std::sync::{RwLock, PoisonError, RwLockWriteGuard, RwLockReadGuard, Arc}; + +#[derive(Clone)] +pub struct IpBanlistSet { + pub ip_banlist: Arc>>, +} + +impl IpBanlistSet { + + pub fn new() -> Self { + Self { + ip_banlist: Arc::new(RwLock::new(HashSet::new())), + } + } + + pub fn replace_list(&self, ip_banlist: HashSet) -> AnyhowResult<()> { + match self.ip_banlist.write() { + Err(_) => Err(anyhow!("Can't write lock")), + Ok(mut lock) => { + *lock = ip_banlist; + Ok(()) + }, + } + } + + pub fn is_banned>(&self, ip_address: S) -> AnyhowResult { + match self.ip_banlist.read() { + Err(_) => Err(anyhow!("Can't read lock")), + Ok(read) => { + let result = read.contains(ip_address.as_ref()); + Ok(result) + }, + } + } + + pub fn num_bans(&self) -> AnyhowResult { + match self.ip_banlist.read() { + Err(_) => Err(anyhow!("Can't read lock")), + Ok(read) => { + let result = read.len(); + Ok(result) + }, + } + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/threads/mod.rs b/crates/service/storyteller_web/src/threads/mod.rs new file mode 100644 index 0000000000..de34ef34d1 --- /dev/null +++ b/crates/service/storyteller_web/src/threads/mod.rs @@ -0,0 +1,3 @@ +pub mod db_health_checker_thread; +pub mod ip_banlist_set; +pub mod poll_ip_banlist_thread; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/threads/poll_ip_banlist_thread.rs b/crates/service/storyteller_web/src/threads/poll_ip_banlist_thread.rs new file mode 100644 index 0000000000..babb3fe49a --- /dev/null +++ b/crates/service/storyteller_web/src/threads/poll_ip_banlist_thread.rs @@ -0,0 +1,44 @@ +use anyhow::Error; +use crate::threads::ip_banlist_set::IpBanlistSet; +use database_queries::queries::ip_bans::list_ip_bans::list_ip_bans; +use log::info; +use log::warn; +use sqlx::MySqlPool; +use std::collections::HashSet; +use std::thread; +use std::time::Duration; + +pub async fn poll_ip_bans( + ip_banlist_set: IpBanlistSet, + mysql_pool: MySqlPool, +) { + loop { + info!("Job fetching IP Address Bans..."); + + let bans = match list_ip_bans(&mysql_pool).await { + Ok(bans) => bans, + Err(e) => { + warn!("Error polling IP bans: {:?}", e); + thread::sleep(Duration::from_millis(30_000)); + continue; + } + }; + + let ip_addresses = bans.iter() + .map(|record| record.ip_address.clone()) + .collect::>(); + + info!("Job found {} IP Address Bans.", ip_addresses.len()); + + match ip_banlist_set.replace_list(ip_addresses) { + Ok(_) => info!("internal banlist updated!"), + Err(e) => { + warn!("error replacing banlist: {:?}", e); + }, + } + + thread::sleep(Duration::from_millis(20_000)); + } + + warn!("Should never happen: IP poller exits"); +} diff --git a/crates/service/storyteller_web/src/tts_download_job.rs b/crates/service/storyteller_web/src/tts_download_job.rs new file mode 100644 index 0000000000..d9ef14342a --- /dev/null +++ b/crates/service/storyteller_web/src/tts_download_job.rs @@ -0,0 +1,485 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate serde_derive; + +pub mod common_env; +pub mod job_queries; +pub mod script_execution; +pub mod util; + +use anyhow::anyhow; +use chrono::Utc; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use crate::common_env::CommonEnv; +use crate::job_queries::tts_download_job_queries::TtsUploadJobRecord; +use crate::job_queries::tts_download_job_queries::grab_job_lock_and_mark_pending; +use crate::job_queries::tts_download_job_queries::insert_tts_model; +use crate::job_queries::tts_download_job_queries::mark_tts_upload_job_done; +use crate::job_queries::tts_download_job_queries::mark_tts_upload_job_failure; +use crate::job_queries::tts_download_job_queries::query_tts_upload_job_records; +use crate::script_execution::google_drive_download_command::GoogleDriveDownloadCommand; +use crate::script_execution::tacotron_model_check_command::TacotronModelCheckCommand; +use crate::script_execution::talknet_model_check_command::TalknetModelCheckCommand; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::buckets::bucket_path_unifier::BucketPathUnifier; +use crate::util::buckets::bucket_paths::hash_to_bucket_path; +use crate::util::filesystem::safe_delete_temp_directory; +use crate::util::filesystem::safe_delete_temp_file; +use crate::util::filesystem::{check_directory_exists, check_file_exists}; +use crate::util::hashing::hash_file_sha2::hash_file_sha2; +use crate::util::noop_logger::NoOpLogger; +use crate::util::redis::redis_job_status_logger::RedisJobStatusLogger; +use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE}; +use database_queries::column_types::tts_model_type::TtsModelType; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use log::{warn, info}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use ring::digest::{Context, Digest, SHA256}; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::fs::File; +use std::fs; +use std::io::{BufReader, Read, Error}; +use std::ops::Deref; +use std::path::{PathBuf, Path}; +use std::process::Command; +use std::time::Duration; +use tempdir::TempDir; + +// Buckets +const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; +const ENV_SECRET_KEY : &'static str = "SECRET_KEY"; +const ENV_REGION_NAME : &'static str = "REGION_NAME"; +const ENV_BUCKET_NAME : &'static str = "TTS_DOWNLOAD_BUCKET_NAME"; +const ENV_BUCKET_ROOT : &'static str = "TTS_DOWNLOAD_BUCKET_ROOT"; + +const DEFAULT_TEMP_DIR: &'static str = "/tmp"; + +struct Downloader { + pub download_temp_directory: PathBuf, + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub bucket_client: BucketClient, + pub firehose_publisher: FirehosePublisher, + pub badge_granter: BadgeGranter, + pub google_drive_downloader: GoogleDriveDownloadCommand, + + pub bucket_path_unifier: BucketPathUnifier, + + pub tacotron_tts_check: TacotronModelCheckCommand, + pub talknet_tts_check: TalknetModelCheckCommand, + + // Command to run + pub download_script: String, + // Root to store TTS results + pub bucket_root_tts_model_uploads: String, + + // Sleep between batches + pub job_batch_wait_millis: u64, + + // How long to wait between log lines + pub no_op_logger_millis: u64, + + // Max job attempts before failure. + // NB: This is an i32 so we don't need to convert to db column type. + pub job_max_attempts: i32, +} + +#[tokio::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + let _ = dotenv::from_filename(".env-secrets").ok(); + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("tts-download-job".to_string()); + + info!("Hostname: {}", &server_hostname); + + // Bucket stuff + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + let bucket_name = easyenv::get_env_string_required(ENV_BUCKET_NAME)?; + let bucket_root = easyenv::get_env_string_required(ENV_BUCKET_ROOT)?; + + let bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &bucket_name, + None, + )?; + + let temp_directory = easyenv::get_env_string_or_default( + "DOWNLOAD_TEMP_DIR", + DEFAULT_TEMP_DIR); + + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "./scripts/download_internet_file.py"); + + let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script); + + let temp_directory = PathBuf::from(temp_directory); + + check_directory_exists(&temp_directory)?; + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + info!("Connecting to database..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + let common_env = CommonEnv::read_from_env()?; + + info!("Connecting to redis..."); + + let redis_manager = + RedisConnectionManager::new(common_env.redis_0_connection_string.deref())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + }; + + let badge_granter = BadgeGranter { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + firehose_publisher: firehose_publisher.clone(), // NB: Also safe + }; + + let tacotron_root_code_directory = easyenv::get_env_string_required("TACOTRON_ROOT_CODE_DIRECTORY")?; + let tacotron_virtual_env_activation_command = easyenv::get_env_string_or_default( + "TACOTRON_VIRTUAL_ENV_ACTIVATION_COMMAND", + "source python-tacotron/bin/activate"); + + let tacotron_model_check_script_name = easyenv::get_env_string_or_default( + "TACOTRON_MODEL_CHECK_SCRIPT_NAME", + "vocodes_model_check_tacotron.py"); + + let tacotron_check_command= TacotronModelCheckCommand::new( + &tacotron_root_code_directory, + &tacotron_virtual_env_activation_command, + &tacotron_model_check_script_name, + ); + + let talknet_root_code_directory = easyenv::get_env_string_required("TALKNET_ROOT_CODE_DIRECTORY")?; + let talknet_check_script_name = easyenv::get_env_string_required("TALKNET_MODEL_CHECK_SCRIPT_NAME")?; + + let talknet_check_command= TalknetModelCheckCommand::new( + &talknet_root_code_directory, + &talknet_check_script_name, + ); + + let downloader = Downloader { + download_temp_directory: temp_directory, + mysql_pool, + redis_pool, + bucket_client, + download_script, + google_drive_downloader, + bucket_path_unifier: BucketPathUnifier::default_paths(), + bucket_root_tts_model_uploads: bucket_root.to_string(), + firehose_publisher, + badge_granter, + tacotron_tts_check: tacotron_check_command, + talknet_tts_check: talknet_check_command, + job_batch_wait_millis: common_env.job_batch_wait_millis, + job_max_attempts: common_env.job_max_attempts as i32, + no_op_logger_millis: common_env.no_op_logger_millis, + }; + + main_loop(downloader).await; + + Ok(()) +} + +const START_TIMEOUT_MILLIS : u64 = 500; +const INCREASE_TIMEOUT_MILLIS : u64 = 1000; + +async fn main_loop(downloader: Downloader) { + let mut error_timeout_millis = START_TIMEOUT_MILLIS; + + let mut noop_logger = NoOpLogger::new(downloader.no_op_logger_millis as i64); + + loop { + let num_records = 1; + let query_result = query_tts_upload_job_records(&downloader.mysql_pool, num_records).await; + + let jobs = match query_result { + Ok(jobs) => jobs, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if jobs.is_empty() { + noop_logger.log_after_awhile(); + + std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis)); + continue; + } + + let result = process_jobs(&downloader, jobs).await; + + match result { + Ok(_) => {}, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + } + + error_timeout_millis = START_TIMEOUT_MILLIS; // reset + + std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis)); + } +} + +async fn process_jobs(downloader: &Downloader, jobs: Vec) -> AnyhowResult<()> { + for job in jobs.into_iter() { + let result = process_job(downloader, &job).await; + match result { + Ok(_) => {}, + Err(e) => { + warn!("Failure to process job: {:?}", e); + let failure_reason = ""; + let _r = mark_tts_upload_job_failure( + &downloader.mysql_pool, + &job, + failure_reason, + downloader.job_max_attempts + ).await; + } + } + } + + Ok(()) +} + +#[derive(Deserialize)] +struct FileMetadata { + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + +async fn process_job(downloader: &Downloader, job: &TtsUploadJobRecord) -> AnyhowResult<()> { + // TODO: 1. Mark processing. (DONE) + // TODO: 2. Download. (DONE) + // TODO: 3. Upload. (DONE) + // TODO: 4. Save record. (DONE) + // TODO: 5. Mark job done. (DONE) + + let mut redis = downloader.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_tts_download( + &mut redis, + &job.token); + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + let lock_acquired = grab_job_lock_and_mark_pending(&downloader.mysql_pool, job).await?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {}", &job.id); + return Ok(()) + } + + // ==================== SETUP TEMP DIRS ==================== // + + let temp_dir = format!("temp_{}", job.id); + let temp_dir = TempDir::new(&temp_dir)?; + + // ==================== DOWNLOAD MODEL FILE ==================== // + + info!("Calling downloader..."); + + redis_logger.log_status("downloading model")?; + + let download_url = job.download_url.as_ref() + .map(|c| c.to_string()) + .unwrap_or("".to_string()); + + let download_filename = match downloader.google_drive_downloader.download_file(&download_url, &temp_dir).await { + Ok(filename) => filename, + Err(e) => { + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + let model_type = if download_filename.to_lowercase().ends_with("zip") { + TtsModelType::Talknet + } else { + TtsModelType::Tacotron2 + }; + + info!("Uploaded model type: {:?}", model_type); + + if model_type == TtsModelType::Talknet { + // TODO: Finish supporting TalkNet + warn!("Unsupported model type"); + return Err(anyhow!("unsupported model type")); + } + + // ==================== RUN MODEL CHECK ==================== // + + info!("Checking that model is valid..."); + + redis_logger.log_status("checking model")?; + + let file_path = PathBuf::from(download_filename.clone()); + + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + + match model_type { + TtsModelType::Tacotron2 => { + let result = downloader.tacotron_tts_check.execute( + &file_path, + &output_metadata_fs_path, + false, + ); + + if let Err(e) = result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + } + }, + TtsModelType::Talknet => { + let result = downloader.talknet_tts_check.execute( + &file_path, + &output_metadata_fs_path, + false, + ); + + if let Err(e) = result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + } + }, + } + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that metadata output file exists..."); + + check_file_exists(&output_metadata_fs_path)?; + + let file_metadata = match read_metadata_file(&output_metadata_fs_path) { + Ok(metadata) => metadata, + Err(e) => { + safe_delete_temp_file(&file_path); + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + // ==================== UPLOAD MODEL FILE ==================== // + + info!("Uploading model to GCS..."); + + let private_bucket_hash = hash_file_sha2(&download_filename)?; + + info!("File hash: {}", private_bucket_hash); + + let synthesizer_model_bucket_path = match model_type { + TtsModelType::Tacotron2 => downloader.bucket_path_unifier.tts_synthesizer_path(&private_bucket_hash), + TtsModelType::Talknet => downloader.bucket_path_unifier.tts_zipped_synthesizer_path(&private_bucket_hash), + }; + + info!("Destination bucket path: {:?}", &synthesizer_model_bucket_path); + + redis_logger.log_status("uploading model")?; + + if let Err(e) = downloader.bucket_client.upload_filename(&synthesizer_model_bucket_path, &file_path).await { + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving model record..."); + let (id, model_token) = insert_tts_model( + &downloader.mysql_pool, + job, + &private_bucket_hash, + synthesizer_model_bucket_path, + file_metadata.file_size_bytes) + .await?; + + info!("Marking job complete..."); + mark_tts_upload_job_done( + &downloader.mysql_pool, + job, + true, + Some(&model_token) + ).await?; + + info!("Saved model record: {}", id); + + downloader.firehose_publisher.publish_tts_model_upload_finished(&job.creator_user_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + downloader.badge_granter.maybe_grant_tts_model_uploads_badge(&job.creator_user_token) + .await + .map_err(|e| { + warn!("error maybe awarding badge: {:?}", e); + anyhow!("error maybe awarding badge") + })?; + + redis_logger.log_status("done")?; + + info!("Job done: {}", job.id); + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/tts_inference_job.rs b/crates/service/storyteller_web/src/tts_inference_job.rs new file mode 100644 index 0000000000..1d952670b6 --- /dev/null +++ b/crates/service/storyteller_web/src/tts_inference_job.rs @@ -0,0 +1,963 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate serde_derive; + +pub mod common_env; +pub mod http_clients; +pub mod job_queries; +pub mod script_execution; +pub mod util; + +use anyhow::{anyhow, Error}; +use chrono::{Utc, DateTime, TimeZone}; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use container_common::token::random_uuid::generate_random_uuid; +use crate::common_env::CommonEnv; +use crate::http_clients::tts_inference_sidecar_client::TtsInferenceSidecarClient; +use crate::script_execution::tacotron_inference_command::TacotronInferenceCommand; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::buckets::bucket_path_unifier::BucketPathUnifier; +use crate::util::buckets::bucket_paths::hash_to_bucket_path; +use crate::util::filesystem::check_directory_exists; +use crate::util::filesystem::check_file_exists; +use crate::util::filesystem::safe_delete_temp_directory; +use crate::util::filesystem::safe_delete_temp_file; +use crate::util::hashing::hash_file_sha2::hash_file_sha2; +use crate::util::hashing::hash_string_sha2::hash_string_sha2; +use crate::util::jobs::cache_miss_strategizer::CacheMissStrategizer; +use crate::util::jobs::cache_miss_strategizer::CacheMissStrategy; +use crate::util::jobs::cache_miss_strategizer_multi::SyncMultiCacheMissStrategizer; +use crate::util::jobs::virtual_lfu_cache::SyncVirtualLfuCache; +use crate::util::noop_logger::NoOpLogger; +use crate::util::redis::redis_job_status_logger::RedisJobStatusLogger; +use crate::util::semi_persistent_cache_dir::SemiPersistentCacheDir; +use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE}; +use database_queries::column_types::vocoder_type::VocoderType; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use database_queries::queries::tts::tts_inference_jobs::list_available_tts_inference_jobs::{AvailableTtsInferenceJob, list_available_tts_inference_jobs}; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_done::mark_tts_inference_job_done; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_failure::mark_tts_inference_job_failure; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_pending_and_grab_lock::mark_tts_inference_job_pending_and_grab_lock; +use database_queries::queries::tts::tts_inference_jobs::mark_tts_inference_job_permanently_dead::mark_tts_inference_job_permanently_dead; +use database_queries::queries::tts::tts_models::get_tts_model_for_inference::{get_tts_model_for_inference, TtsModelForInferenceError, TtsModelForInferenceRecord}; +use database_queries::queries::tts::tts_results::insert_tts_result::insert_tts_result; +use log::{warn, info, error}; +use newrelic_telemetry::Client as NewRelicClient; +use newrelic_telemetry::ClientBuilder; +use newrelic_telemetry::Span; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2::PooledConnection; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use ring::digest::{Context, Digest, SHA256}; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::collections::HashMap; +use std::fs::{File, metadata}; +use std::io::{BufReader, Read}; +use std::ops::Deref; +use std::path::{PathBuf, Path}; +use std::process::Command; +use std::thread; +use std::time::{Duration, SystemTime, UNIX_EPOCH, Instant}; +use tempdir::TempDir; +use tts_common::clean_symbols::clean_symbols; + +// Buckets (shared config) +const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; +const ENV_SECRET_KEY : &'static str = "SECRET_KEY"; +const ENV_REGION_NAME : &'static str = "REGION_NAME"; + +// Bucket names +const ENV_PRIVATE_BUCKET_NAME : &'static str = "PRIVATE_BUCKET_NAME"; +const ENV_PUBLIC_BUCKET_NAME : &'static str = "PUBLIC_BUCKET_NAME"; + +// Where models and other assets get downloaded to. +const ENV_SEMIPERSISTENT_CACHE_DIR : &'static str = "SEMIPERSISTENT_CACHE_DIR"; + +// Python code +const ENV_CODE_DIRECTORY : &'static str = "TTS_CODE_DIRECTORY"; +const ENV_INFERENCE_SCRIPT_NAME : &'static str = "TTS_INFERENCE_SCRIPT_NAME"; + +// HTTP sidecar +const ENV_TTS_INFERENCE_SIDECAR_HOSTNAME: &'static str = "TTS_INFERENCE_SIDECAR_HOSTNAME"; + +const DEFAULT_TEMP_DIR: &'static str = "/tmp"; + +struct Inferencer { + pub download_temp_directory: PathBuf, + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub private_bucket_client: BucketClient, + pub public_bucket_client: BucketClient, + + pub firehose_publisher: FirehosePublisher, + + pub bucket_path_unifier: BucketPathUnifier, + pub semi_persistent_cache: SemiPersistentCacheDir, + + pub tts_inference_command: TacotronInferenceCommand, + pub tts_inference_sidecar_client: TtsInferenceSidecarClient, + + pub newrelic_client: NewRelicClient, + + // Keep tabs of which models to hold in the sidecar memory with this virtual LRU cache + pub virtual_model_lfu: SyncVirtualLfuCache, + pub cache_miss_strategizers: SyncMultiCacheMissStrategizer, + + // Waveglow vocoder filename + pub waveglow_vocoder_model_filename: String, + + // Hifigan vocoder filename + pub hifigan_vocoder_model_filename: String, + + // Hifigan super resolution vocoder filename + pub hifigan_superres_vocoder_model_filename: String, + + // Sleep between batches + pub job_batch_wait_millis: u64, + + // Max job attempts before failure. + // NB: This is an i32 so we don't need to convert to db column type. + pub job_max_attempts: i32, + + // Number of jobs to dequeue at once. + pub job_batch_size: u32, + + // How long to wait between log lines + pub no_op_logger_millis: u64, + + // Maximum number of synthesizer models to hold in memory. + pub sidecar_max_synthesizer_models: usize, + + // Typically we'll sort jobs by priority. Occasionally we introduce a chance for low + // priority jobs to run in the order they were enqueued. + // If this is set to "0", we no longer consider priority + pub low_priority_starvation_prevention_every_nth: usize, +} + +#[tokio::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + // NB: Do not check this secrets-containing dotenv file into VCS. + // This file should only contain *development* secrets, never production. + let _ = dotenv::from_filename(".env-secrets").ok(); + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("tts-inference-job".to_string()); + + info!("Hostname: {}", &server_hostname); + + // Bucket stuff (shared) + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + + // Private and Public Buckets + let private_bucket_name = easyenv::get_env_string_required(ENV_PRIVATE_BUCKET_NAME)?; + let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?; + + let private_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &private_bucket_name, + None, + )?; + + let public_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &public_bucket_name, + None, + )?; + + let py_code_directory = easyenv::get_env_string_required(ENV_CODE_DIRECTORY)?; + let py_script_name = easyenv::get_env_string_required(ENV_INFERENCE_SCRIPT_NAME)?; + + let tts_inference_command = TacotronInferenceCommand::new( + &py_code_directory, + &py_script_name, + ); + + let sidecar_hostname = + easyenv::get_env_string_required(ENV_TTS_INFERENCE_SIDECAR_HOSTNAME)?; + + let tts_inference_sidecar_client = + TtsInferenceSidecarClient::new(&sidecar_hostname); + + let temp_directory = easyenv::get_env_string_or_default( + "DOWNLOAD_TEMP_DIR", + DEFAULT_TEMP_DIR); + + let temp_directory = PathBuf::from(temp_directory); + + check_directory_exists(&temp_directory)?; + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + info!("Connecting to database..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(2) + .connect(&db_connection_string) + .await?; + + let common_env = CommonEnv::read_from_env()?; + + info!("Connecting to redis..."); + + let redis_manager = + RedisConnectionManager::new(common_env.redis_0_connection_string.deref())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let persistent_cache_path = easyenv::get_env_string_or_default( + ENV_SEMIPERSISTENT_CACHE_DIR, + "/tmp"); + + let semi_persistent_cache = + SemiPersistentCacheDir::configured_root(&persistent_cache_path); + + info!("Creating pod semi-persistent cache dirs..."); + semi_persistent_cache.create_tts_synthesizer_model_path()?; + semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?; + + let waveglow_vocoder_model_filename = easyenv::get_env_string_or_default( + "TTS_WAVEGLOW_VOCODER_MODEL_FILENAME", "waveglow.pth"); + + let hifigan_vocoder_model_filename = easyenv::get_env_string_or_default( + "TTS_HIFIGAN_VOCODER_MODEL_FILENAME", "hifigan.pth"); + + let hifigan_superres_vocoder_model_filename = easyenv::get_env_string_or_default( + "TTS_HIFIGAN_SUPERRES_VOCODER_MODEL_FILENAME", "hifigan_superres.pth"); + + let sidecar_max_synthesizer_models = easyenv::get_env_num( + "SIDECAR_MAX_SYNTHESIZER_MODELS", 3)?; + + // Set to "0" to always treat low priority the same as high priority + let low_priority_starvation_prevention_every_nth= easyenv::get_env_num( + "LOW_PRIORITY_STARVATION_PREVENTION_EVERY_NTH", 3)?; + + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: MySqlPool is clone/send/sync safe + }; + + let virtual_lfu_cache = SyncVirtualLfuCache::new(sidecar_max_synthesizer_models)?; + + let cache_miss_strategizers = { + let in_memory_strategizer = CacheMissStrategizer::new( + chrono::Duration::milliseconds( + easyenv::get_env_num("MEMORY_MAX_COLD_DURATION_MILLIS", 5_000)?, + ), + chrono::Duration::milliseconds( + easyenv::get_env_num("MEMORY_CACHE_FORGET_DURATION_MILLIS", 60_000)?, + ), + ); + + let on_disk_strategizer = CacheMissStrategizer::new( + chrono::Duration::milliseconds( + easyenv::get_env_num("DISK_MAX_COLD_DURATION_MILLIS", 20_000)?, + ), + chrono::Duration::milliseconds( + easyenv::get_env_num("DISK_CACHE_FORGET_DURATION_MILLIS", 120_000)?, + ), + ); + + SyncMultiCacheMissStrategizer::new( + in_memory_strategizer, + on_disk_strategizer, + ) + }; + + let license_key = easyenv::get_env_string_required("NEWRELIC_API_KEY")?; + let newrelic_client = ClientBuilder::new(&license_key).build().unwrap(); + + let inferencer = Inferencer { + download_temp_directory: temp_directory, + mysql_pool, + redis_pool, + public_bucket_client, + private_bucket_client, + tts_inference_command, + tts_inference_sidecar_client, + newrelic_client, + virtual_model_lfu: virtual_lfu_cache, + cache_miss_strategizers, + bucket_path_unifier: BucketPathUnifier::default_paths(), + semi_persistent_cache, + firehose_publisher, + waveglow_vocoder_model_filename, + hifigan_vocoder_model_filename, + hifigan_superres_vocoder_model_filename, + job_batch_wait_millis: common_env.job_batch_wait_millis, + job_max_attempts: common_env.job_max_attempts as i32, + job_batch_size: common_env.job_batch_size, + no_op_logger_millis: common_env.no_op_logger_millis, + sidecar_max_synthesizer_models, + low_priority_starvation_prevention_every_nth, + }; + + main_loop(inferencer).await; + + Ok(()) +} + +const START_TIMEOUT_MILLIS : u64 = 500; +const INCREASE_TIMEOUT_MILLIS : u64 = 1000; + +async fn main_loop(inferencer: Inferencer) { + let mut error_timeout_millis = START_TIMEOUT_MILLIS; + + let mut noop_logger = NoOpLogger::new(inferencer.no_op_logger_millis as i64); + + let mut span_batch = Vec::new(); + + let mut sort_by_priority = true; + let mut sort_by_priority_count = 0; + + loop { + let num_records = inferencer.job_batch_size; + + // Don't completely starve low-priority jobs + if sort_by_priority_count >= inferencer.low_priority_starvation_prevention_every_nth { + sort_by_priority_count = 0; + sort_by_priority = false; + } + + let maybe_available_jobs = + list_available_tts_inference_jobs( + &inferencer.mysql_pool, + sort_by_priority, + num_records + ).await; + + sort_by_priority = true; + sort_by_priority_count += 1; + + let jobs = match maybe_available_jobs { + Ok(jobs) => jobs, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if jobs.is_empty() { + noop_logger.log_after_awhile(); + + std::thread::sleep(Duration::from_millis(inferencer.job_batch_wait_millis)); + continue; + } + + let batch_result = process_jobs( + &inferencer, + jobs, + ).await; + + let mut spans = match batch_result { + Ok(spans) => spans, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + span_batch.append(&mut spans); + + if span_batch.len() > 50 { + let spans_to_send = span_batch.split_off(0).into(); + inferencer.newrelic_client.send_spans(spans_to_send).await; + } + + error_timeout_millis = START_TIMEOUT_MILLIS; // reset + + std::thread::sleep(Duration::from_millis(inferencer.job_batch_wait_millis)); + } +} + +/// Process a batch of jobs, returning the count of cold-cache skipped jobs. +async fn process_jobs( + inferencer: &Inferencer, + jobs: Vec, +) -> AnyhowResult> { + + let mut batch_spans = Vec::new(); + + for job in jobs.into_iter() { + let model_state_result = ModelState::query_model_and_check_filesystem( + &job, + &inferencer.mysql_pool, + &inferencer.semi_persistent_cache, + &inferencer.virtual_model_lfu, + ).await; + + let model_state = match model_state_result { + Ok(model_state) => model_state, + Err(e) => { + error!("TTS model fetch and state check error: {}, reason: {:?}", &job.model_token, &e); + + let (failure_reason, permanent_failure) = match e { + ModelStateError::ModelNotFound => ("model was not found", true), + ModelStateError::ModelDeleted => ("model has been deleted", true), + ModelStateError::CacheError { .. } => ("internal cache error", false), + ModelStateError::DatabaseError { .. } => ("unknown database error", false), + }; + + + let mut redis = inferencer.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_tts_inference( + &mut redis, + &job.inference_job_token); + + redis_logger.log_status(failure_reason)?; + + if permanent_failure { + warn!("Marking job permanently dead: {} because: {:?}", job.inference_job_token, &e); + + let _r = mark_tts_inference_job_permanently_dead( + &inferencer.mysql_pool, + job.id, + failure_reason + ).await; + } else { + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + inferencer.job_max_attempts + ).await; + } + + continue; + } + }; + + if !model_state.is_downloaded_to_filesystem || !model_state.is_in_memory_cache { + warn!("Model isn't ready (downloaded = {}), (in memory = {})", + model_state.is_downloaded_to_filesystem, + model_state.is_in_memory_cache); + + let maybe_strategy = if !model_state.is_downloaded_to_filesystem { + inferencer.cache_miss_strategizers.disk_cache_miss(&model_state.model_record.model_token) + } else { + inferencer.cache_miss_strategizers.memory_cache_miss(&model_state.model_record.model_token) + }; + + match maybe_strategy { + Err(err) => { + warn!("Failure to process job: {:?}", err); + let failure_reason = ""; + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + inferencer.job_max_attempts + ).await; + continue; + }, + Ok(CacheMissStrategy::WaitOrSkip) => { + // We're going to skip this for now. + // Maybe another worker has a warm cache and can continue. + warn!("Skipping TTS due to cold cache: {} ({})", + model_state.model_record.model_token, + model_state.model_record.title); + continue; + }, + Ok(CacheMissStrategy::Proceed) => {}, // We're going to go ahead... + } + } + + let result = process_job(inferencer, &job, &model_state.model_record).await; + match result { + Ok((span1, span2)) => { + batch_spans.push(span1); + batch_spans.push(span2); + }, + Err(e) => { + warn!("Failure to process job: {:?}", e); + let failure_reason = ""; + let _r = mark_tts_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + inferencer.job_max_attempts + ).await; + } + } + } + + Ok(batch_spans) +} + +#[derive(Deserialize, Default)] +struct FileMetadata { + pub duration_millis: Option, + pub mimetype: Option, + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + +/// We check both of these in one go so that we can reuse the ModelRecord later +/// without another DB query. +struct ModelState { + pub model_record: TtsModelForInferenceRecord, + pub is_downloaded_to_filesystem: bool, + pub is_in_memory_cache: bool, +} + +#[derive(Debug, Clone)] +enum ModelStateError { + ModelNotFound, + ModelDeleted, + CacheError { reason: String }, + DatabaseError { reason: String }, +} + +impl std::fmt::Display for ModelStateError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ModelStateError::ModelNotFound => write!(f, "ModelNotFound"), + ModelStateError::ModelDeleted => write!(f, "ModelDeleted"), + ModelStateError::CacheError { reason} => write!(f, "Cache error: {:?}", reason), + ModelStateError::DatabaseError { reason} => write!(f, "Database error: {:?}", reason), + } + } +} + +impl From for ModelStateError { + fn from(error: TtsModelForInferenceError) -> Self { + match error { + TtsModelForInferenceError::ModelNotFound => ModelStateError::ModelNotFound, + TtsModelForInferenceError::ModelDeleted => ModelStateError::ModelDeleted, + TtsModelForInferenceError::DatabaseError { reason } => ModelStateError::DatabaseError { reason } + } + } +} + +impl std::error::Error for ModelStateError {} + +impl ModelState { + /// Query the model details and see if the model file is on the filesystem in one go. + pub async fn query_model_and_check_filesystem( + job: &AvailableTtsInferenceJob, + mysql_pool: &MySqlPool, + semi_persistent_cache: &SemiPersistentCacheDir, + virtual_cache: &SyncVirtualLfuCache, + ) -> Result { + info!("Looking up TTS model by token: {}", &job.model_token); + + let tts_model = get_tts_model_for_inference( + &mysql_pool, + &job.model_token + ).await?; + + let tts_synthesizer_fs_path = semi_persistent_cache.tts_synthesizer_model_path( + &tts_model.model_token); + + let is_downloaded_to_filesystem = tts_synthesizer_fs_path.exists(); + + let path = tts_synthesizer_fs_path + .to_str() + .ok_or(ModelStateError::CacheError { reason: "could not make path".to_string() })? + .to_string(); + + let is_in_memory_cache = virtual_cache.in_cache(&path) + .map_err(|e| ModelStateError::CacheError { reason: format!("Model cache error: {:?}", e) })?; + + Ok(Self { + model_record: tts_model, + is_downloaded_to_filesystem, + is_in_memory_cache, + }) + } +} + +async fn process_job( + inferencer: &Inferencer, + job: &AvailableTtsInferenceJob, + model_record: &TtsModelForInferenceRecord, +) -> AnyhowResult<(Span, Span)> { + + // TODO 1. Mark processing (DONE) + + // TODO 2. Check if vocoder model is downloaded / download to stable location (DONE) + + // TODO 3. Query model by token. (DONE) + // TODO 4. Check if model is downloaded, otherwise download to stable location (DONE) + + // TODO 5. Write text to file + // TODO 6. Process Inference + + // TODO 7. Upload Result + // TODO 8. Save record + // TODO 9. Mark job done + + let start = Instant::now(); + + let span_id = generate_random_uuid(); + let trace_id = generate_random_uuid(); + let maybe_user_token = job.maybe_creator_user_token.as_deref().unwrap_or(""); + + let mut job_iteration_span = Span::new(&span_id, &trace_id, get_timestamp_millis()) + .name("single job execution") + .attribute("user_token", maybe_user_token) + .service_name("tts-inference-job"); + + let span_id = generate_random_uuid(); + let trace_id = generate_random_uuid(); + + let created_at_timestamp = (job.created_at.timestamp() as u64) * 1000; + + let mut since_creation_span = Span::new(&span_id, &trace_id, created_at_timestamp) + .name("job since creation") + .attribute("user_token", maybe_user_token) + .service_name("tts-inference-job"); + + let mut redis = inferencer.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_tts_inference( + &mut redis, + &job.inference_job_token); + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + let lock_acquired = + mark_tts_inference_job_pending_and_grab_lock(&inferencer.mysql_pool, job.id).await?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {:?}", &job.id); + let duration = start.elapsed(); + + since_creation_span.set_attribute("status", "failure"); + since_creation_span.set_duration(duration); + + job_iteration_span.set_attribute("status", "failure"); + job_iteration_span.set_duration(duration); + + return Ok((since_creation_span, job_iteration_span)); + } + + // ==================== CONFIRM OR DOWNLOAD WAVEGLOW VOCODER MODEL ==================== // + + let waveglow_vocoder_model_filename = inferencer.waveglow_vocoder_model_filename.clone(); + let waveglow_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_pretrained_vocoder_model_path(&waveglow_vocoder_model_filename); + + if !waveglow_vocoder_model_fs_path.exists() { + warn!("Waveglow vocoder model file does not exist: {:?}", &waveglow_vocoder_model_fs_path); + + redis_logger.log_status("downloading vocoder (1 of 3)")?; + + let waveglow_vocoder_model_object_path = inferencer.bucket_path_unifier + .tts_pretrained_vocoders_path(&waveglow_vocoder_model_filename); + + info!("Download waveglow vocoder from bucket path: {:?}", &waveglow_vocoder_model_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &waveglow_vocoder_model_object_path, + &waveglow_vocoder_model_fs_path + ).await?; + + info!("Downloaded waveglow vocoder model from bucket!"); + } + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== // + + let hifigan_vocoder_model_filename = inferencer.hifigan_vocoder_model_filename.clone(); + let hifigan_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename); + + if !hifigan_vocoder_model_fs_path.exists() { + warn!("Hifigan vocoder model file does not exist: {:?}", &hifigan_vocoder_model_fs_path); + + redis_logger.log_status("downloading vocoder (2 of 3)")?; + + let hifigan_vocoder_model_object_path = inferencer.bucket_path_unifier + .tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename); + + info!("Download hifigan vocoder from bucket path: {:?}", &hifigan_vocoder_model_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &hifigan_vocoder_model_object_path, + &hifigan_vocoder_model_fs_path + ).await?; + + info!("Downloaded hifigan vocoder model from bucket!"); + } + + // ==================== CONFIRM OR DOWNLOAD HIFIGAN (SUPERRES) VOCODER MODEL ==================== // + + let hifigan_superres_vocoder_model_filename = inferencer.hifigan_superres_vocoder_model_filename.clone(); + let hifigan_superres_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_superres_vocoder_model_filename); + + if !hifigan_superres_vocoder_model_fs_path.exists() { + warn!("Hifigan superres vocoder model file does not exist: {:?}", &hifigan_superres_vocoder_model_fs_path); + + redis_logger.log_status("downloading vocoder (3 of 3)")?; + + let hifigan_superres_vocoder_model_object_path = inferencer.bucket_path_unifier + .tts_pretrained_vocoders_path(&hifigan_superres_vocoder_model_filename); + + info!("Download hifigan superres vocoder from bucket path: {:?}", &hifigan_superres_vocoder_model_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &hifigan_superres_vocoder_model_object_path, + &hifigan_superres_vocoder_model_fs_path + ).await?; + + info!("Downloaded hifigan superres vocoder model from bucket!"); + } + +// // ==================== LOOK UP TTS SYNTHESIZER RECORD (WHICH CONTAINS ITS BUCKET PATH) ==================== // +// +// info!("Looking up TTS model by token: {}", &job.model_token); +// +// let query_result = get_tts_model_by_token( +// &inferencer.mysql_pool, +// &job.model_token).await?; +// +// let tts_model = match query_result { +// Some(model) => model, +// None => { +// warn!("TTS model not found: {}", &job.model_token); +// return Err(anyhow!("Model not found!")) +// }, +// }; + + // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== // + + // TODO: Let's just put paths in the db + // TODO: We'll probably need to LRU cache these. + + let tts_synthesizer_fs_path = inferencer.semi_persistent_cache.tts_synthesizer_model_path( + &model_record.model_token); + + if !tts_synthesizer_fs_path.exists() { + info!("TTS synthesizer model file does not exist: {:?}", &tts_synthesizer_fs_path); + + redis_logger.log_status("downloading synthesizer")?; + + let tts_synthesizer_object_path = inferencer.bucket_path_unifier + .tts_synthesizer_path(&model_record.private_bucket_hash); + + info!("Download from template media path: {:?}", &tts_synthesizer_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &tts_synthesizer_object_path, + &tts_synthesizer_fs_path + ).await?; + + info!("Downloaded template media from bucket!"); + } + + // ==================== Preprocess text ==================== // + + let cleaned_inference_text = clean_symbols(&job.raw_inference_text); + + // ==================== WRITE TEXT TO FILE ==================== // + + let temp_dir = format!("temp_tts_inference_{}", job.id.0); + let temp_dir = TempDir::new(&temp_dir)?; // NB: Exists until it goes out of scope. + + let text_input_fs_path = temp_dir.path().join("inference_input.txt"); + + std::fs::write(&text_input_fs_path, &cleaned_inference_text)?; + + // ==================== RUN INFERENCE ==================== // + + redis_logger.log_status("running inference")?; + + // TODO: Fix this. + let maybe_unload_model_path = inferencer + .virtual_model_lfu + .insert_returning_replaced(tts_synthesizer_fs_path.to_str().unwrap_or(""))?; + + if let Some(model_path) = maybe_unload_model_path.as_deref() { + warn!("Remove model from LFU cache: {:?}", model_path); + } + + let output_audio_fs_path = temp_dir.path().join("output.wav"); + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json"); + + info!("Running TTS inference..."); + + info!("Expected output audio filename: {:?}", &output_audio_fs_path); + info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path); + info!("Expected output metadata filename: {:?}", &output_metadata_fs_path); + + if let Some(model_path) = maybe_unload_model_path.as_deref() { + warn!("Unload model from sidecar: {:?}", &model_path); + } + + //inferencer.tts_inference_command.execute( + // &tts_synthesizer_fs_path, + // &tts_vocoder_model_fs_path, + // &text_input_fs_path, + // &output_audio_fs_path, + // &output_spectrogram_fs_path, + // &output_metadata_fs_path, + // false, + //)?; + + let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution; + if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() { + pretrained_vocoder = VocoderType::from_str(default_vocoder)?; + } + + info!("With pretrained vocoder: {:?}", pretrained_vocoder); + + inferencer.tts_inference_sidecar_client.request_inference( + &cleaned_inference_text, + &tts_synthesizer_fs_path, + pretrained_vocoder, + &hifigan_vocoder_model_fs_path, + &hifigan_superres_vocoder_model_fs_path, + &waveglow_vocoder_model_fs_path, + &output_audio_fs_path, + &output_spectrogram_fs_path, + &output_metadata_fs_path, + maybe_unload_model_path, + ).await?; + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that output files exist..."); + + check_file_exists(&output_audio_fs_path)?; + check_file_exists(&output_spectrogram_fs_path)?; + check_file_exists(&output_metadata_fs_path)?; + + let file_metadata = read_metadata_file(&output_metadata_fs_path)?; + + safe_delete_temp_file(&output_metadata_fs_path); + + // ==================== UPLOAD AUDIO TO BUCKET ==================== // + + redis_logger.log_status("uploading result")?; + + let audio_result_object_path = inferencer.bucket_path_unifier.tts_inference_wav_audio_output_path( + &job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Audio destination bucket path: {:?}", &audio_result_object_path); + + info!("Uploading audio..."); + + inferencer.public_bucket_client.upload_filename_with_content_type( + &audio_result_object_path, + &output_audio_fs_path, + "audio/wav") + .await?; + + safe_delete_temp_file(&output_audio_fs_path); + + // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // + + let spectrogram_result_object_path = inferencer.bucket_path_unifier.tts_inference_spectrogram_output_path( + &job.uuid_idempotency_token); // TODO: Don't use this! + + info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); + + info!("Uploading spectrogram..."); + + inferencer.public_bucket_client.upload_filename_with_content_type( + &spectrogram_result_object_path, + &output_spectrogram_fs_path, + "application/json") + .await?; + + safe_delete_temp_file(&output_spectrogram_fs_path); + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + let text_hash = hash_string_sha2(&cleaned_inference_text)?; + + info!("Saving tts inference record..."); + let (id, inference_result_token) = insert_tts_result( + &inferencer.mysql_pool, + job, + &text_hash, + pretrained_vocoder, + &audio_result_object_path, + &spectrogram_result_object_path, + file_metadata.file_size_bytes, + file_metadata.duration_millis.unwrap_or(0)) + .await?; + + info!("Marking job complete..."); + mark_tts_inference_job_done( + &inferencer.mysql_pool, + job.id, + true, + Some(&inference_result_token) + ).await?; + + info!("TTS Done. Original text was: {}", &job.raw_inference_text); + + inferencer.firehose_publisher.tts_inference_finished( + job.maybe_creator_user_token.as_deref(), + &model_record.model_token, + &inference_result_token) + .await + .map_err(|e| { + error!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + redis_logger.log_status("done")?; + + info!("Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}", + job.id, id); + + let duration = start.elapsed(); + + since_creation_span.set_attribute("status", "success"); + since_creation_span.set_duration(duration); + + job_iteration_span.set_attribute("status", "success"); + job_iteration_span.set_duration(duration); + + Ok((since_creation_span, job_iteration_span)) +} + +fn get_timestamp_millis() -> u64 { + SystemTime::now() + .duration_since(UNIX_EPOCH) + .map(|d| d.as_millis() as u64) + .unwrap_or(0) +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/util/buckets/bucket_client.rs b/crates/service/storyteller_web/src/util/buckets/bucket_client.rs new file mode 100644 index 0000000000..7f06de635b --- /dev/null +++ b/crates/service/storyteller_web/src/util/buckets/bucket_client.rs @@ -0,0 +1,215 @@ +use anyhow::anyhow; +use anyhow::bail; +use log::info; +use log::warn; +use s3::bucket::Bucket; +use s3::creds::Credentials; +use s3::region::Region; +use std::io::Read; +use std::path::{PathBuf, Path}; +use std::str::FromStr; +use tokio::fs::File; +use tokio::io::AsyncBufReadExt; +use tokio::io::{AsyncBufRead, BufReader, AsyncReadExt}; + +#[derive(Clone)] +pub struct BucketClient { + bucket: Bucket, + + /// If set, put all files under this root path. + optional_bucket_root: Option, +} + +impl BucketClient { + + pub fn create( + access_key: &str, + secret_key: &str, + region_name: &str, + bucket_name: &str, + optional_bucket_root: Option<&str>) -> anyhow::Result + { + let credentials = Credentials { + access_key: Some(access_key.to_string()), + secret_key: Some(secret_key.to_string()), + security_token: None, + session_token: None, + }; + + // NB: The GCS buckets aren't supported by default. + let region = Region::Custom { + region: region_name.to_owned(), + endpoint: "https://storage.googleapis.com".to_owned(), + }; + let mut bucket = Bucket::new(&bucket_name, region, credentials)?; + + bucket.set_path_style(); + bucket.set_subdomain_style(); + + let optional_bucket_root = optional_bucket_root.map(|s| s.to_string()); + + Ok(Self { + bucket, + optional_bucket_root, + }) + } + + fn get_rooted_object_name(&self, object_name: &str) -> String { + match &self.optional_bucket_root { + None => object_name.to_string(), + Some(root) => format!("{}/{}", root, object_name), + } + } + + pub async fn upload_file(&self, object_name: &str, bytes: &[u8]) -> anyhow::Result<()> { + info!("Filename for bucket: {}", object_name); + + let object_name = self.get_rooted_object_name(object_name); + info!("Rooted filename for bucket: {}", object_name); + + let (body_bytes, code) = self.bucket.put_object(&object_name, bytes).await?; + + info!("upload code: {}", code); + + if code != 200 { + let body = String::from_utf8(body_bytes)?; + warn!("upload body: {}", body); + } + + Ok(()) + } + + pub async fn upload_file_with_content_type(&self, object_name: &str, bytes: &[u8], content_type: &str) -> anyhow::Result<()> { + info!("Filename for bucket: {}", object_name); + + let object_name = self.get_rooted_object_name(object_name); + info!("Rooted filename for bucket: {}", object_name); + + let (body_bytes, code) = self.bucket.put_object_with_content_type(&object_name, bytes, content_type).await?; + + info!("upload code: {}", code); + + if code != 200 { + let body = String::from_utf8(body_bytes)?; + warn!("upload body: {}", body); + } + + Ok(()) + } + + // NB: New version has blocking client rather than blocking calls. + // pub fn upload_file_blocking(&self, object_name: &str, bytes: &[u8]) -> anyhow::Result<()> { + // info!("Filename for bucket: {}", object_name); + // + // let (_, code) = self.bucket.put_object_blocking(object_name, bytes)?; + // + // info!("upload code: {}", code); + // + // Ok(()) + // } + + pub async fn upload_filename, Q: AsRef>( + &self, + object_path: P, + filename: Q + ) -> anyhow::Result<()> { + let object_path_str = object_path.as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("could not convert object path to string"))?; + + // TODO: does a newer version of this crate handle streaming/buffering file contents? + let mut file = File::open(filename).await?; + let mut buffer : Vec = Vec::new(); + file.read_to_end(&mut buffer).await?; + + info!("Uploading..."); + + self.upload_file(&object_path_str, &buffer).await + } + + pub async fn upload_filename_with_content_type, Q: AsRef>( + &self, + object_path: P, + filename: Q, + content_type: &str + ) -> anyhow::Result<()> { + let object_path_str = object_path.as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("could not convert object path to string"))?; + + // TODO: does a newer version of this crate handle streaming/buffering file contents? + let mut file = File::open(filename).await?; + let mut buffer : Vec = Vec::new(); + file.read_to_end(&mut buffer).await?; + + info!("Uploading with content type..."); + + self.upload_file_with_content_type(&object_path_str, &buffer, content_type).await + } + + // NB: New version has blocking client rather than blocking calls. + // pub fn upload_filename_blocking(&self, object_name: &str, filename: &Path) -> anyhow::Result<()> { + // // TODO: does a newer version of this crate handle streaming/buffering file contents? + // let mut file = File::open(filename)?; + // let mut buffer : Vec = Vec::new(); + // file.read_to_end(&mut buffer)?; + // + // self.upload_file_blocking(object_name, &buffer) + // } + + pub async fn download_file(&self, path: &str) -> anyhow::Result> { + info!("downloading from bucket: {}", path); + let (bytes, code) = self.bucket.get_object(path).await?; + + match code { + 404 => bail!("File not found in bucket: {}", path), + _ => {}, + } + + info!("download code: {}", code); + Ok(bytes) + } + + pub async fn download_file_to_disk, Q: AsRef>( + &self, + object_path: P, + filesystem_path: Q, + ) -> anyhow::Result<()> { + let object_path_str = object_path.as_ref() + .to_str() + .map(|s| s.to_string()) + .ok_or(anyhow!("could not convert object path to string"))?; + + info!("downloading from bucket: {:?}", &object_path_str); + + let mut output_file = tokio::fs::File::create(filesystem_path).await?; + + let status_code = self.bucket.get_object_stream(&object_path_str, &mut output_file).await?; + + match status_code { + 404 => bail!("File not found in bucket: {}", &object_path_str), + _ => {}, + } + + info!("download code: {}", status_code); + + Ok(()) + } + + + // NB: New version has blocking client rather than blocking calls. + // pub fn download_file_blocking(&self, path: &str) -> anyhow::Result> { + // info!("downloading from bucket: {}", path); + // let (bytes, code) = self.bucket.get_object_blocking(path)?; + // + // match code { + // 404 => bail!("File not found in bucket: {}", path), + // _ => {}, + // } + // + // info!("download code: {}", code); + // Ok(bytes) + // } +} diff --git a/crates/service/storyteller_web/src/util/buckets/bucket_path_unifier.rs b/crates/service/storyteller_web/src/util/buckets/bucket_path_unifier.rs new file mode 100644 index 0000000000..a5e935f1d9 --- /dev/null +++ b/crates/service/storyteller_web/src/util/buckets/bucket_path_unifier.rs @@ -0,0 +1,308 @@ +use container_common::anyhow_result::AnyhowResult; +use crate::util::buckets::bucket_paths::hash_to_bucket_path; +use std::path::PathBuf; + +/// This is designed to make it centrally configurable where +/// different types of objects are stored. +pub struct BucketPathUnifier { + // TTS + pub tts_pretrained_vocoder_model_root: PathBuf, + pub tts_synthesizer_model_root: PathBuf, + pub tts_inference_output_root: PathBuf, + // W2L + pub user_uploaded_w2l_templates_root: PathBuf, + pub user_uploaded_audio_for_w2l_root: PathBuf, + pub w2l_inference_output_root: PathBuf, + pub w2l_model_root: PathBuf, + pub w2l_end_bump_root: PathBuf, +} + +impl BucketPathUnifier { + + // TODO + //pub fn from_env_vars() -> AnyhowResult { + // Ok(Self { + // user_uploaded_audio_for_w2l_root: , + // user_uploaded_w2l_templates_root: , + // }) + //} + + pub fn default_paths() -> Self { + Self { + // TTS + tts_pretrained_vocoder_model_root: PathBuf::from("/tts_pretrained_vocoders"), + tts_synthesizer_model_root: PathBuf::from("/user_uploaded_tts_synthesizers"), + tts_inference_output_root: PathBuf::from("/tts_inference_output"), + // W2L + user_uploaded_audio_for_w2l_root: PathBuf::from("/user_uploaded_w2l_audio"), + user_uploaded_w2l_templates_root: PathBuf::from("/user_uploaded_w2l_templates"), + w2l_inference_output_root: PathBuf::from("/w2l_inference_output"), + w2l_model_root: PathBuf::from("/w2l_pretrained_models"), + w2l_end_bump_root: PathBuf::from("/w2l_end_bumps"), + } + } + + // ==================== TTS MODELS (SYNTHESIZER + VOCODER) ==================== // + + // TTS pretrained vocoder models. + // For now this will be limited, but once they're user-uploadable, we'll add more. + pub fn tts_pretrained_vocoders_path(&self, tts_vocoder_model_name: &str) -> PathBuf { + self.tts_pretrained_vocoder_model_root.join(tts_vocoder_model_name) + } + + pub fn tts_synthesizer_path(&self, tts_synthesizer_file_hash: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(tts_synthesizer_file_hash); + let model_filename = format!("{}.pt", &tts_synthesizer_file_hash); + + self.tts_synthesizer_model_root + .join(hashed_path) + .join(model_filename) + } + + pub fn tts_zipped_synthesizer_path(&self, tts_synthesizer_file_hash: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(tts_synthesizer_file_hash); + let model_filename = format!("{}.zip", &tts_synthesizer_file_hash); + + self.tts_synthesizer_model_root + .join(hashed_path) + .join(model_filename) + } + + // ==================== TTS INFERENCE OUTPUT ==================== // + + /// This should include the string "vocodes" for downloaders. + pub fn tts_inference_wav_audio_output_path(&self, tts_inference_output_uuid: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(tts_inference_output_uuid); + let audio_filename = format!("vocodes_{}.wav", &tts_inference_output_uuid); + + self.tts_inference_output_root + .join(hashed_path) + .join(audio_filename) + } + + pub fn tts_inference_spectrogram_output_path(&self, tts_inference_output_uuid: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(tts_inference_output_uuid); + let json_filename = format!("{}.json", &tts_inference_output_uuid); + + self.tts_inference_output_root + .join(hashed_path) + .join(json_filename) + } + + // ==================== W2L STATIC RESOURCES ==================== // + + // W2L pretrained models. There are only two. + pub fn w2l_pretrained_models_path(&self, w2l_model_name: &str) -> PathBuf { + self.w2l_model_root.join(w2l_model_name) + } + + // W2L "end bumps" are videos added at the end. + pub fn end_bump_video_for_w2l_path(&self, end_bump_filename: &str) -> PathBuf { + self.w2l_end_bump_root.join(end_bump_filename) + } + + // ==================== W2L USER-UPLOADED RESOURCES ==================== // + + // The video or images uploaded as templates + // eg. /user_uploaded_w2l_templates/1/5/1/151a[...60]... + pub fn media_templates_for_w2l_path(&self, template_file_hash: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(template_file_hash); + self.user_uploaded_w2l_templates_root + .join(hashed_path) + .join(template_file_hash) + } + + // These share the same directory as the uploaded w2l template media. + // eg. /user_uploaded_w2l_templates/1/5/1/151a[...60]_detected_faces.pickle + pub fn precomputed_faces_for_w2l_path(&self, template_file_hash: &str) -> PathBuf { + let faces_filename = format!("{}_detected_faces.pickle", &template_file_hash); + let hashed_path = Self::hashed_directory_path(template_file_hash); + + self.user_uploaded_w2l_templates_root + .join(hashed_path) + .join(faces_filename) + } + + // User-uploaded audio. + // eg. /user_uploaded_w2l_audio/0/0/b/00bcc7a4-bdf5-43a5-9603-a15ca780d866 + pub fn user_audio_for_w2l_inference_path(&self, audio_uuid: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(audio_uuid); + self.user_uploaded_audio_for_w2l_root + .join(hashed_path) + .join(audio_uuid) + } + + // ==================== W2L INFERENCE OUTPUT ==================== // + + // W2L inference output videos + pub fn w2l_inference_video_output_path(&self, w2l_inference_job_token: &str) -> PathBuf { + // NB: We don't want colons from the token in the filename. + if w2l_inference_job_token.contains(":") { + if let Some((token_type, token)) = w2l_inference_job_token.split_once(":") { + let video_filename = w2l_inference_job_token.replace(":", ""); + let video_filename = format!("vocodes_video_{}.mp4", video_filename); + + let hashed_path = Self::hashed_directory_path(token); + let hashed_path = hashed_path.to_lowercase(); + + return self.w2l_inference_output_root + .join(hashed_path) + .join(video_filename); + } + } + + let video_filename = w2l_inference_job_token.replace(":", ""); + let video_filename = format!("vocodes_video_{}.mp4", video_filename); + + let hashed_path = Self::hashed_directory_path(w2l_inference_job_token); + + self.w2l_inference_output_root + .join(hashed_path) + .join(video_filename) + } + + pub fn hashed_directory_path(file_hash: &str) -> String { + match file_hash.len() { + 0 | 1=> "".to_string(), + 2 => format!("{}/", &file_hash[0..1]), + 3 => format!("{}/{}/", &file_hash[0..1], &file_hash[1..2]), + _ => format!("{}/{}/{}/", &file_hash[0..1], &file_hash[1..2], &file_hash[2..3]), + } + } +} + +#[cfg(test)] +mod tests { + use crate::util::buckets::bucket_path_unifier::BucketPathUnifier; + use std::path::PathBuf; + + fn get_instance() -> BucketPathUnifier { + BucketPathUnifier { + // TTS + tts_pretrained_vocoder_model_root: PathBuf::from("/test_path_tts_vocoders"), + tts_synthesizer_model_root: PathBuf::from("/test_path_synthesizers"), + tts_inference_output_root: PathBuf::from("/test_path_tts_output"), + // W2L + user_uploaded_w2l_templates_root: PathBuf::from("/test_path_w2l_templates"), + user_uploaded_audio_for_w2l_root: PathBuf::from("/test_path_w2l_audio"), + w2l_inference_output_root: PathBuf::from("/test_path_w2l_output"), + w2l_model_root: PathBuf::from("/test_path_w2l_pretrained_models"), + w2l_end_bump_root: PathBuf::from("/test_path_w2l_end_bumps"), + } + } + + #[test] + fn test_tts_pretrained_vocoders_path() { + let paths = get_instance(); + assert_eq!(paths.tts_pretrained_vocoders_path("melgan.pth").to_str().unwrap(), + "/test_path_tts_vocoders/melgan.pth"); + } + + #[test] + fn test_tts_synthesizer_path() { + let paths = get_instance(); + assert_eq!(paths.tts_synthesizer_path("foobar").to_str().unwrap(), + "/test_path_synthesizers/f/o/o/foobar.pt"); + } + + #[test] + fn test_tts_zipped_synthesizer_path() { + let paths = get_instance(); + assert_eq!(paths.tts_zipped_synthesizer_path("foobar").to_str().unwrap(), + "/test_path_synthesizers/f/o/o/foobar.zip"); + } + + #[test] + fn test_tts_inference_wav_audio_output_path() { + let paths = get_instance(); + assert_eq!(paths.tts_inference_wav_audio_output_path("foobar").to_str().unwrap(), + "/test_path_tts_output/f/o/o/vocodes_foobar.wav"); + } + + #[test] + fn test_tts_inference_spectrogram_output_path() { + let paths = get_instance(); + assert_eq!(paths.tts_inference_spectrogram_output_path("foobar").to_str().unwrap(), + "/test_path_tts_output/f/o/o/foobar.json"); + } + + #[test] + fn test_w2l_pretrained_models_path() { + let paths = get_instance(); + assert_eq!(paths.w2l_pretrained_models_path("model.pth").to_str().unwrap(), + "/test_path_w2l_pretrained_models/model.pth"); + } + + #[test] + fn test_end_bump_video_for_w2l_path() { + let paths = get_instance(); + assert_eq!(paths.end_bump_video_for_w2l_path("logo.mp4").to_str().unwrap(), + "/test_path_w2l_end_bumps/logo.mp4"); + } + + #[test] + fn test_user_audio_for_w2l_inference_path() { + let paths = get_instance(); + assert_eq!(paths.user_audio_for_w2l_inference_path("foobar").to_str().unwrap(), + "/test_path_w2l_audio/f/o/o/foobar"); + } + + #[test] + fn test_media_templates_for_w2l_path() { + let paths = get_instance(); + assert_eq!(paths.media_templates_for_w2l_path("foobar").to_str().unwrap(), + "/test_path_w2l_templates/f/o/o/foobar"); + } + + #[test] + fn test_precomputed_faces_for_w2l_path() { + let paths = get_instance(); + assert_eq!(paths.precomputed_faces_for_w2l_path("foobar").to_str().unwrap(), + "/test_path_w2l_templates/f/o/o/foobar_detected_faces.pickle"); + } + + #[test] + fn test_w2l_inference_video_output_path() { + let paths = get_instance(); + + // Case 1: Tokens without a "token type" + assert_eq!(paths.w2l_inference_video_output_path("foobar").to_str().unwrap(), + "/test_path_w2l_output/f/o/o/vocodes_video_foobar.mp4"); + + // Case 2: Tokens with a "token type" + // Note: that it removes the token type from dir path and handles the colon: + assert_eq!(paths.w2l_inference_video_output_path("type:abcdef").to_str().unwrap(), + "/test_path_w2l_output/a/b/c/vocodes_video_typeabcdef.mp4"); + + // It also handles capitalization + assert_eq!(paths.w2l_inference_video_output_path("TYPE:ABCDEF").to_str().unwrap(), + "/test_path_w2l_output/a/b/c/vocodes_video_TYPEABCDEF.mp4"); + } + + #[test] + fn hashed_directory_path_length_zero() { + assert_eq!(&BucketPathUnifier::hashed_directory_path(""), ""); + } + + #[test] + fn hashed_directory_path_length_one() { + assert_eq!(&BucketPathUnifier::hashed_directory_path("a"), ""); + } + + #[test] + fn hashed_directory_path_length_two() { + assert_eq!(&BucketPathUnifier::hashed_directory_path("ab"), "a/"); + } + + #[test] + fn hashed_directory_path_length_three() { + assert_eq!(&BucketPathUnifier::hashed_directory_path("abc"), "a/b/"); + } + + #[test] + fn hashed_directory_path_length_more() { + assert_eq!(&BucketPathUnifier::hashed_directory_path("abcd"), "a/b/c/"); + assert_eq!(&BucketPathUnifier::hashed_directory_path("abcde"), "a/b/c/"); + assert_eq!(&BucketPathUnifier::hashed_directory_path("abcdef01234"), "a/b/c/"); + } +} diff --git a/crates/service/storyteller_web/src/util/buckets/bucket_paths.rs b/crates/service/storyteller_web/src/util/buckets/bucket_paths.rs new file mode 100644 index 0000000000..eadedbf56c --- /dev/null +++ b/crates/service/storyteller_web/src/util/buckets/bucket_paths.rs @@ -0,0 +1,51 @@ +use anyhow::anyhow; +use anyhow::bail; +use container_common::anyhow_result::AnyhowResult; +use std::path::PathBuf; + +pub fn hash_to_bucket_path(file_hash: &str, + optional_root_directory: Option<&str>) + -> AnyhowResult +{ + if file_hash.len() < 4 { + bail!("File length is too short"); + } + + let root_directory = optional_root_directory + .map(|root| PathBuf::from(root.to_string())) + .unwrap_or(PathBuf::from("/")); + + let file_path = format!( + "{}/{}/{}/{}", + &file_hash[0..1], + &file_hash[1..2], + &file_hash[2..3], + &file_hash, + ); + + let maybe_path = root_directory + .join(file_path) + .to_str() + .map(|s| s.to_string()); + + maybe_path.ok_or(anyhow!("Path could not be converted")) +} + +#[cfg(test)] +mod tests { + use crate::util::buckets::bucket_paths::hash_to_bucket_path; + + #[test] + fn test_bucket_without_root() { + assert_eq!( + hash_to_bucket_path("abcdefg", None).unwrap(), + "/a/b/c/abcdefg".to_string()); + } + + #[test] + fn test_bucket_with_root() { + assert_eq!( + hash_to_bucket_path("abcdefg", Some("/root")).unwrap(), + "/root/a/b/c/abcdefg".to_string()); + } +} diff --git a/crates/service/storyteller_web/src/util/buckets/mod.rs b/crates/service/storyteller_web/src/util/buckets/mod.rs new file mode 100644 index 0000000000..d4d95f081a --- /dev/null +++ b/crates/service/storyteller_web/src/util/buckets/mod.rs @@ -0,0 +1,3 @@ +pub mod bucket_client; +pub mod bucket_path_unifier; +pub mod bucket_paths; diff --git a/crates/service/storyteller_web/src/util/caching/mod.rs b/crates/service/storyteller_web/src/util/caching/mod.rs new file mode 100644 index 0000000000..cffdc1415f --- /dev/null +++ b/crates/service/storyteller_web/src/util/caching/mod.rs @@ -0,0 +1 @@ +pub mod single_item_ttl_cache; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/util/caching/single_item_ttl_cache.rs b/crates/service/storyteller_web/src/util/caching/single_item_ttl_cache.rs new file mode 100644 index 0000000000..a3ce6117c1 --- /dev/null +++ b/crates/service/storyteller_web/src/util/caching/single_item_ttl_cache.rs @@ -0,0 +1,46 @@ +use anyhow::bail; +use container_common::anyhow_result::AnyhowResult; +use lru_time_cache::LruCache; +use std::sync::{Arc, Mutex, PoisonError, MutexGuard}; +use std::time::Duration; + +/// NB: There's only ONE ITEM of ONE TYPE in the cache. We can use a single key. +const CACHE_KEY : &'static str = "ITEM"; + +/// Store a single payload in the cache +/// There's only ONE ITEM of ONE TYPE in the cache. +/// This is essentially a singleton cache with expiry. +#[derive(Clone)] +pub struct SingleItemTtlCache { + cache: Arc>>, +} + +impl SingleItemTtlCache { + pub fn create_with_duration(time_to_live: Duration) -> Self { + let cache = LruCache::with_expiry_duration(time_to_live); + let cache = Arc::new(Mutex::new(cache)); + Self { + cache, + } + } + + pub fn copy_without_bump_if_unexpired(&self) -> AnyhowResult> { + let maybe_copy = match self.cache.lock() { + Err(e) => bail!("could not unlock mutex to read: {:?}", e), + Ok(mut cache) => { + cache.peek(CACHE_KEY).map(|inner| inner.clone()) + }, + }; + Ok(maybe_copy) + } + + pub fn store_copy(&self, item: &T) -> AnyhowResult<()> { + match self.cache.lock() { + Err(e) => bail!("could not unlock mutex to write: {:?}", e), + Ok(mut cache) => { + cache.insert(CACHE_KEY.into(), item.clone()); + }, + }; + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/util/delete_role_disambiguation.rs b/crates/service/storyteller_web/src/util/delete_role_disambiguation.rs new file mode 100644 index 0000000000..7519a71373 --- /dev/null +++ b/crates/service/storyteller_web/src/util/delete_role_disambiguation.rs @@ -0,0 +1,59 @@ +//! This is meant to disambiguate whether the deletion should occur as a user or moderator +//! for soft-deleted records. All such requests have a similar shape, and this was tested +//! with a truth table. + +#[derive(Copy, Clone, PartialEq, Eq, Debug)] +pub enum DeleteRole { + ErrorDoNotDelete, + AsUser, + AsMod, +} + +pub fn delete_role_disambiguation( + user_is_mod: bool, + user_is_author: bool, + as_mod_flag: Option +) -> DeleteRole { + if !user_is_mod && !user_is_author { + return DeleteRole::ErrorDoNotDelete; + } + // NB: Explored this with a truth table. + let as_mod_flag_value = as_mod_flag.unwrap_or(false); + let delete_as_mod = user_is_mod && !(user_is_author && as_mod_flag.is_some() && !as_mod_flag_value); + if delete_as_mod { DeleteRole::AsMod } else { DeleteRole::AsUser } +} + +#[cfg(test)] +mod tests { + use crate::util::delete_role_disambiguation::{delete_role_disambiguation, DeleteRole}; + + #[test] + fn test_delete_as_mod() { + // User is neither author nor mod. + // Do not delete! + assert_eq!(DeleteRole::ErrorDoNotDelete, delete_role_disambiguation(false, false, None)); + assert_eq!(DeleteRole::ErrorDoNotDelete, delete_role_disambiguation(false, false, Some(false))); + assert_eq!(DeleteRole::ErrorDoNotDelete, delete_role_disambiguation(false, false, Some(true))); + + // User is author. + // Deleted as a user. + assert_eq!(DeleteRole::AsUser, delete_role_disambiguation(false, true, None)); + assert_eq!(DeleteRole::AsUser, delete_role_disambiguation(false, true, Some(false))); + assert_eq!(DeleteRole::AsUser, delete_role_disambiguation(false, true, Some(true))); + + // User is mod. + // Deleted as a mod. + assert_eq!(DeleteRole::AsMod, delete_role_disambiguation(true, false, None)); + assert_eq!(DeleteRole::AsMod, delete_role_disambiguation(true, false, Some(false))); + assert_eq!(DeleteRole::AsMod, delete_role_disambiguation(true, false, Some(true))); + + // User is mod and author. + // Deleting as a user. + assert_eq!(DeleteRole::AsUser, delete_role_disambiguation(true, true, Some(false))); + + // User is mod and author. + // Deleting as a mod. + assert_eq!(DeleteRole::AsMod, delete_role_disambiguation(true, true, None)); + assert_eq!(DeleteRole::AsMod, delete_role_disambiguation(true, true, Some(true))); + } +} diff --git a/crates/service/storyteller_web/src/util/email_to_gravatar.rs b/crates/service/storyteller_web/src/util/email_to_gravatar.rs new file mode 100644 index 0000000000..26f4f0fea8 --- /dev/null +++ b/crates/service/storyteller_web/src/util/email_to_gravatar.rs @@ -0,0 +1,12 @@ +use md5::{Md5, Digest}; + +pub fn email_to_gravatar(email_address: &str) -> String { + let email = email_address.trim().to_lowercase(); + + let mut hasher = Md5::new(); + hasher.update(email); + let hash = hasher.finalize(); + let gravatar_hash = format!("{:x}", hash); + + gravatar_hash +} diff --git a/crates/service/storyteller_web/src/util/encrypted_sort_id.rs b/crates/service/storyteller_web/src/util/encrypted_sort_id.rs new file mode 100644 index 0000000000..56ce681e2b --- /dev/null +++ b/crates/service/storyteller_web/src/util/encrypted_sort_id.rs @@ -0,0 +1,118 @@ +use base64::{Config, CharacterSet}; +use base64; +use container_common::anyhow_result::AnyhowResult; +use magic_crypt::generic_array::typenum::U256; +use magic_crypt::new_magic_crypt; +use magic_crypt::{MagicCryptTrait, MagicCrypt256}; +use rand::RngCore; +use std::io::Cursor; + +// TODO: A protobuf would be more compact! +/// This gets encrypted and sent to the frontend as an opaque handle. +#[derive(Serialize, Deserialize)] +pub struct SortId { + // NB: Entropy is causing React to panic on re-renders + //pub entropy: u32, + pub column_id: u64, +} + +/// This exists so that we don't leak our database IDs. +/// If competitors see our IDs, they'll know how big our database is. +/// We encrypt our IDs so we don't leak "business secret information". +#[derive(Clone)] +pub struct SortKeyCrypto { + crypt: MagicCrypt256, + base64_config: Config, +} + +impl SortKeyCrypto { + pub fn new(secret: &str) -> Self { + let base64_config = Config::new(CharacterSet::UrlSafe, false) + .decode_allow_trailing_bits(true); + Self { + crypt: new_magic_crypt!(secret, 256), + base64_config, + } + } + + pub fn encrypt_id(&self, id: u64) -> AnyhowResult { + let mut rng = rand::thread_rng(); + + let payload = SortId { + //entropy: rng.next_u32(), + column_id: id, + }; + + let payload = serde_json::to_string(&payload)?; + + let mut reader = Cursor::new(payload); + let mut writer = Vec::new(); + + self.crypt.encrypt_reader_to_writer2::(&mut reader, &mut writer)?; + let encoded = base64::encode_config(&writer, self.base64_config.clone()); + + Ok(encoded) + } + + pub fn decrypt_id(&self, base_64_payload: &str) -> AnyhowResult { + //let payload = self.crypt.decrypt_base64_to_string(base_64_payload)?; + + let decoded_bytes= base64::decode_config(base_64_payload, self.base64_config)?; + let decrypted_bytes = self.crypt.decrypt_bytes_to_bytes(&decoded_bytes)?; + + let payload = String::from_utf8(decrypted_bytes)?; + let payload = serde_json::from_str::(&payload)?; + + Ok(payload.column_id) + } +} + +#[cfg(test)] +mod tests { + use crate::util::encrypted_sort_id::SortKeyCrypto; + use std::collections::HashSet; + +// #[test] +// fn encrypt_entropy_means_no_duplicate_values() { +// let sorter = SortKeyCrypto::new("secret"); +// +// let mut encrypted_tokens = HashSet::new(); +// +// for i in 0 .. 1000 { +// let encrypted = sorter.encrypt_id(1234).unwrap(); +// encrypted_tokens.insert(encrypted); +// } +// +// assert_eq!(encrypted_tokens.len(), 1000); +// } + + #[test] + fn encrypt_roundtrip() { + let sorter = SortKeyCrypto::new("secret"); + let encrypted = sorter.encrypt_id(1234).unwrap(); + let decrypted = sorter.decrypt_id(&encrypted).unwrap(); + assert_eq!(decrypted, 1234); + } + + #[test] + fn encrypt_lots_roundtrip_low_numbers() { + let sorter = SortKeyCrypto::new("secret"); + + for i in 0 .. 1000 { + let encrypted = sorter.encrypt_id(i).unwrap(); + let decrypted = sorter.decrypt_id(&encrypted).unwrap(); + assert_eq!(decrypted, i); + } + } + + #[test] + fn encrypt_lots_roundtrip_high_numbers() { + let sorter = SortKeyCrypto::new("secret"); + + for i in 100000000 .. 100005000 { + let encrypted = sorter.encrypt_id(i).unwrap(); + let decrypted = sorter.decrypt_id(&encrypted).unwrap(); + assert_eq!(decrypted, i); + } + } +} diff --git a/crates/service/storyteller_web/src/util/filesystem.rs b/crates/service/storyteller_web/src/util/filesystem.rs new file mode 100644 index 0000000000..b4553ed95a --- /dev/null +++ b/crates/service/storyteller_web/src/util/filesystem.rs @@ -0,0 +1,47 @@ +use anyhow::bail; +use container_common::anyhow_result::AnyhowResult; +use log::info; +use log::warn; +use std::fs; +use std::path::Path; +use std::path::PathBuf; + +pub fn check_directory_exists(path: &PathBuf) -> AnyhowResult<()> { + if !path.exists() { + bail!("Path doesn't exist: {:?}", path); + } + if !path.is_dir() { + bail!("Path isn't a directory: {:?}", path); + } + Ok(()) +} + +pub fn check_file_exists(path: &PathBuf) -> AnyhowResult<()> { + if !path.exists() { + bail!("Path doesn't exist: {:?}", path); + } + if !path.is_file() { + bail!("Path isn't a file: {:?}", path); + } + Ok(()) +} + +pub fn safe_delete_temp_file>(file_path: P) { + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + let printable_name = file_path.as_ref().to_str().unwrap_or("bad filename"); + match fs::remove_file(&file_path) { + Ok(_) => info!("Temp file deleted: {}", printable_name), + Err(e) => warn!("Could not delete temp file {:?} (not a fatal error): {:?}", + printable_name, e), + } +} + +pub fn safe_delete_temp_directory>(directory_path: P) { + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + let printable_name = directory_path.as_ref().to_str().unwrap_or("bad directory"); + match fs::remove_dir_all(&directory_path) { + Ok(_) => info!("Temp directory deleted: {}", printable_name), + Err(e) => warn!("Could not delete temp directory{:?} (not a fatal error): {:?}", + printable_name, e), + } +} diff --git a/crates/service/storyteller_web/src/util/hashing/hash_file_sha2.rs b/crates/service/storyteller_web/src/util/hashing/hash_file_sha2.rs new file mode 100644 index 0000000000..b0a31ba68c --- /dev/null +++ b/crates/service/storyteller_web/src/util/hashing/hash_file_sha2.rs @@ -0,0 +1,16 @@ +use container_common::anyhow_result::AnyhowResult; +use crate::util::hashing::sha256_digest::sha256_digest; +use data_encoding::HEXLOWER_PERMISSIVE; +use ring::digest::{Context, Digest, SHA256}; +use std::fs::File; +use std::io::{Read, BufReader}; + +pub fn hash_file_sha2(filename: &str) -> AnyhowResult { + let input = File::open(filename)?; + let reader = BufReader::new(input); + let digest = sha256_digest(reader)?; + + let hash = HEXLOWER_PERMISSIVE.encode(digest.as_ref()); + Ok(hash) +} + diff --git a/crates/service/storyteller_web/src/util/hashing/hash_string_sha2.rs b/crates/service/storyteller_web/src/util/hashing/hash_string_sha2.rs new file mode 100644 index 0000000000..a6fe989c72 --- /dev/null +++ b/crates/service/storyteller_web/src/util/hashing/hash_string_sha2.rs @@ -0,0 +1,12 @@ +use container_common::anyhow_result::AnyhowResult; +use crate::util::hashing::sha256_digest::sha256_digest; +use data_encoding::HEXLOWER_PERMISSIVE; +use ring::digest::{Context, Digest, SHA256}; +use std::fs::File; +use std::io::{Read, BufReader}; + +pub fn hash_string_sha2(string_input: &str) -> AnyhowResult { + let digest = sha256_digest(string_input.as_bytes())?; + let hash = HEXLOWER_PERMISSIVE.encode(digest.as_ref()); + Ok(hash) +} diff --git a/crates/service/storyteller_web/src/util/hashing/mod.rs b/crates/service/storyteller_web/src/util/hashing/mod.rs new file mode 100644 index 0000000000..a16479be48 --- /dev/null +++ b/crates/service/storyteller_web/src/util/hashing/mod.rs @@ -0,0 +1,3 @@ +pub mod hash_string_sha2; +pub mod sha256_digest; +pub mod hash_file_sha2; diff --git a/crates/service/storyteller_web/src/util/hashing/sha256_digest.rs b/crates/service/storyteller_web/src/util/hashing/sha256_digest.rs new file mode 100644 index 0000000000..d66480ee82 --- /dev/null +++ b/crates/service/storyteller_web/src/util/hashing/sha256_digest.rs @@ -0,0 +1,20 @@ +use container_common::anyhow_result::AnyhowResult; +use data_encoding::HEXLOWER_PERMISSIVE; +use ring::digest::{Context, Digest, SHA256}; +use std::fs::File; +use std::io::{Read, BufReader}; + +pub fn sha256_digest(mut reader: R) -> AnyhowResult { + let mut context = Context::new(&SHA256); + let mut buffer = [0; 1024]; + + loop { + let count = reader.read(&mut buffer)?; + if count == 0 { + break; + } + context.update(&buffer[..count]); + } + + Ok(context.finish()) +} diff --git a/crates/service/storyteller_web/src/util/jobs/cache_miss_strategizer.rs b/crates/service/storyteller_web/src/util/jobs/cache_miss_strategizer.rs new file mode 100644 index 0000000000..333fdbbda7 --- /dev/null +++ b/crates/service/storyteller_web/src/util/jobs/cache_miss_strategizer.rs @@ -0,0 +1,187 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc, Duration}; +use std::collections::HashMap; +use std::hash::Hash; + +/// How to handle cache miss events, as dictated by the `CacheMissStrategizer`. +#[derive(Eq, PartialEq, Copy, Clone, Debug)] +pub enum CacheMissStrategy { + /// Don't handle the cache miss right now. + /// Wait or skip until told otherwise. + WaitOrSkip, + /// Handle the cache miss now. + /// Do the download, calculation, whatever immediately. + Proceed, +} + +struct CacheMissLog { + pub first_miss: DateTime, + pub most_recent_miss: DateTime, +} + +/// CacheMissStrategizer +/// +/// This is the core component of the cold cache system. +/// +/// Keep track of a cache and tell the caller when to proceed with work. +/// If it isn't time yet, skip or wait and let another consumer handle it. +pub struct CacheMissStrategizer { + /// How long it takes after the first cache miss to do work. + max_cold_duration: Duration, + + /// How long it takes to reset the cold cache timer. + /// eg. If we see a cache miss for X and then another one a day later, + /// that shouldn't be cause to do work now. + forget_duration: Duration, + + /// Log of all cache misses. + cache_miss_log: HashMap, + + /// Clock can be replaced for testing. + get_time_function: Box DateTime>, +} + +impl CacheMissLog { + fn new(first_miss: DateTime) -> Self { + Self { + first_miss: first_miss.clone(), + most_recent_miss: first_miss, + } + } + + fn with_most_recent_miss(&self, most_recent_miss: DateTime) -> Self { + Self { + first_miss: self.first_miss.clone(), + most_recent_miss, + } + } +} + +impl CacheMissStrategizer { + pub fn new(max_cold_duration: Duration, forget_duration: Duration) -> Self { + Self { + max_cold_duration, + forget_duration, + cache_miss_log: HashMap::new(), + get_time_function: Box::new(|| Utc::now()), + } + } + + pub fn new_for_testing( + max_cold_duration: Duration, + forget_duration: Duration, + get_time_function: Box DateTime> + ) -> Self { + Self { + max_cold_duration, + forget_duration, + cache_miss_log: HashMap::new(), + get_time_function, + } + } + + // NB: Not threadsafe due to multiple operations against hashes! + #[must_use] + pub fn cache_miss(&mut self, id: T) -> CacheMissStrategy { + let now : DateTime = (self.get_time_function)(); + + if let Some((_id, cache_miss_log)) = self.cache_miss_log.get_key_value(&id) { + let admit_duration = now.signed_duration_since(cache_miss_log.first_miss.clone()); + let forget_duration = now.signed_duration_since(cache_miss_log.most_recent_miss.clone()); + + if forget_duration > self.forget_duration { + // Too much time elapsed. + // Treat this as essentially seeing the miss for the first time. + let cache_miss_log = CacheMissLog::new(now); + self.cache_miss_log.insert(id, cache_miss_log); + CacheMissStrategy::WaitOrSkip + } + else if admit_duration > self.max_cold_duration { + // We're done waiting. Do the work. + self.cache_miss_log.remove(&id); + CacheMissStrategy::Proceed + } + else { + // Still waiting. + let cache_miss_log = cache_miss_log.with_most_recent_miss(now); + self.cache_miss_log.insert(id, cache_miss_log); + CacheMissStrategy::WaitOrSkip + } + } else { + let cache_miss_log = CacheMissLog::new(now); + self.cache_miss_log.insert(id, cache_miss_log); + CacheMissStrategy::WaitOrSkip + } + } + + // NB: Private; for testing only + fn set_time_function(&mut self, get_time_function: Box DateTime>) { + self.get_time_function = get_time_function; + } +} + +#[cfg(test)] +mod tests { + use chrono::{Duration, Utc, DateTime, TimeZone}; + use crate::util::jobs::cache_miss_strategizer::CacheMissStrategizer; + use crate::util::jobs::cache_miss_strategizer::CacheMissStrategy; + + fn get_date(datetime: &str) -> DateTime { + let datetime = DateTime::parse_from_rfc3339(datetime).unwrap(); + let utc : DateTime = DateTime::from(datetime); + utc + } + + #[test] + fn cold_cache_cache_miss_algorithm() { + let mut cache_miss_strategizer = CacheMissStrategizer::new( + Duration::seconds(10), + Duration::seconds(60), + ); + + // First invocation + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:00+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::WaitOrSkip); + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:01+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::WaitOrSkip); + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:05+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::WaitOrSkip); + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:10+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::WaitOrSkip); + // Final invocation after time expires. Proceed. + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:11+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::Proceed); + + // New invocation. + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:15+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::WaitOrSkip); + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:25+00:00"))); + // New ID. + assert_eq!(cache_miss_strategizer.cache_miss(20), CacheMissStrategy::WaitOrSkip); + // Old ID (still wait) + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::WaitOrSkip); + // Old ID is done, new ID is still waiting. + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:30+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(10), CacheMissStrategy::Proceed); + assert_eq!(cache_miss_strategizer.cache_miss(20), CacheMissStrategy::WaitOrSkip); + // Now the new ID is also done. + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:00:41+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(20), CacheMissStrategy::Proceed); + + // New Invocation + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:01:00+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(30), CacheMissStrategy::WaitOrSkip); + // We waited too long... + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:02:01+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(30), CacheMissStrategy::WaitOrSkip); + + // New Invocation + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:04:00+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(30), CacheMissStrategy::WaitOrSkip); // previously seen + assert_eq!(cache_miss_strategizer.cache_miss(40), CacheMissStrategy::WaitOrSkip); // new + // We waited just long enough before expiry + cache_miss_strategizer.set_time_function(Box::new(|| get_date("2021-07-01T13:04:59+00:00"))); + assert_eq!(cache_miss_strategizer.cache_miss(30), CacheMissStrategy::Proceed); + assert_eq!(cache_miss_strategizer.cache_miss(40), CacheMissStrategy::Proceed); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/util/jobs/cache_miss_strategizer_multi.rs b/crates/service/storyteller_web/src/util/jobs/cache_miss_strategizer_multi.rs new file mode 100644 index 0000000000..34dbf66997 --- /dev/null +++ b/crates/service/storyteller_web/src/util/jobs/cache_miss_strategizer_multi.rs @@ -0,0 +1,74 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc, Duration}; +use container_common::anyhow_result::AnyhowResult; +use crate::util::jobs::cache_miss_strategizer::{CacheMissStrategizer, CacheMissStrategy}; +use std::sync::{RwLock, Arc}; + +/// Keep track of multiple caches, each with different time penalties. +/// (This isn't strictly necessary) +pub struct MultiCacheMissStrategizer { + in_memory_log: CacheMissStrategizer, + on_disk_log: CacheMissStrategizer, + // ... future caches ? +} + +/// And the Sync/Send threadsafe + interior mutability version. +#[derive(Clone)] +pub struct SyncMultiCacheMissStrategizer { + multi_cache_log: Arc>, +} + +impl SyncMultiCacheMissStrategizer { + pub fn new( + in_memory_log: CacheMissStrategizer, + on_disk_log: CacheMissStrategizer, + ) -> Self { + let multi_cache_log = MultiCacheMissStrategizer::new( + in_memory_log, + on_disk_log, + ); + Self { + multi_cache_log: Arc::new(RwLock::new(multi_cache_log)), + } + } + + #[must_use] + pub fn memory_cache_miss(&self, token: &str) -> AnyhowResult { + let result = self.multi_cache_log + .write() + .map(|mut cache_logs| cache_logs.memory_cache_miss(token)) + .map_err(|err| anyhow!("mutex error: {:?}", err))?; + Ok(result) + } + + #[must_use] + pub fn disk_cache_miss(&self, token: &str) -> AnyhowResult { + let result = self.multi_cache_log + .write() + .map(|mut cache_logs| cache_logs.disk_cache_miss(token)) + .map_err(|err| anyhow!("mutex error: {:?}", err))?; + Ok(result) + } +} + +impl MultiCacheMissStrategizer { + pub fn new( + in_memory_log: CacheMissStrategizer, + on_disk_log: CacheMissStrategizer, + ) -> Self { + Self { + in_memory_log, + on_disk_log, + } + } + + #[must_use] + pub fn memory_cache_miss(&mut self, token: &str) -> CacheMissStrategy { + self.in_memory_log.cache_miss(token.to_string()) + } + + #[must_use] + pub fn disk_cache_miss(&mut self, token: &str) -> CacheMissStrategy { + self.on_disk_log.cache_miss(token.to_string()) + } +} diff --git a/crates/service/storyteller_web/src/util/jobs/mod.rs b/crates/service/storyteller_web/src/util/jobs/mod.rs new file mode 100644 index 0000000000..b871837a3e --- /dev/null +++ b/crates/service/storyteller_web/src/util/jobs/mod.rs @@ -0,0 +1,3 @@ +pub mod cache_miss_strategizer; +pub mod cache_miss_strategizer_multi; +pub mod virtual_lfu_cache; diff --git a/crates/service/storyteller_web/src/util/jobs/virtual_lfu_cache.rs b/crates/service/storyteller_web/src/util/jobs/virtual_lfu_cache.rs new file mode 100644 index 0000000000..1487630d4e --- /dev/null +++ b/crates/service/storyteller_web/src/util/jobs/virtual_lfu_cache.rs @@ -0,0 +1,205 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use lfu::LFUCache; +use std::borrow::Borrow; +use std::collections::HashSet; +use std::sync::{Arc, RwLock}; + +/// This stands in front of the python HTTP sidecar and controls which +/// models get to remain in memory. The Python sidecar keeps multiple +/// models around, and Rust dictates which ones to keep and which to +/// purge. +pub struct VirtualLfuCache { + cache: LFUCache, +} + +/// And the Sync/Send threadsafe + interior mutability version. +#[derive(Clone)] +pub struct SyncVirtualLfuCache { + cache: Arc::> +} + +impl VirtualLfuCache { + + pub fn new(capacity: usize) -> AnyhowResult { + let cache = LFUCache::with_capacity(capacity) + .map_err(|e| anyhow!("Error creating cache (probably capacity): {:?}", e))?; + Ok(Self { + cache, + }) + } + + pub fn in_cache(&self, model: &String) -> bool { + self.cache.contains(model) + } + + /// Returns the evicted entry. + pub fn insert_returning_replaced(&mut self, path: &str) -> Option { + let initial_keys = self.get_keyset(); + + self.cache.set(path.to_string(), ()); + + let after_keys = self.get_keyset(); + + let difference = initial_keys.difference(&after_keys); + + difference.last() + .map(|item| item.to_string()) + } + + pub fn size(&self) -> usize { + self.cache.len() + } + + fn get_keyset(&self) -> HashSet { + // NB: This only makes sense for our use while we're using small caches. + let mut key_set = HashSet::with_capacity(self.size()); + for (k, _) in self.cache.iter() { + key_set.insert(k.to_string()); + } + key_set + } +} + +impl SyncVirtualLfuCache { + + pub fn new(capacity: usize) -> AnyhowResult { + let cache = VirtualLfuCache::new(capacity)?; + Ok(Self { + cache: Arc::new(RwLock::new(cache)), + }) + } + + pub fn in_cache(&self, model: &String) -> AnyhowResult { + let exists = self.cache + .read() + .map(|cache| cache.in_cache(model)) + .map_err(|err| anyhow!("mutex error: {:?}", err))?; + Ok(exists) + } + + /// Returns the evicted entry. + pub fn insert_returning_replaced(&self, path: &str) -> AnyhowResult> { + let maybe_replaced = self.cache + .write() + .map(|mut cache| cache.insert_returning_replaced(path)) + .map_err(|err| anyhow!("mutex error: {:?}", err))?; + Ok(maybe_replaced) + } + + pub fn size(&self) -> AnyhowResult { + let size = self.cache + .read() + .map(|cache| cache.size()) + .map_err(|err| anyhow!("mutex error: {:?}", err))?; + Ok(size) + } + + fn get_keyset(&self) -> AnyhowResult> { + let key_set = self.cache + .read() + .map(|cache| cache.get_keyset()) + .map_err(|err| anyhow!("mutex error: {:?}", err))?; + Ok(key_set) + } +} + +#[cfg(test)] +pub mod tests { + use crate::util::jobs::virtual_lfu_cache::VirtualLfuCache; + + #[test] + fn large_capacity_does_not_shed() { + let mut cache = VirtualLfuCache::new(10).unwrap(); + cache.insert_returning_replaced("1"); + cache.insert_returning_replaced("2"); + cache.insert_returning_replaced("3"); + cache.insert_returning_replaced("4"); + cache.insert_returning_replaced("5"); + cache.insert_returning_replaced("6"); + cache.insert_returning_replaced("7"); + cache.insert_returning_replaced("8"); + cache.insert_returning_replaced("9"); + assert!(cache.insert_returning_replaced("10").is_none()); // Does not remove + assert!(cache.insert_returning_replaced("11").is_some()); // Removes + assert_eq!(cache.size(), 10); + } + + #[test] + fn insert_beyond_capacity() { + let mut cache = VirtualLfuCache::new(3).unwrap(); + assert_eq!(cache.size(), 0); + cache.insert_returning_replaced("foo"); + assert_eq!(cache.size(), 1); + cache.insert_returning_replaced("bar"); + assert_eq!(cache.size(), 2); + cache.insert_returning_replaced("baz"); + assert_eq!(cache.size(), 3); + cache.insert_returning_replaced("bin"); + assert_eq!(cache.size(), 3); + cache.insert_returning_replaced("bin"); + assert_eq!(cache.size(), 3); + cache.insert_returning_replaced("111111111111111111"); + assert_eq!(cache.size(), 3); + } + + fn repeated_insert_single_key() { + let mut cache = VirtualLfuCache::new(3).unwrap(); + assert_eq!(cache.size(), 0); + for _ in 0..10 { + cache.insert_returning_replaced("foo"); + } + assert_eq!(cache.size(), 1); + } + + #[test] + fn returns_first_for_eviction_when_all_used_once() { + let mut cache = VirtualLfuCache::new(3).unwrap(); + cache.insert_returning_replaced("foo"); + cache.insert_returning_replaced("bar"); + cache.insert_returning_replaced("baz"); + let discarded = cache.insert_returning_replaced("bin"); + assert_eq!(discarded, Some("foo".to_string())); + + assert_eq!(false, cache.in_cache(&"foo".to_string())); + assert_eq!(true, cache.in_cache(&"bar".to_string())); + assert_eq!(true, cache.in_cache(&"baz".to_string())); + assert_eq!(true, cache.in_cache(&"bin".to_string())); + } + + #[test] + fn retains_frequent_value() { + let mut cache = VirtualLfuCache::new(3).unwrap(); + for _ in 0..10 { + cache.insert_returning_replaced("foo"); + } + cache.insert_returning_replaced("bar"); + cache.insert_returning_replaced("baz"); + let discarded = cache.insert_returning_replaced("bin"); + assert_eq!(discarded, Some("bar".to_string())); + + assert_eq!(true, cache.in_cache(&"foo".to_string())); + assert_eq!(false, cache.in_cache(&"bar".to_string())); + assert_eq!(true, cache.in_cache(&"baz".to_string())); + assert_eq!(true, cache.in_cache(&"bin".to_string())); + } + + #[test] + fn retains_second_frequent_value() { + let mut cache = VirtualLfuCache::new(3).unwrap(); + for _ in 0..10 { + cache.insert_returning_replaced("foo"); + } + for _ in 0..2 { + cache.insert_returning_replaced("bar"); + } + cache.insert_returning_replaced("baz"); + let discarded = cache.insert_returning_replaced("bin"); + assert_eq!(discarded, Some("baz".to_string())); + + assert_eq!(true, cache.in_cache(&"foo".to_string())); + assert_eq!(true, cache.in_cache(&"bar".to_string())); + assert_eq!(false, cache.in_cache(&"baz".to_string())); + assert_eq!(true, cache.in_cache(&"bin".to_string())); + } +} diff --git a/crates/service/storyteller_web/src/util/markdown_to_html.rs b/crates/service/storyteller_web/src/util/markdown_to_html.rs new file mode 100644 index 0000000000..2dccc6cafb --- /dev/null +++ b/crates/service/storyteller_web/src/util/markdown_to_html.rs @@ -0,0 +1,60 @@ +use container_common::anyhow_result::AnyhowResult; +use pulldown_cmark::{Options, Parser, html}; + +pub fn markdown_to_html(markdown_input: &str) -> String { + // NB: CommonMark allows for HTML, including "), + "

<script>alert();</script>

\n"); + } + +} diff --git a/crates/service/storyteller_web/src/util/mod.rs b/crates/service/storyteller_web/src/util/mod.rs new file mode 100644 index 0000000000..da7e8cad13 --- /dev/null +++ b/crates/service/storyteller_web/src/util/mod.rs @@ -0,0 +1,12 @@ +pub mod buckets; +pub mod caching; +pub mod delete_role_disambiguation; +pub mod email_to_gravatar; +pub mod encrypted_sort_id; +pub mod filesystem; +pub mod hashing; +pub mod jobs; +pub mod markdown_to_html; +pub mod noop_logger; +pub mod redis; +pub mod semi_persistent_cache_dir; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/util/noop_logger.rs b/crates/service/storyteller_web/src/util/noop_logger.rs new file mode 100644 index 0000000000..17dcc1da93 --- /dev/null +++ b/crates/service/storyteller_web/src/util/noop_logger.rs @@ -0,0 +1,32 @@ +use chrono::{DateTime, Utc, TimeZone, Duration}; +use log::info; + +/// Log after periods of no work, but prevent spamming stdout. +pub struct NoOpLogger { + last_log_time: DateTime, + logless_duration: Duration, + op_count: u64, +} + +impl NoOpLogger { + pub fn new(log_every_millis: i64) -> Self { + Self { + last_log_time: Utc.timestamp(0, 0), + logless_duration: Duration::milliseconds(log_every_millis), + op_count: 0 + } + } + + pub fn log_after_awhile(&mut self) { + let now = Utc::now(); + let delta = now - self.last_log_time; + + if delta > self.logless_duration { + info!("No operations for {:?}. {} operations.", delta, self.op_count); + self.op_count = 0; + self.last_log_time = now; + } + + self.op_count += 1; + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/util/redis/mod.rs b/crates/service/storyteller_web/src/util/redis/mod.rs new file mode 100644 index 0000000000..5cac961a97 --- /dev/null +++ b/crates/service/storyteller_web/src/util/redis/mod.rs @@ -0,0 +1 @@ +pub mod redis_job_status_logger; diff --git a/crates/service/storyteller_web/src/util/redis/redis_job_status_logger.rs b/crates/service/storyteller_web/src/util/redis/redis_job_status_logger.rs new file mode 100644 index 0000000000..c61de40b06 --- /dev/null +++ b/crates/service/storyteller_web/src/util/redis/redis_job_status_logger.rs @@ -0,0 +1,45 @@ +use container_common::anyhow_result::AnyhowResult; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2::PooledConnection; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; + +pub type RedisPool = PooledConnection; + +pub struct RedisJobStatusLogger <'a> { + redis: &'a mut RedisPool, + status_key: String, +} + +impl <'a> RedisJobStatusLogger <'a> { + pub const STATUS_KEY_TTL_SECONDS : usize = 60 * 60; + + pub fn new_tts_download(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self { + let status_key = RedisKeys::tts_download_extra_status_info(w2l_job_token); + Self { redis, status_key } + } + + pub fn new_tts_inference(redis: &'a mut RedisPool, tts_job_token: &str) -> Self { + let status_key = RedisKeys::tts_inference_extra_status_info(tts_job_token); + Self { redis, status_key } + } + + pub fn new_w2l_download(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self { + let status_key = RedisKeys::w2l_download_extra_status_info(w2l_job_token); + Self { redis, status_key } + } + + pub fn new_w2l_inference(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self { + let status_key = RedisKeys::w2l_inference_extra_status_info(w2l_job_token); + Self { redis, status_key } + } + + /// Record the status to Redis so the frontend has the latest state. + pub fn log_status(&mut self, logging_details: &str) -> AnyhowResult<()> { + let _r : String = self.redis // NB: Compiler can't figure out the throwaway result type + .set_ex(&self.status_key, + logging_details, + Self::STATUS_KEY_TTL_SECONDS)?; + Ok(()) + } +} diff --git a/crates/service/storyteller_web/src/util/semi_persistent_cache_dir.rs b/crates/service/storyteller_web/src/util/semi_persistent_cache_dir.rs new file mode 100644 index 0000000000..6c8827553a --- /dev/null +++ b/crates/service/storyteller_web/src/util/semi_persistent_cache_dir.rs @@ -0,0 +1,187 @@ +use container_common::anyhow_result::AnyhowResult; +use std::fs; +use std::path::{PathBuf, Path}; +use tempdir::TempDir; + +/// These are for files on the worker filesystems +pub struct SemiPersistentCacheDir { + cache_root: PathBuf, + + tts_synthesizer_model_root: PathBuf, + tts_pretrained_vocoder_model_root: PathBuf, + + w2l_model_root: PathBuf, + w2l_end_bump_root: PathBuf, + w2l_face_templates_root: PathBuf, + w2l_templates_media_root: PathBuf, + video_asset_root: PathBuf, // end bump, etc. +} + +impl SemiPersistentCacheDir { + + + /// Everything is rooted at `/file_cache`. + pub fn default_paths() -> Self { + // NB: This is the root of the filesystem + Self::configured_root("/file_cache/") + } + + pub fn configured_root(root_path: &str) -> Self { + let cache_root = PathBuf::from(root_path); + Self { + cache_root: cache_root.clone(), + + tts_synthesizer_model_root: cache_root.join("tts/synthesizer_models/"), + tts_pretrained_vocoder_model_root: cache_root.join("tts/vocoder_models_pretrained/"), + + w2l_model_root: cache_root.join("w2l/models/"), + w2l_end_bump_root: cache_root.join("w2l/end_bumps/"), + w2l_face_templates_root: cache_root.join("w2l/face_templates/"), + w2l_templates_media_root: cache_root.join("w2l/template_media/"), + video_asset_root: cache_root.join("static_video_assets/"), + } + } + + // ==================== TTS SYNTHESIZER MODELS ==================== + + /// We cache TTS synthesizer models here. + /// We'll likely need to LRU cache them. + pub fn tts_synthesizer_model_path(&self, model_filename: &str) -> PathBuf { + self.tts_synthesizer_model_root.join(model_filename) + } + + pub fn tts_synthesizer_model_directory(&self) -> &Path { + &self.tts_synthesizer_model_root + } + + pub fn create_tts_synthesizer_model_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.tts_synthesizer_model_directory())?; + Ok(()) + } + + // ==================== TTS PRETRAINED VOCODER MODELS (WAVEGLOW, HIFIGAN, ETC) ==================== + + /// We'll start with just a few vocoders, but this may grow. + pub fn tts_pretrained_vocoder_model_path(&self, model_filename: &str) -> PathBuf { + self.tts_pretrained_vocoder_model_root.join(model_filename) + } + + pub fn tts_pretrained_vocoder_model_directory(&self) -> &Path { + &self.tts_pretrained_vocoder_model_root + } + + pub fn create_tts_pretrained_vocoder_model_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.tts_pretrained_vocoder_model_directory())?; + Ok(()) + } + + // ==================== W2L MODELS (there are only two of them) ==================== + + /// There are only two pretrained W2L models, so we won't run out of space. + pub fn w2l_model_path(&self, model_filename: &str) -> PathBuf { + self.w2l_model_root.join(model_filename) + } + + pub fn w2l_model_directory(&self) -> &Path { + &self.w2l_model_root + } + + pub fn create_w2l_model_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.w2l_model_directory())?; + Ok(()) + } + + // ==================== W2L END BUMPS (short video logos) ==================== + + pub fn w2l_end_bump_path(&self, end_bump_filename: &str) -> PathBuf { + self.w2l_end_bump_root.join(end_bump_filename) + } + + pub fn w2l_end_bump_directory(&self) -> &Path { + &self.w2l_end_bump_root + } + + pub fn create_w2l_end_bump_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.w2l_end_bump_directory())?; + Ok(()) + } + + // ==================== W2L MEDIA ==================== + + /// We cache W2L media here. + /// We'll likely need to LRU cache them. + pub fn w2l_template_media_path(&self, template_private_bucket_hash: &str) -> PathBuf { + self.w2l_templates_media_root.join(template_private_bucket_hash) + } + + pub fn w2l_template_media_directory(&self) -> &Path { + &self.w2l_templates_media_root + } + + pub fn create_w2l_template_media_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.w2l_template_media_directory())?; + Ok(()) + } + + // ==================== W2L CACHED FACES ==================== + + /// We cache W2L faces here. + /// We'll likely need to LRU cache them. + pub fn w2l_face_template_path(&self, template_private_bucket_hash: &str) -> PathBuf { + let filename = format!("{}_detected_faces.pickle", template_private_bucket_hash); + self.w2l_face_templates_root.join(filename) + } + + pub fn w2l_face_template_directory(&self) -> &Path { + &self.w2l_face_templates_root + } + + pub fn create_w2l_face_template_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.w2l_face_template_directory())?; + Ok(()) + } + + // ==================== VIDEO ASSETS (End bump, etc) ==================== + + /// There is only a single end bump, and we'll add a watermark file. + pub fn video_asset_path(&self) -> &Path { + &self.video_asset_root + } + + pub fn create_video_asset_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.video_asset_path())?; + Ok(()) + } + + // ==================== W2L OUTPUT RESULTS ==================== + + // We cache W2L faces here. + // We'll likely need to LRU cache them. + //pub fn w2l_output_results_path(&self, temp_dir: &TempDir, inference_job_token: &str) -> PathBuf { + // // NB: We don't want colons from the token in the filename. + // let filename = inference_job_token.replace(":", ""); + // let filename = format!("{}_result.mp4", filename); + + // temp_dir.path().join(&filename) + //} +} + +//#[cfg(test)] +//mod tests { +// use crate::util::semi_persistent_cache_dir::SemiPersistentCacheDir; +// use tempdir::TempDir; +// +// #[test] +// fn test_w2l_output_results_path() { +// let dirs = SemiPersistentCacheDir::default_paths(); +// let temp_dir = TempDir::new("test").unwrap(); +// +// let filename = dirs.w2l_output_results_path(&temp_dir, "TYPE:TOKEN") +// .to_str() +// .map(|s| s.to_string()) +// .unwrap(); +// +// assert!(&filename.ends_with("TYPETOKEN_result.mp4")); +// } +//} +// \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/cashapp_username.rs b/crates/service/storyteller_web/src/validations/cashapp_username.rs new file mode 100644 index 0000000000..960b087ed9 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/cashapp_username.rs @@ -0,0 +1,64 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_cashapp_username(username: &str) -> Result<(), String> { + lazy_static! { + static ref CASHAPP_USERNAME_REGEX: Regex = { + Regex::new(r"^\$?.{1,20}$").expect("should be valid regex") + }; + } + + if username.starts_with("$") { + if username.len() < 2 { + return Err("cashapp username is too short".to_string()); + } + if username.len() > 21 { + return Err("cashapp username is too long".to_string()); + } + } else { + if username.len() < 1 { + return Err("cashapp username is too short".to_string()); + } + if username.len() > 20 { + return Err("cashapp username is too long".to_string()); + } + } + + if !CASHAPP_USERNAME_REGEX.is_match(username) { + return Err("cashapp username invalid".to_string()); + } + + Ok(()) +} + +/// Remove the leading '$' for consistency and better internal use. +pub fn normalize_cashapp_username_for_storage(username: &str) -> String { + username.replace("$", "") +} + +#[cfg(test)] +mod tests { + use crate::validations::cashapp_username::{validate_cashapp_username, normalize_cashapp_username_for_storage}; + + #[test] + fn valid_cases() { + assert!(validate_cashapp_username("echelon").is_ok()); + assert!(validate_cashapp_username("$echelon").is_ok()); + assert!(validate_cashapp_username("a").is_ok()); + assert!(validate_cashapp_username("12345678901234567890").is_ok()); + assert!(validate_cashapp_username("$12345678901234567890").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_cashapp_username("").is_err()); + assert!(validate_cashapp_username("$").is_err()); + assert!(validate_cashapp_username("123456789012345678901").is_err()); + } + + #[test] + fn test_normalize_twitter_username_for_storage() { + assert_eq!(normalize_cashapp_username_for_storage("echelon"), "echelon".to_string()); + assert_eq!(normalize_cashapp_username_for_storage("$echelon"), "echelon".to_string()); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/discord_username.rs b/crates/service/storyteller_web/src/validations/discord_username.rs new file mode 100644 index 0000000000..1b49ce1fcd --- /dev/null +++ b/crates/service/storyteller_web/src/validations/discord_username.rs @@ -0,0 +1,40 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_discord_username(username: &str) -> Result<(), String> { + lazy_static! { + static ref DISCORD_USERNAME_REGEX: Regex = { + Regex::new(r"^@?[^#@:]{2,32}#[0-9]{4}$").expect("should be valid regex") + }; + } + + if username.len() < 2 { + return Err("discord username is too short".to_string()); + } + + if username.len() > 40 { + return Err("discord username is too long".to_string()); + } + + if !DISCORD_USERNAME_REGEX.is_match(username) { + return Err("discord username invalid".to_string()); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::discord_username::validate_discord_username; + + #[test] + fn valid_cases() { + assert!(validate_discord_username("echelon#0001").is_ok()); + assert!(validate_discord_username("@echelon#0001").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_discord_username("").is_err()); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/github_username.rs b/crates/service/storyteller_web/src/validations/github_username.rs new file mode 100644 index 0000000000..1654b887a1 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/github_username.rs @@ -0,0 +1,41 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_github_username(username: &str) -> Result<(), String> { + lazy_static! { + // TODO: https://github.com/shinnn/github-username-regex + static ref GITHUB_USERNAME_REGEX: Regex = { + Regex::new(r"^[a-zA-Z0-9\-]{1,39}$").expect("should be valid regex") + }; + } + + if username.len() < 1 { + return Err("github username is too short".to_string()); + } + + if username.len() > 39 { + return Err("github username is too long".to_string()); + } + + if !GITHUB_USERNAME_REGEX.is_match(username) { + return Err("github username invalid".to_string()); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::github_username::validate_github_username; + + #[test] + fn valid_cases() { + assert!(validate_github_username("echelon").is_ok()); + assert!(validate_github_username("a-b-c-D-E-F-0-1-2").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_github_username("").is_err()); + } +} diff --git a/crates/service/storyteller_web/src/validations/ip_addresses.rs b/crates/service/storyteller_web/src/validations/ip_addresses.rs new file mode 100644 index 0000000000..20a22d4e45 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/ip_addresses.rs @@ -0,0 +1,44 @@ +use crate::AnyhowResult; +use regex::Regex; + +/// Do not use this for IP addresses from requests. +/// Only for moderator-input IP addresses +pub fn validate_moderator_provided_ip_address(ip_address: &str) -> Result<(), String> { + lazy_static! { + static ref IP_ADDRESS_REGEX: Regex = { + // TODO: This doesn't require valid IP addresses. Just a number/dot format. + Regex::new(r"^(\d{1,3})\.(\d{1,3})\.(\d{1,3})\.(\d{1,3})$").expect("should be valid regex") + }; + } + + if !IP_ADDRESS_REGEX.is_match(ip_address) { + return Err("ip address is invalid".to_string()); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::ip_addresses::validate_moderator_provided_ip_address; + + #[test] + fn valid_cases() { + assert!(validate_moderator_provided_ip_address("0.0.0.0").is_ok()); + assert!(validate_moderator_provided_ip_address("1.1.1.1").is_ok()); + assert!(validate_moderator_provided_ip_address("1.2.3.4").is_ok()); + assert!(validate_moderator_provided_ip_address("111.222.3.0").is_ok()); + assert!(validate_moderator_provided_ip_address("127.0.0.1").is_ok()); + assert!(validate_moderator_provided_ip_address("255.255.255.255").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_moderator_provided_ip_address(" 255.255.255.255 ").is_err()); + assert!(validate_moderator_provided_ip_address("").is_err()); + assert!(validate_moderator_provided_ip_address(".1.1.1.1").is_err()); + assert!(validate_moderator_provided_ip_address("1.1.1").is_err()); + assert!(validate_moderator_provided_ip_address("127.0.0.1notvalid").is_err()); + assert!(validate_moderator_provided_ip_address("127.127.127.127.").is_err()); + } +} diff --git a/crates/service/storyteller_web/src/validations/mod.rs b/crates/service/storyteller_web/src/validations/mod.rs new file mode 100644 index 0000000000..e881b761a1 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/mod.rs @@ -0,0 +1,11 @@ +pub mod cashapp_username; +pub mod discord_username; +pub mod github_username; +pub mod ip_addresses; +pub mod model_uploads; +pub mod passwords; +pub mod twitch_username; +pub mod twitter_username; +pub mod username; +pub mod username_reservations; +pub mod website_url; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/model_uploads.rs b/crates/service/storyteller_web/src/validations/model_uploads.rs new file mode 100644 index 0000000000..c324f5ebc0 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/model_uploads.rs @@ -0,0 +1,14 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_model_title(title: &str) -> Result<(), String> { + if title.len() < 3 { + return Err("title is too short".to_string()); + } + + if title.len() > 200 { + return Err("title is too long".to_string()); + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/validations/passwords.rs b/crates/service/storyteller_web/src/validations/passwords.rs new file mode 100644 index 0000000000..b86477f2cc --- /dev/null +++ b/crates/service/storyteller_web/src/validations/passwords.rs @@ -0,0 +1,12 @@ + +pub fn validate_passwords(password: &str, password_confirmation: &str) -> Result<(), String> { + if password.len() < 6 { + return Err("password is too short".to_string()); + } + + if password != password_confirmation { + return Err("passwords do not match".to_string()); + } + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/validations/twitch_username.rs b/crates/service/storyteller_web/src/validations/twitch_username.rs new file mode 100644 index 0000000000..e93509ffae --- /dev/null +++ b/crates/service/storyteller_web/src/validations/twitch_username.rs @@ -0,0 +1,43 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_twitch_username(username: &str) -> Result<(), String> { + lazy_static! { + static ref TWITCH_USERNAME_REGEX: Regex = { + Regex::new(r"^[a-zA-Z0-9]{4,25}$").expect("should be valid regex") + }; + } + + if username.len() < 4 { + return Err("twitch username is too short".to_string()); + } + + if username.len() > 25 { + return Err("twitch username is too long".to_string()); + } + + if !TWITCH_USERNAME_REGEX.is_match(username) { + return Err("twitch username invalid".to_string()); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::twitch_username::validate_twitch_username; + + #[test] + fn valid_cases() { + assert!(validate_twitch_username("echelon").is_ok()); + assert!(validate_twitch_username("1234").is_ok()); + assert!(validate_twitch_username("1234567890123456789012345").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_twitch_username("").is_err()); + assert!(validate_twitch_username("@=[]@34!").is_err()); + assert!(validate_twitch_username("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa").is_err()); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/twitter_username.rs b/crates/service/storyteller_web/src/validations/twitter_username.rs new file mode 100644 index 0000000000..3deed67fa8 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/twitter_username.rs @@ -0,0 +1,72 @@ +use crate::AnyhowResult; +use regex::Regex; + + +pub fn validate_twitter_username(username: &str) -> Result<(), String> { + lazy_static! { + static ref TWITTER_USERNAME_REGEX: Regex = { + Regex::new("^@?[A-Za-z0-9_]{4,15}$").expect("should be valid regex") + }; + } + + if username.starts_with("@") { + if username.len() < 5 { + return Err("twitter username is too short".to_string()); + } + if username.len() > 16 { + return Err("twitter username is too long".to_string()); + } + } else { + if username.len() < 4 { + return Err("twitter username is too short".to_string()); + } + if username.len() > 15 { + return Err("twitter username is too long".to_string()); + } + } + + if !TWITTER_USERNAME_REGEX.is_match(username) { + return Err("twitter username invalid".to_string()); + } + + Ok(()) +} + +/// Remove the leading '@' for consistency and better internal use. +pub fn normalize_twitter_username_for_storage(username: &str) -> String { + username.replace("@", "") +} + +#[cfg(test)] +mod tests { + use crate::validations::twitter_username::{validate_twitter_username, normalize_twitter_username_for_storage}; + + #[test] + fn valid_cases() { + assert!(validate_twitter_username("@echelon").is_ok()); + assert!(validate_twitter_username("echelon").is_ok()); + assert!(validate_twitter_username("four").is_ok()); + assert!(validate_twitter_username("@four").is_ok()); + assert!(validate_twitter_username("123456789012345").is_ok()); + assert!(validate_twitter_username("@123456789012345").is_ok()); + assert!(validate_twitter_username("a_A_b_B_c_C_d_D").is_ok()); + assert!(validate_twitter_username("@a_A_b_B_c_C_d_D").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_twitter_username("").is_err()); + assert!(validate_twitter_username("@").is_err()); + assert!(validate_twitter_username(" ").is_err()); + assert!(validate_twitter_username("!!!!!!!!!!!").is_err()); + assert!(validate_twitter_username("a").is_err()); + assert!(validate_twitter_username("@a").is_err()); + assert!(validate_twitter_username("aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa").is_err()); + } + + #[test] + fn test_normalize_twitter_username_for_storage() { + assert_eq!(normalize_twitter_username_for_storage("echelon"), "echelon".to_string()); + assert_eq!(normalize_twitter_username_for_storage("@echelon"), "echelon".to_string()); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/username.rs b/crates/service/storyteller_web/src/validations/username.rs new file mode 100644 index 0000000000..c823959b7f --- /dev/null +++ b/crates/service/storyteller_web/src/validations/username.rs @@ -0,0 +1,67 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_username(username: &str) -> Result<(), String> { + lazy_static! { + static ref USERNAME_REGEX: Regex = { + Regex::new(r"^[A-Za-z0-9_\-]{3,16}$").expect("should be valid regex") + }; + } + + if username.len() < 3 { + return Err("username is too short".to_string()); + } + + if username.len() > 16 { + return Err("username is too long".to_string()); + } + + if !USERNAME_REGEX.is_match(username) { + return Err("invalid username characters".to_string()); + } + + Ok(()) +} + +/*pub fn validate_display_name(display_name: &str) -> Result<(), String> { + lazy_static! { + static ref DISPLAY_NAME_REGEX: Regex = { + Regex::new(r"^\w{3,16}$").expect("should be valid regex") + }; + } + + if username.len() < 3 { + return Err("display name is too short".to_string()); + } + + if username.len() > 16 { + return Err("display name is too long".to_string()); + } + + if !DISPLAY_NAME_REGEX.is_match(username) { + return Err("invalid display name".to_string()); + } + + Ok(()) +}*/ + +#[cfg(test)] +mod tests { + use crate::validations::username::validate_username; + + #[test] + fn test_valid_usernames() { + assert!(validate_username("echelon").is_ok()); + assert!(validate_username("dr_mario").is_ok()); + assert!(validate_username("Foo1234").is_ok()); + assert!(validate_username("mr-person").is_ok()); + } + + #[test] + fn test_invalid_usernames() { + assert!(validate_username("").is_err()); + assert!(validate_username("&&&&").is_err()); + assert!(validate_username("........").is_err()); + assert!(validate_username("really-long-username-that-is-too-long").is_err()); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/username_reservations.rs b/crates/service/storyteller_web/src/validations/username_reservations.rs new file mode 100644 index 0000000000..5501341162 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/username_reservations.rs @@ -0,0 +1,72 @@ +use std::collections::HashSet; +use crate::RESERVED_SUBSTRINGS; +use crate::RESERVED_USERNAMES; + +pub fn is_reserved_username(username: &str) -> bool { + lazy_static! { + static ref RESERVED_USERNAMES_SET : HashSet = RESERVED_USERNAMES.lines() + .map(|line| line.trim()) + .filter(|line| !(line.starts_with("#") || line.is_empty())) + .map(|line| line.to_string()) + .collect::>(); + } + + if RESERVED_USERNAMES_SET.contains(username) { + return true; + } + + is_reserved_substring(username) +} + +fn is_reserved_substring(username: &str) -> bool { + lazy_static! { + static ref RESERVED_SUBSTRINGS_LIST : Vec = RESERVED_SUBSTRINGS.lines() + .map(|line| line.trim()) + .filter(|line| !(line.starts_with("#") || line.is_empty())) + .map(|line| line.to_string()) + .collect::>(); + } + + let undashed = username.replace("_", "").replace("-", ""); + + for substr in RESERVED_SUBSTRINGS_LIST.iter() { + if username.contains(substr) || undashed.contains(substr) { + return true; + } + } + + return false; +} + +#[cfg(test)] +mod tests { + use crate::validations::username_reservations::{is_reserved_username}; + + #[test] + fn reserved_usernames() { + assert_eq!(is_reserved_username("vocodes"), true); + assert_eq!(is_reserved_username("user"), true); + assert_eq!(is_reserved_username("username"), true); + } + + #[test] + fn unreserved_usernames() { + assert_eq!(is_reserved_username("echelon"), false); + assert_eq!(is_reserved_username("asdfasdfadsf"), false); + assert_eq!(is_reserved_username("bobdole11"), false); + } + + #[test] + fn reserved_substrings() { + assert_eq!(is_reserved_username("test112345"), true); + assert_eq!(is_reserved_username("12345test"), true); + assert_eq!(is_reserved_username("test"), true); + assert_eq!(is_reserved_username("111vocodes111"), true); + } + + #[test] + fn reserved_substrings_with_dashes() { + assert_eq!(is_reserved_username("t_e_s_t"), true); + assert_eq!(is_reserved_username("-vo-co-de-s--"), true); + } +} \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/website_url.rs b/crates/service/storyteller_web/src/validations/website_url.rs new file mode 100644 index 0000000000..b3b5614b49 --- /dev/null +++ b/crates/service/storyteller_web/src/validations/website_url.rs @@ -0,0 +1,41 @@ +use crate::AnyhowResult; +use regex::Regex; + +pub fn validate_website_url(website_url: &str) -> Result<(), String> { + lazy_static! { + static ref WEBSITE_URL_REGEX: Regex = { + Regex::new(r"^(https?:\\/\\/)?.*$").expect("should be valid regex") + }; + } + + if website_url.len() < 3 { + return Err("website url is too short".to_string()); + } + + if website_url.len() > 200 { + return Err("website url is too long".to_string()); + } + + if !WEBSITE_URL_REGEX.is_match(website_url) { + return Err("website url invalid".to_string()); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::website_url::validate_website_url; + + #[test] + fn valid_cases() { + assert!(validate_website_url("http://vo.codes").is_ok()); + assert!(validate_website_url("https://vocodes.com").is_ok()); + assert!(validate_website_url("www.zombo.com").is_ok()); + } + + #[test] + fn invalid_cases() { + assert!(validate_website_url("a").is_err()); + } +} diff --git a/crates/service/storyteller_web/src/w2l_download_job.rs b/crates/service/storyteller_web/src/w2l_download_job.rs new file mode 100644 index 0000000000..370bbf8783 --- /dev/null +++ b/crates/service/storyteller_web/src/w2l_download_job.rs @@ -0,0 +1,654 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate serde_derive; + +pub mod common_env; +pub mod job_queries; +pub mod script_execution; +pub mod util; + +use anyhow::{anyhow, Error}; +use chrono::Utc; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use crate::common_env::CommonEnv; +use crate::job_queries::w2l_download_job_queries::W2lTemplateUploadJobRecord; +use crate::job_queries::w2l_download_job_queries::grab_job_lock_and_mark_pending; +use crate::job_queries::w2l_download_job_queries::insert_w2l_template; +use crate::job_queries::w2l_download_job_queries::mark_w2l_template_upload_job_done; +use crate::job_queries::w2l_download_job_queries::mark_w2l_template_upload_job_failure; +use crate::job_queries::w2l_download_job_queries::mark_w2l_template_upload_job_permanently_dead; +use crate::job_queries::w2l_download_job_queries::query_w2l_template_upload_job_records; +use crate::script_execution::ffmpeg_generate_preview_image_command::FfmpegGeneratePreviewImageCommand; +use crate::script_execution::ffmpeg_generate_preview_video_command::FfmpegGeneratePreviewVideoCommand; +use crate::script_execution::google_drive_download_command::GoogleDriveDownloadCommand; +use crate::script_execution::imagemagick_generate_preview_image_command::ImagemagickGeneratePreviewImageCommand; +use crate::script_execution::wav2lip_process_upload_command::{Wav2LipPreprocessClient, Wav2LipPreprocessError}; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::buckets::bucket_paths::hash_to_bucket_path; +use crate::util::filesystem::check_directory_exists; +use crate::util::filesystem::check_file_exists; +use crate::util::filesystem::safe_delete_temp_directory; +use crate::util::filesystem::safe_delete_temp_file; +use crate::util::hashing::hash_file_sha2::hash_file_sha2; +use crate::util::noop_logger::NoOpLogger; +use crate::util::redis::redis_job_status_logger::RedisJobStatusLogger; +use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE}; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use log::{warn, info}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use ring::digest::{Context, Digest, SHA256}; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::fs::{File, metadata}; +use std::io::{BufReader, Read}; +use std::ops::Deref; +use std::path::{PathBuf, Path}; +use std::process::Command; +use std::thread; +use std::time::Duration; +use tempdir::TempDir; + +// Buckets (shared config) +const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; +const ENV_SECRET_KEY : &'static str = "SECRET_KEY"; +const ENV_REGION_NAME : &'static str = "REGION_NAME"; + +// Buckets (private data) +const ENV_PRIVATE_BUCKET_NAME : &'static str = "W2L_PRIVATE_DOWNLOAD_BUCKET_NAME"; +// Buckets (public data) +const ENV_PUBLIC_BUCKET_NAME : &'static str = "W2L_PUBLIC_DOWNLOAD_BUCKET_NAME"; + +// Various bucket roots +const ENV_DOWNLOAD_BUCKET_ROOT : &'static str = "W2L_DOWNLOAD_BUCKET_ROOT"; + +// Python code +const ENV_CODE_DIRECTORY : &'static str = "W2L_CODE_DIRECTORY"; +const ENV_MODEL_CHECKPOINT : &'static str = "W2L_MODEL_CHECKPOINT"; +const ENV_DOWNLOAD_SCRIPT_NAME : &'static str = "W2L_DOWNLOAD_SCRIPT_NAME"; + +const DEFAULT_TEMP_DIR: &'static str = "/tmp"; + +struct Downloader { + pub download_temp_directory: PathBuf, + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub private_bucket_client: BucketClient, + pub public_bucket_client: BucketClient, + + pub firehose_publisher: FirehosePublisher, + pub badge_granter: BadgeGranter, + + pub google_drive_downloader: GoogleDriveDownloadCommand, + pub w2l_processor: Wav2LipPreprocessClient, + pub ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand, + pub ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand, + pub imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand, + + // Command to run + pub download_script: String, + + // Root to store W2L templates (public and private) + pub bucket_root_w2l_template_uploads: String, + + // Temporary for debugging + // Arbitrary timeouts can be inserted so we can exec in and poke around. + pub debug_download_sleep_millis: u64, + pub debug_face_detect_sleep_millis: u64, + pub debug_job_end_sleep_millis: u64, + + // Sleep between batches + pub job_batch_wait_millis: u64, + + // How long to wait between log lines + pub no_op_logger_millis: u64, + + // Max job attempts before failure. + // NB: This is an i32 so we don't need to convert to db column type. + pub job_max_attempts: i32, +} + +#[tokio::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + // NB: Do not check this secrets-containing dotenv file into VCS. + // This file should only contain *development* secrets, never production. + let _ = dotenv::from_filename(".env-secrets").ok(); + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("w2l-download-job".to_string()); + + info!("Hostname: {}", &server_hostname); + + // Bucket stuff (shared) + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + + // Private and Public Buckets + let private_bucket_name = easyenv::get_env_string_required(ENV_PRIVATE_BUCKET_NAME)?; + let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?; + + // Bucket roots + let bucket_root = easyenv::get_env_string_required(ENV_DOWNLOAD_BUCKET_ROOT)?; + + let private_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &private_bucket_name, + None, + )?; + + let public_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &public_bucket_name, + None, + )?; + + let py_code_directory = easyenv::get_env_string_required(ENV_CODE_DIRECTORY)?; + let py_script_name = easyenv::get_env_string_required(ENV_DOWNLOAD_SCRIPT_NAME)?; + let py_model_checkpoint = easyenv::get_env_string_required(ENV_MODEL_CHECKPOINT)?; + + let w2l_preprecess_command = Wav2LipPreprocessClient::new( + &py_code_directory, + &py_script_name, + &py_model_checkpoint, + ); + + let temp_directory = easyenv::get_env_string_or_default( + "DOWNLOAD_TEMP_DIR", + DEFAULT_TEMP_DIR); + + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "./scripts/download_internet_file.py"); + + let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script); + + let temp_directory = PathBuf::from(temp_directory); + + check_directory_exists(&temp_directory)?; + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + info!("Connecting to database..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + let common_env = CommonEnv::read_from_env()?; + + info!("Connecting to redis..."); + + let redis_manager = + RedisConnectionManager::new(common_env.redis_0_connection_string.deref())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + }; + + let badge_granter = BadgeGranter { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + firehose_publisher: firehose_publisher.clone(), // NB: Also safe + }; + + let downloader = Downloader { + download_temp_directory: temp_directory, + mysql_pool, + redis_pool, + public_bucket_client, + private_bucket_client, + download_script, + google_drive_downloader, + firehose_publisher, + badge_granter, + ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand {}, + ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand {}, + imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand {}, + w2l_processor: w2l_preprecess_command, + bucket_root_w2l_template_uploads: bucket_root.to_string(), + debug_download_sleep_millis: easyenv::get_env_num("DEBUG_DOWNLOAD_SLEEP_MILLIS", 0)?, + debug_face_detect_sleep_millis: easyenv::get_env_num("DEBUG_FACE_DETECT_SLEEP_MILLIS", 0)?, + job_batch_wait_millis: common_env.job_batch_wait_millis, + job_max_attempts: common_env.job_max_attempts as i32, + no_op_logger_millis: common_env.no_op_logger_millis, + debug_job_end_sleep_millis: common_env.debug_job_end_sleep_millis, + }; + + main_loop(downloader).await; + + Ok(()) +} + +const START_TIMEOUT_MILLIS : u64 = 500; +const INCREASE_TIMEOUT_MILLIS : u64 = 1000; + +async fn main_loop(downloader: Downloader) { + let mut error_timeout_millis = START_TIMEOUT_MILLIS; + + let mut noop_logger = NoOpLogger::new(downloader.no_op_logger_millis as i64); + + loop { + let num_records = 1; + + let query_result = query_w2l_template_upload_job_records( + &downloader.mysql_pool, + num_records) + .await; + + let jobs = match query_result { + Ok(jobs) => jobs, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if jobs.is_empty() { + noop_logger.log_after_awhile(); + + std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis)); + continue; + } + + let result = process_jobs(&downloader, jobs).await; + + match result { + Ok(_) => {}, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + } + + error_timeout_millis = START_TIMEOUT_MILLIS; // reset + + std::thread::sleep(Duration::from_millis(downloader.job_batch_wait_millis)); + } +} + +async fn process_jobs(downloader: &Downloader, jobs: Vec) -> AnyhowResult<()> { + for job in jobs.into_iter() { + let result = process_job(downloader, &job).await; + match result { + Ok(_) => {}, + Err(e) => { + warn!("Failure to process job: {:?}", e); + let failure_reason = ""; + let _r = mark_w2l_template_upload_job_failure( + &downloader.mysql_pool, + &job, + failure_reason, + downloader.job_max_attempts + ).await; + } + } + } + + Ok(()) +} + +#[derive(Deserialize)] +struct FileMetadata { + pub is_video: bool, + pub width: u32, + pub height: u32, + pub num_frames: u64, + pub fps: Option, + pub duration_millis: Option, + pub mimetype: Option, + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &str) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + +async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord) -> AnyhowResult<()> { + // TODO: 1. Mark processing. + // TODO: 2. Download. (DONE) + // TODO: 3. Process template with face detection (DONE) + // TODO: 4. Determine if picture or video + // TODO: 5. Take a screenshot/gif + // TODO: 6. Upload all (partially done). + // TODO: 7. Save record. (DONE) + // TODO: 8. Mark job done. (DONE) + + let mut redis = downloader.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_w2l_download( + &mut redis, + &job.token); + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + let lock_acquired = grab_job_lock_and_mark_pending(&downloader.mysql_pool, job).await?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {}", &job.id); + return Ok(()) + } + + // ==================== SETUP TEMP DIRS ==================== // + + let temp_dir = format!("temp_{}", job.id); + let temp_dir = TempDir::new(&temp_dir)?; + + // ==================== DOWNLOAD FILE ==================== // + + redis_logger.log_status("downloading media")?; + + let download_url = job.download_url.as_ref() + .map(|c| c.to_string()) + .unwrap_or("".to_string()); + + info!("Calling downloader..."); + let download_filename = match downloader.google_drive_downloader.download_file(&download_url, &temp_dir).await { + Ok(filename) => filename, + Err(e) => { + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + if downloader.debug_download_sleep_millis != 0 { + warn!("Debug sleep after download: {} ms", downloader.debug_download_sleep_millis); + thread::sleep(Duration::from_millis(downloader.debug_download_sleep_millis)); + } + + // ==================== PROCESS FACES ==================== // + + redis_logger.log_status("detecting faces")?; + + // This is the Python Pickle file with all the face frames. + // We'll retain it forever since it's expensive to compute. + let cached_faces_filename = format!("{}_detected_faces.pickle", &download_filename); + + // This is a file that we'll use to determine if the file is an image or video. + let output_metadata_filename = format!("{}_metadata.json", &download_filename); + + let is_image = false; // TODO: Don't always treat as video. + let spawn_process = false; + + let face_detect_result = downloader.w2l_processor.execute( + &download_filename, + &cached_faces_filename, + &output_metadata_filename, + is_image, + spawn_process); + + match face_detect_result { + Ok(_) => {}, // Ok to proceed. + Err(Wav2LipPreprocessError::UnknownError) => {} // Ok to proceed. + Err(Wav2LipPreprocessError::NoFacesDetected) => { + // Permanently fail. + warn!("Permanently failed due to no face detection!"); + + safe_delete_temp_directory(&temp_dir); + + mark_w2l_template_upload_job_permanently_dead( + &downloader.mysql_pool, + job, + "faces must be present in every frame" + ).await?; + return Ok(()); + } + } + + if downloader.debug_face_detect_sleep_millis != 0 { + warn!("Debug sleep after face detect: {} ms", downloader.debug_face_detect_sleep_millis); + thread::sleep(Duration::from_millis(downloader.debug_face_detect_sleep_millis)); + } + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + let video_or_image_path = PathBuf::from(&download_filename); + let cached_faces_path = PathBuf::from(&cached_faces_filename); + let output_metadata_path = PathBuf::from(&output_metadata_filename); + + info!("Checking that both files exist (original source + cached faces) ..."); + + check_file_exists(&video_or_image_path)?; + check_file_exists(&cached_faces_path)?; + check_file_exists(&output_metadata_path)?; + + let file_metadata = read_metadata_file(&output_metadata_filename)?; + + // ==================== BASE OBJECT NAMES BASED ON HASH ==================== // + + let private_bucket_hash = hash_file_sha2(&download_filename)?; + + info!("File hash: {}", private_bucket_hash); + + // Full path to video/image + let full_object_path = hash_to_bucket_path( + &private_bucket_hash, + Some(&downloader.bucket_root_w2l_template_uploads))?; + + // ==================== GENERATE VIDEO PREVIEWS ==================== // + + let mut maybe_image_preview_filename : Option = None; + let mut maybe_image_preview_object_name : Option = None; + + let mut maybe_video_preview_filename : Option = None; + let mut maybe_video_preview_object_name : Option = None; + + redis_logger.log_status("generating preview")?; + + if file_metadata.is_video { + let preview_filename = format!("{}_preview.webp", &download_filename); + + let preview_result = downloader.ffmpeg_video_preview_generator.execute( + &download_filename, + &preview_filename, + 500, + 500, + true, + false + ); + + if let Err(e) = preview_result { + safe_delete_temp_file(&download_filename); + safe_delete_temp_file(&preview_filename); + safe_delete_temp_directory(&temp_dir); + } + + let video_preview_path = PathBuf::from(&preview_filename); + check_file_exists(&video_preview_path)?; + + let preview_object_path = format!("{}_preview.webp", full_object_path); + maybe_video_preview_object_name = Some(preview_object_path); + maybe_video_preview_filename = Some(video_preview_path); + + } else { + let preview_filename = format!("{}_preview.webp", &download_filename); + + downloader.imagemagick_image_preview_generator.execute( + &download_filename, + &preview_filename, + 500, + 500, + false + )?; + + let image_preview_path = PathBuf::from(&preview_filename); + check_file_exists(&image_preview_path)?; + + let preview_object_path = format!("{}_preview.webp", full_object_path); + maybe_image_preview_object_name = Some(preview_object_path); + maybe_image_preview_filename = Some(image_preview_path); + } + + // ==================== UPLOAD TO BUCKETS ==================== // + + redis_logger.log_status("uploading results")?; + + info!("Image/video destination bucket path: {}", full_object_path); + + // Full path to cached faces + let full_object_path_cached_faces = format!("{}_detected_faces.pickle", full_object_path); + + info!("Cached faces destination bucket path: {}", full_object_path_cached_faces); + + info!("Uploading image/video..."); + + let original_mime_type = file_metadata.mimetype + .as_deref() + .unwrap_or("application/octet-stream"); + + downloader.private_bucket_client.upload_filename_with_content_type( + &full_object_path, + &video_or_image_path, + original_mime_type).await?; + + downloader.public_bucket_client.upload_filename_with_content_type( + &full_object_path, + &video_or_image_path, + original_mime_type).await?; + + safe_delete_temp_file(&video_or_image_path); + + info!("Uploading cached faces..."); + let path_copy: PathBuf = cached_faces_path.clone(); + downloader.private_bucket_client.upload_filename( + &full_object_path_cached_faces, + &path_copy).await?; + + safe_delete_temp_file(&path_copy); + + // TODO: Fix this ugh. + if let Some(image_preview_filename) = maybe_image_preview_filename.as_deref() { + if let Some(image_preview_object_name) = maybe_image_preview_object_name.as_deref() { + info!("Uploading image preview..."); + downloader.private_bucket_client.upload_filename_with_content_type( + &image_preview_object_name, + image_preview_filename, + "image/webp").await?; + + info!("Uploading image preview..."); + downloader.public_bucket_client.upload_filename_with_content_type( + &image_preview_object_name, + image_preview_filename, + "image/webp").await?; + } + + safe_delete_temp_file(&image_preview_filename); + } + + // TODO: Fix this ugh. + if let Some(video_preview_filename) = maybe_video_preview_filename.as_deref() { + if let Some(video_preview_object_name) = maybe_video_preview_object_name.as_deref() { + info!("Uploading video preview..."); + downloader.private_bucket_client.upload_filename_with_content_type( + &video_preview_object_name, + video_preview_filename, + "image/webp").await?; + + info!("Uploading video preview..."); + downloader.public_bucket_client.upload_filename_with_content_type( + &video_preview_object_name, + video_preview_filename, + "image/webp").await?; + } + + safe_delete_temp_file(&video_preview_filename); + } + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + let template_type = if file_metadata.is_video { "video" } else { "image" }; + + info!("Saving model record..."); + let (id, model_token) = insert_w2l_template( + &downloader.mysql_pool, + template_type, + job, + &private_bucket_hash, + &full_object_path, + &full_object_path_cached_faces, + maybe_image_preview_object_name.as_deref(), + maybe_video_preview_object_name.as_deref(), + file_metadata.file_size_bytes, + file_metadata.mimetype.as_deref(), + file_metadata.width, + file_metadata.height, + file_metadata.num_frames, + file_metadata.fps.unwrap_or(0.0f32), + file_metadata.duration_millis.unwrap_or(0)) + .await?; + + info!("Marking job complete..."); + + mark_w2l_template_upload_job_done( + &downloader.mysql_pool, + job, + true, + Some(&model_token) + ).await?; + + downloader.firehose_publisher.publish_w2l_template_upload_finished(&job.creator_user_token, &model_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + downloader.badge_granter.maybe_grant_w2l_template_uploads_badge(&job.creator_user_token) + .await + .map_err(|e| { + warn!("error maybe awarding badge: {:?}", e); + anyhow!("error maybe awarding badge") + })?; + + if downloader.debug_job_end_sleep_millis != 0 { + warn!("Debug sleep after job end: {} ms", downloader.debug_job_end_sleep_millis); + thread::sleep(Duration::from_millis(downloader.debug_job_end_sleep_millis)); + } + + redis_logger.log_status("done")?; + + info!("Job {} complete success! Downloaded, processed, and uploaded. Saved model record: {}", + job.id, id); + + Ok(()) +} diff --git a/crates/service/storyteller_web/src/w2l_inference_job.rs b/crates/service/storyteller_web/src/w2l_inference_job.rs new file mode 100644 index 0000000000..6c2ce2178d --- /dev/null +++ b/crates/service/storyteller_web/src/w2l_inference_job.rs @@ -0,0 +1,662 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate serde_derive; + +pub mod common_env; +pub mod job_queries; +pub mod script_execution; +pub mod util; + +use anyhow::anyhow; +use chrono::Utc; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use crate::common_env::CommonEnv; +use crate::job_queries::w2l_inference_job_queries::W2lInferenceJobRecord; +use crate::job_queries::w2l_inference_job_queries::get_w2l_template_by_token; +use crate::job_queries::w2l_inference_job_queries::grab_job_lock_and_mark_pending; +use crate::job_queries::w2l_inference_job_queries::insert_w2l_result; +use crate::job_queries::w2l_inference_job_queries::mark_w2l_inference_job_done; +use crate::job_queries::w2l_inference_job_queries::mark_w2l_inference_job_failure; +use crate::job_queries::w2l_inference_job_queries::query_w2l_inference_job_records; +use crate::script_execution::ffmpeg_generate_preview_image_command::FfmpegGeneratePreviewImageCommand; +use crate::script_execution::ffmpeg_generate_preview_video_command::FfmpegGeneratePreviewVideoCommand; +use crate::script_execution::google_drive_download_command::GoogleDriveDownloadCommand; +use crate::script_execution::imagemagick_generate_preview_image_command::ImagemagickGeneratePreviewImageCommand; +use crate::script_execution::wav2lip_inference_command::Wav2LipInferenceCommand; +use crate::util::buckets::bucket_client::BucketClient; +use crate::util::buckets::bucket_path_unifier::BucketPathUnifier; +use crate::util::buckets::bucket_paths::hash_to_bucket_path; +use crate::util::filesystem::check_directory_exists; +use crate::util::filesystem::check_file_exists; +use crate::util::filesystem::safe_delete_temp_directory; +use crate::util::filesystem::safe_delete_temp_file; +use crate::util::hashing::hash_file_sha2::hash_file_sha2; +use crate::util::noop_logger::NoOpLogger; +use crate::util::redis::redis_job_status_logger::RedisJobStatusLogger; +use crate::util::semi_persistent_cache_dir::SemiPersistentCacheDir; +use data_encoding::{HEXUPPER, HEXLOWER, HEXLOWER_PERMISSIVE}; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use log::{warn, info}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2::PooledConnection; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use ring::digest::{Context, Digest, SHA256}; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::fs::{File, metadata}; +use std::io::{BufReader, Read}; +use std::ops::Deref; +use std::path::{PathBuf, Path}; +use std::process::Command; +use std::thread; +use std::time::Duration; +use tempdir::TempDir; + +// Buckets (shared config) +const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; +const ENV_SECRET_KEY : &'static str = "SECRET_KEY"; +const ENV_REGION_NAME : &'static str = "REGION_NAME"; + +// Buckets (private data) +const ENV_PRIVATE_BUCKET_NAME : &'static str = "W2L_PRIVATE_DOWNLOAD_BUCKET_NAME"; +// Buckets (public data) +const ENV_PUBLIC_BUCKET_NAME : &'static str = "W2L_PUBLIC_DOWNLOAD_BUCKET_NAME"; + +// Where models and other assets get downloaded to. +const ENV_SEMIPERSISTENT_CACHE_DIR : &'static str = "SEMIPERSISTENT_CACHE_DIR"; + +// Python code +const ENV_CODE_DIRECTORY : &'static str = "W2L_CODE_DIRECTORY"; +const ENV_MODEL_CHECKPOINT : &'static str = "W2L_MODEL_CHECKPOINT"; +const ENV_INFERENCE_SCRIPT_NAME : &'static str = "W2L_INFERENCE_SCRIPT_NAME"; + +const DEFAULT_TEMP_DIR: &'static str = "/tmp"; + +struct Inferencer { + pub download_temp_directory: PathBuf, + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub private_bucket_client: BucketClient, + pub public_bucket_client: BucketClient, + + pub firehose_publisher: FirehosePublisher, + + pub bucket_path_unifier: BucketPathUnifier, + pub semi_persistent_cache: SemiPersistentCacheDir, + + pub google_drive_downloader: GoogleDriveDownloadCommand, + pub w2l_inference: Wav2LipInferenceCommand, + //pub ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand, + //pub ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand, + //pub imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand, + + // Command to run + pub inference_script: String, + pub w2l_model_filename: String, + pub w2l_end_bump_filename: String, + + // Sleep between batches + pub job_batch_wait_millis: u64, + + // How long to wait between log lines + pub no_op_logger_millis: u64, + + // Max job attempts before failure. + // NB: This is an i32 so we don't need to convert to db column type. + pub job_max_attempts: i32, + + // Temporary for debugging + // Arbitrary timeouts can be inserted so we can exec in and poke around. + pub debug_job_end_sleep_millis: u64, +} + +#[tokio::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + // NB: Do not check this secrets-containing dotenv file into VCS. + // This file should only contain *development* secrets, never production. + let _ = dotenv::from_filename(".env-secrets").ok(); + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("w2l-inference-job".to_string()); + + info!("Hostname: {}", &server_hostname); + + // Bucket stuff (shared) + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + + // Private and Public Buckets + let private_bucket_name = easyenv::get_env_string_required(ENV_PRIVATE_BUCKET_NAME)?; + let public_bucket_name = easyenv::get_env_string_required(ENV_PUBLIC_BUCKET_NAME)?; + + let private_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &private_bucket_name, + None, + )?; + + let public_bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &public_bucket_name, + None, + )?; + + let py_code_directory = easyenv::get_env_string_required(ENV_CODE_DIRECTORY)?; + let py_script_name = easyenv::get_env_string_required(ENV_INFERENCE_SCRIPT_NAME)?; + let py_model_checkpoint = easyenv::get_env_string_required(ENV_MODEL_CHECKPOINT)?; + + let w2l_inference_command = Wav2LipInferenceCommand::new( + &py_code_directory, + &py_script_name, + &py_model_checkpoint, + ); + + let temp_directory = easyenv::get_env_string_or_default( + "DOWNLOAD_TEMP_DIR", + DEFAULT_TEMP_DIR); + + // TODO: In the future, we may want to enable downloading images or audio files. + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "./scripts/download_internet_file.py"); + let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script); + + let temp_directory = PathBuf::from(temp_directory); + + check_directory_exists(&temp_directory)?; + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + info!("Connecting to database..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + let common_env = CommonEnv::read_from_env()?; + + info!("Connecting to redis..."); + + let redis_manager = + RedisConnectionManager::new(common_env.redis_0_connection_string.deref())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let inference_script = "TODO".to_string(); + + let persistent_cache_path = easyenv::get_env_string_or_default( + ENV_SEMIPERSISTENT_CACHE_DIR, + "/tmp"); + + let semi_persistent_cache = SemiPersistentCacheDir::configured_root(&persistent_cache_path); + + info!("Creating pod semi-persistent cache dirs..."); + semi_persistent_cache.create_w2l_model_path()?; + semi_persistent_cache.create_w2l_end_bump_path()?; + semi_persistent_cache.create_w2l_face_template_path()?; + semi_persistent_cache.create_w2l_template_media_path()?; + + let w2l_model_filename = easyenv::get_env_string_or_default( + "W2L_MODEL_FILENAME", "wav2lip_gan.pth"); + + let w2l_end_bump_filename = easyenv::get_env_string_or_default( + "W2L_END_BUMP_FILENAME", "vocodes-short-end-bump.mp4"); + + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: MySqlPool is clone/send/sync safe + }; + + let common_env = CommonEnv::read_from_env()?; + + let inferencer = Inferencer { + download_temp_directory: temp_directory, + mysql_pool, + redis_pool, + public_bucket_client, + private_bucket_client, + inference_script, + google_drive_downloader, + //ffmpeg_image_preview_generator: FfmpegGeneratePreviewImageCommand {}, + //ffmpeg_video_preview_generator: FfmpegGeneratePreviewVideoCommand {}, + //imagemagick_image_preview_generator: ImagemagickGeneratePreviewImageCommand {}, + w2l_inference: w2l_inference_command, + bucket_path_unifier: BucketPathUnifier::default_paths(), + semi_persistent_cache, + firehose_publisher, + w2l_model_filename, + w2l_end_bump_filename, + job_batch_wait_millis: common_env.job_batch_wait_millis, + job_max_attempts: common_env.job_max_attempts as i32, + no_op_logger_millis: common_env.no_op_logger_millis, + debug_job_end_sleep_millis: common_env.debug_job_end_sleep_millis, + }; + + main_loop(inferencer).await; + + Ok(()) +} + +const START_TIMEOUT_MILLIS : u64 = 500; +const INCREASE_TIMEOUT_MILLIS : u64 = 1000; + +async fn main_loop(inferencer: Inferencer) { + let mut error_timeout_millis = START_TIMEOUT_MILLIS; + + let mut noop_logger = NoOpLogger::new(inferencer.no_op_logger_millis as i64); + + loop { + let num_records = 1; + + let query_result = query_w2l_inference_job_records( + &inferencer.mysql_pool, + num_records) + .await; + + let jobs = match query_result { + Ok(jobs) => jobs, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if jobs.is_empty() { + noop_logger.log_after_awhile(); + + std::thread::sleep(Duration::from_millis(inferencer.job_batch_wait_millis)); + continue; + } + + let result = process_jobs(&inferencer, jobs).await; + + match result { + Ok(_) => {}, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + } + + error_timeout_millis = START_TIMEOUT_MILLIS; // reset + + std::thread::sleep(Duration::from_millis(inferencer.job_batch_wait_millis)); + } +} + +async fn process_jobs(inferencer: &Inferencer, jobs: Vec) -> AnyhowResult<()> { + for job in jobs.into_iter() { + let result = process_job(inferencer, &job).await; + match result { + Ok(_) => {}, + Err(e) => { + warn!("Failure to process job: {:?}", e); + let failure_reason = ""; + let _r = mark_w2l_inference_job_failure( + &inferencer.mysql_pool, + &job, + failure_reason, + inferencer.job_max_attempts + ).await; + } + } + } + + Ok(()) +} + +#[derive(Deserialize)] +struct FileMetadata { + pub is_video: bool, + pub width: u32, + pub height: u32, + //pub num_frames: u64, + //pub fps: Option, + pub duration_millis: Option, + pub mimetype: Option, + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + +async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> AnyhowResult<()> { + + // TODO 1. Mark Processing + // + // TODO: 2. Check if w2l model is downloaded / download it to a stable cache location (DONE) + // TODO: 3. Check if w2l template faces are downloaded and download it (done) + // TODO: 4. Download user audio (done) + + // TODO: 5. Process Inference + + // TODO 6. Upload result + // TODO 7. Save record + // TODO 8. Mark job done + + let mut redis = inferencer.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_w2l_inference( + &mut redis, + &job.inference_job_token); + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + let lock_acquired = grab_job_lock_and_mark_pending(&inferencer.mysql_pool, job).await?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {}", &job.id); + return Ok(()) + } + + // ==================== CONFIRM OR DOWNLOAD W2L MODEL ==================== // + + let model_filename = inferencer.w2l_model_filename.clone(); + let model_fs_path = inferencer.semi_persistent_cache.w2l_model_path(&model_filename); + + if !model_fs_path.exists() { + warn!("Model file does not exist: {:?}", &model_fs_path); + + redis_logger.log_status("downloading model")?; + + let model_object_path = inferencer.bucket_path_unifier + .w2l_pretrained_models_path(&model_filename); + + info!("Download from bucket path: {:?}", &model_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &model_object_path, + &model_fs_path + ).await?; + + info!("Downloaded model from bucket!"); + } + + // ==================== CONFIRM OR DOWNLOAD W2L END BUMP ==================== // + + let end_bump_filename = inferencer.w2l_end_bump_filename.clone(); + let end_bump_fs_path = inferencer.semi_persistent_cache.w2l_end_bump_path(&end_bump_filename); + + if !end_bump_fs_path.exists() { + warn!("End bump file does not exist: {:?}", &end_bump_fs_path); + + redis_logger.log_status("downloading assets")?; + + let end_bump_object_path = inferencer.bucket_path_unifier + .end_bump_video_for_w2l_path(&end_bump_filename); + + info!("Download from bucket path: {:?}", &end_bump_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &end_bump_object_path, + &end_bump_fs_path + ).await?; + + info!("Downloaded end bump from bucket!"); + } + + // ==================== LOOK UP TEMPLATE RECORD ==================== // + + let template_token = match &job.maybe_w2l_template_token { + Some(token) => token.to_string(), + None => { + warn!("non-template token based inference not yet supported"); + return Err(anyhow!("non-template token based inference not yet supported")) + }, + }; + + info!("Looking up w2l template by token: {}", &template_token); + + let query_result = get_w2l_template_by_token(&inferencer.mysql_pool, &template_token).await?; + + let w2l_template = match query_result { + Some(template) => template, + None => { + warn!("W2L Template not found: {}", &template_token); + return Err(anyhow!("Template not found!")) + }, + }; + + // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE AUDIO OR VIDEO ==================== // + + // Template is based on the `private_bucket_hash`: + // - private_bucket_hash: 1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 + // - private_bucket_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 + // - private_bucket_cached_faces_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_detected_faces.pickle + // - maybe_public_bucket_preview_image_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_preview.webp + + let template_media_fs_path = inferencer.semi_persistent_cache.w2l_template_media_path( + &w2l_template.private_bucket_hash); + + if !template_media_fs_path.exists() { + info!("W2L template media file does not exist: {:?}", &template_media_fs_path); + + redis_logger.log_status("downloading template")?; + + let template_media_object_path = inferencer.bucket_path_unifier + .media_templates_for_w2l_path(&w2l_template.private_bucket_hash); + + info!("Download from template media path: {:?}", &template_media_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &template_media_object_path, + &template_media_fs_path + ).await?; + + info!("Downloaded template media from bucket!"); + } + + // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE FACE ==================== // + + // Template is based on the `private_bucket_hash`: + // - private_bucket_hash: 1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 + // - private_bucket_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 + // - private_bucket_cached_faces_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_detected_faces.pickle + // - maybe_public_bucket_preview_image_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_preview.webp + + let face_template_fs_path = inferencer.semi_persistent_cache.w2l_face_template_path( + &w2l_template.private_bucket_hash); + + if !face_template_fs_path.exists() { + info!("W2L face template file does not exist: {:?}", &face_template_fs_path); + + redis_logger.log_status("downloading template metadata")?; + + let face_template_object_path = inferencer.bucket_path_unifier + .precomputed_faces_for_w2l_path(&w2l_template.private_bucket_hash); + + info!("Download from face template path: {:?}", &face_template_object_path); + + inferencer.private_bucket_client.download_file_to_disk( + &face_template_object_path, + &face_template_fs_path + ).await?; + + info!("Downloaded face template from bucket!"); + } + + // ==================== DOWNLOAD USER AUDIO ==================== // + + redis_logger.log_status("downloading audio")?; + + let temp_dir = format!("temp_{}", job.id); + let temp_dir = TempDir::new(&temp_dir)?; // NB: Exists until it goes out of scope. + + let audio_bucket_hash = match &job.maybe_public_audio_bucket_hash { + Some(l) => l.clone(), + None => { + warn!("Only W2L jobs with user-uploaded audio are supported right now"); + return Err(anyhow!("Only W2L jobs with user-uploaded audio are supported right now")) + }, + }; + + let audio_fs_path = temp_dir.path().join(&audio_bucket_hash); + + let audio_object_path = inferencer.bucket_path_unifier + .user_audio_for_w2l_inference_path(&audio_bucket_hash); + + inferencer.private_bucket_client.download_file_to_disk( + &audio_object_path, + &audio_fs_path + ).await?; + + + // ==================== RUN INFERENCE ==================== // + + redis_logger.log_status("executing")?; + + let output_video_fs_path = temp_dir.path().join("output.mp4"); + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + + let is_image = w2l_template.template_type.contains("image"); + + info!("Is image? {}", is_image); + info!("Running W2L inference..."); + + let inference_result = inferencer.w2l_inference.execute( + &model_fs_path, + &audio_fs_path, + &end_bump_fs_path, + &template_media_fs_path, + &face_template_fs_path, + &output_metadata_fs_path, + &output_video_fs_path, + false, + false + ); + + if let Err(e) = inference_result { + safe_delete_temp_file(&audio_fs_path); + safe_delete_temp_file(&output_video_fs_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + + info!("Output filename: {:?}", &output_video_fs_path); + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that output files exist..."); + + check_file_exists(&output_video_fs_path)?; + check_file_exists(&output_metadata_fs_path)?; + + let file_metadata = read_metadata_file(&output_metadata_fs_path)?; + + safe_delete_temp_file(&output_metadata_fs_path); + + // ==================== UPLOAD TO BUCKETS ==================== // + + redis_logger.log_status("uploading result")?; + + let result_object_path = inferencer.bucket_path_unifier.w2l_inference_video_output_path( + &job.inference_job_token); + + info!("Image/video destination bucket path: {:?}", &result_object_path); + + info!("Uploading image/video..."); + + let original_mime_type = file_metadata.mimetype + .as_deref() + .unwrap_or("application/octet-stream"); + + let upload_result = inferencer.public_bucket_client.upload_filename_with_content_type( + &result_object_path, + &output_video_fs_path, + original_mime_type) + .await; + + if let Err(e) = upload_result { + safe_delete_temp_file(&audio_fs_path); + safe_delete_temp_file(&output_video_fs_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + + safe_delete_temp_file(&output_video_fs_path); + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving w2l inference record..."); + let (id, inference_result_token) = insert_w2l_result( + &inferencer.mysql_pool, + job, + &result_object_path, + file_metadata.file_size_bytes, + file_metadata.mimetype.as_deref(), + file_metadata.width, + file_metadata.height, + file_metadata.duration_millis.unwrap_or(0)) + .await?; + + info!("Marking job complete..."); + + mark_w2l_inference_job_done( + &inferencer.mysql_pool, + job, + true, + Some(&inference_result_token) + ).await?; + + inferencer.firehose_publisher.w2l_inference_finished( + job.maybe_creator_user_token.as_deref(), + &job.inference_job_token, + &inference_result_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + redis_logger.log_status("done")?; + + if inferencer.debug_job_end_sleep_millis != 0 { + warn!("Debug sleep after job end: {} ms", inferencer.debug_job_end_sleep_millis); + thread::sleep(Duration::from_millis(inferencer.debug_job_end_sleep_millis)); + } + + info!("Job {} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}", + job.id, id); + + Ok(()) +} diff --git a/crates/service/twitch_pubsub_subscriber/Cargo.toml b/crates/service/twitch_pubsub_subscriber/Cargo.toml new file mode 100644 index 0000000000..c9c752bdfc --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/Cargo.toml @@ -0,0 +1,119 @@ +[package] +name = "twitch-pubsub-subscriber" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[[bin]] +name = "twitch-pubsub-subscriber" +path = "src/main.rs" + +[dependencies] + +# Internal +config = { path = "../../lib/config" } +container_common = { path = "../../lib/container_common" } +database_queries = { path = "../../lib/database_queries" } +http_server_common = { path = "../../lib/http_server_common" } +redis_common = { path = "../../lib/redis_common" } +twitch_common = { path = "../../lib/twitch_common" } + +# External +anyhow = "1.0.40" +base64 = "0.13.0" +chrono = { version = "0.4.19", features = ["serde"] } +derive_more = "0.99.14" +dotenv = "0.15.0" +easyenv = "0.0.6" +futures-timer = "3.0.2" +hex = "0.4.3" +hostname = "0.3.1" +lazy_static = "1.4.0" +lfu = "0.2.5" +log = "0.4.14" +rand = "0.8.3" +regex = "1.5.4" +time = "0.2.26" +uuid = { version = "0.8.2", features = ["v4"] } + +# Runtime +# NB: We switched to Tokio from `async-std` for broader ecosystem support. +tokio = { version = "1.10.0", features = ["rt-multi-thread", "macros", "time"] } + +# Futures +futures = { version = "0.3.17", features = ["thread-pool"] } # feature for ThreadPool +futures-util = "0.3.17" +futures_old_for_limiter = { package = "futures", version = "0.1.29" } # Older version of futures for `limiter` crate + +# Threading +crossbeam = "0.8.1" + +# MySql Database +# "offline" supports typesafe builds without a live database +# "rustls" doesn't appear to work with IP addresses, per: +# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828 +#sqlx = { version = "0.5", features = [ "mysql", "runtime-async-std-native-tls", "offline", "chrono" ] } +sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] } +#sqlx = { version = "0.4", features = [ "mysql", "runtime-tokio-native-tls", "offline", "chrono" ] } + +# Redis Database +limitation = "0.1.1" # Redis rate limiter +r2d2_redis = "0.14.0" + +# Twitch +# NB "unsupported" pubsub features may disappear from the Twitch API at any time. It's an unofficial API. +twitch_api2 = { version = "0.6.0-rc.2", features = ["helix", "reqwest_client", "pubsub", "chrono", "unsupported"] } +twitch_oauth2 = "0.6.0-rc.2" + +# HTTP +reqwest = { version = "0.11.4" } +native-tls = "0.2.8" + +# In-memory caching +lru_time_cache = "0.11.11" + +# Web +actix = "0.12.0" +actix-cors = "0.5.4" +actix-files= "0.6.0-beta.8" +actix-http = "3.0.0-beta.11" +actix-rt = "2.4.0" +actix-service = "2.0.1" +actix-utils = "3.0.0" +actix-web = "4.0.0-beta.10" +actix-web-actors = "4.0.0-beta.7" +jwt = "0.13.0" + +# Websockets +tokio-tungstenite = { version = "0.15.0", features = ["rustls", "rustls-tls"] } + +# NB: Hyper is just for HTTP client use in jobs. I'd use Actix client, but I need tokio runtime +# NB: Furthermore, the Tokio runtime is old, so I have to use an older Hyper. +hyper = { version = "0.13.10" } + +# Crypto +bcrypt = "0.8" # password hashes +hmac = "0.10.0" # cookies +magic-crypt = "3.1.8" # encrypting sort IDs for the frontend so we don't leak info +md-5 = "0.9.1" # gravatar +sha2 = "0.9.5" # file hashes, inference text hashes, etc. + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" +toml = "0.5.8" + +# Text processing +pulldown-cmark = "0.8.0" + +# Temporary (just to ask for oauth credentials during dev) +rpassword = "5.0.1" + +[build-dependencies] +# NB: Only using diesel for migrations support +diesel = { version = "1.4.6", features = ["mysql"] } diff --git a/crates/service/twitch_pubsub_subscriber/src/main.rs b/crates/service/twitch_pubsub_subscriber/src/main.rs new file mode 100644 index 0000000000..eca1f18f65 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/main.rs @@ -0,0 +1,135 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate lazy_static; +#[macro_use] extern crate magic_crypt; +#[macro_use] extern crate serde_derive; + +pub mod threads; +pub mod twitch; +pub mod util; + +use config::shared_constants::{DEFAULT_RUST_LOG, DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING, DEFAULT_MYSQL_CONNECTION_STRING}; +use container_common::anyhow_result::AnyhowResult; +use crate::threads::listen_for_active_obs_sessions::listen_for_active_obs_sessions_thread::ListenForActiveObsSessionThread; +use crate::twitch::oauth::oauth_token_refresher::OauthTokenRefresher; +use crate::twitch::websocket_client::TwitchWebsocketClient; +use futures::executor::{ThreadPool, ThreadPoolBuilder}; +use futures::task::SpawnExt; +use log::info; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2::Pool; +use r2d2_redis::r2d2; +use redis_common::redis_keys::RedisKeys; +use sqlx::mysql::MySqlPoolOptions; +use std::sync::Arc; +use std::thread::sleep; +use std::time::Duration; +use tokio::runtime::{Builder, Runtime}; +use twitch_common::twitch_secrets::TwitchSecrets; +use twitch_oauth2::{ClientId, ClientSecret, RefreshToken, AccessToken}; + +#[tokio::main] +pub async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + // NB: Do not check this secrets-containing dotenv file into VCS. + // This file should only contain *development* secrets, never production. + let _ = dotenv::from_filename(".env-secrets").ok(); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("twitch-pubsub-job".to_string()); + + info!("Reading Twitch secrets..."); + + let secrets = TwitchSecrets::from_env()?; + let client_id = ClientId::new(&secrets.app_client_id); + let client_secret = ClientSecret::new(&secrets.app_client_secret); + + let oauth_token_refresher = OauthTokenRefresher::from_secrets(client_id, client_secret)?; + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + let redis_connection_string = + easyenv::get_env_string_or_default( + "REDIS_1_URL", + DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING); + + info!("Connecting to mysql..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + let mysql_pool = Arc::new(mysql_pool); + + info!("Connecting to redis..."); + + let redis_manager = RedisConnectionManager::new(redis_connection_string.clone())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let redis_pool = Arc::new(redis_pool); + + info!("Connecting to pubsub redis..."); + + let redis_pubsub_manager = RedisConnectionManager::new(redis_connection_string.clone())?; + + let redis_pubsub_pool = r2d2::Pool::builder() + .build(redis_pubsub_manager)?; + + let redis_pubsub_pool = Arc::new(redis_pubsub_pool); + + info!("Creating thread pool..."); + + // https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html + let runtime = Builder::new_multi_thread() + .worker_threads(2) + .thread_name("twitch-pubsub-") + .thread_stack_size(3 * 1024 * 1024) + .enable_all() + .build()?; + + let runtime = Arc::new(runtime); + + // https://docs.rs/tokio/latest/tokio/runtime/struct.Builder.html + let runtime_2 = Builder::new_multi_thread() + .worker_threads(8) + .thread_name("twitch-pubsub-") + .thread_stack_size(3 * 1024 * 1024) + .enable_all() + .build()?; + + let runtime_2 = Arc::new(runtime_2); + + info!("Thread pool created"); + + let thread = ListenForActiveObsSessionThread::new( + server_hostname.to_string(), + oauth_token_refresher, + mysql_pool, + redis_pool, + redis_pubsub_pool, + runtime_2, + ); + + runtime.spawn(thread.start_thread()); + + loop { + sleep(Duration::from_millis(60_000)); + } +} diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/listen_for_active_obs_sessions/listen_for_active_obs_sessions_thread.rs b/crates/service/twitch_pubsub_subscriber/src/threads/listen_for_active_obs_sessions/listen_for_active_obs_sessions_thread.rs new file mode 100644 index 0000000000..ea25c146c2 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/listen_for_active_obs_sessions/listen_for_active_obs_sessions_thread.rs @@ -0,0 +1,110 @@ +use container_common::thread::thread_id::ThreadId; +use crate::threads::twitch_pubsub_user_subscriber::twitch_pubsub_user_subscriber_thread::TwitchPubsubUserSubscriberThread; +use crate::twitch::oauth::oauth_token_refresher::OauthTokenRefresher; +use log::error; +use log::info; +use log::warn; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use redis_common::payloads::lease_payload::LeasePayload; +use redis_common::payloads::obs_active_payload::ObsActivePayload; +use redis_common::redis_keys::RedisKeys; +use redis_common::shared_constants::LEASE_TIMEOUT_SECONDS; +use sqlx::MySql; +use std::sync::Arc; +use std::thread::sleep; +use std::time::Duration; +use tokio::runtime::Runtime; +use twitch_common::twitch_user_id::TwitchUserId; + +// TODO: Apart from error handling, this looks mostly good. + + +pub struct ListenForActiveObsSessionThread { + server_hostname: String, + oauth_token_refresher: OauthTokenRefresher, + mysql_pool: Arc>, + redis_pool: Arc>, + redis_pubsub_pool: Arc>, + runtime: Arc, +} + +impl ListenForActiveObsSessionThread { + pub fn new( + server_hostname: String, + oauth_token_refresher: OauthTokenRefresher, + mysql_pool: Arc>, + redis_pool: Arc>, + redis_pubsub_pool: Arc>, + runtime: Arc, + ) -> Self { + Self { + server_hostname, + oauth_token_refresher, + mysql_pool, + redis_pool, + redis_pubsub_pool, + runtime, + } + } + + pub async fn start_thread(mut self) { + // TODO: ERROR HANDLING + // TODO: ERROR HANDLING + // TODO: ERROR HANDLING + + // Setup PubSub. + + let mut pubsub_pool = self.redis_pubsub_pool.get().unwrap(); + let mut pubsub = pubsub_pool.as_pubsub(); + + pubsub.subscribe(RedisKeys::obs_active_sessions_topic()).unwrap(); + + loop { + let message = pubsub.get_message().unwrap(); + let payload : String = message.get_payload().unwrap(); + let payload = ObsActivePayload::from_json_str(&payload).unwrap(); + + let twitch_user_id = TwitchUserId::from_str(&payload.twitch_user_id).unwrap(); + + // See if we have a lease already... + + let mut redis = self.redis_pool.get().unwrap(); + let lease_key = RedisKeys::twitch_pubsub_lease(&payload.twitch_user_id); + let lease_value : Option = redis.get(&lease_key).unwrap(); + + if let Some(value) = lease_value.as_deref() { + let lease = LeasePayload::deserialize(value).unwrap(); + continue; + } + + info!("No existing lease for {:?}...", &lease_key); + + let thread_id = ThreadId::random_id(); + let lease = LeasePayload::from_thread_id(&self.server_hostname, &thread_id); + + let serialized = lease.serialize(); + let _v : Option = redis.set_ex( + &lease_key, + &serialized, + LEASE_TIMEOUT_SECONDS + ).unwrap(); + + // Then launch the thread... + + let redis_pool2 = self.redis_pool.clone(); + let mysql_pool2 = self.mysql_pool.clone(); + + let thread = TwitchPubsubUserSubscriberThread::new( + twitch_user_id, + self.oauth_token_refresher.clone(), + mysql_pool2, + redis_pool2, + &self.server_hostname, + thread_id.clone()); + + self.runtime.spawn(thread.start_thread()); + } + } +} diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/listen_for_active_obs_sessions/mod.rs b/crates/service/twitch_pubsub_subscriber/src/threads/listen_for_active_obs_sessions/mod.rs new file mode 100644 index 0000000000..dfaf27e527 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/listen_for_active_obs_sessions/mod.rs @@ -0,0 +1 @@ +pub mod listen_for_active_obs_sessions_thread; \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/mod.rs b/crates/service/twitch_pubsub_subscriber/src/threads/mod.rs new file mode 100644 index 0000000000..4a9aa6ac56 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/mod.rs @@ -0,0 +1,2 @@ +pub mod listen_for_active_obs_sessions; +pub mod twitch_pubsub_user_subscriber; \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/bits_event_handler.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/bits_event_handler.rs new file mode 100644 index 0000000000..ee95769659 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/bits_event_handler.rs @@ -0,0 +1,133 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::collections::random_from_vec::random_from_vec; +use container_common::numerics::signed_to_unsigned::i64_to_unsigned_zeroing_negatives; +use crate::threads::twitch_pubsub_user_subscriber::subscriber_preferences_cache::{TwitchPubsubCachedState, TwitchEventRuleLight}; +use crate::threads::twitch_pubsub_user_subscriber::tts_writer::TtsWriter; +use database_queries::column_types::twitch_event_category::TwitchEventCategory; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_pubsub::insert_bits::TwitchPubsubBitsInsertBuilder; +use database_queries::queries::twitch::twitch_pubsub::insert_channel_points::TwitchPubsubChannelPointsInsertBuilder; +use log::info; +use log::error; +use r2d2_redis::r2d2; +use rand::seq::SliceRandom; +use sqlx::MySql; +use std::sync::RwLock; +use std::sync::{Arc, PoisonError, RwLockReadGuard}; +use twitch_api2::pubsub::TopicData; +use twitch_api2::pubsub::channel_bits::{ChannelBitsEventsV2, ChannelBitsEventsV2Reply, BitsEventData}; + +pub struct BitsEventHandler { + twitch_subscriber_state: Arc>, + mysql_pool: Arc>, + tts_writer: Arc, +} + +impl BitsEventHandler { + pub fn new( + twitch_subscriber_state: Arc>, + mysql_pool: Arc>, + tts_writer: Arc, + ) -> Self { + Self { + twitch_subscriber_state, + mysql_pool, + tts_writer + } + } + + pub async fn handle(&self, topic: ChannelBitsEventsV2, reply: Box) -> AnyhowResult<()> { + match *reply { + ChannelBitsEventsV2Reply::BitsEvent { data, message_id, version, is_anonymous } => { + self.handle_bits_event(&data).await?; + } + _ => {} // NB: There are no other enum variants as of 2022-02-21 + } + Ok(()) + } + + async fn handle_bits_event(&self, data: &BitsEventData) -> AnyhowResult<()> { + let maybe_rule = self.find_matching_rule(data)?; + if let Some(rule) = maybe_rule { + info!("Bits Rule matched: {}", &rule.token); + self.handle_matched_rule(&rule, data).await?; + self.report_event_for_analytics(&data).await?; // Report event for analytics + } + Ok(()) + } + + async fn report_event_for_analytics(&self, data: &BitsEventData) -> AnyhowResult<()> { + let user_id = data.user_id.to_string(); + let user_name = data.user_name.to_string(); + let mut event_builder = TwitchPubsubBitsInsertBuilder::new(); + let mut event_builder = event_builder.set_sender_twitch_user_id(&user_id) + .set_sender_twitch_username(&user_name) + .set_destination_channel_id(&data.channel_id.to_string()) + .set_destination_channel_name(&data.channel_name.to_string()) + .set_bits_used(data.bits_used as u64) + .set_total_bits_used(data.total_bits_used as u64) + .set_is_anonymous(data.is_anonymous) + .set_chat_message(&data.chat_message); + event_builder.insert(&self.mysql_pool).await?; + Ok(()) + } + + fn find_matching_rule(&self, data: &BitsEventData) -> AnyhowResult> { + return match self.twitch_subscriber_state.read() { + Err(e) => { Err(anyhow!("Lock error: {:?}", e)) }, + Ok(state) => { + info!("Checking bits event against {} rules...", state.event_rules.len()); + let maybe_rule = state.event_rules.iter() + .filter(|rule| !rule.rule_is_disabled) + .filter(|rule| rule.event_category.eq(&TwitchEventCategory::Bits)) + .find(|rule| { + match rule.event_match_predicate { + EventMatchPredicate::NotSet {} => false, // Not set + EventMatchPredicate::ChannelPointsRewardNameExactMatch { .. } => false, // Wrong type + EventMatchPredicate::BitsCheermoteNameExactMatch { ref cheermote_name } => { + let contains_cheermote = data.chat_message.to_lowercase().contains(&cheermote_name.to_lowercase()); + contains_cheermote + }, + EventMatchPredicate::BitsCheermotePrefixSpendThreshold { ref cheermote_prefix, minimum_bits_spent } => { + let contains_cheermote = data.chat_message.to_lowercase().contains(&cheermote_prefix.to_lowercase()); + let spent = i64_to_unsigned_zeroing_negatives(data.bits_used); + contains_cheermote && spent >= minimum_bits_spent + }, + EventMatchPredicate::BitsSpendThreshold { minimum_bits_spent } => { + let spent = i64_to_unsigned_zeroing_negatives(data.bits_used); + spent >= minimum_bits_spent + }, + } + }) + .map(|rule| rule.clone()); + Ok(maybe_rule) + } + }; + } + + //noinspection DuplicatedCode + async fn handle_matched_rule(&self, rule: &TwitchEventRuleLight, data: &BitsEventData) -> AnyhowResult<()> { + match rule.event_response { + EventResponse::NotSet {} => { + info!("Empty event response."); + return Ok(()) + }, + EventResponse::TtsSingleVoice { ref tts_model_token } => { + self.tts_writer.write_tts_with_model(&data.chat_message, tts_model_token).await?; + } + EventResponse::TtsRandomVoice { ref tts_model_tokens } => { + let maybe_token = random_from_vec(tts_model_tokens); + if let Some(token) = maybe_token { + self.tts_writer.write_tts_with_model(&data.chat_message, token.as_str()) + .await?; + } + } + EventResponse::TtsCommandPresets { .. } => {} // TODO + EventResponse::TtsCommandCustom { .. } => {} // TODO + } + + Ok(()) + } +} \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/channel_points_event_handler.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/channel_points_event_handler.rs new file mode 100644 index 0000000000..aafa50bddf --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/channel_points_event_handler.rs @@ -0,0 +1,142 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::collections::random_from_vec::random_from_vec; +use crate::threads::twitch_pubsub_user_subscriber::subscriber_preferences_cache::{TwitchPubsubCachedState, TwitchEventRuleLight}; +use crate::threads::twitch_pubsub_user_subscriber::tts_writer::TtsWriter; +use database_queries::column_types::twitch_event_category::TwitchEventCategory; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::twitch::twitch_pubsub::insert_channel_points::TwitchPubsubChannelPointsInsertBuilder; +use log::info; +use r2d2_redis::r2d2; +use sqlx::MySql; +use std::sync::Arc; +use std::sync::RwLock; +use twitch_api2::pubsub::TopicData; +use twitch_api2::pubsub::channel_points::{ChannelPointsChannelV1, ChannelPointsChannelV1Reply, Redemption}; + +pub struct ChannelPointsEventHandler { + twitch_subscriber_state: Arc>, + mysql_pool: Arc>, + tts_writer: Arc, +} + +impl ChannelPointsEventHandler { + pub fn new( + twitch_subscriber_state: Arc>, + mysql_pool: Arc>, + tts_writer: Arc, + ) -> Self { + Self { + twitch_subscriber_state, + mysql_pool, + tts_writer + } + } + + // NB: &mut is for Redis pool in downstream write_tts. + pub async fn handle(&self, topic: ChannelPointsChannelV1, reply: Box) -> AnyhowResult<()> { + match *reply { + // Unimplemented + ChannelPointsChannelV1Reply::CustomRewardUpdated { .. } => {} + ChannelPointsChannelV1Reply::RedemptionStatusUpdate { .. } => {} + ChannelPointsChannelV1Reply::UpdateRedemptionStatusesFinished { .. } => {} + ChannelPointsChannelV1Reply::UpdateRedemptionStatusProgress { .. } => {} + // Implemented + ChannelPointsChannelV1Reply::RewardRedeemed { timestamp, redemption } => { + self.handle_reward_redeemed_event(&redemption).await?; + } + _ => {}, + } + Ok(()) + } + + async fn handle_reward_redeemed_event(&self, redemption: &Redemption) -> AnyhowResult<()> { + let maybe_rule = self.find_matching_rule(redemption)?; + if let Some(rule) = maybe_rule { + info!("Channel Points rule matched: {}", &rule.token); + self.handle_matched_rule(&rule, redemption).await?; + self.report_event_for_analytics(redemption).await?; // Report event for analytics + } + Ok(()) + } + + async fn report_event_for_analytics(&self, redemption: &Redemption) -> AnyhowResult<()> { + let mut event_builder = TwitchPubsubChannelPointsInsertBuilder::new(); + + let user_id = redemption.user.id.to_string(); + let user_name = redemption.user.login.to_string(); + + let mut event_builder = event_builder.set_sender_twitch_user_id(&user_id) + .set_sender_twitch_username(&user_name) + .set_destination_channel_id(&redemption.channel_id.to_string()) + // TODO: + .set_destination_channel_name("todo: not available") + .set_title(&redemption.reward.title) + .set_prompt(&redemption.reward.prompt) + .set_user_text_input(redemption.user_input.as_deref()) + .set_redemption_id(&redemption.id.to_string()) + .set_reward_id(&redemption.reward.id.to_string()) + .set_is_sub_only(redemption.reward.is_sub_only) + .set_reward_cost(redemption.reward.cost as u64); + // TODO: + // .set_max_per_stream(redemption.reward.max_per_stream as u64) + // .set_max_per_user_per_stream(redemption.reward.max_per_user_per_stream as u64); + event_builder.insert(&self.mysql_pool).await?; + + Ok(()) + } + + fn find_matching_rule(&self, redemption: &Redemption) -> AnyhowResult> { + return match self.twitch_subscriber_state.read() { + Err(e) => { Err(anyhow!("Lock error: {:?}", e)) }, + Ok(state) => { + info!("Checking channel points event against {} rules...", state.event_rules.len()); + let maybe_rule = state.event_rules.iter() + .filter(|rule| !rule.rule_is_disabled) + .filter(|rule| rule.event_category.eq(&TwitchEventCategory::ChannelPoints)) + .find(|rule| { + match rule.event_match_predicate { + EventMatchPredicate::NotSet {} => false, // Not set + EventMatchPredicate::BitsCheermoteNameExactMatch { ..} => false, // Wrong type + EventMatchPredicate::BitsCheermotePrefixSpendThreshold { .. } => false, // Wrong type + EventMatchPredicate::BitsSpendThreshold { .. } => false, // Wrong type + EventMatchPredicate::ChannelPointsRewardNameExactMatch { ref reward_name } => { + reward_name.eq_ignore_ascii_case(&redemption.reward.title) + } + } + }) + .map(|rule| rule.clone()); + Ok(maybe_rule) + } + }; + } + + //noinspection DuplicatedCode + async fn handle_matched_rule(&self, rule: &TwitchEventRuleLight, redemption: &Redemption) -> AnyhowResult<()> { + let message = match redemption.user_input { + None => return Ok(()), // Nothing to do. + Some(ref m) => m.as_str(), + }; + match rule.event_response { + EventResponse::NotSet {} => { + info!("Empty event response."); + return Ok(()) + }, + EventResponse::TtsSingleVoice { ref tts_model_token } => { + self.tts_writer.write_tts_with_model(message, tts_model_token).await?; + } + EventResponse::TtsRandomVoice { ref tts_model_tokens } => { + let maybe_token = random_from_vec(tts_model_tokens); + if let Some(token) = maybe_token { + self.tts_writer.write_tts_with_model(message, token.as_str()) + .await?; + } + } + EventResponse::TtsCommandPresets { .. } => {} // TODO + EventResponse::TtsCommandCustom { .. } => {} // TODO + } + + Ok(()) + } +} \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/mod.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/mod.rs new file mode 100644 index 0000000000..3f37ba9c3a --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/event_handlers/mod.rs @@ -0,0 +1,2 @@ +pub mod bits_event_handler; +pub mod channel_points_event_handler; diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/mod.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/mod.rs new file mode 100644 index 0000000000..d30da1b0c8 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/mod.rs @@ -0,0 +1,4 @@ +pub mod event_handlers; +pub mod subscriber_preferences_cache; +pub mod tts_writer; +pub mod twitch_pubsub_user_subscriber_thread; diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/subscriber_preferences_cache.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/subscriber_preferences_cache.rs new file mode 100644 index 0000000000..f19f244212 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/subscriber_preferences_cache.rs @@ -0,0 +1,34 @@ +use database_queries::queries::twitch::twitch_event_rules::list_twitch_event_rules_for_user::TwitchEventRule; +use database_queries::column_types::twitch_event_category::TwitchEventCategory; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; + +#[derive(Clone)] +pub struct TwitchEventRuleLight { + pub token: String, + pub event_category: TwitchEventCategory, + pub event_match_predicate: EventMatchPredicate, + pub event_response: EventResponse, + pub user_specified_rule_order: u32, + pub rule_is_disabled: bool, +} + +/// The "Twitch Pubsub Subscriber" thread is complicated. +/// Let's maintain a cache we can pass around. +pub struct TwitchPubsubCachedState { + // TODO: Query global settings. + pub global_settings: Option, + + // TODO: Query event rules. + // TODO: Not in order. + pub event_rules: Vec, +} + +impl TwitchPubsubCachedState { + pub fn new() -> Self { + Self { + global_settings: None, + event_rules: Vec::new(), + } + } +} diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/tts_writer.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/tts_writer.rs new file mode 100644 index 0000000000..bfad6c1675 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/tts_writer.rs @@ -0,0 +1,66 @@ +use container_common::anyhow_result::AnyhowResult; +use database_queries::queries::tts::tts_inference_jobs::insert_tts_inference_job::TtsInferenceJobInsertBuilder; +use database_queries::tokens::Tokens; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use redis_common::shared_constants::STREAMER_TTS_JOB_QUEUE_TTL_SECONDS; +use sqlx::MySql; +use std::sync::Arc; +use twitch_common::cheers::remove_cheers; +use twitch_common::twitch_user_id::TwitchUserId; + +/// Twitch users get priority above all FakeYou users. +/// TODO: In the future, there should be higher priority for paid users. +const TWITCH_TTS_PRIORITY_LEVEL : u8 = 10; + +pub struct TtsWriter { + mysql_pool: Arc>, + redis_pool: Arc>, + twitch_user_id: TwitchUserId, +} + +impl TtsWriter { + + pub fn new( + mysql_pool: Arc>, + redis_pool: Arc>, + twitch_user_id: TwitchUserId, + ) -> Self { + Self { + mysql_pool, + redis_pool, + twitch_user_id, + } + } + + pub async fn write_tts(&self, message_text: &str) -> AnyhowResult<()> { + //let model_token = "TM:7wbtjphx8h8v"; // "Mario *" voice (prod) + let model_token = "TM:40m3aqtt41y0"; // "Wakko" voice (dev) + self.write_tts_with_model(message_text, model_token).await + } + + pub async fn write_tts_with_model(&self, message_text: &str, model_token: &str) -> AnyhowResult<()> { + let sanitized_text = remove_cheers(message_text); + let job_token = Tokens::new_tts_inference_job()?; + + let mut builder = TtsInferenceJobInsertBuilder::new_for_internal_tts() + .set_is_for_twitch(true) + .set_priority_level(TWITCH_TTS_PRIORITY_LEVEL) + .set_job_token(&job_token) + .set_model_token(model_token) + .set_raw_inference_text(&sanitized_text); + + let _r = builder.insert(&self.mysql_pool).await?; + + // TODO: Report job token to frontend + let mut redis = self.redis_pool.get()?; + let redis_key = RedisKeys::twitch_tts_job_queue(&self.twitch_user_id.get_str()); + + let _size : Option = redis.rpush(&redis_key, job_token)?; + let _size : Option = redis.expire(&redis_key, STREAMER_TTS_JOB_QUEUE_TTL_SECONDS)?; + + Ok(()) + } +} \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/twitch_pubsub_user_subscriber_thread.rs b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/twitch_pubsub_user_subscriber_thread.rs new file mode 100644 index 0000000000..a53ed2d963 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/threads/twitch_pubsub_user_subscriber/twitch_pubsub_user_subscriber_thread.rs @@ -0,0 +1,684 @@ +use anyhow::{anyhow, Error}; +use container_common::anyhow_result::AnyhowResult; +use container_common::thread::thread_id::ThreadId; +use crate::twitch::constants::TWITCH_PING_CADENCE; +use crate::twitch::oauth::oauth_token_refresher::OauthTokenRefresher; +use crate::twitch::pubsub::build_pubsub_topics_for_user::build_pubsub_topics_for_user; +use crate::twitch::websocket_client::TwitchWebsocketClient; +use database_queries::complex_models::event_match_predicate::EventMatchPredicate; +use database_queries::complex_models::event_responses::EventResponse; +use database_queries::queries::tts::tts_inference_jobs::insert_tts_inference_job::TtsInferenceJobInsertBuilder; +use database_queries::queries::twitch::twitch_event_rules::list_twitch_event_rules_for_user::list_twitch_event_rules_for_user; +use database_queries::queries::twitch::twitch_oauth::find::{TwitchOauthTokenRecord, TwitchOauthTokenFinder}; +use database_queries::queries::twitch::twitch_oauth::insert::TwitchOauthTokenInsertBuilder; +use database_queries::queries::twitch::twitch_pubsub::insert_bits::TwitchPubsubBitsInsertBuilder; +use database_queries::queries::twitch::twitch_pubsub::insert_channel_points::TwitchPubsubChannelPointsInsertBuilder; +use database_queries::tokens::Tokens; +use log::debug; +use log::error; +use log::info; +use log::warn; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use redis_common::payloads::lease_payload::LeasePayload; +use redis_common::redis_keys::RedisKeys; +use redis_common::shared_constants::LEASE_CHECK_PERIOD; +use redis_common::shared_constants::LEASE_RENEW_PERIOD; +use redis_common::shared_constants::LEASE_TIMEOUT_SECONDS; +use redis_common::shared_constants::OBS_ACTIVE_CHECK_PERIOD; +use redis_common::shared_constants::STREAMER_TTS_JOB_QUEUE_TTL_SECONDS; +use sqlx::MySql; +use std::ops::Sub; +use std::sync::{Arc, RwLock, PoisonError, RwLockWriteGuard}; +use std::thread::sleep; +use std::time::Duration; +use time::Duration as TimeDuration; +use time::Instant; +use twitch_api2::pubsub::channel_bits::ChannelBitsEventsV2Reply; +use twitch_api2::pubsub::channel_points::ChannelPointsChannelV1Reply; +use twitch_api2::pubsub::{Response, TwitchResponse, TopicData}; +use twitch_common::cheers::remove_cheers; +use twitch_common::twitch_user_id::TwitchUserId; +use crate::threads::twitch_pubsub_user_subscriber::subscriber_preferences_cache::TwitchPubsubCachedState; +use crate::threads::twitch_pubsub_user_subscriber::subscriber_preferences_cache::TwitchEventRuleLight; +use crate::threads::twitch_pubsub_user_subscriber::event_handlers::channel_points_event_handler::ChannelPointsEventHandler; +use crate::threads::twitch_pubsub_user_subscriber::tts_writer::TtsWriter; +use crate::threads::twitch_pubsub_user_subscriber::event_handlers::bits_event_handler::BitsEventHandler; + +// TODO: Publish events back to OBS thread +// TODO: (cleanup) make the logic clearer to follow. + +pub const REFRESH_WEBSITE_SETTINGS_CADENCE : TimeDuration = TimeDuration::minutes(2); + +// ========================================= +// =============== STAGE ONE =============== +// ========================================= + +pub struct TwitchPubsubUserSubscriberThread { + thread_id: ThreadId, + server_hostname: String, + twitch_user_id: TwitchUserId, + maybe_user_token: Option, // Storyteller user + oauth_token_refresher: OauthTokenRefresher, + mysql_pool: Arc>, + redis_pool: Arc>, + twitch_subscriber_state: Arc>, +} + +impl TwitchPubsubUserSubscriberThread { + pub fn new( + twitch_user_id: TwitchUserId, + oauth_token_refresher: OauthTokenRefresher, + mysql_pool: Arc>, + redis_pool: Arc>, + server_hostname: &str, + thread_id: ThreadId, + ) -> Self { + Self { + thread_id, + oauth_token_refresher, + server_hostname: server_hostname.to_string(), + twitch_user_id, + maybe_user_token: None, + mysql_pool, + redis_pool, + twitch_subscriber_state: Arc::new(RwLock::new(TwitchPubsubCachedState::new())), + } + } + + pub async fn start_thread(mut self) { + // By failing to look this up, the thread will fail fast. + // When the user auths, the thread will be picked back up again. + let lookup_result + = lookup_oauth_record(&self.twitch_user_id, &self.mysql_pool).await; + + let mut record = match lookup_result { + Ok(Some(record)) => record, + Ok(None) => { + error!("No twitch oauth token record"); + return; + } + Err(e) => { + error!("Error looking up twitch oauth token record: {:?}", e); + return; + } + }; + + self.maybe_user_token = record.maybe_user_token.clone(); + + let expected_lease_payload + = LeasePayload::from_thread_id(&self.server_hostname, &self.thread_id); + + loop { + info!("Connecting to Twitch PubSub for user {}...", &record.twitch_username); + let maybe_client = + self.create_subscribed_twitch_client(&record.access_token).await; + + let twitch_websocket_client = match maybe_client { + Ok(client) => client, + Err(e) => { + error!("Error building Twitch client (exiting thread): {:?}", e); + return; + } + }; + + let tts_writer = Arc::new(TtsWriter::new( + self.mysql_pool.clone(), + self.redis_pool.clone(), + self.twitch_user_id.clone(), + )); + + let bits_event_handler = BitsEventHandler::new( + self.twitch_subscriber_state.clone(), + self.mysql_pool.clone(), + tts_writer.clone(), + ); + + let channel_points_event_handler = ChannelPointsEventHandler::new( + self.twitch_subscriber_state.clone(), + self.mysql_pool.clone(), + tts_writer.clone(), + ); + + // NB: All of the timers (thus far) can wait to run. + // We set their first run time to now so we don't have to deal with Option<>. + let now = Instant::now(); + + // NB(2022-02-20): Actually, we can set these in the past with a bit of a hack. + let long_ago = Instant::now().sub(Duration::from_secs(60*60*24)); + + let thread = TwitchPubsubUserSubscriberThreadStageTwo { + thread_id: self.thread_id.clone(), + server_hostname: self.server_hostname.clone(), + twitch_user_id: self.twitch_user_id.clone(), + maybe_user_token: self.maybe_user_token.clone(), + oauth_token_refresher: self.oauth_token_refresher.clone(), + mysql_pool: self.mysql_pool.clone(), + redis_pool: self.redis_pool.clone(), + twitch_websocket_client, + expected_lease_payload: expected_lease_payload.clone(), + twitch_oauth_token_record: record.clone(), + redis_lease_last_renewed_at: now, + redis_lease_last_checked_at: now, + obs_session_active_last_checked_at: now, + twitch_last_pinged_at: now, + twitch_subscriber_state: self.twitch_subscriber_state.clone(), + website_settings_last_refreshed_at: long_ago, + channel_points_event_handler, + bits_event_handler, + }; + + // NB: The following call will run its main loop until/unless the Twitch client + // fails to auth or disconnects. If this happens, we'll try again. + match thread.continue_thread().await { + Ok(LoopEndedReason::ExitThread { reason}) => { + warn!("Thread has ended with reason: {}", reason); + return; + } + Ok(LoopEndedReason::RefreshedOauthToken { token }) => { + warn!("OAuth token was refreshed."); + record = token; + sleep(Duration::from_secs(3)); + } + Err(e) => { + error!("There was an error, restarting thread shortly: {:?}", e); + sleep(Duration::from_secs(15)); + continue; + } + } + } + } + + async fn create_subscribed_twitch_client( + &self, + access_token: &str + ) -> AnyhowResult { + let mut twitch_websocket_client = TwitchWebsocketClient::new()?; + + twitch_websocket_client.connect().await?; + twitch_websocket_client.send_ping().await?; + + // NB: Failure to auth won't be immediate. + let topics = build_pubsub_topics_for_user(self.twitch_user_id.get_numeric()); + twitch_websocket_client.listen(access_token, &topics).await?; + + Ok(twitch_websocket_client) + } +} + +// ========================================= +// =============== STAGE TWO =============== +// ========================================= + +/// The thread is somewhat of a state machine. +/// The first stage of thread startup can end prematurely, which is why +/// this is modeled as two different structs. +struct TwitchPubsubUserSubscriberThreadStageTwo { + thread_id: ThreadId, + server_hostname: String, + twitch_user_id: TwitchUserId, + maybe_user_token: Option, + oauth_token_refresher: OauthTokenRefresher, + mysql_pool: Arc>, + redis_pool: Arc>, + + // ========== CACHE ========== + + twitch_subscriber_state: Arc>, + + // ========== Handlers ========== + + bits_event_handler: BitsEventHandler, + channel_points_event_handler: ChannelPointsEventHandler, + + // ========== Stage Two Thread State ========== + + twitch_websocket_client: TwitchWebsocketClient, + + // We'll use this again and again, so precompute it. + expected_lease_payload: LeasePayload, + + // The user's oauth access and refresh tokens. + twitch_oauth_token_record: TwitchOauthTokenRecord, + + // The thread must renew the lease, or another worker will pick it up. + // If the lease gets taken by another, we abandon our own workload. + redis_lease_last_renewed_at: Instant, + + // If the lease gets taken by another thread, we abandon our own workload. + // This controls when we periodically check. + redis_lease_last_checked_at: Instant, + + // Check if the OBS session is still active. + // If the underlying Redis key dies, we abandon our thread. + obs_session_active_last_checked_at: Instant, + + /// Twitch PubSub requires PINGs at regular intervals, + /// "To keep the server from closing the connection, clients must send a PING + /// command at least once every 5 minutes. If a client does not receive a PONG + /// message within 10 seconds of issuing a PING command, it should reconnect + /// to the server. See details in Handling Connection Failures." + twitch_last_pinged_at: Instant, + + // When we last reloaded the website settings + website_settings_last_refreshed_at: Instant, +} + +impl TwitchPubsubUserSubscriberThreadStageTwo { + + /// This function will loop until it either errors or hits a `LoopEndedReason` condition. + /// The caller will need to handle these cases. + pub async fn continue_thread(mut self) -> AnyhowResult { + loop { + let is_valid = self.maybe_check_redis_lease_is_valid()?; + if !is_valid { + return Ok(LoopEndedReason::ExitThread { reason: "Thread lease taken".to_string() }); + } + + self.maybe_renew_redis_lease()?; + self.maybe_send_twitch_ping().await?; + + let is_active = self.maybe_check_obs_session_active()?; + if !is_active { + return Ok(LoopEndedReason::ExitThread { reason: "OBS session ended".to_string() }); + } + + self.maybe_refresh_website_user_settings().await?; + + // NB: We can't have calls to read the Twitch websocket client block forever, and they + // would do exactly that if not for this code. This is adapted from the very good example + // in the `tokio-tungstenite` repo, which also contains good recipes for splitting sockets + // into two unidirectional streams: + // https://github.com/snapview/tokio-tungstenite/blob/master/examples/interval-server.rs + let mut interval = tokio::time::interval(Duration::from_secs(1)); + tokio::select! { + maybe_event = self.twitch_websocket_client.try_next() => { + if let Some(loop_end_reason) = self.handle_event(maybe_event).await? { + return Ok(loop_end_reason); + } + } + _ = interval.tick() => { + sleep(Duration::from_secs(1)); + } + } + + sleep(Duration::from_secs(1)); + } + } + + // =============== TWITCH PUBSUB EVENTS =============== + + async fn handle_event( + &mut self, + maybe_event: AnyhowResult> + ) -> AnyhowResult> { + + let maybe_event = maybe_event.map_err(|error| { + error!("socket error: {:?}", error); + error + })?; + + let event = match maybe_event { + None => return Ok(None), + Some(event) => event, + }; + + info!("event: {:?}", event); + + match event { + Response::Response(response) => { + // NB: Auth failure might cause the loop to end. + return self.handle_maybe_auth_error_event(response).await; + } + Response::Message { data } => { + self.handle_pubsub_topic_event(data).await?; + } + Response::Pong => {} + Response::Reconnect => {} + } + + Ok(None) // Don't end loop + } + + async fn handle_maybe_auth_error_event(&mut self, response: TwitchResponse) + -> AnyhowResult> + { + let error = match response.error.as_deref() { + None => return Ok(None), + Some(e) => e, + }; + match error { + "" => {}, // No-op + "ERR_BADAUTH" => { + warn!("Invalid token. Bad auth. Need to refresh"); + let token_record = self.refresh_twitch_oauth_token().await?; + return Ok(Some(LoopEndedReason::RefreshedOauthToken { token: token_record })); + } + _ => warn!("Unknown Twitch PubSub error: {:?}", error), + } + Ok(None) + } + + async fn handle_pubsub_topic_event(&mut self, topic_data: TopicData) -> AnyhowResult<()> { + match topic_data { + // Unimplemented + TopicData::AutoModQueue { .. } => {} + TopicData::ChannelBitsBadgeUnlocks { .. } => {} + TopicData::ChatModeratorActions { .. } => {} + TopicData::ChannelSubscribeEventsV1 { .. } => {} + TopicData::CommunityPointsChannelV1 { .. } => {} + TopicData::ChannelCheerEventsPublicV1 { .. } => {} + TopicData::ChannelSubGiftsV1 { .. } => {} + TopicData::VideoPlayback { .. } => {} + TopicData::VideoPlaybackById { .. } => {} + TopicData::HypeTrainEventsV1 { .. } => {} + TopicData::HypeTrainEventsV1Rewards { .. } => {} + TopicData::Following { .. } => {} + TopicData::Raid { .. } => {} + TopicData::UserModerationNotifications { .. } => {} + // Implemented + TopicData::ChannelBitsEventsV2 { topic, reply } => { + self.bits_event_handler.handle(topic, reply).await?; + } + TopicData::ChannelPointsChannelV1 { topic, reply } => { + self.channel_points_event_handler.handle(topic, reply).await?; + } + } + + Ok(()) + } + + // =============== TWITCH PUBSUB KEEPALIVE =============== + + async fn maybe_send_twitch_ping(&mut self) -> AnyhowResult<()> { + let mut should_send_ping = self.twitch_last_pinged_at + .elapsed() + .gt(&TWITCH_PING_CADENCE); + + if should_send_ping { + info!("Sending Twitch ping for user {}", self.twitch_user_id.get_numeric()); + self.twitch_websocket_client.send_ping().await?; + self.twitch_last_pinged_at = Instant::now(); + } + + Ok(()) + } + + // =============== REDIS THREAD LEASE =============== + + fn maybe_check_redis_lease_is_valid(&mut self) -> AnyhowResult { + let mut should_check_lease = self.redis_lease_last_checked_at + .elapsed() + .gt(&LEASE_CHECK_PERIOD); + + if should_check_lease { + info!("Checking Redis Lease for user {}", self.twitch_user_id.get_numeric()); + let is_valid = self.check_redis_lease_is_valid()?; + + if !is_valid { + warn!("Lease got taken by another thread"); + return Ok(false); + } + + self.redis_lease_last_checked_at = Instant::now(); + } + + Ok(true) + } + + fn check_redis_lease_is_valid(&mut self) -> AnyhowResult { + let mut redis = self.redis_pool.get()?; + let lease_key = RedisKeys::twitch_pubsub_lease(self.twitch_user_id.get_str()); + + let payload : Option = redis.get(&lease_key)?; + match payload { + None => { + warn!("Redis lease payload absent. Another thread could be started."); + Ok(true) + } + Some(payload) => { + let actual_payload = LeasePayload::deserialize(&payload)?; + let equals_expected = self.expected_lease_payload.eq(&actual_payload); + Ok(equals_expected) + } + } + } + + fn maybe_renew_redis_lease(&mut self) -> AnyhowResult<()> { + let mut should_renew_lease = self.redis_lease_last_renewed_at + .elapsed() + .gt(&LEASE_RENEW_PERIOD); + + if should_renew_lease { + info!("Renewing Redis Lease for user {}", self.twitch_user_id.get_numeric()); + self.renew_redis_lease()?; + self.redis_lease_last_renewed_at = Instant::now(); + } + + Ok(()) + } + + fn renew_redis_lease(&mut self) -> AnyhowResult<()> { + let mut redis = self.redis_pool.get()?; + + let lease_key = RedisKeys::twitch_pubsub_lease(self.twitch_user_id.get_str()); + let lease_value = self.expected_lease_payload.serialize(); + + let _v : Option = redis.set_ex( + &lease_key, + &lease_value, + LEASE_TIMEOUT_SECONDS + )?; + Ok(()) + } + + // =============== OBS SESSION ACTIVITY =============== + + fn maybe_check_obs_session_active(&mut self) -> AnyhowResult { + let mut should_check_active = self.obs_session_active_last_checked_at + .elapsed() + .gt(&OBS_ACTIVE_CHECK_PERIOD); + + if should_check_active { + info!("Checking OBS active for user {}", self.twitch_user_id.get_numeric()); + let is_active = self.check_obs_session_active()?; + + if !is_active { + warn!("OBS session is no longer active"); + return Ok(false); + } + + self.obs_session_active_last_checked_at = Instant::now(); + } + + Ok(true) + } + + fn check_obs_session_active(&mut self) -> AnyhowResult { + let mut redis = self.redis_pool.get()?; + let key = RedisKeys::obs_active_session_keepalive(self.twitch_user_id.get_str()); + + // The value doesn't matter, just the presence of the key. + let payload : Option = redis.get(&key)?; + match payload { + None => Ok(false), + Some(_payload) => Ok(true), + } + } + + // =============== OAUTH TOKEN LOOKUP AND RENEWAL =============== + + async fn refresh_twitch_oauth_token(&mut self) -> AnyhowResult { + let refresh_token = match self.twitch_oauth_token_record.maybe_refresh_token.as_deref() { + Some(token) => token, + None => { + error!("No refresh token present. Cannot refresh"); + return Err(anyhow!("No refresh token present. Cannot refresh!")); + }, + }; + + let refresh_result = self.oauth_token_refresher.refresh_token(refresh_token) + .await?; + + let access_token = refresh_result.access_token.secret().to_string(); + let refresh_token : Option = refresh_result.maybe_refresh_token + .map(|t| t.secret().to_string()); + let expires_seconds = refresh_result.duration.as_secs() as u32; + + // TODO: Move saving a refreshed record somewhere common + let mut query_builder = TwitchOauthTokenInsertBuilder::new( + &self.twitch_oauth_token_record.twitch_user_id, + &self.twitch_oauth_token_record.twitch_username, + &access_token, + &self.twitch_oauth_token_record.oauth_refresh_grouping_token) + .set_refresh_token(refresh_token.as_deref()) + .set_user_token(self.twitch_oauth_token_record.maybe_user_token.as_deref()) + .set_expires_in_seconds(Some(expires_seconds)) + .set_refresh_count(self.twitch_oauth_token_record.refresh_count.saturating_add(1)) + // NB: We don't get these back from the refresh, but it seems like they would stay the same. + .set_token_type(self.twitch_oauth_token_record.token_type.as_deref()) + .has_bits_read(self.twitch_oauth_token_record.has_bits_read) + .has_channel_read_redemptions(self.twitch_oauth_token_record.has_channel_read_redemptions) + .has_channel_read_subscriptions(self.twitch_oauth_token_record.has_channel_read_subscriptions) + .has_chat_edit(self.twitch_oauth_token_record.has_chat_edit) + .has_chat_read(self.twitch_oauth_token_record.has_chat_read); + + query_builder.insert(&self.mysql_pool).await?; + + let maybe_inserted + = lookup_oauth_record(&self.twitch_user_id, &self.mysql_pool).await?; + + let new_record = match maybe_inserted { + Some(record) => record, + None => { + error!("Did not find oauth token record in database upon refresh"); + return Err(anyhow!("Did not find oauth token record in database upon refresh")); + } + }; + + // NB: Instead of resetting the local loop state, we'll create a fresh new Twitch client. + // TODO: Compare a "refresh_count" + self.twitch_oauth_token_record = new_record.clone(); + + Ok(new_record) + } + + // =============== TTS EVENTS =============== + + async fn write_tts_inference_event(&mut self, tts_text: &str) -> AnyhowResult<()> { + let sanitized_text = remove_cheers(tts_text); + let job_token = Tokens::new_tts_inference_job()?; + //let model_token = "TM:7wbtjphx8h8v"; // "Mario *" voice (prod) + let model_token = "TM:40m3aqtt41y0"; // "Wakko" voice (dev) + + let mut builder = TtsInferenceJobInsertBuilder::new_for_internal_tts() + .set_is_for_twitch(true) + .set_priority_level(1) + .set_job_token(&job_token) + .set_model_token(model_token) + .set_raw_inference_text(&sanitized_text); + + let _r = builder.insert(&self.mysql_pool).await?; + + // TODO: Report job token to frontend + let mut redis = self.redis_pool.get()?; + let redis_key = RedisKeys::twitch_tts_job_queue(&self.twitch_user_id.get_str()); + + let _size : Option = redis.rpush(&redis_key, job_token)?; + let _size : Option = redis.expire(&redis_key, STREAMER_TTS_JOB_QUEUE_TTL_SECONDS)?; + + Ok(()) + } + + // =============== WEBSITE SETTINGS =============== + + async fn maybe_refresh_website_user_settings(&mut self) -> AnyhowResult<()> { + let user_token = match self.maybe_user_token.as_deref() { + Some(token) => token, + None => { + // Only Twitch users with a storyteller account get settings. + return Ok(()); + } + }; + + let mut should_refresh = self.website_settings_last_refreshed_at + .elapsed() + .gt(&REFRESH_WEBSITE_SETTINGS_CADENCE); + + if !should_refresh { + return Ok(()); + } + + info!("Reloading website settings for user {}", self.twitch_user_id.get_numeric()); + + let event_rules = + lookup_event_rules(user_token, &self.mysql_pool).await?; + + match self.twitch_subscriber_state.write() { + Err(e) => { + return Err(anyhow!("Lock error: {:?}", e)) + }, + Ok(mut lock) => { + lock.event_rules = event_rules; + } + } + + self.website_settings_last_refreshed_at = Instant::now(); + + Ok(()) + } +} + +// ==================================== +// =============== MISC =============== +// ==================================== + +async fn lookup_oauth_record( + twitch_user_id: &TwitchUserId, + mysql_pool: &sqlx::Pool +) -> AnyhowResult> { + TwitchOauthTokenFinder::new() + .scope_twitch_user_id(Some(twitch_user_id.get_numeric())) + .allow_expired_tokens(true) + .perform_query(mysql_pool) + .await +} + +async fn lookup_event_rules( + user_token: &str, + mysql_pool: &sqlx::Pool +) -> AnyhowResult> { + let rules = list_twitch_event_rules_for_user(user_token, mysql_pool) + .await?; + let rules = rules.into_iter() + .map(|rule| { + let event_match_predicate = serde_json::from_str(&rule.event_match_predicate) + .unwrap_or_else(|e| { + error!("Issue with deserializing: {}", e); + EventMatchPredicate::NotSet {} + }); + let event_response = serde_json::from_str(&rule.event_response) + .unwrap_or_else(|e| { + error!("Issue with deserializing: {}", e); + EventResponse::NotSet {} + }); + TwitchEventRuleLight { + token: rule.token, + event_category: rule.event_category, + event_match_predicate, + event_response, + user_specified_rule_order: rule.user_specified_rule_order, + rule_is_disabled: rule.rule_is_disabled, + } + }) + .collect(); + Ok(rules) +} + +enum LoopEndedReason { + /// Terminate the thread + ExitThread { reason: String }, + /// A new OAuth token was minted. + RefreshedOauthToken { token: TwitchOauthTokenRecord } +} diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/constants.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/constants.rs new file mode 100644 index 0000000000..02b790c58b --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/constants.rs @@ -0,0 +1,11 @@ +use time::Duration; + +/// Twitch PubSub requires PINGs at regular intervals, +/// +/// "To keep the server from closing the connection, clients must send a PING +/// command at least once every 5 minutes. If a client does not receive a PONG +/// message within 10 seconds of issuing a PING command, it should reconnect +/// to the server. See details in Handling Connection Failures." +/// +/// We'll keep this intentionally short so we don't elapse. +pub const TWITCH_PING_CADENCE: Duration = Duration::minutes(3); diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/mod.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/mod.rs new file mode 100644 index 0000000000..3aaaab19b6 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/mod.rs @@ -0,0 +1,4 @@ +pub mod constants; +pub mod oauth; +pub mod pubsub; +pub mod websocket_client; \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/oauth/mod.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/oauth/mod.rs new file mode 100644 index 0000000000..02d8f8a8c2 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/oauth/mod.rs @@ -0,0 +1 @@ +pub mod oauth_token_refresher; \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/oauth/oauth_token_refresher.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/oauth/oauth_token_refresher.rs new file mode 100644 index 0000000000..ca1330286c --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/oauth/oauth_token_refresher.rs @@ -0,0 +1,60 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use reqwest::Client; +use reqwest::redirect::Policy; +use std::time::Duration; +use twitch_oauth2::{ClientId, ClientSecret, RefreshToken, AccessToken}; + +/// Refresh OAuth tokens when they expire. +/// NB: This can be cloned and passed around. +#[derive(Clone)] +pub struct OauthTokenRefresher { + http_client: Client, + client_id: ClientId, + client_secret: ClientSecret, +} + +pub struct RefreshTokenResult { + pub access_token: AccessToken, + pub duration: Duration, + pub maybe_refresh_token: Option, +} + +impl OauthTokenRefresher { + + pub fn from_secrets(client_id: ClientId, client_secret: ClientSecret) -> AnyhowResult { + let http_client = Client::builder() + .redirect(Policy::none()) + .build()?; + Ok(Self { + http_client, + client_id, + client_secret, + }) + } + + pub fn from_secrets_str(client_id: &str, client_secret: &str) -> AnyhowResult { + let client_id = ClientId::new(client_id); + let client_secret = ClientSecret::new(client_secret); + Self::from_secrets(client_id, client_secret) + } + + /// Perform a token refresh + pub async fn refresh_token(&self, refresh_token: &str) -> AnyhowResult { + let refresh_token = RefreshToken::new(refresh_token); + let result = twitch_oauth2::refresh_token( + &self.http_client, + &refresh_token, + &self.client_id, + &self.client_secret) + .await + .map(|(access_token, duration, maybe_refresh_token)| { + RefreshTokenResult { + access_token, + duration, + maybe_refresh_token, + } + })?; + Ok(result) + } +} \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/pubsub/build_pubsub_topics_for_user.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/pubsub/build_pubsub_topics_for_user.rs new file mode 100644 index 0000000000..b61ae41e70 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/pubsub/build_pubsub_topics_for_user.rs @@ -0,0 +1,19 @@ +use twitch_api2::pubsub::{Topics, Topic}; +use twitch_api2::pubsub; + +pub fn build_pubsub_topics_for_user(user_id: u32) -> Vec { + vec![ + pubsub::channel_bits::ChannelBitsEventsV2 { + channel_id: user_id, + }.into_topic(), + pubsub::channel_points::ChannelPointsChannelV1 { + channel_id: user_id, + }.into_topic(), + pubsub::channel_cheer::ChannelCheerEventsPublicV1 { + channel_id: user_id, + }.into_topic(), + pubsub::channel_subscriptions::ChannelSubscribeEventsV1 { + channel_id: user_id, + }.into_topic(), + ] +} diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/pubsub/mod.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/pubsub/mod.rs new file mode 100644 index 0000000000..e771e309d2 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/pubsub/mod.rs @@ -0,0 +1 @@ +pub mod build_pubsub_topics_for_user; \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/twitch/websocket_client.rs b/crates/service/twitch_pubsub_subscriber/src/twitch/websocket_client.rs new file mode 100644 index 0000000000..f579e5f08f --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/twitch/websocket_client.rs @@ -0,0 +1,135 @@ +use anyhow::anyhow; +use anyhow::bail; +use container_common::anyhow_result::AnyhowResult; +use crate::util::random_nonce::random_nonce; +use futures_util::{SinkExt, StreamExt, TryStreamExt}; +use log::debug; +use reqwest::Url; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::{Arc, Mutex, PoisonError, MutexGuard}; +use std::time::Duration; +use tokio::net::TcpStream; +use tokio; +use tokio_tungstenite::tungstenite::Message; +use tokio_tungstenite::{connect_async, WebSocketStream, MaybeTlsStream}; +use twitch_api2::pubsub::{Response, TopicData, Topic, Topics}; +use twitch_api2::pubsub; + +// Reference javascript Twitch Websocket client: +// https://github.com/twitchdev/pubsub-javascript-sample/blob/main/main.js + +const WEBSOCKET_GATEWAY : &'static str = "wss://pubsub-edge.twitch.tv"; + +const PING_COMMAND : &'static str = "{ \"type\": \"PING\" }"; + +pub struct TwitchWebsocketClient { + socket: Option>>, +} + +impl TwitchWebsocketClient { + pub fn new() -> AnyhowResult { + Ok(Self { + socket: None, + }) + } + + pub async fn connect(&mut self) -> AnyhowResult<()> { + let url = Url::parse(WEBSOCKET_GATEWAY)?; + let (mut socket, _response) = + connect_async(url).await?; + self.socket = Some(socket); + Ok(()) + } + + pub async fn close(&mut self) -> AnyhowResult<()> { + let mut socket = match self.socket { + None => return Ok(()), // Already closed. + Some(ref mut socket) => socket, + }; + + socket.close(None).await?; + self.socket = None; + + Ok(()) + } + + pub async fn listen<'t>(&mut self, auth_token: &str, topics: &'t [Topics]) -> AnyhowResult<()> { + let mut socket = match self.socket { + None => return Err(anyhow!("not connected")), + Some(ref mut socket) => socket, + }; + + let nonce = random_nonce(128); + + let command = pubsub::listen_command( + topics, + auth_token, + nonce.as_ref(), + )?; + + debug!("Sending LISTEN"); + let message = Message::Text(command); + socket.send(message).await?; + + Ok(()) + } + + pub fn is_connected(&self) -> bool { + // TODO: This probably isn't accurate + self.socket.is_some() + } + + pub async fn send_ping(&mut self) -> AnyhowResult<()> { + let mut socket = match self.socket { + None => return Err(anyhow!("not connected")), + Some(ref mut socket) => socket, + }; + + debug!("Sending PING"); + let message = Message::Text(PING_COMMAND.to_string()); + socket.send(message).await?; + + // TODO: Need to read from all call sites due to ordering. + //match self.read_response().await { + // Err(e) => Err(e), + // Ok(Response::Pong) => Ok(()), + // _ => Err(bail!("wrong response type")), + //} + debug!("send ping ok"); + Ok(()) + } + + pub async fn next(&mut self) -> AnyhowResult { + let socket = match self.socket { + None => return Err(anyhow!("not connected")), + Some(ref mut s) => s, + }; + + match socket.next().await { + None => Err(anyhow!("nothing to read")), + Some(Err(e)) => Err(anyhow!("websocket error: {:?}", e)), + Some(Ok(response)) => { + let text = response.to_text()?; + let typed_response = Response::parse(text)?; + Ok(typed_response) + }, + } + } + + pub async fn try_next(&mut self) -> AnyhowResult> { + let socket = match self.socket { + None => return Err(anyhow!("not connected")), + Some(ref mut s) => s, + }; + + match socket.try_next().await { + Err(e) => Err(anyhow!("websocket error: {:?}", e)), + Ok(None) => Ok(None), + Ok(Some(response)) => { + let text = response.to_text()?; + let typed_response = Response::parse(text)?; + Ok(Some(typed_response)) + }, + } + } +} diff --git a/crates/service/twitch_pubsub_subscriber/src/util/mod.rs b/crates/service/twitch_pubsub_subscriber/src/util/mod.rs new file mode 100644 index 0000000000..b0c7022656 --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/util/mod.rs @@ -0,0 +1 @@ +pub mod random_nonce; \ No newline at end of file diff --git a/crates/service/twitch_pubsub_subscriber/src/util/random_nonce.rs b/crates/service/twitch_pubsub_subscriber/src/util/random_nonce.rs new file mode 100644 index 0000000000..f87ab1765e --- /dev/null +++ b/crates/service/twitch_pubsub_subscriber/src/util/random_nonce.rs @@ -0,0 +1,24 @@ +use rand::distributions::Alphanumeric; +use rand::{thread_rng, Rng}; + +/// Returns a random nonce for Twitch +pub fn random_nonce(length: usize) -> String { + thread_rng() + .sample_iter(&Alphanumeric) + .take(length) + .map(char::from) + .collect() +} + +#[cfg(test)] +mod tests { + use crate::util::random_nonce::random_nonce; + + #[test] + fn random_nonce_length() { + assert_eq!(1, random_nonce(1).len()); + assert_eq!(10, random_nonce(10).len()); + assert_eq!(15, random_nonce(15).len()); + assert_eq!(64, random_nonce(64).len()); + } +} \ No newline at end of file diff --git a/crates/service/websocket_gateway/Cargo.toml b/crates/service/websocket_gateway/Cargo.toml new file mode 100644 index 0000000000..9b907e0e1b --- /dev/null +++ b/crates/service/websocket_gateway/Cargo.toml @@ -0,0 +1,123 @@ +[package] +name = "websocket-gateway" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[[bin]] +name = "websocket-gateway" +path = "src/main.rs" + +[[bin]] +name = "twitch-pubsub-job" +path = "src/twitch_pubsub_job.rs" + +[dependencies] + +# Internal +config = { path = "../../lib/config" } +container_common = { path = "../../lib/container_common" } +database_queries = { path = "../../lib/database_queries" } +http_server_common = { path = "../../lib/http_server_common" } +redis_common = { path = "../../lib/redis_common" } +twitch_common = { path = "../../lib/twitch_common" } + +# External +anyhow = "1.0.40" +base64 = "0.13.0" +chrono = { version = "0.4.19", features = ["serde"] } +derive_more = "0.99.14" +dotenv = "0.15.0" +easyenv = "0.0.6" +futures-timer = "3.0.2" +hex = "0.4.3" +hostname = "0.3.1" +lazy_static = "1.4.0" +lfu = "0.2.5" +log = "0.4.14" +rand = "0.8.3" +regex = "1.5.4" +time = "0.2.26" +uuid = { version = "0.8.2", features = ["v4"] } + +# Runtime +# NB: We switched to Tokio from `async-std` for broader ecosystem support. +tokio = { version = "1.10.0", features = ["rt-multi-thread", "macros", "time"] } + +# Futures +futures = { version = "0.3.17", features = ["thread-pool"] } # feature for ThreadPool +futures-util = "0.3.17" +futures_old_for_limiter = { package = "futures", version = "0.1.29" } # Older version of futures for `limiter` crate + +# Threading +crossbeam = "0.8.1" + +# MySql Database +# "offline" supports typesafe builds without a live database +# "rustls" doesn't appear to work with IP addresses, per: +# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828 +#sqlx = { version = "0.5", features = [ "mysql", "runtime-async-std-native-tls", "offline", "chrono" ] } +sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] } +#sqlx = { version = "0.4", features = [ "mysql", "runtime-tokio-native-tls", "offline", "chrono" ] } + +# Redis Database +limitation = "0.1.1" # Redis rate limiter +r2d2_redis = "0.14.0" + +# Twitch +# NB "unsupported" pubsub features may disappear from the Twitch API at any time. It's an unofficial API. +twitch_api2 = { version = "0.6.0-rc.2", features = ["helix", "reqwest_client", "pubsub", "chrono", "unsupported"] } +twitch_oauth2 = "0.6.0-rc.2" + +# HTTP +reqwest = { version = "0.11.4" } +native-tls = "0.2.8" + +# In-memory caching +lru_time_cache = "0.11.11" + +# Web +actix = "0.12.0" +actix-cors = "0.5.4" +actix-files= "0.6.0-beta.8" +actix-http = "3.0.0-beta.11" +actix-rt = "2.4.0" +actix-service = "2.0.1" +actix-utils = "3.0.0" +actix-web = "4.0.0-beta.10" +actix-web-actors = "4.0.0-beta.7" +jwt = "0.13.0" + +# Websockets +tokio-tungstenite = { version = "0.15.0", features = ["rustls", "rustls-tls"] } + +# NB: Hyper is just for HTTP client use in jobs. I'd use Actix client, but I need tokio runtime +# NB: Furthermore, the Tokio runtime is old, so I have to use an older Hyper. +hyper = { version = "0.13.10" } + +# Crypto +bcrypt = "0.8" # password hashes +hmac = "0.10.0" # cookies +magic-crypt = "3.1.8" # encrypting sort IDs for the frontend so we don't leak info +md-5 = "0.9.1" # gravatar +sha2 = "0.9.5" # file hashes, inference text hashes, etc. + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" +toml = "0.5.8" + +# Text processing +pulldown-cmark = "0.8.0" + +# Temporary (just to ask for oauth credentials during dev) +rpassword = "5.0.1" + +[build-dependencies] +# NB: Only using diesel for migrations support +diesel = { version = "1.4.6", features = ["mysql"] } diff --git a/crates/service/websocket_gateway/src/endpoints_ws/mod.rs b/crates/service/websocket_gateway/src/endpoints_ws/mod.rs new file mode 100644 index 0000000000..fa4404b876 --- /dev/null +++ b/crates/service/websocket_gateway/src/endpoints_ws/mod.rs @@ -0,0 +1 @@ +pub mod obs_gateway_websocket_handler; \ No newline at end of file diff --git a/crates/service/websocket_gateway/src/endpoints_ws/obs_gateway_websocket_handler.rs b/crates/service/websocket_gateway/src/endpoints_ws/obs_gateway_websocket_handler.rs new file mode 100644 index 0000000000..7df0547faf --- /dev/null +++ b/crates/service/websocket_gateway/src/endpoints_ws/obs_gateway_websocket_handler.rs @@ -0,0 +1,309 @@ +use actix::prelude::*; +use actix_rt::Runtime; +use actix_web::web::Path; +use actix_web::{middleware, web, App, Error, HttpRequest, HttpResponse, HttpServer}; +use actix_web_actors::ws; +use container_common::anyhow_result::AnyhowResult; +use container_common::token::random_crockford_token::random_crockford_token; +use crate::endpoints_ws::obs_gateway_websocket_handler::ResponseType::TtsEvent; +use crate::server_state::ObsGatewayServerState; +use database_queries::queries::twitch::twitch_oauth::find::{TwitchOauthTokenFinder, TwitchOauthTokenRecord}; +use futures_timer::Delay; +use futures_util::FutureExt; +use http_server_common::error::common_server_error::CommonServerError; +use log::debug; +use log::error; +use log::info; +use log::warn; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2::PooledConnection; +use r2d2_redis::redis::Commands; +use redis_common::payloads::obs_active_payload::ObsActivePayload; +use redis_common::redis_keys::RedisKeys; +use redis_common::shared_constants::OBS_ACTIVE_TTL_SECONDS; +use std::sync::Arc; +use std::thread::sleep; +use std::time::Duration; +use tokio::runtime::Handle; +use twitch_api2::pubsub::Topic; +use twitch_api2::pubsub::TopicData::ChannelPointsChannelV1; +use twitch_api2::pubsub::channel_bits::ChannelBitsEventsV2Reply::BitsEvent; +use twitch_api2::pubsub; +use twitch_common::twitch_user_id::TwitchUserId; + +// TODO: Redis calls are synchronous (but fast), but is there any way to make them async? + +#[derive(Deserialize)] +pub struct PathInfo { + twitch_username: String, +} + +#[derive(Deserialize)] +pub struct QueryData { + voice_token_1: Option, + voice_token_2: Option, + // Other preferences... +} + +/// Sent back to the frontend websocket. +#[derive(Serialize, Copy, Clone)] +pub enum ResponseType { + Pong, + TtsEvent, +} + +/// Sent back to the frontend websocket. +#[derive(Serialize)] +pub struct FrontendEventPayload { + pub response_type: ResponseType, + #[serde(skip_serializing_if = "Option::is_none")] + pub tts_job_tokens: Option>, +} + +/// Endpoint +pub async fn obs_gateway_websocket_handler( + path: Path, + http_request: HttpRequest, + server_state: web::Data>, + stream: web::Payload, +) -> Result { + + let mut finder = TwitchOauthTokenFinder::new() + .allow_expired_tokens(true) + .scope_twitch_username(Some(&path.twitch_username)); + + let lookup_result = finder.perform_query(&server_state.backends.mysql_pool).await; + + let token_record = match lookup_result { + Ok(Some(token_record)) => token_record, + Ok(None) => { + warn!("Could not find Twitch user: {}", &path.twitch_username); + return Err(CommonServerError::NotFound); + }, + Err(e) => { + error!("MySQL Error: {}", e); + return Err(CommonServerError::ServerError); + }, + }; + + let twitch_user_id = TwitchUserId::from_str(&token_record.twitch_user_id) + .map_err(|e| { + error!("Error converting twitch user id: {}, id= {}", e, &token_record.twitch_user_id); + CommonServerError::ServerError + })?; + + let mut redis = server_state.backends + .redis_pool + .get() + .map_err(|err| { + error!("Could not get Redis: {:?}", err); + CommonServerError::ServerError + })?; + + let channel = RedisKeys::obs_active_sessions_topic(); + + let payload = ObsActivePayload::new(&token_record.twitch_user_id); + let json_payload = payload.serialize() + .map_err(|e| { + error!("Could not serialize JSON: {:?}", e); + CommonServerError::ServerError + })?; + + let _count_received : Option = redis.publish(channel, &json_payload) + .map_err(|e| { + warn!("redis error: {:?}", e); + CommonServerError::ServerError + })?; + + let server_state_arc = server_state.get_ref().clone(); + + let websocket = ObsGatewayWebSocket::new( + twitch_user_id.clone(), + server_state_arc + ); + + info!("Begin Javascript WebSocket..."); + + ws::start(websocket, &http_request, stream) + .map_err(|e| { + warn!("Websocket ws::start() error: {}", e); + CommonServerError::ServerError + }) +} + +/// Websocket behavior +struct ObsGatewayWebSocket { + twitch_user_id: TwitchUserId, + server_state: Arc, +} + +impl Actor for ObsGatewayWebSocket { + type Context = ws::WebsocketContext; + + fn started(&mut self, _ctx: &mut Self::Context) { + info!("ObsGatewayWebSocket Actor::started()"); + } +} + +impl StreamHandler> for ObsGatewayWebSocket { + fn handle( + &mut self, + msg: Result, + ctx: &mut Self::Context, + ) { + debug!("Socket Handler::handle()"); + + if let Ok(msg) = msg { + debug!("Socket Handler::handle(): msg = {:?}", msg); + + // TODO: Only send this every 60 seconds. + // TODO: Error handling. + // TODO: Reconnecting. + + let mut redis = self.server_state.backends + .redis_pool + .get() + .map_err(|err| { + error!("Could not get Redis: {:?}", err); + CommonServerError::ServerError + }).unwrap(); // TODO: FIXME + + let channel = RedisKeys::obs_active_sessions_topic(); + + let payload = ObsActivePayload::new(self.twitch_user_id.get_str()); + + let json_payload = payload.serialize() + .map_err(|e| { + error!("Could not serialize JSON: {:?}", e); + CommonServerError::ServerError + }).unwrap(); // TODO: FIXME + + let _count_received : Option = redis.publish(channel, &json_payload) + .map_err(|e| { + warn!("redis error: {:?}", e); + CommonServerError::ServerError + }).unwrap(); // TODO: Fixme + + // TODO: This should be done *BEFORE* PubSub + self.write_obs_active().unwrap(); + + let redis_key = RedisKeys::twitch_tts_job_queue(&self.twitch_user_id.get_str()); +// let values : Vec = redis.lpop((&redis_key, 5)).unwrap(); // TODO: Error handling +// +// if !values.is_empty() { +// info!("Got {} TTS values", values.len()); +// +// let payload = FrontendEventPayload { +// tts_job_tokens: values, +// }; +// +// match serde_json::to_string(&payload) { +// Ok(json) => { +// ctx.text(json); +// }, +// Err(e) => { +// error!("Error with JSON payload: {:?}", e); +// } +// } +// } + + let lookup_result : Option = redis.lpop(&redis_key).unwrap(); // TODO: Error handling + let mut values = Vec::new(); + + if let Some(v) = lookup_result.as_deref() { + values.push(v.to_string()); + let payload = FrontendEventPayload { + response_type: TtsEvent, + tts_job_tokens: Some(values), + }; + + match serde_json::to_string(&payload) { + Ok(json) => { + ctx.text(json); + }, + Err(e) => { + error!("Error with JSON payload: {:?}", e); + } + } + } + + match msg { + ws::Message::Ping(bytes) => { + debug!("Socket Handler::handle(): got ping"); + ctx.pong(&bytes) + }, + ws::Message::Text(text) => { + debug!("Socket Handler::handle(): got text = {:?}", text); + + let payload = FrontendEventPayload { + response_type: ResponseType::Pong, + tts_job_tokens: None, + }; + + match serde_json::to_string(&payload) { + Ok(json) => { + ctx.text(json); + }, + Err(e) => { + error!("Error with JSON PONG payload: {:?}", e); + } + } + + //ctx.text("response") + }, + ws::Message::Binary(bin) => { + debug!("Socket Handler::handle(): got binary..."); + //ctx.binary("response".as_bytes()) + }, + ws::Message::Close(reason) => { + warn!("Socket Handler::handle(): got close, reason = {:?}", reason); + ctx.close(reason); + ctx.stop(); + } + _ => {} + } + } else { + error!(">>>>>> obs streamhandler::STOP"); + ctx.stop(); + } + } +} + +impl ObsGatewayWebSocket { + fn new( + twitch_user_id: TwitchUserId, + server_state: Arc, + ) -> Self { + Self { + twitch_user_id, + server_state, + } + } + + fn get_redis(&self) + -> AnyhowResult> + { + let redis = self.server_state.backends + .redis_pool + .get() + .map_err(|err| { + error!("Could not get Redis: {:?}", err); + CommonServerError::ServerError + }).unwrap(); // TODO: FIXME + + Ok(redis) + } + + fn write_obs_active(&self) -> AnyhowResult<()> { + let mut redis = self.get_redis().unwrap(); + let redis_key = RedisKeys::obs_active_session_keepalive(self.twitch_user_id.get_str()); + let _r : bool = redis.set_ex(redis_key, "1", OBS_ACTIVE_TTL_SECONDS) + .map_err(|e| { + warn!("redis error: {:?}", e); + CommonServerError::ServerError + }).unwrap(); // TODO: Fixme + + Ok(()) + } +} + diff --git a/crates/service/websocket_gateway/src/main.rs b/crates/service/websocket_gateway/src/main.rs new file mode 100644 index 0000000000..3a2b3e4582 --- /dev/null +++ b/crates/service/websocket_gateway/src/main.rs @@ -0,0 +1,180 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate lazy_static; +#[macro_use] extern crate magic_crypt; +#[macro_use] extern crate serde_derive; + +pub mod endpoints_ws; +pub mod server_state; + +use actix_cors::Cors; +use actix_http::http; +use actix_web::middleware::{Logger, DefaultHeaders}; +use actix_web::{HttpServer, web, HttpResponse, App}; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use crate::endpoints_ws::obs_gateway_websocket_handler::obs_gateway_websocket_handler; +use crate::server_state::{ObsGatewayServerState, EnvConfig, TwitchOauthSecrets, BackendsConfig}; +use futures::Future; +use futures::executor::ThreadPool; +use http_server_common::cors::build_common_cors_config; +use http_server_common::endpoints::default_route_404::default_route_404; +use http_server_common::endpoints::root_index::get_root_index; +use limitation::Limiter; +use log::{info}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use r2d2_redis::redis::Commands; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::sync::Arc; +use std::thread::sleep; +use std::time::Duration; +use tokio::runtime::{Handle, Runtime}; +use twitch_api2::pubsub::Topic; +use twitch_api2::pubsub; +use twitch_common::twitch_secrets::TwitchSecrets; +use twitch_oauth2::tokens::UserTokenBuilder; +use twitch_oauth2::{AppAccessToken, Scope, TwitchToken, tokens::errors::AppAccessTokenError, ClientId, ClientSecret}; + +const DEFAULT_BIND_ADDRESS : &'static str = "0.0.0.0:54321"; + +#[actix_web::main] +//#[tokio::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + // NB: Do not check this secrets-containing dotenv file into VCS. + // This file should only contain *development* secrets, never production. + let _ = dotenv::from_filename(".env-secrets").ok(); + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("obs-gateway-server-unknown".to_string()); + + info!("Hostname: {}", &server_hostname); + + info!("Reading Twitch secrets..."); + + let twitch_secrets = TwitchSecrets::from_env()?; + + info!("Reading env vars and setting up utils..."); + + let bind_address = easyenv::get_env_string_or_default("BIND_ADDRESS", DEFAULT_BIND_ADDRESS); + let num_workers = easyenv::get_env_num("NUM_WORKERS", 8)?; + let hmac_secret = easyenv::get_env_string_or_default("COOKIE_SECRET", "notsecret"); + let cookie_domain = easyenv::get_env_string_or_default("COOKIE_DOMAIN", ".vo.codes"); + let cookie_secure = easyenv::get_env_bool_or_default("COOKIE_SECURE", true); + let cookie_http_only = easyenv::get_env_bool_or_default("COOKIE_HTTP_ONLY", true); + let website_homepage_redirect = + easyenv::get_env_string_or_default("WEBSITE_HOMEPAGE_REDIRECT", "https://vo.codes/"); + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + let redis_connection_string = + easyenv::get_env_string_or_default( + "REDIS_1_URL", + DEFAULT_REDIS_DATABASE_1_CONNECTION_STRING); + + info!("Connecting to mysql..."); + + let pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + info!("Connecting to redis..."); + + let redis_manager = RedisConnectionManager::new(redis_connection_string.clone())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let server_state = ObsGatewayServerState { + env_config: EnvConfig { + num_workers, + bind_address, + cookie_domain, + cookie_secure, + cookie_http_only, + website_homepage_redirect, + }, + twitch_oauth_secrets: TwitchOauthSecrets { + client_id: twitch_secrets.app_client_id.clone(), + client_secret: twitch_secrets.app_client_secret.clone(), + }, + hostname: server_hostname, + backends: BackendsConfig { + mysql_pool: pool, + redis_pool, + } + }; + + info!("Starting server..."); + + serve(server_state).await?; + + Ok(()) +} + +pub async fn serve(server_state: ObsGatewayServerState) -> AnyhowResult<()> +{ + let bind_address = server_state.env_config.bind_address.clone(); + let num_workers = server_state.env_config.num_workers.clone(); + let hostname = server_state.hostname.clone(); + + let server_state_arc = web::Data::new(Arc::new(server_state)); + + info!("Starting HTTP service."); + + let log_format = "[%{HOSTNAME}e] %a \"%r\" %s %b \"%{Referer}i\" \"%{User-Agent}i\" %T"; + + HttpServer::new(move || { + App::new() + .app_data(server_state_arc.clone()) + .wrap(build_common_cors_config()) + .wrap(Logger::new(&log_format) + .exclude("/liveness") + .exclude("/readiness")) + .wrap(DefaultHeaders::new() + .header("X-Backend-Hostname", &hostname) + .header("X-Build-Sha", "")) + .service(web::resource("/") + .route(web::get().to(get_root_index)) + ) + // Twitch websocket + .service(web::resource("/obs/{twitch_username}") + .route(web::get().to(obs_gateway_websocket_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + // Local development debugging + .service( + actix_files::Files::new("/static", "static") + .show_files_listing() + .use_last_modified(true), + ) + .default_service( web::route().to(default_route_404)) + }) + .bind(bind_address)? + .workers(num_workers) + .run() + .await?; + + Ok(()) +} diff --git a/crates/service/websocket_gateway/src/server_state.rs b/crates/service/websocket_gateway/src/server_state.rs new file mode 100644 index 0000000000..230ebdcbec --- /dev/null +++ b/crates/service/websocket_gateway/src/server_state.rs @@ -0,0 +1,41 @@ +use r2d2_redis::{r2d2, RedisConnectionManager}; +use sqlx::MySqlPool; +use twitch_oauth2::{ClientId, ClientSecret}; + +/// State that is injected into every endpoint. +#[derive(Clone)] +pub struct ObsGatewayServerState { + /// Configuration from ENV vars. + /// Some of this might not be used. + pub env_config: EnvConfig, + pub hostname: String, + + pub twitch_oauth_secrets: TwitchOauthSecrets, + + pub backends: BackendsConfig, +} + +#[derive(Clone)] +pub struct BackendsConfig { + pub mysql_pool: MySqlPool, + pub redis_pool: r2d2::Pool, +} + +// TODO: Many of these do not need to be passed around past server init. +#[derive(Clone)] +pub struct EnvConfig { + // Number of thread workers. + pub num_workers: usize, + pub bind_address: String, + pub cookie_domain: String, + pub cookie_secure: bool, + pub cookie_http_only: bool, + pub website_homepage_redirect: String, +} + +/// Necessary to run the OAuth flow. +#[derive(Clone)] +pub struct TwitchOauthSecrets { + pub client_id: String, + pub client_secret: String, +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/.github/ISSUE_TEMPLATE/bug_report.md b/crates/vendor/newrelic-telemetry-sdk-rust/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000000..71eb0035eb --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,38 @@ +--- +name: Bug report +about: Describe a scenario in which this project behaves unexpectedly +title: '' +labels: bug, needs-triage +assignees: '' + +--- + +[NOTE]: # ( ^^ Provide a general summary of the issue in the title above. ^^ ) + +## Description + +[NOTE]: # ( Describe the problem you're encountering. ) +[TIP]: # ( Do NOT give us access or passwords to your New Relic account or API keys! ) + +## Steps to Reproduce + +[NOTE]: # ( Please be as specific as possible. ) + +## Expected Behavior + +[NOTE]: # ( Tell us what you expected to happen. ) + +## Relevant Logs / Console output + +[NOTE]: # ( Please provide specifics of the local error logs, Browser Dev Tools console, etc. if appropriate and possible. ) + +## Your Environment + +[TIP]: # ( Include as many relevant details about your environment as possible. ) + +* ex: Browser name and version: +* ex: Operating System and version: + +## Additional context + +[TIP]: # ( Add any other context about the problem here. ) diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/.github/ISSUE_TEMPLATE/enhancement.md b/crates/vendor/newrelic-telemetry-sdk-rust/.github/ISSUE_TEMPLATE/enhancement.md new file mode 100644 index 0000000000..0907a9ff87 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/.github/ISSUE_TEMPLATE/enhancement.md @@ -0,0 +1,27 @@ +--- +name: Enhancement request +about: Suggest an idea for a future version of this project +title: '' +labels: enhancement, needs-triage +assignees: '' + +--- + +[NOTE]: # ( ^^ Provide a general summary of the request in the title above. ^^ ) + +## Summary + +[NOTE]: # ( Provide a brief overview of what the new feature is all about. ) + +## Desired Behavior + +[NOTE]: # ( Tell us how the new feature should work. Be specific. ) +[TIP]: # ( Do NOT give us access or passwords to your New Relic account or API keys! ) + +## Possible Solution + +[NOTE]: # ( Not required. Suggest how to implement the addition or change. ) + +## Additional context + +[TIP]: # ( Why does this feature matter to you? What unique circumstances do you have? ) diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/.github/dependabot.yml b/crates/vendor/newrelic-telemetry-sdk-rust/.github/dependabot.yml new file mode 100644 index 0000000000..f1d26d808b --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/.github/dependabot.yml @@ -0,0 +1,11 @@ +# To get started with Dependabot version updates, you'll need to specify which +# package ecosystems to update and where the package manifests are located. +# Please see the documentation for all configuration options: +# https://help.github.com/github/administering-a-repository/configuration-options-for-dependency-updates + +version: 2 +updates: + - package-ecosystem: "cargo" # See documentation for possible values + directory: "/" # Location of package manifests + schedule: + interval: "daily" diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/.github/workflows/ci.yml b/crates/vendor/newrelic-telemetry-sdk-rust/.github/workflows/ci.yml new file mode 100644 index 0000000000..72a0b5a5ea --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/.github/workflows/ci.yml @@ -0,0 +1,38 @@ +name: CI + +on: + push: + branches: [ main ] + pull_request: + branches: [ main ] + +jobs: + ubuntu_build: + name: Ubuntu Build + runs-on: ubuntu-latest + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Install latest rust toolchain + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + default: true + override: true + - name: run tests + run: cargo test --all-features -- --test-threads=1 + + windows_build: + name: Windows Build + runs-on: windows-latest + steps: + - name: Checkout + uses: actions/checkout@v2 + - name: Install latest rust toolchain + uses: actions-rs/toolchain@v1 + with: + toolchain: stable + default: true + override: true + - name: run tests + run: cargo test --all-features -- --test-threads=1 diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/.gitignore b/crates/vendor/newrelic-telemetry-sdk-rust/.gitignore new file mode 100644 index 0000000000..96ef6c0b94 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/.gitignore @@ -0,0 +1,2 @@ +/target +Cargo.lock diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/CODE_OF_CONDUCT.md b/crates/vendor/newrelic-telemetry-sdk-rust/CODE_OF_CONDUCT.md new file mode 100644 index 0000000000..9503218f9e --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/CODE_OF_CONDUCT.md @@ -0,0 +1,5 @@ +### Hello! We're glad you've joined us. + +We believe participation in our community should be a harassment free experience for everyone. + +Learn more about our guidelines and principles by reading our [Code of Conduct](https://opensource.newrelic.com/code-of-conduct/) on our Open Source Website. diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/CONTRIBUTING.md b/crates/vendor/newrelic-telemetry-sdk-rust/CONTRIBUTING.md new file mode 100644 index 0000000000..e16db5528e --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/CONTRIBUTING.md @@ -0,0 +1,30 @@ +# Contributing + +Contributions are always welcome. Before contributing please read the +[code of conduct](./CODE_OF_CONDUCT.md) and [search the issue tracker](issues); your issue may have already been discussed or fixed in `main`. To contribute, +[fork](https://help.github.com/articles/fork-a-repo/) this repository, commit your changes, and [send a Pull Request](https://help.github.com/articles/using-pull-requests/). + +Note that our [code of conduct](./CODE_OF_CONDUCT.md) applies to all platforms and venues related to this project; please follow it in all your interactions with the project and its participants. + +## Feature Requests + +Feature requests should be submitted in the [Issue tracker](../../issues), with a description of the expected behavior & use case, where they’ll remain closed until sufficient interest, [e.g. :+1: reactions](https://help.github.com/articles/about-discussions-in-issues-and-pull-requests/), has been [shown by the community](../../issues?q=label%3A%22votes+needed%22+sort%3Areactions-%2B1-desc). +Before submitting an Issue, please search for similar ones in the +[closed issues](../../issues?q=is%3Aissue+is%3Aclosed+label%3Aenhancement). + +## Pull Requests + +1. Ensure any install or build dependencies are removed before the end of the layer when doing a build. +2. Increase the version numbers in any examples files and the README.md to the new version that this Pull Request would represent. The versioning scheme we use is [SemVer](http://semver.org/). +3. You may merge the Pull Request in once you have the sign-off of two other developers, or if you do not have permission to do that, you may request the second reviewer to merge it for you. + +## Contributor License Agreement + +Keep in mind that when you submit your Pull Request, you'll need to sign the CLA via the click-through using CLA-Assistant. If you'd like to execute our corporate CLA, or if you have any questions, please drop us an email at opensource@newrelic.com. + +For more information about CLAs, please check out Alex Russell’s excellent post, +[“Why Do I Need to Sign This?”](https://infrequently.org/2008/06/why-do-i-need-to-sign-this/). + +## Slack + +We host a public Slack with a dedicated channel for contributors and maintainers of open source projects hosted by New Relic. If you are contributing to this project, you're welcome to request access to the #oss-contributors channel in the newrelicusers.slack.com workspace. To request access, see https://newrelicusers-signup.herokuapp.com/. diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/Cargo.toml b/crates/vendor/newrelic-telemetry-sdk-rust/Cargo.toml new file mode 100644 index 0000000000..b11c42a815 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "newrelic-telemetry" +version = "0.1.0" +authors = ["New Relic "] +edition = "2018" + +[features] +default = [ "client" ] +client = [ "flate2", "hyper", "hyper-tls" ] +blocking = [ "client", "futures", "tokio" ] + +[dependencies] +anyhow = "1.0" +log = "0.4" +serde = { version = "1.0.57", features = ["derive"] } +serde_json = "1.0.57" + +hyper = { version = "0.14", features = [], optional = true } +hyper-tls = { version = "0.5.0", optional = true } +flate2 = { version = "1.0.16", optional = true } +futures = { version = "0.3.1", optional = true } +tokio = { version = "1.1.0", features = ["macros"], optional = true } +uuid = { version = "0.8.1", features = ["v4"] } + +[dev-dependencies] +rand = "0.7.3" +env_logger = "0.7.1" +futures = { version = "0.3.1" } +tokio = { version = "1.1.0", features = ["macros"] } diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/LICENSE b/crates/vendor/newrelic-telemetry-sdk-rust/LICENSE new file mode 100644 index 0000000000..261eeb9e9f --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/README.md b/crates/vendor/newrelic-telemetry-sdk-rust/README.md new file mode 100644 index 0000000000..08194ef18b --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/README.md @@ -0,0 +1,108 @@ +[![Community Project header](https://github.com/newrelic/open-source-office/raw/master/examples/categories/images/Community_Project.png)](https://github.com/newrelic/open-source-office/blob/master/examples/categories/index.md#community-project) + +# New Relic Rust Telemetry SDK + +![Build status](https://github.com/newrelic/newrelic-telemetry-sdk-rust/workflows/CI/badge.svg) +[![Release](https://img.shields.io/github/v/release/newrelic/newrelic-telemetry-sdk-rust?include_prereleases&style=)](https://github.com/newrelic/newrelic-telemetry-sdk-rust/releases/) + +What is the New Relic Rust Telemetry SDK? + +* It's a helper library that supports sending New Relic data from within your Rust application. +* It's the foundation of [New Relic's C Telemetry SDK](https://github.com/newrelic/newrelic-telemetry-sdk-c). +* It’s an example of "best practices" for sending us data. + +[The Telemetry SDK](https://docs.newrelic.com/docs/telemetry-data-platform/get-started/capabilities/telemetry-sdks-send-custom-telemetry-data-new-relic) provides you, the end-user programmer, with a `Client `that sends `Spans` to New Relic. Individual spans are collected together into batches (via a `SpanBatch` object), and clients send these batches. It serves as a foundation for getting open-standards based telemetry data like [OpenCensus](https://opencensus.io/), [OpenTracing](https://opentracing.io/), and [OpenTelemetry](https://opentelemetry.io/) into New Relic. You can use this to build tracers/exporters, such as ones based on these open standards. + +This SDK currently supports sending spans to the [Trace API](https://docs.newrelic.com/docs/understand-dependencies/distributed-tracing/trace-api/introduction-trace-api). + +**This project is currently in an alpha state.** The API of this crate is +likely to change, therefore there's no API documentation available yet. + +## Getting Started + +In order to send telemetry data to New Relic APIs, you will need an Insert API key. To find out how to generate this key, see our [docs](https://docs.newrelic.com/docs/apis/get-started/intro-apis/types-new-relic-api-keys). + +The Rust Telemetry SDK is not published on [crates.io](https://crates.io) +yet. In order to use it in your Rust application, you have to clone this +repository locally and include the crate via the local path name. If you +cloned the Rust Telemetry SDK into a `vendor` subfolder of your Rust crate, +you can add it as a dependency to your crate by adding the following dependency +to your `Cargo.toml`: + +```toml +[dependencies] +newrelic-telemetry = { path = "vendor/newrelic-telemetry-sdk-rust" } +``` + +This is a simple application that sends a single span via the asynchronous +client: +```rust +use newrelic_telemetry::{ClientBuilder, Span}; +use std::env; +use std::time::{SystemTime, UNIX_EPOCH}; + +#[tokio::main] +async fn main() { + // Obtain a license key from the environment. + let license_key = env::var("NEW_RELIC_API_KEY").unwrap(); + + // Build an asynchronous client. + let client = ClientBuilder::new(&license_key).build().unwrap(); + + // Create a span and a span batch. + let span = Span::new( + "e9f54a2c322d7578", + "1b1bf29379951c1d", + SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_millis() as u64, + ) + .name("/index.html") + .attribute("username", "user") + .service_name("Telemetry Application"); + + let span_batch = vec![span].into(); + + // Send a span batch and await the future. + client.send_spans(span_batch).await; +} +``` + +## Find and use your data + +Tips on how to find and query your data in New Relic: +- [Find trace/span data](https://docs.newrelic.com/docs/understand-dependencies/distributed-tracing/trace-api/introduction-trace-api#view-data) + +For general querying information, see: +- [Query New Relic data](https://docs.newrelic.com/docs/query-your-data/explore-query-data/explore-data/introduction-querying-new-relic-data) +- [Intro to NRQL](https://docs.newrelic.com/docs/query-your-data/nrql-new-relic-query-language/get-started/introduction-nrql-new-relics-query-language) + +## Support + +New Relic hosts and moderates an online forum where customers can interact with +New Relic employees as well as other customers to get help and share best +practices. Like all official New Relic open source projects, there's a related +Community topic in the New Relic Explorers Hub. You can find this project's +topic/threads in the [Telemetry SDK section of Explorers Hub](https://discuss.newrelic.com/t/new-relic-rust-telemetry-sdk/114558) + +## Contributing + +We encourage your contributions to improve the Rust Telemetry SDK! Keep in mind +when you submit your pull request, you'll need to sign the CLA via the +click-through using CLA-Assistant. You only have to sign the CLA one time per +project. If you have any questions, or to execute our corporate CLA, required +if your contribution is on behalf of a company, please drop us an email at +opensource@newrelic.com. + +## License + +The Rust Telemetry SDK is licensed under the [Apache 2.0](http://apache.org/licenses/LICENSE-2.0.txt) +License. + +### Limitations + +The New Relic Telemetry APIs are rate limited. Please reference the +documentation for [New Relic Metric API](https://docs.newrelic.com/docs/introduction-new-relic-metric-api) +and [New Relic Trace API requirements and limits](https://docs.newrelic.com/docs/apm/distributed-tracing/trace-api/trace-api-general-requirements-limits) +on the specifics of the rate limits. diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/cla.md b/crates/vendor/newrelic-telemetry-sdk-rust/cla.md new file mode 100644 index 0000000000..27af9b4dd7 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/cla.md @@ -0,0 +1,16 @@ +# NEW RELIC, INC. +## INDIVIDUAL CONTRIBUTOR LICENSE AGREEMENT +Thank you for your interest in contributing to the open source projects of New Relic, Inc. (“New Relic”). In order to clarify the intellectual property license granted with Contributions from any person or entity, New Relic must have a Contributor License Agreement ("Agreement") on file that has been signed by each Contributor, indicating agreement to the license terms below. This Agreement is for your protection as a Contributor as well as the protection of New Relic; it does not change your rights to use your own Contributions for any other purpose. + +You accept and agree to the following terms and conditions for Your present and future Contributions submitted to New Relic. Except for the licenses granted herein to New Relic and recipients of software distributed by New Relic, You reserve all right, title, and interest in and to Your Contributions. + +## Definitions. +1. "You" (or "Your") shall mean the copyright owner or legal entity authorized by the copyright owner that is entering into this Agreement with New Relic. For legal entities, the entity making a Contribution and all other entities that control, are controlled by, or are under common control with that entity are considered to be a single Contributor. For the purposes of this definition, "control" means (i) the power, direct or indirect, to cause the direction or management of such entity, whether by contract or otherwise, or (ii) ownership of fifty percent (50%) or more of the outstanding shares, or (iii) beneficial ownership of such entity. +2. "Contribution" shall mean any original work of authorship, including any modifications or additions to an existing work, that is intentionally submitted by You to New Relic for inclusion in, or documentation of, any of the products managed or maintained by New Relic (the "Work"). For the purposes of this definition, "submitted" means any form of electronic, verbal, or written communication sent to New Relic or its representatives, including but not limited to communication on electronic mailing lists, source code control systems, and issue tracking systems that are managed by, or on behalf of, New Relic for the purpose of discussing and improving the Work, but excluding communication that is conspicuously marked or otherwise designated in writing by You as "Not a Contribution." +3. Grant of Copyright License. Subject to the terms and conditions of this Agreement, You hereby grant to New Relic and to recipients of software distributed by New Relic a perpetual, worldwide, non-exclusive, no-charge, royalty-free, transferable, irrevocable copyright license to reproduce, prepare derivative works of, publicly display, publicly perform, sublicense, and distribute Your Contributions and such derivative works. +4. Grant of Patent License. Subject to the terms and conditions of this Agreement, You hereby grant to New Relic and to recipients of software distributed by New Relic a perpetual, worldwide, non-exclusive, no-charge, royalty-free, transferable, irrevocable (except as stated in this section) patent license to make, have made, use, offer to sell, sell, import, and otherwise transfer the Work, where such license applies only to those patent claims licensable by You that are necessarily infringed by Your Contributions alone or by combination of Your Contribution(s) with the Work to which such Contribution(s) was submitted. If any entity institutes patent litigation against You or any other entity (including a cross-claim or counterclaim in a lawsuit) alleging that Your Contribution, or the Work to which You have contributed, constitutes direct or contributory patent infringement, then any patent licenses granted to that entity under this Agreement for that Contribution or Work shall terminate as of the date such litigation is filed. +5. You represent that You are legally entitled to grant the above licenses. If Your employer(s) has rights to intellectual property that You create that includes Your Contributions, You represent that You have received permission to make Contributions on behalf of that employer, that Your employer has waived such rights for Your Contributions to New Relic, or that Your employer has executed a separate Agreement with New Relic. +6. You represent that each of Your Contributions is Your original creation (see section 7 for submissions on behalf of others). You represent that Your Contribution submissions include complete details of any third-party license or other restriction (including, but not limited to, related patents and trademarks) of which You are personally aware and which are associated with any part of Your Contributions. +7. You are not expected to provide support for Your Contributions, except to the extent You desire to provide support. You may provide support for free, for a fee, or not at all. Unless required by applicable law or agreed to in writing, You provide Your Contributions on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied, including, without limitation, any warranties or conditions of TITLE, NON- INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A PARTICULAR PURPOSE. +8. Should You wish to submit work that is not Your original creation, You may submit it to New Relic separately from any Contribution, identifying the complete details of its source and of any license or other restriction (including, but not limited to, related patents, trademarks, and license agreements) of which You are personally aware, and conspicuously marking the work as "Submitted on behalf of a third-party: [named here]". +9. You agree to notify New Relic of any facts or circumstances of which You become aware that would make these representations inaccurate in any respect. diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/src/attribute.rs b/crates/vendor/newrelic-telemetry-sdk-rust/src/attribute.rs new file mode 100644 index 0000000000..4b9c787aa0 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/src/attribute.rs @@ -0,0 +1,252 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// + +/// Represents any valid attribute value. +/// +/// According to the [specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/capabilities.md), +/// attribute values can be a string, numeric, or boolean. A numeric value is +/// represented either as a signed integer, an unsigned integer or a float. +#[derive(serde::Serialize, Debug, PartialEq, Clone)] +#[serde(untagged)] +pub enum Value { + /// Represents a signed integer attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::Int(-3); + /// ``` + Int(i64), + + /// Represents an unsigned integer attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::UInt(42); + /// ``` + UInt(u64), + + /// Represents a signed 128 bit integer attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::Int128(-30); + /// ``` + Int128(i128), + + /// Represents an unsigned 128 bit integer attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::UInt128(30); + /// ``` + UInt128(u128), + + /// Represents a string attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::Str(String::from("root")); + /// ``` + Str(String), + + /// Represents a float attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::Float(3.14159); + /// ``` + Float(f64), + + /// Represents a bool attribute value. + /// + /// ``` + /// # use newrelic_telemetry::attribute::Value; + /// # + /// let v = Value::Bool(true); + /// ``` + Bool(bool), +} + +/// Converts an i128 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: i128 = -10; +/// assert_eq!(Value::Int128(-10), v.into()); +/// ``` +impl From for Value { + fn from(value: i128) -> Value { + Value::Int128(value) + } +} + +/// Converts an i64 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: i64 = -5; +/// assert_eq!(Value::Int(-5), v.into()); +/// ``` +impl From for Value { + fn from(value: i64) -> Value { + Value::Int(value) + } +} + +/// Converts an i32 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: i32 = -5; +/// assert_eq!(Value::Int(-5), v.into()); +/// ``` +impl From for Value { + fn from(value: i32) -> Value { + Value::Int(value as i64) + } +} + +/// Converts a u128 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: u128 = 50; +/// assert_eq!(Value::UInt128(50), v.into()); +/// ``` +impl From for Value { + fn from(value: u128) -> Value { + Value::UInt128(value) + } +} + +/// Converts a u64 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: u64 = 5; +/// assert_eq!(Value::UInt(5), v.into()); +/// ``` +impl From for Value { + fn from(value: u64) -> Value { + Value::UInt(value) + } +} + +/// Converts a u32 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: u32 = 5; +/// assert_eq!(Value::UInt(5), v.into()); +/// ``` +impl From for Value { + fn from(value: u32) -> Value { + Value::UInt(value as u64) + } +} + +/// Converts a string to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v = "root"; +/// assert_eq!(Value::Str(String::from("root")), v.into()); +/// ``` +impl From<&str> for Value { + fn from(value: &str) -> Value { + Value::Str(value.to_string()) + } +} + +/// Converts a f64 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: f64 = 3.14159; +/// assert_eq!(Value::Float(v), v.into()); +/// ``` +impl From for Value { + fn from(value: f64) -> Value { + Value::Float(value) + } +} + +/// Converts a f32 to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v: f32 = 3.14159; +/// assert_eq!(Value::Float(v as f64), v.into()); +/// ``` +impl From for Value { + fn from(value: f32) -> Value { + Value::Float(value as f64) + } +} + +/// Converts a bool to an attribute value. +/// +/// ``` +/// # use newrelic_telemetry::attribute::Value; +/// # +/// let v = true; +/// assert_eq!(Value::Bool(true), v.into()); +/// ``` +impl From for Value { + fn from(value: bool) -> Value { + Value::Bool(value) + } +} + +#[cfg(test)] +mod tests { + use super::Value; + use serde_json::json; + + #[test] + fn value_to_json() { + // Attribute values should serialize to plain JSON values. + assert_eq!(json!(Value::Int(-5)), json!(-5)); + assert_eq!(json!(Value::UInt(5)), json!(5)); + assert_eq!(json!(Value::Float(3.14159)), json!(3.14159)); + assert_eq!(json!(Value::Str(String::from("root"))), json!("root")); + assert_eq!(json!(Value::Bool(true)), json!(true)); + } + + #[test] + fn into_value() { + // Should be able to use Value::from or .into() to create Values + assert_eq!(Value::Int(-5), Value::from(-5)); + assert_eq!(Value::Int(-5), (-5 as i32).into()); + + // cast needed because integer types default to i32 + assert_eq!(Value::UInt(5), Value::from(5 as u64)); + assert_eq!(Value::UInt(5), (5 as u64).into()); + + assert_eq!(Value::Float(3.14159), Value::from(3.14159)); + assert_eq!(Value::Float(3.14159), (3.14159 as f64).into()); + + assert_eq!(Value::Str("root".to_string()), Value::from("root")); + assert_eq!(Value::Str("root".to_string()), "root".into()); + + assert_eq!(Value::Bool(true), Value::from(true)); + assert_eq!(Value::Bool(true), true.into()); + } +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/src/client.rs b/crates/vendor/newrelic-telemetry-sdk-rust/src/client.rs new file mode 100644 index 0000000000..0efb8bf4ff --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/src/client.rs @@ -0,0 +1,971 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// +use crate::span::SpanBatch; +use anyhow::{anyhow, Result}; +use flate2::write::GzEncoder; +use flate2::Compression; +use hyper::client::HttpConnector; +use hyper::header::{CONTENT_ENCODING, CONTENT_TYPE, USER_AGENT}; +use hyper::{Body, HeaderMap, Method, Request, Response, Uri}; +use hyper_tls::HttpsConnector; +use log::{debug, error, info}; +use std::future::Future; +use std::io::Write; +use std::pin::Pin; +use std::thread; +use std::time::Duration; + +const VERSION: &'static str = env!("CARGO_PKG_VERSION"); +const TRACE_API_PATH: &'static str = "trace/v1"; + +/// Types that can be sent to a New Relic ingest API +/// +/// New Relic ingest APIs currently accept batches of traces, metrics, events +/// or logs. +pub trait Sendable: std::fmt::Display + Send { + /// Return the uuid for the `Sendable` + /// + /// This method returns a version 4 UUID string which enables the ingest + /// service to identify duplicate requests. + fn uuid(&self) -> &str; + + // Create a payload + // + // This method creates a JSON payload representing the contents of the + // `Sendable` object, conforming to the requirements of a related ingest + // API (traces, metrics, events or logs). + fn marshall(&self) -> Result; + + // Split a `Sendable` + // + // New Relic ingest APIs reject payloads that are too large. In that case, + // a 413 response code is sent, the payload must be split and sent again + // (see [the specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/communication.md#response-codes) + // for further details). + // + // This method removes half of the content of the `Sendable` object and + // puts it into a second `Sendable` object, which is returned. + fn split(&mut self) -> Box; +} + +// Represents a New Relic ingest endpoint. +#[derive(Debug)] +struct Endpoint { + // The host name or address of the endpoint. + host: String, + + // The port of the endpoint. This is optional, if not given it will default + // to the standard HTTPS port. + port: Option, + + // The path for the endpoint. + path: &'static str, +} + +impl Endpoint { + // Create an URI from an endpoint. + // + // This uses the parser of `hyper::Uri` to validate the URI and returns + // `https` or `http` URIs, based on the `use_tls` flag. + fn uri(&self, use_tls: bool) -> Result { + let port_str = match self.port { + Some(p) => format!(":{}", p), + _ => "".to_string(), + }; + + let uri = format!( + "{}://{}{}/{}", + if use_tls { "https" } else { "http" }, + self.host, + port_str, + self.path + ); + + Ok(uri.parse::()?) + } +} + +/// `ClientBuilder` acts as builder for initializing a `Client`. +/// +/// It can be used to customize ingest URLs, the backoff factor, the retry +/// maximum, and the product info. +/// +/// ``` +/// # use anyhow::Result; +/// # use newrelic_telemetry::ClientBuilder; +/// # use std::time::Duration; +/// # fn main() -> Result<()> { +/// # let api_key = ""; +/// let mut builder = ClientBuilder::new(api_key); +/// +/// let client = builder.backoff_factor(Duration::from_secs(10)) +/// .product_info("RustDoc", "1.0") +/// .build()?; +/// # Ok(()) +/// # } +/// ``` +pub struct ClientBuilder { + api_key: String, + backoff_factor: Duration, + retries_max: u32, + endpoint_traces: Endpoint, + product_info: Option<(String, String)>, + blocking_queue_max: usize, + use_tls: bool, +} + +impl ClientBuilder { + /// Initialize the client builder with an API key. + /// + /// Other values will be set to defaults: + /// * The default backoff factor will be 5 seconds. + /// * The default maximum of retries is 8. + /// * The default trace endpoint is `https://trace-api.newrelic.com/trace/v1` on port 80. + /// * By default, product information is empty. + /// * By default, no more than 100 batches are sent in one go in blocking mode. + /// + /// ``` + /// # use newrelic_telemetry::ClientBuilder; + /// # let api_key = ""; + /// let mut builder = ClientBuilder::new(api_key); + /// ``` + pub fn new(api_key: &str) -> Self { + ClientBuilder { + api_key: api_key.to_string(), + backoff_factor: Duration::from_secs(5), + retries_max: 8, + endpoint_traces: Endpoint { + host: "trace-api.newrelic.com".to_string(), + port: None, + path: TRACE_API_PATH, + }, + product_info: None, + blocking_queue_max: 100, + use_tls: true, + } + } + + /// Configures a backoff factor. + /// + /// If a request fails, the SDK retries the request at increasing intervals + /// and eventually drops data if the request cannot be completed. + /// + /// The amount of time to wait after a request can be computed using this + /// logic: + /// + /// `backoff_factor * (2 ^ (number_of_retries - 1))` + /// + /// For a backoff factor of 1 second, and a maximum of 6 retries, the retry + /// delay interval should follow a pattern of [0, 1, 2, 4, 8, 16]. + /// + /// See the [specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/communication.md#graceful-degradation) + /// for further details. + /// + /// ``` + /// # use newrelic_telemetry::ClientBuilder; + /// # use std::time::Duration; + /// # let api_key = ""; + /// let mut builder = + /// ClientBuilder::new(api_key).backoff_factor(Duration::from_secs(10)); + /// ``` + pub fn backoff_factor(mut self, factor: Duration) -> Self { + self.backoff_factor = factor; + self + } + + /// Configures the maximum numbers of retries. + /// + /// If a request fails, the SDK retries the request at increasing intervals + /// and eventually drops data if the request cannot be completed. + /// + /// If zero is given as a maximum, no retries will be made for failed + /// requests. + /// + /// For a backoff factor of 1 second, and a maximum of 6 retries, the retry + /// delay interval should follow a pattern of [0, 1, 2, 4, 8, 16]. + /// + /// See the [specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/communication.md#graceful-degradation) + /// for further details. + /// + /// ``` + /// # use newrelic_telemetry::ClientBuilder; + /// # let api_key = ""; + /// let mut builder = + /// ClientBuilder::new(api_key).retries_max(4); + /// ``` + pub fn retries_max(mut self, retries: u32) -> Self { + self.retries_max = retries; + self + } + + /// Configure the ingest host for traces. + /// + /// Overrides the default ingest host for traces to facilitate communication + /// with alternative New Relic backends. + /// + /// ``` + /// # use newrelic_telemetry::ClientBuilder; + /// # let api_key = ""; + /// let mut builder = + /// ClientBuilder::new(api_key).endpoint_traces("127.0.0.1", None); + /// ``` + pub fn endpoint_traces(mut self, url: &str, port: Option) -> Self { + self.endpoint_traces = Endpoint { + host: url.to_string(), + path: TRACE_API_PATH, + port: port, + }; + self + } + + /// Configure a product and version. + /// + /// The specified product and version will be appended to the `User-Agent` + /// header of payloads. + /// + /// See the [specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/communication.md#user-agent) + /// for further details. + /// + /// ``` + /// # use newrelic_telemetry::ClientBuilder; + /// # let api_key = ""; + /// let mut builder = + /// ClientBuilder::new(api_key).product_info("NewRelic-Cpp-OpenTelemetry", "0.2.1"); + /// ``` + pub fn product_info(mut self, product: &str, version: &str) -> Self { + self.product_info = Some((product.to_string(), version.to_string())); + self + } + + /// Configure the maximum number of batches sent in one go in blocking mode. + /// + /// This configuration has no effect for default non-blocking clients. + /// + /// If the number of batches in the blocking client's batch queue exceeds + /// the maximum given here, the addditional batches will be dropped. This + /// mechanism avoids accumulating back pressure. + /// + /// ``` + /// # use newrelic_telemetry::ClientBuilder; + /// # let api_key = ""; + /// let mut builder = + /// ClientBuilder::new(api_key).blocking_queue_max(10); + /// ``` + pub fn blocking_queue_max(mut self, queue_max: usize) -> Self { + self.blocking_queue_max = queue_max; + self + } + + // Configure TLS usage. + // + // New Relic endpoints exclusively support HTTPS. This is mainly provided + // for testing purposes. + pub fn tls(mut self, tls: bool) -> Self { + self.use_tls = tls; + self + } + + /// Build a client. + /// + /// ``` + /// # use anyhow::Result; + /// # use newrelic_telemetry::ClientBuilder; + /// # fn main() -> Result<()> { + /// # let api_key = ""; + /// let builder = ClientBuilder::new(api_key); + /// + /// let client = builder.build()?; + /// # Ok(()) + /// # } + /// ``` + pub fn build(self) -> Result { + Client::new(self) + } + + /// Build a blocking client. + /// + /// ``` + /// # use anyhow::Result; + /// # use newrelic_telemetry::ClientBuilder; + /// # fn main() -> Result<()> { + /// # let api_key = ""; + /// let builder = ClientBuilder::new(api_key); + /// + /// let client = builder.build_blocking()?; + /// # Ok(()) + /// # } + /// ``` + #[cfg(feature = "blocking")] + pub fn build_blocking(self) -> Result { + blocking::Client::new(self) + } + + fn get_backoff_sequence(&self) -> Vec { + (0..self.retries_max) + .map(|num_retry| { + if num_retry == 0 { + Duration::from_secs(0) + } else { + self.backoff_factor * (2_u32.pow(num_retry - 1)) + } + }) + .collect() + } + + fn get_user_agent_header(&self) -> String { + let product_info = match &self.product_info { + Some(s) => format!(" {}/{}", s.0, s.1), + _ => "".to_string(), + }; + + format!("NewRelic-Rust-TelemetrySDK/{}{}", VERSION, product_info) + } +} + +// An internal enum representing the state of a payload. +#[derive(Debug, PartialEq)] +enum SendableState { + // No retry should be made. + Done, + + // A retry should be made. Either after the given duration, or, if it + // is `None`, according to the backoff sequence. + Retry(Option), + + // The payload should be split and a retry should be made for both + // payloads. + Split, +} + +pub struct Client { + api_key: String, + user_agent: String, + backoff_sequence: Vec, + endpoint_traces: Uri, + client: hyper::Client>, +} + +impl Client { + /// Constructs a `Client` from a `ClientBuilder`. + pub fn new(builder: ClientBuilder) -> Result { + let https = HttpsConnector::new(); + let user_agent = builder.get_user_agent_header(); + let backoff_seq = builder.get_backoff_sequence(); + + Ok(Client { + api_key: builder.api_key, + endpoint_traces: builder.endpoint_traces.uri(builder.use_tls)?, + user_agent: user_agent, + backoff_sequence: backoff_seq, + client: hyper::Client::builder().build::<_, hyper::Body>(https), + }) + } + + /// Sends a span batch. + /// + /// This asynchronously sends a span batch, encapsulating retry and backoff + /// mechanisms defined in the [specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/communication.md) + /// and customized via the `ClientBuilder`. + pub async fn send_spans(&self, batch: SpanBatch) { + self.send(Box::new(batch), &self.endpoint_traces).await + } + + // Returns a gzip compressed version of the given string. + fn to_gzip(text: &String) -> Result> { + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(text.as_bytes())?; + Ok(encoder.finish()?) + } + + // Extract the value of the Retry-After HTTP response header + fn extract_retry_after(headers: &HeaderMap) -> Result { + if let Some(dur) = headers.get("retry-after") { + Ok(Duration::from_secs(dur.to_str()?.parse::()?)) + } else { + Err(anyhow!("missing retry-after header")) + } + } + + // Sends a given `Sendable` asynchronously to a given endpoint. + fn send<'a>( + &'a self, + mut batch: Box, + endpoint: &'a Uri, + ) -> Pin + 'a>> { + Box::pin(async move { + for duration in self.backoff_sequence.iter() { + let request = match self.request(&*batch, endpoint) { + Ok(r) => r, + Err(e) => { + error!("cannot create request for {}, dropping due to {}", batch, e); + return; + } + }; + + let response = match self.client.request(request).await { + Ok(r) => r, + Err(e) => { + error!("cannot send request for {}, dropping due to {}", batch, e); + return; + } + }; + + let status = Self::process_response(&*batch, response); + + let duration = match status { + SendableState::Done => return, + SendableState::Retry(Some(duration)) => duration, + SendableState::Split => { + let batch2 = batch.split(); + self.send(batch, endpoint).await; + self.send(batch2, endpoint).await; + return; + } + _ => *duration, + }; + + thread::sleep(duration); + } + }) + } + + // Create a request from the given batch and endpoint. + fn request<'a>(&self, batch: &(dyn Sendable + 'a), endpoint: &Uri) -> Result> { + let raw = batch.marshall()?; + let gzipped = Self::to_gzip(&raw)?; + + Ok(Request::builder() + .method(Method::POST) + .uri(endpoint) + .header("Api-Key", &self.api_key) + .header("Data-Format", "newrelic") + .header("Data-Format-Version", "1") + .header("x-request-id", batch.uuid()) + .header(USER_AGENT, &self.user_agent) + .header(CONTENT_ENCODING, "gzip") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(gzipped))?) + } + + // Based on the response from an ingest endpoint, decide whether to + // retry or split a payload. + // + // See the [specification](https://github.com/newrelic/newrelic-telemetry-sdk-specs/blob/master/communication.md#response-codes) + // for further details. + fn process_response<'a, T>( + batch: &(dyn Sendable + 'a), + response: Response, + ) -> SendableState { + let status = response.status(); + + match status.as_u16() { + 200..=299 => { + debug!("response {}, successfully sent {}", status, batch); + } + 400 | 401 | 403 | 404 | 405 | 409 | 410 | 411 => { + error!("response {}, dropping {}", status, batch); + } + 413 => { + info!( + "response {}, payload too large, splitting {}", + status, batch + ); + return SendableState::Split; + } + 429 => match Self::extract_retry_after(response.headers()) { + Ok(duration) => { + info!( + "response {}: retry interval {:?}, retrying {}", + status, duration, batch + ); + + return SendableState::Retry(Some(duration)); + } + Err(e) => { + error!("response {}, {}, dropping {}", status, e, batch); + } + }, + _ => { + debug!("response {}, retry {}", status, batch); + return SendableState::Retry(None); + } + } + return SendableState::Done; + } +} + +#[cfg(feature = "blocking")] +pub mod blocking { + use super::{ClientBuilder, SpanBatch}; + use anyhow::Result; + use futures::future; + use log::warn; + use std::sync::mpsc; + use std::sync::Mutex; + use std::thread; + use tokio::runtime::Builder; + + enum SendableType { + Spans(SpanBatch), + } + + pub struct Client { + channel: Mutex>>, + handle: thread::JoinHandle<()>, + } + + impl Client { + pub fn new(builder: ClientBuilder) -> Result { + let (tx, rx) = mpsc::channel::>(); + let mut runtime = Builder::new().threaded_scheduler().enable_all().build()?; + let queue_max = builder.blocking_queue_max; + let client = builder.build()?; + + let handle = thread::spawn(move || loop { + let mut batches = vec![]; + + // Wait until at least one batch is received. + match rx.recv() { + Ok(b) => batches.push(b), + Err(_) => break, + }; + + // Empty the channel. + loop { + match rx.try_recv() { + Ok(b) => batches.push(b), + Err(_) => break, + } + } + + // Drop batches that exceed the maximum defined queue size. + if batches.len() > queue_max { + warn!( + "back pressure, dropping {} span batches", + batches.len() - queue_max + ); + batches.drain(queue_max..); + } + + // Block until all batches are sent. + runtime.block_on(future::join_all(batches.drain(..).map(|b| match *b { + SendableType::Spans(batch) => client.send_spans(batch), + }))); + }); + + Ok(Client { + channel: Mutex::new(tx), + handle, + }) + } + + pub fn send_spans(&self, b: SpanBatch) { + if let Ok(ch) = self.channel.lock() { + if let Err(_) = ch.send(Box::new(SendableType::Spans(b))) {} + } + } + + pub fn shutdown(self) { + drop(self.channel); + + let _ = self.handle.join(); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use anyhow::Result; + use flate2::read::GzDecoder; + use hyper::header::{HeaderValue, CONTENT_ENCODING, CONTENT_TYPE, USER_AGENT}; + use hyper::{Method, Response}; + use std::fmt; + use std::io::Read; + use std::time::Duration; + pub struct TestBatch; + + impl Sendable for TestBatch { + fn uuid(&self) -> &str { + "" + } + + fn marshall(&self) -> Result { + Ok("".to_string()) + } + + fn split(&mut self) -> Box { + Box::new(TestBatch) + } + } + + impl fmt::Display for TestBatch { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!(f, "") + } + } + + #[test] + fn build() -> Result<()> { + let client = ClientBuilder::new("0000") + .backoff_factor(Duration::from_secs(2)) + .retries_max(6) + .endpoint_traces("127.0.0.1", Some(8080)) + .product_info("Test", "1.0") + .build()?; + + assert_eq!(client.api_key, "0000"); + assert_eq!(client.endpoint_traces.host(), Some("127.0.0.1")); + assert_eq!(client.endpoint_traces.port_u16(), Some(8080)); + assert_eq!(client.endpoint_traces.scheme().unwrap().as_str(), "https"); + assert_eq!( + client.user_agent, + format!("NewRelic-Rust-TelemetrySDK/{} Test/1.0", VERSION) + ); + assert_eq!( + client.backoff_sequence, + vec![0, 2, 4, 8, 16, 32] + .into_iter() + .map(|d| Duration::from_secs(d)) + .collect::>() + ); + + Ok(()) + } + + #[test] + fn build_error() { + let client = ClientBuilder::new("0000") + .endpoint_traces(":80", Some(8080)) + .build(); + + assert!(client.is_err()); + } + + #[test] + fn uri_from_endpoint_ok() -> Result<()> { + let endpoint = Endpoint { + host: "host".to_string(), + path: TRACE_API_PATH, + port: Some(80), + }; + + let uri = endpoint.uri(true)?; + assert_eq!(uri.host(), Some("host")); + assert_eq!(uri.port_u16(), Some(80)); + assert_eq!(uri.path(), "/trace/v1"); + assert_eq!(uri.scheme().unwrap().as_str(), "https"); + + Ok(()) + } + + #[test] + fn uri_from_endpoint_no_tls_ok() -> Result<()> { + let endpoint = Endpoint { + host: "host".to_string(), + path: TRACE_API_PATH, + port: Some(80), + }; + + let uri = endpoint.uri(false)?; + assert_eq!(uri.host(), Some("host")); + assert_eq!(uri.port_u16(), Some(80)); + assert_eq!(uri.path(), "/trace/v1"); + assert_eq!(uri.scheme().unwrap().as_str(), "http"); + + Ok(()) + } + + #[test] + fn uri_from_endpoint_error() -> Result<()> { + for endpoint in vec![ + Endpoint { + host: "host:80".to_string(), + path: TRACE_API_PATH, + port: Some(80), + }, + Endpoint { + host: "?".to_string(), + path: TRACE_API_PATH, + port: Some(80), + }, + Endpoint { + host: "".to_string(), + path: TRACE_API_PATH, + port: None, + }, + ] { + let uri = endpoint.uri(true); + + assert!( + uri.is_err(), + format!("Could create an uri from {:?}: {:?}", endpoint, uri) + ); + } + + Ok(()) + } + + #[test] + fn to_gzip() -> Result<()> { + let text = "Text to be encoded".to_string(); + let encoded = Client::to_gzip(&text)?; + + let mut gz = GzDecoder::new(&encoded[..]); + let mut decoded = String::new(); + gz.read_to_string(&mut decoded)?; + + assert_eq!(decoded, text); + + Ok(()) + } + + #[test] + fn extract_retry_after() -> Result<()> { + let mut headers = hyper::HeaderMap::new(); + + let when = Client::extract_retry_after(&headers); + assert!(when.is_err()); + + headers.insert("Retry-after", "7".parse()?); + + let when = Client::extract_retry_after(&headers)?; + assert_eq!(when, Duration::from_secs(7)); + + headers.insert("Retry-after", "seven".parse()?); + + let when = Client::extract_retry_after(&headers); + assert!(when.is_err()); + + Ok(()) + } + + #[test] + fn process_response_success() -> Result<()> { + for code in 200..300 { + let batch = Box::new(TestBatch); + let response = Response::builder().status(code).body(())?; + + assert_eq!( + Client::process_response(&*batch, response), + SendableState::Done + ); + } + + Ok(()) + } + + #[test] + fn process_response_error() -> Result<()> { + for code in vec![400, 401, 403, 404, 405, 409, 410, 411] { + let batch = Box::new(TestBatch); + let response = Response::builder().status(code).body(())?; + + assert_eq!( + Client::process_response(&*batch, response), + SendableState::Done + ); + } + + Ok(()) + } + + #[test] + fn process_response_split() -> Result<()> { + let batch = Box::new(TestBatch); + let response = Response::builder().status(413).body(())?; + + assert_eq!( + Client::process_response(&*batch, response), + SendableState::Split + ); + + Ok(()) + } + + #[test] + fn process_response_retry_from_header() -> Result<()> { + let batch = Box::new(TestBatch); + let response = Response::builder() + .status(429) + .header("retry-after", "7") + .body(())?; + + assert_eq!( + Client::process_response(&*batch, response), + SendableState::Retry(Some(Duration::from_secs(7))) + ); + + Ok(()) + } + + #[test] + fn process_response_retry() -> Result<()> { + let mut codes = vec![402, 406, 407, 408]; + codes.append(&mut (100..200).collect()); + codes.append(&mut (300..400).collect()); + codes.append(&mut (430..600).collect()); + + for code in codes { + let batch = Box::new(TestBatch); + let response = Response::builder().status(code).body(())?; + + assert_eq!( + Client::process_response(&*batch, response), + SendableState::Retry(None), + "expected retry on {}", + code + ); + } + + Ok(()) + } + + #[test] + fn request() -> Result<()> { + let batch = Box::new(TestBatch); + let client = ClientBuilder::new("").build()?; + let endpoint = Endpoint { + host: "host".to_string(), + path: TRACE_API_PATH, + port: None, + }; + + let request = client.request(&*batch, &endpoint.uri(true)?)?; + + assert_eq!(request.uri().port(), None); + assert_eq!(request.uri().host(), Some("host")); + assert_eq!(request.uri().path(), "/trace/v1"); + assert_eq!(request.method(), Method::POST); + + let headers = request.headers(); + + let expected_headers = vec![ + (CONTENT_ENCODING.as_str(), "gzip"), + (CONTENT_TYPE.as_str(), "application/json"), + ("Api-Key", &client.api_key), + ("Data-Format", "newrelic"), + ("Data-Format-Version", "1"), + ("x-request-id", batch.uuid()), + (USER_AGENT.as_str(), &client.user_agent), + ]; + let expected_count = expected_headers.len(); + for (header, expected) in expected_headers { + let value = headers.get(header); + let expected = HeaderValue::from_str(expected)?; + assert_eq!(value, Some(&expected)); + } + + assert_eq!(headers.len(), expected_count, "Unexpected header present!"); + + Ok(()) + } + + #[test] + fn request_port() -> Result<()> { + let batch = Box::new(TestBatch); + let client = ClientBuilder::new("").build()?; + let endpoint = Endpoint { + host: "host".to_string(), + path: TRACE_API_PATH, + port: Some(80), + }; + + let request = client.request(&*batch, &endpoint.uri(true)?)?; + + assert_eq!(request.uri().port().unwrap().as_u16(), 80); + assert_eq!(request.uri().host(), Some("host")); + assert_eq!(request.uri().path(), "/trace/v1"); + + Ok(()) + } + + #[test] + fn builder_default() { + let b = ClientBuilder::new("0000"); + + assert_eq!(b.api_key, "0000"); + assert_eq!(b.backoff_factor, Duration::from_secs(5)); + assert_eq!(b.retries_max, 8); + assert_eq!(b.endpoint_traces.host, "trace-api.newrelic.com"); + assert_eq!(b.endpoint_traces.port, None); + assert_eq!(b.product_info, None); + assert_eq!(b.use_tls, true); + } + + #[test] + fn builder_setters() { + let b = ClientBuilder::new("0000") + .backoff_factor(Duration::from_secs(10)) + .retries_max(10) + .endpoint_traces("127.0.0.1", Some(8080)) + .product_info("Test", "1.0"); + + assert_eq!(b.api_key, "0000"); + assert_eq!(b.backoff_factor, Duration::from_secs(10)); + assert_eq!(b.retries_max, 10); + assert_eq!(b.endpoint_traces.host, "127.0.0.1"); + assert_eq!(b.endpoint_traces.port, Some(8080)); + assert_eq!( + b.product_info, + Some(("Test".to_string(), "1.0".to_string())) + ); + } + + #[test] + fn backoff_sequence_default() { + let seq = ClientBuilder::new("").get_backoff_sequence(); + + assert_eq!( + seq, + vec![0, 5, 10, 20, 40, 80, 160, 320] + .into_iter() + .map(|d| Duration::from_secs(d)) + .collect::>() + ); + } + + #[test] + fn backoff_sequence_no_retry() { + let seq = ClientBuilder::new("").retries_max(0).get_backoff_sequence(); + + assert_eq!(seq, vec![]); + } + + #[test] + fn backoff_sequence_custom() { + let seq = ClientBuilder::new("") + .backoff_factor(Duration::from_secs(2)) + .retries_max(6) + .get_backoff_sequence(); + + assert_eq!( + seq, + vec![0, 2, 4, 8, 16, 32] + .into_iter() + .map(|d| Duration::from_secs(d)) + .collect::>() + ); + } + + #[test] + fn user_agent_header_default() { + let header = ClientBuilder::new("").get_user_agent_header(); + + assert_eq!(header, format!("NewRelic-Rust-TelemetrySDK/{}", VERSION)); + } + + #[test] + fn user_agent_header_custom() { + let header = ClientBuilder::new("") + .product_info("Doc", "1.0") + .get_user_agent_header(); + + assert_eq!( + header, + format!("NewRelic-Rust-TelemetrySDK/{} Doc/1.0", VERSION) + ); + } +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/src/lib.rs b/crates/vendor/newrelic-telemetry-sdk-rust/src/lib.rs new file mode 100644 index 0000000000..0cc8270150 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/src/lib.rs @@ -0,0 +1,18 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// +pub mod attribute; + +pub mod span; +pub use span::{Span, SpanBatch}; + +#[cfg(feature = "client")] +mod client; +#[cfg(feature = "client")] +pub use client::{Client, ClientBuilder}; + +#[cfg(feature = "blocking")] +pub mod blocking { + pub use super::client::blocking::Client; +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/src/span.rs b/crates/vendor/newrelic-telemetry-sdk-rust/src/span.rs new file mode 100644 index 0000000000..29591866b9 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/src/span.rs @@ -0,0 +1,569 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// +use crate::attribute::Value; +use crate::client::Sendable; +use anyhow::Result; +use serde::{Serialize, Serializer}; +use std::collections::HashMap; +use std::fmt; +use std::time::Duration; +use uuid::Uuid; + +/// Represents a distributed tracing span. +#[derive(serde::Serialize, Clone, Debug, PartialEq)] +pub struct Span { + id: String, + + #[serde(rename = "trace.id")] + trace_id: String, + + timestamp: u64, + + #[serde(skip_serializing_if = "HashMap::is_empty")] + attributes: HashMap, +} + +impl Span { + /// Create a new span and assign an unique identifier, trace id and timestamp + pub fn new(id: &str, trace_id: &str, timestamp: u64) -> Span { + Span { + id: id.to_string(), + trace_id: trace_id.to_string(), + timestamp: timestamp, + attributes: HashMap::new(), + } + } + + /// Set a unique identifier for this span. This is a required field. + pub fn id(mut self, id: &str) -> Self { + self.id = id.to_string(); + self + } + + pub fn set_id(&mut self, id: &str) { + self.id = id.to_string(); + } + + /// Set a unique identifier shared by all spans within a single trace. + /// This is a required field. + pub fn trace_id(mut self, trace_id: &str) -> Self { + self.trace_id = trace_id.to_string(); + self + } + + pub fn set_trace_id(&mut self, trace_id: &str) { + self.trace_id = trace_id.to_string(); + } + + /// Set the start time of the span. This is a required field. + pub fn timestamp(mut self, timestamp: u64) -> Self { + self.timestamp = timestamp; + self + } + + pub fn set_timestamp(&mut self, timestamp: u64) { + self.timestamp = timestamp; + } + + /// Set the name of this span. + pub fn name(self, name: &str) -> Self { + self.attribute("name", name) + } + + pub fn set_name(&mut self, name: &str) { + self.set_attribute("name", name); + } + + /// Set the duration (in milliseconds) of this span. + pub fn duration(self, duration: Duration) -> Self { + self.attribute("duration.ms", duration.as_millis()) + } + + pub fn set_duration(&mut self, duration: Duration) { + self.set_attribute("duration.ms", duration.as_millis()); + } + + /// Set the id of the previous caller of this span. + pub fn parent_id(self, parent_id: &str) -> Self { + self.attribute("parent.id", parent_id) + } + + pub fn set_parent_id(&mut self, parent_id: &str) { + self.set_attribute("parent.id", parent_id); + } + + /// Set the name of the service that created this span. + pub fn service_name(self, service_name: &str) -> Self { + self.attribute("service.name", service_name) + } + + pub fn set_service_name(&mut self, service_name: &str) { + self.set_attribute("service.name", service_name); + } + + /// Set an attribute on the span. + pub fn attribute>(mut self, key: &str, value: T) -> Self { + self.attributes.insert(key.to_string(), value.into()); + self + } + + pub fn set_attribute>(&mut self, key: &str, value: T) { + self.attributes.insert(key.to_string(), value.into()); + } +} + +fn serialize_attributes(attrs: &HashMap, s: S) -> Result +where + S: Serializer, +{ + let mut wrapper: HashMap> = HashMap::new(); + wrapper.insert("attributes".to_string(), attrs); + wrapper.serialize(s) +} + +/// Encapsulates a collection of spans and the common data they share +#[derive(serde::Serialize, Debug, PartialEq)] +pub struct SpanBatch { + #[serde(skip_serializing)] + uuid: String, + + spans: Vec, + + #[serde(skip_serializing_if = "HashMap::is_empty")] + #[serde(serialize_with = "serialize_attributes")] + #[serde(rename = "common")] + attributes: HashMap, +} + +impl From> for SpanBatch { + /// Creates a new `SpanBatch` from a `Vec` + fn from(spans: Vec) -> Self { + let mut batch = Self::new(); + + for span in spans { + batch.record(span); + } + + batch + } +} + +impl SpanBatch { + /// Creates an empty `SpanBatch`. + pub fn new() -> Self { + SpanBatch { + uuid: Uuid::new_v4().to_string(), + spans: vec![], + attributes: HashMap::new(), + } + } + + /// Adds the provided span to the batch. + pub fn record(&mut self, span: Span) { + self.spans.push(span); + } + + /// Sets an attribute on the span batch. Returns `self` and can be chained + /// for concise addition of multiple attributes. + pub fn attribute>(mut self, key: &str, value: T) -> Self { + self.set_attribute(key, value); + self + } + + /// Sets an attribute on the span batch. + pub fn set_attribute>(&mut self, key: &str, value: T) { + self.attributes.insert(key.to_string(), value.into()); + } +} + +impl Sendable for SpanBatch { + fn uuid(&self) -> &str { + &self.uuid + } + + /// Returns the span batch encoded as a json string in the format expected + /// by the New Relic Telemetry API + fn marshall(&self) -> Result { + Ok(serde_json::to_string(&vec![self])?) + } + + /// Splits the batch in half. This is mostly used when the API service + /// returns a code indicating that the payload is too large. + fn split(&mut self) -> Box { + let new_batch_size: usize = self.spans.len() / 2; + self.uuid = Uuid::new_v4().to_string(); + + Box::new(SpanBatch { + uuid: Uuid::new_v4().to_string(), + spans: self.spans.drain(new_batch_size..).collect(), + attributes: self.attributes.clone(), + }) + } +} + +impl fmt::Display for SpanBatch { + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + write!( + f, + "", + self.spans.len(), + self.attributes.len(), + ) + } +} + +#[cfg(test)] +mod tests { + use super::{Sendable, Span, SpanBatch}; + use crate::attribute::Value; + use anyhow::Result; + use serde_json::json; + use std::time::Duration; + + macro_rules! assert_json_eq { + ($x: expr, $y: expr) => { + let (left, right) = ($x, $y); + assert!( + serde_json::from_str::(left)? + == serde_json::from_str::(right)?, + "expected {}, got {}", + left, + right + ); + }; + } + + #[test] + fn span_set_id() { + let mut span = Span::new("id1", "traceId1", 1); + assert_eq!(span.id, "id1"); + + span.set_id("id2"); + assert_eq!(span.id, "id2"); + + span = span.id("id3"); + assert_eq!(span.id, "id3"); + } + + #[test] + fn span_set_trace_id() { + let mut span = Span::new("id1", "traceId1", 1); + assert_eq!(span.trace_id, "traceId1"); + + span.set_trace_id("traceId2"); + assert_eq!(span.trace_id, "traceId2"); + + span = span.trace_id("traceId3"); + assert_eq!(span.trace_id, "traceId3"); + } + + #[test] + fn span_set_timestamp() { + let mut span = Span::new("id1", "traceId1", 1); + assert_eq!(span.timestamp, 1); + + span.set_timestamp(2); + assert_eq!(span.timestamp, 2); + + span = span.timestamp(3); + assert_eq!(span.timestamp, 3); + } + + #[test] + fn span_to_json() { + // Check span JSON serialization with empty attribute hashmap. + let mut span = Span::new("id1", "traceId1", 1); + let json_span = json!({"id": "id1", "trace.id": "traceId1", "timestamp": 1}); + + assert_eq!(json!(span), json_span); + + // Check span JSON serialization with populated attribute hashmap. + span.set_name("I have a name"); + let json_span_attribute = json!({"id": "id1", "trace.id": "traceId1", "timestamp": 1, "attributes": {"name": "I have a name"}}); + + assert_eq!(json!(span), json_span_attribute); + } + + #[test] + fn span_attribute() { + let mut span = Span::new("id", "traceId", 1); + + // Test name attribute + span.set_name("name"); + assert_eq!( + span.attributes.get("name"), + Some(&Value::Str(String::from("name"))) + ); + + span = span.name("name2"); + assert_eq!( + span.attributes.get("name"), + Some(&Value::Str(String::from("name2"))) + ); + + // Test duration attribute + span.set_duration(Duration::from_millis(10)); + assert_eq!( + span.attributes.get("duration.ms"), + Some(&Value::UInt128(10)) + ); + + span = span.duration(Duration::from_millis(20)); + assert_eq!( + span.attributes.get("duration.ms"), + Some(&Value::UInt128(20)) + ); + + // Test parent id attribute + span.set_parent_id("parent"); + assert_eq!( + span.attributes.get("parent.id"), + Some(&Value::Str(String::from("parent"))) + ); + + span = span.parent_id("parent2"); + assert_eq!( + span.attributes.get("parent.id"), + Some(&Value::Str(String::from("parent2"))) + ); + + // Test service name attribute + span.set_service_name("serviceName"); + assert_eq!( + span.attributes.get("service.name"), + Some(&Value::Str(String::from("serviceName"))) + ); + + span = span.service_name("serviceName2"); + assert_eq!( + span.attributes.get("service.name"), + Some(&Value::Str(String::from("serviceName2"))) + ); + } + + #[test] + fn span_attribute_type() { + let mut span = Span::new("id", "traceId", 1); + + // Test String attribute + span.set_attribute("attr.str", "str"); + assert_eq!( + span.attributes.get("attr.str"), + Some(&Value::Str(String::from("str"))) + ); + + span = span.attribute("attr.str", "str2"); + assert_eq!( + span.attributes.get("attr.str"), + Some(&Value::Str(String::from("str2"))) + ); + + // Test UInt attribute + let val_u32: u32 = 5; + span.set_attribute("attr.uint", val_u32); + assert_eq!( + span.attributes.get("attr.uint"), + Some(&Value::UInt(val_u32 as u64)) + ); + + let val_u64: u64 = 42; + span = span.attribute("attr.uint", val_u64); + assert_eq!( + span.attributes.get("attr.uint"), + Some(&Value::UInt(val_u64)) + ); + + // Test Int attribute + let val_i32: i32 = -5; + span.set_attribute("attr.int", val_i32); + assert_eq!( + span.attributes.get("attr.int"), + Some(&Value::Int(val_i32 as i64)) + ); + + let val_i64: i64 = -42; + span = span.attribute("attr.int", val_i64); + assert_eq!(span.attributes.get("attr.int"), Some(&Value::Int(val_i64))); + + // Test Float attribute + let val_f32: f32 = 3.14; + span.set_attribute("attr.float", val_f32); + assert_eq!( + span.attributes.get("attr.float"), + Some(&Value::Float(val_f32 as f64)) + ); + + let val_f64: f64 = 6.28; + span = span.attribute("attr.float", val_f64); + assert_eq!( + span.attributes.get("attr.float"), + Some(&Value::Float(val_f64)) + ); + + // Test Bool attribute + span.set_attribute("attr.bool", true); + assert_eq!(span.attributes.get("attr.bool"), Some(&Value::Bool(true))); + + span = span.attribute("attr.bool", false); + assert_eq!(span.attributes.get("attr.bool"), Some(&Value::Bool(false))); + } + + /// Helper function to generate a simple SpanBatch + fn span_vec(count: usize) -> Vec { + let mut vec = Vec::new(); + + for n in 0..count { + let id = format!("id{}", n); + let trace_id = format!("trace_id{}", n); + vec.push(Span::new(id.as_str(), trace_id.as_str(), 1)) + } + + vec + } + + #[test] + fn spanbatch_split_partial() { + // Note: since SpanBatch::split() returns a Box, + // we cannot fully test split with regard to the returned + // SpanBatch, only that the originally was drained as expected + // However, the integration tests cover both sides of this case. + let mut batch = SpanBatch::from(span_vec(2)); + let uuid = batch.uuid().to_string(); + let second_batch = batch.split(); + + let second_uuid = second_batch.uuid(); + assert_eq!(batch.spans.len(), 1); + assert_eq!(batch.spans[0], Span::new("id0", "trace_id0", 1)); + + // confirm the uuid for the second batch is not the same as the first + // and that the first remains unchanged + assert_ne!(uuid, second_uuid); + assert_ne!(uuid, batch.uuid()); + } + + #[test] + fn spanbatch_to_json() -> Result<()> { + // Check span JSON serialization with empty attribute hashmap. + let batch = SpanBatch::from(span_vec(2)).attribute("attr.test", 3); + + // json! macro imposes a sort which is different from the serde-derive + // specified order, therefore a string is used + let expected_string = r#"[{"spans":[ + {"id":"id0","trace.id":"trace_id0","timestamp":1}, + {"id":"id1","trace.id":"trace_id1","timestamp":1}], + "common":{"attributes":{"attr.test":3}}}]"#; + + let marshalled = batch.marshall().unwrap(); + assert_json_eq!(marshalled.as_str(), expected_string); + Ok(()) + } + + #[test] + fn spanbatch_attribute_type() { + let mut batch = SpanBatch::new(); + + // Test String attribute + batch.set_attribute("attr.str", "str"); + assert_eq!( + batch.attributes.get("attr.str"), + Some(&Value::Str(String::from("str"))) + ); + + batch = batch.attribute("attr.str", "str2"); + assert_eq!( + batch.attributes.get("attr.str"), + Some(&Value::Str(String::from("str2"))) + ); + + // Test UInt attribute + let val_u32: u32 = 5; + batch.set_attribute("attr.uint", val_u32); + assert_eq!( + batch.attributes.get("attr.uint"), + Some(&Value::UInt(val_u32 as u64)) + ); + + let val_u64: u64 = 42; + batch = batch.attribute("attr.uint", val_u64); + assert_eq!( + batch.attributes.get("attr.uint"), + Some(&Value::UInt(val_u64)) + ); + + // Test Int attribute + let val_i32: i32 = -5; + batch.set_attribute("attr.int", val_i32); + assert_eq!( + batch.attributes.get("attr.int"), + Some(&Value::Int(val_i32 as i64)) + ); + + let val_i64: i64 = -42; + batch = batch.attribute("attr.int", val_i64); + assert_eq!(batch.attributes.get("attr.int"), Some(&Value::Int(val_i64))); + + // Test Float attribute + let val_f32: f32 = 3.14; + batch.set_attribute("attr.float", val_f32); + assert_eq!( + batch.attributes.get("attr.float"), + Some(&Value::Float(val_f32 as f64)) + ); + + let val_f64: f64 = 6.28; + batch = batch.attribute("attr.float", val_f64); + assert_eq!( + batch.attributes.get("attr.float"), + Some(&Value::Float(val_f64)) + ); + + // Test Bool attribute + batch.set_attribute("attr.bool", true); + assert_eq!(batch.attributes.get("attr.bool"), Some(&Value::Bool(true))); + + batch = batch.attribute("attr.bool", false); + assert_eq!(batch.attributes.get("attr.bool"), Some(&Value::Bool(false))); + } + + #[test] + fn spanbatch_from() { + let vec = span_vec(23); + let batch = SpanBatch::from(vec.clone()); + assert_eq!(batch.spans, vec); + assert_eq!(batch.spans.len(), 23); + } + + #[test] + fn spanbatch_record() { + let mut batch = SpanBatch::new(); + let span = Span::new("id0", "trace_id0", 9); + batch.record(span.clone()); + assert_eq!(batch.spans.len(), 1); + assert_eq!(batch.spans[0], span); + } + + #[test] + fn spanbatch_format() { + let batch = SpanBatch::from(span_vec(23)) + .attribute("one", 1) + .attribute("two", "too") + .attribute("three", 3.0); + + let batch_string = format!("{}", batch); + assert_eq!(batch_string, ""); + } + + #[test] + fn spanbatch_attribute_chain() -> Result<()> { + let batch = SpanBatch::new() + .attribute("bad_dogs", 0) + .attribute("howdy", "y'all"); + let expected_string = + r#"[{"spans":[],"common":{"attributes":{"bad_dogs":0,"howdy":"y'all"}}}]"#; + let marshalled = batch.marshall().unwrap(); + assert_json_eq!(marshalled.as_str(), expected_string); + Ok(()) + } +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/tests/blocking.rs b/crates/vendor/newrelic-telemetry-sdk-rust/tests/blocking.rs new file mode 100644 index 0000000000..214851955a --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/tests/blocking.rs @@ -0,0 +1,252 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// +#[cfg(feature = "blocking")] +#[macro_use] +mod common; + +#[cfg(feature = "blocking")] +mod blocking { + use super::common; + use anyhow::Result; + use common::Endpoint; + use newrelic_telemetry::{blocking::Client, ClientBuilder, Span, SpanBatch}; + use std::thread; + use std::time::Duration; + + pub fn setup() -> Result<(Endpoint, Client)> { + let _ = env_logger::builder().is_test(true).try_init(); + + let endpoint = Endpoint::new(); + let client = ClientBuilder::new(&endpoint.license) + .endpoint_traces(&endpoint.host, Some(endpoint.port)) + .tls(false) + .build_blocking()?; + + Ok((endpoint, client)) + } + + #[test] + fn backoff() -> Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut endpoint = Endpoint::new(); + thread::sleep(Duration::from_secs(2)); + let client = ClientBuilder::new(&endpoint.license) + .retries_max(3) + .backoff_factor(Duration::from_secs(0)) + .endpoint_traces(&endpoint.host, Some(endpoint.port)) + .tls(false) + .build_blocking()?; + + let span_batch = SpanBatch::new(); + + client.send_spans(span_batch); + + // Four payloads should be sent: the initial one and 3 retries. + for num in 1..4 { + endpoint.reply(500)?; + assert!(endpoint.next_payload().is_ok(), "receiving payload {}", num); + } + + assert!(endpoint.next_payload().is_err(), "dropping payload"); + + Ok(()) + } + + #[test] + fn retry_after() -> Result<()> { + let _ = env_logger::builder() + .is_test(true) + .parse_filters("debug") + .try_init(); + + let mut endpoint = Endpoint::new(); + let client = ClientBuilder::new(&endpoint.license) + .retries_max(5) + .backoff_factor(Duration::from_secs(3600)) + .endpoint_traces(&endpoint.host, Some(endpoint.port)) + .tls(false) + .build_blocking()?; + + let span_batch = SpanBatch::new(); + + client.send_spans(span_batch); + + // Six payloads should be sent: the initial one and 5 retries. + // + // If the retry-after header is not read, this should hang for a very long time. + for num in 1..6 { + endpoint.reply_details( + 429, + vec![("Retry-After".to_string(), "0".to_string())], + "{}", + )?; + assert!(endpoint.next_payload().is_ok(), "receiving payload {}", num); + } + + assert!(endpoint.next_payload().is_err(), "dropping payload"); + + Ok(()) + } + + #[test] + fn headers() -> Result<()> { + let (mut endpoint, client) = setup()?; + + let span_batch = SpanBatch::new(); + + client.send_spans(span_batch); + endpoint.reply(202)?; + + let p = endpoint.next_payload()?; + + assert_eq!( + p.headers.get("content-type"), + Some(&"application/json".to_string()) + ); + assert_eq!(p.headers.get("content-encoding"), Some(&"gzip".to_string())); + + assert_eq!(p.headers.get("data-format"), Some(&"newrelic".to_string())); + assert_eq!(p.headers.get("data-format-version"), Some(&"1".to_string())); + + let user_agent = format!("NewRelic-Rust-TelemetrySDK/{}", env!("CARGO_PKG_VERSION")); + assert_eq!(p.headers.get("user-agent"), Some(&user_agent)); + + Ok(()) + } + + #[test] + fn product() -> Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut endpoint = Endpoint::new(); + let client = ClientBuilder::new(&endpoint.license) + .endpoint_traces(&endpoint.host, Some(endpoint.port)) + .product_info("SomeProduct", "3.14.9") + .tls(false) + .build_blocking()?; + + let span_batch = SpanBatch::new(); + + client.send_spans(span_batch); + endpoint.reply(202)?; + + let p = endpoint.next_payload()?; + + let user_agent = format!( + "NewRelic-Rust-TelemetrySDK/{} SomeProduct/3.14.9", + env!("CARGO_PKG_VERSION") + ); + assert_eq!(p.headers.get("user-agent"), Some(&user_agent)); + + Ok(()) + } + + #[test] + fn drop_payload() -> Result<()> { + for code in vec![400, 401, 403, 404, 405, 409, 410, 411] { + let (mut endpoint, client) = setup()?; + + let span_batch = SpanBatch::new(); + + client.send_spans(span_batch); + endpoint.reply(code)?; + + assert!(endpoint.next_payload().is_ok(), "first attempt to send"); + assert!(endpoint.next_payload().is_err(), "payload dropped"); + } + + Ok(()) + } + + #[test] + fn backpressure() -> Result<()> { + let _ = env_logger::builder().is_test(true).try_init(); + + let mut endpoint = Endpoint::new(); + let client = ClientBuilder::new(&endpoint.license) + .endpoint_traces(&endpoint.host, Some(endpoint.port)) + .tls(false) + .blocking_queue_max(1) + .build_blocking()?; + + for _ in 0..10 { + client.send_spans(SpanBatch::new()); + } + + endpoint.reply(202)?; + assert!(endpoint.next_payload().is_ok(), "first batch sent"); + + assert!( + endpoint.next_payload().is_err(), + "additional batches dropped" + ); + + Ok(()) + } + + #[test] + fn split_payload() -> Result<()> { + let (mut endpoint, client) = setup()?; + + let mut span_batch = SpanBatch::new(); + + let span_batch = vec![ + Span::new("id1", "tid1", 1000), + Span::new("id2", "tid2", 2000), + Span::new("id1", "tid1", 1000), + Span::new("id2", "tid2", 2000), + ] + .into(); + + client.send_spans(span_batch); + endpoint.reply(413)?; + endpoint.reply(202)?; + endpoint.reply(202)?; + + assert_json_eq!( + &endpoint.next_payload()?.body, + r#" + [{ + "spans": [ + { + "id": "id1", + "timestamp": 1000, + "trace.id": "tid1" + }, + { + "id": "id2", + "timestamp": 2000, + "trace.id": "tid2" + } + ] + }]"# + ); + + assert_json_eq!( + &endpoint.next_payload()?.body, + r#" + [{ + "spans": [ + { + "id": "id1", + "timestamp": 1000, + "trace.id": "tid1" + }, + { + "id": "id2", + "timestamp": 2000, + "trace.id": "tid2" + } + ] + }]"# + ); + + // Skip the first payload that is rejected. + endpoint.next_payload()?.body; + + Ok(()) + } +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/tests/common/mod.rs b/crates/vendor/newrelic-telemetry-sdk-rust/tests/common/mod.rs new file mode 100644 index 0000000000..cea9ac979a --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/tests/common/mod.rs @@ -0,0 +1,250 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// +use flate2::read::GzDecoder; +use hyper::service::{make_service_fn, service_fn}; +use hyper::{Body, Request, Server}; +use rand::seq::SliceRandom; +use rand::thread_rng; +use std::collections::HashMap; +use std::io::Read; +use std::net::{SocketAddr, TcpListener}; +use std::str; +use std::sync::{Arc, Mutex}; +use std::thread; +use std::time::SystemTime; +use tokio::runtime::Builder; + +macro_rules! assert_json_eq { + ($x: expr, $y: expr) => { + let (left, right) = ($x, $y); + assert!( + serde_json::from_str::(left)? + == serde_json::from_str::(right)?, + "expected {}, got {}", + left, + right + ); + }; +} + +#[derive(Debug)] +pub struct Error { + message: String, +} + +impl Error { + fn new(message: &str) -> Error { + Error { + message: message.to_string(), + } + } +} + +impl std::error::Error for Error {} + +impl std::fmt::Display for Error { + fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { + write!(f, "{}", self.message) + } +} + +impl From for Error { + fn from(error: std::str::Utf8Error) -> Self { + Error::new(&error.to_string()) + } +} + +impl From for Error { + fn from(error: hyper::error::Error) -> Self { + Error::new(&error.to_string()) + } +} + +pub struct Payload { + pub headers: HashMap, + pub body: String, +} + +pub struct Response { + pub code: u16, + pub headers: Vec<(String, String)>, + pub body: String, +} + +pub struct Endpoint { + pub license: String, + pub host: String, + pub port: u16, + pub timeout_ms: u128, + server: Option>, + chan_payloads: Arc>>, + chan_responses: Arc>>, + chan_stop: Option>, +} + +impl Endpoint { + pub fn new() -> Self { + let payloads = Arc::new(Mutex::new(Vec::::new())); + let responses = Arc::new(Mutex::new(Vec::::new())); + + // Three rounds of clones are required to get the Arcs through all the + // closure bounds. + let (p, r) = (payloads.clone(), responses.clone()); + let new_service = make_service_fn(move |_conn| { + let (p, r) = (p.clone(), r.clone()); + async { + Ok::<_, hyper::Error>(service_fn(move |req| { + let (p, r) = (p.clone(), r.clone()); + Endpoint::accept_payload(req, p, r) + })) + } + }); + + let mut runtime = Builder::new() + .threaded_scheduler() + .enable_all() + .build() + .unwrap(); + + let (sender, receiver) = futures::channel::oneshot::channel::<()>(); + + let wrapped_receiver = async { + let _ = receiver.await; + }; + + let port = Endpoint::get_available_port().expect("Cannot aquire port"); + + let handle = thread::spawn(move || { + let addr: SocketAddr = ([127, 0, 0, 1], port).into(); + + let scope = async { + Server::bind(&addr) + .http1_half_close(true) + .serve(new_service) + .with_graceful_shutdown(wrapped_receiver) + .await + .unwrap(); + }; + + runtime.block_on(scope); + }); + + Endpoint { + license: "license".to_string(), + host: "127.0.0.1".to_string(), + port: port, + timeout_ms: 5000, + server: Some(handle), + chan_payloads: payloads, + chan_responses: responses, + chan_stop: Some(sender), + } + } + + pub fn reply(&self, code: u16) -> Result<(), Error> { + self.reply_details(code, vec![], "{}") + } + + pub fn reply_details( + &self, + code: u16, + headers: Vec<(String, String)>, + body: &str, + ) -> Result<(), Error> { + let mut lock = self.chan_responses.lock().unwrap(); + lock.push(Response { + code: code, + headers: headers, + body: body.to_string(), + }); + let len = lock.len(); + drop(lock); + + // Block until the response was consumed. + let start = SystemTime::now(); + loop { + let lock = self.chan_responses.lock().unwrap(); + if lock.len() < len { + return Ok(()); + } + + let duration = SystemTime::now().duration_since(start).unwrap(); + if duration.as_millis() > self.timeout_ms { + return Err(Error::new("Timeout: no request received")); + } + } + } + + pub fn next_payload(&mut self) -> Result { + let mut lock = self.chan_payloads.lock().unwrap(); + + match lock.pop() { + Some(p) => Ok(p), + None => Err(Error::new("No payload received")), + } + } + + async fn accept_payload( + req: Request, + payloads: Arc>>, + responses: Arc>>, + ) -> Result, Error> { + let headers: HashMap = req + .headers() + .iter() + .map(|(key, value)| (key.to_string(), value.to_str().unwrap_or("").to_string())) + .collect(); + + let body = hyper::body::to_bytes(req.into_body()).await?; + let body_bytes: Result, _> = body.bytes().collect(); + let body_bytes = body_bytes.unwrap(); + + let mut decoder = GzDecoder::new(&body_bytes[..]); + let mut body_decoded = String::new(); + let _ = decoder.read_to_string(&mut body_decoded); + + let mut lock = payloads.lock().unwrap(); + lock.push(Payload { + headers: headers, + body: body_decoded, + }); + + let mut lock = responses.lock().unwrap(); + match lock.pop() { + Some(r) => { + let mut resp = hyper::Response::builder(); + + for (k, v) in &r.headers { + resp = resp.header(k, v); + } + + Ok(resp.status(r.code).body(Body::from(r.body)).unwrap()) + } + _ => Err(Error::new("No response given")), + } + } + + fn get_available_port() -> Option { + fn port_is_available(p: u16) -> bool { + match TcpListener::bind(("127.0.0.1", p)) { + Ok(_) => true, + _ => false, + } + } + + let mut ports: Vec = (3000..5000).collect(); + ports.shuffle(&mut thread_rng()); + + ports.into_iter().find(|p| port_is_available(*p)) + } +} + +impl Drop for Endpoint { + fn drop(&mut self) { + if let Ok(_) = self.chan_stop.take().unwrap().send(()) { + let _ = self.server.take().unwrap().join(); + } + } +} diff --git a/crates/vendor/newrelic-telemetry-sdk-rust/tests/spans.rs b/crates/vendor/newrelic-telemetry-sdk-rust/tests/spans.rs new file mode 100644 index 0000000000..7bdd98e354 --- /dev/null +++ b/crates/vendor/newrelic-telemetry-sdk-rust/tests/spans.rs @@ -0,0 +1,219 @@ +/// +/// Copyright 2020 New Relic Corporation. All rights reserved. +/// SPDX-License-Identifier: Apache-2.0 +/// +#[cfg(feature = "client")] +#[macro_use] +mod common; + +#[cfg(feature = "client")] +mod client { + use super::common; + use anyhow::Result; + use common::Endpoint; + use newrelic_telemetry::{Client, ClientBuilder, Span, SpanBatch}; + use std::thread; + use std::time::Duration; + + pub fn setup() -> Result<(Endpoint, Client)> { + let _ = env_logger::builder().is_test(true).try_init(); + + let endpoint = Endpoint::new(); + let client = ClientBuilder::new(&endpoint.license) + .endpoint_traces(&endpoint.host, Some(endpoint.port)) + .tls(false) + .build()?; + + Ok((endpoint, client)) + } + + #[tokio::test(threaded_scheduler)] + async fn empty() -> Result<()> { + let (mut endpoint, client) = setup()?; + + // Assertions are all handled in a separate thread, so we can await the + // future in the main thread. + let handle = thread::spawn(move || -> Result<()> { + endpoint.reply(202)?; + + assert_json_eq!( + &endpoint.next_payload().unwrap().body, + r#"[{ "spans" : [] }]"# + ); + + Ok(()) + }); + + client.send_spans(SpanBatch::new()).await; + + handle.join().expect("error from endpoint thread")?; + + Ok(()) + } + + #[tokio::test(threaded_scheduler)] + async fn simple() -> Result<()> { + let (mut endpoint, client) = setup()?; + + let handle = thread::spawn(move || -> Result<()> { + endpoint.reply(202)?; + + assert_json_eq!( + &endpoint.next_payload()?.body, + r#" + [ + { + "spans": [ + { + "id": "id1", + "timestamp": 1000, + "trace.id": "tid1" + } + ] + } + ]"# + ); + + Ok(()) + }); + + let mut span_batch = SpanBatch::new(); + span_batch.record(Span::new("id1", "tid1", 1000)); + + client.send_spans(span_batch).await; + handle.join().expect("error from endpoint thread")?; + + Ok(()) + } + + #[tokio::test(threaded_scheduler)] + async fn all_api_attrs() -> Result<()> { + let (mut endpoint, client) = setup()?; + + let handle = thread::spawn(move || -> Result<()> { + endpoint.reply(202)?; + + assert_json_eq!( + &endpoint.next_payload()?.body, + r#" + [{ + "spans": [{ + "id": "id1", + "timestamp": 1000, + "trace.id": "tid1", + "attributes": { + "name": "name1", + "duration.ms": 2000, + "parent.id": "pid1", + "service.name": "service1" + } + }] + }]"# + ); + + Ok(()) + }); + + let mut span_batch = SpanBatch::new(); + span_batch.record( + Span::new("id1", "tid1", 1000) + .name("name1") + .duration(Duration::from_millis(2000)) + .parent_id("pid1") + .service_name("service1"), + ); + + client.send_spans(span_batch).await; + handle.join().expect("error from endpoint thread")?; + + Ok(()) + } + + #[tokio::test(threaded_scheduler)] + async fn custom_attrs() -> Result<()> { + let (mut endpoint, client) = setup()?; + + let handle = thread::spawn(move || -> Result<()> { + endpoint.reply(202)?; + + assert_json_eq!( + &endpoint.next_payload()?.body, + r#" + [{ + "spans": [{ + "id": "id1", + "timestamp": 1000, + "trace.id": "tid1", + "attributes": { + "bool_attr": true, + "float_attr": 3.14159, + "str_attr": "string", + "int_attr": 40, + "neg_int_attr": -40, + "string_attr": "Str" + } + }] + }]"# + ); + + Ok(()) + }); + + let mut span_batch = SpanBatch::new(); + span_batch.record( + Span::new("id1", "tid1", 1000) + .attribute("bool_attr", true) + .attribute("float_attr", 3.14159) + .attribute("str_attr", "string") + .attribute("int_attr", 40) + .attribute("neg_int_attr", -40) + .attribute("string_attr", "Str"), + ); + + client.send_spans(span_batch).await; + handle.join().expect("error from endpoint thread")?; + + Ok(()) + } + + #[tokio::test(threaded_scheduler)] + async fn two_spans() -> Result<()> { + let (mut endpoint, client) = setup()?; + + let handle = thread::spawn(move || -> Result<()> { + endpoint.reply(202)?; + + assert_json_eq!( + &endpoint.next_payload()?.body, + r#" + [{ + "spans": [ + { + "id": "id1", + "timestamp": 1000, + "trace.id": "tid1" + }, + { + "id": "id2", + "timestamp": 2000, + "trace.id": "tid2" + } + ] + }]"# + ); + + Ok(()) + }); + + let span_batch = vec![ + Span::new("id1", "tid1", 1000), + Span::new("id2", "tid2", 2000), + ] + .into(); + + client.send_spans(span_batch).await; + handle.join().expect("error from endpoint thread")?; + + Ok(()) + } +} diff --git a/db/banned_slurs.txt b/db/banned_slurs.txt new file mode 100644 index 0000000000..b85aab3faa --- /dev/null +++ b/db/banned_slurs.txt @@ -0,0 +1,113 @@ +# This file serves to remove racial slurs and epithets from usage with vocodes. +# There's probably more than this that we'll want to remove. +# I need to investigate. +# Additional slurs from BibbyBoi 2022-03-04 +arabush +aravush +argie +argies +bakra +balija +bamboula +beaner +beaney +bluegum +bogtrotter +boonga +boonie +bootlip +bougnoule +bozgor +buddhahead +bule +bunga +burrhead +carcamano +chankoro +chink +chinks +choc-ice +coon +coonass +dago +dego +dink +dune-coon +dyke +fag +faggot +fags +greaser +groid +guidette +guido +haole +hebe +heeb +honkey +honkie +hymie +injun +jigaboo +jigarooni +jiggabo +jigger +jijjiboo +kike +kniggers +kyke +maumau +mayomonkey +moolinyan +mulignan +mulignon +neger +negro +negro +negroid +nicc a +nicc er +nicca +niccer +nig +nigar +nigette +niggah +niggar +nigger +niggers +niggertoe +niggress +nigguh +nignog +nigs +nik-a +nika +nikk-a +nikk-er +ofay +poortagee +rastus +sambo +sheboon +sheeny +sooty +spic +spick +spig +spig +spigotty +spik +tarbaby +whigg-a +whigg-er +whigga +whigger +wigg-er +wigger +wikk-a +wikk-er +wikka +wikker +zigabo +zipperhead \ No newline at end of file diff --git a/db/reserved_usernames.txt b/db/reserved_usernames.txt new file mode 100644 index 0000000000..10f6d1ff8e --- /dev/null +++ b/db/reserved_usernames.txt @@ -0,0 +1,1119 @@ +# This is a list of reserved usernames to allow for vanity URLs, etc. +# Taken from the following: +# https://gist.github.com/caseyohara/1453705/bd38325ae2b562572afaa119efa59529c057bc85 +# https://github.com/shouldbee/reserved-usernames +# https://github.com/Mottie/github-reserved-names +# https://github.com/mvila/reserved-usernames + +# Stuff we care about +fake_you +fakeyou +hanashi +monogatari +storied +stories +story +storyteller +tell +teller +vocode +vocoder +vocoders +vocodes +voxelcam + +# Types of models +first_order_model +fom +hifigan +melgan +tacotron +tt2 +tts +w2l + +# Companies we'd love to have use our service +# so we'll reserve them to be safe +amazon +apple +cash +cashapp +epic +epicgames +facebook +github +google +microsoft +netflix +nintendo +paypal +reddit +snap +snapchat +softbank +sony +square +stripe +tiktok +twitch +twitter +unreal +venmo +youtube + +# Top 50 languages by speaking population +chinese +mandarin +spanish +english +bengali +hindi +portuguese +russian +japanese +german +wu +javanese +korean +french +vietnamese +telugu +chinese +marathi +tamil +turkish +urdu +min-nan +jinyu +gujarati +polish +arabic +ukrainian +italian +xiang +malayalam +hakka +kannada +oriya +panjabi +sunda +panjabi +romanian +bhojpuri +azerbaijani +farsi +maithili +hausa +arabic +burmese +serbo-croatian +gan +awadhi +thai +dutch +yoruba +sindhi + +# Generic reserved words +0 +100 +101 +102 +1xx +200 +201 +202 +203 +204 +205 +206 +207 +226 +2xx +300 +301 +302 +303 +304 +305 +307 +308 +3xx +400 +401 +402 +403 +404 +405 +406 +407 +408 +409 +410 +411 +412 +413 +414 +415 +416 +417 +418 +419 +420 +421 +422 +423 +424 +425 +426 +427 +428 +429 +430 +431 +451 +4xx +500 +501 +502 +503 +504 +505 +506 +507 +508 +509 +510 +511 +5xx +7xx +a +about +about_us +aboutus +abuse +access +account +accounts +activate +activities +activity +ad +add +address +adm +admin +administration +administrator +ads +adult +advertising +advisories +affiliate +affiliates +ajax +all +alpha +analysis +analytics +android +anon +anonymous +any +api +app +apps +archive +archives +arpabet +article +asct +asset +assets +atom +attributes +audio +auth +authentication +autoconfig +avatar +avatars +b +backup +balancer-manager +ban +bank +banner +banners +beta +billing +bin +blob +blog +blogs +board +book +bookmark +bookmarks +borrow +bot +bots +bounty +branches +broadcasthost +browse +browser +bug +bugs +business +businesses +c +cache +cadastro +calendar +call +campaign +cancel +captcha +career +careers +cart +case-studies +categories +category +central +certification +cgi +cgi-bin +changelog +chat +check +checking +checkout +checks +cla +client +cliente +clients +cloud +code +codereview +collection +collections +comercial +comment +comments +commit +commits +communities +community +companies +company +compare +compras +config +configuration +connect +contact +contact-us +contact_us +contactus +contest +contribute +contributing +conversations +cookbook +corp +coupons +create +crypt +css +customer +customer-stories +customers +d +dashboard +dashboards +data +datas +datum +db +default +delete +delist +demo +design +designer +destroy +dev +devel +develop +developer +developers +diagram +diary +dict +dictionary +die +diff +dir +direct_messages +directory +discover +discussions +display +display_name +displayname +dist +dms +dns +doc +docker +docs +document +documentation +documents +domain +download +downloads +downtime +e +ecommerce +edit +editor +editors +edu +education +email +employment +empty +end +enterprise +entries +entry +error +errors +eval +event +events +exit +explore +explorer +export +f +facebook +false +faq +favorite +favorites +fbi +feature +featured +features +feed +feedback +feeds +file +files +fire_hose +firehose +firewall +first +fixtures +flash +fleet +fleets +flog +follow +followed +followers +following +forgot +forgot-password +forgot_password +forgotpassword +forked +form +forum +forums +founder +free +friend +friends +ftp +ftps +g +gadget +gadgets +game +games +garage +gdpr +get +ghost +gift +gifts +gist +gists +git +github +graph +graphs +group +groups +guest +guests +guid +guide +guides +h +help +help-wanted +home +homepage +hooks +host +hosting +hostmaster +hostname +hovercards +howto +hpg +href +html +http +httpd +https +i +iamges +icon +icons +id +idea +ideas +identity +image +images +imap +img +inbox +included +index +indice +individual +info +information +ingest +ingestion +inquiry +instagram +integration +integration-test +integration_test +integrationtest +interfaces +intranet +introduction +invalid-email-address +investor +investors +invitations +invite +ip +ipad +iphone +irc +is +isatap +issue +issues +it +item +items +j +java +javascript +job +jobs +join +joinus +journal +journals +js +json +jump +k +key +keys +keyserver +knowledgebase +l +lab +labels +labs +language +languages +last +launch +layouts +ldap-status +learn +learn_more +learnmore +legal +library +license +like +liked +link +links +linux +list +listen +listings +lists +live +livestream +local +localdomain +localhost +log +log-in +log-out +log_in +log_off +log_out +login +logoff +logos +logout +logs +m +mac +mail +mail1 +mail2 +mail3 +mail4 +mail5 +mailer +mailer-daemon +mailing +maintenance +malware +man +manager +manual +map +maps +market +marketing +marketplace +master +match +matches +matchmaking +me +media +member +members +mention +mentioned +mentioning +mentions +message +messages +messenger +microblog +microblogs +migrating +milestones +mine +mirrors +mis +mob +mobile +movie +movies +mp3 +msg +msn +music +musicas +mx +my +mysql +n +name +named +names +namespace +namespaces +nan +navi +navigation +net +network +new +news +newsletter +nick +nickname +no-reply +nobody +noc +nodes +none +nonprofit +nonprofits +noreply +notes +notices +noticias +notification +notifications +notify +ns +ns1 +ns10 +ns2 +ns3 +ns4 +ns5 +ns6 +ns7 +ns8 +ns9 +null +o +oauth +oauth_clients +offer +offers +official +old +online +open-source +openid +operator +ops +order +orders +organisations +organization +organizations +orgs +outbox +overview +owner +owners +p +package +page +pager +pages +panel +partners +passwd +password +patch +payment +payments +pending +perl +personal +pfp +pfps +phone +photo +photoalbum +photos +php +phpmyadmin +phppgadmin +phpredisadmin +pic +pics +picture +ping +plain +plaintext +plan +plans +playlist +playlists +plugin +plugins +pms +policies +policy +pop +pop3 +popular +popularity +portal +post +postfix +postmaster +posts +pr +premium +press +price +pricing +privacy +privacy-policy +privacy_policy +privacypolicy +private +product +products +professional +profile +project +projects +promo +promote +pub +public +pulls +purpose +put +pvp +pw +python +q +query +r +random +ranking +raw +read +readme +receive +recent +recommendations +recover +recover_email +recovery +recruit +recruitment +recycle +reddit +reddits +redeem +redo +ref +reference +register +registration +registry +release +releases +remote +remove +render +replies +reply +report +reports +repositories +repository +req +request +requests +res +resend +reset +reset-password +reset_email +reset_password +resetpassword +resource +resources +restore +revert +roc +root +rss +ruby +rule +rules +rust +s +sag +sale +sales +sample +samples +save +save-net-neutrality +saved +saving +school +scraping +script +scripts +search +seats +secure +security +self +send +server +server-info +server-status +service +services +session +sessions +setting +settings +setup +sftp +share +shared +shareholders +sheet +shop +show +showcases +sign-in +sign-up +sign_in +sign_off +sign_out +sign_up +signin +signoff +signout +signup +site +sitemap +sites +slug +smartphone +smtp +soporte +sort +sorted +source +spam +spec +special +sponsors +sql +src +ssh +ssl +ssladmin +ssladministrator +sslwebmaster +staff +stage +staging +starred +stars +start +stat +state +static +stats +status +statuses +storage +store +stores +stories +story +stream +streaming +style +styleguide +styles +stylesheet +stylesheets +subdomain +subreddit +subreddits +subscribe +subscriptions +suggest +suggestion +suggestions +suporte +support +suspended +svn +swf +sys +sysadmin +sysadministrator +system +t +table +tablet +tablets +tag +tags +talk +talks +task +tasks +teach +teacher +teachers +teaching +team +teams +tech +telnet +template +ten +term +terms +terms-of-service +terms_of_service +termsofservice +test +test1 +test2 +test3 +test4 +test5 +test6 +test7 +test8 +test9 +teste +testing +tests +theme +themes +thread +threads +timeline +title +tls +tmp +todo +token +tokenserver +tool +tools +top +topic +topics +tos +tour +train +training +translations +trash +tree +trending +trends +true +tutorial +tux +tv +tweet +tweets +twitter +type +typescript +u +undef +unfollow +unsubscribe +update +updates +upload +uploads +uptime +url +urls +usage +usenet +user +username +users +usr +usuario +util +uucp +uuid +v +value +vanity +vanity_url +vanityurl +vendas +ver +version +versus +video +videos +view +viewer +visitor +visualization +vpn +w +watch +watches +watching +weather +web +webhook +webhooks +webmail +webmaster +website +websites +welcome +who +widget +widgets +wiki +win +windows +word +work +works +works-with +workshop +wpad +ww +wws +www +www0 +www1 +www2 +www3 +www4 +www5 +www6 +www7 +www8 +www9 +wwws +wwww +x +xfn +xml +xmpp +xpg +xxx +y +yaml +year +yml +you +your +yourdomain +yourname +yours +yoursite +yourusername +z diff --git a/db/reserved_usernames_including.txt b/db/reserved_usernames_including.txt new file mode 100644 index 0000000000..cadc30d175 --- /dev/null +++ b/db/reserved_usernames_including.txt @@ -0,0 +1,10 @@ +# This is a more restrictive list of reserved usernames. +# Any of the substrings may not be used. +# You should check string.remove('_') as well. +fakeyou +storytell +storyteller +test +vocode +vocodes +voxelcam \ No newline at end of file diff --git a/debug/curl/api_token_tests.py b/debug/curl/api_token_tests.py new file mode 100755 index 0000000000..18be92ed8a --- /dev/null +++ b/debug/curl/api_token_tests.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 + +import json +import re +import requests +import uuid + +login_url = 'http://localhost:12345/login' +create_api_token_url = 'http://localhost:12345/api_tokens/create' +edit_api_token_url = 'http://localhost:12345/api_tokens/edit' +delete_api_token_url = 'http://localhost:12345/api_tokens/{}/delete' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + +cookies = { 'session': session_cookie } + +payload = { + 'idempotency_token': str(uuid.uuid4()), +} + +r = requests.post(create_api_token_url, cookies=cookies, json=payload) + +print("===== Created API Token =====") +print('Status: {}'.format(r.status_code)) +print(r.content) + +api_token = json.loads(r.content) + + +payload = { + 'api_token': api_token['api_token'], + 'maybe_short_description': 'this describes my token', +} +r = requests.post(edit_api_token_url, cookies=cookies, json=payload) + +print("===== Edited API Token =====") +print('Status: {}'.format(r.status_code)) +print(r.content) + +url = delete_api_token_url.format(api_token['api_token']) +r = requests.post(url, cookies=cookies) + +print("===== Deleted API Token =====") +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/category_tests.py b/debug/curl/category_tests.py new file mode 100755 index 0000000000..ad315ce1a6 --- /dev/null +++ b/debug/curl/category_tests.py @@ -0,0 +1,150 @@ +#!/usr/bin/env python3 + +import json +import re +import requests +import uuid + +login_url = 'http://localhost:12345/login' +category_create_url = 'http://localhost:12345/category/create' +category_delete_url_format = 'http://localhost:12345/category/{}/delete' +category_edit_url_format = 'http://localhost:12345/category/{}/edit' +tts_model_list_url = 'http://localhost:12345/tts/list' +tts_category_assignment_url = 'http://localhost:12345/category/assign/tts' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== [1] Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== [2] Create Category =====") + +cookies = { 'session': session_cookie } + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'name': 'First Category', + 'model_type': 'tts', # tts or w2l + 'can_directly_have_models': False, + 'can_have_subcategories': True, + 'is_mod_approved': True, +} + +r = requests.post(category_create_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) + +d = json.loads(r.content) +category_token = d['token'] + +print(f'created category token: #{category_token}') + + +print("===== [3] (Soft) Delete Category =====") + +category_delete_url = category_delete_url_format.format(category_token) + +cookies = { 'session': session_cookie } + +payload = { + 'set_delete': True, +} + +r = requests.post(category_delete_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + + +print("===== [4] Undelete Category =====") + +cookies = { 'session': session_cookie } + +payload = { + 'set_delete': False, +} + +r = requests.post(category_delete_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + + +print("===== [5] Edit Category =====") + +category_edit_url = category_edit_url_format.format(category_token) + +cookies = { 'session': session_cookie } + +payload = { + 'name': 'Edited Name', + 'dropdown_name': 'Edited Dropdown Name', + 'can_directly_have_models': True, + 'can_have_subcategories': True, + 'can_only_mods_apply': False, + 'is_mod_approved': True, + 'maybe_mod_comments': 'mod comment', +} + +r = requests.post(category_edit_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + + +print("===== [6] List TTS Models =====") + +r = requests.get(tts_model_list_url, cookies=cookies) + +print('Status: {}'.format(r.status_code)) + +d = json.loads(r.content) + +tts_model_token = d['models'][0]['model_token'] + +print('TTS model token: {}'.format(tts_model_token)) + + +print("===== [7] Create Category Assignment to TTS Model =====") + +payload = { + 'category_token': category_token, + 'tts_model_token': tts_model_token, + 'assign': True, +} + +r = requests.post(tts_category_assignment_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + + +print("===== [8] Delete Category Assignment to TTS Model =====") + +payload = { + 'category_token': category_token, + 'tts_model_token': tts_model_token, + 'assign': False, +} + +r = requests.post(tts_category_assignment_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/create_account_1.sh b/debug/curl/create_account_1.sh new file mode 100755 index 0000000000..85a4d0cb09 --- /dev/null +++ b/debug/curl/create_account_1.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username": "echelon", "password": "testing", "password_confirmation": "testing", "email_address": "echelon@gmail.com"}' \ + http://localhost:12345/create_account diff --git a/debug/curl/create_account_2.sh b/debug/curl/create_account_2.sh new file mode 100755 index 0000000000..457b415272 --- /dev/null +++ b/debug/curl/create_account_2.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username": "ernest", "password": "knowwhatimean", "password_confirmation": "knowwhatimean", "email_address": "vern@knowwhatimean.com"}' \ + http://localhost:12345/create_account diff --git a/debug/curl/create_account_reserved.sh b/debug/curl/create_account_reserved.sh new file mode 100755 index 0000000000..020b8534a7 --- /dev/null +++ b/debug/curl/create_account_reserved.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username": "vocodes", "password": "testing", "password_confirmation": "testing", "email_address": "vocodes@gmail.com"}' \ + http://localhost:12345/create_account diff --git a/debug/curl/infer_tts.py b/debug/curl/infer_tts.py new file mode 100755 index 0000000000..84d660ba23 --- /dev/null +++ b/debug/curl/infer_tts.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 + +import requests +import re + +login_url = 'http://localhost:12345/login' +inference_url = 'http://localhost:12345/tts/inference' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Infer TTS =====") + +cookies = { 'session': session_cookie } + +payload = { + 'tts_model_token': 'TTS_TOKEN', + 'inference_text': 'this is the text to generate', +} + +r = requests.post(inference_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/infer_w2l.py b/debug/curl/infer_w2l.py new file mode 100755 index 0000000000..3fcf417e1d --- /dev/null +++ b/debug/curl/infer_w2l.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 + +import requests +import re + +login_url = 'http://localhost:12345/login' +inference_url = 'http://localhost:12345/w2l/inference' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Infer W2L =====") + +cookies = { 'session': session_cookie } + +payload = { + 'w2l_template_token': 'W2L_TOKEN', + 'tts_inference_result_token': 'TTS_TOKEN', +} + +r = requests.post(inference_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/login_bad_password.sh b/debug/curl/login_bad_password.sh new file mode 100755 index 0000000000..8da6754518 --- /dev/null +++ b/debug/curl/login_bad_password.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username_or_email": "echelon", "password": "BAD_PASSWORD"}' \ + http://localhost:12345/login diff --git a/debug/curl/login_bad_username.sh b/debug/curl/login_bad_username.sh new file mode 100755 index 0000000000..85b70e89f6 --- /dev/null +++ b/debug/curl/login_bad_username.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username_or_email": "nobody", "password": "testing"}' \ + http://localhost:12345/login diff --git a/debug/curl/login_email.sh b/debug/curl/login_email.sh new file mode 100755 index 0000000000..977f07023a --- /dev/null +++ b/debug/curl/login_email.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username_or_email": "echelon@gmail.com", "password": "testing"}' \ + http://localhost:12345/login diff --git a/debug/curl/login_username.sh b/debug/curl/login_username.sh new file mode 100755 index 0000000000..95b72e8857 --- /dev/null +++ b/debug/curl/login_username.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# LOCALHOST ONLY + +curl -vvv \ + -H "content-type: application/json" \ + --data '{"username_or_email": "echelon", "password": "testing"}' \ + http://localhost:12345/login diff --git a/debug/curl/logout_roundtrip.py b/debug/curl/logout_roundtrip.py new file mode 100755 index 0000000000..348445541f --- /dev/null +++ b/debug/curl/logout_roundtrip.py @@ -0,0 +1,38 @@ +#!/usr/bin/env python3 + +import requests +import re + +login_url = 'http://localhost:12345/login' +logout_url = 'http://localhost:12345/logout' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + +cookies = { 'session': session_cookie } + +r = requests.post(logout_url, cookies=cookies) + +print("===== Logout =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + diff --git a/debug/curl/seed.py b/debug/curl/seed.py new file mode 100755 index 0000000000..218e71e563 --- /dev/null +++ b/debug/curl/seed.py @@ -0,0 +1,123 @@ +#!/usr/bin/env python3 + +import requests +import re +import uuid +import socket + +# There are some things here that will cause my laptop to freeze! +hostname = socket.gethostname() + +def print_response(response): + print('\tStatus: {}'.format(response.status_code)) + for k, v in response.headers.items(): + print('\t header - {}: {}'.format(k, v)) + print('\tbody: {}', response.json()) + +# ==================== CREATE ACCOUNTS ==================== + +def create_user_account(username, password, email): + print('Create account: {}'.format(username)) + signup_url = 'http://localhost:12345/create_account' + payload = { + 'username': username, + 'password': password, + 'password_confirmation': password, + 'email_address': email, + } + r = requests.post(signup_url, json=payload) + print_response(r) + +create_user_account('echelon', 'testing', 'echelon@gmail.com') +create_user_account('brandon', 'testing', 'brandon@storyteller.company') +create_user_account('nanashi', 'nanashi', 'somebody@somewhere.com') +create_user_account('other', 'testing', 'ryu@storyteller.company') +create_user_account('raptor9', 'testing', 'raptor9@storyteller.company') +create_user_account('ryu', 'testing', 'ryu@storyteller.company') + +# ==================== LOGIN ==================== + +def login(username, password): + print('Login: {}'.format(username)) + login_url = 'http://localhost:12345/login' + payload = { + 'username_or_email': username, + 'password': password, + } + r = requests.post(login_url, json=payload) + print_response(r) + raw_set_cookie_header = r.headers['set-cookie'] + m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) + session_cookie = m.group(1) + return session_cookie + +session_echelon = login('echelon', 'testing') +session_nanashi = login('nanashi', 'nanashi') +session_brandon = login('brandon', 'testing') +session_other = login('other', 'testing') + +# ==================== UPLOAD W2L TEMPLATES ==================== + +def upload_w2l_template(content_url, title, session_cookie): + print('Upload W2L template') + upload_url = 'http://localhost:12345/w2l/upload' + cookies = { 'session': session_cookie } + payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': title, + 'download_url': content_url, + } + r = requests.post(upload_url, cookies=cookies, json=payload) + print_response(r) + + +upload_w2l_template('https://drive.google.com/file/d/1yoBSb6nwFUZVI4CD-nM2BSofGdBjC8Nn/view?usp=sharing', 'Mr. Musk', session_nanashi) +upload_w2l_template('https://drive.google.com/file/d/1ysBAdRjlSmfWQPFjGlMJpv_zcGmVw7RC/view?usp=sharing', 'short bill clinton clip', session_brandon) +upload_w2l_template('https://i.imgur.com/NneE7Eq.jpg', 'itsa me mario!', session_echelon) +upload_w2l_template('https://i.imgur.com/W1EYiPU.jpg', 'Creepy Sonic does a gasp', session_echelon) +upload_w2l_template('https://i.imgur.com/8U3IdUa.png', 'Creepy Sonic actually talks', session_echelon) +upload_w2l_template('https://i.imgur.com/lKaQ4Er.jpg', 'hasta la vista', session_brandon) +upload_w2l_template('https://i.imgur.com/uoGuTJo.jpeg', 'what do you mean vern?', session_brandon) +upload_w2l_template('https://i.imgur.com/v8k9yau.jpeg', 'Miyamoto-san', session_brandon) + + +if hostname == 'halide': + upload_w2l_template('https://www.youtube.com/watch?v=a7mS9ZdU6k4', 'Steve Jobs talks', session_echelon) + upload_w2l_template('https://www.youtube.com/watch?v=6fg8SxgEUYI', 'Paul Graham lectures', session_echelon) + +# ==================== UPLOAD TTS MODELS ==================== + +def upload_tts_model(content_url, title, session_cookie): + print('Upload TTS model') + upload_url = 'http://localhost:12345/tts/upload' + cookies = { 'session': session_cookie } + payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': title, + 'download_url': content_url, + } + r = requests.post(upload_url, cookies=cookies, json=payload) + print_response(r) + + +# From others' Google Drives +#upload_tts_model('https://drive.google.com/file/d/1grwDGbAsPwCMVfye071TZ6PLUub1KGyR/view?usp=sharing', 'Jorgen', session_nanashi) +#upload_tts_model('https://drive.google.com/file/d/1w3oBk5vzyHurqYwP95UEQv-_NeMNu6pt/view?usp=sharing', 'Noire', session_echelon) + +# Copied into Storyteller Google Drive +upload_tts_model('https://drive.google.com/file/d/1d8KPwYqbRsGiYx1c4TCW4xCtcRLdIU1s/view?usp=sharing', 'Zephyr', session_echelon) +upload_tts_model('https://drive.google.com/file/d/1yhoamPz1NOMy403gu2v8sIx5hPj3LQWs/view?usp=sharing', 'Wakko Warner', session_other) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1I6KGQZFMMB2ZvCzWEqy9xTgGl0L5cmfP/view?usp=sharing', 'Turret-192', session_echelon) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1T_hvgGcNOfofhNmDWUqpVtYDdSXQrWdE/view?usp=sharing', 'technoblade-324', session_brandon) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1NKyRiex3K1J0QyIv2np17D_H0rE98HyK/view?usp=sharing', 'Noire', session_other) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1VDHDOpsUR8C-MD07vEsnUf1Lpcf8hXg4/view?usp=sharing', 'MatPat80', session_other) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1atFn4vm7mjkufXNK9TZS6J8WzISrAn3-/view?usp=sharing', 'King Julien-70', session_other) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1-rlJpi2jAsagn-UCStcCc_ITNycXgn2F/view?usp=sharing', 'Jorgen', session_other) +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1togTNOEKnaarqWWrbZLSwIbZIZ8IT_qw/view?usp=sharing', 'edsheeransinging', session_nanashi) +# TEMP COMMENT DUE TO RATE LIMIT +# TEMP COMMENT DUE TO RATE LIMIT upload_tts_model('https://drive.google.com/file/d/1OAiu0cZDyUJMemyLc0Btqw7zEQgiN81D/view?usp=sharing', 'Wakko Warner v2 ST21', session_nanashi) + +upload_tts_model('https://drive.google.com/file/d/1vRVI2gQ5tAs8SLQPrR_M1J77L05gGBOp/view?usp=sharing', 'Testing voice', session_nanashi) + +#upload_tts_model('', '', session_other) + diff --git a/debug/curl/session_info.py b/debug/curl/session_info.py new file mode 100755 index 0000000000..7326a74da1 --- /dev/null +++ b/debug/curl/session_info.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 + +import requests +import re + +login_url = 'http://localhost:12345/login' +session_info_url = 'http://localhost:12345/session' + + +r = requests.get(session_info_url) + +print("===== Session Info (logged out) =====") +print(r.content) + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + +cookies = { 'session': session_cookie } + +r = requests.get(session_info_url, cookies=cookies) + +print("===== Session Info (logged in) =====") +print(r.content) + diff --git a/debug/curl/twitch_event_rule_tests.py b/debug/curl/twitch_event_rule_tests.py new file mode 100755 index 0000000000..fd6583168e --- /dev/null +++ b/debug/curl/twitch_event_rule_tests.py @@ -0,0 +1,306 @@ +#!/usr/bin/env python3 + +import json +import re +import requests +import uuid +import pprint + +login_url = 'http://localhost:12345/login' +create_url = 'http://localhost:12345/twitch/event_rule/create' +list_url = 'http://localhost:12345/twitch/event_rule/list' +update_url = 'http://localhost:12345/twitch/event_rule/{}/update' +delete_url = 'http://localhost:12345/twitch/event_rule/{}/delete' + + +def delete_rule(token, cookies): + url = delete_url.format(token) + r = requests.delete(url, cookies=cookies) + print("===== Deleted Event Rule Output =====") + print('Status: {}'.format(r.status_code)) + print(r.content) + +def create_rule(payload, cookies): + r = requests.post(create_url, cookies=cookies, json=payload) + print("===== Created Event Rule Output =====") + print('Status: {}'.format(r.status_code)) + print(r.content) + + +# ========== Setup / Login ========== + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + +cookies = { 'session': session_cookie } + +# ========== Delete All Existing ========== + +r = requests.get(list_url, cookies=cookies) + +print("===== List Event Rule Output =====") +print('Status: {}'.format(r.status_code)) +response = json.loads(r.content) + +tokens = [rule['token'] for rule in response['twitch_event_rules']] +for token in tokens: + delete_rule(token, cookies) + +# ========== Create (1) ========== + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_spend_threshold': { + 'minimum_bits_spent': 1, + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +r = requests.post(create_url, cookies=cookies, json=payload) + +print("===== Created Event Rule Output =====") +print('Status: {}'.format(r.status_code)) +print(r.content) + +response = json.loads(r.content) +token = response['twitch_event_rule_token'] + +# ========== Update (1) ========== + +payload = { + 'event_match_predicate': { + 'bits_spend_threshold': { + 'minimum_bits_spent': 12345, + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 123, + 'rule_is_disabled': True, +} +url = update_url.format(token) +r = requests.post(url, cookies=cookies, json=payload) + +print("===== Edited Event Rule Output =====") +print('Status: {}'.format(r.status_code)) +print(r.content) + + +# ========== Create: Bits Spend Threshold ========== + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_spend_threshold': { + 'minimum_bits_spent': 1000, + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': True, +} + +create_rule(payload, cookies) + +# ========== Create: Cheermote Exact Match ========== + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_cheermote_name_exact_match': { + 'cheermote_name': 'Cheer1', + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_cheermote_name_exact_match': { + 'cheermote_name': 'Corgo5000', + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_cheermote_name_exact_match': { + 'cheermote_name': 'ZomboCom5000', + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +# ========== Create: Cheermote Prefix And Spend Threshold ========== + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_cheermote_prefix_spend_threshold': { + 'cheermote_prefix': 'DoodleCheer', + 'minimum_bits_spent': 1000, + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_cheermote_prefix_spend_threshold': { + 'cheermote_prefix': 'Corgo', + 'minimum_bits_spent': 12345, + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'bits', + 'event_match_predicate': { + 'bits_cheermote_prefix_spend_threshold': { + 'cheermote_prefix': 'ZomboCom', + 'minimum_bits_spent': 1000, + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +# ========== Create: Channel Points ========== + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'channel_points', + 'event_match_predicate': { + 'channel_points_reward_name_exact_match': { + 'reward_name': 'My Reward', + }, + }, + 'event_response': { + 'tts_single_voice': { + 'tts_model_token': 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'event_category': 'channel_points', + 'event_match_predicate': { + 'channel_points_reward_name_exact_match': { + 'reward_name': 'My Reward', + }, + }, + 'event_response': { + 'tts_random_voice': { + 'tts_model_tokens': [ + 'TM:40m3aqtt41y0', # "Wakko" voice (dev) + 'TM:4c1hycjj3a3t', # "Zephyr" voice (dev) + ], + } + }, + 'user_specified_rule_order': 0, + 'rule_is_disabled': False, +} + +create_rule(payload, cookies) + +# ========== List ========== + +r = requests.get(list_url, cookies=cookies) + +print("===== List Event Rule Output =====") +print('Status: {}'.format(r.status_code)) +response = json.loads(r.content) +pprint.pprint(response) + diff --git a/debug/curl/upload_tts.py b/debug/curl/upload_tts.py new file mode 100755 index 0000000000..f8da0109fd --- /dev/null +++ b/debug/curl/upload_tts.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 + +import requests +import re +import uuid + +login_url = 'http://localhost:12345/login' +upload_url = 'http://localhost:12345/tts/upload' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Upload TTS=====") + +cookies = { 'session': session_cookie } + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': 'this is a new text to speech model', + 'download_url': 'https://website.com', +} + +r = requests.post(upload_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/upload_tts_gdrive.py b/debug/curl/upload_tts_gdrive.py new file mode 100755 index 0000000000..de74ed3902 --- /dev/null +++ b/debug/curl/upload_tts_gdrive.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 + +import requests +import re +import uuid + +login_url = 'http://localhost:12345/login' +upload_url = 'http://localhost:12345/tts/upload' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Upload TTS=====") + +cookies = { 'session': session_cookie } + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': 'this is a new text to speech model', + 'download_url': 'https://drive.google.com/file/d/15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe/view?usp=sharing', +} + +r = requests.post(upload_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/upload_w2l.py b/debug/curl/upload_w2l.py new file mode 100755 index 0000000000..37b6e02d55 --- /dev/null +++ b/debug/curl/upload_w2l.py @@ -0,0 +1,45 @@ +#!/usr/bin/env python3 + +import re +import requests +import uuid + +login_url = 'http://localhost:12345/login' +upload_url = 'http://localhost:12345/w2l/upload' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Upload W2L =====") + +cookies = { 'session': session_cookie } + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': 'this is a new text to W2L template', + 'download_url': 'https://video.com', +} + +r = requests.post(upload_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/upload_w2l_image_gdrive.py b/debug/curl/upload_w2l_image_gdrive.py new file mode 100755 index 0000000000..250afcd4b1 --- /dev/null +++ b/debug/curl/upload_w2l_image_gdrive.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 + +import requests +import re +import uuid + +login_url = 'http://localhost:12345/login' +upload_url = 'http://localhost:12345/w2l/upload' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Upload W2L =====") + +cookies = { 'session': session_cookie } + +# mario.zip +#download_url = 'https://drive.google.com/file/d/15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe/view?usp=sharing' +# elon_integration_test.jpg +download_url = 'https://drive.google.com/file/d/1yoBSb6nwFUZVI4CD-nM2BSofGdBjC8Nn/view?usp=sharing' + +#download_url = 'https://i.imgur.com/...' + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': 'this is a new image template', + 'download_url': download_url, +} + +r = requests.post(upload_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/curl/upload_w2l_lots_of_data.py b/debug/curl/upload_w2l_lots_of_data.py new file mode 100755 index 0000000000..e960113a74 --- /dev/null +++ b/debug/curl/upload_w2l_lots_of_data.py @@ -0,0 +1,61 @@ +#!/usr/bin/env python3 + +import requests +import re +import uuid + +login_url = 'http://localhost:12345/login' +upload_url = 'http://localhost:12345/w2l/upload' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +def upload(url): + print("===== Upload W2L =====") + + cookies = { 'session': session_cookie } + + payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': 'this is a new image template', + 'download_url': url, + } + + r = requests.post(upload_url, cookies=cookies, json=payload) + + print('Status: {}'.format(r.status_code)) + print(r.content) + + +upload('https://i.imgur.com/W1EYiPU.jpg') + +upload('https://i.imgur.com/NneE7Eq.jpg') + +upload('https://i.imgur.com/v8k9yau.jpeg') + +upload('https://drive.google.com/file/d/1yoBSb6nwFUZVI4CD-nM2BSofGdBjC8Nn/view?usp=sharing') + +upload('https://drive.google.com/file/d/1ysBAdRjlSmfWQPFjGlMJpv_zcGmVw7RC/view?usp=sharing') + +upload('https://i.imgur.com/lKaQ4Er.jpg') + +upload('https://i.imgur.com/uoGuTJo.jpeg') + diff --git a/debug/curl/upload_w2l_video_gdrive.py b/debug/curl/upload_w2l_video_gdrive.py new file mode 100755 index 0000000000..be64549005 --- /dev/null +++ b/debug/curl/upload_w2l_video_gdrive.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 + +import requests +import re +import uuid + +login_url = 'http://localhost:12345/login' +upload_url = 'http://localhost:12345/w2l/upload' + +payload = { + 'username_or_email': 'echelon', + 'password': 'testing', +} + +r = requests.post(login_url, json=payload) + +print("===== Login =====") +print('Status: {}'.format(r.status_code)) +print(r.content) +for k, v in r.headers.items(): + print(' {}: {}'.format(k, v)) + +# NB: May not be set due to cookie domain: +#r.cookies['session'] + +raw_set_cookie_header = r.headers['set-cookie'] +m = re.match(r'session=([^;\s]+);', raw_set_cookie_header) +session_cookie = m.group(1) + + +print("===== Upload W2L =====") + +cookies = { 'session': session_cookie } + +# mario.zip +#download_url = 'https://drive.google.com/file/d/15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe/view?usp=sharing' +# elon_integration_test.jpg +#download_url = 'https://drive.google.com/file/d/1yoBSb6nwFUZVI4CD-nM2BSofGdBjC8Nn/view?usp=sharing' +# bill-integration-test.mp4 +# NB: This freezes up my poor laptop: +###download_url = 'https://drive.google.com/file/d/14yO_60lc1FJCYhDGaeiQabgGvYLa86WI/THIS_IS_TOO_BIGview?usp=sharing' +# bill-integration-small.mp4 (small enough to run on my laptop) +download_url = 'https://drive.google.com/file/d/1ysBAdRjlSmfWQPFjGlMJpv_zcGmVw7RC/view?usp=sharing' + +payload = { + 'idempotency_token': str(uuid.uuid4()), + 'title': 'this is a new image template', + 'download_url': download_url, +} + +r = requests.post(upload_url, cookies=cookies, json=payload) + +print('Status: {}'.format(r.status_code)) +print(r.content) + diff --git a/debug/mysql/badges.sh b/debug/mysql/badges.sh new file mode 100755 index 0000000000..ecc8cad037 --- /dev/null +++ b/debug/mysql/badges.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from badges\G" diff --git a/debug/mysql/ip_bans.sh b/debug/mysql/ip_bans.sh new file mode 100755 index 0000000000..b543e80a6f --- /dev/null +++ b/debug/mysql/ip_bans.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from ip_address_bans\G" diff --git a/debug/mysql/last_firehose_events.sh b/debug/mysql/last_firehose_events.sh new file mode 100755 index 0000000000..59b82af149 --- /dev/null +++ b/debug/mysql/last_firehose_events.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from firehose_entries\G" diff --git a/debug/mysql/last_session.sh b/debug/mysql/last_session.sh new file mode 100755 index 0000000000..149a552384 --- /dev/null +++ b/debug/mysql/last_session.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u root \ + -proot \ + -h localhost \ + -D storyteller \ + -e "select * from user_sessions\G" diff --git a/debug/mysql/last_tts_inference_job.sh b/debug/mysql/last_tts_inference_job.sh new file mode 100755 index 0000000000..f72fc914bc --- /dev/null +++ b/debug/mysql/last_tts_inference_job.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from tts_inference_jobs\G" diff --git a/debug/mysql/last_tts_inference_results.sh b/debug/mysql/last_tts_inference_results.sh new file mode 100755 index 0000000000..fea637e52e --- /dev/null +++ b/debug/mysql/last_tts_inference_results.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV + +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from tts_results\G" + + +#echo '===============================================' +# +#read -r -d '' query << EOF +# +# SELECT +# w2l_results.token as w2l_result_token, +# w2l_results.maybe_tts_inference_result_token, +# w2l_templates.token as maybe_w2l_template_token, +# w2l_templates.template_type, +# users.token as maybe_creator_user_token, +# users.username as maybe_creator_username, +# users.display_name as maybe_creator_display_name, +# w2l_results.file_size_bytes, +# w2l_results.frame_width, +# w2l_results.frame_height, +# w2l_results.duration_millis, +# w2l_results.created_at, +# w2l_results.updated_at +# FROM +# w2l_results +# LEFT OUTER JOIN w2l_templates ON w2l_results.maybe_w2l_template_token = w2l_templates.token +# LEFT OUTER JOIN users ON w2l_results.maybe_creator_user_token = users.token +# WHERE +# w2l_results.deleted_at IS NULL +# AND users.username = 'echelon'\G +#EOF +# +#echo $query +# +#mysql -u root \ +# -proot \ +# -h localhost \ +# -D storyteller \ +# -e "${query}" + diff --git a/debug/mysql/last_tts_model.sh b/debug/mysql/last_tts_model.sh new file mode 100755 index 0000000000..aacd3fc0da --- /dev/null +++ b/debug/mysql/last_tts_model.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from tts_models\G" diff --git a/debug/mysql/last_tts_upload.sh b/debug/mysql/last_tts_upload.sh new file mode 100755 index 0000000000..1772026b04 --- /dev/null +++ b/debug/mysql/last_tts_upload.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u root \ + -proot \ + -h localhost \ + -D storyteller \ + -e "select * from tts_model_upload_jobs\G" diff --git a/debug/mysql/last_user.sh b/debug/mysql/last_user.sh new file mode 100755 index 0000000000..a7c46a877d --- /dev/null +++ b/debug/mysql/last_user.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u root \ + -proot \ + -h localhost \ + -D storyteller \ + -e "select * from users\G" diff --git a/debug/mysql/last_w2l_inference_job.sh b/debug/mysql/last_w2l_inference_job.sh new file mode 100755 index 0000000000..9b53b29216 --- /dev/null +++ b/debug/mysql/last_w2l_inference_job.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u root \ + -proot \ + -h localhost \ + -D storyteller \ + -e "select * from w2l_inference_jobs\G" diff --git a/debug/mysql/last_w2l_inference_results.sh b/debug/mysql/last_w2l_inference_results.sh new file mode 100755 index 0000000000..28478a985b --- /dev/null +++ b/debug/mysql/last_w2l_inference_results.sh @@ -0,0 +1,45 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV + +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from w2l_results\G" + + +echo '===============================================' + +read -r -d '' query << EOF + + SELECT + w2l_results.token as w2l_result_token, + w2l_results.maybe_tts_inference_result_token, + w2l_templates.token as maybe_w2l_template_token, + w2l_templates.template_type, + users.token as maybe_creator_user_token, + users.username as maybe_creator_username, + users.display_name as maybe_creator_display_name, + w2l_results.file_size_bytes, + w2l_results.frame_width, + w2l_results.frame_height, + w2l_results.duration_millis, + w2l_results.created_at, + w2l_results.updated_at + FROM + w2l_results + LEFT OUTER JOIN w2l_templates ON w2l_results.maybe_w2l_template_token = w2l_templates.token + LEFT OUTER JOIN users ON w2l_results.maybe_creator_user_token = users.token + WHERE + w2l_results.deleted_at IS NULL + AND users.username = 'echelon'\G +EOF + +echo $query + +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "${query}" + diff --git a/debug/mysql/last_w2l_template.sh b/debug/mysql/last_w2l_template.sh new file mode 100755 index 0000000000..b71893244e --- /dev/null +++ b/debug/mysql/last_w2l_template.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from w2l_templates\G" diff --git a/debug/mysql/last_w2l_upload.sh b/debug/mysql/last_w2l_upload.sh new file mode 100755 index 0000000000..3d7a20448c --- /dev/null +++ b/debug/mysql/last_w2l_upload.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u root \ + -proot \ + -h localhost \ + -D storyteller \ + -e "select * from w2l_template_upload_jobs\G" diff --git a/debug/mysql/roles.sh b/debug/mysql/roles.sh new file mode 100755 index 0000000000..5bd88a8db6 --- /dev/null +++ b/debug/mysql/roles.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from user_roles\G" diff --git a/debug/mysql/set_admin.sh b/debug/mysql/set_admin.sh new file mode 100755 index 0000000000..72d93980b2 --- /dev/null +++ b/debug/mysql/set_admin.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "update users set user_role_slug='admin' where username='echelon' limit 1\G" diff --git a/debug/mysql/set_regular_user.sh b/debug/mysql/set_regular_user.sh new file mode 100755 index 0000000000..aff8414f96 --- /dev/null +++ b/debug/mysql/set_regular_user.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "update users set user_role_slug='user' where username='echelon' limit 1\G" diff --git a/debug/mysql/tts_results_synthetic_ids.sh b/debug/mysql/tts_results_synthetic_ids.sh new file mode 100755 index 0000000000..4bfbd85f23 --- /dev/null +++ b/debug/mysql/tts_results_synthetic_ids.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV +mysql -u storyteller \ + -ppassword \ + -h localhost \ + -D storyteller \ + -e "select * from tts_result_synthetic_ids\G" diff --git a/dev_database_connect.sh b/dev_database_connect.sh new file mode 100755 index 0000000000..41963f4569 --- /dev/null +++ b/dev_database_connect.sh @@ -0,0 +1,6 @@ +#!/bin/bash + +mysql -u "storyteller" \ + -p"password" \ + -D "storyteller" + diff --git a/diesel.toml b/diesel.toml new file mode 100644 index 0000000000..4e9295fe7c --- /dev/null +++ b/diesel.toml @@ -0,0 +1,13 @@ +# For documentation on how to configure this file, +# see diesel.rs/guides/configuring-diesel-cli + +[print_schema] +file = "_ignored_schema.rs" + +# I want to sort this directory first for easier reference +# Not released yet, but added here: +# https://github.com/diesel-rs/diesel/pull/2179 +# You'll need to install the a non-released diesel CLI +# (See storyteller READMEs/docs on how to install) +[migrations_directory] +dir = "_migrations" diff --git a/loc.sh b/loc.sh new file mode 100755 index 0000000000..50e23caf80 --- /dev/null +++ b/loc.sh @@ -0,0 +1,11 @@ +#!/bin/bash +wc -l \ + Cargo.toml \ + Dockerfile \ + migrations/*/*.sql \ + crates/*/*/*/*.rs \ + crates/*/*/*/*/*.rs \ + crates/*/*/*/*/*/*.rs \ + crates/*/*/*/*/*/*/*.rs \ + crates/*/*/*/*/*/*/*/*.rs \ + | sort -n diff --git a/localdev/nginx-http-config/storyteller b/localdev/nginx-http-config/storyteller new file mode 100644 index 0000000000..9410d80522 --- /dev/null +++ b/localdev/nginx-http-config/storyteller @@ -0,0 +1,57 @@ +# Nginx configs for "jungle.horse" and "api.jungle.horse" (http) +# +# On Ubuntu Linux, this config is located at the following location: +# +# - /etc/nginx/sites-enabled/{name-of-site} +# +# On Mac, this is located at the following locations: +# +# - /opt/homebrew/etc/nginx/servers/{name-of-site} (for nginx installed with homebrew) +# - /usr/local/etc/nginx/sites-enabled/{name-of-site} (for default nginx install) +# + +# API backend over plain HTTP +server { + listen 80; + server_name api.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:12345; + proxy_redirect off; + } +} + +# TODO(2022-02-13): This is new and untested +# WebSocket API backend over plain HTTP +server { + listen 80; + server_name ws.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:54321; + proxy_redirect off; + } +} + +# Javascript frontend over plain HTTP +server { + listen 80; + server_name jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:7000; + proxy_redirect off; + } +} + +# TODO(2022-02-13): This is new and untested +# Javascript frontend #2 (OBS Layer) over plain HTTP +server { + listen 80; + server_name obs.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:7001; + proxy_redirect off; + } +} diff --git a/localdev/nginx-http-config/storyteller-ssl b/localdev/nginx-http-config/storyteller-ssl new file mode 100644 index 0000000000..68fb8d9b24 --- /dev/null +++ b/localdev/nginx-http-config/storyteller-ssl @@ -0,0 +1,158 @@ +# Nginx configs for "jungle.horse" and "api.jungle.horse" (http) +# +# On Ubuntu Linux, this config is located at the following location: +# +# - /etc/nginx/sites-enabled/{name-of-site} +# +# Move this file there, and then restart with, +# +# sudo service nginx restart +# +# On Mac, this is located at the following locations: +# +# - /opt/homebrew/etc/nginx/servers/{name-of-site} (for nginx installed with homebrew) +# - /usr/local/etc/nginx/sites-enabled/{name-of-site} (for default nginx install) +# +# +# On Mac, verify configs and restart with, +# +# sudo nginx -t +# sudo nginx -s reload +# +# ================ +# SSL Instructions +# ================ +# +# To create self-signed certs on Ubuntu Linux, +# +# sudo apt-get install openssl +# +# sudo openssl req -x509 -nodes -days 365 -newkey rsa:2048 -keyout /etc/ssl/private/localhost.key -out /etc/ssl/certs/localhost.crt +# +# ^ most fields can be blank, but the "Common Name" should be `*.jungle.horse` +# +# For more: https://www.cloudsavvyit.com/1306/how-to-create-and-use-self-signed-ssl-on-nginx/ +# +# To create self-signed certs on Mac, +# +# openssl req -x509 -sha256 -nodes -newkey rsa:2048 -days 365 -keyout localhost.key -out localhost.crt +# +# ^ most fields can be blank, but the "Common Name" should be `*.jungle.horse` +# +# On Mac, load the self signed cert into the keychain with the following command: +# +# sudo security add-trusted-cert \ +# -d -r trustRoot \ +# -k /Library/Keychains/System.keychain /etc/ssl/certs/localhost.crt +# +# Self signed certs might be considered insecure by the browser. In Firefox, +# set about:config `security.tls.insecure_fallback_hosts` to include a list of +# permitted domains. You may also need to configure other settings. +# +# + +# API backend over plain HTTP +server { + listen 80; + server_name api.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:12345; + proxy_redirect off; + } +} + +# API backend over SSL +server { + listen 443 ssl; + listen [::]:443 ssl; + ssl_certificate /etc/ssl/certs/localhost.crt; + ssl_certificate_key /etc/ssl/private/localhost.key; + ssl_protocols TLSv1.2 TLSv1.1 TLSv1; + server_name api.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:12345; + proxy_redirect off; + } +} + +# TODO(2022-02-13): This is new and untested +# WebSocket API backend over plain HTTP +server { + listen 80; + server_name ws.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:54321; + proxy_redirect off; + } +} + +# Websocket API backend over SSL +server { + listen 443 ssl; + listen [::]:443 ssl; + ssl_certificate /etc/ssl/certs/localhost.crt; + ssl_certificate_key /etc/ssl/private/localhost.key; + ssl_protocols TLSv1.2 TLSv1.1 TLSv1; + server_name ws.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:54321; + proxy_redirect off; + } +} + +# Javascript frontend over plain HTTP +server { + listen 80; + server_name jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:7000; + proxy_redirect off; + } +} + +# TODO(2022-02-13): This is new and untested +# Javascript frontend (OBS Layer) over plain HTTP +server { + listen 80; + server_name obs.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:7007; + proxy_redirect off; + } +} + +# Javascript frontend over SSL +server { + listen 443 ssl; + listen [::]:443 ssl; + ssl_certificate /etc/ssl/certs/localhost.crt; + ssl_certificate_key /etc/ssl/private/localhost.key; + ssl_protocols TLSv1.2 TLSv1.1 TLSv1; + server_name jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:7000; + proxy_redirect off; + } +} + +# Javascript frontend over SSL (OBS Layer) +server { + listen 443 ssl; + listen [::]:443 ssl; + ssl_certificate /etc/ssl/certs/localhost.crt; + ssl_certificate_key /etc/ssl/private/localhost.key; + ssl_protocols TLSv1.2 TLSv1.1 TLSv1; + server_name obs.jungle.horse; + location / { + proxy_set_header Host $host; + proxy_pass http://127.0.0.1:7007; + proxy_redirect off; + } +} diff --git a/models/Wav2Lip/.gitignore b/models/Wav2Lip/.gitignore new file mode 100644 index 0000000000..56f02807a5 --- /dev/null +++ b/models/Wav2Lip/.gitignore @@ -0,0 +1,16 @@ +*.pkl +*.jpg +*.mp4 +*.pth +*.pyc +__pycache__ +*.h5 +*.avi +*.wav +filelists/*.txt +evaluation/test_filelists/lr*.txt +*.pyc +*.mkv +*.gif +*.webm +*.mp3 diff --git a/models/Wav2Lip/README.md b/models/Wav2Lip/README.md new file mode 100644 index 0000000000..19c43c453c --- /dev/null +++ b/models/Wav2Lip/README.md @@ -0,0 +1,145 @@ +# **Wav2Lip**: *Accurately Lip-syncing Videos In The Wild* + +This code is part of the paper: _A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild_ published at ACM Multimedia 2020. + +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/a-lip-sync-expert-is-all-you-need-for-speech/lip-sync-on-lrs2)](https://paperswithcode.com/sota/lip-sync-on-lrs2?p=a-lip-sync-expert-is-all-you-need-for-speech) +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/a-lip-sync-expert-is-all-you-need-for-speech/lip-sync-on-lrs3)](https://paperswithcode.com/sota/lip-sync-on-lrs3?p=a-lip-sync-expert-is-all-you-need-for-speech) +[![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/a-lip-sync-expert-is-all-you-need-for-speech/lip-sync-on-lrw)](https://paperswithcode.com/sota/lip-sync-on-lrw?p=a-lip-sync-expert-is-all-you-need-for-speech) + +[[Paper]](http://arxiv.org/abs/2008.10010) | [[Project Page]](http://cvit.iiit.ac.in/research/projects/cvit-projects/a-lip-sync-expert-is-all-you-need-for-speech-to-lip-generation-in-the-wild/) | [[Demo Video]](https://youtu.be/0fXaDCZNOJc) | [[Interactive Demo]](https://bhaasha.iiit.ac.in/lipsync) | [[Colab Notebook]](https://colab.research.google.com/drive/1tZpDWXz49W6wDcTprANRGLo2D_EbD5J8?usp=sharing) | [[ReSyncED] (coming soon)](#) + + + +---------- +**Highlights** +---------- + - Lip-sync videos to any target speech with high accuracy. Try our [interactive demo](https://bhaasha.iiit.ac.in/lipsync). + - Works for any identity, voice, and language. Also works for CGI faces and synthetic voices. + - Complete training code, inference code, and pretrained models are available. + - Or, quick-start with the Google Colab Notebook: [Link](https://colab.research.google.com/drive/1tZpDWXz49W6wDcTprANRGLo2D_EbD5J8?usp=sharing). Checkpoints and samples are available in a Google Drive [folder](https://drive.google.com/drive/folders/1I-0dNLfFOSFwrfqjNa-SXuwaURHE5K4k?usp=sharing) as well. + - Several new, reliable evaluation benchmarks and metrics [[`evaluation/` folder of this repo]](https://github.com/Rudrabha/Wav2Lip/tree/master/evaluation) released. + - Code to calculate metrics reported in the paper is also made available. + +-------- +**Disclaimer** +-------- +All results from this open-source code or our [demo website](https://bhaasha.iiit.ac.in/lipsync) should only be used for research/academic/personal purposes only. As the models are trained on the LRS2 dataset, any form of commercial use is strictly prohibhited. Please contact us for all further queries. + +Prerequisites +------------- +- `Python 3.5.2` (code has been tested with this version at our end, but several other users say that `3.6+` is the one that works instead.) +- ffmpeg: `sudo apt-get install ffmpeg` +- Install necessary packages using `pip install -r requirements.txt` +- Face detection [pre-trained model](https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth) should be downloaded to `face_detection/detection/sfd/s3fd.pth`. Alternative [link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/prajwal_k_research_iiit_ac_in/EZsy6qWuivtDnANIG73iHjIBjMSoojcIV0NULXV-yiuiIg?e=qTasa8) if the above does not work. + +Getting the weights +---------- +| Model | Description | Link to the model | +| :-------------: | :---------------: | :---------------: | +| Wav2Lip | Highly accurate lip-sync | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/Eb3LEzbfuKlJiR600lQWRxgBIY27JZg80f7V9jtMfbNDaQ?e=TBFBVW) | +| Wav2Lip + GAN | Slightly inferior lip-sync, but better visual quality | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EdjI7bZlgApMqsVoEUUXpLsBxqXbn5z8VTmoxp55YNDcIA?e=n9ljGW) | +| Expert Discriminator | Weights of the expert discriminator | [Link](https://iiitaphyd-my.sharepoint.com/:u:/g/personal/radrabha_m_research_iiit_ac_in/EQRvmiZg-HRAjvI6zqN9eTEBP74KefynCwPWVmF57l-AYA?e=ZRPHKP) | + +Lip-syncing videos using the pre-trained models (Inference) +------- +You can lip-sync any video to any audio: +```bash +python inference.py --checkpoint_path --face --audio +``` +The result is saved (by default) in `results/result_voice.mp4`. You can specify it as an argument, similar to several other available options. The audio source can be any file supported by `FFMPEG` containing audio data: `*.wav`, `*.mp3` or even a video file, from which the code will automatically extract the audio. + +##### Tips for better results: +- Experiment with the `--pads` argument to adjust the detected face bounding box. Often leads to improved results. You might need to increase the bottom padding to include the chin region. E.g. `--pads 0 20 0 0`. +- If you see the mouth position dislocated or some weird artifacts such as two mouths, then it can be because of over-smoothing the face detections. Use the `--nosmooth` argument and give another try. +- Experiment with the `--resize_factor` argument, to get a lower resolution video. Why? The models are trained on faces which were at a lower resolution. You might get better, visually pleasing results for 720p videos than for 1080p videos (in many cases, the latter works well too). +- The Wav2Lip model without GAN usually needs more experimenting with the above two to get the most ideal results, and sometimes, can give you a better result as well. + +Preparing LRS2 for training +---------- +Our models are trained on LRS2. See [here](#training-on-datasets-other-than-lrs2) for a few suggestions regarding training on other datasets. +##### LRS2 dataset folder structure + +``` +data_root (mvlrs_v1) +├── main, pretrain (we use only main folder in this work) +| ├── list of folders +| │ ├── five-digit numbered video IDs ending with (.mp4) +``` + +Place the LRS2 filelists (train, val, test) `.txt` files in the `filelists/` folder. + +##### Preprocess the dataset for fast training + +```bash +python preprocess.py --data_root data_root/main --preprocessed_root lrs2_preprocessed/ +``` +Additional options like `batch_size` and number of GPUs to use in parallel to use can also be set. + +##### Preprocessed LRS2 folder structure +``` +preprocessed_root (lrs2_preprocessed) +├── list of folders +| ├── Folders with five-digit numbered video IDs +| │ ├── *.jpg +| │ ├── audio.wav +``` + +Train! +---------- +There are two major steps: (i) Train the expert lip-sync discriminator, (ii) Train the Wav2Lip model(s). + +##### Training the expert discriminator +You can download [the pre-trained weights](#getting-the-weights) if you want to skip this step. To train it: +```bash +python color_syncnet_train.py --data_root lrs2_preprocessed/ --checkpoint_dir +``` +##### Training the Wav2Lip models +You can either train the model without the additional visual quality disriminator (< 1 day of training) or use the discriminator (~2 days). For the former, run: +```bash +python wav2lip_train.py --data_root lrs2_preprocessed/ --checkpoint_dir --syncnet_checkpoint_path +``` + +To train with the visual quality discriminator, you should run `hq_wav2lip_train.py` instead. The arguments for both the files are similar. In both the cases, you can resume training as well. Look at `python wav2lip_train.py --help` for more details. You can also set additional less commonly-used hyper-parameters at the bottom of the `hparams.py` file. + +Training on datasets other than LRS2 +------------------------------------ +Training on other datasets might require modifications to the code. Please read the following before you raise an issue: + +- You might not get good results by training/fine-tuning on a few minutes of a single speaker. This is a separate research problem, to which we do not have a solution yet. Thus, we would most likely not be able to resolve your issue. +- You must train the expert discriminator for your own dataset before training Wav2Lip. +- If it is your own dataset downloaded from the web, in most cases, needs to be sync-corrected. +- Be mindful of the FPS of the videos of your dataset. Changes to FPS would need significant code changes. +- The expert discriminator's eval loss should go down to ~0.25 and the Wav2Lip eval sync loss should go down to ~0.2 to get good results. + +When raising an issue on this topic, please let us know that you are aware of all these points. + +Evaluation +---------- +Will be updated. + +License and Citation +---------- +The software can only be used for personal/research/non-commercial purposes. Please cite the following paper if you have use this code: +``` +@inproceedings{10.1145/3394171.3413532, +author = {Prajwal, K R and Mukhopadhyay, Rudrabha and Namboodiri, Vinay P. and Jawahar, C.V.}, +title = {A Lip Sync Expert Is All You Need for Speech to Lip Generation In the Wild}, +year = {2020}, +isbn = {9781450379885}, +publisher = {Association for Computing Machinery}, +address = {New York, NY, USA}, +url = {https://doi.org/10.1145/3394171.3413532}, +doi = {10.1145/3394171.3413532}, +booktitle = {Proceedings of the 28th ACM International Conference on Multimedia}, +pages = {484–492}, +numpages = {9}, +keywords = {lip sync, talking face generation, video generation}, +location = {Seattle, WA, USA}, +series = {MM '20} +} +``` + + +Acknowledgements +---------- +Parts of the code structure is inspired by this [TTS repository](https://github.com/r9y9/deepvoice3_pytorch). We thank the author for this wonderful code. The code for Face Detection has been taken from the [face_alignment](https://github.com/1adrianb/face-alignment) repository. We thank the authors for releasing their code and models. diff --git a/models/Wav2Lip/audio.py b/models/Wav2Lip/audio.py new file mode 100644 index 0000000000..1c5c08332d --- /dev/null +++ b/models/Wav2Lip/audio.py @@ -0,0 +1,136 @@ +import librosa +import librosa.filters +import numpy as np +import tensorflow as tf +from scipy import signal +from scipy.io import wavfile +from hparams import hparams as hp + +def load_wav(path, sr): + return librosa.core.load(path, sr=sr)[0] + +def save_wav(wav, path, sr): + wav *= 32767 / max(0.01, np.max(np.abs(wav))) + #proposed by @dsmiller + wavfile.write(path, sr, wav.astype(np.int16)) + +def save_wavenet_wav(wav, path, sr): + librosa.output.write_wav(path, wav, sr=sr) + +def preemphasis(wav, k, preemphasize=True): + if preemphasize: + return signal.lfilter([1, -k], [1], wav) + return wav + +def inv_preemphasis(wav, k, inv_preemphasize=True): + if inv_preemphasize: + return signal.lfilter([1], [1, -k], wav) + return wav + +def get_hop_size(): + hop_size = hp.hop_size + if hop_size is None: + assert hp.frame_shift_ms is not None + hop_size = int(hp.frame_shift_ms / 1000 * hp.sample_rate) + return hop_size + +def linearspectrogram(wav): + D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize)) + S = _amp_to_db(np.abs(D)) - hp.ref_level_db + + if hp.signal_normalization: + return _normalize(S) + return S + +def melspectrogram(wav): + D = _stft(preemphasis(wav, hp.preemphasis, hp.preemphasize)) + S = _amp_to_db(_linear_to_mel(np.abs(D))) - hp.ref_level_db + + if hp.signal_normalization: + return _normalize(S) + return S + +def _lws_processor(): + import lws + return lws.lws(hp.n_fft, get_hop_size(), fftsize=hp.win_size, mode="speech") + +def _stft(y): + if hp.use_lws: + return _lws_processor(hp).stft(y).T + else: + return librosa.stft(y=y, n_fft=hp.n_fft, hop_length=get_hop_size(), win_length=hp.win_size) + +########################################################## +#Those are only correct when using lws!!! (This was messing with Wavenet quality for a long time!) +def num_frames(length, fsize, fshift): + """Compute number of time frames of spectrogram + """ + pad = (fsize - fshift) + if length % fshift == 0: + M = (length + pad * 2 - fsize) // fshift + 1 + else: + M = (length + pad * 2 - fsize) // fshift + 2 + return M + + +def pad_lr(x, fsize, fshift): + """Compute left and right padding + """ + M = num_frames(len(x), fsize, fshift) + pad = (fsize - fshift) + T = len(x) + 2 * pad + r = (M - 1) * fshift + fsize - T + return pad, pad + r +########################################################## +#Librosa correct padding +def librosa_pad_lr(x, fsize, fshift): + return 0, (x.shape[0] // fshift + 1) * fshift - x.shape[0] + +# Conversions +_mel_basis = None + +def _linear_to_mel(spectogram): + global _mel_basis + if _mel_basis is None: + _mel_basis = _build_mel_basis() + return np.dot(_mel_basis, spectogram) + +def _build_mel_basis(): + assert hp.fmax <= hp.sample_rate // 2 + return librosa.filters.mel(hp.sample_rate, hp.n_fft, n_mels=hp.num_mels, + fmin=hp.fmin, fmax=hp.fmax) + +def _amp_to_db(x): + min_level = np.exp(hp.min_level_db / 20 * np.log(10)) + return 20 * np.log10(np.maximum(min_level, x)) + +def _db_to_amp(x): + return np.power(10.0, (x) * 0.05) + +def _normalize(S): + if hp.allow_clipping_in_normalization: + if hp.symmetric_mels: + return np.clip((2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value, + -hp.max_abs_value, hp.max_abs_value) + else: + return np.clip(hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)), 0, hp.max_abs_value) + + assert S.max() <= 0 and S.min() - hp.min_level_db >= 0 + if hp.symmetric_mels: + return (2 * hp.max_abs_value) * ((S - hp.min_level_db) / (-hp.min_level_db)) - hp.max_abs_value + else: + return hp.max_abs_value * ((S - hp.min_level_db) / (-hp.min_level_db)) + +def _denormalize(D): + if hp.allow_clipping_in_normalization: + if hp.symmetric_mels: + return (((np.clip(D, -hp.max_abs_value, + hp.max_abs_value) + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + + hp.min_level_db) + else: + return ((np.clip(D, 0, hp.max_abs_value) * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db) + + if hp.symmetric_mels: + return (((D + hp.max_abs_value) * -hp.min_level_db / (2 * hp.max_abs_value)) + hp.min_level_db) + else: + return ((D * -hp.min_level_db / hp.max_abs_value) + hp.min_level_db) diff --git a/models/Wav2Lip/checkpoints/README.md b/models/Wav2Lip/checkpoints/README.md new file mode 100644 index 0000000000..80258ec8fb --- /dev/null +++ b/models/Wav2Lip/checkpoints/README.md @@ -0,0 +1 @@ +Place all your checkpoints (.pth files) here. \ No newline at end of file diff --git a/models/Wav2Lip/color_syncnet_train.py b/models/Wav2Lip/color_syncnet_train.py new file mode 100644 index 0000000000..afa0054438 --- /dev/null +++ b/models/Wav2Lip/color_syncnet_train.py @@ -0,0 +1,279 @@ +from os.path import dirname, join, basename, isfile +from tqdm import tqdm + +from models import SyncNet_color as SyncNet +import audio + +import torch +from torch import nn +from torch import optim +import torch.backends.cudnn as cudnn +from torch.utils import data as data_utils +import numpy as np + +from glob import glob + +import os, random, cv2, argparse +from hparams import hparams, get_image_list + +parser = argparse.ArgumentParser(description='Code to train the expert lip-sync discriminator') + +parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True) + +parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str) +parser.add_argument('--checkpoint_path', help='Resumed from this checkpoint', default=None, type=str) + +args = parser.parse_args() + + +global_step = 0 +global_epoch = 0 +use_cuda = torch.cuda.is_available() +print('use_cuda: {}'.format(use_cuda)) + +syncnet_T = 5 +syncnet_mel_step_size = 16 + +class Dataset(object): + def __init__(self, split): + self.all_videos = get_image_list(args.data_root, split) + + def get_frame_id(self, frame): + return int(basename(frame).split('.')[0]) + + def get_window(self, start_frame): + start_id = self.get_frame_id(start_frame) + vidname = dirname(start_frame) + + window_fnames = [] + for frame_id in range(start_id, start_id + syncnet_T): + frame = join(vidname, '{}.jpg'.format(frame_id)) + if not isfile(frame): + return None + window_fnames.append(frame) + return window_fnames + + def crop_audio_window(self, spec, start_frame): + # num_frames = (T x hop_size * fps) / sample_rate + start_frame_num = self.get_frame_id(start_frame) + start_idx = int(80. * (start_frame_num / float(hparams.fps))) + + end_idx = start_idx + syncnet_mel_step_size + + return spec[start_idx : end_idx, :] + + + def __len__(self): + return len(self.all_videos) + + def __getitem__(self, idx): + while 1: + idx = random.randint(0, len(self.all_videos) - 1) + vidname = self.all_videos[idx] + + img_names = list(glob(join(vidname, '*.jpg'))) + if len(img_names) <= 3 * syncnet_T: + continue + img_name = random.choice(img_names) + wrong_img_name = random.choice(img_names) + while wrong_img_name == img_name: + wrong_img_name = random.choice(img_names) + + if random.choice([True, False]): + y = torch.ones(1).float() + chosen = img_name + else: + y = torch.zeros(1).float() + chosen = wrong_img_name + + window_fnames = self.get_window(chosen) + if window_fnames is None: + continue + + window = [] + all_read = True + for fname in window_fnames: + img = cv2.imread(fname) + if img is None: + all_read = False + break + try: + img = cv2.resize(img, (hparams.img_size, hparams.img_size)) + except Exception as e: + all_read = False + break + + window.append(img) + + if not all_read: continue + + try: + wavpath = join(vidname, "audio.wav") + wav = audio.load_wav(wavpath, hparams.sample_rate) + + orig_mel = audio.melspectrogram(wav).T + except Exception as e: + continue + + mel = self.crop_audio_window(orig_mel.copy(), img_name) + + if (mel.shape[0] != syncnet_mel_step_size): + continue + + # H x W x 3 * T + x = np.concatenate(window, axis=2) / 255. + x = x.transpose(2, 0, 1) + x = x[:, x.shape[1]//2:] + + x = torch.FloatTensor(x) + mel = torch.FloatTensor(mel.T).unsqueeze(0) + + return x, mel, y + +logloss = nn.BCELoss() +def cosine_loss(a, v, y): + d = nn.functional.cosine_similarity(a, v) + loss = logloss(d.unsqueeze(1), y) + + return loss + +def train(device, model, train_data_loader, test_data_loader, optimizer, + checkpoint_dir=None, checkpoint_interval=None, nepochs=None): + + global global_step, global_epoch + resumed_step = global_step + + while global_epoch < nepochs: + running_loss = 0. + prog_bar = tqdm(enumerate(train_data_loader)) + for step, (x, mel, y) in prog_bar: + model.train() + optimizer.zero_grad() + + # Transform data to CUDA device + x = x.to(device) + + mel = mel.to(device) + + a, v = model(mel, x) + y = y.to(device) + + loss = cosine_loss(a, v, y) + loss.backward() + optimizer.step() + + global_step += 1 + cur_session_steps = global_step - resumed_step + running_loss += loss.item() + + if global_step == 1 or global_step % checkpoint_interval == 0: + save_checkpoint( + model, optimizer, global_step, checkpoint_dir, global_epoch) + + if global_step % hparams.syncnet_eval_interval == 0: + with torch.no_grad(): + eval_model(test_data_loader, global_step, device, model, checkpoint_dir) + + prog_bar.set_description('Loss: {}'.format(running_loss / (step + 1))) + + global_epoch += 1 + +def eval_model(test_data_loader, global_step, device, model, checkpoint_dir): + eval_steps = 1400 + print('Evaluating for {} steps'.format(eval_steps)) + losses = [] + while 1: + for step, (x, mel, y) in enumerate(test_data_loader): + + model.eval() + + # Transform data to CUDA device + x = x.to(device) + + mel = mel.to(device) + + a, v = model(mel, x) + y = y.to(device) + + loss = cosine_loss(a, v, y) + losses.append(loss.item()) + + if step > eval_steps: break + + averaged_loss = sum(losses) / len(losses) + print(averaged_loss) + + return + +def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch): + + checkpoint_path = join( + checkpoint_dir, "checkpoint_step{:09d}.pth".format(global_step)) + optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None + torch.save({ + "state_dict": model.state_dict(), + "optimizer": optimizer_state, + "global_step": step, + "global_epoch": epoch, + }, checkpoint_path) + print("Saved checkpoint:", checkpoint_path) + +def _load(checkpoint_path): + if use_cuda: + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_checkpoint(path, model, optimizer, reset_optimizer=False): + global global_step + global global_epoch + + print("Load checkpoint from: {}".format(path)) + checkpoint = _load(path) + model.load_state_dict(checkpoint["state_dict"]) + if not reset_optimizer: + optimizer_state = checkpoint["optimizer"] + if optimizer_state is not None: + print("Load optimizer state from {}".format(path)) + optimizer.load_state_dict(checkpoint["optimizer"]) + global_step = checkpoint["global_step"] + global_epoch = checkpoint["global_epoch"] + + return model + +if __name__ == "__main__": + checkpoint_dir = args.checkpoint_dir + checkpoint_path = args.checkpoint_path + + if not os.path.exists(checkpoint_dir): os.mkdir(checkpoint_dir) + + # Dataset and Dataloader setup + train_dataset = Dataset('train') + test_dataset = Dataset('val') + + train_data_loader = data_utils.DataLoader( + train_dataset, batch_size=hparams.syncnet_batch_size, shuffle=True, + num_workers=hparams.num_workers) + + test_data_loader = data_utils.DataLoader( + test_dataset, batch_size=hparams.syncnet_batch_size, + num_workers=8) + + device = torch.device("cuda" if use_cuda else "cpu") + + # Model + model = SyncNet().to(device) + print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad))) + + optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad], + lr=hparams.syncnet_lr) + + if checkpoint_path is not None: + load_checkpoint(checkpoint_path, model, optimizer, reset_optimizer=False) + + train(device, model, train_data_loader, test_data_loader, optimizer, + checkpoint_dir=checkpoint_dir, + checkpoint_interval=hparams.syncnet_checkpoint_interval, + nepochs=hparams.nepochs) diff --git a/models/Wav2Lip/evaluation/README.md b/models/Wav2Lip/evaluation/README.md new file mode 100644 index 0000000000..613e3a3b59 --- /dev/null +++ b/models/Wav2Lip/evaluation/README.md @@ -0,0 +1,59 @@ +# Evaluation of Lip-sync using LSE-D and LSE-C metric. + +We use the pre-trained syncnet model available in this [repository](https://github.com/joonson/syncnet_python). + +### Steps to set-up the evaluation repository for LSE-D and LSE-C metric: +* Clone the SyncNet repository. +``` +git clone https://github.com/joonson/syncnet_python.git +``` +* Follow the procedure given in the above linked [repository](https://github.com/joonson/syncnet_python) to download the pretrained models and set up the dependencies. + * **Note: Please install a separate virtual environment for the evaluation scripts. The versions used by Wav2Lip and the publicly released code of SyncNet is different and can cause version mis-match issues. To avoid this, we suggest the users to install a separate virtual environment for the evaluation scripts** +``` +cd syncnet_python +pip install -r requirements.txt +sh download_model.sh +``` +* The above step should ensure that all the dependencies required by the repository is installed and the pre-trained models are downloaded. + +### Running the evaluation scripts: +* Copy our evaluation scripts given in this folder to the cloned repository. +``` + cd Wav2Lip/evaluation/scores_LSE/ + cp *.py syncnet_python/ + cp *.sh syncnet_python/ +``` +**Note: We will release the test filelists for LRW, LRS2 and LRS3 shortly once we receive permission from the dataset creators. We will also release the Real World Dataset we have collected shortly.** + +* Our evaluation technique does not require ground-truth of any sorts. Given lip-synced videos we can directly calculate the scores from only the generated videos. Please store the generated videos (from our test sets or your own generated videos) in the following folder structure. +``` +video data root (Folder containing all videos) +├── All .mp4 files +``` +* Change the folder back to the cloned repository. +``` +cd syncnet_python +``` +* To run evaluation on the LRW, LRS2 and LRS3 test files, please run the following command: +``` +python calculate_scores_LRS.py --data_root /path/to/video/data/root --tmp_dir tmp_dir/ +``` + +* To run evaluation on the ReSynced dataset or your own generated videos, please run the following command: +``` +sh calculate_scores_real_videos.sh /path/to/video/data/root +``` +* The generated scores will be present in the all_scores.txt generated in the ```syncnet_python/``` folder + +# Evaluation of image quality using FID metric. +We will update the scripts for the FID metric shortly + + +# Opening issues related to evaluation scripts +* Please open the issues with the "Evaluation" label if you face any issues in the evaluation scripts. + +# Acknowledgements +Our evaluation pipeline in based on two existing repositories. LSE metrics are based on the [syncnet_python](https://github.com/joonson/syncnet_python) repository and the FID score is based on [pytorch-fid](https://github.com/mseitzer/pytorch-fid) repository. We thank the authors of both the repositories for releasing their wonderful code. + + + diff --git a/models/Wav2Lip/evaluation/gen_videos_from_filelist.py b/models/Wav2Lip/evaluation/gen_videos_from_filelist.py new file mode 100644 index 0000000000..bd666b9325 --- /dev/null +++ b/models/Wav2Lip/evaluation/gen_videos_from_filelist.py @@ -0,0 +1,238 @@ +from os import listdir, path +import numpy as np +import scipy, cv2, os, sys, argparse +import dlib, json, subprocess +from tqdm import tqdm +from glob import glob +import torch + +sys.path.append('../') +import audio +import face_detection +from models import Wav2Lip + +parser = argparse.ArgumentParser(description='Code to generate results for test filelists') + +parser.add_argument('--filelist', type=str, + help='Filepath of filelist file to read', required=True) +parser.add_argument('--results_dir', type=str, help='Folder to save all results into', + required=True) +parser.add_argument('--data_root', type=str, required=True) +parser.add_argument('--checkpoint_path', type=str, + help='Name of saved checkpoint to load weights from', required=True) + +parser.add_argument('--pads', nargs='+', type=int, default=[0, 0, 0, 0], + help='Padding (top, bottom, left, right)') +parser.add_argument('--face_det_batch_size', type=int, + help='Single GPU batch size for face detection', default=64) +parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip', default=128) + +# parser.add_argument('--resize_factor', default=1, type=int) + +args = parser.parse_args() +args.img_size = 96 + +def get_smoothened_boxes(boxes, T): + for i in range(len(boxes)): + if i + T > len(boxes): + window = boxes[len(boxes) - T:] + else: + window = boxes[i : i + T] + boxes[i] = np.mean(window, axis=0) + return boxes + +def face_detect(images): + batch_size = args.face_det_batch_size + + while 1: + predictions = [] + try: + for i in range(0, len(images), batch_size): + predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size]))) + except RuntimeError: + if batch_size == 1: + raise RuntimeError('Image too big to run face detection on GPU') + batch_size //= 2 + args.face_det_batch_size = batch_size + print('Recovering from OOM error; New batch size: {}'.format(batch_size)) + continue + break + + results = [] + pady1, pady2, padx1, padx2 = args.pads + for rect, image in zip(predictions, images): + if rect is None: + raise ValueError('Face not detected!') + + y1 = max(0, rect[1] - pady1) + y2 = min(image.shape[0], rect[3] + pady2) + x1 = max(0, rect[0] - padx1) + x2 = min(image.shape[1], rect[2] + padx2) + + results.append([x1, y1, x2, y2]) + + boxes = get_smoothened_boxes(np.array(results), T=5) + results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)] + + return results + +def datagen(frames, face_det_results, mels): + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + for i, m in enumerate(mels): + if i >= len(frames): raise ValueError('Equal or less lengths only') + + frame_to_save = frames[i].copy() + face, coords, valid_frame = face_det_results[i].copy() + if not valid_frame: + continue + + face = cv2.resize(face, (args.img_size, args.img_size)) + + img_batch.append(face) + mel_batch.append(m) + frame_batch.append(frame_to_save) + coords_batch.append(coords) + + if len(img_batch) >= args.wav2lip_batch_size: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if len(img_batch) > 0: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + +fps = 25 +mel_step_size = 16 +mel_idx_multiplier = 80./fps +device = 'cuda' if torch.cuda.is_available() else 'cpu' +print('Using {} for inference.'.format(device)) + +detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, + flip_input=False, device=device) + +def _load(checkpoint_path): + if device == 'cuda': + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_model(path): + model = Wav2Lip() + print("Load checkpoint from: {}".format(path)) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + + model = model.to(device) + return model.eval() + +model = load_model(args.checkpoint_path) + +def main(): + assert args.data_root is not None + data_root = args.data_root + + if not os.path.isdir(args.results_dir): os.makedirs(args.results_dir) + + with open(args.filelist, 'r') as filelist: + lines = filelist.readlines() + + for idx, line in enumerate(tqdm(lines)): + audio_src, video = line.strip().split() + + audio_src = os.path.join(data_root, audio_src) + '.mp4' + video = os.path.join(data_root, video) + '.mp4' + + command = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'.format(audio_src, '../temp/temp.wav') + subprocess.call(command, shell=True) + temp_audio = '../temp/temp.wav' + + wav = audio.load_wav(temp_audio, 16000) + mel = audio.melspectrogram(wav) + if np.isnan(mel.reshape(-1)).sum() > 0: + continue + + mel_chunks = [] + i = 0 + while 1: + start_idx = int(i * mel_idx_multiplier) + if start_idx + mel_step_size > len(mel[0]): + break + mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) + i += 1 + + video_stream = cv2.VideoCapture(video) + + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading or len(full_frames) > len(mel_chunks): + video_stream.release() + break + full_frames.append(frame) + + if len(full_frames) < len(mel_chunks): + continue + + full_frames = full_frames[:len(mel_chunks)] + + try: + face_det_results = face_detect(full_frames.copy()) + except ValueError as e: + continue + + batch_size = args.wav2lip_batch_size + gen = datagen(full_frames.copy(), face_det_results, mel_chunks) + + for i, (img_batch, mel_batch, frames, coords) in enumerate(gen): + if i == 0: + frame_h, frame_w = full_frames[0].shape[:-1] + out = cv2.VideoWriter('../temp/result.avi', + cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h)) + + img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) + mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) + + with torch.no_grad(): + pred = model(mel_batch, img_batch) + + + pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. + + for pl, f, c in zip(pred, frames, coords): + y1, y2, x1, x2 = c + pl = cv2.resize(pl.astype(np.uint8), (x2 - x1, y2 - y1)) + f[y1:y2, x1:x2] = pl + out.write(f) + + out.release() + + vid = os.path.join(args.results_dir, '{}.mp4'.format(idx)) + + command = 'ffmpeg -loglevel panic -y -i {} -i {} -strict -2 -q:v 1 {}'.format(temp_audio, + '../temp/result.avi', vid) + subprocess.call(command, shell=True) + +if __name__ == '__main__': + main() diff --git a/models/Wav2Lip/evaluation/real_videos_inference.py b/models/Wav2Lip/evaluation/real_videos_inference.py new file mode 100644 index 0000000000..8c9fb15ef3 --- /dev/null +++ b/models/Wav2Lip/evaluation/real_videos_inference.py @@ -0,0 +1,305 @@ +from os import listdir, path +import numpy as np +import scipy, cv2, os, sys, argparse +import dlib, json, subprocess +from tqdm import tqdm +from glob import glob +import torch + +sys.path.append('../') +import audio +import face_detection +from models import Wav2Lip + +parser = argparse.ArgumentParser(description='Code to generate results on ReSyncED evaluation set') + +parser.add_argument('--mode', type=str, + help='random | dubbed | tts', required=True) + +parser.add_argument('--filelist', type=str, + help='Filepath of filelist file to read', default=None) + +parser.add_argument('--results_dir', type=str, help='Folder to save all results into', + required=True) +parser.add_argument('--data_root', type=str, required=True) +parser.add_argument('--checkpoint_path', type=str, + help='Name of saved checkpoint to load weights from', required=True) +parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], + help='Padding (top, bottom, left, right)') + +parser.add_argument('--face_det_batch_size', type=int, + help='Single GPU batch size for face detection', default=16) + +parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip', default=128) +parser.add_argument('--face_res', help='Approximate resolution of the face at which to test', default=180) +parser.add_argument('--min_frame_res', help='Do not downsample further below this frame resolution', default=480) +parser.add_argument('--max_frame_res', help='Downsample to at least this frame resolution', default=720) +# parser.add_argument('--resize_factor', default=1, type=int) + +args = parser.parse_args() +args.img_size = 96 + +def get_smoothened_boxes(boxes, T): + for i in range(len(boxes)): + if i + T > len(boxes): + window = boxes[len(boxes) - T:] + else: + window = boxes[i : i + T] + boxes[i] = np.mean(window, axis=0) + return boxes + +def rescale_frames(images): + rect = detector.get_detections_for_batch(np.array([images[0]]))[0] + if rect is None: + raise ValueError('Face not detected!') + h, w = images[0].shape[:-1] + + x1, y1, x2, y2 = rect + + face_size = max(np.abs(y1 - y2), np.abs(x1 - x2)) + + diff = np.abs(face_size - args.face_res) + for factor in range(2, 16): + downsampled_res = face_size // factor + if min(h//factor, w//factor) < args.min_frame_res: break + if np.abs(downsampled_res - args.face_res) >= diff: break + + factor -= 1 + if factor == 1: return images + + return [cv2.resize(im, (im.shape[1]//(factor), im.shape[0]//(factor))) for im in images] + + +def face_detect(images): + batch_size = args.face_det_batch_size + images = rescale_frames(images) + + while 1: + predictions = [] + try: + for i in range(0, len(images), batch_size): + predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size]))) + except RuntimeError: + if batch_size == 1: + raise RuntimeError('Image too big to run face detection on GPU') + batch_size //= 2 + print('Recovering from OOM error; New batch size: {}'.format(batch_size)) + continue + break + + results = [] + pady1, pady2, padx1, padx2 = args.pads + for rect, image in zip(predictions, images): + if rect is None: + raise ValueError('Face not detected!') + + y1 = max(0, rect[1] - pady1) + y2 = min(image.shape[0], rect[3] + pady2) + x1 = max(0, rect[0] - padx1) + x2 = min(image.shape[1], rect[2] + padx2) + + results.append([x1, y1, x2, y2]) + + boxes = get_smoothened_boxes(np.array(results), T=5) + results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2), True] for image, (x1, y1, x2, y2) in zip(images, boxes)] + + return results, images + +def datagen(frames, face_det_results, mels): + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + for i, m in enumerate(mels): + if i >= len(frames): raise ValueError('Equal or less lengths only') + + frame_to_save = frames[i].copy() + face, coords, valid_frame = face_det_results[i].copy() + if not valid_frame: + continue + + face = cv2.resize(face, (args.img_size, args.img_size)) + + img_batch.append(face) + mel_batch.append(m) + frame_batch.append(frame_to_save) + coords_batch.append(coords) + + if len(img_batch) >= args.wav2lip_batch_size: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if len(img_batch) > 0: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + +def increase_frames(frames, l): + ## evenly duplicating frames to increase length of video + while len(frames) < l: + dup_every = float(l) / len(frames) + + final_frames = [] + next_duplicate = 0. + + for i, f in enumerate(frames): + final_frames.append(f) + + if int(np.ceil(next_duplicate)) == i: + final_frames.append(f) + + next_duplicate += dup_every + + frames = final_frames + + return frames[:l] + +mel_step_size = 16 +device = 'cuda' if torch.cuda.is_available() else 'cpu' +print('Using {} for inference.'.format(device)) + +detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, + flip_input=False, device=device) + +def _load(checkpoint_path): + if device == 'cuda': + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_model(path): + model = Wav2Lip() + print("Load checkpoint from: {}".format(path)) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + + model = model.to(device) + return model.eval() + +model = load_model(args.checkpoint_path) + +def main(): + if not os.path.isdir(args.results_dir): os.makedirs(args.results_dir) + + if args.mode == 'dubbed': + files = listdir(args.data_root) + lines = ['{} {}'.format(f, f) for f in files] + + else: + assert args.filelist is not None + with open(args.filelist, 'r') as filelist: + lines = filelist.readlines() + + for idx, line in enumerate(tqdm(lines)): + video, audio_src = line.strip().split() + + audio_src = os.path.join(args.data_root, audio_src) + video = os.path.join(args.data_root, video) + + command = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}'.format(audio_src, '../temp/temp.wav') + subprocess.call(command, shell=True) + temp_audio = '../temp/temp.wav' + + wav = audio.load_wav(temp_audio, 16000) + mel = audio.melspectrogram(wav) + + if np.isnan(mel.reshape(-1)).sum() > 0: + raise ValueError('Mel contains nan!') + + video_stream = cv2.VideoCapture(video) + + fps = video_stream.get(cv2.CAP_PROP_FPS) + mel_idx_multiplier = 80./fps + + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + + if min(frame.shape[:-1]) > args.max_frame_res: + h, w = frame.shape[:-1] + scale_factor = min(h, w) / float(args.max_frame_res) + h = int(h/scale_factor) + w = int(w/scale_factor) + + frame = cv2.resize(frame, (w, h)) + full_frames.append(frame) + + mel_chunks = [] + i = 0 + while 1: + start_idx = int(i * mel_idx_multiplier) + if start_idx + mel_step_size > len(mel[0]): + break + mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) + i += 1 + + if len(full_frames) < len(mel_chunks): + if args.mode == 'tts': + full_frames = increase_frames(full_frames, len(mel_chunks)) + else: + raise ValueError('#Frames, audio length mismatch') + + else: + full_frames = full_frames[:len(mel_chunks)] + + try: + face_det_results, full_frames = face_detect(full_frames.copy()) + except ValueError as e: + continue + + batch_size = args.wav2lip_batch_size + gen = datagen(full_frames.copy(), face_det_results, mel_chunks) + + for i, (img_batch, mel_batch, frames, coords) in enumerate(gen): + if i == 0: + frame_h, frame_w = full_frames[0].shape[:-1] + + out = cv2.VideoWriter('../temp/result.avi', + cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h)) + + img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) + mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) + + with torch.no_grad(): + pred = model(mel_batch, img_batch) + + + pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. + + for pl, f, c in zip(pred, frames, coords): + y1, y2, x1, x2 = c + pl = cv2.resize(pl.astype(np.uint8), (x2 - x1, y2 - y1)) + f[y1:y2, x1:x2] = pl + out.write(f) + + out.release() + + vid = os.path.join(args.results_dir, '{}.mp4'.format(idx)) + command = 'ffmpeg -loglevel panic -y -i {} -i {} -strict -2 -q:v 1 {}'.format('../temp/temp.wav', + '../temp/result.avi', vid) + subprocess.call(command, shell=True) + + +if __name__ == '__main__': + main() diff --git a/models/Wav2Lip/evaluation/scores_LSE/SyncNetInstance_calc_scores.py b/models/Wav2Lip/evaluation/scores_LSE/SyncNetInstance_calc_scores.py new file mode 100644 index 0000000000..64906e257b --- /dev/null +++ b/models/Wav2Lip/evaluation/scores_LSE/SyncNetInstance_calc_scores.py @@ -0,0 +1,210 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- +# Video 25 FPS, Audio 16000HZ + +import torch +import numpy +import time, pdb, argparse, subprocess, os, math, glob +import cv2 +import python_speech_features + +from scipy import signal +from scipy.io import wavfile +from SyncNetModel import * +from shutil import rmtree + + +# ==================== Get OFFSET ==================== + +def calc_pdist(feat1, feat2, vshift=10): + + win_size = vshift*2+1 + + feat2p = torch.nn.functional.pad(feat2,(0,0,vshift,vshift)) + + dists = [] + + for i in range(0,len(feat1)): + + dists.append(torch.nn.functional.pairwise_distance(feat1[[i],:].repeat(win_size, 1), feat2p[i:i+win_size,:])) + + return dists + +# ==================== MAIN DEF ==================== + +class SyncNetInstance(torch.nn.Module): + + def __init__(self, dropout = 0, num_layers_in_fc_layers = 1024): + super(SyncNetInstance, self).__init__(); + + self.__S__ = S(num_layers_in_fc_layers = num_layers_in_fc_layers).cuda(); + + def evaluate(self, opt, videofile): + + self.__S__.eval(); + + # ========== ========== + # Convert files + # ========== ========== + + if os.path.exists(os.path.join(opt.tmp_dir,opt.reference)): + rmtree(os.path.join(opt.tmp_dir,opt.reference)) + + os.makedirs(os.path.join(opt.tmp_dir,opt.reference)) + + command = ("ffmpeg -loglevel error -y -i %s -threads 1 -f image2 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'%06d.jpg'))) + output = subprocess.call(command, shell=True, stdout=None) + + command = ("ffmpeg -loglevel error -y -i %s -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 %s" % (videofile,os.path.join(opt.tmp_dir,opt.reference,'audio.wav'))) + output = subprocess.call(command, shell=True, stdout=None) + + # ========== ========== + # Load video + # ========== ========== + + images = [] + + flist = glob.glob(os.path.join(opt.tmp_dir,opt.reference,'*.jpg')) + flist.sort() + + for fname in flist: + img_input = cv2.imread(fname) + img_input = cv2.resize(img_input, (224,224)) #HARD CODED, CHANGE BEFORE RELEASE + images.append(img_input) + + im = numpy.stack(images,axis=3) + im = numpy.expand_dims(im,axis=0) + im = numpy.transpose(im,(0,3,4,1,2)) + + imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float()) + + # ========== ========== + # Load audio + # ========== ========== + + sample_rate, audio = wavfile.read(os.path.join(opt.tmp_dir,opt.reference,'audio.wav')) + mfcc = zip(*python_speech_features.mfcc(audio,sample_rate)) + mfcc = numpy.stack([numpy.array(i) for i in mfcc]) + + cc = numpy.expand_dims(numpy.expand_dims(mfcc,axis=0),axis=0) + cct = torch.autograd.Variable(torch.from_numpy(cc.astype(float)).float()) + + # ========== ========== + # Check audio and video input length + # ========== ========== + + #if (float(len(audio))/16000) != (float(len(images))/25) : + # print("WARNING: Audio (%.4fs) and video (%.4fs) lengths are different."%(float(len(audio))/16000,float(len(images))/25)) + + min_length = min(len(images),math.floor(len(audio)/640)) + + # ========== ========== + # Generate video and audio feats + # ========== ========== + + lastframe = min_length-5 + im_feat = [] + cc_feat = [] + + tS = time.time() + for i in range(0,lastframe,opt.batch_size): + + im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ] + im_in = torch.cat(im_batch,0) + im_out = self.__S__.forward_lip(im_in.cuda()); + im_feat.append(im_out.data.cpu()) + + cc_batch = [ cct[:,:,:,vframe*4:vframe*4+20] for vframe in range(i,min(lastframe,i+opt.batch_size)) ] + cc_in = torch.cat(cc_batch,0) + cc_out = self.__S__.forward_aud(cc_in.cuda()) + cc_feat.append(cc_out.data.cpu()) + + im_feat = torch.cat(im_feat,0) + cc_feat = torch.cat(cc_feat,0) + + # ========== ========== + # Compute offset + # ========== ========== + + #print('Compute time %.3f sec.' % (time.time()-tS)) + + dists = calc_pdist(im_feat,cc_feat,vshift=opt.vshift) + mdist = torch.mean(torch.stack(dists,1),1) + + minval, minidx = torch.min(mdist,0) + + offset = opt.vshift-minidx + conf = torch.median(mdist) - minval + + fdist = numpy.stack([dist[minidx].numpy() for dist in dists]) + # fdist = numpy.pad(fdist, (3,3), 'constant', constant_values=15) + fconf = torch.median(mdist).numpy() - fdist + fconfm = signal.medfilt(fconf,kernel_size=9) + + numpy.set_printoptions(formatter={'float': '{: 0.3f}'.format}) + #print('Framewise conf: ') + #print(fconfm) + #print('AV offset: \t%d \nMin dist: \t%.3f\nConfidence: \t%.3f' % (offset,minval,conf)) + + dists_npy = numpy.array([ dist.numpy() for dist in dists ]) + return offset.numpy(), conf.numpy(), minval.numpy() + + def extract_feature(self, opt, videofile): + + self.__S__.eval(); + + # ========== ========== + # Load video + # ========== ========== + cap = cv2.VideoCapture(videofile) + + frame_num = 1; + images = [] + while frame_num: + frame_num += 1 + ret, image = cap.read() + if ret == 0: + break + + images.append(image) + + im = numpy.stack(images,axis=3) + im = numpy.expand_dims(im,axis=0) + im = numpy.transpose(im,(0,3,4,1,2)) + + imtv = torch.autograd.Variable(torch.from_numpy(im.astype(float)).float()) + + # ========== ========== + # Generate video feats + # ========== ========== + + lastframe = len(images)-4 + im_feat = [] + + tS = time.time() + for i in range(0,lastframe,opt.batch_size): + + im_batch = [ imtv[:,:,vframe:vframe+5,:,:] for vframe in range(i,min(lastframe,i+opt.batch_size)) ] + im_in = torch.cat(im_batch,0) + im_out = self.__S__.forward_lipfeat(im_in.cuda()); + im_feat.append(im_out.data.cpu()) + + im_feat = torch.cat(im_feat,0) + + # ========== ========== + # Compute offset + # ========== ========== + + print('Compute time %.3f sec.' % (time.time()-tS)) + + return im_feat + + + def loadParameters(self, path): + loaded_state = torch.load(path, map_location=lambda storage, loc: storage); + + self_state = self.__S__.state_dict(); + + for name, param in loaded_state.items(): + + self_state[name].copy_(param); diff --git a/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_LRS.py b/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_LRS.py new file mode 100644 index 0000000000..eda02b8fbb --- /dev/null +++ b/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_LRS.py @@ -0,0 +1,53 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- + +import time, pdb, argparse, subprocess +import glob +import os +from tqdm import tqdm + +from SyncNetInstance_calc_scores import * + +# ==================== LOAD PARAMS ==================== + + +parser = argparse.ArgumentParser(description = "SyncNet"); + +parser.add_argument('--initial_model', type=str, default="data/syncnet_v2.model", help=''); +parser.add_argument('--batch_size', type=int, default='20', help=''); +parser.add_argument('--vshift', type=int, default='15', help=''); +parser.add_argument('--data_root', type=str, required=True, help=''); +parser.add_argument('--tmp_dir', type=str, default="data/work/pytmp", help=''); +parser.add_argument('--reference', type=str, default="demo", help=''); + +opt = parser.parse_args(); + + +# ==================== RUN EVALUATION ==================== + +s = SyncNetInstance(); + +s.loadParameters(opt.initial_model); +#print("Model %s loaded."%opt.initial_model); +path = os.path.join(opt.data_root, "*.mp4") + +all_videos = glob.glob(path) + +prog_bar = tqdm(range(len(all_videos))) +avg_confidence = 0. +avg_min_distance = 0. + + +for videofile_idx in prog_bar: + videofile = all_videos[videofile_idx] + offset, confidence, min_distance = s.evaluate(opt, videofile=videofile) + avg_confidence += confidence + avg_min_distance += min_distance + prog_bar.set_description('Avg Confidence: {}, Avg Minimum Dist: {}'.format(round(avg_confidence / (videofile_idx + 1), 3), round(avg_min_distance / (videofile_idx + 1), 3))) + prog_bar.refresh() + +print ('Average Confidence: {}'.format(avg_confidence/len(all_videos))) +print ('Average Minimum Distance: {}'.format(avg_min_distance/len(all_videos))) + + + diff --git a/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_real_videos.py b/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_real_videos.py new file mode 100755 index 0000000000..0962258465 --- /dev/null +++ b/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_real_videos.py @@ -0,0 +1,45 @@ +#!/usr/bin/python +#-*- coding: utf-8 -*- + +import time, pdb, argparse, subprocess, pickle, os, gzip, glob + +from SyncNetInstance_calc_scores import * + +# ==================== PARSE ARGUMENT ==================== + +parser = argparse.ArgumentParser(description = "SyncNet"); +parser.add_argument('--initial_model', type=str, default="data/syncnet_v2.model", help=''); +parser.add_argument('--batch_size', type=int, default='20', help=''); +parser.add_argument('--vshift', type=int, default='15', help=''); +parser.add_argument('--data_dir', type=str, default='data/work', help=''); +parser.add_argument('--videofile', type=str, default='', help=''); +parser.add_argument('--reference', type=str, default='', help=''); +opt = parser.parse_args(); + +setattr(opt,'avi_dir',os.path.join(opt.data_dir,'pyavi')) +setattr(opt,'tmp_dir',os.path.join(opt.data_dir,'pytmp')) +setattr(opt,'work_dir',os.path.join(opt.data_dir,'pywork')) +setattr(opt,'crop_dir',os.path.join(opt.data_dir,'pycrop')) + + +# ==================== LOAD MODEL AND FILE LIST ==================== + +s = SyncNetInstance(); + +s.loadParameters(opt.initial_model); +#print("Model %s loaded."%opt.initial_model); + +flist = glob.glob(os.path.join(opt.crop_dir,opt.reference,'0*.avi')) +flist.sort() + +# ==================== GET OFFSETS ==================== + +dists = [] +for idx, fname in enumerate(flist): + offset, conf, dist = s.evaluate(opt,videofile=fname) + print (str(dist)+" "+str(conf)) + +# ==================== PRINT RESULTS TO FILE ==================== + +#with open(os.path.join(opt.work_dir,opt.reference,'activesd.pckl'), 'wb') as fil: +# pickle.dump(dists, fil) diff --git a/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_real_videos.sh b/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_real_videos.sh new file mode 100644 index 0000000000..bd02c95405 --- /dev/null +++ b/models/Wav2Lip/evaluation/scores_LSE/calculate_scores_real_videos.sh @@ -0,0 +1,7 @@ +rm all_scores.txt +yourfilenames=`ls $1` +for eachfile in $yourfilenames +do + python run_pipeline.py --videofile $eachfile --reference wav2lip --data_dir tmp_dir + python calculate_scores_real_videos.py --videofile $eachfile --reference wav2lip --data_dir /ssd_scratch/cvit/prajwalkr_rudra/tmp_dir >> all_scores.txt +done diff --git a/models/Wav2Lip/evaluation/test_filelists/README.md b/models/Wav2Lip/evaluation/test_filelists/README.md new file mode 100644 index 0000000000..f8033e6978 --- /dev/null +++ b/models/Wav2Lip/evaluation/test_filelists/README.md @@ -0,0 +1 @@ +Place LRS2, LRW, and LRS3 (and any other) test set filelists here. \ No newline at end of file diff --git a/models/Wav2Lip/evaluation/test_filelists/ReSyncED/random_pairs.txt b/models/Wav2Lip/evaluation/test_filelists/ReSyncED/random_pairs.txt new file mode 100644 index 0000000000..ffe2c40e11 --- /dev/null +++ b/models/Wav2Lip/evaluation/test_filelists/ReSyncED/random_pairs.txt @@ -0,0 +1,160 @@ +sachin.mp4 emma_cropped.mp4 +sachin.mp4 mourinho.mp4 +sachin.mp4 elon.mp4 +sachin.mp4 messi2.mp4 +sachin.mp4 cr1.mp4 +sachin.mp4 sachin.mp4 +sachin.mp4 sg.mp4 +sachin.mp4 fergi.mp4 +sachin.mp4 spanish_lec1.mp4 +sachin.mp4 bush_small.mp4 +sachin.mp4 macca_cut.mp4 +sachin.mp4 ca_cropped.mp4 +sachin.mp4 lecun.mp4 +sachin.mp4 spanish_lec0.mp4 +srk.mp4 emma_cropped.mp4 +srk.mp4 mourinho.mp4 +srk.mp4 elon.mp4 +srk.mp4 messi2.mp4 +srk.mp4 cr1.mp4 +srk.mp4 srk.mp4 +srk.mp4 sachin.mp4 +srk.mp4 sg.mp4 +srk.mp4 fergi.mp4 +srk.mp4 spanish_lec1.mp4 +srk.mp4 bush_small.mp4 +srk.mp4 macca_cut.mp4 +srk.mp4 ca_cropped.mp4 +srk.mp4 guardiola.mp4 +srk.mp4 lecun.mp4 +srk.mp4 spanish_lec0.mp4 +cr1.mp4 emma_cropped.mp4 +cr1.mp4 elon.mp4 +cr1.mp4 messi2.mp4 +cr1.mp4 cr1.mp4 +cr1.mp4 spanish_lec1.mp4 +cr1.mp4 bush_small.mp4 +cr1.mp4 macca_cut.mp4 +cr1.mp4 ca_cropped.mp4 +cr1.mp4 lecun.mp4 +cr1.mp4 spanish_lec0.mp4 +macca_cut.mp4 emma_cropped.mp4 +macca_cut.mp4 elon.mp4 +macca_cut.mp4 messi2.mp4 +macca_cut.mp4 spanish_lec1.mp4 +macca_cut.mp4 macca_cut.mp4 +macca_cut.mp4 ca_cropped.mp4 +macca_cut.mp4 spanish_lec0.mp4 +lecun.mp4 emma_cropped.mp4 +lecun.mp4 elon.mp4 +lecun.mp4 messi2.mp4 +lecun.mp4 spanish_lec1.mp4 +lecun.mp4 macca_cut.mp4 +lecun.mp4 ca_cropped.mp4 +lecun.mp4 lecun.mp4 +lecun.mp4 spanish_lec0.mp4 +messi2.mp4 emma_cropped.mp4 +messi2.mp4 elon.mp4 +messi2.mp4 messi2.mp4 +messi2.mp4 spanish_lec1.mp4 +messi2.mp4 macca_cut.mp4 +messi2.mp4 ca_cropped.mp4 +messi2.mp4 spanish_lec0.mp4 +ca_cropped.mp4 emma_cropped.mp4 +ca_cropped.mp4 elon.mp4 +ca_cropped.mp4 spanish_lec1.mp4 +ca_cropped.mp4 ca_cropped.mp4 +ca_cropped.mp4 spanish_lec0.mp4 +spanish_lec1.mp4 spanish_lec1.mp4 +spanish_lec1.mp4 spanish_lec0.mp4 +elon.mp4 elon.mp4 +elon.mp4 spanish_lec1.mp4 +elon.mp4 spanish_lec0.mp4 +guardiola.mp4 emma_cropped.mp4 +guardiola.mp4 mourinho.mp4 +guardiola.mp4 elon.mp4 +guardiola.mp4 messi2.mp4 +guardiola.mp4 cr1.mp4 +guardiola.mp4 sachin.mp4 +guardiola.mp4 sg.mp4 +guardiola.mp4 fergi.mp4 +guardiola.mp4 spanish_lec1.mp4 +guardiola.mp4 bush_small.mp4 +guardiola.mp4 macca_cut.mp4 +guardiola.mp4 ca_cropped.mp4 +guardiola.mp4 guardiola.mp4 +guardiola.mp4 lecun.mp4 +guardiola.mp4 spanish_lec0.mp4 +fergi.mp4 emma_cropped.mp4 +fergi.mp4 mourinho.mp4 +fergi.mp4 elon.mp4 +fergi.mp4 messi2.mp4 +fergi.mp4 cr1.mp4 +fergi.mp4 sachin.mp4 +fergi.mp4 sg.mp4 +fergi.mp4 fergi.mp4 +fergi.mp4 spanish_lec1.mp4 +fergi.mp4 bush_small.mp4 +fergi.mp4 macca_cut.mp4 +fergi.mp4 ca_cropped.mp4 +fergi.mp4 lecun.mp4 +fergi.mp4 spanish_lec0.mp4 +spanish.mp4 emma_cropped.mp4 +spanish.mp4 spanish.mp4 +spanish.mp4 mourinho.mp4 +spanish.mp4 elon.mp4 +spanish.mp4 messi2.mp4 +spanish.mp4 cr1.mp4 +spanish.mp4 srk.mp4 +spanish.mp4 sachin.mp4 +spanish.mp4 sg.mp4 +spanish.mp4 fergi.mp4 +spanish.mp4 spanish_lec1.mp4 +spanish.mp4 bush_small.mp4 +spanish.mp4 macca_cut.mp4 +spanish.mp4 ca_cropped.mp4 +spanish.mp4 guardiola.mp4 +spanish.mp4 lecun.mp4 +spanish.mp4 spanish_lec0.mp4 +bush_small.mp4 emma_cropped.mp4 +bush_small.mp4 elon.mp4 +bush_small.mp4 messi2.mp4 +bush_small.mp4 spanish_lec1.mp4 +bush_small.mp4 bush_small.mp4 +bush_small.mp4 macca_cut.mp4 +bush_small.mp4 ca_cropped.mp4 +bush_small.mp4 lecun.mp4 +bush_small.mp4 spanish_lec0.mp4 +emma_cropped.mp4 emma_cropped.mp4 +emma_cropped.mp4 elon.mp4 +emma_cropped.mp4 spanish_lec1.mp4 +emma_cropped.mp4 spanish_lec0.mp4 +sg.mp4 emma_cropped.mp4 +sg.mp4 mourinho.mp4 +sg.mp4 elon.mp4 +sg.mp4 messi2.mp4 +sg.mp4 cr1.mp4 +sg.mp4 sachin.mp4 +sg.mp4 sg.mp4 +sg.mp4 fergi.mp4 +sg.mp4 spanish_lec1.mp4 +sg.mp4 bush_small.mp4 +sg.mp4 macca_cut.mp4 +sg.mp4 ca_cropped.mp4 +sg.mp4 lecun.mp4 +sg.mp4 spanish_lec0.mp4 +spanish_lec0.mp4 spanish_lec0.mp4 +mourinho.mp4 emma_cropped.mp4 +mourinho.mp4 mourinho.mp4 +mourinho.mp4 elon.mp4 +mourinho.mp4 messi2.mp4 +mourinho.mp4 cr1.mp4 +mourinho.mp4 sachin.mp4 +mourinho.mp4 sg.mp4 +mourinho.mp4 fergi.mp4 +mourinho.mp4 spanish_lec1.mp4 +mourinho.mp4 bush_small.mp4 +mourinho.mp4 macca_cut.mp4 +mourinho.mp4 ca_cropped.mp4 +mourinho.mp4 lecun.mp4 +mourinho.mp4 spanish_lec0.mp4 diff --git a/models/Wav2Lip/evaluation/test_filelists/ReSyncED/tts_pairs.txt b/models/Wav2Lip/evaluation/test_filelists/ReSyncED/tts_pairs.txt new file mode 100644 index 0000000000..b7dc1a8c0b --- /dev/null +++ b/models/Wav2Lip/evaluation/test_filelists/ReSyncED/tts_pairs.txt @@ -0,0 +1,18 @@ +adam_1.mp4 andreng_optimization.wav +agad_2.mp4 agad_2.wav +agad_1.mp4 agad_1.wav +agad_3.mp4 agad_3.wav +rms_prop_1.mp4 rms_prop_tts.wav +tf_1.mp4 tf_1.wav +tf_2.mp4 tf_2.wav +andrew_ng_ai_business.mp4 andrewng_business_tts.wav +covid_autopsy_1.mp4 autopsy_tts.wav +news_1.mp4 news_tts.wav +andrew_ng_fund_1.mp4 andrewng_ai_fund.wav +covid_treatments_1.mp4 covid_tts.wav +pytorch_v_tf.mp4 pytorch_vs_tf_eng.wav +pytorch_1.mp4 pytorch.wav +pkb_1.mp4 pkb_1.wav +ss_1.mp4 ss_1.wav +carlsen_1.mp4 carlsen_eng.wav +french.mp4 french.wav \ No newline at end of file diff --git a/models/Wav2Lip/face_detection/README.md b/models/Wav2Lip/face_detection/README.md new file mode 100644 index 0000000000..c073376e4e --- /dev/null +++ b/models/Wav2Lip/face_detection/README.md @@ -0,0 +1 @@ +The code for Face Detection in this folder has been taken from the wonderful [face_alignment](https://github.com/1adrianb/face-alignment) repository. This has been modified to take batches of faces at a time. \ No newline at end of file diff --git a/models/Wav2Lip/face_detection/__init__.py b/models/Wav2Lip/face_detection/__init__.py new file mode 100644 index 0000000000..4bae29fd5f --- /dev/null +++ b/models/Wav2Lip/face_detection/__init__.py @@ -0,0 +1,7 @@ +# -*- coding: utf-8 -*- + +__author__ = """Adrian Bulat""" +__email__ = 'adrian.bulat@nottingham.ac.uk' +__version__ = '1.0.1' + +from .api import FaceAlignment, LandmarksType, NetworkSize diff --git a/models/Wav2Lip/face_detection/api.py b/models/Wav2Lip/face_detection/api.py new file mode 100644 index 0000000000..cb02d5252d --- /dev/null +++ b/models/Wav2Lip/face_detection/api.py @@ -0,0 +1,79 @@ +from __future__ import print_function +import os +import torch +from torch.utils.model_zoo import load_url +from enum import Enum +import numpy as np +import cv2 +try: + import urllib.request as request_file +except BaseException: + import urllib as request_file + +from .models import FAN, ResNetDepth +from .utils import * + + +class LandmarksType(Enum): + """Enum class defining the type of landmarks to detect. + + ``_2D`` - the detected points ``(x,y)`` are detected in a 2D space and follow the visible contour of the face + ``_2halfD`` - this points represent the projection of the 3D points into 3D + ``_3D`` - detect the points ``(x,y,z)``` in a 3D space + + """ + _2D = 1 + _2halfD = 2 + _3D = 3 + + +class NetworkSize(Enum): + # TINY = 1 + # SMALL = 2 + # MEDIUM = 3 + LARGE = 4 + + def __new__(cls, value): + member = object.__new__(cls) + member._value_ = value + return member + + def __int__(self): + return self.value + +ROOT = os.path.dirname(os.path.abspath(__file__)) + +class FaceAlignment: + def __init__(self, landmarks_type, network_size=NetworkSize.LARGE, + device='cuda', flip_input=False, face_detector='sfd', verbose=False): + self.device = device + self.flip_input = flip_input + self.landmarks_type = landmarks_type + self.verbose = verbose + + network_size = int(network_size) + + if 'cuda' in device: + torch.backends.cudnn.benchmark = True + + # Get the face detector + face_detector_module = __import__('face_detection.detection.' + face_detector, + globals(), locals(), [face_detector], 0) + self.face_detector = face_detector_module.FaceDetector(device=device, verbose=verbose) + + def get_detections_for_batch(self, images): + images = images[..., ::-1] + detected_faces = self.face_detector.detect_from_batch(images.copy()) + results = [] + + for i, d in enumerate(detected_faces): + if len(d) == 0: + results.append(None) + continue + d = d[0] + d = np.clip(d, 0, None) + + x1, y1, x2, y2 = map(int, d[:-1]) + results.append((x1, y1, x2, y2)) + + return results \ No newline at end of file diff --git a/models/Wav2Lip/face_detection/detection/__init__.py b/models/Wav2Lip/face_detection/detection/__init__.py new file mode 100644 index 0000000000..1a6b0402da --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/__init__.py @@ -0,0 +1 @@ +from .core import FaceDetector \ No newline at end of file diff --git a/models/Wav2Lip/face_detection/detection/core.py b/models/Wav2Lip/face_detection/detection/core.py new file mode 100644 index 0000000000..0f8275e8e5 --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/core.py @@ -0,0 +1,130 @@ +import logging +import glob +from tqdm import tqdm +import numpy as np +import torch +import cv2 + + +class FaceDetector(object): + """An abstract class representing a face detector. + + Any other face detection implementation must subclass it. All subclasses + must implement ``detect_from_image``, that return a list of detected + bounding boxes. Optionally, for speed considerations detect from path is + recommended. + """ + + def __init__(self, device, verbose): + self.device = device + self.verbose = verbose + + if verbose: + if 'cpu' in device: + logger = logging.getLogger(__name__) + logger.warning("Detection running on CPU, this may be potentially slow.") + + if 'cpu' not in device and 'cuda' not in device: + if verbose: + logger.error("Expected values for device are: {cpu, cuda} but got: %s", device) + raise ValueError + + def detect_from_image(self, tensor_or_path): + """Detects faces in a given image. + + This function detects the faces present in a provided BGR(usually) + image. The input can be either the image itself or the path to it. + + Arguments: + tensor_or_path {numpy.ndarray, torch.tensor or string} -- the path + to an image or the image itself. + + Example:: + + >>> path_to_image = 'data/image_01.jpg' + ... detected_faces = detect_from_image(path_to_image) + [A list of bounding boxes (x1, y1, x2, y2)] + >>> image = cv2.imread(path_to_image) + ... detected_faces = detect_from_image(image) + [A list of bounding boxes (x1, y1, x2, y2)] + + """ + raise NotImplementedError + + def detect_from_directory(self, path, extensions=['.jpg', '.png'], recursive=False, show_progress_bar=True): + """Detects faces from all the images present in a given directory. + + Arguments: + path {string} -- a string containing a path that points to the folder containing the images + + Keyword Arguments: + extensions {list} -- list of string containing the extensions to be + consider in the following format: ``.extension_name`` (default: + {['.jpg', '.png']}) recursive {bool} -- option wherever to scan the + folder recursively (default: {False}) show_progress_bar {bool} -- + display a progressbar (default: {True}) + + Example: + >>> directory = 'data' + ... detected_faces = detect_from_directory(directory) + {A dictionary of [lists containing bounding boxes(x1, y1, x2, y2)]} + + """ + if self.verbose: + logger = logging.getLogger(__name__) + + if len(extensions) == 0: + if self.verbose: + logger.error("Expected at list one extension, but none was received.") + raise ValueError + + if self.verbose: + logger.info("Constructing the list of images.") + additional_pattern = '/**/*' if recursive else '/*' + files = [] + for extension in extensions: + files.extend(glob.glob(path + additional_pattern + extension, recursive=recursive)) + + if self.verbose: + logger.info("Finished searching for images. %s images found", len(files)) + logger.info("Preparing to run the detection.") + + predictions = {} + for image_path in tqdm(files, disable=not show_progress_bar): + if self.verbose: + logger.info("Running the face detector on image: %s", image_path) + predictions[image_path] = self.detect_from_image(image_path) + + if self.verbose: + logger.info("The detector was successfully run on all %s images", len(files)) + + return predictions + + @property + def reference_scale(self): + raise NotImplementedError + + @property + def reference_x_shift(self): + raise NotImplementedError + + @property + def reference_y_shift(self): + raise NotImplementedError + + @staticmethod + def tensor_or_path_to_ndarray(tensor_or_path, rgb=True): + """Convert path (represented as a string) or torch.tensor to a numpy.ndarray + + Arguments: + tensor_or_path {numpy.ndarray, torch.tensor or string} -- path to the image, or the image itself + """ + if isinstance(tensor_or_path, str): + return cv2.imread(tensor_or_path) if not rgb else cv2.imread(tensor_or_path)[..., ::-1] + elif torch.is_tensor(tensor_or_path): + # Call cpu in case its coming from cuda + return tensor_or_path.cpu().numpy()[..., ::-1].copy() if not rgb else tensor_or_path.cpu().numpy() + elif isinstance(tensor_or_path, np.ndarray): + return tensor_or_path[..., ::-1].copy() if not rgb else tensor_or_path + else: + raise TypeError diff --git a/models/Wav2Lip/face_detection/detection/sfd/__init__.py b/models/Wav2Lip/face_detection/detection/sfd/__init__.py new file mode 100644 index 0000000000..5a63ecd456 --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/sfd/__init__.py @@ -0,0 +1 @@ +from .sfd_detector import SFDDetector as FaceDetector \ No newline at end of file diff --git a/models/Wav2Lip/face_detection/detection/sfd/bbox.py b/models/Wav2Lip/face_detection/detection/sfd/bbox.py new file mode 100644 index 0000000000..4bd7222e5e --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/sfd/bbox.py @@ -0,0 +1,129 @@ +from __future__ import print_function +import os +import sys +import cv2 +import random +import datetime +import time +import math +import argparse +import numpy as np +import torch + +try: + from iou import IOU +except BaseException: + # IOU cython speedup 10x + def IOU(ax1, ay1, ax2, ay2, bx1, by1, bx2, by2): + sa = abs((ax2 - ax1) * (ay2 - ay1)) + sb = abs((bx2 - bx1) * (by2 - by1)) + x1, y1 = max(ax1, bx1), max(ay1, by1) + x2, y2 = min(ax2, bx2), min(ay2, by2) + w = x2 - x1 + h = y2 - y1 + if w < 0 or h < 0: + return 0.0 + else: + return 1.0 * w * h / (sa + sb - w * h) + + +def bboxlog(x1, y1, x2, y2, axc, ayc, aww, ahh): + xc, yc, ww, hh = (x2 + x1) / 2, (y2 + y1) / 2, x2 - x1, y2 - y1 + dx, dy = (xc - axc) / aww, (yc - ayc) / ahh + dw, dh = math.log(ww / aww), math.log(hh / ahh) + return dx, dy, dw, dh + + +def bboxloginv(dx, dy, dw, dh, axc, ayc, aww, ahh): + xc, yc = dx * aww + axc, dy * ahh + ayc + ww, hh = math.exp(dw) * aww, math.exp(dh) * ahh + x1, x2, y1, y2 = xc - ww / 2, xc + ww / 2, yc - hh / 2, yc + hh / 2 + return x1, y1, x2, y2 + + +def nms(dets, thresh): + if 0 == len(dets): + return [] + x1, y1, x2, y2, scores = dets[:, 0], dets[:, 1], dets[:, 2], dets[:, 3], dets[:, 4] + areas = (x2 - x1 + 1) * (y2 - y1 + 1) + order = scores.argsort()[::-1] + + keep = [] + while order.size > 0: + i = order[0] + keep.append(i) + xx1, yy1 = np.maximum(x1[i], x1[order[1:]]), np.maximum(y1[i], y1[order[1:]]) + xx2, yy2 = np.minimum(x2[i], x2[order[1:]]), np.minimum(y2[i], y2[order[1:]]) + + w, h = np.maximum(0.0, xx2 - xx1 + 1), np.maximum(0.0, yy2 - yy1 + 1) + ovr = w * h / (areas[i] + areas[order[1:]] - w * h) + + inds = np.where(ovr <= thresh)[0] + order = order[inds + 1] + + return keep + + +def encode(matched, priors, variances): + """Encode the variances from the priorbox layers into the ground truth boxes + we have matched (based on jaccard overlap) with the prior boxes. + Args: + matched: (tensor) Coords of ground truth for each prior in point-form + Shape: [num_priors, 4]. + priors: (tensor) Prior boxes in center-offset form + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + encoded boxes (tensor), Shape: [num_priors, 4] + """ + + # dist b/t match center and prior's center + g_cxcy = (matched[:, :2] + matched[:, 2:]) / 2 - priors[:, :2] + # encode variance + g_cxcy /= (variances[0] * priors[:, 2:]) + # match wh / prior wh + g_wh = (matched[:, 2:] - matched[:, :2]) / priors[:, 2:] + g_wh = torch.log(g_wh) / variances[1] + # return target for smooth_l1_loss + return torch.cat([g_cxcy, g_wh], 1) # [num_priors,4] + + +def decode(loc, priors, variances): + """Decode locations from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + loc (tensor): location predictions for loc layers, + Shape: [num_priors,4] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded bounding box predictions + """ + + boxes = torch.cat(( + priors[:, :2] + loc[:, :2] * variances[0] * priors[:, 2:], + priors[:, 2:] * torch.exp(loc[:, 2:] * variances[1])), 1) + boxes[:, :2] -= boxes[:, 2:] / 2 + boxes[:, 2:] += boxes[:, :2] + return boxes + +def batch_decode(loc, priors, variances): + """Decode locations from predictions using priors to undo + the encoding we did for offset regression at train time. + Args: + loc (tensor): location predictions for loc layers, + Shape: [num_priors,4] + priors (tensor): Prior boxes in center-offset form. + Shape: [num_priors,4]. + variances: (list[float]) Variances of priorboxes + Return: + decoded bounding box predictions + """ + + boxes = torch.cat(( + priors[:, :, :2] + loc[:, :, :2] * variances[0] * priors[:, :, 2:], + priors[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1])), 2) + boxes[:, :, :2] -= boxes[:, :, 2:] / 2 + boxes[:, :, 2:] += boxes[:, :, :2] + return boxes diff --git a/models/Wav2Lip/face_detection/detection/sfd/detect.py b/models/Wav2Lip/face_detection/detection/sfd/detect.py new file mode 100644 index 0000000000..efef6273ad --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/sfd/detect.py @@ -0,0 +1,112 @@ +import torch +import torch.nn.functional as F + +import os +import sys +import cv2 +import random +import datetime +import math +import argparse +import numpy as np + +import scipy.io as sio +import zipfile +from .net_s3fd import s3fd +from .bbox import * + + +def detect(net, img, device): + img = img - np.array([104, 117, 123]) + img = img.transpose(2, 0, 1) + img = img.reshape((1,) + img.shape) + + if 'cuda' in device: + torch.backends.cudnn.benchmark = True + + img = torch.from_numpy(img).float().to(device) + BB, CC, HH, WW = img.size() + with torch.no_grad(): + olist = net(img) + + bboxlist = [] + for i in range(len(olist) // 2): + olist[i * 2] = F.softmax(olist[i * 2], dim=1) + olist = [oelem.data.cpu() for oelem in olist] + for i in range(len(olist) // 2): + ocls, oreg = olist[i * 2], olist[i * 2 + 1] + FB, FC, FH, FW = ocls.size() # feature map size + stride = 2**(i + 2) # 4,8,16,32,64,128 + anchor = stride * 4 + poss = zip(*np.where(ocls[:, 1, :, :] > 0.05)) + for Iindex, hindex, windex in poss: + axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride + score = ocls[0, 1, hindex, windex] + loc = oreg[0, :, hindex, windex].contiguous().view(1, 4) + priors = torch.Tensor([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]]) + variances = [0.1, 0.2] + box = decode(loc, priors, variances) + x1, y1, x2, y2 = box[0] * 1.0 + # cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1) + bboxlist.append([x1, y1, x2, y2, score]) + bboxlist = np.array(bboxlist) + if 0 == len(bboxlist): + bboxlist = np.zeros((1, 5)) + + return bboxlist + +def batch_detect(net, imgs, device): + imgs = imgs - np.array([104, 117, 123]) + imgs = imgs.transpose(0, 3, 1, 2) + + if 'cuda' in device: + torch.backends.cudnn.benchmark = True + + imgs = torch.from_numpy(imgs).float().to(device) + BB, CC, HH, WW = imgs.size() + with torch.no_grad(): + olist = net(imgs) + + bboxlist = [] + for i in range(len(olist) // 2): + olist[i * 2] = F.softmax(olist[i * 2], dim=1) + olist = [oelem.data.cpu() for oelem in olist] + for i in range(len(olist) // 2): + ocls, oreg = olist[i * 2], olist[i * 2 + 1] + FB, FC, FH, FW = ocls.size() # feature map size + stride = 2**(i + 2) # 4,8,16,32,64,128 + anchor = stride * 4 + poss = zip(*np.where(ocls[:, 1, :, :] > 0.05)) + for Iindex, hindex, windex in poss: + axc, ayc = stride / 2 + windex * stride, stride / 2 + hindex * stride + score = ocls[:, 1, hindex, windex] + loc = oreg[:, :, hindex, windex].contiguous().view(BB, 1, 4) + priors = torch.Tensor([[axc / 1.0, ayc / 1.0, stride * 4 / 1.0, stride * 4 / 1.0]]).view(1, 1, 4) + variances = [0.1, 0.2] + box = batch_decode(loc, priors, variances) + box = box[:, 0] * 1.0 + # cv2.rectangle(imgshow,(int(x1),int(y1)),(int(x2),int(y2)),(0,0,255),1) + bboxlist.append(torch.cat([box, score.unsqueeze(1)], 1).cpu().numpy()) + bboxlist = np.array(bboxlist) + if 0 == len(bboxlist): + bboxlist = np.zeros((1, BB, 5)) + + return bboxlist + +def flip_detect(net, img, device): + img = cv2.flip(img, 1) + b = detect(net, img, device) + + bboxlist = np.zeros(b.shape) + bboxlist[:, 0] = img.shape[1] - b[:, 2] + bboxlist[:, 1] = b[:, 1] + bboxlist[:, 2] = img.shape[1] - b[:, 0] + bboxlist[:, 3] = b[:, 3] + bboxlist[:, 4] = b[:, 4] + return bboxlist + + +def pts_to_bb(pts): + min_x, min_y = np.min(pts, axis=0) + max_x, max_y = np.max(pts, axis=0) + return np.array([min_x, min_y, max_x, max_y]) diff --git a/models/Wav2Lip/face_detection/detection/sfd/net_s3fd.py b/models/Wav2Lip/face_detection/detection/sfd/net_s3fd.py new file mode 100644 index 0000000000..fc64313c27 --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/sfd/net_s3fd.py @@ -0,0 +1,129 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F + + +class L2Norm(nn.Module): + def __init__(self, n_channels, scale=1.0): + super(L2Norm, self).__init__() + self.n_channels = n_channels + self.scale = scale + self.eps = 1e-10 + self.weight = nn.Parameter(torch.Tensor(self.n_channels)) + self.weight.data *= 0.0 + self.weight.data += self.scale + + def forward(self, x): + norm = x.pow(2).sum(dim=1, keepdim=True).sqrt() + self.eps + x = x / norm * self.weight.view(1, -1, 1, 1) + return x + + +class s3fd(nn.Module): + def __init__(self): + super(s3fd, self).__init__() + self.conv1_1 = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1) + self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=1) + + self.conv2_1 = nn.Conv2d(64, 128, kernel_size=3, stride=1, padding=1) + self.conv2_2 = nn.Conv2d(128, 128, kernel_size=3, stride=1, padding=1) + + self.conv3_1 = nn.Conv2d(128, 256, kernel_size=3, stride=1, padding=1) + self.conv3_2 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1) + self.conv3_3 = nn.Conv2d(256, 256, kernel_size=3, stride=1, padding=1) + + self.conv4_1 = nn.Conv2d(256, 512, kernel_size=3, stride=1, padding=1) + self.conv4_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) + self.conv4_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) + + self.conv5_1 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) + self.conv5_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) + self.conv5_3 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=1) + + self.fc6 = nn.Conv2d(512, 1024, kernel_size=3, stride=1, padding=3) + self.fc7 = nn.Conv2d(1024, 1024, kernel_size=1, stride=1, padding=0) + + self.conv6_1 = nn.Conv2d(1024, 256, kernel_size=1, stride=1, padding=0) + self.conv6_2 = nn.Conv2d(256, 512, kernel_size=3, stride=2, padding=1) + + self.conv7_1 = nn.Conv2d(512, 128, kernel_size=1, stride=1, padding=0) + self.conv7_2 = nn.Conv2d(128, 256, kernel_size=3, stride=2, padding=1) + + self.conv3_3_norm = L2Norm(256, scale=10) + self.conv4_3_norm = L2Norm(512, scale=8) + self.conv5_3_norm = L2Norm(512, scale=5) + + self.conv3_3_norm_mbox_conf = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1) + self.conv3_3_norm_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1) + self.conv4_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1) + self.conv4_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1) + self.conv5_3_norm_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1) + self.conv5_3_norm_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1) + + self.fc7_mbox_conf = nn.Conv2d(1024, 2, kernel_size=3, stride=1, padding=1) + self.fc7_mbox_loc = nn.Conv2d(1024, 4, kernel_size=3, stride=1, padding=1) + self.conv6_2_mbox_conf = nn.Conv2d(512, 2, kernel_size=3, stride=1, padding=1) + self.conv6_2_mbox_loc = nn.Conv2d(512, 4, kernel_size=3, stride=1, padding=1) + self.conv7_2_mbox_conf = nn.Conv2d(256, 2, kernel_size=3, stride=1, padding=1) + self.conv7_2_mbox_loc = nn.Conv2d(256, 4, kernel_size=3, stride=1, padding=1) + + def forward(self, x): + h = F.relu(self.conv1_1(x)) + h = F.relu(self.conv1_2(h)) + h = F.max_pool2d(h, 2, 2) + + h = F.relu(self.conv2_1(h)) + h = F.relu(self.conv2_2(h)) + h = F.max_pool2d(h, 2, 2) + + h = F.relu(self.conv3_1(h)) + h = F.relu(self.conv3_2(h)) + h = F.relu(self.conv3_3(h)) + f3_3 = h + h = F.max_pool2d(h, 2, 2) + + h = F.relu(self.conv4_1(h)) + h = F.relu(self.conv4_2(h)) + h = F.relu(self.conv4_3(h)) + f4_3 = h + h = F.max_pool2d(h, 2, 2) + + h = F.relu(self.conv5_1(h)) + h = F.relu(self.conv5_2(h)) + h = F.relu(self.conv5_3(h)) + f5_3 = h + h = F.max_pool2d(h, 2, 2) + + h = F.relu(self.fc6(h)) + h = F.relu(self.fc7(h)) + ffc7 = h + h = F.relu(self.conv6_1(h)) + h = F.relu(self.conv6_2(h)) + f6_2 = h + h = F.relu(self.conv7_1(h)) + h = F.relu(self.conv7_2(h)) + f7_2 = h + + f3_3 = self.conv3_3_norm(f3_3) + f4_3 = self.conv4_3_norm(f4_3) + f5_3 = self.conv5_3_norm(f5_3) + + cls1 = self.conv3_3_norm_mbox_conf(f3_3) + reg1 = self.conv3_3_norm_mbox_loc(f3_3) + cls2 = self.conv4_3_norm_mbox_conf(f4_3) + reg2 = self.conv4_3_norm_mbox_loc(f4_3) + cls3 = self.conv5_3_norm_mbox_conf(f5_3) + reg3 = self.conv5_3_norm_mbox_loc(f5_3) + cls4 = self.fc7_mbox_conf(ffc7) + reg4 = self.fc7_mbox_loc(ffc7) + cls5 = self.conv6_2_mbox_conf(f6_2) + reg5 = self.conv6_2_mbox_loc(f6_2) + cls6 = self.conv7_2_mbox_conf(f7_2) + reg6 = self.conv7_2_mbox_loc(f7_2) + + # max-out background label + chunk = torch.chunk(cls1, 4, 1) + bmax = torch.max(torch.max(chunk[0], chunk[1]), chunk[2]) + cls1 = torch.cat([bmax, chunk[3]], dim=1) + + return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6] diff --git a/models/Wav2Lip/face_detection/detection/sfd/sfd_detector.py b/models/Wav2Lip/face_detection/detection/sfd/sfd_detector.py new file mode 100644 index 0000000000..8fbce15253 --- /dev/null +++ b/models/Wav2Lip/face_detection/detection/sfd/sfd_detector.py @@ -0,0 +1,59 @@ +import os +import cv2 +from torch.utils.model_zoo import load_url + +from ..core import FaceDetector + +from .net_s3fd import s3fd +from .bbox import * +from .detect import * + +models_urls = { + 's3fd': 'https://www.adrianbulat.com/downloads/python-fan/s3fd-619a316812.pth', +} + + +class SFDDetector(FaceDetector): + def __init__(self, device, path_to_detector=os.path.join(os.path.dirname(os.path.abspath(__file__)), 's3fd.pth'), verbose=False): + super(SFDDetector, self).__init__(device, verbose) + + # Initialise the face detector + if not os.path.isfile(path_to_detector): + model_weights = load_url(models_urls['s3fd']) + else: + model_weights = torch.load(path_to_detector) + + self.face_detector = s3fd() + self.face_detector.load_state_dict(model_weights) + self.face_detector.to(device) + self.face_detector.eval() + + def detect_from_image(self, tensor_or_path): + image = self.tensor_or_path_to_ndarray(tensor_or_path) + + bboxlist = detect(self.face_detector, image, device=self.device) + keep = nms(bboxlist, 0.3) + bboxlist = bboxlist[keep, :] + bboxlist = [x for x in bboxlist if x[-1] > 0.5] + + return bboxlist + + def detect_from_batch(self, images): + bboxlists = batch_detect(self.face_detector, images, device=self.device) + keeps = [nms(bboxlists[:, i, :], 0.3) for i in range(bboxlists.shape[1])] + bboxlists = [bboxlists[keep, i, :] for i, keep in enumerate(keeps)] + bboxlists = [[x for x in bboxlist if x[-1] > 0.5] for bboxlist in bboxlists] + + return bboxlists + + @property + def reference_scale(self): + return 195 + + @property + def reference_x_shift(self): + return 0 + + @property + def reference_y_shift(self): + return 0 diff --git a/models/Wav2Lip/face_detection/models.py b/models/Wav2Lip/face_detection/models.py new file mode 100644 index 0000000000..ee2dde32bd --- /dev/null +++ b/models/Wav2Lip/face_detection/models.py @@ -0,0 +1,261 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import math + + +def conv3x3(in_planes, out_planes, strd=1, padding=1, bias=False): + "3x3 convolution with padding" + return nn.Conv2d(in_planes, out_planes, kernel_size=3, + stride=strd, padding=padding, bias=bias) + + +class ConvBlock(nn.Module): + def __init__(self, in_planes, out_planes): + super(ConvBlock, self).__init__() + self.bn1 = nn.BatchNorm2d(in_planes) + self.conv1 = conv3x3(in_planes, int(out_planes / 2)) + self.bn2 = nn.BatchNorm2d(int(out_planes / 2)) + self.conv2 = conv3x3(int(out_planes / 2), int(out_planes / 4)) + self.bn3 = nn.BatchNorm2d(int(out_planes / 4)) + self.conv3 = conv3x3(int(out_planes / 4), int(out_planes / 4)) + + if in_planes != out_planes: + self.downsample = nn.Sequential( + nn.BatchNorm2d(in_planes), + nn.ReLU(True), + nn.Conv2d(in_planes, out_planes, + kernel_size=1, stride=1, bias=False), + ) + else: + self.downsample = None + + def forward(self, x): + residual = x + + out1 = self.bn1(x) + out1 = F.relu(out1, True) + out1 = self.conv1(out1) + + out2 = self.bn2(out1) + out2 = F.relu(out2, True) + out2 = self.conv2(out2) + + out3 = self.bn3(out2) + out3 = F.relu(out3, True) + out3 = self.conv3(out3) + + out3 = torch.cat((out1, out2, out3), 1) + + if self.downsample is not None: + residual = self.downsample(residual) + + out3 += residual + + return out3 + + +class Bottleneck(nn.Module): + + expansion = 4 + + def __init__(self, inplanes, planes, stride=1, downsample=None): + super(Bottleneck, self).__init__() + self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False) + self.bn1 = nn.BatchNorm2d(planes) + self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride, + padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(planes) + self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False) + self.bn3 = nn.BatchNorm2d(planes * 4) + self.relu = nn.ReLU(inplace=True) + self.downsample = downsample + self.stride = stride + + def forward(self, x): + residual = x + + out = self.conv1(x) + out = self.bn1(out) + out = self.relu(out) + + out = self.conv2(out) + out = self.bn2(out) + out = self.relu(out) + + out = self.conv3(out) + out = self.bn3(out) + + if self.downsample is not None: + residual = self.downsample(x) + + out += residual + out = self.relu(out) + + return out + + +class HourGlass(nn.Module): + def __init__(self, num_modules, depth, num_features): + super(HourGlass, self).__init__() + self.num_modules = num_modules + self.depth = depth + self.features = num_features + + self._generate_network(self.depth) + + def _generate_network(self, level): + self.add_module('b1_' + str(level), ConvBlock(self.features, self.features)) + + self.add_module('b2_' + str(level), ConvBlock(self.features, self.features)) + + if level > 1: + self._generate_network(level - 1) + else: + self.add_module('b2_plus_' + str(level), ConvBlock(self.features, self.features)) + + self.add_module('b3_' + str(level), ConvBlock(self.features, self.features)) + + def _forward(self, level, inp): + # Upper branch + up1 = inp + up1 = self._modules['b1_' + str(level)](up1) + + # Lower branch + low1 = F.avg_pool2d(inp, 2, stride=2) + low1 = self._modules['b2_' + str(level)](low1) + + if level > 1: + low2 = self._forward(level - 1, low1) + else: + low2 = low1 + low2 = self._modules['b2_plus_' + str(level)](low2) + + low3 = low2 + low3 = self._modules['b3_' + str(level)](low3) + + up2 = F.interpolate(low3, scale_factor=2, mode='nearest') + + return up1 + up2 + + def forward(self, x): + return self._forward(self.depth, x) + + +class FAN(nn.Module): + + def __init__(self, num_modules=1): + super(FAN, self).__init__() + self.num_modules = num_modules + + # Base part + self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3) + self.bn1 = nn.BatchNorm2d(64) + self.conv2 = ConvBlock(64, 128) + self.conv3 = ConvBlock(128, 128) + self.conv4 = ConvBlock(128, 256) + + # Stacking part + for hg_module in range(self.num_modules): + self.add_module('m' + str(hg_module), HourGlass(1, 4, 256)) + self.add_module('top_m_' + str(hg_module), ConvBlock(256, 256)) + self.add_module('conv_last' + str(hg_module), + nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module('bn_end' + str(hg_module), nn.BatchNorm2d(256)) + self.add_module('l' + str(hg_module), nn.Conv2d(256, + 68, kernel_size=1, stride=1, padding=0)) + + if hg_module < self.num_modules - 1: + self.add_module( + 'bl' + str(hg_module), nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0)) + self.add_module('al' + str(hg_module), nn.Conv2d(68, + 256, kernel_size=1, stride=1, padding=0)) + + def forward(self, x): + x = F.relu(self.bn1(self.conv1(x)), True) + x = F.avg_pool2d(self.conv2(x), 2, stride=2) + x = self.conv3(x) + x = self.conv4(x) + + previous = x + + outputs = [] + for i in range(self.num_modules): + hg = self._modules['m' + str(i)](previous) + + ll = hg + ll = self._modules['top_m_' + str(i)](ll) + + ll = F.relu(self._modules['bn_end' + str(i)] + (self._modules['conv_last' + str(i)](ll)), True) + + # Predict heatmaps + tmp_out = self._modules['l' + str(i)](ll) + outputs.append(tmp_out) + + if i < self.num_modules - 1: + ll = self._modules['bl' + str(i)](ll) + tmp_out_ = self._modules['al' + str(i)](tmp_out) + previous = previous + ll + tmp_out_ + + return outputs + + +class ResNetDepth(nn.Module): + + def __init__(self, block=Bottleneck, layers=[3, 8, 36, 3], num_classes=68): + self.inplanes = 64 + super(ResNetDepth, self).__init__() + self.conv1 = nn.Conv2d(3 + 68, 64, kernel_size=7, stride=2, padding=3, + bias=False) + self.bn1 = nn.BatchNorm2d(64) + self.relu = nn.ReLU(inplace=True) + self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1) + self.layer1 = self._make_layer(block, 64, layers[0]) + self.layer2 = self._make_layer(block, 128, layers[1], stride=2) + self.layer3 = self._make_layer(block, 256, layers[2], stride=2) + self.layer4 = self._make_layer(block, 512, layers[3], stride=2) + self.avgpool = nn.AvgPool2d(7) + self.fc = nn.Linear(512 * block.expansion, num_classes) + + for m in self.modules(): + if isinstance(m, nn.Conv2d): + n = m.kernel_size[0] * m.kernel_size[1] * m.out_channels + m.weight.data.normal_(0, math.sqrt(2. / n)) + elif isinstance(m, nn.BatchNorm2d): + m.weight.data.fill_(1) + m.bias.data.zero_() + + def _make_layer(self, block, planes, blocks, stride=1): + downsample = None + if stride != 1 or self.inplanes != planes * block.expansion: + downsample = nn.Sequential( + nn.Conv2d(self.inplanes, planes * block.expansion, + kernel_size=1, stride=stride, bias=False), + nn.BatchNorm2d(planes * block.expansion), + ) + + layers = [] + layers.append(block(self.inplanes, planes, stride, downsample)) + self.inplanes = planes * block.expansion + for i in range(1, blocks): + layers.append(block(self.inplanes, planes)) + + return nn.Sequential(*layers) + + def forward(self, x): + x = self.conv1(x) + x = self.bn1(x) + x = self.relu(x) + x = self.maxpool(x) + + x = self.layer1(x) + x = self.layer2(x) + x = self.layer3(x) + x = self.layer4(x) + + x = self.avgpool(x) + x = x.view(x.size(0), -1) + x = self.fc(x) + + return x diff --git a/models/Wav2Lip/face_detection/utils.py b/models/Wav2Lip/face_detection/utils.py new file mode 100644 index 0000000000..3dc4cf3e32 --- /dev/null +++ b/models/Wav2Lip/face_detection/utils.py @@ -0,0 +1,313 @@ +from __future__ import print_function +import os +import sys +import time +import torch +import math +import numpy as np +import cv2 + + +def _gaussian( + size=3, sigma=0.25, amplitude=1, normalize=False, width=None, + height=None, sigma_horz=None, sigma_vert=None, mean_horz=0.5, + mean_vert=0.5): + # handle some defaults + if width is None: + width = size + if height is None: + height = size + if sigma_horz is None: + sigma_horz = sigma + if sigma_vert is None: + sigma_vert = sigma + center_x = mean_horz * width + 0.5 + center_y = mean_vert * height + 0.5 + gauss = np.empty((height, width), dtype=np.float32) + # generate kernel + for i in range(height): + for j in range(width): + gauss[i][j] = amplitude * math.exp(-(math.pow((j + 1 - center_x) / ( + sigma_horz * width), 2) / 2.0 + math.pow((i + 1 - center_y) / (sigma_vert * height), 2) / 2.0)) + if normalize: + gauss = gauss / np.sum(gauss) + return gauss + + +def draw_gaussian(image, point, sigma): + # Check if the gaussian is inside + ul = [math.floor(point[0] - 3 * sigma), math.floor(point[1] - 3 * sigma)] + br = [math.floor(point[0] + 3 * sigma), math.floor(point[1] + 3 * sigma)] + if (ul[0] > image.shape[1] or ul[1] > image.shape[0] or br[0] < 1 or br[1] < 1): + return image + size = 6 * sigma + 1 + g = _gaussian(size) + g_x = [int(max(1, -ul[0])), int(min(br[0], image.shape[1])) - int(max(1, ul[0])) + int(max(1, -ul[0]))] + g_y = [int(max(1, -ul[1])), int(min(br[1], image.shape[0])) - int(max(1, ul[1])) + int(max(1, -ul[1]))] + img_x = [int(max(1, ul[0])), int(min(br[0], image.shape[1]))] + img_y = [int(max(1, ul[1])), int(min(br[1], image.shape[0]))] + assert (g_x[0] > 0 and g_y[1] > 0) + image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1] + ] = image[img_y[0] - 1:img_y[1], img_x[0] - 1:img_x[1]] + g[g_y[0] - 1:g_y[1], g_x[0] - 1:g_x[1]] + image[image > 1] = 1 + return image + + +def transform(point, center, scale, resolution, invert=False): + """Generate and affine transformation matrix. + + Given a set of points, a center, a scale and a targer resolution, the + function generates and affine transformation matrix. If invert is ``True`` + it will produce the inverse transformation. + + Arguments: + point {torch.tensor} -- the input 2D point + center {torch.tensor or numpy.array} -- the center around which to perform the transformations + scale {float} -- the scale of the face/object + resolution {float} -- the output resolution + + Keyword Arguments: + invert {bool} -- define wherever the function should produce the direct or the + inverse transformation matrix (default: {False}) + """ + _pt = torch.ones(3) + _pt[0] = point[0] + _pt[1] = point[1] + + h = 200.0 * scale + t = torch.eye(3) + t[0, 0] = resolution / h + t[1, 1] = resolution / h + t[0, 2] = resolution * (-center[0] / h + 0.5) + t[1, 2] = resolution * (-center[1] / h + 0.5) + + if invert: + t = torch.inverse(t) + + new_point = (torch.matmul(t, _pt))[0:2] + + return new_point.int() + + +def crop(image, center, scale, resolution=256.0): + """Center crops an image or set of heatmaps + + Arguments: + image {numpy.array} -- an rgb image + center {numpy.array} -- the center of the object, usually the same as of the bounding box + scale {float} -- scale of the face + + Keyword Arguments: + resolution {float} -- the size of the output cropped image (default: {256.0}) + + Returns: + [type] -- [description] + """ # Crop around the center point + """ Crops the image around the center. Input is expected to be an np.ndarray """ + ul = transform([1, 1], center, scale, resolution, True) + br = transform([resolution, resolution], center, scale, resolution, True) + # pad = math.ceil(torch.norm((ul - br).float()) / 2.0 - (br[0] - ul[0]) / 2.0) + if image.ndim > 2: + newDim = np.array([br[1] - ul[1], br[0] - ul[0], + image.shape[2]], dtype=np.int32) + newImg = np.zeros(newDim, dtype=np.uint8) + else: + newDim = np.array([br[1] - ul[1], br[0] - ul[0]], dtype=np.int) + newImg = np.zeros(newDim, dtype=np.uint8) + ht = image.shape[0] + wd = image.shape[1] + newX = np.array( + [max(1, -ul[0] + 1), min(br[0], wd) - ul[0]], dtype=np.int32) + newY = np.array( + [max(1, -ul[1] + 1), min(br[1], ht) - ul[1]], dtype=np.int32) + oldX = np.array([max(1, ul[0] + 1), min(br[0], wd)], dtype=np.int32) + oldY = np.array([max(1, ul[1] + 1), min(br[1], ht)], dtype=np.int32) + newImg[newY[0] - 1:newY[1], newX[0] - 1:newX[1] + ] = image[oldY[0] - 1:oldY[1], oldX[0] - 1:oldX[1], :] + newImg = cv2.resize(newImg, dsize=(int(resolution), int(resolution)), + interpolation=cv2.INTER_LINEAR) + return newImg + + +def get_preds_fromhm(hm, center=None, scale=None): + """Obtain (x,y) coordinates given a set of N heatmaps. If the center + and the scale is provided the function will return the points also in + the original coordinate frame. + + Arguments: + hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H] + + Keyword Arguments: + center {torch.tensor} -- the center of the bounding box (default: {None}) + scale {float} -- face scale (default: {None}) + """ + max, idx = torch.max( + hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2) + idx += 1 + preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float() + preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1) + preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1) + + for i in range(preds.size(0)): + for j in range(preds.size(1)): + hm_ = hm[i, j, :] + pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1 + if pX > 0 and pX < 63 and pY > 0 and pY < 63: + diff = torch.FloatTensor( + [hm_[pY, pX + 1] - hm_[pY, pX - 1], + hm_[pY + 1, pX] - hm_[pY - 1, pX]]) + preds[i, j].add_(diff.sign_().mul_(.25)) + + preds.add_(-.5) + + preds_orig = torch.zeros(preds.size()) + if center is not None and scale is not None: + for i in range(hm.size(0)): + for j in range(hm.size(1)): + preds_orig[i, j] = transform( + preds[i, j], center, scale, hm.size(2), True) + + return preds, preds_orig + +def get_preds_fromhm_batch(hm, centers=None, scales=None): + """Obtain (x,y) coordinates given a set of N heatmaps. If the centers + and the scales is provided the function will return the points also in + the original coordinate frame. + + Arguments: + hm {torch.tensor} -- the predicted heatmaps, of shape [B, N, W, H] + + Keyword Arguments: + centers {torch.tensor} -- the centers of the bounding box (default: {None}) + scales {float} -- face scales (default: {None}) + """ + max, idx = torch.max( + hm.view(hm.size(0), hm.size(1), hm.size(2) * hm.size(3)), 2) + idx += 1 + preds = idx.view(idx.size(0), idx.size(1), 1).repeat(1, 1, 2).float() + preds[..., 0].apply_(lambda x: (x - 1) % hm.size(3) + 1) + preds[..., 1].add_(-1).div_(hm.size(2)).floor_().add_(1) + + for i in range(preds.size(0)): + for j in range(preds.size(1)): + hm_ = hm[i, j, :] + pX, pY = int(preds[i, j, 0]) - 1, int(preds[i, j, 1]) - 1 + if pX > 0 and pX < 63 and pY > 0 and pY < 63: + diff = torch.FloatTensor( + [hm_[pY, pX + 1] - hm_[pY, pX - 1], + hm_[pY + 1, pX] - hm_[pY - 1, pX]]) + preds[i, j].add_(diff.sign_().mul_(.25)) + + preds.add_(-.5) + + preds_orig = torch.zeros(preds.size()) + if centers is not None and scales is not None: + for i in range(hm.size(0)): + for j in range(hm.size(1)): + preds_orig[i, j] = transform( + preds[i, j], centers[i], scales[i], hm.size(2), True) + + return preds, preds_orig + +def shuffle_lr(parts, pairs=None): + """Shuffle the points left-right according to the axis of symmetry + of the object. + + Arguments: + parts {torch.tensor} -- a 3D or 4D object containing the + heatmaps. + + Keyword Arguments: + pairs {list of integers} -- [order of the flipped points] (default: {None}) + """ + if pairs is None: + pairs = [16, 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0, + 26, 25, 24, 23, 22, 21, 20, 19, 18, 17, 27, 28, 29, 30, 35, + 34, 33, 32, 31, 45, 44, 43, 42, 47, 46, 39, 38, 37, 36, 41, + 40, 54, 53, 52, 51, 50, 49, 48, 59, 58, 57, 56, 55, 64, 63, + 62, 61, 60, 67, 66, 65] + if parts.ndimension() == 3: + parts = parts[pairs, ...] + else: + parts = parts[:, pairs, ...] + + return parts + + +def flip(tensor, is_label=False): + """Flip an image or a set of heatmaps left-right + + Arguments: + tensor {numpy.array or torch.tensor} -- [the input image or heatmaps] + + Keyword Arguments: + is_label {bool} -- [denote wherever the input is an image or a set of heatmaps ] (default: {False}) + """ + if not torch.is_tensor(tensor): + tensor = torch.from_numpy(tensor) + + if is_label: + tensor = shuffle_lr(tensor).flip(tensor.ndimension() - 1) + else: + tensor = tensor.flip(tensor.ndimension() - 1) + + return tensor + +# From pyzolib/paths.py (https://bitbucket.org/pyzo/pyzolib/src/tip/paths.py) + + +def appdata_dir(appname=None, roaming=False): + """ appdata_dir(appname=None, roaming=False) + + Get the path to the application directory, where applications are allowed + to write user specific files (e.g. configurations). For non-user specific + data, consider using common_appdata_dir(). + If appname is given, a subdir is appended (and created if necessary). + If roaming is True, will prefer a roaming directory (Windows Vista/7). + """ + + # Define default user directory + userDir = os.getenv('FACEALIGNMENT_USERDIR', None) + if userDir is None: + userDir = os.path.expanduser('~') + if not os.path.isdir(userDir): # pragma: no cover + userDir = '/var/tmp' # issue #54 + + # Get system app data dir + path = None + if sys.platform.startswith('win'): + path1, path2 = os.getenv('LOCALAPPDATA'), os.getenv('APPDATA') + path = (path2 or path1) if roaming else (path1 or path2) + elif sys.platform.startswith('darwin'): + path = os.path.join(userDir, 'Library', 'Application Support') + # On Linux and as fallback + if not (path and os.path.isdir(path)): + path = userDir + + # Maybe we should store things local to the executable (in case of a + # portable distro or a frozen application that wants to be portable) + prefix = sys.prefix + if getattr(sys, 'frozen', None): + prefix = os.path.abspath(os.path.dirname(sys.executable)) + for reldir in ('settings', '../settings'): + localpath = os.path.abspath(os.path.join(prefix, reldir)) + if os.path.isdir(localpath): # pragma: no cover + try: + open(os.path.join(localpath, 'test.write'), 'wb').close() + os.remove(os.path.join(localpath, 'test.write')) + except IOError: + pass # We cannot write in this directory + else: + path = localpath + break + + # Get path specific for this app + if appname: + if path == userDir: + appname = '.' + appname.lstrip('.') # Make it a hidden directory + path = os.path.join(path, appname) + if not os.path.isdir(path): # pragma: no cover + os.mkdir(path) + + # Done + return path diff --git a/models/Wav2Lip/filelists/README.md b/models/Wav2Lip/filelists/README.md new file mode 100644 index 0000000000..e7d7e7bb3b --- /dev/null +++ b/models/Wav2Lip/filelists/README.md @@ -0,0 +1 @@ +Place LRS2 (and any other) filelists here for training. \ No newline at end of file diff --git a/models/Wav2Lip/hparams.py b/models/Wav2Lip/hparams.py new file mode 100644 index 0000000000..cc2a0d39c0 --- /dev/null +++ b/models/Wav2Lip/hparams.py @@ -0,0 +1,86 @@ +from tensorflow.contrib.training import HParams +from glob import glob +import os + +def get_image_list(data_root, split): + filelist = [] + + with open('filelists/{}.txt'.format(split)) as f: + for line in f: + line = line.strip() + if ' ' in line: line = line.split()[0] + filelist.append(os.path.join(data_root, line)) + + return filelist + +# Default hyperparameters +hparams = HParams( + num_mels=80, # Number of mel-spectrogram channels and local conditioning dimensionality + # network + rescale=True, # Whether to rescale audio prior to preprocessing + rescaling_max=0.9, # Rescaling value + + # Use LWS (https://github.com/Jonathan-LeRoux/lws) for STFT and phase reconstruction + # It"s preferred to set True to use with https://github.com/r9y9/wavenet_vocoder + # Does not work if n_ffit is not multiple of hop_size!! + use_lws=False, + + n_fft=800, # Extra window size is filled with 0 paddings to match this parameter + hop_size=200, # For 16000Hz, 200 = 12.5 ms (0.0125 * sample_rate) + win_size=800, # For 16000Hz, 800 = 50 ms (If None, win_size = n_fft) (0.05 * sample_rate) + sample_rate=16000, # 16000Hz (corresponding to librispeech) (sox --i ) + + frame_shift_ms=None, # Can replace hop_size parameter. (Recommended: 12.5) + + # Mel and Linear spectrograms normalization/scaling and clipping + signal_normalization=True, + # Whether to normalize mel spectrograms to some predefined range (following below parameters) + allow_clipping_in_normalization=True, # Only relevant if mel_normalization = True + symmetric_mels=True, + # Whether to scale the data to be symmetric around 0. (Also multiplies the output range by 2, + # faster and cleaner convergence) + max_abs_value=4., + # max absolute value of data. If symmetric, data will be [-max, max] else [0, max] (Must not + # be too big to avoid gradient explosion, + # not too small for fast convergence) + # Contribution by @begeekmyfriend + # Spectrogram Pre-Emphasis (Lfilter: Reduce spectrogram noise and helps model certitude + # levels. Also allows for better G&L phase reconstruction) + preemphasize=True, # whether to apply filter + preemphasis=0.97, # filter coefficient. + + # Limits + min_level_db=-100, + ref_level_db=20, + fmin=55, + # Set this to 55 if your speaker is male! if female, 95 should help taking off noise. (To + # test depending on dataset. Pitch info: male~[65, 260], female~[100, 525]) + fmax=7600, # To be increased/reduced depending on data. + + ###################### Our training parameters ################################# + img_size=96, + fps=25, + + batch_size=16, + initial_learning_rate=1e-4, + nepochs=200000000000000000, ### ctrl + c, stop whenever eval loss is consistently greater than train loss for ~10 epochs + num_workers=16, + checkpoint_interval=3000, + eval_interval=3000, + save_optimizer_state=True, + + syncnet_wt=0.0, # is initially zero, will be set automatically to 0.03 later. Leads to faster convergence. + syncnet_batch_size=64, + syncnet_lr=1e-4, + syncnet_eval_interval=10000, + syncnet_checkpoint_interval=10000, + + disc_wt=0.07, + disc_initial_learning_rate=1e-4, +) + + +def hparams_debug_string(): + values = hparams.values() + hp = [" %s: %s" % (name, values[name]) for name in sorted(values) if name != "sentences"] + return "Hyperparameters:\n" + "\n".join(hp) diff --git a/models/Wav2Lip/hq_wav2lip_train.py b/models/Wav2Lip/hq_wav2lip_train.py new file mode 100644 index 0000000000..c384ad9d80 --- /dev/null +++ b/models/Wav2Lip/hq_wav2lip_train.py @@ -0,0 +1,443 @@ +from os.path import dirname, join, basename, isfile +from tqdm import tqdm + +from models import SyncNet_color as SyncNet +from models import Wav2Lip, Wav2Lip_disc_qual +import audio + +import torch +from torch import nn +from torch.nn import functional as F +from torch import optim +import torch.backends.cudnn as cudnn +from torch.utils import data as data_utils +import numpy as np + +from glob import glob + +import os, random, cv2, argparse +from hparams import hparams, get_image_list + +parser = argparse.ArgumentParser(description='Code to train the Wav2Lip model WITH the visual quality discriminator') + +parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True, type=str) + +parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str) +parser.add_argument('--syncnet_checkpoint_path', help='Load the pre-trained Expert discriminator', required=True, type=str) + +parser.add_argument('--checkpoint_path', help='Resume generator from this checkpoint', default=None, type=str) +parser.add_argument('--disc_checkpoint_path', help='Resume quality disc from this checkpoint', default=None, type=str) + +args = parser.parse_args() + + +global_step = 0 +global_epoch = 0 +use_cuda = torch.cuda.is_available() +print('use_cuda: {}'.format(use_cuda)) + +syncnet_T = 5 +syncnet_mel_step_size = 16 + +class Dataset(object): + def __init__(self, split): + self.all_videos = get_image_list(args.data_root, split) + + def get_frame_id(self, frame): + return int(basename(frame).split('.')[0]) + + def get_window(self, start_frame): + start_id = self.get_frame_id(start_frame) + vidname = dirname(start_frame) + + window_fnames = [] + for frame_id in range(start_id, start_id + syncnet_T): + frame = join(vidname, '{}.jpg'.format(frame_id)) + if not isfile(frame): + return None + window_fnames.append(frame) + return window_fnames + + def read_window(self, window_fnames): + if window_fnames is None: return None + window = [] + for fname in window_fnames: + img = cv2.imread(fname) + if img is None: + return None + try: + img = cv2.resize(img, (hparams.img_size, hparams.img_size)) + except Exception as e: + return None + + window.append(img) + + return window + + def crop_audio_window(self, spec, start_frame): + if type(start_frame) == int: + start_frame_num = start_frame + else: + start_frame_num = self.get_frame_id(start_frame) + start_idx = int(80. * (start_frame_num / float(hparams.fps))) + + end_idx = start_idx + syncnet_mel_step_size + + return spec[start_idx : end_idx, :] + + def get_segmented_mels(self, spec, start_frame): + mels = [] + assert syncnet_T == 5 + start_frame_num = self.get_frame_id(start_frame) + 1 # 0-indexing ---> 1-indexing + if start_frame_num - 2 < 0: return None + for i in range(start_frame_num, start_frame_num + syncnet_T): + m = self.crop_audio_window(spec, i - 2) + if m.shape[0] != syncnet_mel_step_size: + return None + mels.append(m.T) + + mels = np.asarray(mels) + + return mels + + def prepare_window(self, window): + # 3 x T x H x W + x = np.asarray(window) / 255. + x = np.transpose(x, (3, 0, 1, 2)) + + return x + + def __len__(self): + return len(self.all_videos) + + def __getitem__(self, idx): + while 1: + idx = random.randint(0, len(self.all_videos) - 1) + vidname = self.all_videos[idx] + img_names = list(glob(join(vidname, '*.jpg'))) + if len(img_names) <= 3 * syncnet_T: + continue + + img_name = random.choice(img_names) + wrong_img_name = random.choice(img_names) + while wrong_img_name == img_name: + wrong_img_name = random.choice(img_names) + + window_fnames = self.get_window(img_name) + wrong_window_fnames = self.get_window(wrong_img_name) + if window_fnames is None or wrong_window_fnames is None: + continue + + window = self.read_window(window_fnames) + if window is None: + continue + + wrong_window = self.read_window(wrong_window_fnames) + if wrong_window is None: + continue + + try: + wavpath = join(vidname, "audio.wav") + wav = audio.load_wav(wavpath, hparams.sample_rate) + + orig_mel = audio.melspectrogram(wav).T + except Exception as e: + continue + + mel = self.crop_audio_window(orig_mel.copy(), img_name) + + if (mel.shape[0] != syncnet_mel_step_size): + continue + + indiv_mels = self.get_segmented_mels(orig_mel.copy(), img_name) + if indiv_mels is None: continue + + window = self.prepare_window(window) + y = window.copy() + window[:, :, window.shape[2]//2:] = 0. + + wrong_window = self.prepare_window(wrong_window) + x = np.concatenate([window, wrong_window], axis=0) + + x = torch.FloatTensor(x) + mel = torch.FloatTensor(mel.T).unsqueeze(0) + indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1) + y = torch.FloatTensor(y) + return x, indiv_mels, mel, y + +def save_sample_images(x, g, gt, global_step, checkpoint_dir): + x = (x.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8) + g = (g.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8) + gt = (gt.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8) + + refs, inps = x[..., 3:], x[..., :3] + folder = join(checkpoint_dir, "samples_step{:09d}".format(global_step)) + if not os.path.exists(folder): os.mkdir(folder) + collage = np.concatenate((refs, inps, g, gt), axis=-2) + for batch_idx, c in enumerate(collage): + for t in range(len(c)): + cv2.imwrite('{}/{}_{}.jpg'.format(folder, batch_idx, t), c[t]) + +logloss = nn.BCELoss() +def cosine_loss(a, v, y): + d = nn.functional.cosine_similarity(a, v) + loss = logloss(d.unsqueeze(1), y) + + return loss + +device = torch.device("cuda" if use_cuda else "cpu") +syncnet = SyncNet().to(device) +for p in syncnet.parameters(): + p.requires_grad = False + +recon_loss = nn.L1Loss() +def get_sync_loss(mel, g): + g = g[:, :, :, g.size(3)//2:] + g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1) + # B, 3 * T, H//2, W + a, v = syncnet(mel, g) + y = torch.ones(g.size(0), 1).float().to(device) + return cosine_loss(a, v, y) + +def train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer, + checkpoint_dir=None, checkpoint_interval=None, nepochs=None): + global global_step, global_epoch + resumed_step = global_step + + while global_epoch < nepochs: + print('Starting Epoch: {}'.format(global_epoch)) + running_sync_loss, running_l1_loss, disc_loss, running_perceptual_loss = 0., 0., 0., 0. + running_disc_real_loss, running_disc_fake_loss = 0., 0. + prog_bar = tqdm(enumerate(train_data_loader)) + for step, (x, indiv_mels, mel, gt) in prog_bar: + disc.train() + model.train() + + x = x.to(device) + mel = mel.to(device) + indiv_mels = indiv_mels.to(device) + gt = gt.to(device) + + ### Train generator now. Remove ALL grads. + optimizer.zero_grad() + disc_optimizer.zero_grad() + + g = model(indiv_mels, x) + + if hparams.syncnet_wt > 0.: + sync_loss = get_sync_loss(mel, g) + else: + sync_loss = 0. + + if hparams.disc_wt > 0.: + perceptual_loss = disc.perceptual_forward(g) + else: + perceptual_loss = 0. + + l1loss = recon_loss(g, gt) + + loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \ + (1. - hparams.syncnet_wt - hparams.disc_wt) * l1loss + + loss.backward() + optimizer.step() + + ### Remove all gradients before Training disc + disc_optimizer.zero_grad() + + pred = disc(gt) + disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device)) + disc_real_loss.backward() + + pred = disc(g.detach()) + disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device)) + disc_fake_loss.backward() + + disc_optimizer.step() + + running_disc_real_loss += disc_real_loss.item() + running_disc_fake_loss += disc_fake_loss.item() + + if global_step % checkpoint_interval == 0: + save_sample_images(x, g, gt, global_step, checkpoint_dir) + + # Logs + global_step += 1 + cur_session_steps = global_step - resumed_step + + running_l1_loss += l1loss.item() + if hparams.syncnet_wt > 0.: + running_sync_loss += sync_loss.item() + else: + running_sync_loss += 0. + + if hparams.disc_wt > 0.: + running_perceptual_loss += perceptual_loss.item() + else: + running_perceptual_loss += 0. + + if global_step == 1 or global_step % checkpoint_interval == 0: + save_checkpoint( + model, optimizer, global_step, checkpoint_dir, global_epoch) + save_checkpoint(disc, disc_optimizer, global_step, checkpoint_dir, global_epoch, prefix='disc_') + + + if global_step % hparams.eval_interval == 0: + with torch.no_grad(): + average_sync_loss = eval_model(test_data_loader, global_step, device, model, disc) + + if average_sync_loss < .75: + hparams.set_hparam('syncnet_wt', 0.03) + + prog_bar.set_description('L1: {}, Sync: {}, Percep: {} | Fake: {}, Real: {}'.format(running_l1_loss / (step + 1), + running_sync_loss / (step + 1), + running_perceptual_loss / (step + 1), + running_disc_fake_loss / (step + 1), + running_disc_real_loss / (step + 1))) + + global_epoch += 1 + +def eval_model(test_data_loader, global_step, device, model, disc): + eval_steps = 300 + print('Evaluating for {} steps'.format(eval_steps)) + running_sync_loss, running_l1_loss, running_disc_real_loss, running_disc_fake_loss, running_perceptual_loss = [], [], [], [], [] + while 1: + for step, (x, indiv_mels, mel, gt) in enumerate((test_data_loader)): + model.eval() + disc.eval() + + x = x.to(device) + mel = mel.to(device) + indiv_mels = indiv_mels.to(device) + gt = gt.to(device) + + pred = disc(gt) + disc_real_loss = F.binary_cross_entropy(pred, torch.ones((len(pred), 1)).to(device)) + + g = model(indiv_mels, x) + pred = disc(g) + disc_fake_loss = F.binary_cross_entropy(pred, torch.zeros((len(pred), 1)).to(device)) + + running_disc_real_loss.append(disc_real_loss.item()) + running_disc_fake_loss.append(disc_fake_loss.item()) + + sync_loss = get_sync_loss(mel, g) + + if hparams.disc_wt > 0.: + perceptual_loss = disc.perceptual_forward(g) + else: + perceptual_loss = 0. + + l1loss = recon_loss(g, gt) + + loss = hparams.syncnet_wt * sync_loss + hparams.disc_wt * perceptual_loss + \ + (1. - hparams.syncnet_wt - hparams.disc_wt) * l1loss + + running_l1_loss.append(l1loss.item()) + running_sync_loss.append(sync_loss.item()) + + if hparams.disc_wt > 0.: + running_perceptual_loss.append(perceptual_loss.item()) + else: + running_perceptual_loss.append(0.) + + if step > eval_steps: break + + print('L1: {}, Sync: {}, Percep: {} | Fake: {}, Real: {}'.format(sum(running_l1_loss) / len(running_l1_loss), + sum(running_sync_loss) / len(running_sync_loss), + sum(running_perceptual_loss) / len(running_perceptual_loss), + sum(running_disc_fake_loss) / len(running_disc_fake_loss), + sum(running_disc_real_loss) / len(running_disc_real_loss))) + return sum(running_sync_loss) / len(running_sync_loss) + + +def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch, prefix=''): + checkpoint_path = join( + checkpoint_dir, "{}checkpoint_step{:09d}.pth".format(prefix, global_step)) + optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None + torch.save({ + "state_dict": model.state_dict(), + "optimizer": optimizer_state, + "global_step": step, + "global_epoch": epoch, + }, checkpoint_path) + print("Saved checkpoint:", checkpoint_path) + +def _load(checkpoint_path): + if use_cuda: + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + + +def load_checkpoint(path, model, optimizer, reset_optimizer=False, overwrite_global_states=True): + global global_step + global global_epoch + + print("Load checkpoint from: {}".format(path)) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + if not reset_optimizer: + optimizer_state = checkpoint["optimizer"] + if optimizer_state is not None: + print("Load optimizer state from {}".format(path)) + optimizer.load_state_dict(checkpoint["optimizer"]) + if overwrite_global_states: + global_step = checkpoint["global_step"] + global_epoch = checkpoint["global_epoch"] + + return model + +if __name__ == "__main__": + checkpoint_dir = args.checkpoint_dir + + # Dataset and Dataloader setup + train_dataset = Dataset('train') + test_dataset = Dataset('val') + + train_data_loader = data_utils.DataLoader( + train_dataset, batch_size=hparams.batch_size, shuffle=True, + num_workers=hparams.num_workers) + + test_data_loader = data_utils.DataLoader( + test_dataset, batch_size=hparams.batch_size, + num_workers=4) + + device = torch.device("cuda" if use_cuda else "cpu") + + # Model + model = Wav2Lip().to(device) + disc = Wav2Lip_disc_qual().to(device) + + print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad))) + print('total DISC trainable params {}'.format(sum(p.numel() for p in disc.parameters() if p.requires_grad))) + + optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad], + lr=hparams.initial_learning_rate, betas=(0.5, 0.999)) + disc_optimizer = optim.Adam([p for p in disc.parameters() if p.requires_grad], + lr=hparams.disc_initial_learning_rate, betas=(0.5, 0.999)) + + if args.checkpoint_path is not None: + load_checkpoint(args.checkpoint_path, model, optimizer, reset_optimizer=False) + + if args.disc_checkpoint_path is not None: + load_checkpoint(args.disc_checkpoint_path, disc, disc_optimizer, + reset_optimizer=False, overwrite_global_states=False) + + load_checkpoint(args.syncnet_checkpoint_path, syncnet, None, reset_optimizer=True, + overwrite_global_states=False) + + if not os.path.exists(checkpoint_dir): + os.mkdir(checkpoint_dir) + + # Train! + train(device, model, disc, train_data_loader, test_data_loader, optimizer, disc_optimizer, + checkpoint_dir=checkpoint_dir, + checkpoint_interval=hparams.checkpoint_interval, + nepochs=hparams.nepochs) diff --git a/models/Wav2Lip/inference.py b/models/Wav2Lip/inference.py new file mode 100644 index 0000000000..b47f4af24f --- /dev/null +++ b/models/Wav2Lip/inference.py @@ -0,0 +1,286 @@ +from os import listdir, path +import numpy as np +import scipy, cv2, os, sys, argparse, audio +import json, subprocess, random, string +from tqdm import tqdm +from glob import glob +import torch, face_detection +from models import Wav2Lip + +for k, v in os.environ.items(): + print('{}={}'.format(k, v)) + +print('PyTorch version', torch.__version__) +print('CUDA Available?', torch.cuda.is_available()) +print('CUDA Device count', torch.cuda.device_count()) + +parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models') + +parser.add_argument('--checkpoint_path', type=str, + help='Name of saved checkpoint to load weights from', required=True) + +parser.add_argument('--face', type=str, + help='Filepath of video/image that contains faces to use', required=True) +parser.add_argument('--audio', type=str, + help='Filepath of video/audio file to use as raw audio source', required=True) +parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', + default='results/result_voice.mp4') + +parser.add_argument('--static', type=bool, + help='If True, then use only first video frame for inference', default=False) +parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 25)', + default=25., required=False) + +parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], + help='Padding (top, bottom, left, right). Please adjust to include chin at least') + +parser.add_argument('--face_det_batch_size', type=int, + help='Batch size for face detection', default=16) +parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128) + +parser.add_argument('--resize_factor', default=1, type=int, + help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p') + +parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], + help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' + 'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width') + +parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], + help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.' + 'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).') + +parser.add_argument('--rotate', default=False, action='store_true', + help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.' + 'Use if you get a flipped result, despite feeding a normal looking video') + +parser.add_argument('--nosmooth', default=False, action='store_true', + help='Prevent smoothing face detections over a short temporal window') + +args = parser.parse_args() +args.img_size = 96 + +if os.path.isfile(args.face) and args.face.split('.')[1] in ['jpg', 'png', 'jpeg']: + args.static = True + +def get_smoothened_boxes(boxes, T): + for i in range(len(boxes)): + if i + T > len(boxes): + window = boxes[len(boxes) - T:] + else: + window = boxes[i : i + T] + boxes[i] = np.mean(window, axis=0) + return boxes + +def face_detect(images): + detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, + flip_input=False, device=device) + + batch_size = args.face_det_batch_size + + while 1: + predictions = [] + try: + for i in tqdm(range(0, len(images), batch_size)): + predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size]))) + except RuntimeError: + if batch_size == 1: + raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument') + batch_size //= 2 + print('Recovering from OOM error; New batch size: {}'.format(batch_size)) + continue + break + + results = [] + pady1, pady2, padx1, padx2 = args.pads + for rect, image in zip(predictions, images): + if rect is None: + cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected. + raise ValueError('Face not detected! Ensure the video contains a face in all the frames.') + + y1 = max(0, rect[1] - pady1) + y2 = min(image.shape[0], rect[3] + pady2) + x1 = max(0, rect[0] - padx1) + x2 = min(image.shape[1], rect[2] + padx2) + + results.append([x1, y1, x2, y2]) + + boxes = np.array(results) + if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5) + results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)] + + del detector + return results + +def datagen(frames, mels): + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if args.box[0] == -1: + if not args.static: + face_det_results = face_detect(frames) # BGR2RGB for CNN face detection + else: + face_det_results = face_detect([frames[0]]) + else: + print('Using the specified bounding box instead of face detection...') + y1, y2, x1, x2 = args.box + face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames] + + for i, m in enumerate(mels): + idx = 0 if args.static else i%len(frames) + frame_to_save = frames[idx].copy() + face, coords = face_det_results[idx].copy() + + face = cv2.resize(face, (args.img_size, args.img_size)) + + img_batch.append(face) + mel_batch.append(m) + frame_batch.append(frame_to_save) + coords_batch.append(coords) + + if len(img_batch) >= args.wav2lip_batch_size: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if len(img_batch) > 0: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + +mel_step_size = 16 +device = 'cuda' if torch.cuda.is_available() else 'cpu' +print('Using {} for inference.'.format(device)) + +def _load(checkpoint_path): + if device == 'cuda': + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_model(path): + model = Wav2Lip() + print("Load checkpoint from: {}".format(path)) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + + model = model.to(device) + return model.eval() + +def main(): + if not os.path.isfile(args.face): + raise ValueError('--face argument must be a valid path to video/image file') + + elif args.face.split('.')[1] in ['jpg', 'png', 'jpeg']: + full_frames = [cv2.imread(args.face)] + fps = args.fps + + else: + video_stream = cv2.VideoCapture(args.face) + fps = video_stream.get(cv2.CAP_PROP_FPS) + + print('Reading video frames...') + + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + if args.resize_factor > 1: + frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor)) + + if args.rotate: + frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE) + + y1, y2, x1, x2 = args.crop + if x2 == -1: x2 = frame.shape[1] + if y2 == -1: y2 = frame.shape[0] + + frame = frame[y1:y2, x1:x2] + + full_frames.append(frame) + + print ("Number of frames available for inference: "+str(len(full_frames))) + + if not args.audio.endswith('.wav'): + print('Extracting raw audio...') + command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, 'temp/temp.wav') + + subprocess.call(command, shell=True) + args.audio = 'temp/temp.wav' + + wav = audio.load_wav(args.audio, 16000) + mel = audio.melspectrogram(wav) + print(mel.shape) + + if np.isnan(mel.reshape(-1)).sum() > 0: + raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again') + + mel_chunks = [] + mel_idx_multiplier = 80./fps + i = 0 + while 1: + start_idx = int(i * mel_idx_multiplier) + if start_idx + mel_step_size > len(mel[0]): + mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:]) + break + mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) + i += 1 + + print("Length of mel chunks: {}".format(len(mel_chunks))) + + full_frames = full_frames[:len(mel_chunks)] + + batch_size = args.wav2lip_batch_size + gen = datagen(full_frames.copy(), mel_chunks) + + for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, + total=int(np.ceil(float(len(mel_chunks))/batch_size)))): + if i == 0: + model = load_model(args.checkpoint_path) + print ("Model loaded") + + frame_h, frame_w = full_frames[0].shape[:-1] + out = cv2.VideoWriter('temp/result.avi', + cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h)) + + img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) + mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) + + with torch.no_grad(): + pred = model(mel_batch, img_batch) + + pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. + + for p, f, c in zip(pred, frames, coords): + y1, y2, x1, x2 = c + p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1)) + + f[y1:y2, x1:x2] = p + out.write(f) + + out.release() + + command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, 'temp/result.avi', args.outfile) + subprocess.call(command, shell=True) + +if __name__ == '__main__': + main() diff --git a/models/Wav2Lip/models/__init__.py b/models/Wav2Lip/models/__init__.py new file mode 100644 index 0000000000..4374370494 --- /dev/null +++ b/models/Wav2Lip/models/__init__.py @@ -0,0 +1,2 @@ +from .wav2lip import Wav2Lip, Wav2Lip_disc_qual +from .syncnet import SyncNet_color \ No newline at end of file diff --git a/models/Wav2Lip/models/conv.py b/models/Wav2Lip/models/conv.py new file mode 100644 index 0000000000..ed83da00cb --- /dev/null +++ b/models/Wav2Lip/models/conv.py @@ -0,0 +1,44 @@ +import torch +from torch import nn +from torch.nn import functional as F + +class Conv2d(nn.Module): + def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs): + super().__init__(*args, **kwargs) + self.conv_block = nn.Sequential( + nn.Conv2d(cin, cout, kernel_size, stride, padding), + nn.BatchNorm2d(cout) + ) + self.act = nn.ReLU() + self.residual = residual + + def forward(self, x): + out = self.conv_block(x) + if self.residual: + out += x + return self.act(out) + +class nonorm_Conv2d(nn.Module): + def __init__(self, cin, cout, kernel_size, stride, padding, residual=False, *args, **kwargs): + super().__init__(*args, **kwargs) + self.conv_block = nn.Sequential( + nn.Conv2d(cin, cout, kernel_size, stride, padding), + ) + self.act = nn.LeakyReLU(0.01, inplace=True) + + def forward(self, x): + out = self.conv_block(x) + return self.act(out) + +class Conv2dTranspose(nn.Module): + def __init__(self, cin, cout, kernel_size, stride, padding, output_padding=0, *args, **kwargs): + super().__init__(*args, **kwargs) + self.conv_block = nn.Sequential( + nn.ConvTranspose2d(cin, cout, kernel_size, stride, padding, output_padding), + nn.BatchNorm2d(cout) + ) + self.act = nn.ReLU() + + def forward(self, x): + out = self.conv_block(x) + return self.act(out) diff --git a/models/Wav2Lip/models/syncnet.py b/models/Wav2Lip/models/syncnet.py new file mode 100644 index 0000000000..e773cdca67 --- /dev/null +++ b/models/Wav2Lip/models/syncnet.py @@ -0,0 +1,66 @@ +import torch +from torch import nn +from torch.nn import functional as F + +from .conv import Conv2d + +class SyncNet_color(nn.Module): + def __init__(self): + super(SyncNet_color, self).__init__() + + self.face_encoder = nn.Sequential( + Conv2d(15, 32, kernel_size=(7, 7), stride=1, padding=3), + + Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=1), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(64, 128, kernel_size=3, stride=2, padding=1), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(128, 256, kernel_size=3, stride=2, padding=1), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(256, 512, kernel_size=3, stride=2, padding=1), + Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(512, 512, kernel_size=3, stride=2, padding=1), + Conv2d(512, 512, kernel_size=3, stride=1, padding=0), + Conv2d(512, 512, kernel_size=1, stride=1, padding=0),) + + self.audio_encoder = nn.Sequential( + Conv2d(1, 32, kernel_size=3, stride=1, padding=1), + Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(64, 128, kernel_size=3, stride=3, padding=1), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(256, 512, kernel_size=3, stride=1, padding=0), + Conv2d(512, 512, kernel_size=1, stride=1, padding=0),) + + def forward(self, audio_sequences, face_sequences): # audio_sequences := (B, dim, T) + face_embedding = self.face_encoder(face_sequences) + audio_embedding = self.audio_encoder(audio_sequences) + + audio_embedding = audio_embedding.view(audio_embedding.size(0), -1) + face_embedding = face_embedding.view(face_embedding.size(0), -1) + + audio_embedding = F.normalize(audio_embedding, p=2, dim=1) + face_embedding = F.normalize(face_embedding, p=2, dim=1) + + + return audio_embedding, face_embedding diff --git a/models/Wav2Lip/models/wav2lip.py b/models/Wav2Lip/models/wav2lip.py new file mode 100644 index 0000000000..ae5d691916 --- /dev/null +++ b/models/Wav2Lip/models/wav2lip.py @@ -0,0 +1,184 @@ +import torch +from torch import nn +from torch.nn import functional as F +import math + +from .conv import Conv2dTranspose, Conv2d, nonorm_Conv2d + +class Wav2Lip(nn.Module): + def __init__(self): + super(Wav2Lip, self).__init__() + + self.face_encoder_blocks = nn.ModuleList([ + nn.Sequential(Conv2d(6, 16, kernel_size=7, stride=1, padding=3)), # 96,96 + + nn.Sequential(Conv2d(16, 32, kernel_size=3, stride=2, padding=1), # 48,48 + Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True)), + + nn.Sequential(Conv2d(32, 64, kernel_size=3, stride=2, padding=1), # 24,24 + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True)), + + nn.Sequential(Conv2d(64, 128, kernel_size=3, stride=2, padding=1), # 12,12 + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True)), + + nn.Sequential(Conv2d(128, 256, kernel_size=3, stride=2, padding=1), # 6,6 + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True)), + + nn.Sequential(Conv2d(256, 512, kernel_size=3, stride=2, padding=1), # 3,3 + Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),), + + nn.Sequential(Conv2d(512, 512, kernel_size=3, stride=1, padding=0), # 1, 1 + Conv2d(512, 512, kernel_size=1, stride=1, padding=0)),]) + + self.audio_encoder = nn.Sequential( + Conv2d(1, 32, kernel_size=3, stride=1, padding=1), + Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(32, 32, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(32, 64, kernel_size=3, stride=(3, 1), padding=1), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(64, 128, kernel_size=3, stride=3, padding=1), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(128, 256, kernel_size=3, stride=(3, 2), padding=1), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + + Conv2d(256, 512, kernel_size=3, stride=1, padding=0), + Conv2d(512, 512, kernel_size=1, stride=1, padding=0),) + + self.face_decoder_blocks = nn.ModuleList([ + nn.Sequential(Conv2d(512, 512, kernel_size=1, stride=1, padding=0),), + + nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=1, padding=0), # 3,3 + Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),), + + nn.Sequential(Conv2dTranspose(1024, 512, kernel_size=3, stride=2, padding=1, output_padding=1), + Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(512, 512, kernel_size=3, stride=1, padding=1, residual=True),), # 6, 6 + + nn.Sequential(Conv2dTranspose(768, 384, kernel_size=3, stride=2, padding=1, output_padding=1), + Conv2d(384, 384, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(384, 384, kernel_size=3, stride=1, padding=1, residual=True),), # 12, 12 + + nn.Sequential(Conv2dTranspose(512, 256, kernel_size=3, stride=2, padding=1, output_padding=1), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(256, 256, kernel_size=3, stride=1, padding=1, residual=True),), # 24, 24 + + nn.Sequential(Conv2dTranspose(320, 128, kernel_size=3, stride=2, padding=1, output_padding=1), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(128, 128, kernel_size=3, stride=1, padding=1, residual=True),), # 48, 48 + + nn.Sequential(Conv2dTranspose(160, 64, kernel_size=3, stride=2, padding=1, output_padding=1), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True), + Conv2d(64, 64, kernel_size=3, stride=1, padding=1, residual=True),),]) # 96,96 + + self.output_block = nn.Sequential(Conv2d(80, 32, kernel_size=3, stride=1, padding=1), + nn.Conv2d(32, 3, kernel_size=1, stride=1, padding=0), + nn.Sigmoid()) + + def forward(self, audio_sequences, face_sequences): + # audio_sequences = (B, T, 1, 80, 16) + B = audio_sequences.size(0) + + input_dim_size = len(face_sequences.size()) + if input_dim_size > 4: + audio_sequences = torch.cat([audio_sequences[:, i] for i in range(audio_sequences.size(1))], dim=0) + face_sequences = torch.cat([face_sequences[:, :, i] for i in range(face_sequences.size(2))], dim=0) + + audio_embedding = self.audio_encoder(audio_sequences) # B, 512, 1, 1 + + feats = [] + x = face_sequences + for f in self.face_encoder_blocks: + x = f(x) + feats.append(x) + + x = audio_embedding + for f in self.face_decoder_blocks: + x = f(x) + try: + x = torch.cat((x, feats[-1]), dim=1) + except Exception as e: + print(x.size()) + print(feats[-1].size()) + raise e + + feats.pop() + + x = self.output_block(x) + + if input_dim_size > 4: + x = torch.split(x, B, dim=0) # [(B, C, H, W)] + outputs = torch.stack(x, dim=2) # (B, C, T, H, W) + + else: + outputs = x + + return outputs + +class Wav2Lip_disc_qual(nn.Module): + def __init__(self): + super(Wav2Lip_disc_qual, self).__init__() + + self.face_encoder_blocks = nn.ModuleList([ + nn.Sequential(nonorm_Conv2d(3, 32, kernel_size=7, stride=1, padding=3)), # 48,96 + + nn.Sequential(nonorm_Conv2d(32, 64, kernel_size=5, stride=(1, 2), padding=2), # 48,48 + nonorm_Conv2d(64, 64, kernel_size=5, stride=1, padding=2)), + + nn.Sequential(nonorm_Conv2d(64, 128, kernel_size=5, stride=2, padding=2), # 24,24 + nonorm_Conv2d(128, 128, kernel_size=5, stride=1, padding=2)), + + nn.Sequential(nonorm_Conv2d(128, 256, kernel_size=5, stride=2, padding=2), # 12,12 + nonorm_Conv2d(256, 256, kernel_size=5, stride=1, padding=2)), + + nn.Sequential(nonorm_Conv2d(256, 512, kernel_size=3, stride=2, padding=1), # 6,6 + nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1)), + + nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=2, padding=1), # 3,3 + nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=1),), + + nn.Sequential(nonorm_Conv2d(512, 512, kernel_size=3, stride=1, padding=0), # 1, 1 + nonorm_Conv2d(512, 512, kernel_size=1, stride=1, padding=0)),]) + + self.binary_pred = nn.Sequential(nn.Conv2d(512, 1, kernel_size=1, stride=1, padding=0), nn.Sigmoid()) + self.label_noise = .0 + + def get_lower_half(self, face_sequences): + return face_sequences[:, :, face_sequences.size(2)//2:] + + def to_2d(self, face_sequences): + B = face_sequences.size(0) + face_sequences = torch.cat([face_sequences[:, :, i] for i in range(face_sequences.size(2))], dim=0) + return face_sequences + + def perceptual_forward(self, false_face_sequences): + false_face_sequences = self.to_2d(false_face_sequences) + false_face_sequences = self.get_lower_half(false_face_sequences) + + false_feats = false_face_sequences + for f in self.face_encoder_blocks: + false_feats = f(false_feats) + + false_pred_loss = F.binary_cross_entropy(self.binary_pred(false_feats).view(len(false_feats), -1), + torch.ones((len(false_feats), 1)).cuda()) + + return false_pred_loss + + def forward(self, face_sequences): + face_sequences = self.to_2d(face_sequences) + face_sequences = self.get_lower_half(face_sequences) + + x = face_sequences + for f in self.face_encoder_blocks: + x = f(x) + + return self.binary_pred(x).view(len(x), -1) diff --git a/models/Wav2Lip/preprocess.py b/models/Wav2Lip/preprocess.py new file mode 100644 index 0000000000..5322012ac6 --- /dev/null +++ b/models/Wav2Lip/preprocess.py @@ -0,0 +1,113 @@ +import sys + +if sys.version_info[0] < 3 and sys.version_info[1] < 2: + raise Exception("Must be using >= Python 3.2") + +from os import listdir, path + +if not path.isfile('face_detection/detection/sfd/s3fd.pth'): + raise FileNotFoundError('Save the s3fd model to face_detection/detection/sfd/s3fd.pth \ + before running this script!') + +import multiprocessing as mp +from concurrent.futures import ThreadPoolExecutor, as_completed +import numpy as np +import argparse, os, cv2, traceback, subprocess +from tqdm import tqdm +from glob import glob +import audio +from hparams import hparams as hp + +import face_detection + +parser = argparse.ArgumentParser() + +parser.add_argument('--ngpu', help='Number of GPUs across which to run in parallel', default=1, type=int) +parser.add_argument('--batch_size', help='Single GPU Face detection batch size', default=32, type=int) +parser.add_argument("--data_root", help="Root folder of the LRS2 dataset", required=True) +parser.add_argument("--preprocessed_root", help="Root folder of the preprocessed dataset", required=True) + +args = parser.parse_args() + +fa = [face_detection.FaceAlignment(face_detection.LandmarksType._2D, flip_input=False, + device='cuda:{}'.format(id)) for id in range(args.ngpu)] + +template = 'ffmpeg -loglevel panic -y -i {} -strict -2 {}' +# template2 = 'ffmpeg -hide_banner -loglevel panic -threads 1 -y -i {} -async 1 -ac 1 -vn -acodec pcm_s16le -ar 16000 {}' + +def process_video_file(vfile, args, gpu_id): + video_stream = cv2.VideoCapture(vfile) + + frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + frames.append(frame) + + vidname = os.path.basename(vfile).split('.')[0] + dirname = vfile.split('/')[-2] + + fulldir = path.join(args.preprocessed_root, dirname, vidname) + os.makedirs(fulldir, exist_ok=True) + + batches = [frames[i:i + args.batch_size] for i in range(0, len(frames), args.batch_size)] + + i = -1 + for fb in batches: + preds = fa[gpu_id].get_detections_for_batch(np.asarray(fb)) + + for j, f in enumerate(preds): + i += 1 + if f is None: + continue + + x1, y1, x2, y2 = f + cv2.imwrite(path.join(fulldir, '{}.jpg'.format(i)), fb[j][y1:y2, x1:x2]) + +def process_audio_file(vfile, args): + vidname = os.path.basename(vfile).split('.')[0] + dirname = vfile.split('/')[-2] + + fulldir = path.join(args.preprocessed_root, dirname, vidname) + os.makedirs(fulldir, exist_ok=True) + + wavpath = path.join(fulldir, 'audio.wav') + + command = template.format(vfile, wavpath) + subprocess.call(command, shell=True) + + +def mp_handler(job): + vfile, args, gpu_id = job + try: + process_video_file(vfile, args, gpu_id) + except KeyboardInterrupt: + exit(0) + except: + traceback.print_exc() + +def main(args): + print('Started processing for {} with {} GPUs'.format(args.data_root, args.ngpu)) + + filelist = glob(path.join(args.data_root, '*/*.mp4')) + + jobs = [(vfile, args, i%args.ngpu) for i, vfile in enumerate(filelist)] + p = ThreadPoolExecutor(args.ngpu) + futures = [p.submit(mp_handler, j) for j in jobs] + _ = [r.result() for r in tqdm(as_completed(futures), total=len(futures))] + + print('Dumping audios...') + + for vfile in tqdm(filelist): + try: + process_audio_file(vfile, args) + except KeyboardInterrupt: + exit(0) + except: + traceback.print_exc() + continue + +if __name__ == '__main__': + main(args) \ No newline at end of file diff --git a/models/Wav2Lip/requirements.txt b/models/Wav2Lip/requirements.txt new file mode 100644 index 0000000000..e0422532b0 --- /dev/null +++ b/models/Wav2Lip/requirements.txt @@ -0,0 +1,16 @@ +librosa==0.7.0 +numba==0.48 +numpy==1.17.1 +opencv-contrib-python==4.2.0.34 +opencv-python==4.1.0.25 +python-magic==0.4.22 +tensorflow==1.12.0 +torch==1.1.0 +torchvision==0.3.0 +tqdm==4.45.0 + +# To download (w2l-download-job) +# Note: these will need frequent updates! +gdown==3.13.0 +youtube-dl==2021.5.16 + diff --git a/models/Wav2Lip/results/README.md b/models/Wav2Lip/results/README.md new file mode 100644 index 0000000000..b1bbfd53fd --- /dev/null +++ b/models/Wav2Lip/results/README.md @@ -0,0 +1 @@ +Generated results will be placed in this folder by default. \ No newline at end of file diff --git a/models/Wav2Lip/temp/README.md b/models/Wav2Lip/temp/README.md new file mode 100644 index 0000000000..04c9104993 --- /dev/null +++ b/models/Wav2Lip/temp/README.md @@ -0,0 +1 @@ +Temporary files at the time of inference/testing will be saved here. You can ignore them. \ No newline at end of file diff --git a/models/Wav2Lip/vocodes_inference.py b/models/Wav2Lip/vocodes_inference.py new file mode 100644 index 0000000000..7fd4a2032d --- /dev/null +++ b/models/Wav2Lip/vocodes_inference.py @@ -0,0 +1,458 @@ +from os import listdir, path +import numpy as np +import scipy, cv2, os, sys, argparse, audio +import json, subprocess, random, string +from tqdm import tqdm +from glob import glob +import torch, face_detection +from models import Wav2Lip +import tempfile +import shutil +import datetime +import pickle +from PIL import Image + +if 'PRINT_ENV' in os.environ and os.environ['PRINT_ENV']: + for k, v in os.environ.items(): + print('{}={}'.format(k, v)) + +print('========================================') +print('Python interpreter', sys.executable) +print('PyTorch version', torch.__version__) +print('CUDA Available?', torch.cuda.is_available()) +print('CUDA Device count', torch.cuda.device_count()) +print('========================================', flush=True) + +parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models') + +parser.add_argument('--checkpoint_path', type=str, + help='Name of saved checkpoint to load weights from', required=True) + +parser.add_argument('--face', type=str, + help='Filepath of video/image that contains faces to use', required=True) +parser.add_argument('--audio', type=str, + help='Filepath of video/audio file to use as raw audio source', required=True) +parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', + default='results/result_voice.mp4') + +parser.add_argument('--static', type=bool, + help='If True, then use only first video frame for inference', default=False) + +# NB: I'm adjusting the FPS to match that of the end bump I made. The old default FPS was "25." (float 25) +parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 29.97)', + default=29.97, required=False) + +parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], + help='Padding (top, bottom, left, right). Please adjust to include chin at least') + +parser.add_argument('--face_det_batch_size', type=int, + help='Batch size for face detection', default=16) +parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128) + +parser.add_argument('--resize_factor', default=1, type=int, + help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p') + +parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], + help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' + 'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width') + +parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], + help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.' + 'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).') + +parser.add_argument('--rotate', default=False, action='store_true', + help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.' + 'Use if you get a flipped result, despite feeding a normal looking video') + +parser.add_argument('--nosmooth', default=False, action='store_true', + help='Prevent smoothing face detections over a short temporal window') + +parser.add_argument('--audio_start_pad_millis', default=0, type=int, + help='Seconds to pad the start of the audio') + +parser.add_argument('--audio_end_pad_millis', default=0, type=int, + help='Seconds to pad the end of the audio') + +parser.add_argument('--end_bump_file', type=str, + help='Video file to concatenate at the end') + +# This is a little more deliberate than "static" and causes us to pick an FPS, etc. +# This is useful when the filename doesn't have an extension. +parser.add_argument('--is_image', default=False, action='store_true', + help='Denote that the input is a single-frame image, not a video') + +# Purely for debugging on the host machine: +parser.add_argument('--preserve_tempdir', default=False, action='store_true', + help='Keep the tempdir arround for debugging') + +args = parser.parse_args() +args.img_size = 96 + +if os.path.isfile(args.face) and os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg']: + args.static = True +elif args.is_image: + args.static = True + +def get_smoothened_boxes(boxes, T): + for i in range(len(boxes)): + if i + T > len(boxes): + window = boxes[len(boxes) - T:] + else: + window = boxes[i : i + T] + boxes[i] = np.mean(window, axis=0) + return boxes + +def face_detect(images): + detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, + flip_input=False, device=device) + + batch_size = args.face_det_batch_size + + while 1: + predictions = [] + try: + for i in tqdm(range(0, len(images), batch_size)): + predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size]))) + except RuntimeError: + if batch_size == 1: + raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument') + batch_size //= 2 + print('Recovering from OOM error; New batch size: {}'.format(batch_size)) + continue + break + + results = [] + pady1, pady2, padx1, padx2 = args.pads + for rect, image in zip(predictions, images): + if rect is None: + cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected. + raise ValueError('Face not detected! Ensure the video contains a face in all the frames.') + + y1 = max(0, rect[1] - pady1) + y2 = min(image.shape[0], rect[3] + pady2) + x1 = max(0, rect[0] - padx1) + x2 = min(image.shape[1], rect[2] + padx2) + + results.append([x1, y1, x2, y2]) + + boxes = np.array(results) + if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5) + results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)] + + del detector + return results + +def detect_faces_in_frames(frames, video_faces_pickle_file): + t0 = datetime.datetime.now() + if args.box[0] == -1: + if not args.static: + face_det_results = face_detect(frames) # BGR2RGB for CNN face detection + else: + face_det_results = face_detect([frames[0]]) + else: + print('Using the specified bounding box instead of face detection...') + y1, y2, x1, x2 = args.box + face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames] + t1 = datetime.datetime.now() + + d = t1 - t0 + print('Total seconds to detect faces in {} frames: {}'.format(len(frames), d.total_seconds())) + + with open(video_faces_pickle_file, 'wb') as f: + print('Saving pickle file: {}'.format(video_faces_pickle_file)) + pickle.dump(face_det_results, f) + + return face_det_results + +#def datagen(frames, mels): +def datagen(frames, face_det_results, mels): + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + #face_det_results = detect_faces_in_frames(frames) + + for i, m in enumerate(mels): + idx = 0 if args.static else i%len(frames) + frame_to_save = frames[idx].copy() + face, coords = face_det_results[idx].copy() + + face = cv2.resize(face, (args.img_size, args.img_size)) + + img_batch.append(face) + mel_batch.append(m) + frame_batch.append(frame_to_save) + coords_batch.append(coords) + + if len(img_batch) >= args.wav2lip_batch_size: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if len(img_batch) > 0: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + +mel_step_size = 16 +device = 'cuda' if torch.cuda.is_available() else 'cpu' +print('Using {} for inference.'.format(device)) + +def _load(checkpoint_path): + if device == 'cuda': + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_model(path): + model = Wav2Lip() + print("Load checkpoint from: {}".format(path), flush=True) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + + model = model.to(device) + return model.eval() + +def maybe_pad_audio_file(tempdir, args): + if not args.audio_start_pad_millis and not args.audio_end_pad_millis: + return + + print('Padding audio with {} millis at start, {} millis at end.'.format( + args.audio_start_pad_millis, args.audio_end_pad_millis), flush=True) + + padded_wav_filename = os.path.join(tempdir, 'padded_audio.wav') + # ffmpeg padding: https://superuser.com/a/579110 + # Structure of filters: https://stackoverflow.com/a/55463101 + command = ' '.join([ + 'ffmpeg', + '-i {}'.format(args.audio), + '-filter_complex "[0] adelay={}ms:all=true [a] ; [a] apad=pad_dur={}ms"'.format(args.audio_start_pad_millis, args.audio_end_pad_millis), + '{}'.format(padded_wav_filename) + ]) + print('command:', command, flush=True) + subprocess.call(command, shell=True) + args.audio = padded_wav_filename + + +def maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h): + if not args.end_bump_file: + print('No end bump file.', flush=True) + return + + print('Adding end bump file: {}'.format(args.end_bump_file), flush=True) + + bumpfile = args.end_bump_file + + # May need to resize our bump file... + # After many attempts to get this into the same ffmpeg pipeline, + # I'm giving up and just adding another step. + # + # NOTE: The FPS of the bump needs to match the video. Especially + # important for videos generated from still photos, because that defaults + # to the odd choice of "25 fps". + if frame_w is not 1920 or frame_h is not 1080: + print('Need to resize bump.', flush=True) + bumpfile = os.path.join(tempdir, 'end_bump_resized.mp4') + + # Scale / aspect ratio / sar / resizing / padding: + # https://stackoverflow.com/a/51946719 + command = ' '.join([ + 'ffmpeg', + '-i {}'.format(args.end_bump_file), + '-vf "scale={}:{}:force_original_aspect_ratio=decrease,pad={}:{}:(ow-iw)/2:(oh-ih)/2,setsar=1"'.format(frame_w, frame_h, frame_w, frame_h), + '{}'.format(bumpfile) + ]) + print('command:', command, flush=True) + failure_code = subprocess.call(command, shell=True) + + if failure_code: + print('Failed to resize end bump') + return + + padded_wav_filename = os.path.join(tempdir, 'output_with_end_bump.mp4') + # ffmpeg video concatenation: https://stackoverflow.com/a/11175851 + # Structure of filters: https://stackoverflow.com/a/55463101 + # More info on filters: https://stackoverflow.com/a/22958746 + command = ' '.join([ + 'ffmpeg', + '-i {}'.format(args.outfile), + '-i {}'.format(bumpfile), + '-filter_complex "concat=n=2:v=1:a=1"', + '{}'.format(padded_wav_filename) + ]) + print('command:', command, flush=True) + failure_code = subprocess.call(command, shell=True) + + if failure_code: + print('Failed to concatenate end bump') + else: + print('Done concatenating end bump; renaming file...', flush=True) + os.rename(padded_wav_filename, args.outfile) + + +def main(tempdir): + video_faces_pickle_file = args.face + '.faces' + print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True) + + frame_w = 0 + frame_h = 0 + + if not os.path.isfile(args.face): + raise ValueError('--face argument must be a valid path to video/image file') + + elif os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg'] \ + or args.is_image: + full_frames = [cv2.imread(args.face)] + fps = args.fps + frame_h, frame_w = full_frames[0].shape[:-1] + + else: + video_stream = cv2.VideoCapture(args.face) + fps = video_stream.get(cv2.CAP_PROP_FPS) + + print('Reading video frames...') + + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + if args.resize_factor > 1: + frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor)) + + if args.rotate: + frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE) + + y1, y2, x1, x2 = args.crop + if x2 == -1: x2 = frame.shape[1] + if y2 == -1: y2 = frame.shape[0] + + frame = frame[y1:y2, x1:x2] + + full_frames.append(frame) + + frame_h, frame_w = full_frames[0].shape[:-1] + + print("Number of frames available for inference: "+str(len(full_frames)), flush=True) + print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True) + + if not args.audio.endswith('.wav'): + print('Extracting raw audio...') + temp_wav_filename = os.path.join(tempdir, 'temp.wav') # previously 'temp/temp.wav' + command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio, temp_wav_filename) + + print('command:', command, flush=True) + subprocess.call(command, shell=True) + args.audio = temp_wav_filename + + maybe_pad_audio_file(tempdir, args) + + wav = audio.load_wav(args.audio, 16000) + mel = audio.melspectrogram(wav) + print(mel.shape) + + if np.isnan(mel.reshape(-1)).sum() > 0: + raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again') + + mel_chunks = [] + mel_idx_multiplier = 80./fps + i = 0 + while 1: + start_idx = int(i * mel_idx_multiplier) + if start_idx + mel_step_size > len(mel[0]): + mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:]) + break + mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) + i += 1 + + print("Length of mel chunks: {}".format(len(mel_chunks)), flush=True) + + if len(mel_chunks) > 2000: + # TODO: Catch that the file is too long earlier on. + raise Exception("Too many mel chunks: {}".format(len(mel_chunks))) + + full_frames = full_frames[:len(mel_chunks)] + + # TODO: Do this offline. + print('Face detection time...', flush=True) + if not os.path.isfile(video_faces_pickle_file): + print('Detecting faces...', flush=True) + face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file) + else: + with open(video_faces_pickle_file, 'rb') as f: + face_det_results = pickle.load(f) + if len(face_det_results) < len(full_frames): + print("Face detect results have {} frames, but we need {}".format(len(face_det_results), len(full_frames))) + print('Detecting faces again...', flush=True) + face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file) + else: + print("We don't need to find faces! Woo!", flush=True) + + face_det_results = face_det_results[:len(mel_chunks)] + + batch_size = args.wav2lip_batch_size + #gen = datagen(full_frames.copy(), mel_chunks) + gen = datagen(full_frames.copy(), face_det_results.copy(), mel_chunks) + + output_video_filename = os.path.join(tempdir, 'result.avi') # previously 'temp/result.avi' + + for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, + total=int(np.ceil(float(len(mel_chunks))/batch_size)))): + if i == 0: + model = load_model(args.checkpoint_path) + print ("Model loaded") + + #frame_h, frame_w = full_frames[0].shape[:-1] + out = cv2.VideoWriter(output_video_filename, + cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h)) + + img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) + mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) + + with torch.no_grad(): + pred = model(mel_batch, img_batch) + + pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. + + for p, f, c in zip(pred, frames, coords): + y1, y2, x1, x2 = c + p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1)) + + f[y1:y2, x1:x2] = p + out.write(f) + + out.release() + + command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio, output_video_filename, args.outfile) + print('command:', command, flush=True) + subprocess.call(command, shell=True) + + maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h) + +if __name__ == '__main__': + tempdir = tempfile.mkdtemp() + print('Using Python tempdir: {}'.format(tempdir), flush=True) + try: + main(tempdir) + finally: + if not args.preserve_tempdir: + shutil.rmtree(tempdir) + diff --git a/models/Wav2Lip/vocodes_process_inference.py b/models/Wav2Lip/vocodes_process_inference.py new file mode 100644 index 0000000000..86e891c059 --- /dev/null +++ b/models/Wav2Lip/vocodes_process_inference.py @@ -0,0 +1,433 @@ +from os import listdir, path +import numpy as np +import scipy, cv2, os, sys, argparse, audio +import json, subprocess, random, string +from tqdm import tqdm +from glob import glob +import torch, face_detection +from models import Wav2Lip +import tempfile +import shutil +import datetime +import pickle +import magic +from PIL import Image + +if 'PRINT_ENV' in os.environ and os.environ['PRINT_ENV']: + for k, v in os.environ.items(): + print('{}={}'.format(k, v)) + +print('========================================') +print('Python interpreter', sys.executable) +print('PyTorch version', torch.__version__) +print('CUDA Available?', torch.cuda.is_available()) +print('CUDA Device count', torch.cuda.device_count()) +print('========================================', flush=True) + +parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models') + +parser.add_argument('--checkpoint_path', type=str, + help='Name of saved checkpoint to load weights from', required=True) + +parser.add_argument('--image_or_video_filename', type=str, + help='Filepath of video/image that contains faces to use', required=True) + +#parser.add_argument('--image_or_video_type', type=str, +# help='either "image" or "video', required=True) + +# This is a little more deliberate than "static" and causes us to pick an FPS, etc. +# This is useful when the filename doesn't have an extension. +parser.add_argument('--is_image', default=False, action='store_true', + help='Denote that the input is a single-frame image, not a video') + +parser.add_argument('--cached_faces_filename', type=str, + help='Filename for the cached faces file to use', required=True) + +#parser.add_argument('--face', type=str, +# help='Filepath of video/image that contains faces to use', required=True) +parser.add_argument('--audio_filename', type=str, + help='Filepath of video/audio file to use as raw audio source', required=True) + +parser.add_argument('--end_bump_file', type=str, + help='Video file to concatenate at the end') + +parser.add_argument('--output_video_filename', type=str, + help='Output filename for the final video', required=True) + +parser.add_argument('--output_metadata_filename', type=str, + help='Output filename for the JSON containing width, height, etc.', required=True) + +#parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', +# default='results/result_voice.mp4') + +parser.add_argument('--static', type=bool, + help='If True, then use only first video frame for inference', default=False) + +# NB: I'm adjusting the FPS to match that of the end bump I made. The old default FPS was "25." (float 25) +parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 29.97)', + default=29.97, required=False) + +parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], + help='Padding (top, bottom, left, right). Please adjust to include chin at least') + +parser.add_argument('--face_det_batch_size', type=int, + help='Batch size for face detection', default=16) +parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128) + +parser.add_argument('--resize_factor', default=1, type=int, + help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p') + +parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], + help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' + 'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width') + +parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], + help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.' + 'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).') + +parser.add_argument('--rotate', default=False, action='store_true', + help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.' + 'Use if you get a flipped result, despite feeding a normal looking video') + +parser.add_argument('--nosmooth', default=False, action='store_true', + help='Prevent smoothing face detections over a short temporal window') + +parser.add_argument('--audio_start_pad_millis', default=0, type=int, + help='Seconds to pad the start of the audio') + +parser.add_argument('--audio_end_pad_millis', default=0, type=int, + help='Seconds to pad the end of the audio') + +# Purely for debugging on the host machine: +parser.add_argument('--preserve_tempdir', default=False, action='store_true', + help='Keep the tempdir arround for debugging') + +args = parser.parse_args() +args.img_size = 96 + +def datagen(frames, face_det_results, mels): + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + for i, m in enumerate(mels): + idx = 0 if args.static else i%len(frames) + frame_to_save = frames[idx].copy() + face, coords = face_det_results[idx].copy() + + face = cv2.resize(face, (args.img_size, args.img_size)) + + img_batch.append(face) + mel_batch.append(m) + frame_batch.append(frame_to_save) + coords_batch.append(coords) + + if len(img_batch) >= args.wav2lip_batch_size: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if len(img_batch) > 0: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + +mel_step_size = 16 +device = 'cuda' if torch.cuda.is_available() else 'cpu' +print('Using {} for inference.'.format(device)) + +def _load(checkpoint_path): + if device == 'cuda': + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_model(path): + model = Wav2Lip() + print("Load checkpoint from: {}".format(path), flush=True) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + + model = model.to(device) + return model.eval() + +def maybe_pad_audio_file(tempdir, args): + if not args.audio_start_pad_millis and not args.audio_end_pad_millis: + return + + print('Padding audio with {} millis at start, {} millis at end.'.format( + args.audio_start_pad_millis, args.audio_end_pad_millis), flush=True) + + padded_wav_filename = os.path.join(tempdir, 'padded_audio.wav') + # ffmpeg padding: https://superuser.com/a/579110 + # Structure of filters: https://stackoverflow.com/a/55463101 + command = ' '.join([ + 'ffmpeg', + '-i {}'.format(args.audio_filename), + '-filter_complex "[0] adelay={}ms:all=true [a] ; [a] apad=pad_dur={}ms"'.format(args.audio_start_pad_millis, args.audio_end_pad_millis), + '{}'.format(padded_wav_filename) + ]) + print('command:', command, flush=True) + subprocess.call(command, shell=True) + args.audio_filename = padded_wav_filename + + +def maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h): + if not args.end_bump_file: + print('No end bump file.', flush=True) + return + + print('Adding end bump file: {}'.format(args.end_bump_file), flush=True) + + bumpfile = args.end_bump_file + + # May need to resize our bump file... + # After many attempts to get this into the same ffmpeg pipeline, + # I'm giving up and just adding another step. + # + # NOTE: The FPS of the bump needs to match the video. Especially + # important for videos generated from still photos, because that defaults + # to the odd choice of "25 fps". + if frame_w is not 1920 or frame_h is not 1080: + print('Need to resize bump.', flush=True) + bumpfile = os.path.join(tempdir, 'end_bump_resized.mp4') + + # Scale / aspect ratio / sar / resizing / padding: + # https://stackoverflow.com/a/51946719 + command = ' '.join([ + 'ffmpeg', + '-i {}'.format(args.end_bump_file), + '-vf "scale={}:{}:force_original_aspect_ratio=decrease,pad={}:{}:(ow-iw)/2:(oh-ih)/2,setsar=1"'.format(frame_w, frame_h, frame_w, frame_h), + '{}'.format(bumpfile) + ]) + print('command:', command, flush=True) + failure_code = subprocess.call(command, shell=True) + + if failure_code: + print('Failed to resize end bump') + return + + padded_wav_filename = os.path.join(tempdir, 'output_with_end_bump.mp4') + # ffmpeg video concatenation: https://stackoverflow.com/a/11175851 + # Structure of filters: https://stackoverflow.com/a/55463101 + # More info on filters: https://stackoverflow.com/a/22958746 + command = ' '.join([ + 'ffmpeg', + '-i {}'.format(args.output_video_filename), + '-i {}'.format(bumpfile), + '-filter_complex "concat=n=2:v=1:a=1"', + '{}'.format(padded_wav_filename) + ]) + print('command:', command, flush=True) + failure_code = subprocess.call(command, shell=True) + + if failure_code: + print('Failed to concatenate end bump') + else: + print('Done concatenating end bump; renaming file...', flush=True) + os.rename(padded_wav_filename, args.output_video_filename) + + +def main(tempdir): + frame_w = 0 + frame_h = 0 + fps = 0 + + if not os.path.isfile(args.image_or_video_filename): + raise ValueError('image_or_video_filename is not a file') + + if not os.path.isfile(args.cached_faces_filename): + raise ValueError('cached_faces_filename is not a file') + + if not os.path.isfile(args.audio_filename): + raise ValueError('audio_filename is not a file') + + if os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg'] \ + or args.is_image: + full_frames = [cv2.imread(args.image_or_video_filename)] + fps = args.fps + frame_h, frame_w = full_frames[0].shape[:-1] + + else: + video_stream = cv2.VideoCapture(args.image_or_video_filename) + fps = video_stream.get(cv2.CAP_PROP_FPS) + + print('Reading video frames...') + + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + if args.resize_factor > 1: + frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor)) + + if args.rotate: + frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE) + + y1, y2, x1, x2 = args.crop + if x2 == -1: x2 = frame.shape[1] + if y2 == -1: y2 = frame.shape[0] + + frame = frame[y1:y2, x1:x2] + + full_frames.append(frame) + + frame_h, frame_w = full_frames[0].shape[:-1] + + print("Number of frames available for inference: "+str(len(full_frames)), flush=True) + print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True) + + # TODO(bt): This could be more efficient + if not args.audio_filename.endswith('.wav'): + print('Extracting raw audio...') + temp_wav_filename = os.path.join(tempdir, 'temp.wav') # previously 'temp/temp.wav' + command = 'ffmpeg -y -i {} -strict -2 {}'.format(args.audio_filename, temp_wav_filename) + + print('command:', command, flush=True) + subprocess.call(command, shell=True) + args.audio_filename = temp_wav_filename + + maybe_pad_audio_file(tempdir, args) + + # TODO(bt): Save the spectrogram? + wav = audio.load_wav(args.audio_filename, 16000) + mel = audio.melspectrogram(wav) + print(mel.shape) + + if np.isnan(mel.reshape(-1)).sum() > 0: + raise ValueError('Mel contains nan! Using a TTS voice? Add a small epsilon noise to the wav file and try again') + + mel_chunks = [] + mel_idx_multiplier = 80./fps + i = 0 + while 1: + start_idx = int(i * mel_idx_multiplier) + if start_idx + mel_step_size > len(mel[0]): + mel_chunks.append(mel[:, len(mel[0]) - mel_step_size:]) + break + mel_chunks.append(mel[:, start_idx : start_idx + mel_step_size]) + i += 1 + + print("Length of mel chunks: {}".format(len(mel_chunks)), flush=True) + + if len(mel_chunks) > 2000: + # TODO(bt): Catch that the file is too long earlier on. + raise Exception("Too many mel chunks: {}".format(len(mel_chunks))) + + full_frames = full_frames[:len(mel_chunks)] + + # Load face detection results from pickle file + video_faces_pickle_file = args.cached_faces_filename + with open(video_faces_pickle_file, 'rb') as f: + face_det_results = pickle.load(f) + + face_det_results = face_det_results[:len(mel_chunks)] + + batch_size = args.wav2lip_batch_size + gen = datagen(full_frames.copy(), face_det_results.copy(), mel_chunks) + + output_video_filename = os.path.join(tempdir, 'result.avi') # previously 'temp/result.avi' + + for i, (img_batch, mel_batch, frames, coords) in enumerate(tqdm(gen, + total=int(np.ceil(float(len(mel_chunks))/batch_size)))): + if i == 0: + model = load_model(args.checkpoint_path) + print ("Model loaded") + + #frame_h, frame_w = full_frames[0].shape[:-1] + out = cv2.VideoWriter(output_video_filename, + cv2.VideoWriter_fourcc(*'DIVX'), fps, (frame_w, frame_h)) + + img_batch = torch.FloatTensor(np.transpose(img_batch, (0, 3, 1, 2))).to(device) + mel_batch = torch.FloatTensor(np.transpose(mel_batch, (0, 3, 1, 2))).to(device) + + with torch.no_grad(): + pred = model(mel_batch, img_batch) + + pred = pred.cpu().numpy().transpose(0, 2, 3, 1) * 255. + + for p, f, c in zip(pred, frames, coords): + y1, y2, x1, x2 = c + p = cv2.resize(p.astype(np.uint8), (x2 - x1, y2 - y1)) + + f[y1:y2, x1:x2] = p + out.write(f) + + out.release() + + command = 'ffmpeg -y -i {} -i {} -strict -2 -q:v 1 {}'.format(args.audio_filename, output_video_filename, args.output_video_filename) + print('command:', command, flush=True) + subprocess.call(command, shell=True) + + maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h) + + # ==== METADATA (1) ==== + + command = [ + "ffprobe", + "-loglevel", "quiet", + "-print_format", "json", + "-show_format", + "-show_streams", + args.output_video_filename + ] + + pipe = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.STDOUT) + out, err = pipe.communicate() + ffmpeg_metadata = json.loads(out) + + # ==== METADATA (2) ==== + mime_type = magic.from_file(args.output_video_filename, mime=True) + file_size_bytes = os.path.getsize(args.output_video_filename) + + print(ffmpeg_metadata) + + # We can't really get frame count anymore. The original source is invalid. + # ffmpeg reports fps, but it's a ratio, eg. 25/1 + #frame_count = len(full_frames) + #duration_seconds = frame_count / fps + + duration_millis = int(float(ffmpeg_metadata['format']['duration']) * 1000) + + metadata = { + 'is_video': True, + 'width': frame_w, + 'height': frame_h, + #'num_frames': frame_count, + #'fps': fps, + 'duration_millis': duration_millis, + 'mimetype': mime_type, + 'file_size_bytes': file_size_bytes, + } + + with open(args.output_metadata_filename, 'w') as json_file: + json.dump(metadata, json_file) + +if __name__ == '__main__': + tempdir = tempfile.mkdtemp() + print('Using Python tempdir: {}'.format(tempdir), flush=True) + try: + main(tempdir) + finally: + if not args.preserve_tempdir: + shutil.rmtree(tempdir) + diff --git a/models/Wav2Lip/vocodes_process_upload.py b/models/Wav2Lip/vocodes_process_upload.py new file mode 100644 index 0000000000..924ce674d5 --- /dev/null +++ b/models/Wav2Lip/vocodes_process_upload.py @@ -0,0 +1,378 @@ +from os import listdir, path +import numpy as np +import scipy, cv2, os, sys, argparse, audio +import json, subprocess, random, string +from tqdm import tqdm +from glob import glob +import torch, face_detection +from models import Wav2Lip +import tempfile +import shutil +import datetime +import pickle +import magic +from PIL import Image + +""" +Current worker filesystem architecture: + +/tmp +/tmp/models +/tmp/models/wav2lip_gan.pth +/tmp/.tmppd1Kvq +/tmp/.tmppd1Kvq/input_audio_file +/tmp/templates_images +/tmp/templates_images/elon-musk.jpg +/tmp/templates_videos +/tmp/templates_videos/ben-shapiro-baby.mp4 +/tmp/templates_videos/ben-shapiro-baby.mp4.faces +/tmp/templates_videos/dr-phil-bubble.mp4 +/tmp/tmp7um0yxjo +/tmp/tmp7um0yxjo/padded_audio.wav +/tmp/tmp7um0yxjo/temp.wav +/tmp/assets +/tmp/assets/vocodes-short-end-bump.mp4 + +""" + +if 'PRINT_ENV' in os.environ and os.environ['PRINT_ENV']: + for k, v in os.environ.items(): + print('{}={}'.format(k, v)) + +print('========================================') +print('Python interpreter', sys.executable) +print('PyTorch version', torch.__version__) +print('CUDA Available?', torch.cuda.is_available()) +print('CUDA Device count', torch.cuda.device_count()) +print('========================================', flush=True) + +parser = argparse.ArgumentParser(description='Inference code to lip-sync videos in the wild using Wav2Lip models') + +parser.add_argument('--checkpoint_path', type=str, + help='Name of saved checkpoint to load weights from', required=True) + +parser.add_argument('--image_or_video_filename', type=str, + help='Filepath of video/image that contains faces to use', required=True) + +parser.add_argument('--output_cached_faces_filename', type=str, + help='Output filename for the cached faces file', required=True) + +parser.add_argument('--output_metadata_filename', type=str, + help='Output filename for the JSON containing width, height, etc.', required=True) + +# NB: Not needed for processing the upload file for faces: +# +# parser.add_argument('--audio', type=str, +# help='Filepath of video/audio file to use as raw audio source', required=True) +# parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.', +# default='results/result_voice.mp4') +# + +parser.add_argument('--static', type=bool, + help='If True, then use only first video frame for inference', default=False) + +# NB: I'm adjusting the FPS to match that of the end bump I made. The old default FPS was "25." (float 25) +parser.add_argument('--fps', type=float, help='Can be specified only if input is a static image (default: 29.97)', + default=29.97, required=False) + +parser.add_argument('--pads', nargs='+', type=int, default=[0, 10, 0, 0], + help='Padding (top, bottom, left, right). Please adjust to include chin at least') + +parser.add_argument('--face_det_batch_size', type=int, + help='Batch size for face detection', default=16) +parser.add_argument('--wav2lip_batch_size', type=int, help='Batch size for Wav2Lip model(s)', default=128) + +parser.add_argument('--resize_factor', default=1, type=int, + help='Reduce the resolution by this factor. Sometimes, best results are obtained at 480p or 720p') + +parser.add_argument('--crop', nargs='+', type=int, default=[0, -1, 0, -1], + help='Crop video to a smaller region (top, bottom, left, right). Applied after resize_factor and rotate arg. ' + 'Useful if multiple face present. -1 implies the value will be auto-inferred based on height, width') + +parser.add_argument('--box', nargs='+', type=int, default=[-1, -1, -1, -1], + help='Specify a constant bounding box for the face. Use only as a last resort if the face is not detected.' + 'Also, might work only if the face is not moving around much. Syntax: (top, bottom, left, right).') + +parser.add_argument('--rotate', default=False, action='store_true', + help='Sometimes videos taken from a phone can be flipped 90deg. If true, will flip video right by 90deg.' + 'Use if you get a flipped result, despite feeding a normal looking video') + +parser.add_argument('--nosmooth', default=False, action='store_true', + help='Prevent smoothing face detections over a short temporal window') + +parser.add_argument('--audio_start_pad_millis', default=0, type=int, + help='Seconds to pad the start of the audio') + +parser.add_argument('--audio_end_pad_millis', default=0, type=int, + help='Seconds to pad the end of the audio') + +parser.add_argument('--end_bump_file', type=str, + help='Video file to concatenate at the end') + +# This is a little more deliberate than "static" and causes us to pick an FPS, etc. +# This is useful when the filename doesn't have an extension. +parser.add_argument('--is_image', default=False, action='store_true', + help='Denote that the input is a single-frame image, not a video') + +# Purely for debugging on the host machine: +parser.add_argument('--preserve_tempdir', default=False, action='store_true', + help='Keep the tempdir arround for debugging') + +args = parser.parse_args() +args.img_size = 96 + +if os.path.isfile(args.image_or_video_filename) \ + and os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg']: + args.static = True +elif args.is_image: + args.static = True + + +class FacelessFrameError(Exception): + """Throw this when there are frames without faces.""" + pass + +def get_smoothened_boxes(boxes, T): + for i in range(len(boxes)): + if i + T > len(boxes): + window = boxes[len(boxes) - T:] + else: + window = boxes[i : i + T] + boxes[i] = np.mean(window, axis=0) + return boxes + +def face_detect(images): + detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D, + flip_input=False, device=device) + + batch_size = args.face_det_batch_size + + while 1: + predictions = [] + try: + for i in tqdm(range(0, len(images), batch_size)): + predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size]))) + except RuntimeError: + if batch_size == 1: + raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument') + batch_size //= 2 + print('Recovering from OOM error; New batch size: {}'.format(batch_size)) + continue + break + + results = [] + pady1, pady2, padx1, padx2 = args.pads + for rect, image in zip(predictions, images): + if rect is None: + cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected. + raise FacelessFrameError('Face not detected! Ensure the video contains a face in all the frames.') + + y1 = max(0, rect[1] - pady1) + y2 = min(image.shape[0], rect[3] + pady2) + x1 = max(0, rect[0] - padx1) + x2 = min(image.shape[1], rect[2] + padx2) + + results.append([x1, y1, x2, y2]) + + boxes = np.array(results) + if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5) + results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)] + + del detector + return results + +def detect_faces_in_frames(frames, video_faces_pickle_file): + t0 = datetime.datetime.now() + if args.box[0] == -1: + if not args.static: + face_det_results = face_detect(frames) # BGR2RGB for CNN face detection + else: + face_det_results = face_detect([frames[0]]) + else: + print('Using the specified bounding box instead of face detection...') + y1, y2, x1, x2 = args.box + face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames] + t1 = datetime.datetime.now() + + d = t1 - t0 + print('Total seconds to detect faces in {} frames: {}'.format(len(frames), d.total_seconds())) + + with open(video_faces_pickle_file, 'wb') as f: + print('Saving pickle file: {}'.format(video_faces_pickle_file)) + pickle.dump(face_det_results, f) + + return face_det_results + +#def datagen(frames, mels): +def datagen(frames, face_det_results, mels): + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + #face_det_results = detect_faces_in_frames(frames) + + for i, m in enumerate(mels): + idx = 0 if args.static else i%len(frames) + frame_to_save = frames[idx].copy() + face, coords = face_det_results[idx].copy() + + face = cv2.resize(face, (args.img_size, args.img_size)) + + img_batch.append(face) + mel_batch.append(m) + frame_batch.append(frame_to_save) + coords_batch.append(coords) + + if len(img_batch) >= args.wav2lip_batch_size: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + img_batch, mel_batch, frame_batch, coords_batch = [], [], [], [] + + if len(img_batch) > 0: + img_batch, mel_batch = np.asarray(img_batch), np.asarray(mel_batch) + + img_masked = img_batch.copy() + img_masked[:, args.img_size//2:] = 0 + + img_batch = np.concatenate((img_masked, img_batch), axis=3) / 255. + mel_batch = np.reshape(mel_batch, [len(mel_batch), mel_batch.shape[1], mel_batch.shape[2], 1]) + + yield img_batch, mel_batch, frame_batch, coords_batch + +mel_step_size = 16 +device = 'cuda' if torch.cuda.is_available() else 'cpu' +print('Using {} for inference.'.format(device)) + +#def _load(checkpoint_path): +# if device == 'cuda': +# checkpoint = torch.load(checkpoint_path) +# else: +# checkpoint = torch.load(checkpoint_path, +# map_location=lambda storage, loc: storage) +# return checkpoint + +#def load_model(path): +# model = Wav2Lip() +# print("Load checkpoint from: {}".format(path), flush=True) +# checkpoint = _load(path) +# s = checkpoint["state_dict"] +# new_s = {} +# for k, v in s.items(): +# new_s[k.replace('module.', '')] = v +# model.load_state_dict(new_s) +# +# model = model.to(device) +# return model.eval() + + +def main(tempdir): + video_faces_pickle_file = args.output_cached_faces_filename + print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True) + + is_video = False + frame_w = 0 + frame_h = 0 + frame_count = 0 + fps = 0.0 + duration_millis = 0.0 + + if not os.path.isfile(args.image_or_video_filename): + raise ValueError('--image_or_video_filename argument must be a valid path to video/image file') + + elif os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg'] \ + or args.is_image: + full_frames = [cv2.imread(args.image_or_video_filename)] + fps = args.fps + frame_h, frame_w = full_frames[0].shape[:-1] + is_video = False + frame_count = len(full_frames) + + else: + video_stream = cv2.VideoCapture(args.image_or_video_filename) + fps = video_stream.get(cv2.CAP_PROP_FPS) + + print('Reading video frames...') + + full_frames = [] + while 1: + still_reading, frame = video_stream.read() + if not still_reading: + video_stream.release() + break + if args.resize_factor > 1: + frame = cv2.resize(frame, (frame.shape[1]//args.resize_factor, frame.shape[0]//args.resize_factor)) + + if args.rotate: + frame = cv2.rotate(frame, cv2.cv2.ROTATE_90_CLOCKWISE) + + y1, y2, x1, x2 = args.crop + if x2 == -1: x2 = frame.shape[1] + if y2 == -1: y2 = frame.shape[0] + + frame = frame[y1:y2, x1:x2] + + full_frames.append(frame) + + frame_h, frame_w = full_frames[0].shape[:-1] + + if len(full_frames) < 2: + is_video = False + frame_count = len(full_frames) + else: + is_video = True + frame_count = len(full_frames) + duration_seconds = frame_count / fps + duration_millis = int(duration_seconds * 1000) + + print("Number of frames available for inference: "+str(len(full_frames)), flush=True) + print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True) + + # NB: Instead of truncating, let's detect the entire file. + # full_frames = full_frames[:len(mel_chunks)] + + print('Face detection time...', flush=True) + print('Number of frames to detect faces in: {}'.format(len(full_frames)), flush=True) + + print('Detecting faces...', flush=True) + try: + _face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file) + except FacelessFrameError as e: + print('One or more frames without faces were encountered.') + print(e) + sys.exit(5) + + print('Done detecting faces!', flush=True) + + mime_type = magic.from_file(args.image_or_video_filename, mime=True) + file_size_bytes = os.path.getsize(args.image_or_video_filename) + + metadata = { + 'is_video': is_video, + 'width': frame_w, + 'height': frame_h, + 'num_frames': len(full_frames), + 'mimetype': mime_type, + 'file_size_bytes': file_size_bytes, + } + + if is_video: + metadata['fps'] = fps + metadata['duration_millis'] = duration_millis + + with open(args.output_metadata_filename, 'w') as json_file: + json.dump(metadata, json_file) + +if __name__ == '__main__': + tempdir = tempfile.mkdtemp() + print('Using Python tempdir: {}'.format(tempdir), flush=True) + try: + main(tempdir) + finally: + if not args.preserve_tempdir: + shutil.rmtree(tempdir) + diff --git a/models/Wav2Lip/wav2lip_train.py b/models/Wav2Lip/wav2lip_train.py new file mode 100644 index 0000000000..6e0811808a --- /dev/null +++ b/models/Wav2Lip/wav2lip_train.py @@ -0,0 +1,374 @@ +from os.path import dirname, join, basename, isfile +from tqdm import tqdm + +from models import SyncNet_color as SyncNet +from models import Wav2Lip as Wav2Lip +import audio + +import torch +from torch import nn +from torch import optim +import torch.backends.cudnn as cudnn +from torch.utils import data as data_utils +import numpy as np + +from glob import glob + +import os, random, cv2, argparse +from hparams import hparams, get_image_list + +parser = argparse.ArgumentParser(description='Code to train the Wav2Lip model without the visual quality discriminator') + +parser.add_argument("--data_root", help="Root folder of the preprocessed LRS2 dataset", required=True, type=str) + +parser.add_argument('--checkpoint_dir', help='Save checkpoints to this directory', required=True, type=str) +parser.add_argument('--syncnet_checkpoint_path', help='Load the pre-trained Expert discriminator', required=True, type=str) + +parser.add_argument('--checkpoint_path', help='Resume from this checkpoint', default=None, type=str) + +args = parser.parse_args() + + +global_step = 0 +global_epoch = 0 +use_cuda = torch.cuda.is_available() +print('use_cuda: {}'.format(use_cuda)) + +syncnet_T = 5 +syncnet_mel_step_size = 16 + +class Dataset(object): + def __init__(self, split): + self.all_videos = get_image_list(args.data_root, split) + + def get_frame_id(self, frame): + return int(basename(frame).split('.')[0]) + + def get_window(self, start_frame): + start_id = self.get_frame_id(start_frame) + vidname = dirname(start_frame) + + window_fnames = [] + for frame_id in range(start_id, start_id + syncnet_T): + frame = join(vidname, '{}.jpg'.format(frame_id)) + if not isfile(frame): + return None + window_fnames.append(frame) + return window_fnames + + def read_window(self, window_fnames): + if window_fnames is None: return None + window = [] + for fname in window_fnames: + img = cv2.imread(fname) + if img is None: + return None + try: + img = cv2.resize(img, (hparams.img_size, hparams.img_size)) + except Exception as e: + return None + + window.append(img) + + return window + + def crop_audio_window(self, spec, start_frame): + if type(start_frame) == int: + start_frame_num = start_frame + else: + start_frame_num = self.get_frame_id(start_frame) # 0-indexing ---> 1-indexing + start_idx = int(80. * (start_frame_num / float(hparams.fps))) + + end_idx = start_idx + syncnet_mel_step_size + + return spec[start_idx : end_idx, :] + + def get_segmented_mels(self, spec, start_frame): + mels = [] + assert syncnet_T == 5 + start_frame_num = self.get_frame_id(start_frame) + 1 # 0-indexing ---> 1-indexing + if start_frame_num - 2 < 0: return None + for i in range(start_frame_num, start_frame_num + syncnet_T): + m = self.crop_audio_window(spec, i - 2) + if m.shape[0] != syncnet_mel_step_size: + return None + mels.append(m.T) + + mels = np.asarray(mels) + + return mels + + def prepare_window(self, window): + # 3 x T x H x W + x = np.asarray(window) / 255. + x = np.transpose(x, (3, 0, 1, 2)) + + return x + + def __len__(self): + return len(self.all_videos) + + def __getitem__(self, idx): + while 1: + idx = random.randint(0, len(self.all_videos) - 1) + vidname = self.all_videos[idx] + img_names = list(glob(join(vidname, '*.jpg'))) + if len(img_names) <= 3 * syncnet_T: + continue + + img_name = random.choice(img_names) + wrong_img_name = random.choice(img_names) + while wrong_img_name == img_name: + wrong_img_name = random.choice(img_names) + + window_fnames = self.get_window(img_name) + wrong_window_fnames = self.get_window(wrong_img_name) + if window_fnames is None or wrong_window_fnames is None: + continue + + window = self.read_window(window_fnames) + if window is None: + continue + + wrong_window = self.read_window(wrong_window_fnames) + if wrong_window is None: + continue + + try: + wavpath = join(vidname, "audio.wav") + wav = audio.load_wav(wavpath, hparams.sample_rate) + + orig_mel = audio.melspectrogram(wav).T + except Exception as e: + continue + + mel = self.crop_audio_window(orig_mel.copy(), img_name) + + if (mel.shape[0] != syncnet_mel_step_size): + continue + + indiv_mels = self.get_segmented_mels(orig_mel.copy(), img_name) + if indiv_mels is None: continue + + window = self.prepare_window(window) + y = window.copy() + window[:, :, window.shape[2]//2:] = 0. + + wrong_window = self.prepare_window(wrong_window) + x = np.concatenate([window, wrong_window], axis=0) + + x = torch.FloatTensor(x) + mel = torch.FloatTensor(mel.T).unsqueeze(0) + indiv_mels = torch.FloatTensor(indiv_mels).unsqueeze(1) + y = torch.FloatTensor(y) + return x, indiv_mels, mel, y + +def save_sample_images(x, g, gt, global_step, checkpoint_dir): + x = (x.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8) + g = (g.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8) + gt = (gt.detach().cpu().numpy().transpose(0, 2, 3, 4, 1) * 255.).astype(np.uint8) + + refs, inps = x[..., 3:], x[..., :3] + folder = join(checkpoint_dir, "samples_step{:09d}".format(global_step)) + if not os.path.exists(folder): os.mkdir(folder) + collage = np.concatenate((refs, inps, g, gt), axis=-2) + for batch_idx, c in enumerate(collage): + for t in range(len(c)): + cv2.imwrite('{}/{}_{}.jpg'.format(folder, batch_idx, t), c[t]) + +logloss = nn.BCELoss() +def cosine_loss(a, v, y): + d = nn.functional.cosine_similarity(a, v) + loss = logloss(d.unsqueeze(1), y) + + return loss + +device = torch.device("cuda" if use_cuda else "cpu") +syncnet = SyncNet().to(device) +for p in syncnet.parameters(): + p.requires_grad = False + +recon_loss = nn.L1Loss() +def get_sync_loss(mel, g): + g = g[:, :, :, g.size(3)//2:] + g = torch.cat([g[:, :, i] for i in range(syncnet_T)], dim=1) + # B, 3 * T, H//2, W + a, v = syncnet(mel, g) + y = torch.ones(g.size(0), 1).float().to(device) + return cosine_loss(a, v, y) + +def train(device, model, train_data_loader, test_data_loader, optimizer, + checkpoint_dir=None, checkpoint_interval=None, nepochs=None): + + global global_step, global_epoch + resumed_step = global_step + + while global_epoch < nepochs: + print('Starting Epoch: {}'.format(global_epoch)) + running_sync_loss, running_l1_loss = 0., 0. + prog_bar = tqdm(enumerate(train_data_loader)) + for step, (x, indiv_mels, mel, gt) in prog_bar: + model.train() + optimizer.zero_grad() + + # Move data to CUDA device + x = x.to(device) + mel = mel.to(device) + indiv_mels = indiv_mels.to(device) + gt = gt.to(device) + + g = model(indiv_mels, x) + + if hparams.syncnet_wt > 0.: + sync_loss = get_sync_loss(mel, g) + else: + sync_loss = 0. + + l1loss = recon_loss(g, gt) + + loss = hparams.syncnet_wt * sync_loss + (1 - hparams.syncnet_wt) * l1loss + loss.backward() + optimizer.step() + + if global_step % checkpoint_interval == 0: + save_sample_images(x, g, gt, global_step, checkpoint_dir) + + global_step += 1 + cur_session_steps = global_step - resumed_step + + running_l1_loss += l1loss.item() + if hparams.syncnet_wt > 0.: + running_sync_loss += sync_loss.item() + else: + running_sync_loss += 0. + + if global_step == 1 or global_step % checkpoint_interval == 0: + save_checkpoint( + model, optimizer, global_step, checkpoint_dir, global_epoch) + + if global_step == 1 or global_step % hparams.eval_interval == 0: + with torch.no_grad(): + average_sync_loss = eval_model(test_data_loader, global_step, device, model, checkpoint_dir) + + if average_sync_loss < .75: + hparams.set_hparam('syncnet_wt', 0.01) # without image GAN a lesser weight is sufficient + + prog_bar.set_description('L1: {}, Sync Loss: {}'.format(running_l1_loss / (step + 1), + running_sync_loss / (step + 1))) + + global_epoch += 1 + + +def eval_model(test_data_loader, global_step, device, model, checkpoint_dir): + eval_steps = 700 + print('Evaluating for {} steps'.format(eval_steps)) + sync_losses, recon_losses = [], [] + step = 0 + while 1: + for x, indiv_mels, mel, gt in test_data_loader: + step += 1 + model.eval() + + # Move data to CUDA device + x = x.to(device) + gt = gt.to(device) + indiv_mels = indiv_mels.to(device) + mel = mel.to(device) + + g = model(indiv_mels, x) + + sync_loss = get_sync_loss(mel, g) + l1loss = recon_loss(g, gt) + + sync_losses.append(sync_loss.item()) + recon_losses.append(l1loss.item()) + + if step > eval_steps: + averaged_sync_loss = sum(sync_losses) / len(sync_losses) + averaged_recon_loss = sum(recon_losses) / len(recon_losses) + + print('L1: {}, Sync loss: {}'.format(averaged_recon_loss, averaged_sync_loss)) + + return averaged_sync_loss + +def save_checkpoint(model, optimizer, step, checkpoint_dir, epoch): + + checkpoint_path = join( + checkpoint_dir, "checkpoint_step{:09d}.pth".format(global_step)) + optimizer_state = optimizer.state_dict() if hparams.save_optimizer_state else None + torch.save({ + "state_dict": model.state_dict(), + "optimizer": optimizer_state, + "global_step": step, + "global_epoch": epoch, + }, checkpoint_path) + print("Saved checkpoint:", checkpoint_path) + + +def _load(checkpoint_path): + if use_cuda: + checkpoint = torch.load(checkpoint_path) + else: + checkpoint = torch.load(checkpoint_path, + map_location=lambda storage, loc: storage) + return checkpoint + +def load_checkpoint(path, model, optimizer, reset_optimizer=False, overwrite_global_states=True): + global global_step + global global_epoch + + print("Load checkpoint from: {}".format(path)) + checkpoint = _load(path) + s = checkpoint["state_dict"] + new_s = {} + for k, v in s.items(): + new_s[k.replace('module.', '')] = v + model.load_state_dict(new_s) + if not reset_optimizer: + optimizer_state = checkpoint["optimizer"] + if optimizer_state is not None: + print("Load optimizer state from {}".format(path)) + optimizer.load_state_dict(checkpoint["optimizer"]) + if overwrite_global_states: + global_step = checkpoint["global_step"] + global_epoch = checkpoint["global_epoch"] + + return model + +if __name__ == "__main__": + checkpoint_dir = args.checkpoint_dir + + # Dataset and Dataloader setup + train_dataset = Dataset('train') + test_dataset = Dataset('val') + + train_data_loader = data_utils.DataLoader( + train_dataset, batch_size=hparams.batch_size, shuffle=True, + num_workers=hparams.num_workers) + + test_data_loader = data_utils.DataLoader( + test_dataset, batch_size=hparams.batch_size, + num_workers=4) + + device = torch.device("cuda" if use_cuda else "cpu") + + # Model + model = Wav2Lip().to(device) + print('total trainable params {}'.format(sum(p.numel() for p in model.parameters() if p.requires_grad))) + + optimizer = optim.Adam([p for p in model.parameters() if p.requires_grad], + lr=hparams.initial_learning_rate) + + if args.checkpoint_path is not None: + load_checkpoint(args.checkpoint_path, model, optimizer, reset_optimizer=False) + + load_checkpoint(args.syncnet_checkpoint_path, syncnet, None, reset_optimizer=True, overwrite_global_states=False) + + if not os.path.exists(checkpoint_dir): + os.mkdir(checkpoint_dir) + + # Train! + train(device, model, train_data_loader, test_data_loader, optimizer, + checkpoint_dir=checkpoint_dir, + checkpoint_interval=hparams.checkpoint_interval, + nepochs=hparams.nepochs) diff --git a/models/tts/ControllableTalkNet/Dockerfile b/models/tts/ControllableTalkNet/Dockerfile new file mode 100644 index 0000000000..be6e6b1481 --- /dev/null +++ b/models/tts/ControllableTalkNet/Dockerfile @@ -0,0 +1,29 @@ +FROM nvidia/cuda:11.0.3-base-ubuntu20.04 +ENV TZ=Etc/GMT +RUN ln -snf /usr/share/zoneinfo/$TZ /etc/localtime && echo $TZ > /etc/timezone +RUN apt-get update && apt-get install -y git ffmpeg python3.8 python3-pip curl +RUN git clone -q https://github.com/SortAnon/ControllableTalkNet /talknet +RUN git clone -q https://github.com/SortAnon/hifi-gan /talknet/hifi-gan +RUN python3.8 -m pip --no-cache-dir install -r "/talknet/requirements.txt" -f https://download.pytorch.org/whl/torch_stable.html +RUN python3.8 -m pip --no-cache-dir install git+https://github.com/SortAnon/NeMo.git +RUN python3.8 -m pip uninstall -y pesq +RUN python3.8 -m pip install pesq==0.0.2 +RUN touch /talknet/is_docker +WORKDIR /talknet +EXPOSE 8050 + +RUN printf "#!/bin/bash \ + \necho Updating TalkNet... \ + \ngit -C /talknet reset --hard origin/main -q \ + \ngit -C /talknet pull origin main -q \ + \necho Updating HiFi-GAN... \ + \ngit -C /talknet reset --hard origin/main -q \ + \ngit -C /talknet pull origin main -q \ + \necho Updating Python dependencies... \ + \npython3.8 -m pip --quiet --no-cache-dir install -r /talknet/requirements.txt -f https://download.pytorch.org/whl/torch_stable.html \ + \npython3.8 -m pip --quiet --no-cache-dir install git+https://github.com/SortAnon/NeMo.git \ + \necho Launching TalkNet... \ + \npython3.8 talknet_offline.py\n" > /talknet/docker_launch.sh + +RUN chmod +x /talknet/docker_launch.sh +CMD ["./docker_launch.sh"] diff --git a/models/tts/ControllableTalkNet/LICENSE b/models/tts/ControllableTalkNet/LICENSE new file mode 100644 index 0000000000..0e259d42c9 --- /dev/null +++ b/models/tts/ControllableTalkNet/LICENSE @@ -0,0 +1,121 @@ +Creative Commons Legal Code + +CC0 1.0 Universal + + CREATIVE COMMONS CORPORATION IS NOT A LAW FIRM AND DOES NOT PROVIDE + LEGAL SERVICES. DISTRIBUTION OF THIS DOCUMENT DOES NOT CREATE AN + ATTORNEY-CLIENT RELATIONSHIP. CREATIVE COMMONS PROVIDES THIS + INFORMATION ON AN "AS-IS" BASIS. CREATIVE COMMONS MAKES NO WARRANTIES + REGARDING THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS + PROVIDED HEREUNDER, AND DISCLAIMS LIABILITY FOR DAMAGES RESULTING FROM + THE USE OF THIS DOCUMENT OR THE INFORMATION OR WORKS PROVIDED + HEREUNDER. + +Statement of Purpose + +The laws of most jurisdictions throughout the world automatically confer +exclusive Copyright and Related Rights (defined below) upon the creator +and subsequent owner(s) (each and all, an "owner") of an original work of +authorship and/or a database (each, a "Work"). + +Certain owners wish to permanently relinquish those rights to a Work for +the purpose of contributing to a commons of creative, cultural and +scientific works ("Commons") that the public can reliably and without fear +of later claims of infringement build upon, modify, incorporate in other +works, reuse and redistribute as freely as possible in any form whatsoever +and for any purposes, including without limitation commercial purposes. +These owners may contribute to the Commons to promote the ideal of a free +culture and the further production of creative, cultural and scientific +works, or to gain reputation or greater distribution for their Work in +part through the use and efforts of others. + +For these and/or other purposes and motivations, and without any +expectation of additional consideration or compensation, the person +associating CC0 with a Work (the "Affirmer"), to the extent that he or she +is an owner of Copyright and Related Rights in the Work, voluntarily +elects to apply CC0 to the Work and publicly distribute the Work under its +terms, with knowledge of his or her Copyright and Related Rights in the +Work and the meaning and intended legal effect of CC0 on those rights. + +1. Copyright and Related Rights. A Work made available under CC0 may be +protected by copyright and related or neighboring rights ("Copyright and +Related Rights"). Copyright and Related Rights include, but are not +limited to, the following: + + i. the right to reproduce, adapt, distribute, perform, display, + communicate, and translate a Work; + ii. moral rights retained by the original author(s) and/or performer(s); +iii. publicity and privacy rights pertaining to a person's image or + likeness depicted in a Work; + iv. rights protecting against unfair competition in regards to a Work, + subject to the limitations in paragraph 4(a), below; + v. rights protecting the extraction, dissemination, use and reuse of data + in a Work; + vi. database rights (such as those arising under Directive 96/9/EC of the + European Parliament and of the Council of 11 March 1996 on the legal + protection of databases, and under any national implementation + thereof, including any amended or successor version of such + directive); and +vii. other similar, equivalent or corresponding rights throughout the + world based on applicable law or treaty, and any national + implementations thereof. + +2. Waiver. To the greatest extent permitted by, but not in contravention +of, applicable law, Affirmer hereby overtly, fully, permanently, +irrevocably and unconditionally waives, abandons, and surrenders all of +Affirmer's Copyright and Related Rights and associated claims and causes +of action, whether now known or unknown (including existing as well as +future claims and causes of action), in the Work (i) in all territories +worldwide, (ii) for the maximum duration provided by applicable law or +treaty (including future time extensions), (iii) in any current or future +medium and for any number of copies, and (iv) for any purpose whatsoever, +including without limitation commercial, advertising or promotional +purposes (the "Waiver"). Affirmer makes the Waiver for the benefit of each +member of the public at large and to the detriment of Affirmer's heirs and +successors, fully intending that such Waiver shall not be subject to +revocation, rescission, cancellation, termination, or any other legal or +equitable action to disrupt the quiet enjoyment of the Work by the public +as contemplated by Affirmer's express Statement of Purpose. + +3. Public License Fallback. Should any part of the Waiver for any reason +be judged legally invalid or ineffective under applicable law, then the +Waiver shall be preserved to the maximum extent permitted taking into +account Affirmer's express Statement of Purpose. In addition, to the +extent the Waiver is so judged Affirmer hereby grants to each affected +person a royalty-free, non transferable, non sublicensable, non exclusive, +irrevocable and unconditional license to exercise Affirmer's Copyright and +Related Rights in the Work (i) in all territories worldwide, (ii) for the +maximum duration provided by applicable law or treaty (including future +time extensions), (iii) in any current or future medium and for any number +of copies, and (iv) for any purpose whatsoever, including without +limitation commercial, advertising or promotional purposes (the +"License"). The License shall be deemed effective as of the date CC0 was +applied by Affirmer to the Work. Should any part of the License for any +reason be judged legally invalid or ineffective under applicable law, such +partial invalidity or ineffectiveness shall not invalidate the remainder +of the License, and in such case Affirmer hereby affirms that he or she +will not (i) exercise any of his or her remaining Copyright and Related +Rights in the Work or (ii) assert any associated claims and causes of +action with respect to the Work, in either case contrary to Affirmer's +express Statement of Purpose. + +4. Limitations and Disclaimers. + + a. No trademark or patent rights held by Affirmer are waived, abandoned, + surrendered, licensed or otherwise affected by this document. + b. Affirmer offers the Work as-is and makes no representations or + warranties of any kind concerning the Work, express, implied, + statutory or otherwise, including without limitation warranties of + title, merchantability, fitness for a particular purpose, non + infringement, or the absence of latent or other defects, accuracy, or + the present or absence of errors, whether or not discoverable, all to + the greatest extent permissible under applicable law. + c. Affirmer disclaims responsibility for clearing rights of other persons + that may apply to the Work or any use thereof, including without + limitation any person's Copyright and Related Rights in the Work. + Further, Affirmer disclaims responsibility for obtaining any necessary + consents, permissions or other rights required for any use of the + Work. + d. Affirmer understands and acknowledges that Creative Commons is not a + party to this document and has no duty or obligation with respect to + this CC0 or use of the Work. diff --git a/models/tts/ControllableTalkNet/README.md b/models/tts/ControllableTalkNet/README.md new file mode 100644 index 0000000000..20f0d0429a --- /dev/null +++ b/models/tts/ControllableTalkNet/README.md @@ -0,0 +1,26 @@ +# Controllable TalkNet +Controllable TalkNet is a web application that lets you synthesize speech, +which mimics the pitch and pacing of an existing audio clip. It's based on [NVIDIA's implementation](https://github.com/NVIDIA/NeMo) +of TalkNet 2, with some changes to support singing synthesis and higher audio quality. + +## Requirements +* A Google account to run Colab, or... +* An NVIDIA GPU with 4+ GB of VRAM +* 10 GB of free space + +## How to run +### Google Colab +* [Go to the Colab notebook](https://colab.research.google.com/drive/1aj6Jk8cpRw7SsN3JSYCv57CrR6s0gYPB) and follow the instructions. + +### TalkNet Offline (Windows) +* [Download the setup script](https://github.com/SortAnon/ControllableTalkNet/releases/latest/download/TalkNetOffline.zip) +and extract it to a folder. +* Run setup.bat. The initial setup will take about 20 minutes. +* When it's done, run talknet.bat to start TalkNet on http://127.0.0.1:8050/. To download updates, run update.bat. + +### Docker (Linux) +* Install Docker and NVIDIA Container Toolkit. +* [Download the Dockerfile.](https://raw.githubusercontent.com/SortAnon/ControllableTalkNet/main/Dockerfile) +Open a terminal, and navigate to the directory where you saved it. +* Run ```docker build -t talknet-offline .``` to build the image. Add ```sudo``` if you're not using rootless Docker. +* Run ```docker run -it --gpus all -p 8050:8050 talknet-offline``` to start TalkNet on http://127.0.0.1:8050/. diff --git a/models/tts/ControllableTalkNet/TalkNet_Training_Offline.ipynb b/models/tts/ControllableTalkNet/TalkNet_Training_Offline.ipynb new file mode 100644 index 0000000000..44b8cda03f --- /dev/null +++ b/models/tts/ControllableTalkNet/TalkNet_Training_Offline.ipynb @@ -0,0 +1,925 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "source": [ + "# TalkNet Training (offline)\n", + "\n", + "To train a 22KHz TalkNet, run the cells below and follow the instructions.\n", + "\n", + "The notebook will automatically resume training any models from the last saved checkpoint. If you're resuming from a new session, always re-run step 1 first." + ], + "metadata": { + "id": "Gss5Ox_RNiba" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 1:** Configure training data paths. Create the following and change the paths below:\n", + "* A dataset of .wav files, packaged as a .zip or .tar file\n", + "* Training and validation filelists, in LJSpeech format with relative paths (note: ARPABET transcripts are not supported)\n", + "* An output path for checkpoints" + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "### CHANGE THESE ###\n", + "\n", + "dataset = \"/absolute/path_to_dataset.zip\"\n", + "train_filelist = \"/absolute/train_filelist.txt\"\n", + "val_filelist = \"/absolute/val_filelist.txt\"\n", + "output_dir = \"/absolute/name_of_character\"\n", + "\n", + "### ------------ ###\n", + "\n", + "import os\n", + "import gdown\n", + "import platform\n", + "\n", + "assert os.path.exists(dataset), \"Cannot find dataset\"\n", + "assert os.path.exists(train_filelist), \"Cannot find training filelist\"\n", + "assert os.path.exists(val_filelist), \"Cannot find validation filelist\"\n", + "if not os.path.exists(output_dir):\n", + " os.makedirs(output_dir)\n", + "\n", + "if os.getcwd() != output_dir:\n", + " cwd = os.getcwd()\n", + "else:\n", + " os.chdir(cwd)\n", + " \n", + "# Download pre-trained models\n", + "if not os.path.exists(\"talknet_spect.nemo\"):\n", + " print(\"Downloading pre-trained models...\")\n", + " zip_path = \"TalkNet_Pretrained.zip\"\n", + " if not os.path.exists(zip_path) or os.stat(zip_path).st_size < 100:\n", + " d = 'https://drive.google.com/uc?id='\n", + " gdown.download(d+\"19wSym9mNEnmzLS9XdPlfNAW9_u-mP1hR\", zip_path, quiet=False)\n", + " print(platform.system())\n", + " if \"Windows\" in platform.system():\n", + " !tar -xf {zip_path}\n", + " else:\n", + " !unzip -qo {zip_path}\n", + " os.remove(zip_path)\n", + "\n", + "print(\"OK\")" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "nfSawDUD5tqv" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 2:** Dataset processing, part 1.\n", + "\n", + "If this step fails, try the following:\n", + "* Make sure your filelists are correct. They should have relative \n", + "paths that match the contents of the archive." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "import os\n", + "import shutil\n", + "import sys\n", + "import json\n", + "import nemo\n", + "import torch\n", + "import torchaudio\n", + "import numpy as np\n", + "#from pysptk import sptk\n", + "from pathlib import Path\n", + "from tqdm.notebook import tqdm\n", + "import ffmpeg\n", + "\n", + "def fix_transcripts(inpath):\n", + " found_arpabet = False\n", + " found_grapheme = False\n", + " with open(inpath, \"r\", encoding=\"utf8\") as f:\n", + " lines = f.readlines()\n", + " with open(inpath, \"w\", encoding=\"utf8\") as f:\n", + " for l in lines:\n", + " if l.strip() == \"\":\n", + " continue\n", + " if \"{\" in l:\n", + " if not found_arpabet:\n", + " print(\"Warning: Skipping ARPABET lines (not supported).\")\n", + " found_arpabet = True\n", + " else:\n", + " f.write(l)\n", + " found_grapheme = True\n", + " assert found_grapheme, \"No non-ARPABET lines found in \" + inpath\n", + "\n", + "def generate_json(inpath, outpath):\n", + " output = \"\"\n", + " sample_rate = 22050\n", + " with open(inpath, \"r\", encoding=\"utf8\") as f:\n", + " for l in f.readlines():\n", + " lpath = l.split(\"|\")[0].strip()\n", + " if lpath[:5] != \"wavs/\":\n", + " lpath = \"wavs/\" + lpath\n", + " lpath = os.path.join(output_dir, \"_training_data\", lpath)\n", + " size = os.stat(lpath).st_size\n", + " x = {\n", + " \"audio_filepath\": lpath,\n", + " \"duration\": size / (sample_rate * 2),\n", + " \"text\": l.split(\"|\")[1].strip(),\n", + " }\n", + " output += json.dumps(x) + \"\\n\"\n", + " with open(outpath, \"w\", encoding=\"utf8\") as w:\n", + " w.write(output)\n", + "\n", + "def convert_to_22k(inpath):\n", + " if inpath.strip()[-4:].lower() != \".wav\":\n", + " print(\"Warning: \" + inpath.strip() + \" is not a .wav file!\")\n", + " return\n", + " ffmpeg.input(inpath).output(\n", + " inpath + \"_22k.wav\",\n", + " ar=\"22050\",\n", + " ac=\"1\",\n", + " acodec=\"pcm_s16le\",\n", + " map_metadata=\"-1\",\n", + " fflags=\"+bitexact\",\n", + " ).overwrite_output().run(quiet=True)\n", + " os.remove(inpath)\n", + " os.rename(inpath + \"_22k.wav\", inpath)\n", + "\n", + "# Extract dataset\n", + "os.chdir(output_dir)\n", + "\n", + "if os.path.exists(\"_training_data\"):\n", + " shutil.rmtree(\"_training_data\")\n", + "os.mkdir(\"_training_data\")\n", + "os.chdir(\"_training_data\")\n", + "os.mkdir(\"wavs\")\n", + "os.chdir(\"wavs\")\n", + "if dataset[-4:] == \".zip\":\n", + " if \"Windows\" in platform.system():\n", + " !tar -xf \"{dataset}\"\n", + " else:\n", + " !unzip -qo \"{dataset}\"\n", + "elif dataset[-4:] == \".tar\":\n", + " !tar -xf \"{dataset}\"\n", + "else:\n", + " raise Exception(\"Unknown extension for dataset\")\n", + "if os.path.exists(os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\")):\n", + " shutil.move(\n", + " os.path.join(output_dir, \"_training_data\", \"wavs\", \"wavs\"), \n", + " os.path.join(output_dir, \"_training_data\", \"tempwavs\")\n", + " )\n", + " shutil.rmtree(os.path.join(output_dir, \"_training_data\", \"wavs\"))\n", + " shutil.move(\n", + " os.path.join(output_dir, \"_training_data\", \"tempwavs\"), \n", + " os.path.join(output_dir, \"_training_data\", \"wavs\")\n", + " )\n", + "# Filelist for preprocessing\n", + "os.chdir(output_dir)\n", + "shutil.copy(train_filelist, \"trainfiles.txt\")\n", + "shutil.copy(val_filelist, \"valfiles.txt\")\n", + "fix_transcripts(\"trainfiles.txt\")\n", + "fix_transcripts(\"valfiles.txt\")\n", + "seen_files = []\n", + "with open(\"trainfiles.txt\") as f:\n", + " t = f.read().split(\"\\n\")\n", + "with open(\"valfiles.txt\") as f:\n", + " v = f.read().split(\"\\n\")\n", + " all_filelist = t[:] + v[:]\n", + "with open(\"allfiles.txt\", \"w\") as f:\n", + " for x in all_filelist:\n", + " if x.strip() == \"\":\n", + " continue\n", + " if x.split(\"|\")[0] not in seen_files:\n", + " seen_files.append(x.split(\"|\")[0])\n", + " f.write(x.strip() + \"\\n\")\n", + "\n", + "# Ensure audio is 22k\n", + "print(\"Converting audio...\")\n", + "for r, _, f in os.walk(os.path.join(\"_training_data\", \"wavs\")):\n", + " for name in tqdm(f):\n", + " convert_to_22k(os.path.join(r, name))\n", + "\n", + "# Convert to JSON\n", + "generate_json(\"trainfiles.txt\", \"trainfiles.json\")\n", + "generate_json(\"valfiles.txt\", \"valfiles.json\")\n", + "generate_json(\"allfiles.txt\", \"allfiles.json\")\n", + "\n", + "print(\"OK\")" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "bxFr3Fdi_kOC" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 3:** Dataset processing, part 2. This takes a while, but\n", + "you only have to run this once per dataset (results are saved to Drive).\n", + "\n", + "If this step fails, try the following:\n", + "* Make sure your dataset only contains WAV files." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "# Extract phoneme duration\n", + "\n", + "import json\n", + "from nemo.collections.asr.models import EncDecCTCModel\n", + "asr_model = EncDecCTCModel.from_pretrained(model_name=\"asr_talknet_aligner\").cpu().eval()\n", + "\n", + "def forward_extractor(tokens, log_probs, blank):\n", + " \"\"\"Computes states f and p.\"\"\"\n", + " n, m = len(tokens), log_probs.shape[0]\n", + " # `f[s, t]` -- max sum of log probs for `s` first codes\n", + " # with `t` first timesteps with ending in `tokens[s]`.\n", + " f = np.empty((n + 1, m + 1), dtype=float)\n", + " f.fill(-(10 ** 9))\n", + " p = np.empty((n + 1, m + 1), dtype=int)\n", + " f[0, 0] = 0.0 # Start\n", + " for s in range(1, n + 1):\n", + " c = tokens[s - 1]\n", + " for t in range((s + 1) // 2, m + 1):\n", + " f[s, t] = log_probs[t - 1, c]\n", + " # Option #1: prev char is equal to current one.\n", + " if s == 1 or c == blank or c == tokens[s - 3]:\n", + " options = f[s : (s - 2 if s > 1 else None) : -1, t - 1]\n", + " else: # Is not equal to current one.\n", + " options = f[s : (s - 3 if s > 2 else None) : -1, t - 1]\n", + " f[s, t] += np.max(options)\n", + " p[s, t] = np.argmax(options)\n", + " return f, p\n", + "\n", + "\n", + "def backward_extractor(f, p):\n", + " \"\"\"Computes durs from f and p.\"\"\"\n", + " n, m = f.shape\n", + " n -= 1\n", + " m -= 1\n", + " durs = np.zeros(n, dtype=int)\n", + " if f[-1, -1] >= f[-2, -1]:\n", + " s, t = n, m\n", + " else:\n", + " s, t = n - 1, m\n", + " while s > 0:\n", + " durs[s - 1] += 1\n", + " s -= p[s, t]\n", + " t -= 1\n", + " assert durs.shape[0] == n\n", + " assert np.sum(durs) == m\n", + " assert np.all(durs[1::2] > 0)\n", + " return durs\n", + "\n", + "def preprocess_tokens(tokens, blank):\n", + " new_tokens = [blank]\n", + " for c in tokens:\n", + " new_tokens.extend([c, blank])\n", + " tokens = new_tokens\n", + " return tokens\n", + "\n", + "data_config = {\n", + " 'manifest_filepath': \"allfiles.json\",\n", + " 'sample_rate': 22050,\n", + " 'labels': asr_model.decoder.vocabulary,\n", + " 'batch_size': 1,\n", + "}\n", + "\n", + "parser = nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset.make_vocab(\n", + " notation='phonemes', punct=True, spaces=True, stresses=False, add_blank_at=\"last\"\n", + ")\n", + "\n", + "dataset = nemo.collections.asr.data.audio_to_text._AudioTextDataset(\n", + " manifest_filepath=data_config['manifest_filepath'], sample_rate=data_config['sample_rate'], parser=parser,\n", + ")\n", + "\n", + "dl = torch.utils.data.DataLoader(\n", + " dataset=dataset, batch_size=data_config['batch_size'], collate_fn=dataset.collate_fn, shuffle=False,\n", + ")\n", + "\n", + "blank_id = asr_model.decoder.num_classes_with_blank - 1\n", + "\n", + "if os.path.exists(os.path.join(output_dir, \"durations.pt\")):\n", + " print(\"durations.pt already exists; skipping\")\n", + "else:\n", + " dur_data = {}\n", + " for sample_idx, test_sample in tqdm(enumerate(dl), total=len(dl)):\n", + " log_probs, _, greedy_predictions = asr_model(\n", + " input_signal=test_sample[0], input_signal_length=test_sample[1]\n", + " )\n", + "\n", + " log_probs = log_probs[0].cpu().detach().numpy()\n", + " seq_ids = test_sample[2][0].cpu().detach().numpy()\n", + "\n", + " target_tokens = preprocess_tokens(seq_ids, blank_id)\n", + "\n", + " f, p = forward_extractor(target_tokens, log_probs, blank_id)\n", + " durs = backward_extractor(f, p)\n", + "\n", + " dur_key = Path(dl.dataset.collection[sample_idx].audio_file).stem\n", + " dur_data[dur_key] = {\n", + " 'blanks': torch.tensor(durs[::2], dtype=torch.long).cpu().detach(), \n", + " 'tokens': torch.tensor(durs[1::2], dtype=torch.long).cpu().detach()\n", + " }\n", + "\n", + " del test_sample\n", + "\n", + " torch.save(dur_data, os.path.join(output_dir, \"durations.pt\"))\n", + "\n", + "#Extract F0 (pitch)\n", + "import crepe\n", + "from scipy.io import wavfile\n", + "\n", + "def crepe_f0(audio_file, hop_length=256):\n", + " sr, audio = wavfile.read(audio_file)\n", + " audio_x = np.arange(0, len(audio)) / 22050.0\n", + " time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True)\n", + "\n", + " x = np.arange(0, len(audio), hop_length) / 22050.0\n", + " freq_interp = np.interp(x, time, frequency)\n", + " conf_interp = np.interp(x, time, confidence)\n", + " audio_interp = np.interp(x, audio_x, np.absolute(audio)) / 32768.0\n", + " weights = [0.5, 0.25, 0.25]\n", + " audio_smooth = np.convolve(audio_interp, np.array(weights)[::-1], \"same\")\n", + "\n", + " conf_threshold = 0.25\n", + " audio_threshold = 0.0005\n", + " for i in range(len(freq_interp)):\n", + " if conf_interp[i] < conf_threshold:\n", + " freq_interp[i] = 0.0\n", + " if audio_smooth[i] < audio_threshold:\n", + " freq_interp[i] = 0.0\n", + "\n", + " # Hack to make f0 and mel lengths equal\n", + " if len(audio) % hop_length == 0:\n", + " freq_interp = np.pad(freq_interp, pad_width=[0, 1])\n", + " return torch.from_numpy(freq_interp.astype(np.float32))\n", + "\n", + "if os.path.exists(os.path.join(output_dir, \"f0s.pt\")):\n", + " print(\"f0s.pt already exists; skipping\")\n", + "else:\n", + " f0_data = {}\n", + " with open(\"allfiles.json\") as f:\n", + " for i, l in enumerate(f.readlines()):\n", + " print(str(i))\n", + " audio_path = json.loads(l)[\"audio_filepath\"]\n", + " f0_data[Path(audio_path).stem] = crepe_f0(audio_path)\n", + "\n", + " # calculate f0 stats (mean & std) only for train set\n", + " with open(\"trainfiles.json\") as f:\n", + " train_ids = {Path(json.loads(l)[\"audio_filepath\"]).stem for l in f}\n", + " all_f0 = torch.cat([f0[f0 >= 1e-5] for f0_id, f0 in f0_data.items() if f0_id in train_ids])\n", + "\n", + " F0_MEAN, F0_STD = all_f0.mean().item(), all_f0.std().item() \n", + " print(\"F0_MEAN: \" + str(F0_MEAN) + \", F0_STD: \" + str(F0_STD))\n", + " torch.save(f0_data, os.path.join(output_dir, \"f0s.pt\"))\n", + " with open(os.path.join(output_dir, \"f0_info.json\"), \"w\") as f:\n", + " f.write(json.dumps({\"FO_MEAN\": F0_MEAN, \"F0_STD\": F0_STD}))\n", + "\n", + "print(\"OK\")" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "sos9vsxPkIN7" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 4:** Train duration predictor.\n", + "\n", + "If CUDA runs out of memory, try the following:\n", + "* Click on Kernel -> Restart, re-run step 1, and try again.\n", + "* If that doesn't help, reduce the batch size (default 64)." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "batch_size = 64\n", + "\n", + "import os\n", + "from hydra.experimental import compose, initialize\n", + "from hydra.core.global_hydra import GlobalHydra\n", + "from omegaconf import OmegaConf\n", + "import pytorch_lightning as pl\n", + "from nemo.collections.common.callbacks import LogEpochTimeCallback\n", + "from nemo.collections.tts.models import TalkNetDursModel\n", + "from nemo.core.config import hydra_runner\n", + "from nemo.utils.exp_manager import exp_manager\n", + "\n", + "epochs = 20\n", + "learning_rate = 1e-3\n", + "min_learning_rate = 3e-6\n", + "load_checkpoints = True\n", + "\n", + "def train(cfg):\n", + " cfg.sample_rate = 22050\n", + " cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n", + " cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n", + " cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n", + " cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n", + " cfg.trainer.accelerator = \"dp\"\n", + " cfg.trainer.max_epochs = epochs\n", + " cfg.trainer.check_val_every_n_epoch = 5\n", + " cfg.model.train_ds.dataloader_params.batch_size = batch_size\n", + " cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n", + " cfg.model.optim.lr = learning_rate\n", + " cfg.model.optim.sched.min_lr = min_learning_rate\n", + " cfg.exp_manager.exp_dir = output_dir\n", + "\n", + " # Find checkpoints\n", + " ckpt_path = \"\"\n", + " if load_checkpoints:\n", + " path0 = os.path.join(output_dir, \"TalkNetDurs\")\n", + " if os.path.exists(path0):\n", + " path1 = sorted(os.listdir(path0))\n", + " for i in range(len(path1)):\n", + " path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n", + " if os.path.exists(path2):\n", + " match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n", + " if len(match) > 0:\n", + " ckpt_path = os.path.join(path2, match[0])\n", + " print(\"Resuming training from \" + match[0])\n", + " break\n", + " \n", + " if ckpt_path != \"\":\n", + " trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n", + " model = TalkNetDursModel(cfg=cfg.model, trainer=trainer)\n", + " else:\n", + " warmstart_path = os.path.join(cwd, \"talknet_durs.nemo\")\n", + " trainer = pl.Trainer(**cfg.trainer)\n", + " model = TalkNetDursModel.restore_from(warmstart_path, override_config_path=cfg)\n", + " model.set_trainer(trainer)\n", + " model.setup_training_data(cfg.model.train_ds)\n", + " model.setup_validation_data(cfg.model.validation_ds)\n", + " model.setup_optimization(cfg.model.optim)\n", + " print(\"Warm-starting from \" + warmstart_path)\n", + " exp_manager(trainer, cfg.get('exp_manager', None))\n", + " trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n", + " trainer.fit(model)\n", + "\n", + "os.chdir(cwd)\n", + "GlobalHydra().clear()\n", + "initialize(config_path=\"conf\")\n", + "cfg = compose(config_name=\"talknet-durs\")\n", + "train(cfg)\n" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "nM7-bMpKO7U2" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 5:** Train pitch predictor.\n", + "\n", + "If CUDA runs out of memory, try the following:\n", + "* Click on Kernel -> Restart, re-run step 1, and try again.\n", + "* If that doesn't help, reduce the batch size (default 64)." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "batch_size = 64\n", + "epochs = 50\n", + "\n", + "import json\n", + "\n", + "with open(os.path.join(output_dir, \"f0_info.json\"), \"r\") as f:\n", + " f0_info = json.load(f)\n", + " f0_mean = f0_info[\"FO_MEAN\"]\n", + " f0_std = f0_info[\"F0_STD\"]\n", + "\n", + "learning_rate = 1e-3\n", + "min_learning_rate = 3e-6\n", + "load_checkpoints = True\n", + "\n", + "import os\n", + "from hydra.experimental import compose, initialize\n", + "from hydra.core.global_hydra import GlobalHydra\n", + "from omegaconf import OmegaConf\n", + "import pytorch_lightning as pl\n", + "from nemo.collections.common.callbacks import LogEpochTimeCallback\n", + "from nemo.collections.tts.models import TalkNetPitchModel\n", + "from nemo.core.config import hydra_runner\n", + "from nemo.utils.exp_manager import exp_manager\n", + "\n", + "def train(cfg):\n", + " cfg.sample_rate = 22050\n", + " cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n", + " cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n", + " cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n", + " cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n", + " cfg.trainer.accelerator = \"dp\"\n", + " cfg.trainer.max_epochs = epochs\n", + " cfg.trainer.check_val_every_n_epoch = 5\n", + " cfg.model.f0_mean=f0_mean\n", + " cfg.model.f0_std=f0_std\n", + " cfg.model.train_ds.dataloader_params.batch_size = batch_size\n", + " cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n", + " cfg.model.optim.lr = learning_rate\n", + " cfg.model.optim.sched.min_lr = min_learning_rate\n", + " cfg.exp_manager.exp_dir = output_dir\n", + "\n", + " # Find checkpoints\n", + " ckpt_path = \"\"\n", + " if load_checkpoints:\n", + " path0 = os.path.join(output_dir, \"TalkNetPitch\")\n", + " if os.path.exists(path0):\n", + " path1 = sorted(os.listdir(path0))\n", + " for i in range(len(path1)):\n", + " path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n", + " if os.path.exists(path2):\n", + " match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n", + " if len(match) > 0:\n", + " ckpt_path = os.path.join(path2, match[0])\n", + " print(\"Resuming training from \" + match[0])\n", + " break\n", + " \n", + " if ckpt_path != \"\":\n", + " trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n", + " model = TalkNetPitchModel(cfg=cfg.model, trainer=trainer)\n", + " else:\n", + " warmstart_path = os.path.join(cwd, \"talknet_pitch.nemo\")\n", + " trainer = pl.Trainer(**cfg.trainer)\n", + " model = TalkNetPitchModel.restore_from(warmstart_path, override_config_path=cfg)\n", + " model.set_trainer(trainer)\n", + " model.setup_training_data(cfg.model.train_ds)\n", + " model.setup_validation_data(cfg.model.validation_ds)\n", + " model.setup_optimization(cfg.model.optim)\n", + " print(\"Warm-starting from \" + warmstart_path)\n", + " exp_manager(trainer, cfg.get('exp_manager', None))\n", + " trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n", + " trainer.fit(model)\n", + "\n", + "os.chdir(cwd)\n", + "GlobalHydra().clear()\n", + "initialize(config_path=\"conf\")\n", + "cfg = compose(config_name=\"talknet-pitch\")\n", + "train(cfg)" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "JLfm00NuJfon" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 6:** Train spectrogram generator. 200+ epochs are recommended. \n", + "\n", + "This is the slowest of the three models to train, and the hardest to\n", + "get good results from. If your character sounds noisy or robotic,\n", + "try improving the dataset, or adjusting the epochs and learning rate.\n", + "\n", + "If CUDA runs out of memory, try the following:\n", + "* Click on Kernel -> Restart, re-run step 1, and try again.\n", + "* If that doesn't help, reduce the batch size (default 32)." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "epochs = 200\n", + "batch_size = 32\n", + "\n", + "# Advanced settings. You can probably leave these at their defaults (1e-3, 3e-6, empty, checked).\n", + "learning_rate = 1e-3\n", + "min_learning_rate = 3e-6\n", + "pretrained_path = \"\"\n", + "load_checkpoints = True\n", + "\n", + "import os\n", + "from hydra.experimental import compose, initialize\n", + "from hydra.core.global_hydra import GlobalHydra\n", + "from omegaconf import OmegaConf\n", + "import pytorch_lightning as pl\n", + "from nemo.collections.common.callbacks import LogEpochTimeCallback\n", + "from nemo.collections.tts.models import TalkNetSpectModel\n", + "from nemo.core.config import hydra_runner\n", + "from nemo.utils.exp_manager import exp_manager\n", + "\n", + "def train(cfg):\n", + " cfg.sample_rate = 22050\n", + " cfg.train_dataset = os.path.join(output_dir, \"trainfiles.json\")\n", + " cfg.validation_datasets = os.path.join(output_dir, \"valfiles.json\")\n", + " cfg.durs_file = os.path.join(output_dir, \"durations.pt\")\n", + " cfg.f0_file = os.path.join(output_dir, \"f0s.pt\")\n", + " cfg.trainer.accelerator = \"dp\"\n", + " cfg.trainer.max_epochs = epochs\n", + " cfg.trainer.check_val_every_n_epoch = 5\n", + " cfg.model.train_ds.dataloader_params.batch_size = batch_size\n", + " cfg.model.validation_ds.dataloader_params.batch_size = batch_size\n", + " cfg.model.optim.lr = learning_rate\n", + " cfg.model.optim.sched.min_lr = min_learning_rate\n", + " cfg.exp_manager.exp_dir = output_dir\n", + "\n", + " # Find checkpoints\n", + " ckpt_path = \"\"\n", + " if load_checkpoints:\n", + " path0 = os.path.join(output_dir, \"TalkNetSpect\")\n", + " if os.path.exists(path0):\n", + " path1 = sorted(os.listdir(path0))\n", + " for i in range(len(path1)):\n", + " path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n", + " if os.path.exists(path2):\n", + " match = [x for x in os.listdir(path2) if \"last.ckpt\" in x]\n", + " if len(match) > 0:\n", + " ckpt_path = os.path.join(path2, match[0])\n", + " print(\"Resuming training from \" + match[0])\n", + " break\n", + " \n", + " if ckpt_path != \"\":\n", + " trainer = pl.Trainer(**cfg.trainer, resume_from_checkpoint = ckpt_path)\n", + " model = TalkNetSpectModel(cfg=cfg.model, trainer=trainer)\n", + " else:\n", + " if pretrained_path != \"\":\n", + " warmstart_path = pretrained_path\n", + " else:\n", + " warmstart_path = os.path.join(cwd, \"talknet_spect.nemo\")\n", + " trainer = pl.Trainer(**cfg.trainer)\n", + " model = TalkNetSpectModel.restore_from(warmstart_path, override_config_path=cfg)\n", + " model.set_trainer(trainer)\n", + " model.setup_training_data(cfg.model.train_ds)\n", + " model.setup_validation_data(cfg.model.validation_ds)\n", + " model.setup_optimization(cfg.model.optim)\n", + " print(\"Warm-starting from \" + warmstart_path)\n", + " exp_manager(trainer, cfg.get('exp_manager', None))\n", + " trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa\n", + " trainer.fit(model)\n", + "\n", + "os.chdir(cwd)\n", + "GlobalHydra().clear()\n", + "initialize(config_path=\"conf\")\n", + "cfg = compose(config_name=\"talknet-spect\")\n", + "train(cfg)" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "N9hh4WPHbCcn" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 7:** Generate GTA spectrograms. This will help HiFi-GAN learn what your TalkNet model sounds like.\n", + "\n", + "If this step fails, make sure you've finished training the spectrogram generator." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "import sys\n", + "import os\n", + "import torch\n", + "import numpy as np\n", + "from tqdm import tqdm\n", + "from nemo.collections.tts.models import TalkNetSpectModel\n", + "import shutil\n", + "\n", + "def fix_paths(inpath):\n", + " output = \"\"\n", + " with open(inpath, \"r\", encoding=\"utf8\") as f:\n", + " for l in f.readlines():\n", + " if l[:5].lower() != \"wavs/\":\n", + " output += \"wavs/\" + l\n", + " else:\n", + " output += l\n", + " with open(inpath, \"w\", encoding=\"utf8\") as w:\n", + " w.write(output)\n", + "\n", + "shutil.copyfile(train_filelist, os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n", + "shutil.copyfile(val_filelist, os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n", + "fix_paths(os.path.join(cwd, \"hifi-gan\", \"training.txt\"))\n", + "fix_paths(os.path.join(cwd, \"hifi-gan\", \"validation.txt\"))\n", + "fix_paths(os.path.join(output_dir, \"allfiles.txt\"))\n", + "\n", + "os.chdir(cwd)\n", + "outdir = os.path.join(output_dir, \"_training_data\", \"wavs\")\n", + "if not os.path.exists(outdir):\n", + " os.mkdir(outdir)\n", + "\n", + "model_path = \"\"\n", + "path0 = os.path.join(output_dir, \"TalkNetSpect\")\n", + "if os.path.exists(path0):\n", + " path1 = sorted(os.listdir(path0))\n", + " for i in range(len(path1)):\n", + " path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n", + " if os.path.exists(path2):\n", + " match = [x for x in os.listdir(path2) if \"TalkNetSpect.nemo\" in x]\n", + " if len(match) > 0:\n", + " model_path = os.path.join(path2, match[0])\n", + " break\n", + "assert model_path != \"\", \"TalkNetSpect.nemo not found\"\n", + "\n", + "dur_path = os.path.join(output_dir, \"durations.pt\")\n", + "f0_path = os.path.join(output_dir, \"f0s.pt\")\n", + "\n", + "model = TalkNetSpectModel.restore_from(model_path)\n", + "model.eval()\n", + "with open(os.path.join(output_dir, \"allfiles.txt\"), \"r\", encoding=\"utf-8\") as f:\n", + " dataset = f.readlines()\n", + "durs = torch.load(dur_path)\n", + "f0s = torch.load(f0_path)\n", + "\n", + "for x in tqdm(dataset):\n", + " x_name = os.path.splitext(os.path.basename(x.split(\"|\")[0].strip()))[0]\n", + " x_tokens = model.parse(text=x.split(\"|\")[1].strip())\n", + " x_durs = (\n", + " torch.stack(\n", + " (\n", + " durs[x_name][\"blanks\"],\n", + " torch.cat((durs[x_name][\"tokens\"], torch.zeros(1).int())),\n", + " ),\n", + " dim=1,\n", + " )\n", + " .view(-1)[:-1]\n", + " .view(1, -1)\n", + " .to(\"cuda:0\")\n", + " )\n", + " x_f0s = f0s[x_name].view(1, -1).to(\"cuda:0\")\n", + " x_spect = model.force_spectrogram(tokens=x_tokens, durs=x_durs, f0=x_f0s)\n", + " rel_path = os.path.splitext(x.split(\"|\")[0].strip())[0][5:]\n", + " abs_dir = os.path.join(outdir, os.path.dirname(rel_path))\n", + " if abs_dir != \"\" and not os.path.exists(abs_dir):\n", + " os.makedirs(abs_dir, exist_ok=True)\n", + " np.save(os.path.join(outdir, rel_path + \".npy\"), x_spect.detach().cpu().numpy())\n" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "Ajfyfz2p9Ior" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 8:** Train HiFi-GAN. 2,000+ steps are recommended.\n", + "Stop this cell to finish training the model.\n", + "\n", + "* Click on Kernel -> Restart, re-run step 1, and try again.\n", + "If this step still fails to start, make sure step 7 finished successfully.\n", + "\n", + "Note: If the training process starts at step 2500000, delete the HiFiGAN folder and try again." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "import gdown\n", + "d = 'https://drive.google.com/uc?id='\n", + "\n", + "os.chdir(os.path.join(cwd, \"hifi-gan\"))\n", + "\n", + "if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\")):\n", + " os.makedirs(os.path.join(output_dir, \"HiFiGAN\"))\n", + "if not os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\")):\n", + " print(\"Downloading universal model...\")\n", + " gdown.download(d+\"1qpgI41wNXFcH-iKq1Y42JlBC9j0je8PW\", os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), quiet=False)\n", + " gdown.download(d+\"1O63eHZR9t1haCdRHQcEgMfMNxiOciSru\", os.path.join(output_dir, \"HiFiGAN\", \"do_00000000\"), quiet=False)\n", + " start_from_universal = \"--warm_start True \"\n", + "else:\n", + " start_from_universal = \"\"\n", + "\n", + "hifi_train = os.path.join(cwd, \"hifi-gan\", \"training.txt\")\n", + "hifi_val = os.path.join(cwd, \"hifi-gan\", \"validation.txt\")\n", + "hifi_wavs = os.path.join(output_dir, \"_training_data\")\n", + "\n", + "!python train.py --fine_tuning True --config config_v1b.json \\\n", + "{start_from_universal} \\\n", + "--checkpoint_interval 250 --checkpoint_path \"{os.path.join(output_dir, 'HiFiGAN')}\" \\\n", + "--input_training_file \"{hifi_train}\" \\\n", + "--input_validation_file \"{hifi_val}\" \\\n", + "--input_wavs_dir \"{hifi_wavs}\"\n" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "yVBjGhRB9hUJ" + } + }, + { + "cell_type": "markdown", + "source": [ + "**Step 9:** Package the models. They'll be saved to the output directory as [character_name]_TalkNet.zip.\n", + "\n", + "When done, upload it to Google Drive, with permissions set to \"Anyone with the link\". \n", + "You can then use it with TalkNet by selecting \"Custom model\" as your character. \n", + "\n", + "This cell will also delete the training checkpoints and logs.\n", + "That should free up roughly 2 GB of space.\n", + "If you wish to keep them, set delete_checkpoints to False." + ], + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": null, + "source": [ + "character_name = \"Character\"\n", + "delete_checkpoints = True\n", + "\n", + "import shutil\n", + "from zipfile import ZipFile\n", + "\n", + "def find_talknet(model_dir):\n", + " ckpt_path = \"\"\n", + " path0 = os.path.join(output_dir, model_dir)\n", + " if os.path.exists(path0):\n", + " path1 = sorted(os.listdir(path0))\n", + " for i in range(len(path1)):\n", + " path2 = os.path.join(path0, path1[-(1+i)], \"checkpoints\")\n", + " if os.path.exists(path2):\n", + " match = [x for x in os.listdir(path2) if \".nemo\" in x]\n", + " if len(match) > 0:\n", + " ckpt_path = os.path.join(path2, match[0])\n", + " break\n", + " assert ckpt_path != \"\", \"Couldn't find \" + model_dir\n", + " return ckpt_path\n", + "\n", + "durs_path = find_talknet(\"TalkNetDurs\")\n", + "pitch_path = find_talknet(\"TalkNetPitch\")\n", + "spect_path = find_talknet(\"TalkNetSpect\")\n", + "assert os.path.exists(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\")), \"Couldn't find HiFi-GAN\"\n", + "\n", + "zip = ZipFile(os.path.join(output_dir, character_name + \"_TalkNet.zip\"), 'w')\n", + "zip.write(durs_path, \"TalkNetDurs.nemo\")\n", + "zip.write(pitch_path, \"TalkNetPitch.nemo\")\n", + "zip.write(spect_path, \"TalkNetSpect.nemo\")\n", + "zip.write(os.path.join(output_dir, \"HiFiGAN\", \"g_00000000\"), \"hifiganmodel\")\n", + "zip.write(os.path.join(output_dir, \"HiFiGAN\", \"config.json\"), \"config.json\")\n", + "zip.write(os.path.join(output_dir, \"f0_info.json\"), \"f0_info.json\")\n", + "zip.close()\n", + "print(\"Archived model to \" + os.path.join(output_dir, character_name + \"_TalkNet.zip\"))\n", + "\n", + "if delete_checkpoints:\n", + " shutil.rmtree((os.path.join(output_dir, \"TalkNetDurs\")))\n", + " shutil.rmtree((os.path.join(output_dir, \"TalkNetPitch\")))\n", + " shutil.rmtree((os.path.join(output_dir, \"TalkNetSpect\")))\n", + " shutil.rmtree((os.path.join(output_dir, \"HiFiGAN\")))\n" + ], + "outputs": [], + "metadata": { + "cellView": "form", + "id": "5OtwfrTT-blU" + } + } + ], + "metadata": { + "accelerator": "GPU", + "colab": { + "collapsed_sections": [], + "name": "TalkNet_Training_Offline.ipynb", + "provenance": [] + }, + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.10" + } + }, + "nbformat": 4, + "nbformat_minor": 2 +} \ No newline at end of file diff --git a/models/tts/ControllableTalkNet/__init__.py b/models/tts/ControllableTalkNet/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/models/tts/ControllableTalkNet/assets/Equestria.woff b/models/tts/ControllableTalkNet/assets/Equestria.woff new file mode 100644 index 0000000000..6234e41d0c Binary files /dev/null and b/models/tts/ControllableTalkNet/assets/Equestria.woff differ diff --git a/models/tts/ControllableTalkNet/assets/dash_customfont.css b/models/tts/ControllableTalkNet/assets/dash_customfont.css new file mode 100644 index 0000000000..c9855c1568 --- /dev/null +++ b/models/tts/ControllableTalkNet/assets/dash_customfont.css @@ -0,0 +1,420 @@ +@font-face { + font-family: 'EquestriaWebfont'; + src: url('Equestria.woff') format('woff'); +} + +/* Table of contents +–––––––––––––––––––––––––––––––––––––––––––––––––– +- Plotly.js +- Grid +- Base Styles +- Typography +- Links +- Buttons +- Forms +- Lists +- Code +- Tables +- Spacing +- Utilities +- Clearing +- Media Queries +*/ + +/* PLotly.js +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +/* plotly.js's modebar's z-index is 1001 by default + * https://github.com/plotly/plotly.js/blob/7e4d8ab164258f6bd48be56589dacd9bdd7fded2/src/css/_modebar.scss#L5 + * In case a dropdown is above the graph, the dropdown's options + * will be rendered below the modebar + * Increase the select option's z-index + */ + +/* This was actually not quite right - + dropdowns were overlapping each other (edited October 26) + +.Select { + z-index: 1002; +}*/ + + +/* Grid +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +.container { + position: relative; + width: 100%; + max-width: 960px; + margin: 0 auto; + padding: 0 20px; + box-sizing: border-box; } +.column, +.columns { + width: 100%; + float: left; + box-sizing: border-box; } + +/* For devices larger than 400px */ +@media (min-width: 400px) { + .container { + width: 85%; + padding: 0; } +} + +/* For devices larger than 550px */ +@media (min-width: 550px) { + .container { + width: 80%; } + .column, + .columns { + margin-left: 4%; } + .column:first-child, + .columns:first-child { + margin-left: 0; } + + .one.column, + .one.columns { width: 4.66666666667%; } + .two.columns { width: 13.3333333333%; } + .three.columns { width: 22%; } + .four.columns { width: 30.6666666667%; } + .five.columns { width: 39.3333333333%; } + .six.columns { width: 48%; } + .seven.columns { width: 56.6666666667%; } + .eight.columns { width: 65.3333333333%; } + .nine.columns { width: 74.0%; } + .ten.columns { width: 82.6666666667%; } + .eleven.columns { width: 91.3333333333%; } + .twelve.columns { width: 100%; margin-left: 0; } + + .one-third.column { width: 30.6666666667%; } + .two-thirds.column { width: 65.3333333333%; } + + .one-half.column { width: 48%; } + + /* Offsets */ + .offset-by-one.column, + .offset-by-one.columns { margin-left: 8.66666666667%; } + .offset-by-two.column, + .offset-by-two.columns { margin-left: 17.3333333333%; } + .offset-by-three.column, + .offset-by-three.columns { margin-left: 26%; } + .offset-by-four.column, + .offset-by-four.columns { margin-left: 34.6666666667%; } + .offset-by-five.column, + .offset-by-five.columns { margin-left: 43.3333333333%; } + .offset-by-six.column, + .offset-by-six.columns { margin-left: 52%; } + .offset-by-seven.column, + .offset-by-seven.columns { margin-left: 60.6666666667%; } + .offset-by-eight.column, + .offset-by-eight.columns { margin-left: 69.3333333333%; } + .offset-by-nine.column, + .offset-by-nine.columns { margin-left: 78.0%; } + .offset-by-ten.column, + .offset-by-ten.columns { margin-left: 86.6666666667%; } + .offset-by-eleven.column, + .offset-by-eleven.columns { margin-left: 95.3333333333%; } + + .offset-by-one-third.column, + .offset-by-one-third.columns { margin-left: 34.6666666667%; } + .offset-by-two-thirds.column, + .offset-by-two-thirds.columns { margin-left: 69.3333333333%; } + + .offset-by-one-half.column, + .offset-by-one-half.columns { margin-left: 52%; } + +} + + +/* Base Styles +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +/* NOTE +html is set to 62.5% so that all the REM measurements throughout Skeleton +are based on 10px sizing. So basically 1.5rem = 15px :) */ +html { + font-size: 62.5%; } +body { + font-size: 1.5em; /* currently ems cause chrome bug misinterpreting rems on body element */ + line-height: 1.6; + font-weight: 400; + font-family: "Open Sans", "HelveticaNeue", "Helvetica Neue", Helvetica, Arial, sans-serif; + color: rgb(50, 50, 50); } + + +/* Typography +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +h1, h2, h3, h4, h5, h6 { + margin-top: 0; + margin-bottom: 0; + font-weight: 300; } +h1 { font-size: 4.5rem; line-height: 1.2; letter-spacing: -.1rem; margin-bottom: 2rem; } +h2 { font-size: 3.6rem; line-height: 1.25; letter-spacing: -.1rem; margin-bottom: 1.8rem; margin-top: 1.8rem;} +h3 { font-size: 3.0rem; line-height: 1.3; letter-spacing: -.1rem; margin-bottom: 1.5rem; margin-top: 1.5rem;} +h4 { font-size: 2.6rem; line-height: 1.35; letter-spacing: -.08rem; margin-bottom: 1.2rem; margin-top: 1.2rem;} +h5 { font-size: 2.2rem; line-height: 1.5; letter-spacing: -.05rem; margin-bottom: 0.6rem; margin-top: 0.6rem;} +h6 { font-size: 2.0rem; line-height: 1.6; letter-spacing: 0; margin-bottom: 0.75rem; margin-top: 0.75rem;} + +p { + margin-top: 0; } + + +/* Blockquotes +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +blockquote { + border-left: 4px lightgrey solid; + padding-left: 1rem; + margin-top: 2rem; + margin-bottom: 2rem; + margin-left: 0rem; +} + + +/* Links +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +a { + color: #1EAEDB; + text-decoration: underline; + cursor: pointer;} +a:hover { + color: #0FA0CE; } + + +/* Buttons +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +.button, +button, +input[type="submit"], +input[type="reset"], +input[type="button"] { + display: inline-block; + height: 38px; + padding: 0 30px; + color: #555; + text-align: center; + font-size: 11px; + font-weight: 600; + line-height: 38px; + letter-spacing: .1rem; + text-transform: uppercase; + text-decoration: none; + white-space: nowrap; + background-color: transparent; + border-radius: 4px; + border: 1px solid #bbb; + cursor: pointer; + box-sizing: border-box; } +.button:hover, +button:hover, +input[type="submit"]:hover, +input[type="reset"]:hover, +input[type="button"]:hover, +.button:focus, +button:focus, +input[type="submit"]:focus, +input[type="reset"]:focus, +input[type="button"]:focus { + color: #333; + border-color: #888; + outline: 0; } +.button.button-primary, +button.button-primary, +input[type="submit"].button-primary, +input[type="reset"].button-primary, +input[type="button"].button-primary { + color: #FFF; + background-color: #33C3F0; + border-color: #33C3F0; } +.button.button-primary:hover, +button.button-primary:hover, +input[type="submit"].button-primary:hover, +input[type="reset"].button-primary:hover, +input[type="button"].button-primary:hover, +.button.button-primary:focus, +button.button-primary:focus, +input[type="submit"].button-primary:focus, +input[type="reset"].button-primary:focus, +input[type="button"].button-primary:focus { + color: #FFF; + background-color: #1EAEDB; + border-color: #1EAEDB; } + + +/* Forms +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +input[type="email"], +input[type="number"], +input[type="search"], +input[type="text"], +input[type="tel"], +input[type="url"], +input[type="password"], +textarea, +select { + height: 38px; + padding: 6px 10px; /* The 6px vertically centers text on FF, ignored by Webkit */ + background-color: #fff; + border: 1px solid #D1D1D1; + border-radius: 4px; + box-shadow: none; + box-sizing: border-box; + font-family: inherit; + font-size: inherit; /*https://stackoverflow.com/questions/6080413/why-doesnt-input-inherit-the-font-from-body*/} +/* Removes awkward default styles on some inputs for iOS */ +input[type="email"], +input[type="number"], +input[type="search"], +input[type="text"], +input[type="tel"], +input[type="url"], +input[type="password"], +textarea { + -webkit-appearance: none; + -moz-appearance: none; + appearance: none; } +textarea { + min-height: 65px; + padding-top: 6px; + padding-bottom: 6px; } +input[type="email"]:focus, +input[type="number"]:focus, +input[type="search"]:focus, +input[type="text"]:focus, +input[type="tel"]:focus, +input[type="url"]:focus, +input[type="password"]:focus, +textarea:focus, +select:focus { + border: 1px solid #33C3F0; + outline: 0; } +label, +legend { + display: block; + margin-bottom: 0px; } +fieldset { + padding: 0; + border-width: 0; } +input[type="checkbox"], +input[type="radio"] { + display: inline; } +label > .label-body { + display: inline-block; + margin-left: .5rem; + font-weight: normal; } + + +/* Lists +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +ul { + list-style: circle inside; } +ol { + list-style: decimal inside; } +ol, ul { + padding-left: 0; + margin-top: 0; } +ul ul, +ul ol, +ol ol, +ol ul { + margin: 1.5rem 0 1.5rem 3rem; + font-size: 90%; } +li { + margin-bottom: 1rem; } + + +/* Tables +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +table { + border-collapse: collapse; +} +th:not(.CalendarDay), +td:not(.CalendarDay) { + padding: 12px 15px; + text-align: left; + border-bottom: 1px solid #E1E1E1; } +th:first-child:not(.CalendarDay), +td:first-child:not(.CalendarDay) { + padding-left: 0; } +th:last-child:not(.CalendarDay), +td:last-child:not(.CalendarDay) { + padding-right: 0; } + + +/* Spacing +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +button, +.button { + margin-bottom: 0rem; } +input, +textarea, +select, +fieldset { + margin-bottom: 0rem; } +pre, +dl, +figure, +table, +form { + margin-bottom: 0rem; } +p, +ul, +ol { + margin-bottom: 0.75rem; } + +/* Utilities +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +.u-full-width { + width: 100%; + box-sizing: border-box; } +.u-max-full-width { + max-width: 100%; + box-sizing: border-box; } +.u-pull-right { + float: right; } +.u-pull-left { + float: left; } + + +/* Misc +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +hr { + margin-top: 3rem; + margin-bottom: 3.5rem; + border-width: 0; + border-top: 1px solid #E1E1E1; } + + +/* Clearing +–––––––––––––––––––––––––––––––––––––––––––––––––– */ + +/* Self Clearing Goodness */ +.container:after, +.row:after, +.u-cf { + content: ""; + display: table; + clear: both; } + + +/* Media Queries +–––––––––––––––––––––––––––––––––––––––––––––––––– */ +/* +Note: The best way to structure the use of media queries is to create the queries +near the relevant code. For example, if you wanted to change the styles for buttons +on small devices, paste the mobile query code up in the buttons section and style it +there. +*/ + + +/* Larger than mobile */ +@media (min-width: 400px) {} + +/* Larger than phablet (also point when grid becomes active) */ +@media (min-width: 550px) {} + +/* Larger than tablet */ +@media (min-width: 750px) {} + +/* Larger than desktop */ +@media (min-width: 1000px) {} + +/* Larger than Desktop HD */ +@media (min-width: 1200px) {} \ No newline at end of file diff --git a/models/tts/ControllableTalkNet/controllable_talknet.py b/models/tts/ControllableTalkNet/controllable_talknet.py new file mode 100644 index 0000000000..f258054d9f --- /dev/null +++ b/models/tts/ControllableTalkNet/controllable_talknet.py @@ -0,0 +1,1097 @@ +import sys +import os +import base64 +import dash +from jupyter_dash import JupyterDash +import dash_core_components as dcc +import dash_html_components as html +import dash_bootstrap_components as dbc +from dash.exceptions import PreventUpdate +import torch +import numpy as np +import tensorflow as tf +import crepe +import scipy +from scipy.io import wavfile +import psola +import io +import nemo +from nemo.collections.asr.models import EncDecCTCModel +from nemo.collections.tts.models import TalkNetSpectModel +from nemo.collections.tts.models import TalkNetPitchModel +from nemo.collections.tts.models import TalkNetDursModel +from ControllableTalkNet.talknet_singer import TalkNetSingerModel +import json +from tqdm import tqdm +import gdown +import zipfile +import resampy +import traceback +import ffmpeg +import time +import uuid + +sys.path.append("hifi-gan") +from hifigan.env import AttrDict +from hifigan.meldataset import mel_spectrogram, MAX_WAV_VALUE +from hifigan.models import Generator +from hifigan.denoiser import Denoiser + +app = JupyterDash(__name__) +DEVICE = "cuda:0" +CPU_PITCH = False +RUN_PATH = os.path.dirname(os.path.realpath(__file__)) +if RUN_PATH == "/content": + UI_MODE = "colab" +elif os.path.exists("/talknet/is_docker"): + UI_MODE = "docker" +else: + UI_MODE = "offline" +torch.set_grad_enabled(False) +if CPU_PITCH: + tf.config.set_visible_devices([], "GPU") + +app.title = "Controllable TalkNet" +app.layout = html.Div( + children=[ + html.H1( + id="header", + children="Controllable TalkNet", + style={ + "font-family": "EquestriaWebfont", + "color": "#ed3c96", + "font-size": "4em", + "text-align": "center", + "margin-top": "0em", + "margin-bottom": "0em", + }, + ), + html.Label("Character selection", htmlFor="model-dropdown"), + dbc.Select( + id="model-dropdown", + options=[ + { + "label": "Custom model", + "value": "Custom", + }, + { + "label": "--- ERROR LOADING MODEL LISTS ---", + "value": "", + "disabled": True, + }, + ], + value=None, + style={ + "max-width": "90vw", + "width": "35em", + "margin-bottom": "0.7em", + }, + ), + html.Div( + children=[ + dcc.Input( + id="drive-id", + type="text", + placeholder="Drive ID for custom model", + style={"width": "22em"}, + ), + ], + id="custom-model", + style={ + "display": "none", + }, + ), + html.Label( + "Upload reference audio to " + RUN_PATH, + htmlFor="reference-dropdown", + id="upload-label", + ), + dcc.Upload( + id="upload-audio", + children=html.Div(["Drag and drop or click to select a file to upload."]), + style={ + "display": "none", + }, + multiple=True, + ), + dcc.Store(id="current-f0s"), + dcc.Store(id="current-f0s-nosilence"), + dcc.Store(id="current-filename"), + dcc.Loading( + id="audio-loading", + children=[ + html.Div( + [ + html.Button( + "Update file list", + id="update-button", + style={ + "margin-right": "10px", + }, + ), + dbc.Select( + id="reference-dropdown", + options=[], + value=None, + style={ + "max-width": "80vw", + "width": "30em", + }, + disabled=False, + ), + dcc.Store(id="pitch-clicks"), + html.Button( + "Debug pitch", + id="pitch-button", + style={ + "margin-left": "10px", + }, + disabled=False, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "row", + "margin-left": "50px", + "vertical-align": "middle", + }, + ), + html.Audio( + id="pitch-out", + controls=True, + style={"display": "none"}, + ), + html.Div( + id="audio-loading-output", + style={ + "font-style": "italic", + "margin-bottom": "0.7em", + "text-align": "center", + }, + ), + ], + type="default", + ), + html.Div( + [ + dcc.Checklist( + id="pitch-options", + options=[ + {"label": "Change input pitch", "value": "pf"}, + {"label": "Auto-tune output", "value": "pc"}, + {"label": "Disable reference audio", "value": "dra"}, + ], + value=[], + ), + html.Div( + [ + html.Label("Semitones", htmlFor="pitch-factor"), + dcc.Input( + id="pitch-factor", + type="number", + value="0", + style={"width": "7em"}, + min=-11, + max=11, + step=1, + disabled=True, + ), + ], + style={ + "flex-direction": "column", + "margin-left": "10px", + "margin-bottom": "0.7em", + }, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "row", + "margin-left": "50px", + "margin-bottom": "0.7em", + }, + ), + html.Label("Transcript", htmlFor="transcript-input"), + dcc.Textarea( + id="transcript-input", + value="", + style={ + "max-width": "90vw", + "width": "50em", + "height": "8em", + "margin-bottom": "0.7em", + }, + ), + dcc.Loading( + html.Div( + [ + html.Button( + "Generate", + id="gen-button", + ), + html.Audio( + id="audio-out", + controls=True, + style={ + "display": "none", + }, + ), + html.Div( + id="generated-info", + style={ + "font-style": "italic", + }, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "column", + }, + ) + ), + html.Footer( + children=""" + Presented by the Pony Preservation Project. + """, + style={"margin-top": "2em", "font-size": "0.7em"}, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "column", + "background-color": "#FFF", + }, +) + +upload_display = { + "width": "100%", + "height": "60px", + "lineHeight": "60px", + "borderWidth": "1px", + "borderStyle": "dashed", + "borderRadius": "5px", + "textAlign": "center", + "margin": "10px", +} + +playback_style = { + "margin-top": "0.3em", + "margin-bottom": "0.3em", + "display": "block", + "width": "600px", + "max-width": "90vw", +} + +playback_hide = { + "display": "none", +} + + +@app.callback( + [ + dash.dependencies.Output("model-dropdown", "options"), + dash.dependencies.Output("upload-audio", "style"), + ], + dash.dependencies.Input("header", "children"), +) +def init_dropdown(value): + if UI_MODE == "docker": + upload_style = upload_display + else: + upload_style = playback_hide + + dropdown = [ + { + "label": "Custom model", + "value": "Custom|default", + } + ] + prev_values = ["Custom|default"] + + def add_to_dropdown(entry): + if entry["value"] in prev_values: + return + dropdown.append(entry) + prev_values.append(entry["value"]) + + all_dict = {} + for filename in os.listdir("model_lists"): + if len(filename) < 5 or filename[-5:].lower() != ".json": + continue + with open(os.path.join("model_lists", filename)) as f: + j = json.load(f) + for s in j: + for c in s["characters"]: + c["source_file"] = filename[:-5] + if s["source"] not in all_dict: + all_dict[s["source"]] = s["characters"] + else: + all_dict[s["source"]].extend(s["characters"]) + for k in sorted(all_dict): + seen_chars = [] + seen_ids = [] + characters = {} + characters_sing = {} + has_singers = False + for c in all_dict[k]: + if c["drive_id"] in seen_ids: + continue + seen_ids.append(c["drive_id"]) + # Handle duplicate names + if c["name"] in seen_chars: + if c["name"] in characters: + rename = ( + c["name"] + " [" + characters[c["name"]]["source_file"] + "]" + ) + characters[rename] = characters[c["name"]] + del characters[c["name"]] + c["name"] = c["name"] + " [" + c["source_file"] + "]" + else: + seen_chars.append(c["name"]) + + characters[c["name"]] = { + "drive_id": c["drive_id"], + "is_singing": c["is_singing"], + "source_file": c["source_file"], + } + if c["is_singing"]: + has_singers = True + if has_singers: + for ck in sorted(characters): + if characters[ck]["is_singing"]: + characters_sing[ck] = characters[ck] + del characters[ck] + separator = "--- " + k.strip().upper() + " MODELS (TALKING) ---" + else: + separator = "--- " + k.strip().upper() + " MODELS ---" + if len(characters) > 0: + add_to_dropdown( + { + "label": separator, + "value": str(uuid.uuid4()) + "|default", + "disabled": True, + } + ) + for ck in sorted(characters): + add_to_dropdown( + { + "label": ck, + "value": characters[ck]["drive_id"] + "|default", + } + ) + if has_singers: + separator = "--- " + k.strip().upper() + " MODELS (SINGING) ---" + add_to_dropdown( + { + "label": separator, + "value": str(uuid.uuid4()) + "|default", + "disabled": True, + } + ) + for ck in sorted(characters_sing): + add_to_dropdown( + { + "label": ck, + "value": characters_sing[ck]["drive_id"] + "|singing", + } + ) + if len(all_dict) == 0: + add_to_dropdown( + { + "label": "--- NO MODEL LISTS FOUND ---", + "value": str(uuid.uuid4()) + "|default", + "disabled": True, + } + ) + return [dropdown, upload_style] + + +def load_hifigan(model_name, conf_name): + # Load HiFi-GAN + conf = os.path.join("hifi-gan", conf_name + ".json") + with open(conf) as f: + json_config = json.loads(f.read()) + h = AttrDict(json_config) + torch.manual_seed(h.seed) + hifigan = Generator(h).to(torch.device(DEVICE)) + state_dict_g = torch.load(model_name, map_location=torch.device(DEVICE)) + hifigan.load_state_dict(state_dict_g["generator"]) + hifigan.eval() + hifigan.remove_weight_norm() + denoiser = Denoiser(hifigan, mode="normal") + return hifigan, h, denoiser + + +def generate_json(input, outpath): + output = "" + sample_rate = 22050 + lpath = input.split("|")[0].strip() + size = os.stat(lpath).st_size + x = { + "audio_filepath": lpath, + "duration": size / (sample_rate * 2), + "text": input.split("|")[1].strip(), + } + output += json.dumps(x) + "\n" + with open(outpath, "w", encoding="utf8") as w: + w.write(output) + + +asr_model = ( + EncDecCTCModel.from_pretrained(model_name="asr_talknet_aligner").cpu().eval() +) + + +def forward_extractor(tokens, log_probs, blank): + """Computes states f and p.""" + n, m = len(tokens), log_probs.shape[0] + # `f[s, t]` -- max sum of log probs for `s` first codes + # with `t` first timesteps with ending in `tokens[s]`. + f = np.empty((n + 1, m + 1), dtype=float) + f.fill(-(10 ** 9)) + p = np.empty((n + 1, m + 1), dtype=int) + f[0, 0] = 0.0 # Start + for s in range(1, n + 1): + c = tokens[s - 1] + for t in range((s + 1) // 2, m + 1): + f[s, t] = log_probs[t - 1, c] + # Option #1: prev char is equal to current one. + if s == 1 or c == blank or c == tokens[s - 3]: + options = f[s : (s - 2 if s > 1 else None) : -1, t - 1] + else: # Is not equal to current one. + options = f[s : (s - 3 if s > 2 else None) : -1, t - 1] + f[s, t] += np.max(options) + p[s, t] = np.argmax(options) + return f, p + + +def backward_extractor(f, p): + """Computes durs from f and p.""" + n, m = f.shape + n -= 1 + m -= 1 + durs = np.zeros(n, dtype=int) + if f[-1, -1] >= f[-2, -1]: + s, t = n, m + else: + s, t = n - 1, m + while s > 0: + durs[s - 1] += 1 + s -= p[s, t] + t -= 1 + assert durs.shape[0] == n + assert np.sum(durs) == m + assert np.all(durs[1::2] > 0) + return durs + + +def preprocess_tokens(tokens, blank): + new_tokens = [blank] + for c in tokens: + new_tokens.extend([c, blank]) + tokens = new_tokens + return tokens + + +parser = ( + nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset.make_vocab( + notation="phonemes", + punct=True, + spaces=True, + stresses=False, + add_blank_at="last", + ) +) + + +def arpa_parse(input, model): + z = [] + space = parser.labels.index(" ") + while "{" in input: + if "}" not in input: + input.replace("{", "") + else: + pre = input[: input.find("{")] + if pre.strip() != "": + x = model.parse(text=pre.strip()) + seq_ids = x.squeeze(0).cpu().detach().numpy() + z.extend(seq_ids) + z.append(space) + + arpaword = input[input.find("{") + 1 : input.find("}")] + arpaword = ( + arpaword.replace("0", "") + .replace("1", "") + .replace("2", "") + .strip() + .split(" ") + ) + + seq_ids = [] + for x in arpaword: + if x == "": + continue + if x.replace("_", " ") not in parser.labels: + continue + seq_ids.append(parser.labels.index(x.replace("_", " "))) + seq_ids.append(space) + z.extend(seq_ids) + input = input[input.find("}") + 1 :] + if input != "": + x = model.parse(text=input.strip()) + seq_ids = x.squeeze(0).cpu().detach().numpy() + z.extend(seq_ids) + if z[-1] == space: + z = z[:-1] + if z[0] == space: + z = z[1:] + return [ + z[i] for i in range(len(z)) if (i == 0) or (z[i] != z[i - 1]) or (z[i] != space) + ] + + +def to_arpa(input): + arpa = "" + z = [] + space = parser.labels.index(" ") + while space in input: + z.append(input[: input.index(space)]) + input = input[input.index(space) + 1 :] + z.append(input) + for y in z: + if len(y) == 0: + continue + + arpaword = " {" + for s in y: + if parser.labels[s] == " ": + arpaword += "_ " + else: + arpaword += parser.labels[s] + " " + arpaword += "} " + if not arpaword.replace("{", "").replace("}", "").replace(" ", "").isalnum(): + arpaword = arpaword.replace("{", "").replace(" }", "") + arpa += arpaword + return arpa.replace(" ", " ").replace(" }", "}").strip() + + +def get_duration(wav_name, transcript, tokens): + if not os.path.exists(os.path.join(RUN_PATH, "temp")): + os.mkdir(os.path.join(RUN_PATH, "temp")) + if "_" not in transcript: + generate_json( + os.path.join(RUN_PATH, "temp", wav_name + "_conv.wav") + + "|" + + transcript.strip(), + os.path.join(RUN_PATH, "temp", wav_name + ".json"), + ) + else: + generate_json( + os.path.join(RUN_PATH, "temp", wav_name + "_conv.wav") + "|" + "dummy", + os.path.join(RUN_PATH, "temp", wav_name + ".json"), + ) + + data_config = { + "manifest_filepath": os.path.join(RUN_PATH, "temp", wav_name + ".json"), + "sample_rate": 22050, + "batch_size": 1, + } + + dataset = nemo.collections.asr.data.audio_to_text._AudioTextDataset( + manifest_filepath=data_config["manifest_filepath"], + sample_rate=data_config["sample_rate"], + parser=parser, + ) + + dl = torch.utils.data.DataLoader( + dataset=dataset, + batch_size=data_config["batch_size"], + collate_fn=dataset.collate_fn, + shuffle=False, + ) + + blank_id = asr_model.decoder.num_classes_with_blank - 1 + + for sample_idx, test_sample in tqdm(enumerate(dl), total=len(dl)): + log_probs, _, greedy_predictions = asr_model( + input_signal=test_sample[0], input_signal_length=test_sample[1] + ) + + log_probs = log_probs[0].cpu().detach().numpy() + target_tokens = preprocess_tokens(tokens, blank_id) + + f, p = forward_extractor(target_tokens, log_probs, blank_id) + durs = backward_extractor(f, p) + + del test_sample + return durs + return None + + +def crepe_f0(wav_path, hop_length=256): + # sr, audio = wavfile.read(io.BytesIO(wav_data)) + sr, audio = wavfile.read(wav_path) + audio_x = np.arange(0, len(audio)) / 22050.0 + f0time, frequency, confidence, activation = crepe.predict(audio, sr, viterbi=True) + + x = np.arange(0, len(audio), hop_length) / 22050.0 + freq_interp = np.interp(x, f0time, frequency) + conf_interp = np.interp(x, f0time, confidence) + audio_interp = np.interp(x, audio_x, np.absolute(audio)) / 32768.0 + weights = [0.5, 0.25, 0.25] + audio_smooth = np.convolve(audio_interp, np.array(weights)[::-1], "same") + + conf_threshold = 0.25 + audio_threshold = 0.0005 + for i in range(len(freq_interp)): + if conf_interp[i] < conf_threshold: + freq_interp[i] = 0.0 + if audio_smooth[i] < audio_threshold: + freq_interp[i] = 0.0 + + # Hack to make f0 and mel lengths equal + if len(audio) % hop_length == 0: + freq_interp = np.pad(freq_interp, pad_width=[0, 1]) + return ( + torch.from_numpy(freq_interp.astype(np.float32)), + torch.from_numpy(frequency.astype(np.float32)), + ) + + +def f0_to_audio(f0s): + volume = 0.2 + sr = 22050 + freq = 440.0 + base_audio = ( + np.sin(2 * np.pi * np.arange(256.0 * len(f0s)) * freq / sr) * volume + ).astype(np.float32) + shifted_audio = psola.vocode(base_audio, sr, target_pitch=f0s) + for i in range(len(f0s)): + if f0s[i] == 0.0: + shifted_audio[i * 256 : (i + 1) * 256] = 0.0 + print(type(shifted_audio[0])) + buffer = io.BytesIO() + wavfile.write(buffer, sr, shifted_audio.astype(np.float32)) + b64 = base64.b64encode(buffer.getvalue()) + sound = "data:audio/x-wav;base64," + b64.decode("ascii") + return sound + + +@app.callback( + dash.dependencies.Output("custom-model", "style"), + dash.dependencies.Input("model-dropdown", "value"), +) +def update_model(model): + if model is not None and model.split("|")[0] == "Custom": + style = {"margin-bottom": "0.7em", "display": "block"} + else: + style = {"display": "none"} + return style + + +@app.callback( + [ + dash.dependencies.Output("pitch-factor", "disabled"), + dash.dependencies.Output("reference-dropdown", "disabled"), + dash.dependencies.Output("pitch-button", "disabled"), + ], + [ + dash.dependencies.Input("pitch-options", "value"), + ], +) +def update_pitch_options(value): + return ["pf" not in value, "dra" in value, "dra" in value] + + +@app.callback( + dash.dependencies.Output("upload-label", "children"), + [ + dash.dependencies.Input("upload-audio", "filename"), + dash.dependencies.Input("upload-audio", "contents"), + ], +) +def save_upload(uploaded_filenames, uploaded_file_contents): + try: + if uploaded_filenames is not None and uploaded_file_contents is not None: + for name, content in zip(uploaded_filenames, uploaded_file_contents): + if name.strip() == "": + continue + data = content.encode("utf8").split(b";base64,")[1] + with open(os.path.join(RUN_PATH, name), "wb") as fp: + fp.write(base64.decodebytes(data)) + except Exception as e: + return str(e) + return "Uploaded " + str(len(uploaded_filenames)) + " file(s)" + + +@app.callback( + dash.dependencies.Output("reference-dropdown", "options"), + [ + dash.dependencies.Input("update-button", "n_clicks"), + ], +) +def update_filelist(n_clicks): + filelist = [] + supported_formats = [".wav", ".ogg", ".mp3", "flac", ".aac"] + for x in sorted(os.listdir(RUN_PATH)): + if x[-4:].lower() in supported_formats: + filelist.append({"label": x, "value": x}) + return filelist + + +@app.callback( + [ + dash.dependencies.Output("audio-loading-output", "children"), + dash.dependencies.Output("current-f0s", "data"), + dash.dependencies.Output("current-f0s-nosilence", "data"), + dash.dependencies.Output("current-filename", "data"), + ], + [ + dash.dependencies.Input("reference-dropdown", "value"), + ], +) +def select_file(dropdown_value): + if dropdown_value is not None: + if not os.path.exists(os.path.join(RUN_PATH, "temp")): + os.mkdir(os.path.join(RUN_PATH, "temp")) + ffmpeg.input(os.path.join(RUN_PATH, dropdown_value)).output( + os.path.join(RUN_PATH, "temp", dropdown_value + "_conv.wav"), + ar="22050", + ac="1", + acodec="pcm_s16le", + map_metadata="-1", + fflags="+bitexact", + ).overwrite_output().run(quiet=True) + fo_with_silence, f0_wo_silence = crepe_f0( + os.path.join(RUN_PATH, "temp", dropdown_value + "_conv.wav") + ) + return [ + "Analyzed " + dropdown_value, + fo_with_silence, + f0_wo_silence, + dropdown_value, + ] + else: + return ["No audio analyzed", None, None] + + +@app.callback( + [ + dash.dependencies.Output("pitch-out", "src"), + dash.dependencies.Output("pitch-out", "style"), + dash.dependencies.Output("pitch-clicks", "data"), + ], + [ + dash.dependencies.Input("pitch-button", "n_clicks"), + dash.dependencies.Input("pitch-clicks", "data"), + dash.dependencies.Input("current-f0s", "data"), + ], +) +def debug_pitch(n_clicks, pitch_clicks, current_f0s): + if not n_clicks or current_f0s is None or n_clicks <= pitch_clicks: + if n_clicks is not None: + pitch_clicks = n_clicks + else: + pitch_clicks = 0 + return [ + None, + playback_hide, + pitch_clicks, + ] + pitch_clicks = n_clicks + return [f0_to_audio(current_f0s), playback_style, pitch_clicks] + + +hifigan_sr = None + + +def download_model(model, custom_model): + try: + global hifigan_sr, h2, denoiser_sr + d = "https://drive.google.com/uc?id=" + if model == "Custom": + drive_id = custom_model + else: + drive_id = model + if drive_id == "" or drive_id is None: + return ("Missing Drive ID", None, None) + if not os.path.exists(os.path.join(RUN_PATH, "models")): + os.mkdir(os.path.join(RUN_PATH, "models")) + if not os.path.exists(os.path.join(RUN_PATH, "models", drive_id)): + os.mkdir(os.path.join(RUN_PATH, "models", drive_id)) + zip_path = os.path.join(RUN_PATH, "models", drive_id, "model.zip") + gdown.download( + d + drive_id, + zip_path, + quiet=False, + ) + if not os.path.exists(zip_path): + os.rmdir(os.path.join(RUN_PATH, "models", drive_id)) + return ("Model download failed", None, None) + if os.stat(zip_path).st_size < 16: + os.remove(zip_path) + os.rmdir(os.path.join(RUN_PATH, "models", drive_id)) + return ("Model zip is empty", None, None) + with zipfile.ZipFile(zip_path, "r") as zip_ref: + zip_ref.extractall(os.path.join(RUN_PATH, "models", drive_id)) + os.remove(zip_path) + + # Download super-resolution HiFi-GAN + sr_path = "hifi-gan/hifisr" + if not os.path.exists(sr_path): + gdown.download( + d + "14fOprFAIlCQkVRxsfInhEPG0n-xN4QOa", sr_path, quiet=False + ) + if not os.path.exists(sr_path): + raise Exception("HiFI-GAN model failed to download!") + if hifigan_sr is None: + hifigan_sr, h2, denoiser_sr = load_hifigan(sr_path, "config_32k") + + return ( + None, + os.path.join(RUN_PATH, "models", drive_id, "TalkNetSpect.nemo"), + os.path.join(RUN_PATH, "models", drive_id, "hifiganmodel"), + ) + except Exception as e: + return (str(e), None, None) + + +tnmodel, tnpath, tndurs, tnpitch = None, None, None, None +hifigan, h, denoiser, hifipath = None, None, None, None + + +@app.callback( + [ + dash.dependencies.Output("audio-out", "src"), + dash.dependencies.Output("generated-info", "children"), + dash.dependencies.Output("audio-out", "style"), + dash.dependencies.Output("audio-out", "title"), + ], + [dash.dependencies.Input("gen-button", "n_clicks")], + [ + dash.dependencies.State("model-dropdown", "value"), + dash.dependencies.State("drive-id", "value"), + dash.dependencies.State("transcript-input", "value"), + dash.dependencies.State("pitch-options", "value"), + dash.dependencies.State("pitch-factor", "value"), + dash.dependencies.State("current-filename", "data"), + dash.dependencies.State("current-f0s", "data"), + dash.dependencies.State("current-f0s-nosilence", "data"), + ], +) +def generate_audio( + n_clicks, + model, + custom_model, + transcript, + pitch_options, + pitch_factor, + wav_name, + f0s, + f0s_wo_silence, +): + global tnmodel, tnpath, tndurs, tnpitch, hifigan, h, denoiser, hifipath + + if n_clicks is None: + raise PreventUpdate + if model is None: + return [None, "No character selected", playback_hide, None] + if transcript is None or transcript.strip() == "": + return [ + None, + "No transcript entered", + playback_hide, + None, + ] + if wav_name is None and "dra" not in pitch_options: + return [ + None, + "No reference audio selected", + playback_hide, + None, + ] + load_error, talknet_path, hifigan_path = download_model( + model.split("|")[0], custom_model + ) + if load_error is not None: + return [ + None, + load_error, + playback_hide, + None, + ] + + try: + with torch.no_grad(): + if tnpath != talknet_path: + singer_path = os.path.join( + os.path.dirname(talknet_path), "TalkNetSinger.nemo" + ) + if os.path.exists(singer_path): + tnmodel = TalkNetSingerModel.restore_from(singer_path) + else: + tnmodel = TalkNetSpectModel.restore_from(talknet_path) + durs_path = os.path.join( + os.path.dirname(talknet_path), "TalkNetDurs.nemo" + ) + pitch_path = os.path.join( + os.path.dirname(talknet_path), "TalkNetPitch.nemo" + ) + if os.path.exists(durs_path): + tndurs = TalkNetDursModel.restore_from(durs_path) + tnmodel.add_module("_durs_model", tndurs) + tnpitch = TalkNetPitchModel.restore_from(pitch_path) + tnmodel.add_module("_pitch_model", tnpitch) + else: + tndurs = None + tnpitch = None + tnmodel.eval() + tnpath = talknet_path + + token_list = arpa_parse(transcript, tnmodel) + tokens = torch.IntTensor(token_list).view(1, -1).to(DEVICE) + arpa = to_arpa(token_list) + + if "dra" in pitch_options: + if tndurs is None or tnpitch is None: + return [ + None, + "Model doesn't support pitch prediction", + playback_hide, + None, + ] + spect = tnmodel.generate_spectrogram(tokens=tokens) + else: + durs = get_duration(wav_name, transcript, token_list) + + # Change pitch + if "pf" in pitch_options: + f0_factor = np.power(np.e, (0.0577623 * float(pitch_factor))) + f0s = [x * f0_factor for x in f0s] + f0s_wo_silence = [x * f0_factor for x in f0s_wo_silence] + + spect = tnmodel.force_spectrogram( + tokens=tokens, + durs=torch.from_numpy(durs) + .view(1, -1) + .type(torch.LongTensor) + .to(DEVICE), + f0=torch.FloatTensor(f0s).view(1, -1).to(DEVICE), + ) + + if hifipath != hifigan_path: + hifigan, h, denoiser = load_hifigan(hifigan_path, "config_v1") + hifipath = hifigan_path + + y_g_hat = hifigan(spect.float()) + audio = y_g_hat.squeeze() + audio = audio * MAX_WAV_VALUE + audio_denoised = denoiser(audio.view(1, -1), strength=35)[:, 0] + audio_np = ( + audio_denoised.detach().cpu().numpy().reshape(-1).astype(np.int16) + ) + + # Auto-tuning + if "pc" in pitch_options and "dra" not in pitch_options: + _, output_freq, _, _ = crepe.predict(audio_np, 22050, viterbi=True) + output_pitch = torch.from_numpy(output_freq.astype(np.float32)) + target_pitch = torch.FloatTensor(f0s_wo_silence) + factor = torch.mean(output_pitch) / torch.mean(target_pitch) + + octaves = [0.125, 0.25, 0.5, 1.0, 2.0, 4.0, 8.0] + nearest_octave = min(octaves, key=lambda x: abs(x - factor)) + target_pitch *= nearest_octave + if len(target_pitch) < len(output_pitch): + target_pitch = torch.nn.functional.pad( + target_pitch, + (0, list(output_pitch.shape)[0] - list(target_pitch.shape)[0]), + "constant", + 0, + ) + if len(target_pitch) > len(output_pitch): + target_pitch = target_pitch[0 : list(output_pitch.shape)[0]] + + audio_np = psola.vocode( + audio_np, 22050, target_pitch=target_pitch + ).astype(np.float32) + normalize = (1.0 / np.max(np.abs(audio_np))) ** 0.9 + audio_np = audio_np * normalize * MAX_WAV_VALUE + audio_np = audio_np.astype(np.int16) + + # Resample to 32k + wave = resampy.resample( + audio_np, + h.sampling_rate, + h2.sampling_rate, + filter="sinc_window", + window=scipy.signal.windows.hann, + num_zeros=8, + ) + wave_out = wave.astype(np.int16) + + # HiFi-GAN super-resolution + wave = wave / MAX_WAV_VALUE + wave = torch.FloatTensor(wave).to(torch.device(DEVICE)) + new_mel = mel_spectrogram( + wave.unsqueeze(0), + h2.n_fft, + h2.num_mels, + h2.sampling_rate, + h2.hop_size, + h2.win_size, + h2.fmin, + h2.fmax, + ) + y_g_hat2 = hifigan_sr(new_mel) + audio2 = y_g_hat2.squeeze() + audio2 = audio2 * MAX_WAV_VALUE + audio2_denoised = denoiser(audio2.view(1, -1), strength=35)[:, 0] + + # High-pass filter, mixing and denormalizing + audio2_denoised = audio2_denoised.detach().cpu().numpy().reshape(-1) + b = scipy.signal.firwin( + 101, cutoff=10500, fs=h2.sampling_rate, pass_zero=False + ) + y = scipy.signal.lfilter(b, [1.0], audio2_denoised) + y *= 4.0 # superres strength + y_out = y.astype(np.int16) + y_padded = np.zeros(wave_out.shape) + y_padded[: y_out.shape[0]] = y_out + sr_mix = wave_out + y_padded + + buffer = io.BytesIO() + wavfile.write(buffer, 32000, sr_mix.astype(np.int16)) + b64 = base64.b64encode(buffer.getvalue()) + sound = "data:audio/x-wav;base64," + b64.decode("ascii") + + output_name = "TalkNet_" + str(int(time.time())) + return [sound, arpa, playback_style, output_name] + except Exception: + return [ + None, + str(traceback.format_exc()), + playback_hide, + None, + ] + + +if __name__ == "__main__": + app.run_server( + mode="external", + debug=True, + dev_tools_ui=True, + dev_tools_hot_reload=True, + threaded=True, + ) diff --git a/models/tts/ControllableTalkNet/diffsvc_gui.py b/models/tts/ControllableTalkNet/diffsvc_gui.py new file mode 100644 index 0000000000..837bd4e759 --- /dev/null +++ b/models/tts/ControllableTalkNet/diffsvc_gui.py @@ -0,0 +1,368 @@ +import sys +import os +import io +import base64 +import dash +from jupyter_dash import JupyterDash +import dash_core_components as dcc +import dash_html_components as html +from dash.exceptions import PreventUpdate +import gdown +import traceback +from scipy.io import wavfile +import numpy as np +import torch + +sys.path.append("DiffSVC_inference_only") +from end2end import load_e2e_diffsvc, endtoend_from_path, write_to_file + +app = JupyterDash(__name__) +UPLOAD_DIRECTORY = "/content" +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +app.layout = html.Div( + children=[ + html.H1( + children="DiffSVC", + style={ + "font-family": "EquestriaWebfont", + "color": "#280e5f", + "font-size": "4em", + "text-align": "center", + "margin-top": "0em", + "margin-bottom": "0em", + }, + ), + html.Label("Character selection", htmlFor="speaker-dropdown"), + dcc.Dropdown( + id="speaker-dropdown", + options=[ + { + "label": "Twilight Sparkle", + "value": "Twilight", + }, + { + "label": "Discord", + "value": "Discord", + }, + { + "label": "Pinkie Pie", + "value": "Pinkie", + }, + {"label": "Nancy", "value": "Nancy"}, + ], + value=None, + style={ + "max-width": "90vw", + "width": "20em", + "margin-bottom": "0.7em", + }, + ), + html.Label( + "Upload reference audio to " + UPLOAD_DIRECTORY, + htmlFor="reference-dropdown", + ), + dcc.Store(id="current-f0s"), + dcc.Store(id="current-filename"), + dcc.Loading( + id="audio-loading", + children=[ + html.Div( + [ + html.Button( + "Update file list", + id="update-button", + style={ + "margin-right": "10px", + }, + ), + dcc.Dropdown( + id="reference-dropdown", + options=[], + value=None, + style={ + "max-width": "80vw", + "width": "30em", + }, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "row", + "margin-left": "50px", + "vertical-align": "middle", + }, + ), + html.Div( + id="audio-loading-output", + style={ + "font-style": "italic", + "margin-bottom": "0.7em", + "text-align": "center", + }, + ), + ], + type="default", + ), + html.Div( + [ + dcc.Checklist( + id="svc-options", + options=[ + {"label": "Singing mode", "value": "no_pc"}, + ], + value=[], + ) + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "row", + "margin-left": "50px", + "margin-bottom": "0.7em", + }, + ), + html.Div( + [ + html.Label("Steps", htmlFor="linsteps-input"), + dcc.Input( + id="linsteps-input", + type="number", + value="1000", + style={"width": "7em", "margin-left": "10px"}, + min=1, + max=1000, + step=1, + ), + html.Label("Step strength", htmlFor="linend-input"), + dcc.Input( + id="linend-input", + type="number", + value="0.06", + style={"width": "7em", "margin-left": "10px"}, + min=0.01, + max=0.3, + step=0.01, + ), + ], + style={ + "display": "grid", + "grid-template-columns": "auto 100px", + "margin-bottom": "0.7em", + }, + ), + dcc.Loading( + html.Div( + [ + html.Button( + "Generate", + id="gen-button", + ), + html.Audio( + id="audio-out", + controls=True, + style={ + "display": "none", + }, + ), + html.Div( + id="generated-info", + style={ + "font-style": "italic", + }, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "column", + }, + ) + ), + html.Footer( + children=""" + Presented by the Pony Preservation Project. Models by Cookie. + """, + style={"margin-top": "2em", "font-size": "0.7em"}, + ), + ], + style={ + "width": "100%", + "display": "flex", + "align-items": "center", + "justify-content": "center", + "flex-direction": "column", + "background-color": "#FFF", + }, +) + + +playback_style = { + "margin-top": "0.3em", + "margin-bottom": "0.3em", + "display": "block", + "width": "600px", + "max-width": "90vw", +} + +playback_hide = { + "display": "none", +} + + +@app.callback( + dash.dependencies.Output("reference-dropdown", "options"), + [ + dash.dependencies.Input("update-button", "n_clicks"), + ], +) +def update_filelist(n_clicks): + filelist = [] + supported_formats = [".wav", ".ogg", ".mp3", "flac"] + for x in os.listdir(UPLOAD_DIRECTORY): + if x[-4:].lower() in supported_formats: + filelist.append({"label": x, "value": x}) + return filelist + + +@app.callback( + [ + dash.dependencies.Output("audio-loading-output", "children"), + dash.dependencies.Output("current-filename", "data"), + ], + [ + dash.dependencies.Input("reference-dropdown", "value"), + ], +) +def select_file(dropdown_value): + if dropdown_value is not None: + return [ + "Selected " + dropdown_value, + dropdown_value, + ] + else: + return ["No audio selected", None] + + +def download_model(drive_id, outname): + d = "https://drive.google.com/uc?id=" + model_dir = os.path.join(UPLOAD_DIRECTORY, "models") + out_fullpath = os.path.join(model_dir, outname) + out_dir = os.path.dirname(out_fullpath) + + if not os.path.exists(model_dir): + os.mkdir(model_dir) + if not os.path.exists(out_fullpath): + os.makedirs(out_dir, exist_ok=True) + gdown.download( + d + drive_id, + out_fullpath, + quiet=False, + ) + return out_fullpath + + +svc_model, svc_name, last_end, last_step = None, None, None, None + + +@app.callback( + [ + dash.dependencies.Output("audio-out", "src"), + dash.dependencies.Output("generated-info", "children"), + dash.dependencies.Output("audio-out", "style"), + ], + [dash.dependencies.Input("gen-button", "n_clicks")], + [ + dash.dependencies.State("speaker-dropdown", "value"), + dash.dependencies.State("svc-options", "value"), + dash.dependencies.State("current-filename", "data"), + dash.dependencies.State("linend-input", "value"), + dash.dependencies.State("linsteps-input", "value"), + ], +) +def generate_audio(n_clicks, speaker, svc_options, wav_name, lin_end, lin_n_steps): + global svc_model, svc_name, last_end, last_step + + if n_clicks is None: + raise PreventUpdate + if speaker is None: + return [None, "No character selected", playback_hide] + if wav_name is None: + return [None, "No reference audio selected", playback_hide] + + try: + diffsvc_id = "1Uh17L1JtynFDgX9X6jKLPtD_t4IcTEp9" + asr_id = "1qt0pGhCbH0TltFFwSP2Oy3NClthcpr7N" + hifigan_id = "1QQT0HjMhGgDuPhyesZYbvck2lEY7GHAB" + hifigan_config_id = "1mi_O54zi6nW2eU6VuAO_Y4gv_Df-H9Yv" + + diffsvc_path = download_model( + diffsvc_id, + os.path.join(diffsvc_id, "diffsvc_model"), + ) + asr_path = download_model( + asr_id, + os.path.join(asr_id, "asr_model"), + ) + hifigan_path = download_model( + hifigan_id, + os.path.join(hifigan_id, "hifigan_model"), + ) + download_model( + hifigan_config_id, + os.path.join(hifigan_id, "config.json"), + ) + if svc_name != diffsvc_id: + svc_model = load_e2e_diffsvc( + diffsvc_path=diffsvc_path, + dilated_asr_path=asr_path, + hifigan_path=hifigan_path, + device=DEVICE, + ) + svc_name = diffsvc_id + if last_end != lin_end or last_step != lin_n_steps: + svc_model[0].generator.diffusion.set_noise_schedule( + 1e-4, float(lin_end), int(lin_n_steps), device=DEVICE + ) + last_end = lin_end + last_step = lin_n_steps + + pred_audio = endtoend_from_path( + *svc_model, + os.path.join(UPLOAD_DIRECTORY, wav_name), + speaker, + "no_pc" not in svc_options, + "no_pc" not in svc_options, + "no_pc" not in svc_options, + t_max_step=int(lin_n_steps), + ) + pred_audio *= 32768.0 + + buffer = io.BytesIO() + wavfile.write( + buffer, + svc_model[4].sampling_rate, + pred_audio.squeeze().cpu().numpy().astype(np.int16), + ) + b64 = base64.b64encode(buffer.getvalue()) + sound = "data:audio/x-wav;base64," + b64.decode("ascii") + return [sound, "Conversion complete", playback_style] + except Exception: + return [None, str(traceback.format_exc()), playback_hide] + + +if __name__ == "__main__": + app.run_server( + mode="external", + debug=True, + dev_tools_ui=True, + dev_tools_hot_reload=True, + threaded=True, + ) diff --git a/models/tts/ControllableTalkNet/model_lists/sortanon_models.json b/models/tts/ControllableTalkNet/model_lists/sortanon_models.json new file mode 100644 index 0000000000..3d7033c307 --- /dev/null +++ b/models/tts/ControllableTalkNet/model_lists/sortanon_models.json @@ -0,0 +1,165 @@ +[ + {"source": "MLP", + "characters": [ + { + "name": "Applejack", + "drive_id": "1kpEjZ3YqMN3chKSXODOqayEm581rxj4r", + "is_singing": false + }, + { + "name": "Celestia", + "drive_id": "1whXXcnXu9XPcI60xIkTEofpaDYDOw5yB", + "is_singing": false + }, + { + "name": "Chrysalis", + "drive_id": "1bb5jKAcQcEQbx1feVwT1UmAEgocINh-E", + "is_singing": false + }, + { + "name": "Discord", + "drive_id": "1Cg9Oc_K9UDe5WgVDAcaCSbbBoo-Npj1E", + "is_singing": false + }, + { + "name": "Fluttershy", + "drive_id": "1KgVnjrnxZTXgjnI56ilkq5G4UJCbbwZZ", + "is_singing": false + }, + { + "name": "Luna", + "drive_id": "1_ztAbe5YArCMwyyQ_G9lUiz74ym5xJKC", + "is_singing": false + }, + { + "name": "Maud Pie", + "drive_id": "132G6oD0HHPPn4t1H6IkYv18_F0UVLWgi", + "is_singing": false + }, + { + "name": "Pinkie Pie", + "drive_id": "1CdYZ2r52mtgJsFs88U0ZViMSnzpQ_HRp", + "is_singing": false + }, + { + "name": "Rainbow Dash", + "drive_id": "1k3EMXxLC0fLvfxzGbeP6B6plgu9hqCSx", + "is_singing": false + }, + { + "name": "Rarity", + "drive_id": "1QWBvQSso4guc1LRUD40WRJ8DY2CfqHGK", + "is_singing": false + }, + { + "name": "Spike", + "drive_id": "1TKFdmFLttjjzByj2fZW8J70ZHjR-RTwc", + "is_singing": false + }, + { + "name": "Starlight Glimmer", + "drive_id": "1M1AMBq_xjwGTNzRUCXtSLIDJHbcSs3zR", + "is_singing": false + }, + { + "name": "Sunset Shimmer", + "drive_id": "1x1aJt06lBvzUWRlxJ9CEKcFHxQxZPpST", + "is_singing": false + }, + { + "name": "Trixie", + "drive_id": "1a3CYt0-oTTSFjxtZvAVMpClTmQteYua5", + "is_singing": false + }, + { + "name": "Twilight Sparkle", + "drive_id": "1QnOliOAmerMUNuo2wXoH-YoainoSjZen", + "is_singing": false + }, + { + "name": "Twilight Sparkle (whispering)", + "drive_id": "14_TUQVirITdyBh9etfNV8KFFhi_PUs30", + "is_singing": false + }, + { + "name": "Zecora", + "drive_id": "1gL0hqqB7952Q1S185moQd_DRCFfIa3_g", + "is_singing": false + }, + { + "name": "Twilight Sparkle (singing)", + "drive_id": "10CENYWV5ugTXZbnsldN6OKR7wkDEe7V7", + "is_singing": true + }, + { + "name": "Pinkie Pie (singing)", + "drive_id": "1lo-8wI8iB3GArrzIDlaNTA1If8WX0OEf", + "is_singing": true + }, + { + "name": "Applejack (singing)", + "drive_id": "1kegd0J2G4nogaX4IOB2O_vd8msaqMx8M", + "is_singing": true + }, + { + "name": "Fluttershy (singing)", + "drive_id": "1BBdTHis91MwnHTt7tD_xtZ-nQ9SgvqD6", + "is_singing": true + }, + { + "name": "Rarity (singing)", + "drive_id": "1EI7npfQeQYuP6b_Gy8iwA5zAg9Yj5MJA", + "is_singing": true + }, + { + "name": "Rainbow Dash (singing)", + "drive_id": "1q4RgCEVJCSaR_1Dh0rxs6JVRJt30hQRK", + "is_singing": true + }, + { + "name": "Apple Bloom", + "drive_id": "1Wm-7gqws0B3j8mtSQa77Tdk3RZwHEwT5", + "is_singing": false + }, + { + "name": "Scootaloo", + "drive_id": "1YkV1VtP1w5XOx3jYYarrCKSzXCB_FLCy", + "is_singing": false + }, + { + "name": "Sweetie Belle", + "drive_id": "1jLX0Py6j8uY93Fjf2l0HOZQYXiShfWUO", + "is_singing": false + }, + { + "name": "Cozy Glow", + "drive_id": "1cUQQ_4KXTXbmH3MqwwyZu77SpMr20HAL", + "is_singing": false + }, + { + "name": "Big McIntosh", + "drive_id": "1UzL5vo_ykmRdbi-5JOlsQzNKkoV4tJBP", + "is_singing": false + }, + { + "name": "Cadance", + "drive_id": "1w98tMJPfGhWvl797gImMfErPXDz_hA-6", + "is_singing": false + }, + { + "name": "Granny Smith", + "drive_id": "1FJuxLB6MhcHLtmEnOIxhN-yxJ85kmacu", + "is_singing": false + }, + { + "name": "Shining Armor", + "drive_id": "1PUUS71w2ik0uuycNB30nXFze8C7O8OzY", + "is_singing": false + }, + { + "name": "Tirek", + "drive_id": "1rcPDqgDeCIHGDdvfOo-fxfA1XeM4g3CB", + "is_singing": false + } + ]} +] diff --git a/models/tts/ControllableTalkNet/mycroft_talknet.py b/models/tts/ControllableTalkNet/mycroft_talknet.py new file mode 100644 index 0000000000..d94eb8ed04 --- /dev/null +++ b/models/tts/ControllableTalkNet/mycroft_talknet.py @@ -0,0 +1,114 @@ +import os +from flask import Flask, request, send_file +import sys +import torch +import numpy as np +from scipy.io import wavfile +import io +from nemo.collections.tts.models import TalkNetSpectModel +from nemo.collections.tts.models import TalkNetPitchModel +from nemo.collections.tts.models import TalkNetDursModel +import json + +sys.path.append("hifi-gan") +from env import AttrDict +from meldataset import MAX_WAV_VALUE +from models import Generator +from denoiser import Denoiser + + +app = Flask(__name__) +RUN_PATH = os.path.dirname(os.path.realpath(__file__)) +DEVICE = "cuda:0" + + +def load_hifigan(model_name, conf_name): + # Load HiFi-GAN + conf = os.path.join("hifi-gan", conf_name + ".json") + with open(conf) as f: + json_config = json.loads(f.read()) + h = AttrDict(json_config) + torch.manual_seed(h.seed) + hifigan = Generator(h).to(torch.device(DEVICE)) + state_dict_g = torch.load(model_name, map_location=torch.device(DEVICE)) + hifigan.load_state_dict(state_dict_g["generator"]) + hifigan.eval() + hifigan.remove_weight_norm() + denoiser = Denoiser(hifigan, mode="normal") + return hifigan, h, denoiser + + +def generate_json(input, outpath): + output = "" + sample_rate = 22050 + lpath = input.split("|")[0].strip() + size = os.stat(lpath).st_size + x = { + "audio_filepath": lpath, + "duration": size / (sample_rate * 2), + "text": input.split("|")[1].strip(), + } + output += json.dumps(x) + "\n" + with open(outpath, "w", encoding="utf8") as w: + w.write(output) + + +def load_talknet(talknet_path): + with torch.no_grad(): + tnmodel = TalkNetSpectModel.restore_from(talknet_path) + durs_path = os.path.join(os.path.dirname(talknet_path), "TalkNetDurs.nemo") + tndurs = TalkNetDursModel.restore_from(durs_path) + tnmodel.add_module("_durs_model", tndurs) + pitch_path = os.path.join(os.path.dirname(talknet_path), "TalkNetPitch.nemo") + tnpitch = TalkNetPitchModel.restore_from(pitch_path) + tnmodel.add_module("_pitch_model", tnpitch) + tnmodel.eval() + return tnmodel + + +def generate_audio(transcript, tnmodel, hifigan, denoiser): + with torch.no_grad(): + tokens = tnmodel.parse(text=transcript.strip()) + spect = tnmodel.generate_spectrogram(tokens=tokens) + + y_g_hat = hifigan(spect.float()) + audio = y_g_hat.squeeze() + audio = audio * MAX_WAV_VALUE + audio_denoised = denoiser(audio.view(1, -1), strength=35)[:, 0] + audio_np = audio_denoised.detach().cpu().numpy().reshape(-1).astype(np.int16) + + buffer = io.BytesIO() + wavfile.write(buffer, 22050, audio_np) + return buffer + + +@app.route("/", methods=["GET"]) +def get_check(): + return "TalkNet server online" + + +@app.route("/api/tts", methods=["GET"]) +def get_tts(): + if "text" not in request.args: + return "" + transcript = request.args.get("text") + return send_file( + generate_audio(transcript, tnmodel, hifigan, denoiser), + attachment_filename="audio.wav", + mimetype="audio/x-wav", + ) + + +if __name__ == "__main__": + hifigan, h, denoiser = load_hifigan( + os.path.join( + RUN_PATH, "models", "1QnOliOAmerMUNuo2wXoH-YoainoSjZen", "hifiganmodel" + ), + "config_v1", + ) + tnmodel = load_talknet( + os.path.join( + RUN_PATH, "models", "1QnOliOAmerMUNuo2wXoH-YoainoSjZen", "TalkNetSpect.nemo" + ) + ) + app.run(debug=False) diff --git a/models/tts/ControllableTalkNet/requirements.txt b/models/tts/ControllableTalkNet/requirements.txt new file mode 100644 index 0000000000..f7270b68cb --- /dev/null +++ b/models/tts/ControllableTalkNet/requirements.txt @@ -0,0 +1,28 @@ +tensorflow==2.4.1 +tensorboard==2.4.1 +dash==1.21.0 +dash-bootstrap-components==0.13.0 +jupyter-dash==0.4.0 +psola==0.0.1 +wget==3.2 +unidecode==1.2.0 +pysptk==0.1.18 +frozendict==2.0.3 +torch==1.8.1+cu111 +torchvision==0.9.1+cu111 +torchaudio==0.8.1 +torchtext==0.9.1 +torch_stft==0.1.4 +kaldiio==2.17.2 +pydub==0.25.1 +pyannote.audio==1.1.2 +g2p_en==2.1.0 +pesq==0.0.2 +pystoi==0.3.3 +crepe==0.0.12 +resampy==0.2.2 +ffmpeg-python==0.2.0 +tqdm +gdown==3.13.0 +editdistance==0.5.3 +ipywidgets==7.6.3 diff --git a/models/tts/ControllableTalkNet/talknet_offline.py b/models/tts/ControllableTalkNet/talknet_offline.py new file mode 100644 index 0000000000..b2b37a6499 --- /dev/null +++ b/models/tts/ControllableTalkNet/talknet_offline.py @@ -0,0 +1,17 @@ +if __name__ == "__main__": + import os + from controllable_talknet import * + + if os.path.exists("/talknet/is_docker"): + app.run_server( + host="0.0.0.0", + mode="external", + debug=False, + threaded=True, + ) + else: + app.run_server( + mode="external", + debug=False, + threaded=True, + ) diff --git a/models/tts/ControllableTalkNet/talknet_singer.py b/models/tts/ControllableTalkNet/talknet_singer.py new file mode 100644 index 0000000000..34f573dca7 --- /dev/null +++ b/models/tts/ControllableTalkNet/talknet_singer.py @@ -0,0 +1,225 @@ +from collections import OrderedDict +from typing import List + +import torch +from hydra.utils import instantiate +from omegaconf import DictConfig +from pytorch_lightning import Trainer +from torch import nn +from torch.nn import functional as F +import numpy as np + +from nemo.collections.asr.data.audio_to_text import AudioToCharWithDursF0Dataset +from nemo.collections.tts.helpers.helpers import get_mask_from_lengths +from nemo.collections.tts.models.base import SpectrogramGenerator +from nemo.collections.tts.modules.talknet import ( + GaussianEmbedding, + MaskedInstanceNorm1d, + StyleResidual, +) +from nemo.core import Exportable +from nemo.core.classes import ModelPT, PretrainedModelInfo, typecheck +from nemo.core.neural_types import MelSpectrogramType, NeuralType + + +class TalkNetSingerModel(SpectrogramGenerator, Exportable): + """TalkNet's mel spectrogram prediction pipeline.""" + + @property + def output_types(self): + return OrderedDict( + {"mel-spectrogram": NeuralType(("B", "D", "T"), MelSpectrogramType())} + ) + + def __init__(self, cfg: DictConfig, trainer: "Trainer" = None): + super().__init__(cfg=cfg, trainer=trainer) + typecheck.set_typecheck_enabled(enabled=False) + + cfg = self._cfg + self.vocab = AudioToCharWithDursF0Dataset.make_vocab( + **cfg.train_ds.dataset.vocab + ) + self.blanking = cfg.train_ds.dataset.blanking + self.preprocessor = instantiate(cfg.preprocessor) + self.embed = GaussianEmbedding(self.vocab, cfg.d_char) + self.norm_f0 = MaskedInstanceNorm1d(1) + self.res_f0 = StyleResidual(cfg.d_char, 1, kernel_size=3) + self.model = instantiate(cfg.model) + d_out = cfg.model.jasper[-1].filters + self.proj = nn.Conv1d(d_out, cfg.n_mels, kernel_size=1) + + def forward(self, text, text_len, durs, f0): + x, x_len = self.embed(text, durs).transpose(1, 2), durs.sum(-1) + f0_sin = np.nan_to_num(np.sin(9.06472 * np.log(0.0363636 * f0.cpu()))) + f0_cos = np.nan_to_num(np.cos(9.06472 * np.log(0.0363636 * f0.cpu()))) + f0_note = torch.stack( + (torch.from_numpy(f0_sin), torch.from_numpy(f0_cos)), 1 + ).to(self.device) + f0, f0_mask = f0.clone(), f0 > 0.0 + f0 = self.norm_f0(f0.unsqueeze(1), f0_mask) + f0[~f0_mask.unsqueeze(1)] = 0.0 + x = self.res_f0(x, f0) + x = torch.cat((x, f0_note), 1) + y, _ = self.model(x, x_len) + mel = self.proj(y) + return mel + + @staticmethod + def _metrics(true_mel, true_mel_len, pred_mel): + loss = F.mse_loss(pred_mel, true_mel, reduction="none").mean(dim=-2) + mask = get_mask_from_lengths(true_mel_len) + loss *= mask.float() + loss = loss.sum() / mask.sum() + return loss + + def training_step(self, batch, batch_idx): + audio, audio_len, text, text_len, durs, f0, f0_mask = batch + mel, mel_len = self.preprocessor(audio, audio_len) + pred_mel = self(text=text, text_len=text_len, durs=durs, f0=f0) + loss = self._metrics(true_mel=mel, true_mel_len=mel_len, pred_mel=pred_mel) + train_log = {"train_loss": loss} + return {"loss": loss, "progress_bar": train_log, "log": train_log} + + def validation_step(self, batch, batch_idx): + audio, audio_len, text, text_len, durs, f0, f0_mask = batch + mel, mel_len = self.preprocessor(audio, audio_len) + pred_mel = self(text=text, text_len=text_len, durs=durs, f0=f0) + loss = self._metrics(true_mel=mel, true_mel_len=mel_len, pred_mel=pred_mel) + val_log = {"val_loss": loss} + self.log_dict( + val_log, prog_bar=False, on_epoch=True, logger=True, sync_dist=True + ) + + @staticmethod + def _loader(cfg): + dataset = instantiate(cfg.dataset) + return torch.utils.data.DataLoader( # noqa + dataset=dataset, + collate_fn=dataset.collate_fn, + **cfg.dataloader_params, + ) + + def setup_training_data(self, cfg): + self._train_dl = self._loader(cfg) + + def setup_validation_data(self, cfg): + self._validation_dl = self._loader(cfg) + + def setup_test_data(self, cfg): + """Omitted.""" + pass + + def parse(self, text: str, **kwargs) -> torch.Tensor: + return torch.tensor(self.vocab.encode(text)).long().unsqueeze(0).to(self.device) + + def generate_spectrogram(self, tokens: torch.Tensor, **kwargs) -> torch.Tensor: + assert hasattr(self, "_durs_model") and hasattr(self, "_pitch_model") + + if self.blanking: + tokens = [ + AudioToCharWithDursF0Dataset.interleave( + x=torch.empty(len(t) + 1, dtype=torch.long, device=t.device).fill_( + self.vocab.blank + ), + y=t, + ) + for t in tokens + ] + tokens = AudioToCharWithDursF0Dataset.merge( + tokens, value=self.vocab.pad, dtype=torch.long + ) + + text_len = torch.tensor(tokens.shape[-1], dtype=torch.long).unsqueeze(0) + durs = self._durs_model(tokens, text_len) + durs = durs.exp() - 1 + durs[durs < 0.0] = 0.0 + durs = durs.round().long() + + # Pitch + f0_sil, f0_body = self._pitch_model(tokens, text_len, durs) + sil_mask = f0_sil.sigmoid() > 0.5 + f0 = f0_body * self._pitch_model.f0_std + self._pitch_model.f0_mean + f0 = (~sil_mask * f0).float() + + # Spect + mel = self(tokens, text_len, durs, f0) + + return mel + + def forward_for_export(self, tokens: torch.Tensor, text_len: torch.Tensor): + durs = self._durs_model(tokens, text_len) + durs = durs.exp() - 1 + durs[durs < 0.0] = 0.0 + durs = durs.round().long() + + # Pitch + f0_sil, f0_body = self._pitch_model(tokens, text_len, durs) + sil_mask = f0_sil.sigmoid() > 0.5 + f0 = f0_body * self._pitch_model.f0_std + self._pitch_model.f0_mean + f0 = (~sil_mask * f0).float() + + # Spect + x, x_len = self.embed(tokens, durs).transpose(1, 2), durs.sum(-1) + f0_sin = np.nan_to_num(np.sin(9.06472 * np.log(0.0363636 * f0.cpu()))) + f0_cos = np.nan_to_num(np.cos(9.06472 * np.log(0.0363636 * f0.cpu()))) + f0_note = torch.stack( + (torch.from_numpy(f0_sin), torch.from_numpy(f0_cos)), 1 + ).to(self.device) + f0, f0_mask = f0.clone(), f0 > 0.0 + f0 = self.norm_f0(f0.unsqueeze(1), f0_mask) + f0[~f0_mask.unsqueeze(1)] = 0.0 + x = self.res_f0(x, f0) + x = torch.cat((x, f0_note), 1) + y, _ = self.model(x, x_len) + mel = self.proj(y) + + return mel + + def force_spectrogram( + self, tokens: torch.Tensor, durs: torch.Tensor, f0: torch.Tensor, **kwargs + ) -> torch.Tensor: + if self.blanking: + tokens = [ + AudioToCharWithDursF0Dataset.interleave( + x=torch.empty(len(t) + 1, dtype=torch.long, device=t.device).fill_( + self.vocab.blank + ), + y=t, + ) + for t in tokens + ] + tokens = AudioToCharWithDursF0Dataset.merge( + tokens, value=self.vocab.pad, dtype=torch.long + ) + + text_len = torch.tensor(tokens.shape[-1], dtype=torch.long).unsqueeze(0) + durs_len = torch.tensor(durs.shape[-1], dtype=torch.long).unsqueeze(0) + assert text_len == durs_len + + # Spect + mel = self(tokens, text_len, durs, f0) + return mel + + @classmethod + def list_available_models(cls) -> "List[PretrainedModelInfo]": + """ + This method returns a list of pre-trained model which can be instantiated directly from NVIDIA's NGC cloud. + Returns: + List of available pre-trained models. + """ + list_of_models = [] + model = PretrainedModelInfo( + pretrained_model_name="tts_en_talknet", + location=( + "https://api.ngc.nvidia.com/v2/models/nvidia/nemo/tts_en_talknet/versions/1.0.0rc1/files" + "/talknet_spect.nemo" + ), + description=( + "This model is trained on LJSpeech sampled at 22050Hz, and can be used to generate female " + "English voices with an American accent." + ), + class_=cls, # noqa + aliases=["TalkNet-22050Hz"], + ) + list_of_models.append(model) + return list_of_models diff --git a/models/tts/NeMo/.dockerignore b/models/tts/NeMo/.dockerignore new file mode 100644 index 0000000000..14f5114d01 --- /dev/null +++ b/models/tts/NeMo/.dockerignore @@ -0,0 +1,19 @@ +__pycache__ +*.pyc +*.pyo +*.pyd +.Python +env +pip-log.txt +pip-delete-this-directory.txt +.tox +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*,cover +*.log +.git +**/*.nemo +**/*.ckpt diff --git a/models/tts/NeMo/.github/ISSUE_TEMPLATE/bug_report.md b/models/tts/NeMo/.github/ISSUE_TEMPLATE/bug_report.md new file mode 100644 index 0000000000..5aedacf07f --- /dev/null +++ b/models/tts/NeMo/.github/ISSUE_TEMPLATE/bug_report.md @@ -0,0 +1,42 @@ +--- +name: Bug report +about: Create a report to help us improve +title: '' +labels: bug +assignees: '' + +--- + +**Describe the bug** + +A clear and concise description of what the bug is. + +**Steps/Code to reproduce bug** + +Please list *minimal* steps or code snippet for us to be able to reproduce the bug. + +A helpful guide on on how to craft a minimal bug report http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports. + + +**Expected behavior** + +A clear and concise description of what you expected to happen. + +**Environment overview (please complete the following information)** + + - Environment location: [Bare-metal, Docker, Cloud(specify cloud provider - AWS, Azure, GCP, Collab)] + - Method of NeMo install: [pip install or from source]. Please specify exact commands you used to install. + - If method of install is [Docker], provide `docker pull` & `docker run` commands used + +**Environment details** + +If NVIDIA docker image is used you don't need to specify these. +Otherwise, please provide: +- OS version +- PyTorch version +- Python version + +**Additional context** + +Add any other context about the problem here. +Example: GPU model diff --git a/models/tts/NeMo/.github/ISSUE_TEMPLATE/feature_request.md b/models/tts/NeMo/.github/ISSUE_TEMPLATE/feature_request.md new file mode 100644 index 0000000000..e56d0d05e0 --- /dev/null +++ b/models/tts/NeMo/.github/ISSUE_TEMPLATE/feature_request.md @@ -0,0 +1,25 @@ +--- +name: Feature request +about: Suggest an idea for this project +title: '' +labels: feature request +assignees: okuchaiev + +--- + +**Is your feature request related to a problem? Please describe.** + +A clear and concise description of what the problem is. Ex. I'm always frustrated when [...] + +**Describe the solution you'd like** + +A clear and concise description of what you want to happen. +Provide a code snippet on how new APIs/changes would be used by others. + +**Describe alternatives you've considered** + +A clear and concise description of any alternative solutions or features you've considered. + +**Additional context** + +Add any other context or screenshots about the feature request here. diff --git a/models/tts/NeMo/.gitignore b/models/tts/NeMo/.gitignore new file mode 100644 index 0000000000..ff04a26262 --- /dev/null +++ b/models/tts/NeMo/.gitignore @@ -0,0 +1,176 @@ +# log and data files +*.model +*.pkl +#*.ipynb +output +result +*.pt +tests/data/asr +.DS_Store +bert.pt.json +work +runs +fastspeech_output +.hydra + +# Byte-compiled / optimized / DLL files +__pycache__/ +*.py[cod] +*$py.class +**.pyc + +# C extensions +*.so + +# Distribution / packaging +.idea +.Python +wandb +build/ +develop-eggs/ +dist/ +downloads/ +eggs/ +.eggs/ +lib/ +lib64/ +#parts/ +sdist/ +var/ +wheels/ +pip-wheel-metadata/ +share/python-wheels/ +*.egg-info/ +.installed.cfg +*.egg +MANIFEST + +# PyInstaller +# Usually these files are written by a python script from a template +# before PyInstaller builds the exe, so as to inject date/other infos into it. +*.manifest +*.spec + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +htmlcov/ +.tox/ +.nox/ +.coverage +.coverage.* +.cache +nosetests.xml +coverage.xml +*.cover +.hypothesis/ +.pytest_cache/ + +# Translations +*.mo +*.pot + +# Django stuff: +*.log +local_settings.py +db.sqlite3 + +# Flask stuff: +instance/ +.webassets-cache + +# Scrapy stuff: +.scrapy + +# Sphinx documentation +docs/build + +# PyBuilder +target/ + +# Jupyter Notebook +.ipynb_checkpoints + +# IPython +profile_default/ +ipython_config.py + +# pyenv +.python-version + +# pipenv +# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. +# However, in case of collaboration, if having platform-specific dependencies or dependencies +# having no cross-platform support, pipenv may install dependencies that don’t work, or not +# install all needed dependencies. +#Pipfile.lock + +# celery beat schedule file +celerybeat-schedule + +# SageMath parsed files +*.sage.py + +# Environments +.env +.venv +env/ +venv/ +ENV/ +env.bak/ +venv.bak/ + +# VSCode project settins +.vscode/ + +# Spyder project settings +.spyderproject +.spyproject + +# Rope project settings +.ropeproject + +# mkdocs documentation +/site +/docs/html +/docs/docs_zh/zh + +# mypy +.mypy_cache/ +.dmypy.json +dmypy.json + +# Pyre type checker +.pyre/ + +# Emacs backup files +*~ + +cifar-10-batches-py +*.tar.gz + +# Test data. +tests/.data +tests/data + +# outputs folder +examples/*/outputs +examples/*/NeMo_experiments +examples/*/nemo_experiments +examples/*/.hydra +examples/*/wandb +examples/*/data +wandb +dump.py + +docs/sources/source/test_build/ + +# Checkpoints, config files and temporary files created in tutorials. +examples/neural_graphs/*.chkpt +examples/neural_graphs/*.yml + +.hydra/ +nemo_experiments/ + diff --git a/models/tts/NeMo/.readthedocs.yml b/models/tts/NeMo/.readthedocs.yml new file mode 100644 index 0000000000..9b10e64855 --- /dev/null +++ b/models/tts/NeMo/.readthedocs.yml @@ -0,0 +1,31 @@ +# ============================================================================= +# Copyright (c) 2020 NVIDIA. 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. +# ============================================================================= + +# Read the Docs configuration file +# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details + +# Required field. +version: 2 + +# Build documentation in the docs/ directory with Sphinx. +sphinx: + configuration: docs/source/conf.py + +# Set the version of Python and requirements required to build your docs +python: + version: 3.7 + install: + - requirements: requirements/requirements_docs.txt diff --git a/models/tts/NeMo/CONTRIBUTING.md b/models/tts/NeMo/CONTRIBUTING.md new file mode 100644 index 0000000000..394ed9278b --- /dev/null +++ b/models/tts/NeMo/CONTRIBUTING.md @@ -0,0 +1,68 @@ +# Contributions are welcome! + +We do all of NeMo's development in the open. Contributions from NeMo community are welcome. + + +# Pull Requests (PR) Guidelines + +**Send your PRs to the `main` branch** + +1) Make sure your PR does one thing. Have a clear answer to "What does this PR do?". +2) Read General Principles and style guide below +3) Make sure you sign your commits. E.g. use ``git commit -s`` when before your commit +4) Make sure all unittests finish successfully before sending PR ``pytest`` or (if yor dev box does not have GPU) ``pytest --cpu`` from NeMo's root folder +5) Send your PR and request a review + +## Unit tests +Quick tests (locally, while developing) +``` +pytest +# If you don't have NVIDIA GPU do: +# pytest --cpu +``` +Full tests, including pre-trained model downloads +``` +pytest --with_downloads +``` + +## Whom should you ask for review: +1. For changes to NeMo's core: @ericharper, @titu1994, @blisc, or @okuchaiev +1. For changes to NeMo's ASR collection: @titu1994, @redoctopus, @jbalam-nv, or @okuchaiev +1. For changes to NeMo's NLP collection: @MaximumEntropy, @ericharper, @ekmb, @yzhang123, @VahidooX, @vladgets, or @okuchaiev +1. For changes to NeMo's TTS collection: @blisc or @stasbel, or @okuchaiev + +Note that some people may self-assign to review your PR - in which case, please wait for them to add a review. + +Your pull requests must pass all checks and peer-review before they can be merged. + +# General principles +1. **User-oriented**: make it easy for end users, even at the cost of writing more code in the background +1. **Robust**: make it hard for users to make mistakes. +1. **Well-tested**: please add simple, fast unittests. Consider adding CI tests for end-to-end functionality. +1. **Reusable**: for every piece of code, think about how it can be reused in the future and make it easy to be reused. +1. **Readable**: code should be easier to read. +1. **Legal**: if you copy even one line of code from the Internet, make sure that the code allows the license that NeMo supports. Give credit and link back to the code. +1. **Sensible**: code should make sense. If you think a piece of code might be confusing, write comments. + +## Python style +We use ``black`` as our style guide. To check whether your code will pass style check (from the NeMo's repo folder) run: +``python setup.py style`` and if it does not pass run ``python setup.py style --fix``. + +1. Include docstrings for every class and method exposed to the user. +1. Use Python 3 type hints for every class and method exposed to the user. +1. Avoid wild import: ``from X import *`` unless in ``X.py``, ``__all__`` is defined. +1. Minimize the use of ``**kwargs``. +1. ``RaiseError`` is preferred to ``assert``. Write: ```if X: raise Error``` instead of ```assert X```. +1. Classes are preferred to standalone methods. +1. Methods should be atomic. A method shouldn't be longer than 75 lines, e.g. can be fit into the computer screen without scrolling. +1. If a method has arguments that don't fit into one line, each argument should be in its own line for readability. +1. Add ``__init__.py`` for every folder. +1. F-strings are prefered to formatted strings. +1. Loggers are preferred to print. In NeMo, you can use logger from ``from nemo.utils import logging`` +1. Private functions (functions start with ``_``) shouldn't be called outside its host file. +1. If a comment lasts multiple lines, use ``'''`` instead of ``#``. + +# Collections +Collection is a logical grouping of related Neural Modules. It is a grouping of modules that share a domain area or semantics. +When contributing module to a collection, please make sure it belongs to that category. +If you would like to start a new one and contribute back to the platform, you are very welcome to do so. diff --git a/models/tts/NeMo/Dockerfile b/models/tts/NeMo/Dockerfile new file mode 100644 index 0000000000..9e364b4b3e --- /dev/null +++ b/models/tts/NeMo/Dockerfile @@ -0,0 +1,104 @@ +# syntax=docker/dockerfile:experimental + +# Copyright (c) 2019, NVIDIA CORPORATION. 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. + +ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:21.05-py3 + + +# build an image that includes only the nemo dependencies, ensures that dependencies +# are included first for optimal caching, and useful for building a development +# image (by specifying build target as `nemo-deps`) +FROM ${BASE_IMAGE} as nemo-deps + +# Ensure apt-get won't prompt for selecting options +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && \ + apt-get install -y \ + libsndfile1 sox \ + libfreetype6 \ + python-setuptools swig \ + python-dev ffmpeg && \ + rm -rf /var/lib/apt/lists/* + +# uninstall stuff from base container +RUN pip uninstall -y sacrebleu torchtext + +# build torchaudio (change latest release version to match pytorch) +WORKDIR /tmp/torchaudio_build +RUN git clone --depth 1 --branch release/0.7 https://github.com/pytorch/audio.git && \ + cd audio && \ + BUILD_SOX=1 python setup.py install && \ + cd .. && rm -r audio + +# TODO: remove when 21.04 container is released +# build torchtext +WORKDIR /tmp/torchtext_build +RUN git clone --branch v0.8.1 https://github.com/pytorch/text.git && \ + cd text && \ + git submodule update --init --recursive && \ + python setup.py clean install && \ + cd .. && rm -r text + +#install TRT tools: PT quantization support and ONNX graph optimizer +WORKDIR /tmp/trt_build +RUN git clone https://github.com/NVIDIA/TensorRT.git && \ + cd TensorRT/tools/onnx-graphsurgeon && python setup.py install && \ + cd ../pytorch-quantization && \ + python setup.py install && \ + rm -fr /tmp/trt_build + +# install nemo dependencies +WORKDIR /tmp/nemo +COPY requirements . +RUN for f in $(ls requirements*.txt); do pip install --disable-pip-version-check --no-cache-dir -r $f; done + +# install nemo_text_processing dependencies +COPY nemo_text_processing /tmp/nemo/nemo_text_processing/ +RUN /bin/bash /tmp/nemo/nemo_text_processing/setup.sh + +# copy nemo source into a scratch image +FROM scratch as nemo-src +COPY . . + +# start building the final container +FROM nemo-deps as nemo +ARG NEMO_VERSION=1.0.2 + +# Check that NEMO_VERSION is set. Build will fail without this. Expose NEMO and base container +# version information as runtime environment variable for introspection purposes +RUN /usr/bin/test -n "$NEMO_VERSION" && \ + /bin/echo "export NEMO_VERSION=${NEMO_VERSION}" >> /root/.bashrc && \ + /bin/echo "export BASE_IMAGE=${BASE_IMAGE}" >> /root/.bashrc +RUN --mount=from=nemo-src,target=/tmp/nemo cd /tmp/nemo && pip install ".[all]" && \ + python -c "import nemo.collections.asr as nemo_asr" && \ + python -c "import nemo.collections.nlp as nemo_nlp" && \ + python -c "import nemo.collections.tts as nemo_tts" && \ + python -c "import nemo_text_processing.text_normalization as text_normalization" + +# TODO: Remove once 21.04 container is base container +# install latest numba version +RUN conda update -c numba numba -y + +# copy scripts/examples/tests into container for end user +WORKDIR /workspace/nemo +COPY scripts /workspace/nemo/scripts +COPY examples /workspace/nemo/examples +COPY tests /workspace/nemo/tests +COPY tutorials /workspace/nemo/tutorials +# COPY README.rst LICENSE /workspace/nemo/ + +RUN printf "#!/bin/bash\njupyter lab --no-browser --allow-root --ip=0.0.0.0" >> start-jupyter.sh && \ + chmod +x start-jupyter.sh + diff --git a/models/tts/NeMo/Jenkinsfile b/models/tts/NeMo/Jenkinsfile new file mode 100644 index 0000000000..2f4c453955 --- /dev/null +++ b/models/tts/NeMo/Jenkinsfile @@ -0,0 +1,1451 @@ +pipeline { + agent { + docker { + image 'nvcr.io/nvidia/pytorch:21.05-py3' + args '--device=/dev/nvidia0 --gpus all --user 0:128 -v /home/TestData:/home/TestData -v $HOME/.cache/torch:/root/.cache/torch --shm-size=8g' + } + } + options { + timeout(time: 1, unit: 'HOURS') + disableConcurrentBuilds() + } + stages { + stage('PyTorch version') { + steps { + sh 'python -c "import torch; print(torch.__version__)"' + sh 'python -c "import torchvision; print(torchvision.__version__)"' + } + } + + stage('Uninstall torchtext') { + steps { + sh 'pip uninstall -y torchtext' + } + } + + stage('Install test requirements') { + steps { + sh 'apt-get update && apt-get install -y bc && pip install -r requirements/requirements_test.txt' + } + } + + stage('Copyright Headers check') { + steps { + sh 'python /home/TestData/check_copyright_header.py --dir .' + } + } + + stage('PyTorch STFT Patch check') { + steps { + sh 'python /home/TestData/check_stft_patch.py --dir .' + } + } + + stage('Code formatting checks') { + steps { + sh 'python setup.py style' + } + } + stage('Installation') { + steps { + sh './reinstall.sh release' + } + } + + stage('PyTorch Lightning version') { + steps { + sh 'python -c "import pytorch_lightning; print(pytorch_lightning.__version__)"' + } + } + + stage('PyTorch Lightning DDP Checks') { + steps { + sh 'python "tests/core_ptl/check_for_ranks.py"' + } + } + + stage('L0: Unit Tests GPU') { + steps { + sh 'pytest -m "not pleasefixme" --with_downloads --relax_numba_compat' + } + } + + stage('L0: Unit Tests CPU') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + steps { + sh 'CUDA_VISIBLE_DEVICES="" pytest -m "not pleasefixme" --cpu --with_downloads --relax_numba_compat' + } + } + + stage('L0: Computer Vision Integration') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage ('MNIST image classification with LeNet-5 Integration Test - on CPU') { + steps { + sh 'cd examples/cv && \ + python mnist_lenet5_image_classification_pure_lightning.py trainer.gpus=0 \ + trainer.accelerator=null \ + trainer.fast_dev_run=true model.dataset.data_folder=/home/TestData \ + && rm -rf outputs' + } + } + } + } + + // We have no integration tests, please enable this when one is added + // stage('L0: Integration Tests GPU') { + // steps { + // sh 'pytest -s -m "integration and not skipduringci and not pleasefixme"' + // } + // } + + // stage('L0: Integration Tests CPU') { + // when { + // anyOf{ + // branch 'main' + // changeRequest target: 'main' + // } + // } + // steps { + // sh 'pytest -s -m "integration and not pleasefixme" --cpu' + // } + // } + + // We have no system tests, please enable this when one is added + // stage('L1: System Tests GPU') { + // steps { + // sh 'pytest -m "system and not skipduringci and not pleasefixme"' + // } + // } + + // stage('L1: System Tests CPU') { + // when { + // anyOf{ + // branch 'dev + // changeRequest target: 'main' + // } + // } + // steps { + // sh 'pytest -m "system and not pleasefixme" --cpu' + // } + // } + + stage('L2: NeMo text processing') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: TN') { + steps { + sh 'cd tools/text_processing_deployment && python pynini_export.py --output=/home/TestData/nlp/text_norm/output/ --grammars=tn_grammars && ls -R /home/TestData/nlp/text_norm/output/ && echo ".far files created "|| exit 1' + sh 'cd nemo_text_processing/text_normalization/ && python run_predict.py --input=/home/TestData/nlp/text_norm/ci/test.txt --input_case="lower_cased" --output=/home/TestData/nlp/text_norm/output/test.pynini.txt --verbose' + sh 'cmp --silent /home/TestData/nlp/text_norm/output/test.pynini.txt /home/TestData/nlp/text_norm/ci/test_goal_py.txt || exit 1' + sh 'rm -rf /home/TestData/nlp/text_norm/output/*' + } + } + + stage('L2: ITN export') { + steps { + sh 'cd tools/text_processing_deployment && python pynini_export.py --output=/home/TestData/nlp/text_denorm/output/ --grammars=itn_grammars && ls -R /home/TestData/nlp/text_denorm/output/ && echo ".far files created "|| exit 1' + sh 'cd nemo_text_processing/inverse_text_normalization/ && python run_predict.py --input=/home/TestData/nlp/text_denorm/ci/test.txt --output=/home/TestData/nlp/text_denorm/output/test.pynini.txt --verbose' + sh 'cmp --silent /home/TestData/nlp/text_denorm/output/test.pynini.txt /home/TestData/nlp/text_denorm/ci/test_goal_py.txt || exit 1' + sh 'rm -rf /home/TestData/nlp/text_denorm/output/*' + } + } + stage('L2: TN with Audio (audio and raw text)') { + steps { + sh 'cd nemo_text_processing/text_normalization && \ + python normalize_with_audio.py --text "The total amounts to \\$4.76." \ + --audio_data /home/TestData/nlp/text_norm/audio_based/audio.wav | tail -n2 | head -n1 > /home/TestData/nlp/text_norm/audio_based/output/out_raw.txt 2>&1 && \ + cmp --silent /home/TestData/nlp/text_norm/audio_based/output/out_raw.txt /home/TestData/nlp/text_norm/audio_based/result.txt || exit 1' + sh 'rm -rf /home/TestData/nlp/text_norm/audio_based/output/out_raw.txt' + } + } + stage('L2: TN with Audio (audio and text file)') { + steps { + sh 'cd nemo_text_processing/text_normalization && \ + python normalize_with_audio.py --text /home/TestData/nlp/text_norm/audio_based/text.txt \ + --audio_data /home/TestData/nlp/text_norm/audio_based/audio.wav | tail -n2 | head -n1 > /home/TestData/nlp/text_norm/audio_based/output/out_file.txt 2>&1 && \ + cmp --silent /home/TestData/nlp/text_norm/audio_based/output/out_file.txt /home/TestData/nlp/text_norm/audio_based/result.txt || exit 1' + sh 'rm -rf /home/TestData/nlp/text_norm/audio_based/output/out_file.txt' + } + } + stage('L2: TN with Audio (manifest)') { + steps { + sh 'cd nemo_text_processing/text_normalization && \ + python normalize_with_audio.py --audio_data /home/TestData/nlp/text_norm/audio_based/manifest.json --n_tagged=120 && \ + cmp --silent /home/TestData/nlp/text_norm/audio_based/manifest_normalized.json /home/TestData/nlp/text_norm/audio_based/manifest_result.json || exit 1' + sh 'rm -rf /home/TestData/nlp/text_norm/audio_based/manifest_normalized.json' + } + } + } + } + + stage('L2: ASR dev run') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('Speech to Text') { + steps { + sh 'python examples/asr/speech_to_text.py \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/asr/speech_to_text_results' + sh 'rm -rf examples/asr/speech_to_text_results' + } + } + + stage('Speech to Label') { + steps { + sh 'python examples/asr/speech_to_label.py \ + model.train_ds.manifest_filepath=/home/TestData/speech_commands/train_manifest.json \ + model.validation_ds.manifest_filepath=/home/TestData/speech_commands/test_manifest.json \ + model.test_ds.manifest_filepath=/home/TestData/speech_commands/test_manifest.json \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=True \ + model.preprocessor._target_=nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor \ + ~model.preprocessor.window_size \ + ~model.preprocessor.window_stride \ + ~model.preprocessor.window \ + ~model.preprocessor.n_mels \ + ~model.preprocessor.n_mfcc \ + ~model.preprocessor.n_fft \ + exp_manager.exp_dir=examples/asr/speech_to_label_results' + sh 'rm -rf examples/asr/speech_to_label_results' + } + } + + stage('Speaker Recognition') { + steps { + sh 'python examples/speaker_recognition/speaker_reco.py \ + model.train_ds.batch_size=10 \ + model.validation_ds.batch_size=2 \ + model.train_ds.manifest_filepath=/home/TestData/an4_speaker/train.json \ + model.validation_ds.manifest_filepath=/home/TestData/an4_speaker/dev.json \ + model.test_ds.manifest_filepath=/home/TestData/an4_speaker/test.json \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/speaker_recognition/speaker_recognition_results' + sh 'rm -rf examples/speaker_recognition/speaker_recognition_results' + } + } + + stage('Speaker Diarization Inference') { + steps { + sh 'python examples/speaker_recognition/speaker_diarize.py \ + diarizer.oracle_num_speakers=2 \ + diarizer.paths2audio_files=/home/TestData/an4_diarizer/audio_files.scp \ + diarizer.path2groundtruth_rttm_files=/home/TestData/an4_diarizer/rttm_files.scp \ + diarizer.speaker_embeddings.model_path=/home/TestData/an4_diarizer/spkr.nemo \ + diarizer.vad.model_path=/home/TestData/an4_diarizer/MatchboxNet_VAD_3x2.nemo \ + diarizer.out_dir=examples/speaker_recognition/speaker_diarization_results' + sh 'rm -rf examples/speaker_recognition/speaker_diarization_results' + } + } + + stage('L2: Speech to Text WPE - CitriNet') { + steps { + sh 'python examples/asr/speech_to_text_bpe.py \ + --config-path="conf/citrinet/" --config-name="config_bpe" \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + model.tokenizer.dir="/home/TestData/asr_tokenizers/an4_wpe_128/" \ + model.tokenizer.type="wpe" \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/asr/speech_to_text_wpe_results' + sh 'rm -rf examples/asr/speech_to_text_wpe_results' + } + } + + stage('L2: Speech to Text WPE - Conformer') { + steps { + sh 'python examples/asr/speech_to_text_bpe.py \ + --config-path="conf/conformer" --config-name="conformer_ctc_bpe" \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + model.tokenizer.dir="/home/TestData/asr_tokenizers/an4_wpe_128/" \ + model.tokenizer.type="wpe" \ + model.train_ds.batch_size=4 \ + model.validation_ds.batch_size=4 \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/asr/speech_to_text_wpe_conformer_results' + sh 'rm -rf examples/asr/speech_to_text_wpe_conformer_results' + } + } + } + } + + stage('L2: ASR DALI dev run') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('Speech to Text - DALI AudioToMelSpectrogramPreprocessor') { + steps { + sh 'python examples/asr/speech_to_text.py \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + +model.train_ds.use_dali=True \ + model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + +model.validation_ds.use_dali=True \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/asr/speech_to_text_results' + sh 'rm -rf examples/asr/speech_to_text_results' + } + } + // TODO: This would fail due to an unnecessary torchaudio import. + // To be enabled once torchaudio is available in the container used for CI + // stage('Speech to Text - DALI AudioToMFCCPreprocessor') { + // steps { + // sh 'python examples/asr/speech_to_text.py \ + // model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + // +model.train_ds.use_dali=True \ + // model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + // +model.validation_ds.use_dali=True \ + // model.preprocessor._target_=nemo.collections.asr.modules.AudioToMFCCPreprocessor \ + // ~model.preprocessor.normalize \ + // ~model.preprocessor.features \ + // ~model.preprocessor.frame_splicing \ + // ~model.preprocessor.dither \ + // ~model.preprocessor.stft_conv \ + // +model.n_mels=64 \ + // +model.n_mfcc=64 \ + // trainer.gpus=[0] \ + // +trainer.fast_dev_run=True \ + // exp_manager.exp_dir=examples/asr/speech_to_text_results' + // sh 'rm -rf examples/asr/speech_to_text_results' + // } + // } + } + } + +// TODO: UNCOMMENT TESTS AFTER 21.04 release (numba 0.53 min requirement) + stage('L2: ASR RNNT dev run') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('Speech to Text - RNNT') { + steps { + sh 'STRICT_NUMBA_COMPAT_CHECK=false python examples/asr/speech_to_text_rnnt.py \ + --config-path="experimental/contextnet_rnnt/" --config-name="config_rnnt.yaml" \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + model.train_ds.batch_size=2 \ + model.validation_ds.batch_size=2 \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/asr/speech_to_text_rnnt_results' + sh 'rm -rf examples/asr/speech_to_text_rnnt_results' + } + } + stage('L2: Speech to Text RNNT WPE') { + steps { + sh 'STRICT_NUMBA_COMPAT_CHECK=false python examples/asr/speech_to_text_rnnt_bpe.py \ + --config-path="experimental/contextnet_rnnt/" --config-name="config_rnnt_bpe.yaml" \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + model.validation_ds.manifest_filepath=/home/TestData/an4_dataset/an4_val.json \ + model.train_ds.batch_size=2 \ + model.validation_ds.batch_size=2 \ + model.tokenizer.dir="/home/TestData/asr_tokenizers/an4_wpe_128/" \ + model.tokenizer.type="wpe" \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=True \ + exp_manager.exp_dir=examples/asr/speech_to_text_rnnt_wpe_results' + sh 'rm -rf examples/asr/speech_to_text_rnnt_wpe_results' + } + } + } + } + + stage('L2: ASR Multi-dataloader dev run') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('Speech to Text multi-dataloader') { + steps { + sh 'python examples/asr/speech_to_text.py \ + model.train_ds.manifest_filepath=/home/TestData/an4_dataset/an4_train.json \ + model.validation_ds.manifest_filepath=[/home/TestData/an4_dataset/an4_val.json,/home/TestData/an4_dataset/an4_val.json] \ + trainer.gpus=[0] \ + trainer.max_epochs=1 \ + +trainer.max_steps=1 \ + +trainer.num_sanity_val_steps=1 \ + exp_manager.exp_dir=examples/asr/speech_to_text_results' + sh 'rm -rf examples/asr/speech_to_text_results' + } + } + + stage('Speech to Label multi-dataloader') { + steps { + sh 'python examples/asr/speech_to_label.py \ + model.train_ds.manifest_filepath=/home/TestData/speech_commands/train_manifest.json \ + model.validation_ds.manifest_filepath=[/home/TestData/speech_commands/test_manifest.json,/home/TestData/speech_commands/test_manifest.json] \ + trainer.gpus=[1] \ + trainer.max_epochs=1 \ + +trainer.max_steps=1 \ + +trainer.num_sanity_val_steps=1 \ + model.preprocessor._target_=nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor \ + ~model.preprocessor.window_size \ + ~model.preprocessor.window_stride \ + ~model.preprocessor.window \ + ~model.preprocessor.n_mels \ + ~model.preprocessor.n_mfcc \ + ~model.preprocessor.n_fft \ + exp_manager.exp_dir=examples/asr/speech_to_label_results' + sh 'rm -rf examples/asr/speech_to_label_results' + } + } + } + } + + stage('L2: Speech Transcription') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('Speech to Text Transcribe') { + steps { + sh 'python examples/asr/transcribe_speech.py \ + pretrained_name="QuartzNet15x5Base-En" \ + audio_dir="/home/TestData/an4_transcribe/test_subset/" \ + output_filename="stt_test_res.json" \ + cuda=true \ + amp=true' + sh 'rm -rf stt_test_res.json' + } + } + } + } + + stage('L2: Segmentation Tool') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + stages { + stage('Install ctc_segmentation requirements') { + steps { + sh 'cd tools/ctc_segmentation && \ + pip install -r requirements.txt && \ + DEBIAN_FRONTEND=noninteractive apt-get install -y --no-install-recommends ffmpeg' + } + } + + stage('Parallel ctc_segmentation test') { + failFast true + parallel { + stage('L2: Eng QN with .wav') { + steps { + sh 'cd tools/ctc_segmentation && \ + TIME=`date +"%Y-%m-%d-%T"` && \ + /bin/bash run_sample.sh \ + --MODEL_NAME_OR_PATH=QuartzNet15x5Base-En \ + --DATA_DIR=/home/TestData/ctc_segmentation/eng \ + --OUTPUT_DIR=/home/TestData/ctc_segmentation/eng/output${TIME} \ + --LANGUAGE=eng \ + --OFFSET=0 \ + --CUT_PREFIX=0 \ + --MIN_SEGMENT_LEN=0 \ + --AUDIO_FORMAT=.wav && \ + python /home/TestData/ctc_segmentation/verify_alignment.py \ + -r /home/TestData/ctc_segmentation/eng/eng_valid_segments.txt \ + -g /home/TestData/ctc_segmentation/eng/output${TIME}/verified_segments/nv_test_segments.txt && \ + rm -rf /home/TestData/ctc_segmentation/eng/output${TIME}' + } + } + stage('L2: Ru QN with .mp3') { + steps { + sh 'cd tools/ctc_segmentation && \ + TIME=`date +"%Y-%m-%d-%T"` && \ + /bin/bash run_sample.sh \ + --MODEL_NAME_OR_PATH=/home/TestData/ctc_segmentation/QuartzNet15x5-Ru-e512-wer14.45.nemo \ + --DATA_DIR=/home/TestData/ctc_segmentation/ru \ + --OUTPUT_DIR=/home/TestData/ctc_segmentation/ru/output${TIME} \ + --LANGUAGE=ru \ + --OFFSET=0 \ + --CUT_PREFIX=0 \ + --MIN_SEGMENT_LEN=0 \ + --AUDIO_FORMAT=.mp3 \ + --ADDITIONAL_SPLIT_SYMBOLS=";" && \ + python /home/TestData/ctc_segmentation/verify_alignment.py \ + -r /home/TestData/ctc_segmentation/ru/valid_ru_segments.txt \ + -g /home/TestData/ctc_segmentation/ru/output${TIME}/verified_segments/ru_segments.txt && \ + rm -rf /home/TestData/ctc_segmentation/ru/output${TIME}' + } + } + } + } + } + } + + stage('L2: Multi-GPU Megatron finetuning') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: Cased Megatron finetuning on MRPC') { + steps { + sh 'cd examples/nlp/glue_benchmark && \ + python glue_benchmark.py \ + model.dataset.data_dir=/home/TestData/nlp/glue_fake/MRPC \ + trainer.gpus=[0,1] \ + +trainer.fast_dev_run=true \ + model.dataset.use_cache=false \ + model.language_model.pretrained_model_name=megatron-bert-345m-cased \ + trainer.accelerator=ddp \ + exp_manager=null' + } + } + } + } + + stage('L2: SGD-QA') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: SGD-QA') { + steps { + sh 'cd examples/nlp/dialogue_state_tracking && \ + python sgd_qa.py \ + model.dataset.data_dir=/home/TestData/nlp/sgd_small \ + model.dataset.dialogues_example_dir=sgd_outputs \ + model.dataset.task_name=debug_sample \ + trainer.max_steps=1 \ + trainer.max_epochs=1 \ + model.train_ds.batch_size=2 \ + model.validation_ds.batch_size=2 \ + model.test_ds.batch_size=2 \ + model.nemo_path=null \ + trainer.val_check_interval=0.0 \ + trainer.gpus=[0,1] \ + model.dataset.use_cache=false \ + model.language_model.pretrained_model_name=bert-base-cased \ + trainer.accelerator=ddp \ + exp_manager=null && \ + rm -rf sgd_outputs' + } + } + stage('GLUE STS-b with AlBERT') { + steps { + sh 'python examples/nlp/glue_benchmark/glue_benchmark.py \ + model.dataset.use_cache=false \ + model.task_name=sts-b \ + model.dataset.data_dir=/home/TestData/nlp/glue_fake/STS-B \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=True \ + model.language_model.pretrained_model_name=albert-base-v1 \ + exp_manager=null' + } + } + } + } + + stage('L2: Parallel BERT SQUAD v1.1 / v2.0') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('BERT SQUAD 1.1') { + // Cannot do fast_dev_run because squad needs whole dev dataset + steps { + sh 'cd examples/nlp/question_answering && \ + python question_answering_squad.py \ + model.train_ds.file=/home/TestData/nlp/squad_mini/v1.1/train-v1.1.json \ + model.dataset.use_cache=false \ + model.validation_ds.file=/home/TestData/nlp/squad_mini/v1.1/dev-v1.1.json \ + model.test_ds.file=/home/TestData/nlp/squad_mini/v1.1/dev-v1.1.json \ + model.train_ds.batch_size=2 \ + model.train_ds.num_samples=2 \ + model.validation_ds.batch_size=2 \ + model.validation_ds.num_samples=2 \ + model.test_ds.num_samples=2 \ + model.test_ds.batch_size=2 \ + trainer.max_epochs=1 \ + +trainer.max_steps=1 \ + model.language_model.pretrained_model_name=bert-base-uncased \ + model.dataset.version_2_with_negative=false \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + trainer.gpus=[0] \ + exp_manager=null' + } + } + stage('BERT SQUAD 2.0') { + // Cannot do fast_dev_run because squad needs whole dev dataset + steps { + sh 'cd examples/nlp/question_answering && \ + python question_answering_squad.py \ + model.train_ds.file=/home/TestData/nlp/squad_mini/v2.0/train-v2.0.json \ + model.dataset.use_cache=false \ + model.train_ds.batch_size=2 \ + model.train_ds.num_samples=2 \ + model.validation_ds.batch_size=2 \ + model.validation_ds.num_samples=2 \ + trainer.max_epochs=1 \ + +trainer.max_steps=1 \ + model.validation_ds.file=/home/TestData/nlp/squad_mini/v2.0/dev-v2.0.json \ + model.language_model.pretrained_model_name=bert-base-uncased \ + model.dataset.version_2_with_negative=true \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + trainer.gpus=[1] \ + exp_manager=null' + } + } + } + } + // Runs out of memory on the 12G TITAN V (GPU 0 on main CI) + stage('L2: MegaBERT Token Classification') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + steps { + sh 'cd examples/nlp/token_classification && \ + python token_classification_train.py \ + model.dataset.data_dir=/home/TestData/nlp/token_classification_punctuation/ \ + model.language_model.pretrained_model_name=megatron-bert-345m-uncased \ + model.train_ds.batch_size=10 \ + model.dataset.max_seq_length=50 \ + model.dataset.use_cache=false \ + trainer.accelerator=ddp \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=true \ + exp_manager=null' + } + } + stage('L2: Parallel SQUAD v1.1 & v2.0') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('SQUAD v2.0 with Megatron with ckpt & config') { + // Cannot do fast_dev_run because squad needs whole dev dataset + steps { + sh 'cd examples/nlp/question_answering && \ + python question_answering_squad.py \ + model.train_ds.file=/home/TestData/nlp/squad_mini/v2.0/train-v2.0.json \ + model.dataset.use_cache=false \ + model.train_ds.batch_size=1 \ + model.train_ds.num_samples=1 \ + model.validation_ds.batch_size=1 \ + model.validation_ds.num_samples=1 \ + trainer.accelerator=ddp \ + trainer.max_epochs=1 \ + +trainer.max_steps=1 \ + model.validation_ds.file=/home/TestData/nlp/squad_mini/v2.0/dev-v2.0.json \ + model.language_model.pretrained_model_name=megatron-bert-uncased \ + model.language_model.lm_checkpoint=/home/TestData/nlp/megatron_345m_uncased/model_optim_rng.pt \ + model.language_model.config_file=/home/TestData/nlp/megatron_345m_uncased/345m_config.json \ + model.dataset.version_2_with_negative=true \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + trainer.gpus=[1] \ + exp_manager=null' + } + } + stage('RoBERTa SQUAD 1.1') { + // Cannot do fast_dev_run because squad needs whole dev dataset + steps { + sh 'cd examples/nlp/question_answering && \ + python question_answering_squad.py \ + model.train_ds.file=/home/TestData/nlp/squad_mini/v1.1/train-v1.1.json \ + model.dataset.use_cache=false \ + model.train_ds.batch_size=2 \ + model.train_ds.num_samples=2 \ + model.validation_ds.batch_size=2 \ + model.validation_ds.num_samples=2 \ + trainer.max_epochs=1 \ + +trainer.max_steps=1 \ + model.validation_ds.file=/home/TestData/nlp/squad_mini/v1.1/dev-v1.1.json \ + model.language_model.pretrained_model_name=roberta-base \ + model.dataset.version_2_with_negative=false \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + trainer.gpus=[0] \ + exp_manager=null' + } + } + stage ('Text Classification with BERT Test') { + steps { + sh 'cd examples/nlp/text_classification && \ + python text_classification_with_bert.py \ + model.dataset.num_classes=6 \ + model.train_ds.file_path=/home/TestData/nlp/retail_text_classification/train.tsv \ + model.validation_ds.file_path=/home/TestData/nlp/retail_text_classification/dev.tsv \ + model.language_model.pretrained_model_name=bert-base-uncased \ + model.train_ds.batch_size=10 \ + model.dataset.max_seq_length=50 \ + model.dataset.use_cache=false \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + exp_manager=null' + } + } + stage('L2: Intent and Slot Classification') { + steps { + sh 'cd examples/nlp/intent_slot_classification && \ + python intent_slot_classification.py \ + model.data_dir=/home/TestData/nlp/retail \ + model.validation_ds.prefix=dev \ + model.test_ds.prefix=dev \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + exp_manager.exp_dir=checkpoints' + sh 'rm -rf checkpoints' + } + } + } + } + + stage('L2: Model Parallel Size 2 Megatron Text Classification') { + when { + anyOf{ + branch 'main' + changeRequest target: 'main' + } + } + failFast true + steps{ + sh 'cd examples/nlp/text_classification && \ + python text_classification_with_bert.py \ + trainer.gpus=[0,1] \ + trainer.num_nodes=1 \ + trainer.precision=16 \ + trainer.gradient_clip_val=1.0 \ + ~trainer.amp_level \ + +trainer.fast_dev_run=true \ + model.dataset.num_classes=6 \ + model.train_ds.file_path=/home/TestData/nlp/retail_text_classification/train.tsv \ + model.train_ds.batch_size=4 \ + model.language_model.pretrained_model_name=megatron-bert-uncased \ + model.language_model.config_file=/home/TestData/nlp/mp_2_bert_toy/config.json \ + model.language_model.lm_checkpoint=/home/TestData/nlp/mp_2_bert_toy/iter_2000000 \ + model.nemo_path=null \ + ~model.infer_samples \ + exp_manager=null' + } + } + + stage('L2: Model Parallel Size 2 Megatron Autoresume') { + when { + anyOf{ + branch 'main' + changeRequest target: 'main' + } + } + failFast true + steps{ + sh 'cd examples/nlp/text_classification && \ + python text_classification_with_bert.py \ + trainer.gpus=[0,1] \ + trainer.num_nodes=1 \ + trainer.precision=16 \ + trainer.gradient_clip_val=1.0 \ + trainer.max_epochs=1 \ + ~trainer.amp_level \ + +trainer.fast_dev_run=true \ + model.dataset.num_classes=6 \ + model.train_ds.file_path=/home/TestData/nlp/retail_text_classification/train.tsv \ + model.train_ds.batch_size=4 \ + model.language_model.pretrained_model_name=megatron-bert-uncased \ + model.language_model.config_file=/home/TestData/nlp/mp_2_bert_toy/config.json \ + model.language_model.lm_checkpoint=/home/TestData/nlp/mp_2_bert_toy/iter_2000000 \ + model.nemo_path=null \ + ~model.infer_samples \ + +exp_manager.explicit_log_dir=/home/TestData/nlp/mp_autoresume \ + +exp_manager.resume_if_exists=true' + } + } + + stage('L2: Model Parallel Size 2 Megatron Evaluation from .nemo') { + when { + anyOf{ + branch 'main' + changeRequest target: 'main' + } + } + failFast true + steps{ + sh 'cd examples/nlp/text_classification && \ + python model_parallel_text_classification_evaluation.py \ + trainer.gpus=[0,1] \ + trainer.num_nodes=1 \ + model.dataset.num_classes=6 \ + model.test_ds.file_path=/home/TestData/nlp/retail_text_classification/dev.tsv \ + model.nemo_path=/home/TestData/nlp/mp_2_nemo/retail_text_class_350M.nemo \ + exp_manager=null' + } + } + + stage('L2: Model Parallel Size 2 Megatron Train from .nemo') { + when { + anyOf{ + branch 'main' + changeRequest target: 'main' + } + } + failFast true + steps{ + sh 'cd examples/nlp/token_classification && \ + python token_classification_train.py \ + pretrained_model=/home/TestData/nlp/mp_2_nemo/ner_350M.nemo \ + model.dataset.data_dir=/home/TestData/nlp/ner/ \ + model.train_ds.batch_size=2 \ + model.dataset.use_cache=false \ + trainer.gpus=[0,1] \ + +trainer.fast_dev_run=true \ + model.dataset.class_balancing="weighted_loss" \ + exp_manager=null' + } + } + + stage('L2: Parallel NLP Examples 2') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage ('NER finetuning from pretrained Test') { + steps { + sh 'cd examples/nlp/token_classification && \ + python token_classification_train.py \ + pretrained_model=ner_en_bert \ + model.dataset.data_dir=/home/TestData/nlp/ner/ \ + model.train_ds.batch_size=2 \ + model.dataset.use_cache=false \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + model.dataset.class_balancing="weighted_loss" \ + exp_manager.exp_dir=null' + } + } + stage ('Punctuation and capitalization finetuning from pretrained test') { + steps { + sh 'cd examples/nlp/token_classification && \ + python punctuation_capitalization_train.py \ + pretrained_model=punctuation_en_bert \ + model.dataset.data_dir=/home/TestData/nlp/token_classification_punctuation/ \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=true \ + model.dataset.use_cache=false \ + exp_manager.exp_dir=null' + } + } + stage ('NER with TurkuNLP/bert-base-finnish-cased-v1') { + steps { + sh 'cd examples/nlp/token_classification && \ + python token_classification_train.py \ + model.dataset.data_dir=/home/TestData/nlp/token_classification_punctuation/ \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + model.dataset.use_cache=false \ + model.language_model.pretrained_model_name="TurkuNLP/bert-base-finnish-cased-v1" \ + exp_manager.exp_dir=null' + } + } + stage('Evaluation script for Token Classification') { + steps { + sh 'python examples/nlp/token_classification/token_classification_evaluate.py \ + model.dataset.data_dir=/home/TestData/nlp/ner/ \ + pretrained_model=/home/TestData/nlp/pretrained_models/NER_Model_with_BERT_base_uncased.nemo && \ + rm -rf nemo_experiments' + } + } + stage('Evaluation script for Punctuation') { + steps { + sh 'python examples/nlp/token_classification/punctuation_capitalization_evaluate.py \ + model.dataset.data_dir=/home/TestData/nlp/token_classification_punctuation/ \ + pretrained_model=/home/TestData/nlp/pretrained_models/Punctuation_Capitalization_with_DistilBERT_base_uncased.nemo && \ + rm -rf nemo_experiments' + } + } + stage('L2: Punctuation & Capitalization, 2GPUs with DistilBERT') { + steps { + sh 'cd examples/nlp/token_classification && \ + python punctuation_capitalization_train.py \ + model.dataset.data_dir=/home/TestData/nlp/token_classification_punctuation/ \ + model.language_model.pretrained_model_name=distilbert-base-uncased \ + model.dataset.use_cache=false \ + trainer.gpus=[0,1] \ + trainer.accelerator=ddp \ + trainer.max_epochs=1 \ + +exp_manager.explicit_log_dir=/home/TestData/nlp/token_classification_punctuation/output && \ + python punctuation_capitalization_evaluate.py \ + pretrained_model=/home/TestData/nlp/token_classification_punctuation/output/checkpoints/Punctuation_and_Capitalization.nemo && \ + rm -rf /home/TestData/nlp/token_classification_punctuation/output/*' + } + } + } + } + + stage('L2: Parallel Pretraining BERT pretraining from Text/Preprocessed') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: Pretraining BERT pretraining from Text') { + steps { + sh 'cd examples/nlp/language_modeling && \ + python bert_pretraining.py \ + --config-name=bert_pretraining_from_text_config.yaml \ + trainer.gpus=[0] \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + +trainer.fast_dev_run=true \ + model.train_ds.data_file=/home/TestData/nlp/wikitext-2/train.txt \ + model.train_ds.batch_size=32 \ + model.validation_ds.data_file=/home/TestData/nlp/wikitext-2/valid.txt \ + model.validation_ds.batch_size=32 \ + model.language_model.config_file=/home/TestData/nlp/bert_configs/bert_3200.json \ + model.optim.lr=0.01 \ + model.optim.sched.warmup_ratio=0.1 \ + model.tokenizer.tokenizer_name=sentencepiece \ + model.tokenizer.tokenizer_model=/home/TestData/nlp/wikitext-2/tokenizer_bpe_v3193/tokenizer.model \ + model.mask_prob=0.15 \ + model.short_seq_prob=0.1 \ + exp_manager.exp_dir=PretrainingBERTFromText \ + ' + sh 'rm -f /home/TestData/nlp/wikitext-2/*.pkl' + sh 'rm -rf examples/nlp/language_modeling/PretrainingBERTFromText' + sh 'ls -lha examples/nlp/language_modeling' + } + } + stage('L2: Pretraining BERT from Preprocessed') { + steps { + sh 'cd examples/nlp/language_modeling && \ + python bert_pretraining.py \ + --config-name=bert_pretraining_from_preprocessed_config.yaml \ + trainer.gpus=[1] \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + +trainer.fast_dev_run=true \ + model.train_ds.data_file=/home/TestData/nlp/wiki_book_mini/training \ + model.train_ds.batch_size=8 \ + model.language_model.lm_checkpoint=/home/TestData/nlp/bert_ckpts/nemo1.0/bert_base_uncased_mlm_final_1074591_nemo1.0.pt \ + model.language_model.config_file=/home/TestData/nlp/bert_configs/uncased_L-12_H-768_A-12.json \ + model.optim.lr=0.875e-4 \ + model.optim.weight_decay=0.01 \ + model.optim.sched.warmup_ratio=0.01 \ + exp_manager.exp_dir=PretrainingBERTFromPreprocessed \ + exp_manager.create_checkpoint_callback=False \ + ' + sh 'rm -rf examples/nlp/language_modeling/PretrainingBERTFromPreprocessed' + sh 'ls -lha examples/nlp/language_modeling' + } + } + stage('L2: Pretraining BERT pretraining from Text with char tokenizer') { + steps { + sh 'cd examples/nlp/language_modeling && \ + python bert_pretraining.py \ + --config-name=bert_pretraining_from_text_config.yaml \ + trainer.gpus=[0] \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + +trainer.fast_dev_run=true \ + model.train_ds.data_file=/home/TestData/nlp/wikitext-2/train.txt \ + model.train_ds.batch_size=32 \ + model.validation_ds.data_file=/home/TestData/nlp/wikitext-2/valid.txt \ + model.validation_ds.batch_size=32 \ + model.language_model.config_file=/home/TestData/nlp/bert_configs/bert_3200.json \ + model.optim.lr=0.01 \ + model.optim.sched.warmup_ratio=0.1 \ + model.tokenizer.tokenizer_name=char \ + model.tokenizer.vocab_file=/home/TestData/nlp/vocabs/mini_vocab.txt \ + model.mask_prob=0.15 \ + model.short_seq_prob=0.1 \ + exp_manager.exp_dir=PretrainingBERTFromTextchartok \ + ' + sh 'rm -rf examples/nlp/language_modeling/PretrainingBERTFromTextchartok' + } + } + stage('L2: Pretraining BERT pretraining from Text with word tokenizer') { + steps { + sh 'cd examples/nlp/language_modeling && \ + python bert_pretraining.py \ + --config-name=bert_pretraining_from_text_config.yaml \ + trainer.gpus=[1] \ + trainer.precision=16 \ + trainer.amp_level=O1 \ + +trainer.fast_dev_run=true \ + model.train_ds.data_file=/home/TestData/nlp/wikitext-2/train.txt \ + model.train_ds.batch_size=32 \ + model.validation_ds.data_file=/home/TestData/nlp/wikitext-2/valid.txt \ + model.validation_ds.batch_size=32 \ + model.language_model.config_file=/home/TestData/nlp/bert_configs/bert_3200.json \ + model.optim.lr=0.01 \ + model.optim.sched.warmup_ratio=0.1 \ + model.tokenizer.tokenizer_name=word \ + model.tokenizer.vocab_file=/home/TestData/nlp/vocabs/mini_vocab.txt \ + model.mask_prob=0.15 \ + model.short_seq_prob=0.1 \ + exp_manager.exp_dir=PretrainingBERTFromTextwordtok \ + ' + sh 'rm -rf examples/nlp/language_modeling/PretrainingBERTFromTextwordtok' + } + } + } + } + + stage('L2: Entity Linking') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage ('Self Alignment Pretraining BERT') { + steps { + sh 'cd examples/nlp/entity_linking && \ + python self_alignment_pretraining.py \ + project_dir=. \ + trainer.val_check_interval=3 \ + model.raw_data=None \ + model.train_ds.data_file=/home/TestData/nlp/entity_linking/tiny_example_train_pairs.tsv \ + model.validation_ds.data_file=/home/TestData/nlp/entity_linking/tiny_example_validation_pairs.tsv \ + model.train_ds.batch_size=8 \ + model.validation_ds.batch_size=8 \ + exp_manager.exp_dir=null' + } + } + } + } + + stage('L2: NMT Attention is All You Need Training') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: NMT Training Post-LN') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=aayn_base \ + do_testing=true \ + model.train_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.train_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.ref \ + model.validation_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.validation_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.encoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + model.decoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + +trainer.limit_test_batches=2 \ + exp_manager=null \ + ' + } + } + + stage('L2: NMT Training Pre-LN') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=aayn_base \ + do_testing=true \ + model.train_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.train_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.ref \ + model.validation_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.validation_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.encoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + model.decoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + model.encoder.pre_ln=true \ + model.decoder.pre_ln=true \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=true \ + +trainer.limit_test_batches=2 \ + exp_manager=null \ + ' + } + } + stage('L2: NMT Multi-Validation') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=aayn_base \ + do_testing=true \ + model.train_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-en-de.src \ + model.train_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-en-de.ref \ + model.validation_ds.src_file_name=[/home/TestData/nlp/nmt/toy_data/wmt13-en-de.src,/home/TestData/nlp/nmt/toy_data/wmt14-en-de.src] \ + model.validation_ds.tgt_file_name=[/home/TestData/nlp/nmt/toy_data/wmt13-en-de.ref,/home/TestData/nlp/nmt/toy_data/wmt14-en-de.ref] \ + model.test_ds.src_file_name=[/home/TestData/nlp/nmt/toy_data/wmt13-en-de.src,/home/TestData/nlp/nmt/toy_data/wmt14-en-de.src] \ + model.test_ds.tgt_file_name=[/home/TestData/nlp/nmt/toy_data/wmt13-en-de.ref,/home/TestData/nlp/nmt/toy_data/wmt14-en-de.ref] \ + model.encoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + model.decoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + +trainer.limit_test_batches=2 \ + exp_manager=null \ + ' + } + } + } + } + stage('L2: NMT Attention is All You Need Inference') { + when { + anyOf { + branch 'v1.0.0' + changeRequest target: 'v1.0.0' + } + } + failFast true + parallel { + stage('L2: NMT Inference - PostLN') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python nmt_transformer_infer.py \ + --model=/home/TestData/nlp/nmt/toy_data/TransformerLargeDe-En.nemo \ + --srctext=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.test.src \ + --tgtout=/home/TestData/nlp/nmt/toy_data/out.txt \ + --target_lang en \ + --source_lang de \ + ' + } + } + stage('L2: NMT Inference - Pre-LN') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python nmt_transformer_infer.py \ + --model=/home/TestData/nlp/nmt/toy_data/en_de_24x6_preln.nemo \ + --srctext=/home/TestData/nlp/nmt/toy_data/wmt14-en-de.test.src \ + --tgtout=/home/TestData/nlp/nmt/toy_data/out.txt \ + --target_lang de \ + --source_lang en \ + ' + } + } + } + } + stage('L2: NMT with HuggingFace') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: NMT Pretrained HF Encoder') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=huggingface \ + model.shared_tokenizer=False \ + model.encoder_tokenizer.library=huggingface \ + model.encoder.library=huggingface \ + model.encoder.model_name=bert-base-cased \ + model.encoder.pretrained=true \ + model.train_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.train_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.ref \ + model.validation_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.validation_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.decoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + exp_manager=null \ + ' + } + } + + stage('L2: NMT Custom HF Encoder') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=huggingface \ + model.shared_tokenizer=True \ + model.encoder_tokenizer.library=yttm \ + model.encoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + model.encoder.library=huggingface \ + model.encoder.model_name=null \ + model.encoder.pretrained=false \ + +model.encoder._target_=transformers.BertConfig \ + +model.encoder.hidden_size=1536 \ + model.train_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.train_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.ref \ + model.validation_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.validation_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.test_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.decoder_tokenizer.tokenizer_model=/home/TestData/nlp/nmt/toy_data/tt_tokenizer.BPE.4096.model \ + trainer.gpus=[1] \ + +trainer.fast_dev_run=true \ + exp_manager=null \ + ' + } + } + } + } + + stage('L2: NMT Tarred Dataset Creation') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + parallel { + stage('L2: NMT Auto Tarred Dataset Creation') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=aayn_base \ + do_training=false \ + model.preproc_out_dir=$PWD/preproc_out_dir \ + model.train_ds.use_tarred_dataset=true \ + model.train_ds.n_preproc_jobs=2 \ + model.train_ds.lines_per_dataset_fragment=500 \ + model.train_ds.num_batches_per_tarfile=10 \ + model.train_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.train_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.ref \ + model.validation_ds.src_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.validation_ds.tgt_file_name=/home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + model.encoder_tokenizer.vocab_size=2000 \ + model.decoder_tokenizer.vocab_size=2000 \ + ~model.test_ds \ + trainer.gpus=[0] \ + +trainer.fast_dev_run=true \ + exp_manager=null \ + ' + } + } + + stage('L2: NMT Script Tarred Dataset Creation') { + steps { + sh 'cd examples/nlp/machine_translation && \ + python create_tarred_parallel_dataset.py \ + --src_fname /home/TestData/nlp/nmt/toy_data/wmt14-de-en.src \ + --tgt_fname /home/TestData/nlp/nmt/toy_data/wmt14-de-en.ref \ + --out_dir $PWD/out_dir \ + --encoder_tokenizer_vocab_size=2000 \ + --decoder_tokenizer_vocab_size=2000 \ + --tokens_in_batch=1000 \ + --lines_per_dataset_fragment=500 \ + --num_batches_per_tarfile=10 \ + --n_preproc_jobs=2 \ + ' + } + } + } + } + + stage('L2: TTS Fast dev runs 1') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + parallel { + stage('Tacotron 2') { + steps { + sh 'python examples/tts/tacotron2.py \ + train_dataset=/home/TestData/an4_dataset/an4_train.json \ + validation_datasets=/home/TestData/an4_dataset/an4_val.json \ + trainer.gpus="[0]" \ + +trainer.fast_dev_run=True \ + trainer.accelerator=null \ + trainer.max_epochs=-1 \ + model.train_ds.dataloader_params.batch_size=12 \ + model.validation_ds.dataloader_params.batch_size=12 \ + ~trainer.check_val_every_n_epoch' + } + } + // stage('FastPitch') { + // steps { + // sh 'python examples/tts/fastpitch.py \ + // train_dataset=/home/TestData/an4_dataset/an4_train.json \ + // validation_datasets=/home/TestData/an4_dataset/an4_val.json \ + // trainer.gpus="[0]" \ + // +trainer.fast_dev_run=True \ + // trainer.accelerator=null \ + // trainer.max_epochs=-1 \ + // model.train_ds.batch_size=12 \ + // model.train_ds.num_workers=1 \ + // model.validation_ds.batch_size=12 \ + // model.validation_ds.num_workers=1 \ + // ~trainer.check_val_every_n_epoch' + // } + // } + stage('WaveGlow') { + steps { + sh 'python examples/tts/waveglow.py \ + train_dataset=/home/TestData/an4_dataset/an4_train.json \ + validation_datasets=/home/TestData/an4_dataset/an4_val.json \ + trainer.gpus="[1]" \ + +trainer.fast_dev_run=True \ + trainer.accelerator=null \ + trainer.max_epochs=-1 \ + model.train_ds.dataloader_params.batch_size=4 \ + model.validation_ds.dataloader_params.batch_size=4 \ + ~trainer.check_val_every_n_epoch' + } + } + } + } + + stage('L2: TTS Fast dev runs 2') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + + parallel { + stage('MelGAN') { + steps { + sh 'python examples/tts/melgan.py \ + train_dataset=/home/TestData/an4_dataset/an4_train.json \ + validation_datasets=/home/TestData/an4_dataset/an4_val.json \ + trainer.gpus="[0]" \ + +trainer.fast_dev_run=True \ + trainer.accelerator=ddp \ + trainer.max_epochs=-1 \ + model.train_ds.dataloader_params.batch_size=4 \ + model.validation_ds.dataloader_params.batch_size=4 \ + ~trainer.check_val_every_n_epoch' + } + } + stage('SqueezeWave') { + steps { + sh 'python examples/tts/squeezewave.py \ + train_dataset=/home/TestData/an4_dataset/an4_train.json \ + validation_datasets=/home/TestData/an4_dataset/an4_val.json \ + trainer.gpus="[0]" \ + +trainer.fast_dev_run=True \ + trainer.accelerator=null \ + trainer.max_epochs=-1 \ + model.train_ds.dataloader_params.batch_size=4 \ + model.validation_ds.dataloader_params.batch_size=4 \ + ~trainer.check_val_every_n_epoch' + } + } + stage('GlowTTS') { + steps { + sh 'python examples/tts/glow_tts.py \ + train_dataset=/home/TestData/an4_dataset/an4_train.json \ + validation_datasets=/home/TestData/an4_dataset/an4_val.json \ + trainer.gpus="[1]" \ + +trainer.fast_dev_run=True \ + trainer.accelerator=null \ + trainer.max_epochs=-1 \ + model.train_ds.batch_size=4 \ + model.validation_ds.batch_size=4 \ + ~trainer.check_val_every_n_epoch' + } + } + } + } + + stage('L??: Speech Checkpoints tests') { + when { + anyOf { + branch 'main' + changeRequest target: 'main' + } + } + failFast true + steps { + sh 'CUDA_VISIBLE_DEVICES=0 python examples/asr/speech_to_text_infer.py --asr_model QuartzNet15x5Base-En --dataset /home/TestData/librispeech/librivox-dev-other.json --wer_tolerance 0.1012 --batch_size 64' + } + } + } + + post { + always { + sh 'chmod -R 777 .' + cleanWs() + } + } +} diff --git a/models/tts/NeMo/LICENSE b/models/tts/NeMo/LICENSE new file mode 100644 index 0000000000..f49a4e16e6 --- /dev/null +++ b/models/tts/NeMo/LICENSE @@ -0,0 +1,201 @@ + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + 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. \ No newline at end of file diff --git a/models/tts/NeMo/README.rst b/models/tts/NeMo/README.rst new file mode 100644 index 0000000000..585064f17e --- /dev/null +++ b/models/tts/NeMo/README.rst @@ -0,0 +1,174 @@ + +|status| |documentation| |license| |lgtm_grade| |lgtm_alerts| |black| + +.. |status| image:: http://www.repostatus.org/badges/latest/active.svg + :target: http://www.repostatus.org/#active + :alt: Project Status: Active – The project has reached a stable, usable state and is being actively developed. + +.. |documentation| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=main + :alt: Documentation + :target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/ + +.. |license| image:: https://img.shields.io/badge/License-Apache%202.0-brightgreen.svg + :target: https://github.com/NVIDIA/NeMo/blob/master/LICENSE + :alt: NeMo core license and license for collections in this repo + +.. |lgtm_grade| image:: https://img.shields.io/lgtm/grade/python/g/NVIDIA/NeMo.svg?logo=lgtm&logoWidth=18 + :target: https://lgtm.com/projects/g/NVIDIA/NeMo/context:python + :alt: Language grade: Python + +.. |lgtm_alerts| image:: https://img.shields.io/lgtm/alerts/g/NVIDIA/NeMo.svg?logo=lgtm&logoWidth=18 + :target: https://lgtm.com/projects/g/NVIDIA/NeMo/alerts/ + :alt: Total alerts + +.. |black| image:: https://img.shields.io/badge/code%20style-black-000000.svg + :target: https://github.com/psf/black + :alt: Code style: black + +.. _main-readme: +**NVIDIA NeMo** +=============== + +Introduction +------------ + +NVIDIA NeMo is a conversational AI toolkit built for researchers working on automatic speech recognition (ASR), natural language processing (NLP), and text-to-speech synthesis (TTS). +The primary objective of NeMo is to help researchers from industry and academia to reuse prior work (code and pretrained models and make it easier to create new `conversational AI models `_. + + +`Introductory video. `_ + +Key Features +------------ + +* Speech processing + * `Automatic Speech recognition (ASR) `_: Jasper, QuartzNet, CitriNet, Conformer + * `Speech Classification `_: MatchboxNet (command recognition), MarbleNet (voice activity detection) + * `Speaker Recognition `_: SpeakerNet + * `Speaker Diarization `_: MarbleNet + SpeakerNet + * `NGC collection of pre-trained speech processing models. `_ +* Natural Language Processing + * `Compatible with Hugging Face Transformers and NVIDIA Megatron `_ + * `Neural Machine Translation (NMT) `_ + * `Punctuation and Capitalization `_ + * `Token classification (named entity recognition) `_ + * `Text classification `_ + * `Joint Intent and Slot Classification `_ + * `BERT pre-training `_ + * `Question answering `_ + * `GLUE benchmark `_ + * `Information retrieval `_ + * `NGC collection of pre-trained NLP models. `_ +* `Speech synthesis (TTS) `_ + * Spectrogram generation: Tacotron2, GlowTTS, FastSpeech2, FastPitch, + * Vocoders: WaveGlow, SqueezeWave, UniGlow, MelGAN, HiFiGAN + * End-to-end speech generation: FastPitch_HifiGan_E2E, FastSpeech2_HifiGan_E2E + * `NGC collection of pre-trained TTS models. `_ + + +Built for speed, NeMo can utilize NVIDIA's Tensor Cores and scale out training to multiple GPUs and multiple nodes. + +Requirements +------------ + +1) Python 3.6, 3.7 or 3.8 +2) Pytorch 1.8.1 or above +3) NVIDIA GPU for training + +Documentation +------------- + +.. |main| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=main + :alt: Documentation Status + :scale: 100% + :target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/ + +.. |v1.0.2| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=v1.0.2 + :alt: Documentation Status + :scale: 100% + :target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/v1.0.2/ + +.. |stable| image:: https://readthedocs.com/projects/nvidia-nemo/badge/?version=stable + :alt: Documentation Status + :scale: 100% + :target: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/ + ++---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+ +| Version | Status | Description | ++=========+=============+==================================================================================================================================+ +| Latest | |main| | `Documentation of the latest (i.e. main) branch. `_ | ++---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+ +| Next | |v1.0.2| | `Documentation of the most recent release: v1.0.2 `_ | ++---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+ +| Stable | |stable| | `Documentation of the stable (i.e. stable) branch. `_ | ++---------+-------------+----------------------------------------------------------------------------------------------------------------------------------+ + +Tutorials +--------- +A great way to start with NeMo is by checking `one of our tutorials `_. + +Getting help with NeMo +---------------------- +FAQ can be found on NeMo's `Discussions board `_. You are welcome to ask questions or start discussions there. + + +Installation +------------ + +Pip +~~~ +Use this installation mode if you want the latest released version. + +.. code-block:: bash + + apt-get update && apt-get install -y libsndfile1 ffmpeg + pip install Cython + pip install nemo_toolkit['all'] + +Pip from source +~~~~~~~~~~~~~~~ +Use this installation mode if you want the a version from particular GitHub branch (e.g main). + +.. code-block:: bash + + apt-get update && apt-get install -y libsndfile1 ffmpeg + pip install Cython + python -m pip install git+https://github.com/NVIDIA/NeMo.git@{BRANCH}#egg=nemo_toolkit[all] + + +From source +~~~~~~~~~~~ +Use this installation mode if you are contributing to NeMo. + +.. code-block:: bash + + apt-get update && apt-get install -y libsndfile1 ffmpeg + git clone https://github.com/NVIDIA/NeMo + cd NeMo + ./reinstall.sh + +Docker containers: +~~~~~~~~~~~~~~~~~~ + +If you chose to work with main branch, we recommend using NVIDIA's PyTorch container version 21.05-py3 and then installing from GitHub. + +.. code-block:: bash + + docker run --gpus all -it --rm -v :/NeMo --shm-size=8g \ + -p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit \ + stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:21.05-py3 + +Examples +-------- + +Many example can be found under `"Examples" `_ folder. + + +Contributing +------------ + +We welcome community contributions! Please refer to the `CONTRIBUTING.md `_ CONTRIBUTING.md for the process. + +License +------- +NeMo is under `Apache 2.0 license `_. diff --git a/models/tts/NeMo/__init__.py b/models/tts/NeMo/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/models/tts/NeMo/docs/.nojekyll b/models/tts/NeMo/docs/.nojekyll new file mode 100644 index 0000000000..e69de29bb2 diff --git a/models/tts/NeMo/docs/Makefile b/models/tts/NeMo/docs/Makefile new file mode 100644 index 0000000000..417fe2a0b1 --- /dev/null +++ b/models/tts/NeMo/docs/Makefile @@ -0,0 +1,216 @@ +# Makefile for Sphinx documentation +# + +# You can set these variables from the command line. +SPHINXOPTS = +SPHINXBUILD = sphinx-build +PAPER = +BUILDDIR = build + +# User-friendly check for sphinx-build +ifeq ($(shell which $(SPHINXBUILD) >/dev/null 2>&1; echo $$?), 1) +$(error The '$(SPHINXBUILD)' command was not found. Make sure you have Sphinx installed, then set the SPHINXBUILD environment variable to point to the full path of the '$(SPHINXBUILD)' executable. Alternatively you can add the directory with the executable to your PATH. If you don't have Sphinx installed, grab it from http://sphinx-doc.org/) +endif + +# Internal variables. +PAPEROPT_a4 = -D latex_paper_size=a4 +PAPEROPT_letter = -D latex_paper_size=letter +ALLSPHINXOPTS = -d $(BUILDDIR)/doctrees $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) source +# the i18n builder cannot share the environment and doctrees with the others +I18NSPHINXOPTS = $(PAPEROPT_$(PAPER)) $(SPHINXOPTS) source + +.PHONY: help +help: + @echo "Please use \`make ' where is one of" + @echo " html to make standalone HTML files" + @echo " dirhtml to make HTML files named index.html in directories" + @echo " singlehtml to make a single large HTML file" + @echo " pickle to make pickle files" + @echo " json to make JSON files" + @echo " htmlhelp to make HTML files and a HTML help project" + @echo " qthelp to make HTML files and a qthelp project" + @echo " applehelp to make an Apple Help Book" + @echo " devhelp to make HTML files and a Devhelp project" + @echo " epub to make an epub" + @echo " latex to make LaTeX files, you can set PAPER=a4 or PAPER=letter" + @echo " latexpdf to make LaTeX files and run them through pdflatex" + @echo " latexpdfja to make LaTeX files and run them through platex/dvipdfmx" + @echo " text to make text files" + @echo " man to make manual pages" + @echo " texinfo to make Texinfo files" + @echo " info to make Texinfo files and run them through makeinfo" + @echo " gettext to make PO message catalogs" + @echo " changes to make an overview of all changed/added/deprecated items" + @echo " xml to make Docutils-native XML files" + @echo " pseudoxml to make pseudoxml-XML files for display purposes" + @echo " linkcheck to check all external links for integrity" + @echo " doctest to run all doctests embedded in the documentation (if enabled)" + @echo " coverage to run coverage check of the documentation (if enabled)" + +.PHONY: clean +clean: + rm -rf $(BUILDDIR)/* + +.PHONY: html +html: + $(SPHINXBUILD) -b html $(ALLSPHINXOPTS) $(BUILDDIR)/html + @echo + @echo "Build finished. The HTML pages are in $(BUILDDIR)/html." + +.PHONY: dirhtml +dirhtml: + $(SPHINXBUILD) -b dirhtml $(ALLSPHINXOPTS) $(BUILDDIR)/dirhtml + @echo + @echo "Build finished. The HTML pages are in $(BUILDDIR)/dirhtml." + +.PHONY: singlehtml +singlehtml: + $(SPHINXBUILD) -b singlehtml $(ALLSPHINXOPTS) $(BUILDDIR)/singlehtml + @echo + @echo "Build finished. The HTML page is in $(BUILDDIR)/singlehtml." + +.PHONY: pickle +pickle: + $(SPHINXBUILD) -b pickle $(ALLSPHINXOPTS) $(BUILDDIR)/pickle + @echo + @echo "Build finished; now you can process the pickle files." + +.PHONY: json +json: + $(SPHINXBUILD) -b json $(ALLSPHINXOPTS) $(BUILDDIR)/json + @echo + @echo "Build finished; now you can process the JSON files." + +.PHONY: htmlhelp +htmlhelp: + $(SPHINXBUILD) -b htmlhelp $(ALLSPHINXOPTS) $(BUILDDIR)/htmlhelp + @echo + @echo "Build finished; now you can run HTML Help Workshop with the" \ + ".hhp project file in $(BUILDDIR)/htmlhelp." + +.PHONY: qthelp +qthelp: + $(SPHINXBUILD) -b qthelp $(ALLSPHINXOPTS) $(BUILDDIR)/qthelp + @echo + @echo "Build finished; now you can run "qcollectiongenerator" with the" \ + ".qhcp project file in $(BUILDDIR)/qthelp, like this:" + @echo "# qcollectiongenerator $(BUILDDIR)/qthelp/OpenSeq2Seq.qhcp" + @echo "To view the help file:" + @echo "# assistant -collectionFile $(BUILDDIR)/qthelp/OpenSeq2Seq.qhc" + +.PHONY: applehelp +applehelp: + $(SPHINXBUILD) -b applehelp $(ALLSPHINXOPTS) $(BUILDDIR)/applehelp + @echo + @echo "Build finished. The help book is in $(BUILDDIR)/applehelp." + @echo "N.B. You won't be able to view it unless you put it in" \ + "~/Library/Documentation/Help or install it in your application" \ + "bundle." + +.PHONY: devhelp +devhelp: + $(SPHINXBUILD) -b devhelp $(ALLSPHINXOPTS) $(BUILDDIR)/devhelp + @echo + @echo "Build finished." + @echo "To view the help file:" + @echo "# mkdir -p $$HOME/.local/share/devhelp/OpenSeq2Seq" + @echo "# ln -s $(BUILDDIR)/devhelp $$HOME/.local/share/devhelp/OpenSeq2Seq" + @echo "# devhelp" + +.PHONY: epub +epub: + $(SPHINXBUILD) -b epub $(ALLSPHINXOPTS) $(BUILDDIR)/epub + @echo + @echo "Build finished. The epub file is in $(BUILDDIR)/epub." + +.PHONY: latex +latex: + $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex + @echo + @echo "Build finished; the LaTeX files are in $(BUILDDIR)/latex." + @echo "Run \`make' in that directory to run these through (pdf)latex" \ + "(use \`make latexpdf' here to do that automatically)." + +.PHONY: latexpdf +latexpdf: + $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex + @echo "Running LaTeX files through pdflatex..." + $(MAKE) -C $(BUILDDIR)/latex all-pdf + @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." + +.PHONY: latexpdfja +latexpdfja: + $(SPHINXBUILD) -b latex $(ALLSPHINXOPTS) $(BUILDDIR)/latex + @echo "Running LaTeX files through platex and dvipdfmx..." + $(MAKE) -C $(BUILDDIR)/latex all-pdf-ja + @echo "pdflatex finished; the PDF files are in $(BUILDDIR)/latex." + +.PHONY: text +text: + $(SPHINXBUILD) -b text $(ALLSPHINXOPTS) $(BUILDDIR)/text + @echo + @echo "Build finished. The text files are in $(BUILDDIR)/text." + +.PHONY: man +man: + $(SPHINXBUILD) -b man $(ALLSPHINXOPTS) $(BUILDDIR)/man + @echo + @echo "Build finished. The manual pages are in $(BUILDDIR)/man." + +.PHONY: texinfo +texinfo: + $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo + @echo + @echo "Build finished. The Texinfo files are in $(BUILDDIR)/texinfo." + @echo "Run \`make' in that directory to run these through makeinfo" \ + "(use \`make info' here to do that automatically)." + +.PHONY: info +info: + $(SPHINXBUILD) -b texinfo $(ALLSPHINXOPTS) $(BUILDDIR)/texinfo + @echo "Running Texinfo files through makeinfo..." + make -C $(BUILDDIR)/texinfo info + @echo "makeinfo finished; the Info files are in $(BUILDDIR)/texinfo." + +.PHONY: gettext +gettext: + $(SPHINXBUILD) -b gettext $(I18NSPHINXOPTS) $(BUILDDIR)/locale + @echo + @echo "Build finished. The message catalogs are in $(BUILDDIR)/locale." + +.PHONY: changes +changes: + $(SPHINXBUILD) -b changes $(ALLSPHINXOPTS) $(BUILDDIR)/changes + @echo + @echo "The overview file is in $(BUILDDIR)/changes." + +.PHONY: linkcheck +linkcheck: + $(SPHINXBUILD) -b linkcheck $(ALLSPHINXOPTS) $(BUILDDIR)/linkcheck + @echo + @echo "Link check complete; look for any errors in the above output " \ + "or in $(BUILDDIR)/linkcheck/output.txt." + +.PHONY: doctest +doctest: + $(SPHINXBUILD) -b doctest $(ALLSPHINXOPTS) $(BUILDDIR)/doctest + @echo "Testing of doctests in the sources finished, look at the " \ + "results in $(BUILDDIR)/doctest/output.txt." + +.PHONY: coverage +coverage: + $(SPHINXBUILD) -b coverage $(ALLSPHINXOPTS) $(BUILDDIR)/coverage + @echo "Testing of coverage in the sources finished, look at the " \ + "results in $(BUILDDIR)/coverage/python.txt." + +.PHONY: xml +xml: + $(SPHINXBUILD) -b xml $(ALLSPHINXOPTS) $(BUILDDIR)/xml + @echo + @echo "Build finished. The XML files are in $(BUILDDIR)/xml." + +.PHONY: pseudoxml +pseudoxml: + $(SPHINXBUILD) -b pseudoxml $(ALLSPHINXOPTS) $(BUILDDIR)/pseudoxml + @echo + @echo "Build finished. The pseudo-XML files are in $(BUILDDIR)/pseudoxml." diff --git a/models/tts/NeMo/docs/source/_static/css/custom.css b/models/tts/NeMo/docs/source/_static/css/custom.css new file mode 100644 index 0000000000..f168711c29 --- /dev/null +++ b/models/tts/NeMo/docs/source/_static/css/custom.css @@ -0,0 +1,199 @@ + +/* Width of template */ + +.wy-nav-content { + max-width: 1200px !important; +} + + + +/* Standard Text Formatting */ + +h1 +{ + color: #76b900; + text-align: center; + background-color: #333333; +} + +h2 +{ + color: #ffffff; + background-color: #76b900; + Padding: 5px; +} + +h3 +{ + padding-top: 0px; + border-top: solid 3px #76b900; + border-bottom: solid 3px #76b900; +} + +p +{ + margin-bottom: 24px; +} +/* Link Colors */ +a +{ + color: #76b900; +} + +a:visited +{ + color: #218219; +} + + +/* ----------------------------------------------TABLES--------------------------------------- */ + +/* Table head Color */ +thead td +{ + background-color: #333333 !important; +} + +.row-odd p +{ + /*padding-bottom: 0px;*/ + /*margin-bottom: 0px;*/ +} +/* even rows*/ + +.row-even tr +{ + background-color: #e5f1e6 !important; +} + +/* odd rows*/ + + +.wy-table-responsive table tr { + background-color: #ffffff !important; +} + + + +.wy-table-responsive table td { + white-space: normal; +} + + +/* Removes bottom margin in tables*/ + +.rst-content .line-block { + margin-bottom: 0px; +} + +.wy-table-responsive { + overflow: visible !important; +} + +/* reduces the size of text in multiline table columns. */ + +.rst-content table.docutils td +{ + font-size: 80%; +} + +.rst-content dl:not(.docutils) dt +{ + + background-color: inherit; + color: #000000; + border-top: solid 0px #000000; + +} + +.rst-content dl:not(.docutils) dt:before { + color: #333333; +} + +.rst-content .line-block { + margin-bottom: 0px; +} + +.wy-side-nav-search, .wy-nav-top + { + background-color: #000000; + padding: 0; + } + +.wy-side-nav-search img + { + padding: 0px; + padding: 0px 0px; + margin-bottom: 0; + } + +.wy-side-nav-search input[type=text] + { + border-radius: 0px; + } + + +.wy-menu-vertical p.caption + { + color: #76b900; + } + + +.wy-side-nav-search>a img.logo, .wy-side-nav-search .wy-dropdown>a img.logo + { + margin: 0px 0px 0px 0px; + } + +.wy-nav-content + { + margin: 0; + min-height: 100%; + height: 100%; + background: #ffffff; + } + + /* List (numbered, bulleted) padding Fix */ + + +.wy-plain-list-decimal li +{ + margin-top: -6px; + margin-bottom: -6px; +} + +.rst-content .section ol.loweralpha +{ + margin-top: -6px; + margin-bottom: 12px; +} + +.wy-plain-list-disc, .rst-content .toctree-wrapper ul, article ul +{ + margin-top: 0px !important; + margin-bottom: 12px; +} + + /* Alert Boxes */ + /* Background color of Alert Box Title */ + +.rst-content .section ul +{ + margin-top: -12px; + margin-bottom: 16px; +} + +.wy-alert.wy-alert-info .wy-alert-title, .rst-content .note .wy-alert-title, .rst-content .wy-alert-info.attention .wy-alert-title, .rst-content .wy-alert-info.caution .wy-alert-title, .rst-content .wy-alert-info.danger .wy-alert-title, .rst-content .wy-alert-info.error .wy-alert-title, .rst-content .wy-alert-info.hint .wy-alert-title, .rst-content .wy-alert-info.important .wy-alert-title, .rst-content .wy-alert-info.tip .wy-alert-title, .rst-content .wy-alert-info.warning .wy-alert-title, .rst-content .seealso .wy-alert-title, .rst-content .wy-alert-info.admonition-todo .wy-alert-title, .rst-content .wy-alert-info.admonition .wy-alert-title, .wy-alert.wy-alert-info .rst-content .admonition-title, .rst-content .wy-alert.wy-alert-info .admonition-title, .rst-content .note .admonition-title, .rst-content .wy-alert-info.attention .admonition-title, .rst-content .wy-alert-info.caution .admonition-title, .rst-content .wy-alert-info.danger .admonition-title, .rst-content .wy-alert-info.error .admonition-title, .rst-content .wy-alert-info.hint .admonition-title, .rst-content .wy-alert-info.important .admonition-title, .rst-content .wy-alert-info.tip .admonition-title, .rst-content .wy-alert-info.warning .admonition-title, .rst-content .seealso .admonition-title, .rst-content .wy-alert-info.admonition-todo .admonition-title, .rst-content .wy-alert-info.admonition .admonition-title + { + background: #76b900; + } + + /* Background and Font Color of Alert Box Main Body*/ +.wy-alert.wy-alert-info, .rst-content .note, .rst-content .wy-alert-info.attention, .rst-content .wy-alert-info.caution, .rst-content .wy-alert-info.danger, .rst-content .wy-alert-info.error, .rst-content .wy-alert-info.hint, .rst-content .wy-alert-info.important, .rst-content .wy-alert-info.tip, .rst-content .wy-alert-info.warning, .rst-content .seealso, .rst-content .wy-alert-info.admonition-todo, .rst-content .wy-alert-info.admonition { + background: #333333; + color: #999999; + } + +.section +{ + margin-top: 50px; +} diff --git a/models/tts/NeMo/docs/source/_static/js/pk_scripts.js b/models/tts/NeMo/docs/source/_static/js/pk_scripts.js new file mode 100644 index 0000000000..23c74982d3 --- /dev/null +++ b/models/tts/NeMo/docs/source/_static/js/pk_scripts.js @@ -0,0 +1,19 @@ +document.addEventListener("DOMContentLoaded", function () { + var params = window.location.search.substring(1).split("&").reduce(function (params, param) { + if (!param) { + return params; + } + + var values = param.split("="); + var name = values[0]; + var value = values[1]; + params[name] = value; + return params; + }, {}); + + var form = document.getElementById("feedback-form"); + for (var name in params) { + var input = form.querySelector("[name=" + name + "]"); + input.value = params[name]; + } +}); \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/api.rst b/models/tts/NeMo/docs/source/asr/api.rst new file mode 100644 index 0000000000..174e0530d9 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/api.rst @@ -0,0 +1,203 @@ +NeMo ASR collection API +======================= + + +Model Classes +------------- + +.. autoclass:: nemo.collections.asr.models.EncDecCTCModel + :show-inheritance: + :members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + + +.. autoclass:: nemo.collections.asr.models.EncDecCTCModelBPE + :show-inheritance: + :members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + + +.. autoclass:: nemo.collections.asr.models.EncDecRNNTModel + :show-inheritance: + :members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + + +.. autoclass:: nemo.collections.asr.models.EncDecRNNTBPEModel + :show-inheritance: + :members: transcribe, change_vocabulary, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + + +.. autoclass:: nemo.collections.asr.models.EncDecClassificationModel + :show-inheritance: + :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + + +.. autoclass:: nemo.collections.asr.models.EncDecSpeakerLabelModel + :show-inheritance: + :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + + +Modules +------- + +.. autoclass:: nemo.collections.asr.modules.ConvASREncoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.modules.ConvASRDecoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.modules.ConvASRDecoderClassification + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.modules.SpeakerDecoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.modules.ConformerEncoder + :show-inheritance: + :members: + + +Parts +----- + +.. autoclass:: nemo.collections.asr.parts.submodules.jasper.JasperBlock + :show-inheritance: + :members: + + +Mixins +------ + +.. autoclass:: nemo.collections.asr.parts.mixins.mixins.ASRBPEMixin + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.mixins.mixins.ASRModuleMixin + :show-inheritance: + :members: + +Datasets +-------- + +Character Encoding Datasets +~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: nemo.collections.asr.data.audio_to_text.AudioToCharDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.data.audio_to_text.TarredAudioToCharDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset + :show-inheritance: + :members: + +Subword Encoding Datasets +~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. autoclass:: nemo.collections.asr.data.audio_to_text.AudioToBPEDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.data.audio_to_text.TarredAudioToBPEDataset + :show-inheritance: + :members: + +Audio Preprocessors +------------------- + +.. autoclass:: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.modules.AudioToMFCCPreprocessor + :show-inheritance: + :members: + +Audio Augmentors +---------------- + +.. autoclass:: nemo.collections.asr.modules.SpectrogramAugmentation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.SpeedPerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.TimeStretchPerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.GainPerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.ImpulsePerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.ShiftPerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.NoisePerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.WhiteNoisePerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.RirAndNoisePerturbation + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.preprocessing.perturb.TranscodePerturbation + :show-inheritance: + :members: + +Miscellaneous Classes +--------------------- + +RNNT Decoding +~~~~~~~~~~~~~ + +.. autoclass:: nemo.collections.asr.metrics.rnnt_wer.RNNTDecoding + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.metrics.rnnt_wer_bpe.RNNTBPEDecoding + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.submodules.rnnt_greedy_decoding.GreedyRNNTInfer + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.submodules.rnnt_greedy_decoding.GreedyBatchedRNNTInfer + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.parts.submodules.rnnt_beam_decoding.BeamRNNTInfer + :show-inheritance: + :members: + +Hypotheses +~~~~~~~~~~ + +.. autoclass:: nemo.collections.asr.parts.utils.rnnt_utils.Hypothesis + :show-inheritance: + :no-members: + +.. autoclass:: nemo.collections.asr.parts.utils.rnnt_utils.NBestHypotheses + :show-inheritance: + :no-members: diff --git a/models/tts/NeMo/docs/source/asr/asr_all.bib b/models/tts/NeMo/docs/source/asr/asr_all.bib new file mode 100644 index 0000000000..7008355c79 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/asr_all.bib @@ -0,0 +1,980 @@ +@article{matchboxnet, + title={{MatchboxNet}: 1D Time-Channel Separable Convolutional Neural Network Architecture for Speech Commands Recognition}, + author={Majumdar, Somshubra and Ginsburg, Boris}, + journal={Proc. Interspeech 2020}, + year={2020} +} + +@article{marblenet, + title={MarbleNet: Deep 1D Time-Channel Separable Convolutional Neural Network for Voice Activity Detection}, + author={Jia, Fei and Majumdar, Somshubra and Ginsburg, Boris}, + journal={arXiv preprint arXiv:2010.13886}, + year={2020} +} + +@inproceedings{panayotov2015librispeech, + title={Librispeech: an ASR corpus based on public domain audio books}, + author={Panayotov, Vassil and Chen, Guoguo and Povey, Daniel and Khudanpur, Sanjeev}, + booktitle={Acoustics, Speech and Signal Processing (ICASSP), 2015 IEEE International Conference on}, + pages={5206--5210}, + year={2015}, + organization={IEEE} +} + +@article{luong17, + author = {Minh{-}Thang Luong and Eugene Brevdo and Rui Zhao}, + title = {Neural Machine Translation (seq2seq) Tutorial}, + journal = {https://github.com/tensorflow/nmt}, + year = {2017}, +} + +@INPROCEEDINGS{LaurentSeqWiseBN, +author={C. {Laurent} and G. {Pereyra} and P. {Brakel} and Y. {Zhang} and Y. {Bengio}}, +booktitle={2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)}, +title={Batch normalized recurrent neural networks}, +year={2016}, +volume={}, +number={}, +pages={2657-2661}, +keywords={feedforward neural nets;learning (artificial intelligence);recurrent neural nets;speech recognition;batch normalized recurrent neural networks;RNN;sequential data;long-term dependency learning;convergence rate improvement;intermediate representation normalization;feedforward neural networks;speech recognition task;language modeling;training criterion;Training;Recurrent neural networks;Convergence;Speech recognition;Computer architecture;Speech;batch normalization;RNN;LSTM;optimization}, +doi={10.1109/ICASSP.2016.7472159}, +ISSN={2379-190X}, +month={March},} + +@article{graves2005, + author = {Alex Graves and Jurgen Schmidhuber}, + title = {Framewise phoneme classification with bidirectional LSTM and other neural network architectures}, + journal = {Neural Networks, vol. 18}, + pages={602–-610}, + year = {2005}, +} + +@inproceedings{graves2006, + title={Connectionist temporal classification: labelling unsegmented sequence data with recurrent neural networks}, + author={Graves, Alex and Fern{\'a}ndez, Santiago and Gomez, Faustino and Schmidhuber, J{\"u}rgen}, + booktitle={Proceedings of the 23rd international conference on Machine learning}, + pages={369--376}, + year={2006}, + organization={ACM} +} + +@article{li2019jasper, + title={Jasper: An End-to-End Convolutional Neural Acoustic Model}, + author={Li, Jason and Lavrukhin, Vitaly and Ginsburg, Boris and Leary, Ryan and Kuchaiev, Oleksii and Cohen, Jonathan M and Nguyen, Huyen and Gadde, Ravi Teja}, + journal={arXiv preprint arXiv:1904.03288}, + year={2019} +} + +@misc{ardila2019common, + title={Common Voice: A Massively-Multilingual Speech Corpus}, + author={Rosana Ardila and Megan Branson and Kelly Davis and Michael Henretty and Michael Kohler and Josh Meyer and Reuben Morais and Lindsay Saunders and Francis M. Tyers and Gregor Weber}, + year={2019}, + eprint={1912.06670}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} + +@article{graves2012, + title={Sequence Transduction with Recurrent Neural Networks}, + author={Graves, Alex}, + journal={arXiv preprint arXiv:1211.3711}, + year={2012} +} + + +@article{graves2013, + title={Generating sequences with recurrent neural networks}, + author={Graves, Alex}, + journal={arXiv preprint arXiv:1308.0850}, + year={2013} +} + +@article{sergeev2018horovod, + title={Horovod: fast and easy distributed deep learning in TensorFlow}, + author={Sergeev, Alexander and Del Balso, Mike}, + journal={arXiv preprint arXiv:1802.05799}, + year={2018} +} + +@misc{NVVolta, + title = {NVIDIA TESLA V100 GPU ARCHITECTURE}, + howpublished = {\url{http://images.nvidia.com/content/volta-architecture/pdf/volta-architecture-whitepaper.pdf}}, + note = {Accessed: 2018-10-09} +} + +@article{NVTuring, + title = {NVIDIA TURING GPU ARCHITECTURE}, + howpublished = {\url{https://www.nvidia.com/content/dam/en-zz/Solutions/design-visualization/technologies/turing-architecture/NVIDIA-Turing-Architecture-Whitepaper.pdf}}, + author = {NVIDIA}, + year = {2018}, + note = {Accessed: 2018-10-09} +} + +@misc{Rygaard2015, + title = {Using Synthesized Speech to Improve Speech Recognition for Low-Resource Languages}, + author = {Luise Valentin Rygaard}, + howpublished = {\url{https://parasol.tamu.edu/dreu2015/Rygaard/report.pdf}}, + year = {2015}, +} + +@misc{OpenSeq2Seq, + title = {OpenSeq2Seq: extensible toolkit for distributed and mixed precision training of sequence-to-sequence models}, + author = {Kuchaiev, Oleksii and Ginsburg, Boris and Gitman, Igor and Lavrukhin,Vitaly and Case, Carl and Micikevicius, Paulius}, + howpublished = {\url{https://arxiv.org/abs/1805.10387}}, + year = {2018}, +} + +@misc{MPGuide, + title = {Training with Mixed Precision}, + howpublished = {\url{http://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/}}, + note = {Accessed: 2018-04-06}, +} + +@misc{Mozilla, + title = {Mozilla: A Journey to less than 10\% Word Error Rate}, + howpublished = {\url{https://hacks.mozilla.org/2017/11/a-journey-to-10-word-error-rate/}}, + note = {Accessed: 2018-04-06}, +} + +@article{Waibel1989, + title={A time-delay neural network architecture for isolated word recognition}, + author={Waibel, Alexander, and Hanazawa, Toshiyki and Hinton,Geoffrey and Shirano, Kiyohiro and Lang, Kevin }, + journal={IEEE Trans. on Acoustics, Speech and Signal Processing}, + year={1989} +} + +@article{Lang1990, + title={A time-delay neural network architecture for isolated word recognition}, + author={Lang, Kevin and Waibel, Alexander, and Hinton,Geoffrey }, + journal={Neural Networks}, + year={1990} +} + +@book{Bengio1996, + Author = {Bengio, Y.}, + Publisher = {International Thomson Computer Press}, + Title = {Neural Networks for Speech and Sequence Recognition}, + Year = {1996} +} + +@article{Bengio1992, + title={Global optimization of a neural network-hidden Markov model hybrid}, + author={Bengio, Y., and De Mori, R., and Flammia, G., and Kompe, R. }, + journal={IEEE Transactions on Neural Networks, 3(2), 252–259}, + year={1992} +} + +@article{Bourlard1994, + title={Connectionist speech recognition: a hybrid approach}, + author={Bourlard, H. A. and Morgan, N.}, + journal={volume 247 Springer }, + year={1994} +} + +@article{srivastava14a, + author = {Nitish Srivastava, and Geoffrey Hinton, and Alex Krizhevsky, and Ilya Sutskever, and Ruslan Salakhutdinov}, + title = {Dropout: A Simple Way to Prevent Neural Networks from Overfitting}, + journal = {Journal of Machine Learning Research}, + year = {2014}, + volume = {15}, + pages = {1929-1958}, + url = {http://jmlr.org/papers/v15/srivastava14a.html} +} + + +@article{Hinton2012, + title={Deep Neural Networks for Acoustic Modeling in Speech Recognition}, + author={ Hinton,Geoffrey and Deng, Li and Yu, Dong and Dahl,George and Mohamed,Abdel-rahman and Jaitly, Navdeep and Senior, Andrew and Vanhoucke, Vincent and Nguyen, Patrick and Kingsbury, Brian and Sainath, Tara}, + journal={IEEE Signal Processing Magazine}, + year={2012} +} + +@article{Graves2014, + title={Towards End-to-End Speech Recognition with Recurrent Neural Networks}, + author={Graves, Alex and Jaitly, Navdeep}, + journal={International Conference on Machine Learning}, + year={2014} +} + +@article{Chorowski2014, + title={End-to-end Continuous Speech Recognition using Attention-based Recurrent NN: First Results}, + author={ Chorowski, Jan, and Bahdanau, Dzmitry , and Cho, Kyunghyun , and Bengio, Yoshua }, + journal={Neural Information Processing Systems: Workshop Deep Learning and Representation Learning Workshop }, + year={2014} +} + +@article{Sak2014, + title={Long short-term memory recurrent neural network architectures for large scale acoustic modeling}, + author={Sak, Hasim and Senior, Andrew and Beaufays, Francoise }, + journal={Interspeech 2014}, + year={2014} +} + +@article{Ko2015, + title={Audio Augmentation for Speech Recognition}, + author={Tom, Ko and Vijayaditya, Peddinti and Daniel, Povey + and Sanjeev, Khudanpur }, + journal={Interspeech 2015}, + year={2015} +} + +@article{Tjandra2017, + title={Listening while Speaking: Speech Chain by Deep Learning}, + author={Andros, Tjandra and Sakriani, Sakti and Satoshi, Nakamura }, + journal={ASRU 2017}, + year={2017} +} + +@article{Tjandra2018, + title={Machine Speech Chain with One-shot Speaker Adaptation}, + author={Andros, Tjandra and Sakriani, Sakti and Satoshi, Nakamura }, + journal={Interspeech 2018}, + year={2018} +} + +@article{bahdanau2014neural, + title={Neural machine translation by jointly learning to align and translate}, + author={Bahdanau, Dzmitry and Cho, Kyunghyun and Bengio, Yoshua}, + journal={arXiv preprint arXiv:1409.0473}, + year={2014} +} + +@article{cho2014learning, + title={Learning phrase representations using RNN encoder-decoder for statistical machine translation}, + author={Cho, Kyunghyun and Van Merri{\"e}nboer, Bart and Gulcehre, Caglar and Bahdanau, Dzmitry and Bougares, Fethi and Schwenk, Holger and Bengio, Yoshua}, + journal={arXiv preprint arXiv:1406.1078}, + year={2014} +} + +@article{rush2015neural, + title={A neural attention model for abstractive sentence summarization}, + author={Rush, Alexander M and Chopra, Sumit and Weston, Jason}, + journal={arXiv preprint arXiv:1509.00685}, + year={2015} +} + +@article{micikevicius2017mixed, + title={Mixed precision training}, + author={Micikevicius, Paulius and Narang, Sharan and Alben, Jonah and Diamos, Gregory and Elsen, Erich and Garcia, David and Ginsburg, Boris and Houston, Michael and Kuchaev, Oleksii and Venkatesh, Ganesh and others}, + journal={arXiv preprint arXiv:1710.03740}, + year={2017} +} + +@ARTICLE{Britz:2017, + author = {{Britz}, Denny and {Goldie}, Anna and {Luong}, Thang and {Le}, Quoc}, + title = "{Massive Exploration of Neural Machine Translation Architectures}", + journal = {ArXiv e-prints arXiv:1703.03906}, + archivePrefix = "arXiv", + eprinttype = {arxiv}, + eprint = {1703.03906}, + primaryClass = "cs.CL", + keywords = {Computer Science - Computation and Language}, + year = 2017, + month = mar, +} + +@inproceedings{vaswani2017attention, + title={Attention is all you need}, + author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, {\L}ukasz and Polosukhin, Illia}, + booktitle={Advances in Neural Information Processing Systems}, + pages={6000--6010}, + year={2017} +} + +@inproceedings{abadi2016tensorflow, + title={TensorFlow: A System for Large-Scale Machine Learning.}, + author={Abadi, Mart{\'\i}n and Barham, Paul and Chen, Jianmin and Chen, Zhifeng and Davis, Andy and Dean, Jeffrey and Devin, Matthieu and Ghemawat, Sanjay and Irving, Geoffrey and Isard, Michael and others}, + booktitle={OSDI}, + volume={16}, + pages={265--283}, + year={2016} +} + +@article{tensor2tensor, + author = {Ashish Vaswani and Samy Bengio and Eugene Brevdo and Francois Chollet and Aidan N. Gomez and Stephan Gouws and Llion Jones and \L{}ukasz Kaiser and Nal Kalchbrenner and Niki Parmar and Ryan Sepassi and + Noam Shazeer and Jakob Uszkoreit}, + title = {Tensor2Tensor for Neural Machine Translation}, + journal = {CoRR}, + volume = {abs/1803.07416}, + year = {2018}, + url = {http://arxiv.org/abs/1803.07416}, +} + +@article{gehring2017convs2s, + author = {Gehring, Jonas, and Auli, Michael and Grangier, David and Yarats, Denis and Dauphin, Yann N}, + title = "{Convolutional Sequence to Sequence Learning}", + journal = {ArXiv e-prints arXiv:1705.03122}, + archivePrefix = "arXiv", + eprinttype = {arxiv}, + eprint = {1705.03122}, + primaryClass = "cs.CL", + keywords = {Computer Science - Computation and Language}, + year = 2017, + month = May, +} + +@inproceedings{chan2015, + title={Listen, attend and spell}, + author={Chan, William and Jaitly, Navdeep and Le, Quoc V and Vinyals, Oriol}, + booktitle={Acoustics, Speech and Signal Processing (ICASSP), 2016 IEEE International Conference on}, + pages={5206--5210}, + year={2016}, + organization={IEEE} +} + +@inproceedings{xu2015show, + title={Show, attend and tell: Neural image caption generation with visual attention}, + author={Xu, Kelvin and Ba, Jimmy and Kiros, Ryan and Cho, Kyunghyun and Courville, Aaron and Salakhudinov, Ruslan and Zemel, Rich and Bengio, Yoshua}, + booktitle={International Conference on Machine Learning}, + pages={2048--2057}, + year={2015} +} + +@incollection{Sutskever2014, + title = {Sequence to Sequence Learning with Neural Networks}, + author = {Sutskever, Ilya and Vinyals, Oriol and Le, Quoc V}, + booktitle = {Advances in Neural Information Processing Systems 27}, + editor = {Z. Ghahramani and M. Welling and C. Cortes and N. D. Lawrence and K. Q. Weinberger}, + pages = {3104--3112}, + year = {2014}, + publisher = {Curran Associates, Inc.}, + url = {http://papers.nips.cc/paper/5346-sequence-to-sequence-learning-with-neural-networks.pdf} +} + +@article{DeepSpeech2014, + title = {Deep Speech: Scaling up end-to-end speech recognition}, + author = {Awni Y. Hannun and Carl Case and Jared Casper and Bryan Catanzaro and Greg Diamos and Erich Elsen and Ryan Prenger and Sanjeev Satheesh and Shubho Sengupta and Adam Coates and Andrew Y. Ng}, + journal = {CoRR}, + volume = {abs/1412.5567}, + year = {2014}, + url = {http://arxiv.org/abs/1412.5567}, + archivePrefix = {arXiv}, + eprint = {1412.5567}, + timestamp = {Mon, 13 Aug 2018 16:48:07 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/HannunCCCDEPSSCN14}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@inproceedings{DeepSpeech2, + author = {Amodei, Dario and Ananthanarayanan, Sundaram and Anubhai, Rishita and Bai, Jingliang and Battenberg, Eric and Case, Carl and Casper, Jared and Catanzaro, Bryan and Cheng, Qiang and Chen, Guoliang and Chen, Jie and Chen, Jingdong and Chen, Zhijie and Chrzanowski, Mike and Coates, Adam and Diamos, Greg and Ding, Ke and Du, Niandong and Elsen, Erich and Engel, Jesse and Fang, Weiwei and Fan, Linxi and Fougner, Christopher and Gao, Liang and Gong, Caixia and Hannun, Awni and Han, Tony and Johannes, Lappi Vaino and Jiang, Bing and Ju, Cai and Jun, Billy and LeGresley, Patrick and Lin, Libby and Liu, Junjie and Liu, Yang and Li, Weigao and Li, Xiangang and Ma, Dongpeng and Narang, Sharan and Ng, Andrew and Ozair, Sherjil and Peng, Yiping and Prenger, Ryan and Qian, Sheng and Quan, Zongfeng and Raiman, Jonathan and Rao, Vinay and Satheesh, Sanjeev and Seetapun, David and Sengupta, Shubho and Srinet, Kavya and Sriram, Anuroop and Tang, Haiyuan and Tang, Liliang and Wang, Chong and Wang, Jidong and Wang, Kaifu and Wang, Yi and Wang, Zhijian and Wang, Zhiqian and Wu, Shuang and Wei, Likai and Xiao, Bo and Xie, Wen and Xie, Yan and Yogatama, Dani and Yuan, Bin and Zhan, Jun and Zhu, Zhenyao}, + title = {Deep Speech 2: End-to-end Speech Recognition in English and Mandarin}, + booktitle = {Proceedings of the 33rd International Conference on International Conference on Machine Learning - Volume 48}, + series = {ICML'16}, + year = {2016}, + location = {New York, NY, USA}, + pages = {173--182}, + numpages = {10}, + url = {http://dl.acm.org/citation.cfm?id=3045390.3045410}, + acmid = {3045410}, + publisher = {JMLR.org}, +} + +@inproceedings{prabhavalkar2017comparison, + title={A comparison of sequence-to-sequence models for speech recognition}, + author={Prabhavalkar, Rohit and Rao, Kanishka and Sainath, Tara N and Li, Bo and Johnson, Leif and Jaitly, Navdeep}, + booktitle={Proc. Interspeech}, + pages={939--943}, + year={2017} +} + +@article{chiu2017state, + title={State-of-the-art speech recognition with sequence-to-sequence models}, + author={Chiu, Chung-Cheng and Sainath, Tara N and Wu, Yonghui and Prabhavalkar, Rohit and Nguyen, Patrick and Chen, Zhifeng and Kannan, Anjuli and Weiss, Ron J and Rao, Kanishka and Gonina, Katya and others}, + journal={arXiv preprint arXiv:1712.01769}, + year={2017} +} + +@misc{NVMixed, + title = {{NVIDA's Mixed-Precision Training - TensorFlow example}}, + howpublished = {\url{https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/#example_tensorflow}}, + author={NVIDIA}, + note = {Accessed: 2018-10-09}, + year={2018} +} + +@article{gehring2017, + title={Convolutional sequence to sequence learning}, + author={Gehring, Jonas and Auli, Michael and Grangier, David and Yarats, Denis and Dauphin, Yann N}, + journal={arXiv preprint arXiv:1705.03122}, + year={2017} +} + +@article{collobert2016, + title={Wav2letter: an end-to-end convnet-based speech recognition system}, + author={Collobert, Ronan and Puhrsch, Christian and Synnaeve, Gabriel}, + journal={arXiv preprint arXiv:1609.03193}, + year={2016} +} + +@inproceedings{Zhang2016, +author={Ying Zhang and Mohammad Pezeshki and Philémon Brakel and Saizheng Zhang and César Laurent and Yoshua Bengio and Aaron Courville}, +title={Towards End-to-End Speech Recognition with Deep Convolutional Neural Networks}, +year=2016, +booktitle={Interspeech 2016}, +doi={10.21437/Interspeech.2016-1446}, +url={http://dx.doi.org/10.21437/Interspeech.2016-1446}, +pages={410--414} +} + +@inproceedings{Zhang2017, + title={Very deep convolutional networks for end-to-end speech recognition}, + author={Zhang, Yu, and Chan, William, and Jaitly, Navdeep}, + booktitle={Acoustics, Speech and Signal Processing (ICASSP), 2017 IEEE International Conference on}, + year={2017}, + organization={IEEE} +} + + +@article{Wang2017, + title={Tacotron: Towards End-to-End Speech Synthesis}, + author={ Wang, Yuxuan, and Skerry-Ryan, RJ, and Stanton, Daisy and Wu, Yonghui and Weiss, Ron, and Jaitly, Navdeep and Yang, Zongheng and Xiao, Ying and Chen,Zhifeng and Bengio, Samy and Le, Quoc and Agiomyrgiannakis, Yannis and Clark,Rob and Saurous, Rif A.}, + journal={arXiv preprint arXiv:1703.10135}, + year={2017} +} + +@inproceedings{shen2018natural, + title={Natural TTS synthesis by conditioning wavenet on mel spectrogram predictions}, + author={Shen, Jonathan and Pang, Ruoming and Weiss, Ron J and Schuster, Mike and Jaitly, Navdeep and Yang, Zongheng and Chen, Zhifeng and Zhang, Yu and Wang, Yuxuan and Skerrv-Ryan, Rj and others}, + booktitle={2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)}, + pages={4779--4783}, + year={2018}, + organization={IEEE} +} + +@article{griffin1984signal, + title={Signal estimation from modified short-time Fourier transform}, + author={Griffin, Daniel and Lim, Jae}, + journal={IEEE Transactions on Acoustics, Speech, and Signal Processing}, + volume={32}, + number={2}, + pages={236--243}, + year={1984}, + publisher={IEEE} +} + +@misc{ito2017lj, + title={The LJ speech dataset}, + author={Ito, Keith and others}, + year={2017} +} + +@misc{mailabs, + title = {{The M-AILABS Speech Dataset}}, + howpublished = {\url{http://www.m-ailabs.bayern/en/the-mailabs-speech-dataset/}}, + author={M-AILABS}, + note = {Accessed: 2018-10-09}, + year={2018} +} + +@article{merity2016pointer, + title={Pointer sentinel mixture models}, + author={Merity, Stephen and Xiong, Caiming and Bradbury, James and Socher, Richard}, + journal={arXiv preprint arXiv:1609.07843}, + year={2016} +} + +@inproceedings{socher2013recursive, + title={Recursive deep models for semantic compositionality over a sentiment treebank}, + author={Socher, Richard and Perelygin, Alex and Wu, Jean and Chuang, Jason and Manning, Christopher D and Ng, Andrew and Potts, Christopher}, + booktitle={Proceedings of the 2013 conference on empirical methods in natural language processing}, + pages={1631--1642}, + year={2013} +} + +@InProceedings{maas-EtAl:2011:ACL-HLT2011, + author = {Maas, Andrew L. and Daly, Raymond E. and Pham, Peter T. and Huang, Dan and Ng, Andrew Y. and Potts, Christopher}, + title = {Learning Word Vectors for Sentiment Analysis}, + booktitle = {Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies}, + month = {June}, + year = {2011}, + address = {Portland, Oregon, USA}, + publisher = {Association for Computational Linguistics}, + pages = {142--150}, + url = {http://www.aclweb.org/anthology/P11-1015} +} + +@inproceedings{Povey2018SemiOrthogonalLM, + title={Semi-Orthogonal Low-Rank Matrix Factorization for Deep Neural Networks}, + author={Daniel Povey and Gaofeng Cheng and Yiming Wang and Ke Li and Hainan Xu and Mahsa Yarmohammadi and Sanjeev Khudanpur}, + booktitle={Interspeech}, + year={2018} +} + +@article{CAPIO2017, + author = {Kyu J. Han and Akshay Chandrashekaran and Jungsuk Kim and Ian R. Lane}, + title = {The {CAPIO} 2017 Conversational Speech Recognition System}, + journal = {CoRR}, + volume = {abs/1801.00059}, + year = {2018}, + url = {http://arxiv.org/abs/1801.00059}, + archivePrefix = {arXiv}, + eprint = {1801.00059}, + timestamp = {Mon, 13 Aug 2018 16:49:10 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/abs-1801-00059}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{WaveNet, + author = {A{\"{a}}ron van den Oord and Sander Dieleman and Heiga Zen and Karen Simonyan and Oriol Vinyals and Alex Graves and Nal Kalchbrenner and Andrew W. Senior and Koray Kavukcuoglu}, + title = {WaveNet: {A} Generative Model for Raw Audio}, + journal = {CoRR}, + volume = {abs/1609.03499}, + year = {2016}, + url = {http://arxiv.org/abs/1609.03499}, + archivePrefix = {arXiv}, + eprint = {1609.03499}, + timestamp = {Mon, 13 Aug 2018 16:49:15 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/OordDZSVGKSK16}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{FacebookGERENGBackTranslation, + author = {Rico Sennrich and Barry Haddow and Alexandra Birch}, + title = {Improving Neural Machine Translation Models with Monolingual Data}, + journal = {CoRR}, + volume = {abs/1511.06709}, + year = {2015}, + url = {http://arxiv.org/abs/1511.06709}, + archivePrefix = {arXiv}, + eprint = {1511.06709}, + timestamp = {Mon, 13 Aug 2018 16:47:05 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/SennrichHB15a}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{GlobalStyleTokens, + author = {Yuxuan Wang and Daisy Stanton and Yu Zhang and R. J. Skerry{-}Ryan and Eric Battenberg and Joel Shor and Ying Xiao and Fei Ren and Ye Jia and Rif A. Saurous}, + title = {Style Tokens: Unsupervised Style Modeling, Control and Transfer in End-to-End Speech Synthesis}, + journal = {CoRR}, + volume = {abs/1803.09017}, + year = {2018}, + url = {http://arxiv.org/abs/1803.09017}, + archivePrefix = {arXiv}, + eprint = {1803.09017}, + timestamp = {Mon, 13 Aug 2018 16:46:53 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/abs-1803-09017}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{IoffeS15BatchNorm, + author = {Sergey Ioffe and Christian Szegedy}, + title = {Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift}, + journal = {CoRR}, + volume = {abs/1502.03167}, + year = {2015}, + url = {http://arxiv.org/abs/1502.03167}, + archivePrefix = {arXiv}, + eprint = {1502.03167}, + timestamp = {Mon, 13 Aug 2018 16:47:06 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/IoffeS15}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{kingma, + author = {Diederik P. Kingma and + Jimmy Ba}, + title = {Adam: {A} Method for Stochastic Optimization}, + journal = {CoRR}, + volume = {abs/1412.6980}, + year = {2014}, + url = {http://arxiv.org/abs/1412.6980}, + archivePrefix = {arXiv}, + eprint = {1412.6980}, + timestamp = {Mon, 13 Aug 2018 01:00:00 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/KingmaB14}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@incollection{Salimans2016WeightNorm, + title = {Weight Normalization: A Simple Reparameterization to Accelerate Training of Deep Neural Networks}, + author = {Salimans, Tim and Kingma, Durk P}, + booktitle = {Advances in Neural Information Processing Systems 29}, + editor = {D. D. Lee and M. Sugiyama and U. V. Luxburg and I. Guyon and R. Garnett}, + pages = {901--909}, + year = {2016}, + publisher = {Curran Associates, Inc.}, + url = {http://papers.nips.cc/paper/6114-weight-normalization-a-simple-reparameterization-to-accelerate-training-of-deep-neural-networks.pdf} +} + +@article{wu2016google, + title={Google's neural machine translation system: Bridging the gap between human and machine translation}, + author={Wu, Yonghui and Schuster, Mike and Chen, Zhifeng and Le, Quoc V and Norouzi, Mohammad and Macherey, Zolfgang and Krikun, Maxim and Cao, Yuan and Gao, Qin and Macherey, Klaus and others}, + journal={arXiv preprint arXiv:1609.08144}, + year={2016} +} + +@inproceedings{opennmt, + author = {Guillaume Klein and Yoon Kim and Yuntian Deng and Jean Senellart and Alexander M. Rush}, + title = {OpenNMT: Open-Source Toolkit for Neural Machine Translation}, + booktitle = {Proc. ACL}, + year = {2017}, + url = {https://doi.org/10.18653/v1/P17-4012}, + doi = {10.18653/v1/P17-4012} +} + +@article{paszke2017automatic, + title={Automatic differentiation in PyTorch}, + author={Paszke, Adam and Gross, Sam and Chintala, Soumith and Chanan, Gregory and Yang, Edward and DeVito, Zachary and Lin, Zeming and Desmaison, Alban and Antiga, Luca and Lerer, Adam}, + year={2017} +} + +@article{yu2014introduction, + title={An introduction to computational networks and the computational network toolkit}, + author={Yu, Dong and Eversole, Adam and Seltzer, Mike and Yao, Kaisheng and Huang, Zhiheng and Guenter, Brian and Kuchaiev, Oleksii and Zhang, Yu and Seide, Frank and Wang, Huaming and others}, + journal={Microsoft Technical Report MSR-TR-2014--112}, + year={2014} +} + +@article{nvidia2017v100, + title={V100 GPU architecture. The world’s most advanced data center GPU. Version WP-08608-001\_v1. 1}, + author={NVIDIA, Tesla}, + journal={NVIDIA. Aug}, + pages={108}, + year={2017} +} + +@article{post2018call, + title={A call for clarity in reporting bleu scores}, + author={Post, Matt}, + journal={arXiv preprint arXiv:1804.08771}, + year={2018} +} + +@article{Ba2016LayerNorm, + author = {Jimmy Lei Ba and Jamie Ryan Kiros and Geoffrey E Hinton}, + title = {Layer normalization}, + journal = {CoRR}, + volume = {abs/1607.06450}, + year = {2016}, + url = {http://arxiv.org/abs/1607.06450}, + archivePrefix = {arXiv}, +} + +@inproceedings{Dauphin2017GLU, + author = {Dauphin, Yann N. and Fan, Angela and Auli, Michael and Grangier, David}, + title = {Language Modeling with Gated Convolutional Networks}, + booktitle = {Proceedings of the 34th International Conference on Machine Learning - Volume 70}, + series = {ICML'17}, + year = {2017}, + location = {Sydney, NSW, Australia}, + pages = {933--941}, + numpages = {9}, + url = {http://dl.acm.org/citation.cfm?id=3305381.3305478}, + acmid = {3305478}, + publisher = {JMLR.org}, +} + +@incollection{Oord2016PixelCNN, +title = {Conditional Image Generation with PixelCNN Decoders}, +author = {van den Oord, Aaron and Kalchbrenner, Nal and Espeholt, Lasse and kavukcuoglu, koray and Vinyals, Oriol and Graves, Alex}, +booktitle = {Advances in Neural Information Processing Systems 29}, +editor = {D. D. Lee and M. Sugiyama and U. V. Luxburg and I. Guyon and R. Garnett}, +pages = {4790--4798}, +year = {2016}, +publisher = {Curran Associates, Inc.}, +url = {http://papers.nips.cc/paper/6527-conditional-image-generation-with-pixelcnn-decoders.pdf} +} + +@article{he2015, + title={Deep residual learning for image recognition}, + author={K. He, and X. Zhang, and S. Ren, and J. Sun}, + journal={arXiv preprint arXiv:1512.03385}, + year={2015} +} + +@article{huang2016, + title={Densely Connected Convolutional Networks}, + author={Gao Huang, and Zhuang Liu, and Laurens van der Maaten, and Kilian Q. Weinberger}, + journal={arXiv preprint arXiv:1608.06993}, + year={2016} +} + +@inproceedings{heafield2011kenlm, + title={KenLM: Faster and smaller language model queries}, + author={Heafield, Kenneth}, + booktitle={Proceedings of the sixth workshop on statistical machine translation}, + pages={187--197}, + year={2011}, + organization={Association for Computational Linguistics} +} + +@article{dai2018transformer, + title={Transformer-XL: Language Modeling with Longer-Term Dependency}, + author={Dai, Zihang and Yang, Zhilin and Yang, Yiming and Cohen, William W and Carbonell, Jaime and Le, Quoc V and Salakhutdinov, Ruslan}, + year={2018}, + journal = {CoRR}, + volume = {abs/1901.02860}, + url = {http://arxiv.org/abs/1901.02860}, + archivePrefix = {arXiv}, + eprint = {1901.02860}, + timestamp = {Fri, 01 Feb 2019 13:39:59 +0100}, + biburl = {https://dblp.org/rec/bib/journals/corr/abs-1901-02860}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@inproceedings{Saon+2016, +author={George Saon and Tom Sercu and Steven Rennie and Hong-Kwang J. Kuo}, +title={The IBM 2016 English Conversational Telephone Speech Recognition System}, +year=2016, +booktitle={Interspeech 2016}, +doi={10.21437/Interspeech.2016-1460}, +url={http://dx.doi.org/10.21437/Interspeech.2016-1460}, +pages={7--11} +} + +@INPROCEEDINGS{Sercu-2016, +author={T. {Sercu} and C. {Puhrsch} and B. {Kingsbury} and Y. {LeCun}}, +booktitle={2016 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)}, +title={Very deep multilingual convolutional neural networks for LVCSR}, +year={2016}, +volume={}, +number={}, +pages={4955-4959}, +keywords={natural language processing;neural nets;speech recognition;very deep multilingual convolutional neural networks;LVCSR;CNN;large vocabulary continuous speech recognition systems;word error rate;Training;Context;Hidden Markov models;Neural networks;Computer architecture;Kernel;Training data;Convolutional Networks;Multilingual;Acoustic Modeling;Speech Recognition;Neural Networks}, +doi={10.1109/ICASSP.2016.7472620}, +ISSN={2379-190X}, +month={March},} + + +@inproceedings{Sercu+2016, +author={Tom Sercu and Vaibhava Goel}, +title={Advances in Very Deep Convolutional Neural Networks for LVCSR}, +year=2016, +booktitle={Interspeech 2016}, +doi={10.21437/Interspeech.2016-1033}, +url={http://dx.doi.org/10.21437/Interspeech.2016-1033}, +pages={3429--3433} +} + +@INPROCEEDINGS{Xiong-2018, +author={W. {Xiong} and L. {Wu} and F. {Alleva} and J. {Droppo} and X. {Huang} and A. {Stolcke}}, +booktitle={2018 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)}, +title={The Microsoft 2017 Conversational Speech Recognition System}, +year={2018}, +volume={}, +number={}, +pages={5934-5938}, +keywords={convolution;feedforward neural nets;natural language processing;speaker recognition;speech processing;language model rescoring step;senone level;switchboard domains;character-based LSTM language models;NIST 2000 switchboard test set;frame level;word-level voting;acoustic model posteriors;dialog session aware LSTM language models;CNN-BLSTM acoustic model;Microsoft 2017 conversational speech recognition system;Acoustics;Error analysis;Training;Speech recognition;Switches;Computational modeling;Context modeling;Conversational speech recognition;CNN;LACE;BLSTM;LSTM-LM;system combination;human parity}, +doi={10.1109/ICASSP.2018.8461870}, +ISSN={2379-190X}, +month={April},} + +@inproceedings{zeyer2018improved, + author={Albert Zeyer and Kazuki Irie and Ralf Schlüter and Hermann Ney}, + title={Improved Training of End-to-end Attention Models for Speech Recognition}, + year=2018, + booktitle={Proc. Interspeech 2018}, + pages={7--11}, + doi={10.21437/Interspeech.2018-1616}, + url={http://dx.doi.org/10.21437/Interspeech.2018-1616} +} + +@article{Wav2LetterV2, + author = {Vitaliy Liptchinsky and + Gabriel Synnaeve and + Ronan Collobert}, + title = {Letter-Based Speech Recognition with Gated ConvNets}, + journal = {CoRR}, + volume = {abs/1712.09444}, + year = {2017}, + url = {http://arxiv.org/abs/1712.09444}, + archivePrefix = {arXiv}, + eprint = {1712.09444}, + timestamp = {Mon, 13 Aug 2018 16:46:33 +0200}, + biburl = {https://dblp.org/rec/bib/journals/corr/abs-1712-09444}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{zeghidour2018, + author = {Neil Zeghidour and + Qiantong Xu and + Vitaliy Liptchinsky and + Nicolas Usunier and + Gabriel Synnaeve and + Ronan Collobert}, + title = {Fully Convolutional Speech Recognition}, + journal = {CoRR}, + volume = {abs/1812.06864}, + year = {2018}, + url = {http://arxiv.org/abs/1812.06864}, + archivePrefix = {arXiv}, + eprint = {1812.06864}, + timestamp = {Tue, 01 Jan 2019 15:01:25 +0100}, + biburl = {https://dblp.org/rec/bib/journals/corr/abs-1812-06864}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@inproceedings{Hadian2018, + author={Hossein Hadian and Hossein Sameti and Daniel Povey and Sanjeev Khudanpur}, + title={End-to-end Speech Recognition Using Lattice-free MMI}, + year=2018, + booktitle={Proc. Interspeech 2018}, + pages={12--16}, + doi={10.21437/Interspeech.2018-1423}, + url={http://dx.doi.org/10.21437/Interspeech.2018-1423} +} + +@inproceedings{Tang2018, + author={Jian Tang and Yan Song and Lirong Dai and Ian McLoughlin}, + title={Acoustic Modeling with Densely Connected Residual Network for Multichannel Speech Recognition}, + year=2018, + booktitle={Proc. Interspeech 2018}, + pages={1783--1787}, + doi={10.21437/Interspeech.2018-1089}, + url={http://dx.doi.org/10.21437/Interspeech.2018-1089} +} + +@article{Kurata2017LanguageMW, + title={Language modeling with highway LSTM}, + author={Gakuto Kurata and Bhuvana Ramabhadran and George Saon and Abhinav Sethy}, + journal={2017 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)}, + year={2017}, + pages={244-251} +} + +@inproceedings{Saon2017, + author={George Saon and Gakuto Kurata and Tom Sercu and Kartik Audhkhasi and Samuel Thomas and Dimitrios Dimitriadis and Xiaodong Cui and Bhuvana Ramabhadran and Michael Picheny and Lynn-Li Lim and Bergul Roomi and Phil Hall}, + title={English Conversational Telephone Speech Recognition by Humans and Machines}, + year=2017, + booktitle={Proc. Interspeech 2017}, + pages={132--136}, + doi={10.21437/Interspeech.2017-405}, + url={http://dx.doi.org/10.21437/Interspeech.2017-405} +} + +@inproceedings{Povey+2016, +author={Daniel Povey and Vijayaditya Peddinti and Daniel Galvez and Pegah Ghahremani and Vimal Manohar and Xingyu Na and Yiming Wang and Sanjeev Khudanpur}, +title={Purely Sequence-Trained Neural Networks for ASR Based on Lattice-Free MMI}, +year=2016, +booktitle={Interspeech 2016}, +doi={10.21437/Interspeech.2016-595}, +url={http://dx.doi.org/10.21437/Interspeech.2016-595}, +pages={2751--2755} +} + +@article{Yang2018, + author = {Xuerui Yang and + Jiwei Li and + Xi Zhou}, + title = {A novel pyramidal-FSMN architecture with lattice-free {MMI} for speech + recognition}, + journal = {CoRR}, + volume = {abs/1810.11352}, + year = {2018}, + url = {http://arxiv.org/abs/1810.11352}, + archivePrefix = {arXiv}, + eprint = {1810.11352}, + timestamp = {Wed, 31 Oct 2018 14:24:29 +0100}, + biburl = {https://dblp.org/rec/bib/journals/corr/abs-1810-11352}, + bibsource = {dblp computer science bibliography, https://dblp.org} +} + +@article{liptchinsky2017based, + title={Letter-Based Speech Recognition with Gated ConvNets}, + author={Liptchinsky, Vitaliy and Synnaeve, Gabriel and Collobert, Ronan}, + journal={arXiv preprint arXiv:1712.09444}, + year={2017} +} + +@inproceedings{Weng2018, + author={Chao Weng and Jia Cui and Guangsen Wang and Jun Wang and Chengzhu Yu and Dan Su and Dong Yu}, + title={Improving Attention Based Sequence-to-Sequence Models for End-to-End English Conversational Speech Recognition}, + year=2018, + booktitle={Proc. Interspeech 2018}, + pages={761--765}, + doi={10.21437/Interspeech.2018-1030}, + url={http://dx.doi.org/10.21437/Interspeech.2018-1030} +} + +@INPROCEEDINGS{Battenberg2017, +author={E. {Battenberg} and J. {Chen} and R. {Child} and A. {Coates} and Y. G. Y. {Li} and H. {Liu} and S. {Satheesh} and A. {Sriram} and Z. {Zhu}}, +booktitle={2017 IEEE Automatic Speech Recognition and Understanding Workshop (ASRU)}, +title={Exploring neural transducers for end-to-end speech recognition}, +year={2017}, +volume={}, +number={}, +pages={206-213}, +keywords={recurrent neural nets;speech recognition;Hub500 benchmark;CTC models;speech recognition pipeline;RNN-Transducer models;language model;Seq2Seq models;end-to-end speech recognition;neural transducers;Decoding;Hidden Markov models;Transducers;Task analysis;Speech;Mathematical model;Neural networks}, +doi={10.1109/ASRU.2017.8268937}, +ISSN={}, +month={Dec}, +} + +@inproceedings{ +loshchilov2018, +title={Decoupled Weight Decay Regularization}, +author={Ilya Loshchilov and Frank Hutter}, +booktitle={International Conference on Learning Representations}, +year={2019}, +url={https://openreview.net/forum?id=Bkg6RiCqY7}, +} + +@article{zhang2017ndadam, + author = {Zijun Zhang and Lin Ma and Zongpeng Li and Chuan Wu}, + title = {Normalized Direction-preserving Adam}, + journal = {arXiv e-prints arXiv:1709.04546}, + year = {2017}, +} + +@article{park2019, + author = {{Park}, Daniel S. and {Chan}, William and {Zhang}, Yu and + {Chiu}, Chung-Cheng and {Zoph}, Barret and {Cubuk}, Ekin D. and + {Le}, Quoc V.}, + title = "{SpecAugment: A Simple Data Augmentation Method for Automatic Speech Recognition}", + journal = {arXiv e-prints}, + year = "2019", + eid = {arXiv:1904.08779}, + eprint = {1904.08779}, +} + +@article{novograd2019, + author = {{Ginsburg}, Boris and {Castonguay}, Patrice and {Hrinchuk}, Oleksii and + {Kuchaiev}, Oleksii and {Lavrukhin}, Vitaly and {Leary}, Ryan and + {Li}, Jason and {Nguyen}, Huyen and {Cohen}, Jonathan M.}, + title = "{Stochastic Gradient Methods with Layer-wise Adaptive Moments for Training of Deep Networks}", + journal = {arXiv e-prints}, + year = "2019", + eid = {arXiv:1905.11286}, + eprint = {1905.11286}, +} + +@article{kriman2019quartznet, + title={Quartznet: {Deep} automatic speech recognition with 1d time-channel separable convolutions}, + author={Kriman, Samuel and Beliaev, Stanislav and Ginsburg, Boris and Huang, Jocelyn and Kuchaiev, Oleksii and Lavrukhin, Vitaly and Leary, Ryan and Li, Jason and Zhang, Yang}, + journal={arXiv preprint arXiv:1910.10261}, + year={2019} +} + +@misc{itu1988g711, + title={{ITU-T} {G.711} - {Pulse} code modulation ({PCM}) of voice frequencies}, + author={ITU-T Geneva Switzerland}, + year={1988}, +} + +@article{han2020contextnet, + title={ContextNet: Improving convolutional neural networks for automatic speech recognition with global context}, + author={Han, Wei and Zhang, Zhengdong and Zhang, Yu and Yu, Jiahui and Chiu, Chung-Cheng and Qin, James and Gulati, Anmol and Pang, Ruoming and Wu, Yonghui}, + journal={arXiv:2005.03191}, + year={2020} +} + +@inproceedings{hu2018squeeze, + title={Squeeze-and-excitation networks}, + author={Hu, Jie and Shen, Li and Sun, Gang}, + booktitle={ICVPR}, + year={2018} +} + +@article{koluguri2020speakernet, + title={SpeakerNet: 1D Depth-wise Separable Convolutional Network for Text-Independent Speaker Recognition and Verification}, + author={Koluguri, Nithin Rao and Li, Jason and Lavrukhin, Vitaly and Ginsburg, Boris}, + journal={arXiv preprint arXiv:2010.12653}, + year={2020} +} + +@article{gulati2020conformer, + title={Conformer: Convolution-augmented transformer for speech recognition}, + author={Gulati, Anmol and Qin, James and Chiu, Chung-Cheng and Parmar, Niki and Zhang, Yu and Yu, Jiahui and Han, Wei and Wang, Shibo and Zhang, Zhengdong and Wu, Yonghui and others}, + journal={arXiv preprint arXiv:2005.08100}, + year={2020} +} \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/asr_language_modeling.rst b/models/tts/NeMo/docs/source/asr/asr_language_modeling.rst new file mode 100644 index 0000000000..2df9994c4a --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/asr_language_modeling.rst @@ -0,0 +1,292 @@ +##################### +ASR Language Modeling +##################### + +Language models have shown to help the accuracy of ASR models. NeMo supports the following two approaches to incorporate language models into the ASR models: + +* :ref:`ngram_modeling` +* :ref:`neural_rescoring` + +It is possible to use both approaches on the same ASR model. + + +.. _ngram_modeling: + +************************ +N-gram Language Modeling +************************ + +In this approach, an N-gram LM is trained on text data, then it is used in fusion with beam search decoding to find the +best candidates. The beam search decoders in NeMo support language models trained with KenLM library ( +`https://github.com/kpu/kenlm `__). +The beam search decoders and KenLM library are not installed by default in NeMo, and you need to install them to be +able to use beam search decoding and N-gram LM. +Please refer to `scripts/ngram_lm/install_beamsearch_decoders.sh` on how to install them. + +NeMo supports both character-based and BPE-based models for N-gram LMs. An N-gram LM can be used with beam search +decoders on top of the ASR models to produce more accurate candidates. The beam search decoder would incorporate +the scores produced by the N-gram LM into its score calculations as the following: + +.. code:: + + final_score = acoustic_score + beam_alpha*lm_score + beam_beta*seq_length + +where acoustic_score is the score predicted by the acoustic encoder and lm_score is the one estimated by the LM. +Parameter 'beam_alpha' specifies amount of importance to place on the N-gram language model, and 'beam_beta' is a +penalty term to consider the sequence length in the scores. Larger alpha means more importance on the LM and less +importance on the acoustic model. Negative values for beta will give penalty to longer sequences and make the decoder +to prefer shorter predictions, while positive values would result in longer candidates. + + +Train N-gram LM +=============== + +The script to train an N-gram language model with KenLM can be found at +`scripts/asr_language_modeling/ngram_lm/train_kenlm.py `__. + +This script would train an N-gram language model with KenLM library which can be used with the beam search decoders +on top of the ASR models. This script supports both character level and BPE level encodings and models which is +detected automatically from the type of the model. + + +You may train the N-gram model as the following: + +.. code:: + + python train_kenlm.py --nemo_model_file \ + --train_file \ + --kenlm_model_file \ + --ngram_length + +The train file specified by `--train_file` can be a text file or JSON manifest. If the file's extension is anything +other than `.json`, it assumes that data format is plain text. For plain text format, each line should contain one +sample. For JSON manifest file, the file need to contain json formatted samples per each line like this: + +.. code:: + + {"audio_filepath": "/data_path/file1.wav", "text": "The transcript of the audio file."} + +It just extracts the `text` field from each line to create the training text file. After the N-gram model is trained, +it is stored at the path specified by `--kenlm_model_file`. + +The following is the list of the arguments for the training script: + ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| **Argument** | **Type** | **Default** | **Description** | ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| nemo_model_file | str | Required | The path of the `.nemo` file of the ASR model. It is needed to extract the tokenizer. | ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| train_file | str | Required | Path to the training file, it can be a text file or JSON manifest. | ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| kenlm_model_file | str | Required | The path to store the KenLM binary model file. | ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| kenlm_bin_path | str | Required | The path to the bin folder of KenLM. It is a folder named `bin` under where KenLM is installed.| ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| ngram_length** | int | Required | Specifies order of N-gram LM. | ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ +| do_lower_case | bool | ``False`` | Whether to make the training text all lower case. | ++------------------+----------+-------------+------------------------------------------------------------------------------------------------+ + +** Note: Recommend to use 6 as the order of the N-gram model for BPE-based models. Higher orders may need the re-compilation of KenLM to support it. + +Evaluate by Beam Search Decoding and N-gram LM +============================================== + +NeMo's beam search decoders are capable of using the KenLM's N-gram models to find the best candidates. +The script to evaluate an ASR model with beam search decoding and N-gram models can be found at +`scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py `__. + +You may evaluate an ASR model as the following: + +.. code:: + + python eval_beamsearch_ngram.py --nemo_model_file \ + --input_manifest \ + --acoustic_batch_size \ + --beam_width \ + --beam_alpha \ + --beam_width \ + --preds_output_folder \ + --decoding_mode beam_search_ngram + +It can evaluate a model in the three following modes by setting the argument `--decoding_mode`: + +* greedy: Just greedy decoding is done, and no beam search decoding is performed. +* beamsearch: The beam search decoding is done but without using the N-gram language model, final results would be equivalent to setting the weight of LM (beam_beta) to zero. +* beamsearch_ngram: The beam search decoding is done with N-gram LM. + +The `beamsearch` mode would evaluate by beam search decoding without any language model. +It would report the performances in terms of Word Error Rate (WER) and Character Error Rate (CER). Moreover, +the WER/CER of the model when the best candidate is selected among the candidates is also reported as the best WER/CER. +It can be an indicator of how good the predicted candidates are. + +The script would initially load the ASR model and predict the outputs of the model's encoder as log probabilities. +This part would be computed in batches on a device selected by `--device`, which can be CPU (`--device=cpu`) or a +single GPU (`--device=cuda:0`). The batch size of this part can get specified by `--acoustic_batch_size`. You may use +the largest batch size feasible to speed up the step of calculating the log probabilities. You may also use `--use_amp` +to speed up the calculation of log probabilities and make it possible to use larger sizes for `--acoustic_batch_size`. +Currently multi-GPU is not supported for calculating the log probabilities, but using `--probs_cache_file` can help. +It stores the log probabilities produced from the model's encoder into a pickle file so that next time the first step +can get skipped. + +The following is the list of the arguments for the evaluation script: + ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| **Argument** |**Type**| **Default** | **Description** | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| nemo_model_file | str | Required | The path of the `.nemo` file of the ASR model to extract the tokenizer. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| input_manifest | str | Required | Path to the training file, it can be a text file or JSON manifest. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| kenlm_model_file | str | Required | The path to store the KenLM binary model file. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| preds_output_folder | str | None | The path to an optional folder to store the predictions. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| probs_cache_file | str | None | The cache file for storing the outputs of the model. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| acoustic_batch_size | int | 16 | The batch size to calculate log probabilities. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| use_amp | bool | ``False`` | Whether to use AMP if available to calculate log probabilities. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| device | str | cuda | The device to load the model onto to calculate log probabilities. | +| | | | It can `cpu`, `cuda`, `cuda:0`, `cuda:1`, ... | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| decoding_mode | str | beamsearch_ngram | The decoding scheme to be used for evaluation. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beam_width | float | Required | The width or list of the widths of the beam search decoding. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beam_alpha | float | Required | The alpha parameter or list of the alphas for the beam search decoding. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beam_beta | float | Required | The beta parameter or list of the betas for the beam search decoding. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beam_batch_size | int | 128 | The batch size to be used for beam search decoding. | +| | | | Larger batch size can be a little faster, but uses larger memory. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ + +Width of the beam search (`--beam_width`) specifies the number of top candidates/predictions the beam search decoder +would search for. Larger beams result in more accurate but slower predictions. + +There is also a tutorial to learn more about evaluating the ASR models with N-gram LM here: +`Offline ASR Inference with Beam Search and External Language Model Rescoring `_ + +Hyperparameter Grid Search +-------------------------- + +Beam search decoding with N-gram LM has three main hyperparameters: `beam_width`, `beam_alpha`, and `beam_beta`. +The accuracy of the model is dependent to the values of these parameters, specially beam_alpha and beam_beta. +You may specify a single or list of values for each of these parameters to perform grid search. It would perform the +beam search decoding on all the combinations of the these three hyperparameters. +For instance, the following set of parameters would results in 2*1*2=4 beam search decodings: + +.. code:: + + python eval_beamsearch_ngram.py ... \ + --beam_width 64 128 \ + --beam_alpha 1.0 \ + --beam_beta 1.0 0.5 + + +.. _neural_rescoring: + +**************** +Neural Rescoring +**************** + +In this approach a neural network is used which can gives scores to a candidate. A candidate is the text transcript predicted by the decoder of the ASR model. +The top K candidates produced by the beam search decoding (beam width of K) are given to a neural language model to rank them. +Ranking can be done by a language model which gives a score to each candidate. +This score is usually combined with the scores from the beam search decoding to produce the final scores and rankings. + +Train Neural Rescorer +===================== + +An example script to train such a language model with Transformer can be found at `examples/nlp/language_modeling/transformer_lm.py `__. +It trains a TransformerLMModel which can be used as a neural rescorer for an ASR system. + + +Evaluation +========== + +Given a trained TransformerLMModel `.nemo` file, the script available at +`scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py `__ +can be used to re-score beams obtained with ASR model. You need the `.tsv` file containing the candidates produced +by the acoustic model and the beam search decoding to use this script. The candidates can be the result of just the beam +search decoding or the result of fusion with an N-gram LM. You may generate this file by specifying `--preds_output_folder' for +`scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py `__. + +The neural rescorer would rescore the beams/candidates by using two parameters of `rescorer_alpha` and `rescorer_beta` as the following: + +.. code:: + + final_score = beam_search_score + rescorer_alpha*neural_rescorer_score + rescorer_beta*seq_length + +Parameter `rescorer_alpha` specifies amount of importance to place on the neural rescorer model, and `rescorer_beta` is +a penalty term to consider the sequence length in the scores. They have similar effects like the parameters +`beam_alpha` and `beam_beta` of beam search decoder and N-gram LM. + +You may follow the following steps to evaluate a neural LM: + +#. Obtain `.tsv` file with beams and their corresponding scores. Scores can be from a regular beam search decoder or + in fusion with an N-gram LM scores. For a given beam size `beam_size` and a number of examples + for evaluation `num_eval_examples`, it should contain (`num_eval_examples` x `beam_size`) lines of + form `beam_candidate_text \t score`. This file can be generated by `scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py `__ + +#. Rescore the candidates by `scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py `__. + +.. code:: + python eval_neural_rescorer.py + --lm_model=[path to .nemo file of the LM] + --beams_file=[path to beams .tsv file] + --beam_size=[size of the beams] + --eval_manifest=[path to eval manifest .json file] + --batch_size=[batch size used for inference on the LM model] + --alpha=[the value for the parameter rescorer_alpha] + --beta=[the value for the parameter rescorer_beta] + --scores_output_file=[the optional path to store the rescored candidates] + +The candidates along with their new scores would be stored at the file specified by `--scores_output_file`. + +The following is the list of the arguments for the evaluation script: + ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| **Argument** |**Type**| **Default** | **Description** | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| lm_model | str | Required | The path of the '.nemo' file of the ASR model to extract the tokenizer. | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| eval_manifest | str | Required | Path to the evaluation manifest file (.json manifest file) | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beams_file | str | Required | path to beams file (.tsv) containing the candidates and their scores | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beam_size | int | Required | The width of the beams (number of candidates) generated by the decoder | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| alpha | float | None | The value for parameter rescorer_alpha | +| | | | Not passing value would enable linear search for rescorer_alpha | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| beta | float | None | The value for parameter rescorer_beta | +| | | | Not passing value would enable linear search for rescorer_beta | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| batch_size | int | 16 | The batch size used to calculate the scores | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| scores_output_file | str | None | The optional file to store the rescored beams | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| use_amp | bool | ``False`` | Whether to use AMP if available calculate the scores | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ +| device | str | cuda | The device to load LM model onto to calculate the scores | +| | | | It can be 'cpu', 'cuda', 'cuda:0', 'cuda:1', ... | ++---------------------+--------+------------------+-------------------------------------------------------------------------+ + + +Hyperparameter Linear Search +---------------------------- + +This script also supports linear search for parameters `rescorer_alpha` and `rescorer_beta`. If any of the two is not +provided, a linear search is performed to find the best value for that parameter. When linear search is used, initially +`rescorer_beta` is set to zero and the best value for `rescorer_alpha` is found, then `rescorer_alpha` is fixed with +that value and another linear search is done to find the best value for `rescorer_beta`. +If any of the of these two parameters is already specified, then search for that one is skipped. After each search for a +parameter, the plot of WER% for different values of the parameter is also shown. + +It is recommended to first use the linear search for both parameters on a validation set by not providing any values for `--alpha` and `--beta`. +Then check the WER curves and decide on the best values for each parameter. Finally, evaluate the best values on the test set. diff --git a/models/tts/NeMo/docs/source/asr/configs.rst b/models/tts/NeMo/docs/source/asr/configs.rst new file mode 100644 index 0000000000..4c93043fcd --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/configs.rst @@ -0,0 +1,528 @@ +NeMo ASR Configuration Files +============================ + +This section describes the NeMo configuration file setup that is specific to models in the ASR collection. For general information +about how to set up and run experiments that is common to all NeMo models (e.g. Experiment Manager and PyTorch Lightning trainer +parameters), see the :doc:`../core/core` section. + +The model section of the NeMo ASR configuration files generally requires information about the dataset(s) being used, the preprocessor +for audio files, parameters for any augmentation being performed, as well as the model architecture specification. The sections on +this page cover each of these in more detail. + +Example configuration files for all of the NeMo ASR scripts can be found in the +`config directory of the examples `_. + + +Dataset Configuration +--------------------- + +Training, validation, and test parameters are specified using the ``train_ds``, ``validation_ds``, and +``test_ds`` sections in the configuration file, respectively. Depending on the task, there may be arguments specifying the sample rate +of the audio files, the vocabulary of the dataset (for character prediction), whether or not to shuffle the dataset, and so on. You may +also decide to leave fields such as the ``manifest_filepath`` blank, to be specified via the command-line at runtime. + +Any initialization parameter that is accepted for the Dataset class used in the experiment can be set in the config file. +Refer to the `Datasets <./api.html#Datasets>`__ section of the API for a list of Datasets and their respective parameters. + +An example ASR train and validation configuration should look similar to the following: + +.. code-block:: yaml + + # Specified at the beginning of the config file + labels: &labels [" ", "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", "'"] + + model: + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels # Uses the labels above + batch_size: 32 + trim_silence: True + max_duration: 16.7 + shuffle: True + is_tarred: False # If set to true, uses the tarred version of the Dataset + tarred_audio_filepaths: null # Not used if is_tarred is false + tarred_shard_strategy: "scatter" # Not used if is_tarred is false + num_workers: 8 + pin_memory: true + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels # Uses the labels above + batch_size: 32 + shuffle: False # No need to shuffle the validation data + num_workers: 8 + pin_memory: true + + +Preprocessor Configuration +-------------------------- + +If you are loading audio files for your experiment, you will likely want to use a preprocessor to convert from the +raw audio signal to features (e.g. mel-spectrogram or MFCC). The ``preprocessor`` section of the config specifies the audio +preprocessor to be used via the ``_target_`` field, as well as any initialization parameters for that preprocessor. + +An example of specifying a preprocessor is as follows: + +.. code-block:: yaml + + model: + ... + preprocessor: + # _target_ is the audio preprocessor module you want to use + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + ... + # Other parameters for the preprocessor + +Refer to the `Audio Preprocessors <./api.html#Audio Preprocessors>`__ API section for the preprocessor options, expected arguments, +and defaults. + +Augmentation Configurations +--------------------------- + +There are a few on-the-fly spectrogram augmentation options for NeMo ASR, which can be specified by the +configuration file using a ``spec_augment`` section. + +For example, there are options for `Cutout `_ and +`SpecAugment `_ available via the ``SpectrogramAugmentation`` module. + +The following example sets up both ``Cutout`` (via the ``rect_*`` parameters) and ``SpecAugment`` (via the ``freq_*`` +and ``time_*`` parameters). + +.. code-block:: yaml + + model: + ... + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + # Cutout parameters + rect_masks: 5 # Number of rectangles to cut from any given spectrogram + rect_freq: 50 # Max cut of size 50 along the frequency dimension + rect_time: 120 # Max cut of size 120 along the time dimension + # SpecAugment parameters + freq_masks: 2 # Cut two frequency bands + freq_width: 15 # ... of width 15 at maximum + time_masks: 5 # Cut out 10 time bands + time_width: 25 # ... of width 25 at maximum + +You can use any combination of ``Cutout``, frequency/time ``SpecAugment``, or neither of them. + +With NeMo ASR, you can also add augmentation pipelines that can be used to simulate various kinds of noise +added to audio in the channel. Augmentors in a pipeline are applied on the audio data read in the data layer. Online +augmentors can be specified in the config file using an ``augmentor`` section in ``train_ds``. The following example +adds an augmentation pipeline that first adds white noise to an audio sample with a probability of 0.5 and at a level +randomly picked between -50 dB and -10 dB and then passes the resultant samples through a room impulse response randomly +picked from the manifest file provided for ``impulse`` augmentation in the config file. + +.. code-block:: yaml + + model: + ... + train_ds: + ... + augmentor: + white_noise: + prob: 0.5 + min_level: -50 + max_level: -10 + impulse: + prob: 0.3 + manifest_path: /path/to/impulse_manifest.json + +Refer to the `Audio Augmentors <./api.html#Audio Augmentors>`__ API section for more details. + +Tokenizer Configurations +------------------------ + +Some models utilize sub-word encoding via an external tokenizer instead of explicitly defining their vocabulary. + +For such models, a ``tokenizer`` section is added to the model config. ASR models currently support two types of +custom tokenizers: + +- Google Sentencepiece tokenizers (tokenizer type of ``bpe`` in the config) +- HuggingFace WordPiece tokenizers (tokenizer type of ``wpe`` in the config) + +In order to build custom tokenizers, refer to the ``ASR_with_Subword_Tokenization`` notebook available in the +ASR tutorials directory. + +The following example sets up a ``SentencePiece Tokenizer`` at a path specified by the user: + +.. code-block:: yaml + + model: + ... + tokenizer: + dir: "" + type: "bpe" # can be "bpe" or "wpe" + +For models which utilize sub-word tokenization, we share the decoder module (``ConvASRDecoder``) with character tokenization models. +All parameters are shared, but for models which utilize sub-word encoding, there are minor differences when setting up the config. For +such models, the tokenizer is utilized to fill in the missing information when the model is constructed automatically. + +For example, a decoder config corresponding to a sub-word tokenization model should look similar to the following: + +.. code-block:: yaml + + model: + ... + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_final + num_classes: -1 # filled with vocabulary size from tokenizer at runtime + vocabulary: [] # filled with vocabulary from tokenizer at runtime + + +Model Architecture Configurations +--------------------------------- + +Each configuration file should describe the model architecture being used for the experiment. Models in the NeMo ASR collection need +an ``encoder`` section and a ``decoder`` section, with the ``_target_`` field specifying the module to use for each. + +Here is the list of the parameters in the model section which are shared among most of the ASR models: + ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ +| **Parameter** | **Datatype** | **Description** | **Supported Values** | ++=========================+==================+===============================================================================================================+=================================+ +| :code:`log_prediction` | bool | Whether a random sample should be printed in the output at each step, along with its predicted transcript. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ +| :code:`ctc_reduction` | string | Specifies the reduction type of CTC loss. Defaults to ``mean_batch`` which would take the average over the | :code:`none`, | +| | | batch after taking the average over the length of each sample. | :code:`mean_batch` | +| | | | :code:`mean`, :code:`sum` | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ + +The following sections go into more detail about the specific configurations of each model architecture. + +For more information about the ASR models, refer to the :doc:`Models <./models>` section. + +Jasper and QuartzNet +~~~~~~~~~~~~~~~~~~~~ + +The `Jasper <./models.html#Jasper>`__ and `QuartzNet <./models.html#QuartzNet>`__ models are very similar, and as such the components in their +configs are very similar as well. + +Both architectures use the ``ConvASREncoder`` for the ``encoder``, with parameters detailed in the table below. The encoder parameters +include details about the Jasper/QuartzNet ``[BxR]`` encoder architecture, including how many blocks to use (``B``), how many times +to repeat each sub-block (``R``), and the convolution parameters for each block. + +The number of blocks ``B`` is determined by the number of list elements under ``jasper`` minus the one prologue and two epilogue blocks. +The number of sub-blocks ``R`` is determined by setting the ``repeat`` parameter. + +To use QuartzNet (which uses more compact time-channel separable convolutions) instead of Jasper, add :code:`separable: true` to all +but the last block in the architecture. + +Change the parameter name ``jasper``. + ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+ +| **Parameter** | **Datatype** | **Description** | **Supported Values** | ++=========================+==================+===============================================================================================================+=====================================+ +| :code:`feat_in` | int | The number of input features. Should be equal to :code:`features` in the preprocessor parameters. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+ +| :code:`activation` | string | Which activation function to use in the encoder. | :code:`hardtanh`, :code:`relu`, | +| | | | :code:`selu`, :code:`swish` | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+ +| :code:`conv_mask` | bool | Whether to use masked convolutions in the encoder. Defaults to ``true``. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+ +| :code:`jasper` | | A list of blocks that specifies your encoder architecture. Each entry in this list represents one block in | | +| | | the architecture and contains the parameters for that block, including convolution parameters, dropout, and | | +| | | the number of times the block is repeated. Refer to the `Jasper `_ and | | +| | | `QuartzNet `_ papers for details about specific model configurations. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+-------------------------------------+ + +A QuartzNet 15x5 (fifteen blocks, each sub-block repeated five times) encoder configuration should look similar to the following example: + +.. code-block:: yaml + + # Specified at the beginning of the file for convenience + n_mels: &n_mels 64 # Used for both the preprocessor and encoder as number of input features + repeat: &repeat 5 # R=5 + dropout: &dropout 0.0 + separable: &separable true # Set to true for QN. Set to false for Jasper. + + model: + ... + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels # Should match "features" in the preprocessor. + activation: relu + conv_mask: true + + jasper: # This field name should be "jasper" for both types of models. + + # Prologue block + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: 1 # Prologue block is not repeated. + residual: false + separable: *separable + stride: [2] + + # Block 1 + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + ... # Entries for blocks 2~14 + + # Block 15 + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + # Two epilogue blocks + - dilation: [2] + dropout: *dropout + filters: 512 + kernel: [87] + repeat: 1 # Epilogue blocks are not repeated + residual: false + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: &enc_filters 1024 + kernel: [1] + repeat: 1 # Epilogue blocks are not repeated + residual: false + stride: [1] + +Both Jasper and QuartzNet use the ``ConvASRDecoder`` as the decoder. The decoder parameters are detailed in the following table. + ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ +| **Parameter** | **Datatype** | **Description** | **Supported Values** | ++=========================+==================+===============================================================================================================+=================================+ +| :code:`feat_in` | int | The number of input features to the decoder. Should be equal to the number of filters in the last block of | | +| | | the encoder. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ +| :code:`vocabulary` | list | A list of the valid output characters for your model. For example, for an English dataset, this could be a | | +| | | list of all lowercase letters, space, and apostrophe. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ +| :code:`num_classes` | int | Number of output classes, i.e. the length of :code:`vocabulary`. | | ++-------------------------+------------------+---------------------------------------------------------------------------------------------------------------+---------------------------------+ + +For example, a decoder config corresponding to the encoder above should look similar to the following: + +.. code-block:: yaml + + model: + ... + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_filters + vocabulary: *labels + num_classes: 28 # Length of the vocabulary list + +Citrinet +~~~~~~~~ + +The `Citrinet <./models.html#Citrinet>`__ and `QuartzNet <./models.html#QuartzNet>`__ models are very similar, and as such the +components in their configs are very similar as well. Citrinet utilizes Squeeze and Excitation, as well as sub-word tokenization, in +contrast to QuartzNet. Depending on the dataset, we utilize different tokenizers. For Librispeech, we utilize the HuggingFace WordPiece +tokenizer, and for all other datasets we utilize the Google Sentencepiece tokenizer - usually the ``unigram`` tokenizer type. + +Both architectures use the ``ConvASREncoder`` for the ``encoder``, with parameters detailed above. The encoder parameters include +details about the Citrinet-C encoder architecture, including how many filters are used per channel (``C``). The Citrinet-C +configuration is a shortform notation for Citrinet-21x5xC, such that ``B = 21`` and ``R = 5`` are the default and should generally +not be changed. + +To use Citrinet instead of QuartzNet, refer to the ``citrinet_512.yaml`` configuration found inside the ``examples/asr/conf/citrinet`` +directory. Citrinet is primarily comprised of the same :class:`~nemo.collections.asr.parts.submodules.jasper.JasperBlock` as ``Jasper`` or +``QuartzNet`. + +While the configs for Citrinet and QuartzNet are similar, we note the additional flags used for Citrinet below. Refer to the +``JasperBlock`` documentation for the meaning of these arguments. + ++---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+ +| **Parameter** | **Datatype** | **Description** | **Supported Values** | ++===========================+==================+===========================================================================================================+===================================+ +| :code:`se` | bool | Whether to apply squeeze-and-excitation mechanism or not. | :code:`true` or :code:`false` | ++---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+ +| :code:`se_context_size` | int | SE context size. -1 means global context. | :code:`-1` or :code:`+ve int` | ++---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+ +| :code:`stride_last` | bool | Stride on the final repeated block or all repeated blocks. | :code:`true` or :code:`false` | ++---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+ +| :code:`residual_mode` | str | Type of residual branch to construct. | :code:`"add"` or | +| | | Can be pointwise residual addition or pointwise strided residual attention | :code:`"stride_add"` | ++---------------------------+------------------+-----------------------------------------------------------------------------------------------------------+-----------------------------------+ + +A Citrinet-512 config should look similar to the following: + +.. code-block:: yaml + + model: + ... + # Specify some defaults across the entire model + model_defaults: + repeat: 5 + dropout: 0.1 + separable: true + se: true + se_context_size: -1 + ... + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels # Should match "features" in the preprocessor. + activation: relu + conv_mask: true + + jasper: # This field name should be "jasper" for the JasperBlock (which constructs Citrinet). + + # Prologue block + - filters: 512 + repeat: 1 + kernel: [5] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + # Block 1 + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [11] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + ... # Entries for blocks 2~21 + + # Block 22 + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [39] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + # Epilogue block + + - filters: &enc_final 640 + repeat: 1 + kernel: [41] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + +As mentioned above, Citrinet uses the ``ConvASRDecoder`` as the decoder layer similar to QuartzNet. Only the configuration must be +changed slightly as Citrinet utilizes sub-word tokenization. + +.. note:: + The following information is relevant to any of the above models that implements its encoder as an :class:`~nemo.collections.asr.modules.conv_asr.ConvASREncoder`, and utilizes the ``SqueezeExcite`` mechanism. + +The ``SqueezeExcite`` block within a :class:`~nemo.collections.asr.modules.conv_asr.ConvASREncoder` network can be modified to utilize a different context window after the model has been instantiated (even after the model has been trained) so as to evaluate the model with limited context. This can be achieved using the :meth:`~nemo.collections.asr.parts.mixins.mixins.ASRModuleMixin.change_conv_asr_se_context_window` + +.. code-block:: python + + # Here, model can be any model that has a `ConvASREncoder` as its encoder, and utilized `SqueezeExcite` blocks + # `context_window` : It is an integer representing the number of timeframes (each corresponding to some window stride). + # `update_config` : Bool flag which determines whether the config of the model should be updated to reflect the new context window. + + # Here, we specify that 128 timeframes of 0.01s stride should be the context window + # This is equivalent to 128 * 0.01s context window for `SqueezeExcite` + model.change_conv_asr_se_context_window(context_window=128, update_config=True) + +Conformer-CTC +~~~~~~~~~~~~~ + +The config files for Conformer-CTC model contain character-based encoding and sub-word encoding at +``/examples/asr/conf/conformer/conformer_ctc_char.yaml`` and ``/examples/asr/conf/conformer/conformer_ctc_bpe.yaml`` +respectively. Some components of the configs of `Conformer-CTC <./models.html#Conformer-CTC>`__ include the following datasets: + +- ``train_ds``, ``validation_ds``, and ``test_ds`` +- opimizer (``optim``) +- augmentation (``spec_augment``) +- ``decoder`` +- ``trainer`` +``exp_manager`` + +These datasets are similar to other ASR models like `QuartzNet <./models.html#QuartzNet>`__. There should be a tokenizer section where you can +specify the tokenizer if you want to use sub-word encoding instead of character-based encoding. + +The encoder section includes the details about the Conformer-CTC encoder architecture. You may find more information in the +config files and also :doc:`nemo.collections.asr.modules.ConformerEncoder<./api.html#nemo.collections.asr.modules.ConformerEncoder>`. + + +Fine-tuning Configurations +------------------------- + +All ASR scripts support easy fine-tuning by partially/fully loading the pretrained weights from a checkpoint into the currently instantiated model. Pre-trained weights can be provided in multiple ways - + +1) Providing a path to a NeMo model (via ``init_from_nemo_model``) +2) Providing a name of a pretrained NeMo model (which will be downloaded via the cloud) (via ``init_from_pretrained_model``) +3) Providing a path to a Pytorch Lightning checkpoint file (via ``init_from_ptl_ckpt``) + +Fine-tuning via a NeMo model +~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: sh + + python examples/asr/script_to_.py \ + --config-path= \ + --config-name=) \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath="" \ + trainer.gpus=-1 \ + trainer.max_epochs=50 \ + +init_from_nemo_model="" + + +Fine-tuning via a NeMo pretrained model name +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: sh + + python examples/asr/script_to_.py \ + --config-path= \ + --config-name=) \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath="" \ + trainer.gpus=-1 \ + trainer.max_epochs=50 \ + +init_from_pretrained_model="" + +Fine-tuning via a Pytorch Lightning checkpoint +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. code-block:: sh + + python examples/asr/script_to_.py \ + --config-path= \ + --config-name=) \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath="" \ + trainer.gpus=-1 \ + trainer.max_epochs=50 \ + +init_from_ptl_ckpt="" \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/data/asr_results.csv b/models/tts/NeMo/docs/source/asr/data/asr_results.csv new file mode 100644 index 0000000000..3b179b5c59 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/asr_results.csv @@ -0,0 +1,21 @@ +Model Name,Model Base Class,Model Card +QuartzNet15x5Base-En,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels" +stt_en_jasper10x5dr,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr" +stt_en_citrinet_256,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256" +stt_en_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512" +stt_en_citrinet_1024,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024" +stt_ca_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5" +stt_it_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5" +stt_fr_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_fr_quartznet15x5" +stt_es_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5" +stt_de_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5" +stt_pl_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5" +stt_ru_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5" +stt_zh_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5" +stt_zh_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512" +stt_en_conformer_ctc_small,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small" +stt_en_conformer_ctc_medium,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium" +stt_en_conformer_ctc_large,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large" +stt_en_conformer_ctc_small_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small_ls" +stt_en_conformer_ctc_medium_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium_ls" +stt_en_conformer_ctc_large_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls" diff --git a/models/tts/NeMo/docs/source/asr/data/asrlm_results.csv b/models/tts/NeMo/docs/source/asr/data/asrlm_results.csv new file mode 100644 index 0000000000..d9a395cb8b --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/asrlm_results.csv @@ -0,0 +1,2 @@ +Model Name,Model Base Class,Model Card +asrlm_en_transformer_large_ls,TransformerLMModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:asrlm_en_transformer_large_ls" diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_ca.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_ca.csv new file mode 100644 index 0000000000..ef5220fe6e --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_ca.csv @@ -0,0 +1,2 @@ +Model,Model Base Class,Model Card +stt_ca_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ca_quartznet15x5" diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_de.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_de.csv new file mode 100644 index 0000000000..f0f86121fe --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_de.csv @@ -0,0 +1,3 @@ +Model,Model Base Class,Model Card +stt_de_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_de_quartznet15x5" + diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_en.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_en.csv new file mode 100644 index 0000000000..ccdf05b806 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_en.csv @@ -0,0 +1,14 @@ +Model Name,Model Base Class,Model Card +QuartzNet15x5Base-En,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemospeechmodels" +stt_zh_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5" +stt_en_jasper10x5dr,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_jasper10x5dr" +stt_en_citrinet_256,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_256" +stt_en_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_512" +stt_en_citrinet_1024,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_citrinet_1024" +stt_zh_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512" +stt_en_conformer_ctc_small,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small" +stt_en_conformer_ctc_medium,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium" +stt_en_conformer_ctc_large,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large" +stt_en_conformer_ctc_small_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_small_ls" +stt_en_conformer_ctc_medium_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_medium_ls" +stt_en_conformer_ctc_large_ls,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls" diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_es.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_es.csv new file mode 100644 index 0000000000..2997fa360c --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_es.csv @@ -0,0 +1,3 @@ +Model,Model Base Class,Model Card +stt_es_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_es_quartznet15x5" + diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_fr.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_fr.csv new file mode 100644 index 0000000000..994ab0d54d --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_fr.csv @@ -0,0 +1,3 @@ +Model,Model Base Class,Model Card +stt_fr_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_fr_quartznet15x5" + diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_it.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_it.csv new file mode 100644 index 0000000000..d605b68809 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_it.csv @@ -0,0 +1,3 @@ +Model,Model Base Class,Model Card +stt_it_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_it_quartznet15x5" + diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_pl.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_pl.csv new file mode 100644 index 0000000000..bf646e1073 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_pl.csv @@ -0,0 +1,2 @@ +Model,Model Base Class,Model Card +stt_pl_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_pl_quartznet15x5" diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_ru.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_ru.csv new file mode 100644 index 0000000000..b46d4d9ca6 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_ru.csv @@ -0,0 +1,3 @@ +Model,Model Base Class,Model Card +stt_ru_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_ru_quartznet15x5" + diff --git a/models/tts/NeMo/docs/source/asr/data/benchmark_zh.csv b/models/tts/NeMo/docs/source/asr/data/benchmark_zh.csv new file mode 100644 index 0000000000..1fbc49879a --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/data/benchmark_zh.csv @@ -0,0 +1,3 @@ +Model,Model Base Class,Model Card +stt_zh_quartznet15x5,EncDecCTCModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_quartznet15x5" +stt_zh_citrinet_512,EncDecCTCBPEModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_zh_citrinet_512" diff --git a/models/tts/NeMo/docs/source/asr/datasets.rst b/models/tts/NeMo/docs/source/asr/datasets.rst new file mode 100644 index 0000000000..d4b2a2062f --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/datasets.rst @@ -0,0 +1,268 @@ +Datasets +======== + +NeMo has scripts to convert several common ASR datasets into the format expected by the ``nemo_asr`` collection. You can get started +with those datasets by following the instructions to run those scripts in the section appropriate to each dataset below. + +If the user has their own data and want to preprocess it to use with NeMo ASR models, refer to the `Preparing Custom ASR Data`_ section. + +If the user already has a dataset that you want to convert to a tarred format, refer to the `Tarred Datasets`_ section. + +.. _LibriSpeech_dataset: + +LibriSpeech +----------- + +Run the following scripts to download the LibriSpeech data and convert it into the format expected by `nemo_asr`. At least 250GB free +space is required. + +.. code-block:: bash + + # install sox + sudo apt-get install sox + mkdir data + python get_librispeech_data.py --data_root=data --data_set=ALL + +After this, the ``data`` folder should contain wav files and ``.json`` manifests for NeMo ASR datalayer. + +Each line is a training example. ``audio_filepath`` contains the path to the wav file, ``duration`` is the duration in seconds, and ``text`` is the transcript: + +.. code-block:: json + + {"audio_filepath": "/1355-39947-0000.wav", "duration": 11.3, "text": "psychotherapy and the community both the physician and the patient find their place in the community the life interests of which are superior to the interests of the individual"} + {"audio_filepath": "/1355-39947-0001.wav", "duration": 15.905, "text": "it is an unavoidable question how far from the higher point of view of the social mind the psychotherapeutic efforts should be encouraged or suppressed are there any conditions which suggest suspicion of or direct opposition to such curative work"} + +Fisher English Training Speech +------------------------------ + +Run these scripts to convert the Fisher English Training Speech data into a format expected by the ``nemo_asr`` collection. + +In brief, the following scripts convert the ``.sph`` files to ``.wav``, slices those files into smaller audio samples, matches the +smaller slices with their corresponding transcripts, and splits the resulting audio segments into train, validation, and test sets +(with one manifest each). + +.. note:: + - 106 GB of space is required to run the ``.wav`` conversion + - additional 105 GB is required for the slicing and matching + - ``sph2pipe`` is required in order to run the ``.wav`` conversion + +**Instructions** + +The following scripts assume that you already have the Fisher dataset from the Linguistic Data Consortium, with a directory structure +that looks similar to the following: + +.. code-block:: bash + + FisherEnglishTrainingSpeech/ + ├── LDC2004S13-Part1 + │   ├── fe_03_p1_transcripts + │   ├── fisher_eng_tr_sp_d1 + │   ├── fisher_eng_tr_sp_d2 + │   ├── fisher_eng_tr_sp_d3 + │   └── ... + └── LDC2005S13-Part2 + ├── fe_03_p2_transcripts + ├── fe_03_p2_sph1 + ├── fe_03_p2_sph2 + ├── fe_03_p2_sph3 + └── ... + +The transcripts that will be used are located in the ``fe_03_p<1,2>_transcripts/data/trans`` directory. The audio files (``.sph``) +are located in the remaining directories in an ``audio`` subdirectory. + +#. Convert the audio files from ``.sph`` to ``.wav`` by running: + + .. code-block:: bash + + cd /scripts + python fisher_audio_to_wav.py \ + --data_root= --dest_root= + + This will place the unsliced ``.wav`` files in ``/LDC200[4,5]S13-Part[1,2]/audio-wav/``. It will take several + minutes to run. + +#. Process the transcripts and slice the audio data. + + .. code-block:: bash + + python process_fisher_data.py \ + --audio_root= --transcript_root= \ + --dest_root= \ + --remove_noises + + This script splits the full dataset into train, validation, test sets, and places the audio slices in the corresponding folders + in the destination directory. One manifest is written out per set, which includes each slice's transcript, duration, and path. + + This will likely take around 20 minutes to run. Once finished, delete the 10 minute long ``.wav`` files. + +2000 HUB5 English Evaluation Speech +----------------------------------- + +Run the following script to convert the HUB5 data into a format expected by the ``nemo_asr`` collection. + +Similarly, to the Fisher dataset processing scripts, this script converts the ``.sph`` files to ``.wav``, slices the audio files and +transcripts into utterances, and combines them into segments of some minimum length (default is 10 seconds). The resulting segments +are all written out to an audio directory and the corresponding transcripts are written to a manifest JSON file. + +.. note:: + - 5 GB of free space is required to run this script + - ``sph2pipe`` is also required to be installed + +This script assumes you already have the 2000 HUB5 dataset from the Linguistic Data Consortium. + +Run the following command to process the 2000 HUB5 English Evaluation Speech samples: + +.. code-block:: bash + + python process_hub5_data.py \ + --data_root= \ + --dest_root= + +You can optionally include ``--min_slice_duration=`` if you would like to change the minimum audio segment duration. + +AN4 Dataset +----------- + +This is a small dataset recorded and distributed by Carnegie Mellon University. It consists of recordings of people spelling out +addresses, names, etc. Information about this dataset can be found on the `official CMU site `_. + +#. `Download and extract the dataset `_ (which is labeled "NIST's Sphere audio (.sph) format (64M)". + +#. Convert the ``.sph`` files to ``.wav`` using sox, and build one training and one test manifest. + + .. code-block:: bash + + python process_an4_data.py --data_root= + +After the script finishes, the ``train_manifest.json`` and ``test_manifest.json`` can be found in the ``/an4/`` directory. + +Aishell-1 +--------- + +To download the Aishell-1 data and convert it into a format expected by ``nemo_asr``, run: + +.. code-block:: bash + + # install sox + sudo apt-get install sox + mkdir data + python get_aishell_data.py --data_root=data + +After the script finishes, the ``data`` folder should contain a ``data_aishell`` folder which contains a wav file, a transcript folder, and related ``.json`` and ``vocab.txt`` files. + +Aishell-2 +--------- + +To process the AIShell-2 dataset, in the command below, set the data folder of AIShell-2 using ``--audio_folder`` and where to push +these files using ``--dest_folder``. In order to generate files in the supported format of ``nemo_asr``, run: + +.. code-block:: bash + + python process_aishell2_data.py --audio_folder= --dest_folder= + +After the script finishes, the ``train.json``, ``dev.json``, ``test.json``, and ``vocab.txt`` files can be found in the ``dest_folder`` directory. + +Preparing Custom ASR Data +------------------------- + +The ``nemo_asr`` collection expects each dataset to consist of a set of utterances in individual audio files plus +a manifest that describes the dataset, with information about one utterance per line (``.json``). +The audio files can be of any format supported by `Pydub `_, though we recommend +WAV files as they are the default and have been most thoroughly tested. + +There should be one manifest file per dataset that will be passed in, therefore, if the user wants separate training and validation +datasets, they should also have separate manifests. Otherwise, thay will be loading validation data with their training data and vice +versa. + +Each line of the manifest should be in the following format: + +.. code:: + + {"audio_filepath": "/path/to/audio.wav", "text": "the transcription of the utterance", "duration": 23.147} + +The :code:`audio_filepath` field should provide an absolute path to the ``.wav`` file corresponding to the utterance. +The :code:`text` field should contain the full transcript for the utterance, and the :code:`duration` field should +reflect the duration of the utterance in seconds. + +Each entry in the manifest (describing one audio file) should be bordered by '{' and '}' and must +be contained on one line. The fields that describe the file should be separated by commas, and have the form :code:`"field_name": value`, +as shown above. There should be no extra lines in the manifest, i.e. there should be exactly as many lines in the manifest as +there are audio files in the dataset. + +Since the manifest specifies the path for each utterance, the audio files do not have to be located +in the same directory as the manifest, or even in any specific directory structure. + +Once there is a manifest that describes each audio file in the dataset, use the dataset by passing +in the manifest file path in the experiment config file, e.g. as ``training_ds.manifest_filepath=``. + +Tarred Datasets +--------------- + +If experiments are run on a cluster with datasets stored on a distributed file system, the user will likely +want to avoid constantly reading multiple small files and would prefer tarring their audio files. +There are tarred versions of some NeMo ASR dataset classes for this case, such as the ``TarredAudioToCharDataset`` +(corresponding to the ``AudioToCharDataset``) and the ``TarredAudioToBPEDataset`` (corresponding to the +``AudioToBPEDataset``). The tarred audio dataset classes in NeMo use `WebDataset `_. + +To use an existing tarred dataset instead of a non-tarred dataset, set ``is_tarred: true`` in +the experiment config file. Then, pass in the paths to all of the audio tarballs in ``tarred_audio_filepaths``, either as a list +of filepaths, e.g. ``['/data/shard1.tar', '/data/shard2.tar']``, or in a single brace-expandable string, e.g. +``'/data/shard_{1..64}.tar'`` or ``'/data/shard__OP_1..64_CL_'`` (recommended, see note below). + +.. note:: + For brace expansion, there may be cases where ``{x..y}`` syntax cannot be used due to shell interference. This occurs most commonly + inside SLURM scripts. Therefore, we provide a few equivalent replacements. Supported opening braces (equivalent to ``{``) are ``(``, + ``[``, ``<`` and the special tag ``_OP_``. Supported closing braces (equivalent to ``}``) are ``)``, ``]``, ``>`` and the special + tag ``_CL_``. For SLURM based tasks, we suggest the use of the special tags for ease of use. + +As with non-tarred datasets, the manifest file should be passed in ``manifest_filepath``. The dataloader assumes that the length +of the manifest after filtering is the correct size of the dataset for reporting training progress. + +The ``tarred_shard_strategy`` field of the config file can be set if you have multiple shards and are running an experiment with +multiple workers. It defaults to ``scatter``, which preallocates a set of shards per worker which do not change during runtime. + +For more information about the individual tarred datasets and the parameters available, including shuffling options, +see the corresponding class APIs in the `Datasets <./api.html#Datasets>`__ section. + +.. warning:: + If using multiple workers, the number of shards should be divisible by the world size to ensure an even + split among workers. If it is not divisible, logging will give a warning but training will proceed, but likely hang at the last epoch. + In addition, if using distributed processing, each shard must have the same number of entries after filtering is + applied such that each worker ends up with the same number of files. We currently do not check for this in any dataloader, but the user's + program may hang if the shards are uneven. + +Conversion to Tarred Datasets +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +You can easily convert your existing NeMo-compatible ASR datasets using the +`conversion script here `_. + +.. code:: + + python convert_to_tarred_audio_dataset.py \ + --manifest_path= \ + --target_dir= \ + --num_shards= + --max_duration= \ + --min_duration= \ + --shuffle --shuffle_seed=0 + +This script shuffles the entries in the given manifest (if ``--shuffle`` is set, which we recommend), filter +audio files according to ``min_duration`` and ``max_duration``, and tar the remaining audio files to the directory +``--target_dir`` in ``n`` shards, along with separate manifest and metadata files. + +The files in the target directory should look similar to the following: + +.. code:: + + target_dir/ + ├── audio_1.tar + ├── audio_2.tar + ├── ... + ├── metadata.yaml + └── tarred_audio_manifest.json + +Note that file structures are flattened such that all audio files are at the top level in each tarball. This ensures that +filenames are unique in the tarred dataset and the filepaths do not contain "-sub" and forward slashes in each ``audio_filepath`` are +simply converted to underscores. For example, a manifest entry for ``/data/directory1/file.wav`` would be ``_data_directory1_file.wav`` +in the tarred dataset manifest, and ``/data/directory2/file.wav`` would be converted to ``_data_directory2_file.wav``. diff --git a/models/tts/NeMo/docs/source/asr/images/citrinet_vertical.png b/models/tts/NeMo/docs/source/asr/images/citrinet_vertical.png new file mode 100644 index 0000000000..a9a7f8c2e8 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/images/citrinet_vertical.png differ diff --git a/models/tts/NeMo/docs/source/asr/images/conformer_ctc.png b/models/tts/NeMo/docs/source/asr/images/conformer_ctc.png new file mode 100644 index 0000000000..e491856502 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/images/conformer_ctc.png differ diff --git a/models/tts/NeMo/docs/source/asr/images/ctc_asr.png b/models/tts/NeMo/docs/source/asr/images/ctc_asr.png new file mode 100644 index 0000000000..56ef705602 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/images/ctc_asr.png differ diff --git a/models/tts/NeMo/docs/source/asr/images/jasper_layers.png b/models/tts/NeMo/docs/source/asr/images/jasper_layers.png new file mode 100644 index 0000000000..af88379af9 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/images/jasper_layers.png differ diff --git a/models/tts/NeMo/docs/source/asr/images/jasper_vertical.png b/models/tts/NeMo/docs/source/asr/images/jasper_vertical.png new file mode 100644 index 0000000000..13ba93c2c7 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/images/jasper_vertical.png differ diff --git a/models/tts/NeMo/docs/source/asr/images/quartz_vertical.png b/models/tts/NeMo/docs/source/asr/images/quartz_vertical.png new file mode 100644 index 0000000000..4cbede9077 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/images/quartz_vertical.png differ diff --git a/models/tts/NeMo/docs/source/asr/intro.rst b/models/tts/NeMo/docs/source/asr/intro.rst new file mode 100644 index 0000000000..193e544dd7 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/intro.rst @@ -0,0 +1,59 @@ +Automatic Speech Recognition (ASR) +================================== + +ASR, or Automatic Speech Recognition, refers to the problem of getting a program to automatically transcribe spoken language +(speech-to-text). Our goal is usually to have a model that minimizes the Word Error Rate (WER) metric when transcribing speech input. +In other words, given some audio file (e.g. a WAV file) containing speech, how do we transform this into the corresponding text with +as few errors as possible? + +Traditional speech recognition takes a generative approach, modeling the full pipeline of how speech sounds are produced in order to +evaluate a speech sample. We would start from a language model that encapsulates the most likely orderings of words that are generated +(e.g. an n-gram model), to a pronunciation model for each word in that ordering (e.g. a pronunciation table), to an acoustic model that +translates those pronunciations to audio waveforms (e.g. a Gaussian Mixture Model). + +Then, if we receive some spoken input, our goal would be to find the most likely sequence of text that would result in the given audio +according to our generative pipeline of models. Overall, with traditional speech recognition, we try to model ``Pr(audio|transcript)*Pr(transcript)``, +and take the argmax of this over possible transcripts. + +Over time, neural nets advanced to the point where each component of the traditional speech recognition model could be replaced by a +neural model that had better performance and that had a greater potential for generalization. For example, we could replace an n-gram +model with a neural language model, and replace a pronunciation table with a neural pronunciation model, and so on. However, each of +these neural models need to be trained individually on different tasks, and errors in any model in the pipeline could throw off the +whole prediction. + +Thus, we can see the appeal of end-to-end ASR architectures: discriminative models that simply take an audio input and give a textual +output, and in which all components of the architecture are trained together towards the same goal. The model's encoder would be +akin to an acoustic model for extracting speech features, which can then be directly piped to a decoder which outputs text. If desired, +we could integrate a language model that would improve our predictions, as well. + +And the entire end-to-end ASR model can be trained at once--a much easier pipeline to handle! + +The full documentation tree is as follows: + +.. toctree:: + :maxdepth: 8 + + models + datasets + asr_language_modeling + results + configs + api + +Resources and Documentation +--------------------------- + +Hands-on speech recognition tutorial notebooks can be found under `the ASR tutorials folder `_. +If you are a beginner to NeMo, consider trying out the `ASR with NeMo `_ tutorial. +This and most other tutorials can be run on Google Colab by specifying the link to the notebooks' GitHub pages on Colab. + +If you are looking for information about a particular ASR model, or would like to find out more about the model +architectures available in the `nemo_asr` collection, refer to the :doc:`Models <./models>` section. + +NeMo includes preprocessing scripts for several common ASR datasets. The :doc:`Datasets <./datasets>` section contains instructions on +running those scripts. It also includes guidance for creating your own NeMo-compatible dataset, if you have your own data. + +Information about how to load model checkpoints (either local files or pretrained ones from NGC), as well as a list of the checkpoints +available on NGC are located on the :doc:`Checkpoints <./results>` section. + +Documentation regarding the configuration files specific to the ``nemo_asr`` models can be found on the :doc:`Configuration Files <./configs>` section. diff --git a/models/tts/NeMo/docs/source/asr/models.rst b/models/tts/NeMo/docs/source/asr/models.rst new file mode 100644 index 0000000000..0b9f071268 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/models.rst @@ -0,0 +1,97 @@ +Models +====== + +This section gives a brief overview of the models that NeMo's ASR collection currently supports. + +Each of these models can be used with the example ASR scripts (in the ``/examples/asr`` directory) by +specifying the model architecture in the config file used. Examples of config files for each model can be found in +the ``/examples/asr/conf`` directory. + +For more information about the config files and how they should be structured, refer to the :doc:`./configs` section. + +Pretrained checkpoints for all of these models, as well as instructions on how to load them, can be found in the :doc:`./results` +section. You can use the available checkpoints for immediate inference, or fine-tune them on your own datasets. The checkpoints section +also contains benchmark results for the available ASR models. + +.. _Jasper_model: + +Jasper +------ + +Jasper ("Just Another Speech Recognizer") :cite:`asr-models-li2019jasper` is a deep time delay neural network (TDNN) comprising of +blocks of 1D-convolutional layers. The Jasper family of models are denoted as ``Jasper_[BxR]`` where ``B`` is the number of blocks +and ``R`` is the number of convolutional sub-blocks within a block. Each sub-block contains a 1-D convolution, batch normalization, +ReLU, and dropout: + + .. image:: images/jasper_vertical.png + :align: center + :alt: japer model + :scale: 50% + +Jasper models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModel` class. + +QuartzNet +--------- + +QuartzNet :cite:`asr-models-kriman2019quartznet` is a version of Jasper :cite:`asr-models-li2019jasper` model with separable +convolutions and larger filters. It can achieve performance similar to Jasper but with an order of magnitude fewer parameters. +Similarly to Jasper, the QuartzNet family of models are denoted as ``QuartzNet_[BxR]`` where ``B`` is the number of blocks and ``R`` +is the number of convolutional sub-blocks within a block. Each sub-block contains a 1-D *separable* convolution, batch normalization, +ReLU, and dropout: + + .. image:: images/quartz_vertical.png + :align: center + :alt: quartznet model + :scale: 40% + +QuartzNet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModel` class. + +Citrinet +-------- + +Citrinet is a version of QuartzNet :cite:`asr-models-kriman2019quartznet` that extends ContextNet :cite:`asr-models-han2020contextnet`, +utilizing subword encoding (via Word Piece tokenization) and Squeeze-and-Excitation mechanism :cite:`asr-models-hu2018squeeze` to +obtain highly accurate audio transcripts while utilizing a non-autoregressive CTC based decoding scheme for efficient inference. + + .. image:: images/citrinet_vertical.png + :align: center + :alt: citrinet model + :scale: 50% + +Citrinet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModelBPE` class. + +.. _Conformer-CTC_model: + +Conformer-CTC +------------- + +Conformer-CTC is a CTC-based variant of the Conformer model introduced in :cite:`asr-models-gulati2020conformer`. Conformer-CTC has a +similar encoder as the original Conformer but uses CTC loss and decoding instead of RNNT loss, which makes it a non-autoregressive model. +We also drop the LSTM decoder and instead use a linear decoder on the top of the encoder. This model uses the combination of +self-attention and convolution modules to achieve the best of the two approaches, the self-attention layers can learn the global +interaction while the convolutions efficiently capture the local correlations. The self-attention modules support both regular +self-attention with absolute positional encoding, and also Transformer-XL's self-attention with relative positional encodings. + +Here is the overall architecture of the encoder of Conformer-CTC: + + .. image:: images/conformer_ctc.png + :align: center + :alt: Conformer-CTC Model + :scale: 50% + +This model supports both the sub-word level and character level encodings. You can find more details on the config files for the +Conformer-CTC models at `Conformer-CTC <./configs.html#conformer-ctc>`. The variant with sub-word encoding is a BPE-based model +which can be instantiated using the :class:`~nemo.collections.asr.models.EncDecCTCModelBPE` class, while the +character-based variant is based on :class:`~nemo.collections.asr.models.EncDecCTCModel`. + +You may find the example config files of Conformer-CTC model with character-based encoding at +``/examples/asr/conf/conformer/conformer_ctc_char.yaml`` and +with sub-word encoding at ``/examples/asr/conf/conformer/conformer_ctc_bpe.yaml``. + +References +---------- + +.. bibliography:: asr_all.bib + :style: plain + :labelprefix: ASR-MODELS + :keyprefix: asr-models- diff --git a/models/tts/NeMo/docs/source/asr/results.rst b/models/tts/NeMo/docs/source/asr/results.rst new file mode 100644 index 0000000000..df9fa83748 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/results.rst @@ -0,0 +1,181 @@ +Checkpoints +=========== + +There are two main ways to load pretrained checkpoints in NeMo: + +* Using the :code:`restore_from()` method to load a local checkpoint file (``.nemo``), or +* Using the :code:`from_pretrained()` method to download and set up a checkpoint from NGC. + +Refer to the following sections for instructions and examples for each. + +Note that these instructions are for loading fully trained checkpoints for evaluation or fine-tuning. For resuming an unfinished +training experiment, use the Experiment Manager to do so by setting the ``resume_if_exists`` flag to ``True``. + +Loading Local Checkpoints +------------------------- + +NeMo automatically saves checkpoints of a model that is trained in a ``.nemo`` format. Alternatively, to manually save the model at any +point, issue :code:`model.save_to(.nemo)`. + +If there is a local ``.nemo`` checkpoint that you'd like to load, use the :code:`restore_from()` method: + +.. code-block:: python + + import nemo.collections.asr as nemo_asr + model = nemo_asr.models..restore_from(restore_path="") + +Where the model base class is the ASR model class of the original checkpoint, or the general ``ASRModel`` class. + +NGC Pretrained Checkpoints +-------------------------- + +The ASR collection has checkpoints of several models trained on various datasets for a variety of tasks. These checkpoints are +obtainable via NGC `NeMo Automatic Speech Recognition collection `_. +The model cards on NGC contain more information about each of the checkpoints available. + +The tables below list the ASR models available from NGC. The models can be accessed via the :code:`from_pretrained()` method inside +the ASR Model class. In general, you can load any of these models with code in the following format: + +.. code-block:: python + + import nemo.collections.asr as nemo_asr + model = nemo_asr.models.ASRModel.from_pretrained(model_name="") + +Where the model name is the value under "Model Name" entry in the tables below. + +For example, to load the base English QuartzNet model for speech recognition, run: + +.. code-block:: python + + model = nemo_asr.models.ASRModel.from_pretrained(model_name="QuartzNet15x5Base-En") + +You can also call :code:`from_pretrained()` from the specific model class (such as :code:`EncDecCTCModel` +for QuartzNet) if you need to access a specific model functionality. + +If you would like to programatically list the models available for a particular base class, you can use the +:code:`list_available_models()` method. + +.. code-block:: python + + nemo_asr.models..list_available_models() + +Transcribing/Inference +^^^^^^^^^^^^^^^^^^^^^^ + +To perform inference and transcribe a sample of speech after loading the model, use the ``transcribe()`` method: + +.. code-block:: python + + model.transcribe(paths2audio_files=[list of audio files], batch_size=BATCH_SIZE, logprobs=False) + +Setting the argument ``logprobs`` to ``True`` returns the log probabilities instead of transcriptions. For more information, see :doc:`./api.html#modules`. +The audio files should be 16KHz monochannel wav files. + +Fine-tuning on Different Datasets +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +There are multiple ASR tutorials provided in the :ref:`Tutorials ` section. Most of these tutorials explain how to instantiate a pre-trained model, prepare the model for fine-tuning on some dataset (in the same language) as a demonstration. + + +Automatic Speech Recognition Models +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. csv-table:: + :file: data/asr_results.csv + :align: left + :widths: 30, 30, 40 + :header-rows: 1 + +Language Models for ASR +^^^^^^^^^^^^^^^^^^^^^^^ + +.. csv-table:: + :file: data/asrlm_results.csv + :align: left + :widths: 30, 30, 40 + :header-rows: 1 + + +Speech Recognition (Languages) +------------------------------ + +English +^^^^^^^ +.. csv-table:: + :file: data/benchmark_en.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + +----------------------------- + +Mandarin +^^^^^^^^ +.. csv-table:: + :file: data/benchmark_zh.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + +----------------------------- + +German +^^^^^^ +.. csv-table:: + :file: data/benchmark_de.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + +----------------------------- + +Polish +^^^^^^ +.. csv-table:: + :file: data/benchmark_pl.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + +----------------------------- + +Italian +^^^^^^^ +.. csv-table:: + :file: data/benchmark_it.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + +----------------------------- + +Russian +^^^^^^^ +.. csv-table:: + :file: data/benchmark_ru.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + +----------------------------- + +Spanish +^^^^^^^ +.. csv-table:: + :file: data/benchmark_es.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + + +----------------------------- + +Catalan +^^^^^^^ +.. csv-table:: + :file: data/benchmark_ca.csv + :align: left + :widths: 40, 10, 50 + :header-rows: 1 + + diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/api.rst b/models/tts/NeMo/docs/source/asr/speaker_diarization/api.rst new file mode 100644 index 0000000000..a820e1dc7b --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/api.rst @@ -0,0 +1,17 @@ +NeMo Speaker Diarization API +============================= + + +Model Classes +------------- +.. autoclass:: nemo.collections.asr.models.ClusteringDiarizer + :show-inheritance: + :members: init_speaker_model, set_vad_model, init_vad_model, setup_vad_test_data, setup_spkr_test_data, run_vad, extract_embeddings, path2audio_files_to_manifest, diarize, make_nemo_file_from_folder, save_to, unpack_nemo_file, restore_from + + +Mixins +------ + +.. autoclass:: nemo.collections.asr.parts.mixins.mixins.DiarizationMixin + :show-inheritance: + :members: diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/configs.rst b/models/tts/NeMo/docs/source/asr/speaker_diarization/configs.rst new file mode 100644 index 0000000000..f7910be8a3 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/configs.rst @@ -0,0 +1,65 @@ +NeMo Speaker Diarization Configuration Files +============================================ + +Since speaker diarization model here is not a fully-trainble End-to-End model but an inference pipeline, we use **diarizer** instead of **model** which is used in other tasks. + +The diarizer section will generally require information about the dataset(s) being +used, models used in this pipline, as well as inference related parameters such as postprocessing of each models. +The sections on this page cover each of these in more detail. + + +Example configuration files for speaker diarization can be found in the +``/examples/speaker_recognition/conf/speaker_diarization.yaml>`` + +.. note:: + For model details and deep understanding about configs, finetuning, tuning threshold, and evaluation, + please refer to ``/tutorials/speaker_recognition/Speaker_Diarization_Inference.ipynb``; + for other applications such as possible integration with ASR, have a look at ``/tutorials/speaker_recognition/ASR_with_SpeakerDiarization.ipynb``. + + +Dataset Configuration +----------------------- + +In contrast to other ASR related tasks or models in NeMo, speaker diarization supported in NeMo is a modular inference pipeline. +Datasets here denotes the data you would like to perform speaker diarization on. + +An example Speaker Diarization dataset configuration could look like: + +.. code-block:: yaml + + diarizer: + oracle_num_speakers: ??? # if number of speakers is known input it or pass null if not known. Accepts int or path to file containing uniq-id and num of speakers of that session + max_num_speakers: 8 # max number of speakers for each recording. If oracle num speakers is passed this value is ignored + out_dir: ??? + paths2audio_files: ??? # file containing paths to audio files for diarzation + path2groundtruth_rttm_files: null # (Optional) file containing paths to rttm files (this can be passed if we need to calculate DER rate based on ground truth rttm files). + ... + +.. note:: + We expect audio and the corresponding RTTM to have the same base name and the name should be unique. + + +Diarizer Architecture Configurations +------------------------------------- + +.. code-block:: yaml + + diarizer: + ... + vad: + model_path: null #.nemo local model path or pretrained model name or none + window_length_in_sec: 0.15 + shift_length_in_sec: 0.01 + threshold: 0.5 # tune threshold on dev set. Check /scripts/voice_activity_detection/vad_tune_threshold.py + vad_decision_smoothing: True + smoothing_params: + method: "median" + overlap: 0.875 + + speaker_embeddings: + oracle_vad_manifest: null # leave it null if to perform diarization with above VAD model else path to manifest file genrerated as shown in Datasets section + model_path: ??? #.nemo local model path or pretrained model name + window_length_in_sec: 1.5 # for better performace try with 3 sec and 1.5 sec + shift_length_in_sec: 0.75 # for better performace try with 1.5 sec and 0.75 sec + + diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/data/diarization_results.csv b/models/tts/NeMo/docs/source/asr/speaker_diarization/data/diarization_results.csv new file mode 100644 index 0000000000..fba1a2eda3 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/data/diarization_results.csv @@ -0,0 +1,4 @@ +Model Name,Model Base Class,Model Card +vad_marblenet,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:vad_marblenet" +vad_telephony_marblenet,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:vad_telephony_marblenet" +speakerdiarization_speakernet,EncDecSpeakerLabelModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerdiarization_speakernet" diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/datasets.rst b/models/tts/NeMo/docs/source/asr/speaker_diarization/datasets.rst new file mode 100644 index 0000000000..f178ad44d3 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/datasets.rst @@ -0,0 +1,72 @@ +Datasets +======== +Check out page :doc:`Speech Classification Datasets <../speech_classification/datasets>` and :doc:`Speaker Recogniton Datasets <../speaker_recognition/datasets>` +for preparing datasets for training and validating VAD and speaker embedding models respectively. + +For Speaker Diarization inference, ``diarizer`` expects following arguments : + + - paths2audio_files + - oracle_num_speakers + - path2groundtruth_rttm_files [Optional] + - oracle_vad_manifest (if vad/rttm files are known)[Optional but suggested for good performace. This avoids FPR and TPR] + +paths2audio_files: + - a file containing absolute paths to audio files or list of paths to audio files. + For generating a file that contains paths to audio files, you can simply use ``find`` bash command as shown below: +.. code-block:: bash + + find $PWD/{data_dir} -iname '*.wav' > path_to_audiofiles.txt + +.. note:: + We expect audio and the corresponding RTTM files to have the same base name and the name should be unique (uniq-id). + +oracle_num_speakers: + - if number of speakers is known input it or pass null if not known. Accepts int or path to file containing uniq-id with num of speakers of that session + + Sample line (uniq-id `iaaa` session has 2 unique speakers ):: + + iaaa 2 + +path2groundtruth_rttm_files [Optional]: + - To evaluate the diarizer system with known rttm files, One needs to provide a txt file (path2groundtruth_rttm_files) for groundtruth label files. + use above mentioned ``find`` command to get all reference rttm files (use '*.rttm' as search pattern) + +Each groundtruth label file should be in NIST Rich Transcription Time Marked (RTTM) format. Take one line from a RTTM file for example: + +.. code-block:: bash + + SPEAKER TS3012d.Mix-Headset 1 331.573 0.671 MTD046ID + + +oracle_vad_manifest [Optional but suggested for good performace. This reduces FPR and TPR]: + - To perform just oracle diarization, that is taking speech activity time stamps from groundtruth RTTMs or another VAD from wild instead of NeMo VAD output, ``diarizer`` expects an oracle manifest json file that contains paths to audio files with offset for start time and duration of segment. + +To prepare an oracle manifest file, use the script from ``scripts`` folder as shown below: + +.. code-block:: bash + + python $NeMo/scripts/ --paths2rttm_files= --paths2audio_files= --manifest_file= + +Here ``paths2audio_files`` and ``path2groundtruth_rttm_files`` are files containing paths to audio files as shown above. + +AMI Meeting Corpus +------------------ + +The following are the suggested parameters for getting Speaker Error Rate (SER) of 4.1% on AMI Lapel test set corpus: + - diarizer.oracle_num_speakers = 4 (since there are exactly 4 speakers per each Lapel test set session) + - diarizer.speaker_embeddings.model_path = ``speakerverification_speakernet`` (This model is trained on voxceleb dataset. ``Use this model for simialr non-telephonic speech datasets``) + - diarizer.speaker_embeddings.window_length_in_sec = 3 + - diarizer.speaker_embeddings.shift_length_in_sec = 1.5 + +Input paths2audio_files, paths2rttm_files and oracle_vad_manifest by following steps as shown above + +CallHome LDC97S42 (CH109) +------------------------- + +The following are the suggested parameters for getting Speaker Error Rate (SER) of 5.4% on CH109 set: + - diarizer.oracle_num_speakers = 2 (since there are exactly 2 speakers per each ch109 session) + - diarizer.speaker_embeddings.model_path = ``speakerdiarization_speakernet`` (This model is trained on voxceleb and telephonic speech Fisher and SWBD. ``Use this model for similar telephonic speech datasets``) + - diarizer.speaker_embeddings.window_length_in_sec = 1.5 + - diarizer.speaker_embeddings.shift_length_in_sec = 0.75 + +Input paths2audio_files, paths2rttm_files and oracle_vad_manifest by following steps as shown above \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/images/sd_pipeline.png b/models/tts/NeMo/docs/source/asr/speaker_diarization/images/sd_pipeline.png new file mode 100644 index 0000000000..258781a2ae Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/speaker_diarization/images/sd_pipeline.png differ diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/intro.rst b/models/tts/NeMo/docs/source/asr/speaker_diarization/intro.rst new file mode 100644 index 0000000000..7ec454ba9d --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/intro.rst @@ -0,0 +1,52 @@ +Speaker Diarization +================================== + +Speaker Diarization (SD) is the task of segmenting audio recordings by speaker labels, that is Who Speaks When? + +.. image:: images/sd_pipeline.png + :align: center + :alt: Speaker Diarization pipeline [todo] + :scale: 50% + + +A diarization system consists of a **Voice Activity Detection (VAD)** model to get the time stamps of audio where speech is being spoken while ignoring the background noise and +a **Speaker Embeddings** model to get speaker embeddings on speech segments obtained from VAD time stamps. +These speaker embeddings would then be clustered into clusters based on number of speakers present in the audio recording. + +In NeMo we support both **oracle VAD** and **non-oracle VAD** diarization. + + +The full documentation tree is as follows: + +.. toctree:: + :maxdepth: 8 + + models + datasets + results + configs + api + +Resource and Documentation Guide +-------------------------------- + +Hands-on speaker diarization tutorial notebooks can be found under ``/tutorials/speaker_recognition/``. + +There are tutorials for peformming inference using :ref:`MarbleNet_model` and :ref:`SpeakerNet_model`, +and how one can get ASR transcriptions combined with Speaker labels along with voice activity time stamps with NeMo asr collections. + +Most of the tutorials can be run on Google Colab by specifying the link to the notebooks' GitHub pages on Colab. + +If you are looking for information about a particular model used for speaker diarization inference, or would like to find out more about the model +architectures available in the `nemo_asr` collection, check out the :doc:`Models <./models>` page. + +Documentation on dataset preprocessing can be found on the :doc:`Datasets <./datasets>` page. +NeMo includes preprocessing scripts for several common ASR datasets, and this page contains instructions on running +those scripts. +It also includes guidance for creating your own NeMo-compatible dataset, if you have your own data. + +Information about how to load model checkpoints (either local files or pretrained ones from NGC), perform inference, as well as a list +of the checkpoints available on NGC are located on the :doc:`Checkpoints <./results>` page. + +Documentation for configuration files specific to the ``nemo_asr`` models can be found on the +:doc:`Configuration Files <./configs>` page. diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/models.rst b/models/tts/NeMo/docs/source/asr/speaker_diarization/models.rst new file mode 100644 index 0000000000..7a543e92f4 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/models.rst @@ -0,0 +1,5 @@ +Models +====== + + +Currenlty NeMo's Speaker Diarization pipeline uses `MarbleNet <../speech_classification/models.html#marblenet-vad>`__ model for Voice Activity Detection (VAD) and `SpeakerNet <../speaker_recognition/models.html#speakernet>`__ model for Speaker Embedding Extraction. \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speaker_diarization/results.rst b/models/tts/NeMo/docs/source/asr/speaker_diarization/results.rst new file mode 100644 index 0000000000..c20be139f7 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_diarization/results.rst @@ -0,0 +1,76 @@ +Checkpoints +=========== + +There are two main ways to load pretrained checkpoints in NeMo as introduced in `loading ASR checkpoints <../results.html#checkpoints>`__. +In speaker diarization, the diarizer loads checkpoints that are passed through the config file. For example: + +Loading Local Checkpoints +--------------------------- + +Load VAD model + +.. code-block:: bash + + pretrained_vad_model='/path/to/vad_marblenet.nemo' # local .nemo or pretrained vad model name + ... + # pass with hydra config + config.diarizer.vad.model_path=pretrained_vad_model + + +Load Speaker Embedding model + +.. code-block:: bash + + pretrained_speaker_model='/path/to/speakerdiarization_speakernet.nemo' # local .nemo or pretrained speakernet model name + ... + # pass with hydra config + config.diarizer.speaker_embeddings.model_path=pretrained_speaker_model + + +NeMo will automatically save checkpoints of a model you are training in a `.nemo` format. +You can also manually save your models at any point using :code:`model.save_to(.nemo)`. + + +Inference +---------- + +.. note:: + For details and deep understanding, please refer to ``/tutorials/speaker_recognition/Speaker_Diarization_Inference.ipynb``. + +Check out :doc:`Datasets <./datasets>` for preparing audio files and optional label files. + +Run and evaluate speaker diarizer with below command: + +.. code-block:: bash + + # Have a look at the instruction inside the script and pass the arguments you might need. + python /examples/speaker_recognition/speaker_diarize.py + + +NGC Pretrained Checkpoints +---------------------------- + +The ASR collection has checkpoints of several models trained on various datasets for a variety of tasks. +These checkpoints are obtainable via NGC `NeMo Automatic Speech Recognition collection `_. +The model cards on NGC contain more information about each of the checkpoints available. + +In general, you can load models with model name in the following format, + +.. code-block:: python + + pretrained_vad_model='vad_telephony_marblenet' + pretrained_speaker_model='speakerdiarization_speakernet' + ... + config.diarizer.vad.model_path=retrained_vad_model \ + config.diarizer.speaker_embeddings.model_path=pretrained_speaker_model + +where the model name is the value under "Model Name" entry in the tables below. + +Models for Speaker Diarization +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. csv-table:: + :file: data/diarization_results.csv + :align: left + :widths: 30, 30, 40 + :header-rows: 1 diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/configs.rst b/models/tts/NeMo/docs/source/asr/speaker_recognition/configs.rst new file mode 100644 index 0000000000..7b84f474e2 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_recognition/configs.rst @@ -0,0 +1,117 @@ +NeMo ASR Configuration Files +============================ + +This page covers NeMo configuration file setup that is specific to speaker recognition models. +For general information about how to set up and run experiments that is common to all NeMo models (e.g. +experiment manager and PyTorch Lightning trainer parameters), see the :doc:`../../core/core` page. + +The model section of NeMo ASR configuration files will generally require information about the dataset(s) being +used, the preprocessor for audio files, parameters for any augmentation being performed, as well as the +model architecture specification. +The sections on this page cover each of these in more detail. + +Example configuration files for all of the Speaker related scripts can be found in the +config directory of the examples ``{NEMO_ROOT/examples/speaker_recognition/conf}``. + + +Dataset Configuration +--------------------- + +Training, validation, and test parameters are specified using the ``train_ds``, ``validation_ds``, and +``test_ds`` sections of your configuration file, respectively. +Depending on the task, you may have arguments specifying the sample rate of your audio files, max time length to consider for each audio file , whether or not to shuffle the dataset, and so on. +You may also decide to leave fields such as the ``manifest_filepath`` blank, to be specified via the command line +at runtime. + +Any initialization parameters that are accepted for the Dataset class used in your experiment +can be set in the config file. + +An example SpeakerNet train and validation configuration could look like: + +.. code-block:: yaml + + model: + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: None # finds labels based on manifest file + batch_size: 32 + trim_silence: False + time_length: 8 + shuffle: True + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: None # Keep None, to match with labels extracted during training + batch_size: 32 + shuffle: False # No need to shuffle the validation data + + +If you would like to use tarred dataset, have a look at `Datasets Configuration <../configs.html#dataset-configuration>`__. + + +Preprocessor Configuration +-------------------------- +Preprocessor helps to compute MFCC or mel spectrogram features that are given as inputs to model. +For details on how to write this section, refer to `Preprocessor Configuration <../configs.html#preprocessor-configuration>`__ + + +Augmentation Configurations +--------------------------- + +For SpeakerNet training we use on-the fly augmentations with MUSAN and RIR impulses using ``noise`` augmentor section + +The following example sets up musan augmentation with audio files taken from manifest path and +minimum and maximum SNR specified with min_snr and max_snr respectively. This section can be added to +``train_ds`` part in model + +.. code-block:: yaml + + model: + ... + train_ds: + ... + augmentor: + noise: + manifest_path: /path/to/musan/manifest_file + prob: 0.2 # probability to augment the incoming batch audio with augmentor data + min_snr_db: 5 + max_snr_db: 15 + + +See the :class:`nemo.collections.asr.parts.preprocessing.perturb.AudioAugmentor` API section for more details. + + +Model Architecture Configurations +--------------------------------- + +Each configuration file should describe the model architecture being used for the experiment. +Models in the NeMo ASR collection need a ``encoder`` section and a ``decoder`` section, with the ``_target_`` field +specifying the module to use for each. + +The following sections go into more detail about the specific configurations of each model architecture. + +For more information about the SpeakerNet Encoder models, see the :doc:`Models <./models>` page and at `Jasper and QuartzNet <../configs.html#jasper-and-quartznet>`__ + +Decoder Configurations +------------------------ + +After features have been computed from speakernet encoder, we pass these features to the decoder to compute embeddings and then to compute log probabilities +for training models. + +.. code-block:: yaml + + model: + ... + decoder: + _target_: nemo.collections.asr.modules.SpeakerDecoder + feat_in: *enc_final + num_classes: 7205 # Total number of classes in training manifest file + pool_mode: xvector # xvector for variance and mean bases statistics pooling + emb_sizes: 256 # number of inermediate emb layers. can be comma separated for additional layers like 512,512 + angular: true # if true then loss will be changed to angular softmax loss and consider scale and margin from loss section else train with cross-entrophy loss + + loss: + scale: 30 + margin 0.2 \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/data/speaker_results.csv b/models/tts/NeMo/docs/source/asr/speaker_recognition/data/speaker_results.csv new file mode 100644 index 0000000000..450c0a1da2 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_recognition/data/speaker_results.csv @@ -0,0 +1,4 @@ +Model Name,Model Base Class,Model Card +speakerverification_speakernet,EncDecSpeakerLabelModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerverification_speakernet" +speakerrecognition_speakernet,EncDecSpeakerLabelModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerrecognition_speakernet" +speakerdiarization_speakernet,EncDecSpeakerLabelModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:speakerdiarization_speakernet" diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/datasets.rst b/models/tts/NeMo/docs/source/asr/speaker_recognition/datasets.rst new file mode 100644 index 0000000000..7439d7263c --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_recognition/datasets.rst @@ -0,0 +1,57 @@ +Datasets +======== + +.. _HI-MIA: + +HI-MIA +-------- + +Run the script to download and process ``hi-mia`` dataset in order to generate files in the supported format of ``nemo_asr``. You should set the data folder of +hi-mia using ``--data_root``. These scripts are present in ``/scripts`` + +.. code-block:: bash + + python get_hi-mia_data.py --data_root= + +After download and conversion, your `data` folder should contain directories with following set of files as: + +* `data//train.json` +* `data//dev.json` +* `data//{set}_all.json` +* `data//utt2spk` + + +All-other Datasets +------------------ + +These methods can be applied to any dataset to get similar training manifest files. + +First we prepare scp file(s) containing absolute paths to all the wav files required for each of the train, dev, and test set. This can be easily prepared by using ``find`` bash command as follows: + +.. code-block:: bash + + !find {data_dir}/{train_dir} -iname "*.wav" > data/train_all.scp + !head -n 3 data/train_all.scp + + +Since we created the scp file for the train, we use `scp_to_manifest.py` to convert this scp file to a manifest file and then optionally split the files to train \& dev for evaluating the models while training by using the --split flag. +We wont be needing the --split option for the test folder. Accordingly please mention the id number, which is the field num separated by / to be considered as the speaker label. +After the download and conversion, your data folder should contain directories with manifest files as: + +* `data//train.json` +* `data//dev.json` +* `data//train_all.json` + +Each line in the manifest file describes a training sample - audio_filepath contains the path to the wav file, duration it's duration in seconds, and label is the speaker class label: + +.. code-block:: json + + {"audio_filepath": "/audio_file.wav", "duration": 3.9, "label": "speaker_id"} + + +Tarred Datasets +--------------- + +Similarly to ASR, you can tar your audio files and use ASR Dataset class ``TarredAudioToSpeechLabelDataset`` (corresponding to the ``AudioToSpeechLabelDataset``) for this case. + +If you want to use tarred dataset, have a look at `ASR Tarred Datasets <../datasets.html#tarred-datasets>`__. \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/images/ICASPP_SpeakerNet.png b/models/tts/NeMo/docs/source/asr/speaker_recognition/images/ICASPP_SpeakerNet.png new file mode 100644 index 0000000000..ecdccccb69 Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/speaker_recognition/images/ICASPP_SpeakerNet.png differ diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/intro.rst b/models/tts/NeMo/docs/source/asr/speaker_recognition/intro.rst new file mode 100644 index 0000000000..f1709a99e4 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_recognition/intro.rst @@ -0,0 +1,48 @@ +Speaker Recognition (SR) +======================== + +Speaker Recognition (SR) is a broad research area which solves two major tasks: speaker identification (who is speaking?) and speaker verification (is the speaker who they claim to be?). +We focus on far-field, text-independent speaker recognition when the identity of the speaker is based on how the speech is spoken, not necessarily in what is being said. +Typically such SR systems operate on unconstrained speech utterances, which are converted into vectors of fixed length, called speaker embeddings. +Speaker embeddings can also be used in automatic speech recognition (ASR) and speech synthesis. + +As the goal of most speaker recognition systems is to get good speaker level embeddings that could help distinguish from other speakers, we shall first train these embeddings in end-to-end manner optimizing the QuatzNet based encoder model on cross-entropy loss. +We modify the decoder to get these fixed size embeddings irrespective of the length of the input audio and employ a mean and variance based statistics pooling method to grab these embeddings. + +In Speaker Identification we typically train on a larger training set with cross-entrophy loss and finetune later on your preferred set of labels where one would want to classify only known set of speakers. +And in Speaker verification we train with Angular Softmax loss and compare embedings extracted from one audio file coming from a single speaker with +embeddings extracted from another file of same or another speaker by employing backend scoring techniques like cosine similarity. + + +The full documentation tree is as follows: + +.. toctree:: + :maxdepth: 8 + + models + configs + datasets + results + +Resource and Documentation Guide +-------------------------------- + +Hands-on speaker recognition tutorial notebooks can be found under +`the speaker recognition tutorials folder `_. This and most other tutorials can be run on Google Colab by specifying the link to the notebooks' GitHub pages on Colab. + +If you are looking for information about a particular SpeakerNet model, or would like to find out more about the model +architectures available in the ``nemo_asr`` collection, check out the :doc:`Models <./models>` page. + +Documentation on dataset preprocessing can be found on the :doc:`Datasets <./datasets>` page. +NeMo includes preprocessing and other scripts for speaker_recognition in folder, and this page contains instructions on running +those scripts. It also includes guidance for creating your own NeMo-compatible dataset, if you have your own data. + +Information about how to load model checkpoints (either local files or pretrained ones from NGC), perform inference, as well as a list +of the checkpoints available on NGC are located on the :doc:`Checkpoints <./results>` page. + +Documentation for configuration files specific to the ``nemo_asr`` models can be found on the +:doc:`Configuration Files <./configs>` page. + + +For a clear step-by-step tutorial we advice you to refer tutorials found in `folder `_. + diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/models.rst b/models/tts/NeMo/docs/source/asr/speaker_recognition/models.rst new file mode 100644 index 0000000000..ebfc5ba42c --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_recognition/models.rst @@ -0,0 +1,52 @@ +Models +====== + +Examples of config files for the following quartznet model can be found in the ``/examples/speaker_recogniton/conf`` directory. + +For more information about the config files and how they should be structured, see the :doc:`./configs` page. + +Pretrained checkpoints for all of these models, as well as instructions on how to load them, can be found on the :doc:`./results` page. +You can use the available checkpoints for immediate inference, or fine-tune them on your own datasets. +The Checkpoints page also contains benchmark results for the available speaker recognition models. + +.. _SpeakerNet_model: + +SpeakerNet +----------- + +The model is based on the QuartzNet ASR architecture :cite:`sr-models-koluguri2020speakernet` +comprising of an encoder and decoder structure. We use the encoder of the QuartzNet model as a top-level feature extractor, and feed the output to the statistics pooling layer, where +we compute the mean and variance across channel dimensions to capture the time-independent utterance-level speaker features. + +The QuartzNet encoder used for speaker embeddings shown in figure below has the following structure: a QuartzNet BxR +model has B blocks, each with R sub-blocks. Each sub-block applies the following operations: a 1D convolution, batch norm, ReLU, and dropout. All sub-blocks in a block have the same number of output channels. These blocks are connected with residual connections. We use QuartzNet with 3 blocks, 2 sub-blocks, and 512 channels, as the Encoder for Speaker Embeddings. All conv layers have stride 1 and dilation 1. + + + .. image:: images/ICASPP_SpeakerNet.png + :align: center + :alt: speakernet model + :scale: 40% + +Top level acoustic Features, obtained from the output of +encoder are used to compute intermediate features that are +then passed to the decoder for getting utterance level speaker +embeddings. The intermediate time-independent features are +computed using a statistics pooling layer, where we compute the mean and standard deviation of features across +time-channels, to get a time-independent feature representation S of size Batch_size × 3000. +The intermediate features, S are passed through the Decoder consisting of two layers each of output size 512 for a +linear transformation from S to the final number of classes +N for the larger (L) model, and a single linear layer of output size 256 to the final number of classes N for the medium +(M) model. We extract q-vectors after the final linear layer +of fixed size 512, 256 for SpeakerNet-L and SpeakerNet-M +models respectively. + +SpeakerNet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecSpeakerLabelModel` class. + + +References +----------- + +.. bibliography:: ../asr_all.bib + :style: plain + :labelprefix: SR-MODELS + :keyprefix: sr-models- diff --git a/models/tts/NeMo/docs/source/asr/speaker_recognition/results.rst b/models/tts/NeMo/docs/source/asr/speaker_recognition/results.rst new file mode 100644 index 0000000000..eb2472e6d9 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speaker_recognition/results.rst @@ -0,0 +1,85 @@ +Checkpoints +=========== + +There are two main ways to load pretrained checkpoints in NeMo: + +* Using the :code:`restore_from()` method to load a local checkpoint file (`.nemo`), or +* Using the :code:`from_pretrained()` method to download and set up a checkpoint from NGC. + +See the following sections for instructions and examples for each. + +Note that these instructions are for loading fully trained checkpoints for evaluation or fine-tuning. +For resuming an unfinished training experiment, please use the experiment manager to do so by setting the +``resume_if_exists`` flag to True. + +Loading Local Checkpoints +------------------------- + +NeMo will automatically save checkpoints of a model you are training in a `.nemo` format. +You can also manually save your models at any point using :code:`model.save_to(.nemo)`. + +If you have a local ``.nemo`` checkpoint that you'd like to load, simply use the :code:`restore_from()` method: + +.. code-block:: python + + import nemo.collections.asr as nemo_asr + model = nemo_asr.models..restore_from(restore_path="") + +Where the model base class is the ASR model class of the original checkpoint, or the general `ASRModel` class. + + +Inference +---------- + +The audio files should be 16KHz monochannel wav files. + +Write audio files to a ``manifest.json`` file with lines as in format: + +.. code-block:: json + + {"audio_filepath": "/audio_file.wav", "duration": "duration of file in sec", "label": "speaker_id"} + +This python call will download best pretrained model from NGC and writes embeddings pickle file to current working directory + +.. code-block:: bash + + python examples/speaker_recognition/extract_speaker_embeddings.py --manifest=manifest.json + + +NGC Pretrained Checkpoints +-------------------------- + +The SpeakerNet-ASR collection has checkpoints of several models trained on various datasets for a variety of tasks. +`Speaker_Recognition `_ and `Speaker_Verification `_ model cards on NGC contain more information about each of the checkpoints available. + +The tables below list the SpeakerNet models available from NGC, and the models can be accessed via the +:code:`from_pretrained()` method inside the EncDecSpeakerLabelModel Model class. + +.. note:: While loading, remember to use EncDecSpeakerLabelModel for recognition tasks and ExtractSpeakerEmbeddingsModel while extracting embeddings. + +In general, you can load any of these models with code in the following format. + +.. code-block:: python + + import nemo.collections.asr as nemo_asr + model = nemo_asr.models..from_pretrained(model_name="") + +Where the model name is the value under "Model Name" entry in the tables below. + +If you would like to programatically list the models available for a particular base class, you can use the +:code:`list_available_models()` method. + +.. code-block:: python + + nemo_asr.models..list_available_models() + + +Speaker Recognition Models +^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. csv-table:: + :file: data/speaker_results.csv + :align: left + :widths: 30, 30, 40 + :header-rows: 1 + diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/configs.rst b/models/tts/NeMo/docs/source/asr/speech_classification/configs.rst new file mode 100644 index 0000000000..f0b17c5651 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speech_classification/configs.rst @@ -0,0 +1,104 @@ +NeMo Speech Classification Configuration Files +================================================ + +This page covers NeMo configuration file setup that is specific to models in the Speech Classification collection. +For general information about how to set up and run experiments that is common to all NeMo models (e.g. +experiment manager and PyTorch Lightning trainer parameters), see the :doc:`../../core/core` page. + +The model section of NeMo Speech Classification configuration files will generally require information about the dataset(s) being +used, the preprocessor for audio files, parameters for any augmentation being performed, as well as the +model architecture specification. +The sections on this page cover each of these in more detail. + +Example configuration files for all of the NeMo ASR scripts can be found in the +``/examples/asr/conf>``. + + +Dataset Configuration +--------------------- + +Training, validation, and test parameters are specified using the ``train_ds``, ``validation_ds``, and +``test_ds`` sections of your configuration file, respectively. +Depending on the task, you may have arguments specifying the sample rate of your audio files, labels, whether or not to shuffle the dataset, and so on. +You may also decide to leave fields such as the ``manifest_filepath`` blank, to be specified via the command line +at runtime. + +Any initialization parameters that are accepted for the Dataset class used in your experiment +can be set in the config file. +See the `Datasets <../api.html#Datasets>`__ section of the API for a list of Datasets and their respective parameters. + +An example Speech Classification train and validation configuration could look like: + +.. code-block:: yaml + + model: + sample_rate: 16000 + repeat: 2 # number of convolutional sub-blocks within a block, R in _[BxRxC] + dropout: 0.0 + kernel_size_factor: 1.0 + labels: ['bed', 'bird', 'cat', 'dog', 'down', 'eight', 'five', 'four', 'go', 'happy', 'house', 'left', 'marvin', + 'nine', 'no', 'off', 'on', 'one', 'right', 'seven', 'sheila', 'six', 'stop', 'three', 'tree', 'two', 'up', + 'wow', 'yes', 'zero'] + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} # Uses the labels above + batch_size: 128 + shuffle: True + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} # Uses the labels above + batch_size: 128 + shuffle: False # No need to shuffle the validation data + +If you would like to use tarred dataset, have a look at `Datasets Configuration <../configs.html#dataset-configuration>`__. + + +Preprocessor Configuration +-------------------------- +Preprocessor helps to compute MFCC or mel spectrogram features that are given as inputs to model. +For details on how to write this section, refer to `Preprocessor Configuration <../configs.html#preprocessor-configuration>`__ + +Check config yaml files in ``/examples/asr/conf`` to find the processors been used by speech classification models. + + +Augmentation Configurations +--------------------------- + +There are a few on-the-fly spectrogram augmentation options for NeMo ASR, which can be specified by the +configuration file using the ``augmentor`` and ``spec_augment`` section. +For details on how to write this section, refer to `Augmentation Configuration <../configs.html#augmentation-configurations>`__ + +Check config yaml files in ``/tutorials/asr/conf`` to find the processors been used by speech classification models. + + +Model Architecture Configurations +--------------------------------- + +Each configuration file should describe the model architecture being used for the experiment. +Models in the NeMo ASR collection need a ``encoder`` section and a ``decoder`` section, with the ``_target_`` field +specifying the module to use for each. + +The following sections go into more detail about the specific configurations of each model architecture. + +The `MatchboxNet <./models.html#matchboxnet-speech-commands>`__ and `MarbleNet <./models.html#marblenet-vad>`__ models are very similar, and they are based on `QuartzNet <../models.html#quartznet>`__ and as such the components in their +configs are very similar as well. + +Decoder Configurations +------------------------ + +After features have been computed from ConvASREncoder, we pass the features to decoder to compute embeddings and then to compute log_probs +for training models. + +.. code-block:: yaml + + model: + ... + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoderClassification + feat_in: *enc_final_filters + return_logits: true # return logits if true, else return softmax output + pooling_type: 'avg' # AdaptiveAvgPool1d 'avg' or AdaptiveMaxPool1d 'max' \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/data/classification_results.csv b/models/tts/NeMo/docs/source/asr/speech_classification/data/classification_results.csv new file mode 100644 index 0000000000..25e3c844a4 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speech_classification/data/classification_results.csv @@ -0,0 +1,9 @@ +Model Name,Model Base Class,Model Card +vad_marblenet,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:vad_marblenet" +vad_telephony_marblenet,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:vad_telephony_marblenet" +commandrecognition_en_matchboxnet3x1x64_v1,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x1x64_v1" +commandrecognition_en_matchboxnet3x2x64_v1,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x2x64_v1" +commandrecognition_en_matchboxnet3x1x64_v2,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x1x64_v2" +commandrecognition_en_matchboxnet3x2x64_v2,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x2x64_v2" +commandrecognition_en_matchboxnet3x1x64_v2_subset_task,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x1x64_v2_subset_task" +commandrecognition_en_matchboxnet3x2x64_v2_subset_task,EncDecClassificationModel,"https://ngc.nvidia.com/catalog/models/nvidia:nemo:commandrecognition_en_matchboxnet3x2x64_v2_subset_task" \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/datasets.rst b/models/tts/NeMo/docs/source/asr/speech_classification/datasets.rst new file mode 100644 index 0000000000..3198954a66 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speech_classification/datasets.rst @@ -0,0 +1,131 @@ +Datasets +======== + +NeMo has scripts to convert several common ASR datasets into the format expected by the `nemo_asr` collection. +You can get started with those datasets by following the instructions to run those scripts in the section appropriate to each dataset below. + +If you have your own data and want to preprocess it to use with NeMo ASR models, check out the `Preparing Custom Speech Classification Data`_ section at the bottom of the page. + +.. _Freesound-dataset: + +Freesound +----------- + +`Freesound `_ is a website that aims to create a huge open collaborative database of audio snippets, samples, recordings, bleeps. +Most audio samples are released under Creative Commons licenses that allow their reuse. +Researchers and developers can access Freesound content using the Freesound API to retrieve meaningful sound information such as metadata, analysis files, and the sounds themselves. + +**Instructions** + +Go to ``/scripts/freesound_download_resample`` and follow the below steps to download and convert freedsound data into a format expected by the `nemo_asr` collection. + +1. We will need some required libraries including freesound, requests, requests_oauthlib, joblib, librosa and sox. If they are not installed, please run `pip install -r freesound_requirements.txt` +2. Create an API key for freesound.org at https://freesound.org/help/developers/ +3. Create a python file called `freesound_private_apikey.py` and add lined `api_key = and client_id = ` +4. Authorize by run `python freesound_download.py --authorize` and visit the website and paste response code +5. Feel free to change any arguments in `download_resample_freesound.sh` such as max_samples and max_filesize +6. Run `bash download_resample_freesound.sh ` . For example: + +.. code-block:: bash + + bash download_resample_freesound.sh 4000 ./freesound ./freesound_resampled_background + +Note that downloading this dataset may take hours. Change categories in download_resample_freesound.sh to include other (speech) categories audio files. +Then, you should have 16khz mono wav files in ``. + + +.. _Google-Speech-Commands-Dataset: + +Google Speech Commands Dataset +------------------------------ + +Google released two versions of the dataset with the first version containing 65k samples over 30 classes and the second containing 110k samples over 35 classes. +We refer to these datasets as `v1` and `v2` respectively. + +Run the script `process_speech_commands_data.py` to process Google Speech Commands dataset in order to generate files in the supported format of `nemo_asr`, +which can be found in ``/scripts/dataset_processing/``. +You should set the data folder of Speech Commands using :code:`--data_root` and the version of the dataset using :code:`--data_version` as an int. + +You can further rebalance the train set by randomly oversampling files inside the manifest by passing the `--rebalance` flag. + +.. code-block:: bash + + python process_speech_commands_data.py --data_root= --data_version=<1 or 2> {--rebalance} + + +Then, you should have `train_manifest.json`, `validation_manifest.json` and `test_manifest.json` +in the directory `{data_root}/google_speech_recognition_v{1/2}`. + +.. note:: + You should have at least 4GB or 6GB of disk space available if you use v1 or v2 respectively. + Also, it will take some time to download and process, so go grab a coffee. + +Each line is a training example. + +.. code-block:: bash + + {"audio_filepath": "/two/8aa35b0c_nohash_0.wav", "duration": 1.0, "label": "two"} + {"audio_filepath": "/two/ec5ab5d5_nohash_2.wav", "duration": 1.0, "label": "two"} + + + +Speech Command & Freesound for VAD +------------------------------------ +Speech Command & Freesound (SCF) dataset is used to train MarbleNet in the `paper `_. Here we show how to download and process it. +This script assumes that you already have the Freesound dataset, if not, have a look at :ref:`Freesound-dataset`. +We will use the open-source :ref:`Google-Speech-Commands-Dataset` (we will use V2 of the dataset for SCF dataset, but require very minor changes to support V1 dataset) as our speech data. + +These scripts below will download the Google Speech Commands v2 dataset and convert speech and background data to a format suitable for use with nemo_asr. + +.. note:: + You may additionally pass :code:`--test_size` or :code:`--val_size` flag for splitting train val and test data. + + You may additionally pass :code:`--seg_len` flag for indicating the segment length. Default is 0.63s. + + You may additionally pass a :code:`-rebalance_method='fixed|over|under'` at the end of the script to rebalance the class samples in the manifest. + + + +* `'fixed'`: Fixed number of sample for each class. Train 5000, val 1000, and test 1000. (Change number in script if you want) +* `'over'`: Oversampling rebalance method +* `'under'`: Undersampling rebalance method + + +.. code-block:: bash + + mkdir './google_dataset_v2' + python process_vad_data.py --out_dir='./manifest/' --speech_data_root='./google_dataset_v2'--background_data_root= --log --rebalance_method='fixed' + + +After download and conversion, your `manifest` folder should contain a few json manifest files: + +* `(balanced_)background_testing_manifest.json` +* `(balanced_)background_training_manifest.json` +* `(balanced_)background_validation_manifest.json` +* `(balanced_)speech_testing_manifest.json` +* `(balanced_)speech_training_manifest.json` +* `(balanced_)speech_validation_manifest.json` + +Each line is a training example. `audio_filepath` contains path to the wav file, `duration` is duration in seconds, `offset` is offset in seconds, and `label` is label (class): + +.. code-block:: bash + + {"audio_filepath": "/two/8aa35b0c_nohash_0.wav", "duration": 0.63, "label": "speech", "offset": 0.0} + {"audio_filepath": "/Emergency_vehicle/id_58368 simambulance.wav", "duration": 0.63, "label": "background", "offset": 4.0} + + + +Preparing Custom Speech Classification Data +-------------------------------------------- + +Preparing Custom Speech Classification Data is almost identical to `Preparing Custom ASR Data <../datasets.html#preparing-custom-asr-data>`__. + +Instead of :code:`text` entry in manifest, you need :code:`label` to determine class of this sample + + +Tarred Datasets +--------------- + +Similarly to ASR, you can tar your audio files and use ASR Dataset class ``TarredAudioToClassificationLabelDataset`` (corresponding to the ``AudioToClassificationLabelDataset``) for this case. + +If you would like to use tarred dataset, have a look at `ASR Tarred Datasets <../datasets.html#tarred-datasets>`__. \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/images/marblenet_vertical.png b/models/tts/NeMo/docs/source/asr/speech_classification/images/marblenet_vertical.png new file mode 100644 index 0000000000..3174d1ef5e Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/speech_classification/images/marblenet_vertical.png differ diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/images/matchboxnet_vertical.png b/models/tts/NeMo/docs/source/asr/speech_classification/images/matchboxnet_vertical.png new file mode 100644 index 0000000000..2d69c94ecd Binary files /dev/null and b/models/tts/NeMo/docs/source/asr/speech_classification/images/matchboxnet_vertical.png differ diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/intro.rst b/models/tts/NeMo/docs/source/asr/speech_classification/intro.rst new file mode 100644 index 0000000000..6a1a2f9ee7 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speech_classification/intro.rst @@ -0,0 +1,47 @@ +Speech Classification +================================== +Speech Classification refers to a set of tasks or problems of getting a program to automatically classify input utterance or audio segment into categories, +such as Speech Command Recognition (multi-class), Voice Activity Detection (binary or multi-class), and Audio Sentiment Classification (typically multi-class), etc. + +**Speech Command Recognition** is the task of classifying an input audio pattern into a discrete set of classes. +It is a subset of Automatic Speech Recognition, sometimes referred to as Key Word Spotting, in which a model is constantly analyzing speech patterns to detect certain "command" classes. +Upon detection of these commands, a specific action can be taken by the system. +It is often the objective of command recognition models to be small and efficient so that they can be deployed onto low-power sensors and remain active for long durations of time. + + +**Voice Activity Detection (VAD)** also known as speech activity detection or speech detection, is the task of predicting which parts of input audio contain speech versus background noise. +It is an essential first step for a variety of speech-based applications including Automatic Speech Recognition. +It serves to determine which samples to be sent to the model and when to close the microphone. + + +The full documentation tree is as follows: + +.. toctree:: + :maxdepth: 8 + + models + datasets + results + configs + + +Resource and Documentation Guide +-------------------------------- + +Hands-on speech classification tutorial notebooks can be found under ``/tutorials/asr/``. +There are training and offline & online microphone inference tutorials for Speech Command Detection and Voice Activity Detection tasks. +This and most other tutorials can be run on Google Colab by specifying the link to the notebooks' GitHub pages on Colab. + +If you are looking for information about a particular Speech Classification model or would like to find out more about the model +architectures available in the `nemo_asr` collection, check out the :doc:`Models <./models>` page. + +Documentation on dataset preprocessing can be found on the :doc:`Datasets <./datasets>` page. +NeMo includes preprocessing scripts for several common ASR datasets, and this page contains instructions on running +those scripts. +It also includes guidance for creating your own NeMo-compatible dataset, if you have your own data. + +Information about how to load model checkpoints (either local files or pretrained ones from NGC), perform inference, as well as a list +of the checkpoints available on NGC are located on the :doc:`Checkpoints <./results>` page. + +Documentation for configuration files specific to the ``nemo_asr`` models can be found on the +:doc:`Configuration Files <./configs>` page. diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/models.rst b/models/tts/NeMo/docs/source/asr/speech_classification/models.rst new file mode 100644 index 0000000000..b8c0248ec5 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speech_classification/models.rst @@ -0,0 +1,71 @@ +Models +====== + +This page gives a brief overview of the models that NeMo's Speech Classification collection currently supports. +For Speech Classification, we support Speech Command (Keyword) Detection and Voice Activity Detection (VAD). + +Each of these models can be used with the example ASR scripts (in the ``/examples/asr`` directory) by +specifying the model architecture in the config file used. +Examples of config files for each model can be found in the ``/examples/asr/conf`` directory. + +For more information about the config files and how they should be structured, see the :doc:`./configs` page. + +Pretrained checkpoints for all of these models, as well as instructions on how to load them, can be found on the :doc:`./results` page. +You can use the available checkpoints for immediate inference, or fine-tune them on your own datasets. +The Checkpoints page also contains benchmark results for the available ASR models. + +.. _MatchboxNet_model: + +MatchboxNet (Speech Commands) +------------------------------ + +MatchboxNet :cite:`sc-models-matchboxnet` is an end-to-end neural network for speech command recognition based on `QuartzNet <../models.html#QuartzNet>`__. + +Similarly to QuartzNet, the MatchboxNet family of models are denoted as MatchBoxNet_[BxRxC] where B is the number of blocks, and R is the number of convolutional sub-blocks within a block, and C is the number of channels. Each sub-block contains a 1-D *separable* convolution, batch normalization, ReLU, and dropout: + + .. image:: images/matchboxnet_vertical.png + :align: center + :alt: MatchboxNet model + :scale: 50% + +It can reach state-of-the art accuracy on the Google Speech Commands dataset while having significantly fewer parameters than similar models. +The `_v1` and `_v2` are denoted for models trained on `v1` (30-way classification) and `v2` (35-way classification) datasets; +And we use _subset_task to represent (10+2)-way subset (10 specific classes + other remaining classes + silence) classification task. + +MatchboxNet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecClassificationModel` class. + +.. note:: + For model details and deep understanding about Speech Command Detedction training, inference, finetuning and etc., + please refer to ``/tutorials/asr/03_Speech_Commands.ipynb`` and ``/tutorials/asr/04_Online_Offline_Speech_Commands_Demo.ipynb``. + + + +.. _MarbleNet_model: + +MarbleNet (VAD) +------------------ + +MarbleNet :cite:`sc-models-marblenet` an end-to-end neural network for speech command recognition based on :ref:`MatchboxNet_model`, + +Similarly to MatchboxNet, the MarbleNet family of models are denoted as MarbleNet_[BxRxC] where B is the number of blocks, and R is the number of convolutional sub-blocks within a block, and C is the number of channels. Each sub-block contains a 1-D *separable* convolution, batch normalization, ReLU, and dropout: + + .. image:: images/marblenet_vertical.png + :align: center + :alt: MarbleNet model + :scale: 30% + +It can reach state-of-the art performance on the difficult `AVA speech dataset `_ while having significantly fewer parameters than similar models even training on simple data. +MarbleNet models can be instantiated using the :class:`~nemo.collections.asr.models.EncDecClassificationModel` class. + +.. note:: + For model details and deep understanding about VAD training, inference, postprocessing, threshold tuning and etc., + please refer to ``/tutorials/asr/06_Voice_Activiy_Detection.ipynb`` and ``/tutorials/asr/07_Online_Offline_Microphone_VAD_Demo.ipynb``. + + +References +---------------- + +.. bibliography:: ../asr_all.bib + :style: plain + :labelprefix: SC-MODELS + :keyprefix: sc-models- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/asr/speech_classification/results.rst b/models/tts/NeMo/docs/source/asr/speech_classification/results.rst new file mode 100644 index 0000000000..6dd65b9906 --- /dev/null +++ b/models/tts/NeMo/docs/source/asr/speech_classification/results.rst @@ -0,0 +1,105 @@ +Checkpoints +=========== + +There are two main ways to load pretrained checkpoints in NeMo: + +* Using the :code:`restore_from()` method to load a local checkpoint file (`.nemo`), or +* Using the :code:`from_pretrained()` method to download and set up a checkpoint from NGC. + +See the following sections for instructions and examples for each. + +Note that these instructions are for loading fully trained checkpoints for evaluation or fine-tuning. +For resuming an unfinished training experiment, please use the experiment manager to do so by setting the +``resume_if_exists`` flag to True. + +Loading Local Checkpoints +------------------------- + +NeMo will automatically save checkpoints of a model you are training in a `.nemo` format. +You can also manually save your models at any point using :code:`model.save_to(.nemo)`. + +If you have a local ``.nemo`` checkpoint that you'd like to load, simply use the :code:`restore_from()` method: + +.. code-block:: python + + import nemo.collections.asr as nemo_asr + model = nemo_asr.models..restore_from(restore_path="") + +Where the model base class is the ASR model class of the original checkpoint, or the general `ASRModel` class. + + +Transcribing/Inference +----------------------- + +The audio files should be 16KHz monochannel wav files. + +**Transcribe speech command segment:** + +You may perform inference and transcribe a sample of speech after loading the model by using its 'transcribe()' method: + +.. code-block:: python + + mbn_model = nemo_asr.models.EncDecClassificationModel.from_pretrained(model_name="") + mbn_model.transcribe([list of audio files], batch_size=BATCH_SIZE, logprobs=False) + +Setting argument ``logprobs`` to True would return the log probabilities instead of transcriptions. You may find more details in `Modules <../api.html#modules>`__. + +Learn how to fine tune on your own data or on subset classes in ``/tutorials/asr/03_Speech_Commands.ipynb`` +f + +**Run VAD inference:** + +.. code-block:: bash + + python examples/speaker_recognition/vad_infer.py --vad_model="vad_marblenet" --dataset= --out_dir='frame/demo' --time_length=0.63 + +Have a look at scripts under ``/scripts/voice_activity_detection`` for posterior processing and threshold tuning. + + +NGC Pretrained Checkpoints +-------------------------- + +The Speech Classification collection has checkpoints of several models trained on various datasets for a variety of tasks. +These checkpoints are obtainable via NGC `NeMo Automatic Speech Recognition collection `_. +The model cards on NGC contain more information about each of the checkpoints available. + +The tables below list the Speech Classification models available from NGC, and the models can be accessed via the +:code:`from_pretrained()` method inside the ASR Model class. + +In general, you can load any of these models with code in the following format. + +.. code-block:: python + + import nemo.collections.asr as nemo_asr + model = nemo_asr.models.EncDecClassificationModel.from_pretrained(model_name="") + +Where the model name is the value under "Model Name" entry in the tables below. + +For example, to load the MatchboxNet3x2x64_v1 model for speech command detection, run: + +.. code-block:: python + + model = nemo_asr.models.EncDecClassificationModel.from_pretrained(model_name="commandrecognition_en_matchboxnet3x2x64_v1") + +You can also call :code:`from_pretrained()` from the specific model class (such as :code:`EncDecClassificationModel` +for MatchboxNet and MarbleNet) if you will need to access specific model functionality. + +If you would like to programatically list the models available for a particular base class, you can use the +:code:`list_available_models()` method. + +.. code-block:: python + + nemo_asr.models..list_available_models() + + +Speech Classification Models +^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +.. tabularcolumns:: 30 30 40 + +.. csv-table:: + :file: data/classification_results.csv + :header-rows: 1 + :class: longtable + :widths: 1 1 1 + diff --git a/models/tts/NeMo/docs/source/common/intro.rst b/models/tts/NeMo/docs/source/common/intro.rst new file mode 100644 index 0000000000..61d160b5c6 --- /dev/null +++ b/models/tts/NeMo/docs/source/common/intro.rst @@ -0,0 +1,39 @@ +Common Collection +================= + +The common collection contains things that could be used across all collections. + +Tokenizers +---------- +.. automodule:: nemo.collections.common.tokenizers.AutoTokenizer + :special-members: __init__ +.. automodule:: nemo.collections.common.tokenizers.SentencePieceTokenizer + :special-members: __init__ +.. automodule:: nemo.collections.common.tokenizers.TokenizerSpec + :special-members: __init__ + + +Losses +------ +.. automodule:: nemo.collections.common.losses.AggregatorLoss + :special-members: __init__ + +.. automodule:: nemo.collections.common.losses.CrossEntropyLoss + :special-members: __init__ + +.. automodule:: nemo.collections.common.losses.MSELoss + :special-members: __init__ + +.. automodule:: nemo.collections.common.losses.SmoothedCrossEntropyLoss + :special-members: __init__ +.. automodule:: nemo.collections.common.losses.SpanningLoss + :special-members: __init__ + + +Metrics +------- + +.. autoclass:: nemo.collections.common.metrics.Perplexity + :show-inheritance: + :members: + :undoc-members: diff --git a/models/tts/NeMo/docs/source/conf.py b/models/tts/NeMo/docs/source/conf.py new file mode 100644 index 0000000000..7dfba6e3e4 --- /dev/null +++ b/models/tts/NeMo/docs/source/conf.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import re +import sys +import glob + +import sphinx_rtd_theme + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. + +sys.path.insert(0, os.path.abspath("../..")) +sys.path.insert(0, os.path.abspath("../../nemo")) +sys.path.insert(0, os.path.abspath("../../nemo_text_processing")) + +from package_info import __version__ + + +autodoc_mock_imports = [ + 'torch', + 'torch.nn', + 'torch.utils', + 'torch.optim', + 'torch.utils.data', + 'torch.utils.data.sampler', + 'torchvision', + 'torchvision.models', + 'torchtext', + 'ruamel.yaml', # ruamel.yaml has ., which is troublesome for this regex + 'hydra', # hydra-core in requirements, hydra during import + 'dateutil', # part of core python + 'transformers.tokenization_bert', # has ., troublesome for this regex + 'megatron', # megatron-lm in requirements, megatron in import + 'sklearn', + 'nemo_text_processing.inverse_text_normalization', # Not installed automatically + 'nemo_text_processing.text_normalization', # Not installed automatically + 'attr', # attrdict in requirements, attr in import + 'torchmetrics', # inherited from PTL +] + +_skipped_autodoc_mock_imports = ['wrapt', 'numpy'] + +for req_path in sorted(list(glob.glob("../../requirements/*.txt"))): + if "docs.txt" in req_path: + continue + + req_file = os.path.abspath(os.path.expanduser(req_path)) + with open(req_file, 'r') as f: + for line in f: + line = line.replace("\n", "") + req = re.search(r"([a-zA-Z0-9-_]*)", line) + if req: + req = req.group(1) + req = req.replace("-", "_") + + if req not in autodoc_mock_imports: + if req in _skipped_autodoc_mock_imports: + print(f"Skipping req : `{req}` (lib {line})") + continue + + autodoc_mock_imports.append(req) + print(f"Adding req : `{req}` to autodoc mock requirements (lib {line})") + else: + print(f"`{req}` already added to autodoc mock requirements (lib {line})") + +# +# -- General configuration ------------------------------------------------ + +# If your documentation needs a minimal Sphinx version, state it here. +# +# needs_sphinx = '1.0' + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. + +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.todo", + "sphinx.ext.coverage", + "sphinx.ext.mathjax", + "sphinx.ext.ifconfig", + "sphinx.ext.viewcode", + "sphinx.ext.napoleon", + "sphinx.ext.githubpages", + "sphinxcontrib.bibtex", + "sphinx.ext.inheritance_diagram", + "sphinx.ext.intersphinx", + "sphinx.ext.autosectionlabel", + "sphinxcontrib.bibtex", + "sphinx_rtd_theme", +] + +bibtex_bibfiles = [ + 'asr/asr_all.bib', + 'nlp/nlp_all.bib', + 'tools/tools_all.bib', + 'nemo_text_processing/textprocessing_all.bib', + 'tts_all.bib', +] + +intersphinx_mapping = { + 'pytorch': ('https://pytorch.org/docs/stable', None), + 'pytorch-lightning': ('https://pytorch-lightning.readthedocs.io/en/latest/', None), +} + +# Set default flags for all classes. +autodoc_default_options = {'members': None, 'undoc-members': None, 'show-inheritance': True} + +locale_dirs = ['locale/'] # path is example but recommended. +gettext_compact = False # optional. + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +# +# source_suffix = ['.rst', '.md'] +source_suffix = ".rst" + +# The master toctree document. +master_doc = "index" + +# General information about the project. +project = "NVIDIA NeMo" +copyright = "2021-, NVIDIA CORPORATION" +author = "NVIDIA CORPORATION" + +# The version info for the project you're documenting, acts as replacement for +# |version| and |release|, also used in various other places throughout the +# built documents. + + +# The short X.Y version. +# version = "0.10.0" +version = __version__ +# The full version, including alpha/beta/rc tags. +# release = "0.9.0" +release = __version__ + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = None + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This patterns also effect to html_static_path and html_extra_path +exclude_patterns = [] + +# The name of the Pygments (syntax highlighting) style to use. +pygments_style = "default" + +### Previous NeMo theme +# # NVIDIA theme settings. +# html_theme = 'nvidia_theme' + +# html_theme_path = ["."] + +# html_theme_options = { +# 'display_version': True, +# 'project_version': version, +# 'project_name': project, +# 'logo_path': None, +# 'logo_only': True, +# } +# html_title = 'Introduction' + +# html_logo = html_theme_options["logo_path"] + +# -- Options for HTMLHelp output ------------------------------------------ + +# Output file base name for HTML help builder. +htmlhelp_basename = "nemodoc" + +### from TLT conf.py +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# + +# html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] + +html_theme = "sphinx_rtd_theme" +html_logo = os.path.join('nv_logo.png') + +html_theme_options = { + 'logo_only': True, + 'display_version': True, + 'prev_next_buttons_location': 'bottom', + 'style_external_links': False, + 'style_nav_header_background': '#000000', + # Toc options + 'collapse_navigation': False, + 'sticky_navigation': False, + # 'navigation_depth': 10, + 'includehidden': False, + 'titles_only': False, +} + + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". + +html_favicon = 'favicon.ico' + +html_static_path = ['_static'] + +html_last_updated_fmt = '' + + +def setup(app): + app.add_css_file('css/custom.css') + app.add_js_file('js/pk_scripts.js') + + +# html_css_files = [ +# './custom.css', +# ] + +# html_js_files = [ +# './pk_scripts.js', +# ] diff --git a/models/tts/NeMo/docs/source/core/core.rst b/models/tts/NeMo/docs/source/core/core.rst new file mode 100644 index 0000000000..1e053bd4dd --- /dev/null +++ b/models/tts/NeMo/docs/source/core/core.rst @@ -0,0 +1,963 @@ +NeMo Models +=========== + +Basics +------ + +NeMo models contain everything needed to train and reproduce Conversational AI models: + +- neural network architectures +- datasets/data loaders +- data preprocessing/postprocessing +- data augmentors +- optimizers and schedulers +- tokenizers +- language models + +NeMo uses `Hydra `_ for configuring both NeMo models and the PyTorch Lightning Trainer. + +.. note:: Every NeMo model has an example configuration file and training script that can be found `here `_. + +The end result of using NeMo, `Pytorch Lightning `_, and Hydra is that NeMo models all have the same look and feel and are also fully compatible with the PyTorch ecosystem. + +Pretrained +---------- + +NeMo comes with many pretrained models for each of our collections: ASR, NLP, and TTS. Every pretrained NeMo model can be downloaded +and used with the ``from_pretrained()`` method. + +As an example, we can instantiate QuartzNet with the following: + +.. code-block:: Python + + import nemo.collections.asr as nemo_asr + + model = nemo_asr.models.EncDecCTCModel.from_pretrained(model_name="QuartzNet15x5Base-En") + +To see all available pretrained models for a specific NeMo model, use the ``list_available_models()`` method. + +.. code-block:: Python + + nemo_asr.model.EncDecCTCModel.list_available_models() + +For detailed information on the available pretrained models, refer to the collections documentation: + +- :ref:`Automatic Speech Recognition (ASR)` +- :ref:`Natural Language Processing (NLP)` +- :ref:`Speech Synthesis (TTS)` + +Training +-------- + +NeMo leverages `PyTorch Lightning `_ for model training. PyTorch Lightning lets NeMo decouple the +conversational AI code from the PyTorch training code. This means that NeMo users can focus on their domain (ASR, NLP, TTS) and +build complex AI applications without having to rewrite boiler plate code for PyTorch training. + +When using PyTorch Lightning, NeMo users can automatically train with: + +- multi-GPU/multi-node +- mixed precision +- model checkpointing +- logging +- early stopping +- and more + +The two main aspects of the Lightning API are the `LightningModule `_ +and the `Trainer `_. + +PyTorch Lightning ``LightningModule`` +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +Every NeMo model is a ``LightningModule`` which is an ``nn.module``. This means that NeMo models are compatible with the PyTorch +ecosystem and can be plugged into existing PyTorch workflows. + +Creating a NeMo model is similar to any other PyTorch workflow. We start by initializing our model architecture, then define the forward pass: + +.. code-block:: python + + class TextClassificationModel(NLPModel, Exportable): + ... + def __init__(self, cfg: DictConfig, trainer: Trainer = None): + """Initializes the BERTTextClassifier model.""" + ... + super().__init__(cfg=cfg, trainer=trainer) + + # instantiate a BERT based encoder + self.bert_model = get_lm_model( + pretrained_model_name=cfg.language_model.pretrained_model_name, + config_file=cfg.language_model.config_file, + config_dict=cfg.language_model.config, + checkpoint_file=cfg.language_model.lm_checkpoint, + vocab_file=cfg.tokenizer.vocab_file, + ) + + # instantiate the FFN for classification + self.classifier = SequenceClassifier( + hidden_size=self.bert_model.config.hidden_size, + num_classes=cfg.dataset.num_classes, + num_layers=cfg.classifier_head.num_output_layers, + activation='relu', + log_softmax=False, + dropout=cfg.classifier_head.fc_dropout, + use_transformer_init=True, + idx_conditioned_on=0, + ) + +.. code-block:: python + + def forward(self, input_ids, token_type_ids, attention_mask): + """ + No special modification required for Lightning, define it as you normally would + in the `nn.Module` in vanilla PyTorch. + """ + hidden_states = self.bert_model( + input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=attention_mask + ) + logits = self.classifier(hidden_states=hidden_states) + return logits + +The ``LightningModule`` organizes PyTorch code so that across all NeMo models we have a similar look and feel. +For example, the training logic can be found in ``training_step``: + +.. code-block:: python + + def training_step(self, batch, batch_idx): + """ + Lightning calls this inside the training loop with the data from the training dataloader + passed in as `batch`. + """ + # forward pass + input_ids, input_type_ids, input_mask, labels = batch + logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask) + + train_loss = self.loss(logits=logits, labels=labels) + + lr = self._optimizer.param_groups[0]['lr'] + + self.log('train_loss', train_loss) + self.log('lr', lr, prog_bar=True) + + return { + 'loss': train_loss, + 'lr': lr, + } + +While validation logic can be found in ``validation_step``: + +.. code-block:: python + + def validation_step(self, batch, batch_idx): + """ + Lightning calls this inside the validation loop with the data from the validation dataloader + passed in as `batch`. + """ + if self.testing: + prefix = 'test' + else: + prefix = 'val' + + input_ids, input_type_ids, input_mask, labels = batch + logits = self.forward(input_ids=input_ids, token_type_ids=input_type_ids, attention_mask=input_mask) + + val_loss = self.loss(logits=logits, labels=labels) + + preds = torch.argmax(logits, axis=-1) + + tp, fn, fp, _ = self.classification_report(preds, labels) + + return {'val_loss': val_loss, 'tp': tp, 'fn': fn, 'fp': fp} + +PyTorch Lightning then handles all of the boiler plate code needed for training. Virtually any aspect of training can be customized +via PyTorch Lightning `hooks `_, +`Plugins `_, +`callbacks `_, or by overriding `methods `_. + +For more domain-specific information, see: + +- :ref:`Automatic Speech Recognition (ASR)` +- :ref:`Natural Language Processing (NLP)` +- :ref:`Speech Synthesis (TTS)` + +PyTorch Lightning Trainer +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Since every NeMo model is a ``LightningModule``, we can automatically take advantage of the PyTorch Lightning ``Trainer``. Every NeMo +`example `_ training script uses the ``Trainer`` object to fit the model. + +First, instantiate the model and trainer, then call ``.fit``: + +.. code-block:: python + + # We first instantiate the trainer based on the model configuration. + # See the model configuration documentation for details. + trainer = pl.Trainer(**cfg.trainer) + + # Then pass the model configuration and trainer object into the NeMo model + model = TextClassificationModel(cfg.model, trainer=trainer) + + # Now we can train with by calling .fit + trainer.fit(model) + + # Or we can run the test loop on test data by calling + trainer.test(model=model) + +All `trainer flags `_ can be set from from the +NeMo configuration. + + +Configuration +------------- + +Hydra is an open-source Python framework that simplifies configuration for complex applications that must bring together many different +software libraries. Conversational AI model training is a great example of such an application. To train a conversational AI model, we +must be able to configure: + +- neural network architectures +- training and optimization algorithms +- data pre/post processing +- data augmentation +- experiment logging/visualization +- model checkpointing + +For an introduction to using Hydra, refer to the `Hydra Tutorials `_. + +With Hydra, we can configure everything needed for NeMo with three interfaces: + +- Command Line (CLI) +- Configuration Files (YAML) +- Dataclasses (Python) + +YAML +~~~~ + +NeMo provides YAML configuration files for all of our `example `_ training scripts. +YAML files make it easy to experiment with different model and training configurations. + +Every NeMo example YAML has the same underlying configuration structure: + +- trainer +- exp_manager +- model + +Model configuration always contain ``train_ds``, ``validation_ds``, ``test_ds``, and ``optim``. Model architectures vary across +domains, therefore, refer to the ASR, NLP, and TTS Collections documentation for more detailed information on Model architecture configuration. + +A NeMo configuration file should look similar to the following: + +.. code-block:: yaml + + # PyTorch Lightning Trainer configuration + # any argument of the Trainer object can be set here + trainer: + gpus: 1 # number of gpus per node + num_nodes: 1 # number of nodes + max_epochs: 10 # how many training epochs to run + val_check_interval: 1.0 # run validation after every epoch + + # Experiment logging configuration + exp_manager: + exp_dir: /path/to/my/nemo/experiments + name: name_of_my_experiment + create_tensorboard_logger: True + create_wandb_logger: True + + # Model configuration + # model network architecture, train/val/test datasets, data augmentation, and optimization + model: + train_ds: + manifest_filepath: /path/to/my/train/manifest.json + batch_size: 256 + shuffle: True + validation_ds: + manifest_filepath: /path/to/my/validation/manifest.json + batch_size: 32 + shuffle: False + test_ds: + manifest_filepath: /path/to/my/test/manifest.json + batch_size: 32 + shuffle: False + optim: + name: novograd + lr: .01 + betas: [0.8, 0.5] + weight_decay: 0.001 + # network architecture can vary greatly depending on the domain + encoder: + ... + decoder: + ... + +More specific details about configuration files for each collection can be found on the following pages: + +:ref:`NeMo ASR Configuration Files` + +CLI +~~~ + +With NeMo and Hydra, every aspect of model training can be modified from the command-line. This is extremely helpful for running lots +of experiments on compute clusters or for quickly testing parameters while developing. + +All NeMo `examples `_ come with instructions on how to +run the training/inference script from the command-line (see `here `_ +for an example). + +With Hydra, arguments are set using the ``=`` operator: + +.. code-block:: bash + + python examples/asr/speech_to_text.py \ + model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \ + model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \ + trainer.gpus=2 \ + trainer.max_epochs=50 + +We can use the ``+`` operator to add arguments from the CLI: + +.. code-block:: bash + + python examples/asr/speech_to_text.py \ + model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \ + model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \ + trainer.gpus=2 \ + trainer.max_epochs=50 \ + +trainer.fast_dev_run=true + +We can use the ``~`` operator to remove configurations: + +.. code-block:: bash + + python examples/asr/speech_to_text.py \ + model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \ + model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \ + ~model.test_ds \ + trainer.gpus=2 \ + trainer.max_epochs=50 \ + +trainer.fast_dev_run=true + +We can specify configuration files using the ``--config-path`` and ``--config-name`` flags: + +.. code-block:: bash + + python examples/asr/speech_to_text.py \ + --config-path=conf/quartznet \ + --config-name=quartznet_15x5 \ + model.train_ds.manifest_filepath=/path/to/my/train/manifest.json \ + model.validation_ds.manifest_filepath=/path/to/my/validation/manifest.json \ + ~model.test_ds \ + trainer.gpus=2 \ + trainer.max_epochs=50 \ + +trainer.fast_dev_run=true + +Dataclasses +~~~~~~~~~~~ + +Dataclasses allow NeMo to ship model configurations as part of the NeMo library and also enables pure Python configuration of NeMo models. +With Hydra, dataclasses can be used to create `structured configs `_ for the conversational AI application. + +As an example, refer to the code block below for an *Attenion is All You Need* machine translation model. The model configuration can +be instantiated and modified like any Python `Dataclass `_. + +.. code-block:: Python + + from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import AAYNBaseConfig + + cfg = AAYNBaseConfig() + + # modify the number of layers in the encoder + cfg.encoder.num_layers = 8 + + # modify the training batch size + cfg.train_ds.tokens_in_batch = 8192 + +.. note:: Configuration with Hydra always has the following precedence CLI > YAML > Dataclass + +Optimization +------------ + +Optimizers and learning rate schedules are configurable across all NeMo models and have their own namespace. Here is a sample YAML +configuration for a Novograd optimizer with Cosine Annealing learning rate schedule. + +.. code-block:: yaml + + optim: + name: novograd + lr: 0.01 + + # optimizer arguments + betas: [0.8, 0.25] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # Optional arguments + max_steps: null # computed at runtime or explicitly set here + monitor: val_loss + reduce_on_plateau: false + + # scheduler config override + warmup_steps: 1000 + warmup_ratio: null + min_lr: 1e-9: + +.. note:: `NeMo Examples `_ has optimizer and scheduler configurations for +every NeMo model. + +Optimizers can be configured from the CLI as well: + +.. code-block:: bash + + python examples/asr/speech_to_text.py \ + --config-path=conf/quartznet \ + --config-name=quartznet_15x5 \ + ... + # train with the adam optimizer + model.optim=adam \ + # change the learning rate + model.optim.lr=.0004 \ + # modify betas + model.optim.betas=[.8, .5] + +Optimizers +~~~~~~~~~~ + +``name`` corresponds to the lowercase name of the optimizer. To view a list of available optimizers, run: + +.. code-block:: Python + + from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS + + for name, opt in AVAILABLE_OPTIMIZERS.items(): + print(f'name: {name}, opt: {opt}') + +.. code-block:: bash + + name: sgd opt: + name: adam opt: + name: adamw opt: + name: adadelta opt: + name: adamax opt: + name: adagrad opt: + name: rmsprop opt: + name: rprop opt: + name: novograd opt: + +Optimizer Params +~~~~~~~~~~~~~~~~ + +Optimizer params can vary between optimizers but the ``lr`` param is required for all optimizers. To see the available params for an +optimizer, we can look at its corresponding dataclass. + +.. code-block:: python + + from nemo.core.config.optimizers import NovogradParams + + print(NovogradParams()) + +.. code-block:: bash + + NovogradParams(lr='???', betas=(0.95, 0.98), eps=1e-08, weight_decay=0, grad_averaging=False, amsgrad=False, luc=False, luc_trust=0.001, luc_eps=1e-08) + +``'???'`` indicates that the lr argument is required. + +Register Optimizer +~~~~~~~~~~~~~~~~~~ + +To register a new optimizer to be used with NeMo, run: + +.. autofunction:: nemo.core.optim.optimizers.register_optimizer + +Learning Rate Schedulers +~~~~~~~~~~~~~~~~~~~~~~~~ + +Learning rate schedulers can be optionally configured under the ``optim.sched`` namespace. + +``name`` corresponds to the name of the learning rate schedule. To view a list of available schedulers, run: + +.. code-block:: Python + + from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS + + for name, opt in AVAILABLE_SCHEDULERS.items(): + print(f'name: {name}, schedule: {opt}') + +.. code-block:: bash + + name: WarmupPolicy, schedule: + name: WarmupHoldPolicy, schedule: + name: SquareAnnealing, schedule: + name: CosineAnnealing, schedule: + name: NoamAnnealing, schedule: + name: WarmupAnnealing, schedule: + name: InverseSquareRootAnnealing, schedule: + name: SquareRootAnnealing, schedule: + name: PolynomialDecayAnnealing, schedule: + name: PolynomialHoldDecayAnnealing, schedule: + name: StepLR, schedule: + name: ExponentialLR, schedule: + name: ReduceLROnPlateau, schedule: + name: CyclicLR, schedule: + +Scheduler Params +~~~~~~~~~~~~~~~~ + +To see the available params for a scheduler, we can look at its corresponding dataclass: + +.. code-block:: Python + + from nemo.core.config.schedulers import CosineAnnealingParams + + print(CosineAnnealingParams()) + +.. code-block:: bash + + CosineAnnealingParams(last_epoch=-1, warmup_steps=None, warmup_ratio=None, min_lr=0.0) + +Register scheduler +~~~~~~~~~~~~~~~~~~ + +To register a new scheduler to be used with NeMo, run: + +.. autofunction:: nemo.core.optim.lr_scheduler.register_scheduler + +Save and Restore +---------------- + +NeMo models all come with ``.save_to`` and ``.restore_from`` methods. + +Save +~~~~ + +To save a NeMo model, run: + +.. code-block:: Python + + model.save_to('/path/to/model.nemo') + +Everything needed to use the trained model is packaged and saved in the ``.nemo`` file. For example, in the NLP domain, ``.nemo`` files +include the necessary tokenizer models and/or vocabulary files, etc. + +.. note:: A ``.nemo`` file is simply an archive like any other ``.tar`` file. + +Restore +~~~~~~~ + +To restore a NeMo model, run: + +.. code-block:: Python + + model.restore_from('/path/to/model.nemo') + +When using the PyTorch Lightning Trainer, a PyTorch Lightning checkpoint is created. These are mainly used within NeMo to auto-resume +training. Since NeMo models are ``LightningModules``, the PyTorch Lightning method ``load_from_checkpoint`` is available. Note that +``load_from_checkpoint`` won't necessarily work out-of-the-box for all models as some models require more artifacts than just the +checkpoint to be restored. For these models, the user will have to override ``load_from_checkpoint`` if they want to use it. + +It's highly recommended to use ``restore_from`` to load NeMo models. + +Register Artifacts +------------------ + +Conversational AI models can be complicated to restore as more information is needed than just the checkpoint weights in order to use the model. +NeMo models can save additional artifacts in the .nemo file by calling ``.register_artifact``. +When restoring NeMo models using ``.restore_from`` or ``.from_pretrained``, any artifacts that were registered will be available automatically. + +As an example, consider an NLP model that requires a trained tokenizer model. +The tokenizer model file can be automatically added to the .nemo file with the following: + +.. code-block:: python + + self.encoder_tokenizer = get_nmt_tokenizer( + ... + tokenizer_model=self.register_artifact(config_path='encoder_tokenizer.tokenizer_model', + src='/path/to/tokenizer.model', + verify_src_exists=True), + ) + +By default, ``.register_artifact`` will always return a path. If the model is being restored from a .nemo file, +then that path will be to the artifact in the .nemo file. Otherwise, ``.register_artifact`` will return the local path specified by the user. + +``config_path`` is the artifact key. It usually corresponds to a model configuration but does not have to. +The model config that is packaged with the .nemo file will be updated according to the ``config_path`` key. +In the above example, the model config will have + +.. code-block:: YAML + + encoder_tokenizer: + ... + tokenizer_model: nemo:4978b28103264263a03439aaa6560e5e_tokenizer.model + +``src`` is the path to the artifact and the base-name of the path will be used when packaging the artifact in the .nemo file. +Each artifact will have a hash prepended to the basename of ``src`` in the .nemo file. This is to prevent collisions with basenames +base-names that are identical (say when there are two or more tokenizers, both called `tokenizer.model`). +The resulting .nemo file will then have the following file: + +.. code-block:: bash + + 4978b28103264263a03439aaa6560e5e_tokenizer.model + +If ``verify_src_exists`` is set to ``False``, then the artifact is optional. This means that ``.register_artifact`` will return ``None`` +if the ``src`` cannot be found. + + + +Experiment Manager +================== + +NeMo's Experiment Manager leverages PyTorch Lightning for model checkpointing, TensorBoard Logging, and Weights and Biases logging. The +Experiment Manager is included by default in all NeMo example scripts. + +To use the experiment manager simply call :class:`~nemo.utils.exp_manager.exp_manager` and pass in the PyTorch Lightning ``Trainer``. + +.. code-block:: python + + exp_manager(trainer, cfg.get("exp_manager", None)) + +And is configurable via YAML with Hydra. + +.. code-block:: bash + + exp_manager: + exp_dir: /path/to/my/experiments + name: my_experiment_name + create_tensorboard_logger: True + create_checkpoint_callback: True + +Optionally, launch TensorBoard to view the training results in ``./nemo_experiments`` (by default). + +.. code-block:: bash + + tensorboard --bind_all --logdir nemo_experiments + +.. + +If ``create_checkpoint_callback`` is set to ``True``, then NeMo automatically creates checkpoints during training +using PyTorch Lightning's `ModelCheckpoint `_. +We can configure the ``ModelCheckpoint`` via YAML or CLI. + +.. code-block:: yaml + + exp_manager: + ... + # configure the PyTorch Lightning ModelCheckpoint using checkpoint_call_back_params + # any ModelCheckpoint argument can be set here + + # save the best checkpoints based on this metric + checkpoint_callback_params.monitor=val_loss + + # choose how many total checkpoints to save + checkpoint_callback_params.save_top_k=5 + +We can auto-resume training as well by configuring the ``exp_manager``. Being able to auto-resume is important when doing long training +runs that are premptible or may be shut down before the training procedure has completed. To auto-resume training, set the following +via YAML or CLI: + +.. code-block:: yaml + + exp_manager: + ... + # resume training if checkpoints already exist + resume_if_exists: True + + # to start training with no existing checkpoints + resume_ignore_no_checkpoint: True + + # by default experiments will be versioned by datetime + # we can set our own version with + exp_manager.version: my_experiment_version + +Neural Modules +============== + +NeMo is built around Neural Modules, conceptual blocks of neural networks that take typed inputs and produce typed outputs. Such +modules typically represent data layers, encoders, decoders, language models, loss functions, or methods of combining activations. +NeMo makes it easy to combine and re-use these building blocks while providing a level of semantic correctness checking via its neural +type system. + +.. note:: *All Neural Modules inherit from ``torch.nn.Module`` and are therefore compatible with the PyTorch ecosystem.* + +There are 3 types on Neural Modules: + + - Regular modules + - Dataset/IterableDataset + - Losses + +Every Neural Module in NeMo must inherit from `nemo.core.classes.module.NeuralModule` class. + +.. autoclass:: nemo.core.classes.module.NeuralModule + +Every Neural Modules inherits the ``nemo.core.classes.common.Typing`` interface and needs to define neural types for its inputs and outputs. +This is done by defining two properties: ``input_types`` and ``output_types``. Each property should return an ordered dictionary of +"port name"->"port neural type" pairs. Here is the example from :class:`~nemo.collections.asr.modules.ConvASREncoder` class: + +.. code-block:: python + + @property + def input_types(self): + return OrderedDict( + { + "audio_signal": NeuralType(('B', 'D', 'T'), SpectrogramType()), + "length": NeuralType(tuple('B'), LengthsType()), + } + ) + + @property + def output_types(self): + return OrderedDict( + { + "outputs": NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()), + "encoded_lengths": NeuralType(tuple('B'), LengthsType()), + } + ) + + @typecheck() + def forward(self, audio_signal, length=None): + ... + +The code snippet above means that ``nemo.collections.asr.modules.conv_asr.ConvASREncoder`` expects two arguments: + * First one, named ``audio_signal`` of shape ``[batch, dimension, time]`` with elements representing spectrogram values. + * Second one, named ``length`` of shape ``[batch]`` with elements representing lengths of corresponding signals. + +It also means that ``.forward(...)`` and ``__call__(...)`` methods each produce two outputs: + * First one, of shape ``[batch, dimension, time]`` but with elements representing encoded representation (``AcousticEncodedRepresentation`` class). + * Second one, of shape ``[batch]``, corresponding to their lengths. + +.. tip:: It is a good practice to define types and add ``@typecheck()`` decorator to your ``.forward()`` method after your module is ready for use by others. + +.. note:: The outputs of ``.forward(...)`` method will always be of type ``torch.Tensor`` or container of tensors and will work with any other Pytorch code. The type information is attached to every output tensor. +If tensors without types is passed to your module, it will not fail, however the types will not be checked. Thus, it is recommended to define input/output types for all your modules, starting with +data layers and add ``@typecheck()`` decorator to them. + +.. note:: To temporarily disable typechecking, you can enclose your code in ```with typecheck.disable_checks():``` statement. + +Neural Types +============ + +Motivation +---------- + +Neural Types describe the semantics, axis order, and dimensions of a tensor. The purpose of this type system is to catch semantic and +dimensionality errors during model creation and facilitate module re-use. + +.. image:: whyntypes.gif + :width: 900 + :alt: Neural Types Motivation + +``NeuralType`` class +-------------------- + +Neural Types perform semantic checks for modules and models inputs/outputs. They contain information about: + + - Semantics of what is stored in the tensors. For example, logits, logprobs, audiosignal, embeddings, etc. + - Axes layout, semantic and (optionally) dimensionality. For example: ``[Batch, Time, Channel]`` + +Types are implemented in ``nemo.core.neural_types.NeuralType`` class. When you instantiate an instance of this class, you +are expected to include both *axes* information and *element type* information. + +.. autoclass:: nemo.core.neural_types.NeuralType + +Type Comparison Results +----------------------- + +When comparing two neural types, the following comparison results are generated. + +.. autoclass:: nemo.core.neural_types.NeuralTypeComparisonResult + +Examples +-------- + +Long vs short notation +~~~~~~~~~~~~~~~~~~~~~~ + +NeMo's ``NeuralType`` class allows you to express axis semantics information in long and short form. Consider these two equivalent types. Both encoder 3 dimensional tensors and both contain elements of type ``AcousticEncodedRepresentation`` (this type is a typical output of ASR encoders). + +.. code-block:: python + + long_version = NeuralType( + axes=(AxisType(AxisKind.Batch, None), AxisType(AxisKind.Dimension, None), AxisType(AxisKind.Time, None)), + elements_type=AcousticEncodedRepresentation(), + ) + short_version = NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()) + assert long_version.compare(short_version) == NeuralTypeComparisonResult.SAME + +Transpose same +~~~~~~~~~~~~~~ + +Often it is useful to know if a simple transposition will solve type incompatibility. This is the case if the comparison result of two types equals ``nemo.core.neural_types.NeuralTypeComparisonResult.TRANSPOSE_SAME``. + +.. code-block:: python + + type1 = NeuralType(axes=('B', 'T', 'C')) + type2 = NeuralType(axes=('T', 'B', 'C')) + assert type1.compare(type2) == NeuralTypeComparisonResult.TRANSPOSE_SAME + assert type2.compare(type1) == NeuralTypeComparisonResult.TRANSPOSE_SAME + +Note that in this example, we dropped ``elements_type`` argument of ``NeuralType`` constructor. If not supplied, the element type is ``VoidType``. + +``VoidType`` for elements +~~~~~~~~~~~~~~~~~~~~~~~~~ + +Sometimes it is useful to express that elements' types don't matter but axes layout do. ``VoidType`` for elements can be used to express this. + +.. note:: ``VoidType`` is compatible with every other elements' type but not the other way around. See the following code snippet below for details. + +.. code-block:: python + + btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType()) + btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType()) + # Note the VoidType for elements here + btc_void = NeuralType(('B', 'T', 'C'), VoidType()) + + # This is true because VoidType is compatible with every other element type (SpectrogramType in this case) + # And axes layout between btc_void and btc_spctr is the same + assert btc_void.compare(btc_spctr) == NeuralTypeComparisonResult.SAME + # These two types are incompatible because even though VoidType is used for elements on one side, + # the axes layout is different + assert btc_void.compare(btc_spct_bad) == NeuralTypeComparisonResult.INCOMPATIBLE + # Note that even though VoidType is compatible with every other type, other types are not compatible with VoidType! + # It is one-way compatibility + assert btc_spctr.compare(btc_void) == NeuralTypeComparisonResult.INCOMPATIBLE + +Element type inheritance +~~~~~~~~~~~~~~~~~~~~~~~~ + +Neural types in NeMo support Python inheritance between element types. Consider an example where you want to develop a Neural Module which performs data augmentation for all kinds of spectrograms. +In ASR, two types of spectrograms are frequently used: mel and mfcc. To express this, we will create 3 classes to express +element's types: ``SpectrogramType``, ``MelSpectrogramType(SpectrogramType)``, ``MFCCSpectrogramType(SpectrogramType)``. + +.. code-block:: python + + input = NeuralType(('B', 'D', 'T'), SpectrogramType()) + out1 = NeuralType(('B', 'D', 'T'), MelSpectrogramType()) + out2 = NeuralType(('B', 'D', 'T'), MFCCSpectrogramType()) + + # MelSpectrogram and MFCCSpectrogram are not interchangeable. + assert out1.compare(out2) == NeuralTypeComparisonResult.INCOMPATIBLE + assert out2.compare(out1) == NeuralTypeComparisonResult.INCOMPATIBLE + # Type comparison detects that MFCC/MelSpectrogramType is a kind of SpectrogramType and can be accepted. + assert input.compare(out1) == NeuralTypeComparisonResult.GREATER + assert input.compare(out2) == NeuralTypeComparisonResult.GREATER + +Custom element types +~~~~~~~~~~~~~~~~~~~~ + +It is possible to create user-defined element types to express the semantics of elements in your tensors. To do so, the user will need to inherit and implement abstract methods of the ``nemo.core.neural_types.elements.ElementType`` class + +.. autoclass:: nemo.core.neural_types.elements.ElementType + +Note that element types can be parametrized. Consider this example where it distinguishes between audio sampled at 8Khz and 16Khz. + +.. code-block:: python + + audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000)) + audio8K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(8000)) + + assert audio8K.compare(audio16K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS + assert audio16K.compare(audio8K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS + +Enforcing dimensions +~~~~~~~~~~~~~~~~~~~~ + +In addition to specifying tensor layout and elements' semantics, neural types also allow you to enforce tensor dimensions. +The user will have to use long notations to specify dimensions. Short notations only allows you to specify axes semantics and assumes +arbitrary dimensions. + +.. code-block:: python + + type1 = NeuralType( + (AxisType(AxisKind.Batch, 64), AxisType(AxisKind.Time, 10), AxisType(AxisKind.Dimension, 128)), + SpectrogramType(), + ) + type2 = NeuralType(('B', 'T', 'C'), SpectrogramType()) + + # type2 will accept elements of type1 because their axes semantics match and type2 does not care about dimensions + assert type2.compare(type1), NeuralTypeComparisonResult.SAME + # type1 will not accept elements of type2 because it need dimensions to match strictly. + assert type1.compare(type2), NeuralTypeComparisonResult.DIM_INCOMPATIBLE + +Generic Axis kind +~~~~~~~~~~~~~~~~~ + +Sometimes (especially in the case of loss modules) it is useful to be able to specify a "generic" axis kind which will make it +compatible with any other kind of axis. This is easy to express with Neural Types by using ``nemo.core.neural_types.axes.AxisKind.Any`` for axes. + +.. code-block:: python + + type1 = NeuralType(('B', 'Any', 'Any'), SpectrogramType()) + type2 = NeuralType(('B', 'T', 'C'), SpectrogramType()) + type3 = NeuralType(('B', 'C', 'T'), SpectrogramType()) + + # type1 will accept elements of type2 and type3 because it only cares about element kind (SpectrogramType) + # number of axes (3) and that first one corresponds to batch + assert type1.compare(type2) == NeuralTypeComparisonResult.SAME + assert type1.compare(type3) == NeuralTypeComparisonResult.INCOMPATIBLE + +Container types +~~~~~~~~~~~~~~~ + +The NeMo-type system understands Python containers (lists). If your module returns a nested list of typed tensors, the way to express it is by +using Python list notation and Neural Types together when defining your input/output types. + +The example below shows how to express that your module returns single output ("out") which is list of lists of two dimensional tensors of shape ``[batch, dimension]`` containing logits. + +.. code-block:: python + + @property + def output_types(self): + return { + "out": [[NeuralType(('B', 'D'), LogitsType())]], + } + +Core APIs +========= + +Base class for all NeMo models +------------------------------ + +.. autoclass:: nemo.core.ModelPT + :show-inheritance: + :members: + :member-order: bysource + :undoc-members: cfg, num_weights + :exclude-members: set_eff_save, use_eff_save, teardown + +Base Neural Module class +------------------------ + +.. autoclass:: nemo.core.NeuralModule + :show-inheritance: + :members: + :member-order: bysource + +Neural Type classes +------------------- + +.. autoclass:: nemo.core.neural_types.NeuralType + :show-inheritance: + :members: + :member-order: bysource + +.. autoclass:: nemo.core.neural_types.axes.AxisType + :show-inheritance: + :members: + :member-order: bysource + +.. autoclass:: nemo.core.neural_types.elements.ElementType + :show-inheritance: + :members: + :member-order: bysource + +.. autoclass:: nemo.core.neural_types.comparison.NeuralTypeComparisonResult + :show-inheritance: + :members: + :member-order: bysource + +Experiment manager +------------------ + +.. autoclass:: nemo.utils.exp_manager.exp_manager + :show-inheritance: + :members: + :member-order: bysource + diff --git a/models/tts/NeMo/docs/source/core/whyntypes.gif b/models/tts/NeMo/docs/source/core/whyntypes.gif new file mode 100644 index 0000000000..56b5ab1543 Binary files /dev/null and b/models/tts/NeMo/docs/source/core/whyntypes.gif differ diff --git a/models/tts/NeMo/docs/source/favicon.ico b/models/tts/NeMo/docs/source/favicon.ico new file mode 100644 index 0000000000..424df87200 Binary files /dev/null and b/models/tts/NeMo/docs/source/favicon.ico differ diff --git a/models/tts/NeMo/docs/source/index.rst b/models/tts/NeMo/docs/source/index.rst new file mode 100644 index 0000000000..8bf17eb486 --- /dev/null +++ b/models/tts/NeMo/docs/source/index.rst @@ -0,0 +1,67 @@ +NVIDIA NeMo User Guide +====================== + +.. toctree:: + :maxdepth: 2 + :caption: Getting Started + :name: starthere + + starthere/intro + starthere/tutorials + starthere/best-practices + + +.. toctree:: + :maxdepth: 2 + :caption: NeMo Core + :name: core + + core/core + + +.. toctree:: + :maxdepth: 2 + :caption: Speech Processing + :name: Speech Processing + + asr/intro + asr/speech_classification/intro + asr/speaker_recognition/intro + asr/speaker_diarization/intro + +.. toctree:: + :maxdepth: 2 + :caption: Natural Language Processing + :name: Natural Language Processing + + nlp/intro + +.. toctree:: + :maxdepth: 2 + :caption: Text To Speech + :name: Text To Speech + :includehidden: + :titlesonly: + + tts/intro + +.. toctree:: + :maxdepth: 2 + :caption: Common + :name: Common + + common/intro + +.. toctree:: + :maxdepth: 2 + :caption: Text Processing + :name: Text Processing + + nemo_text_processing/intro + +.. toctree:: + :maxdepth: 2 + :caption: Tools + :name: Tools + + tools/intro diff --git a/models/tts/NeMo/docs/source/nemo_text_processing/api.rst b/models/tts/NeMo/docs/source/nemo_text_processing/api.rst new file mode 100755 index 0000000000..2634f15ade --- /dev/null +++ b/models/tts/NeMo/docs/source/nemo_text_processing/api.rst @@ -0,0 +1,35 @@ +NeMo Text Processing API +======================== + +Model Classes +------------- + +.. autoclass:: nemo_text_processing.text_normalization.GraphFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.text_normalization.ClassifyFst + :show-inheritance: + :members: + + +.. autoclass:: nemo_text_processing.text_normalization.VerbalizeFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.text_normalization.VerbalizeFinalFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.inverse_text_normalization.ClassifyFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.inverse_text_normalization.VerbalizeFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.inverse_text_normalization.VerbalizeFinalFst + :show-inheritance: + :members: + diff --git a/models/tts/NeMo/docs/source/nemo_text_processing/intro.rst b/models/tts/NeMo/docs/source/nemo_text_processing/intro.rst new file mode 100644 index 0000000000..c3be9a8964 --- /dev/null +++ b/models/tts/NeMo/docs/source/nemo_text_processing/intro.rst @@ -0,0 +1,17 @@ +Text Processing +=============== + +`nemo_text_processing` is a python package that is installed with the `nemo_toolkit`. + +See :doc:`NeMo Introduction <../starthere/intro>` for installation details. + +Additional requirements can be found in `setup.sh `_. + +.. toctree:: + :maxdepth: 1 + + text_normalization + inverse_text_normalization + api + + diff --git a/models/tts/NeMo/docs/source/nemo_text_processing/inverse_text_normalization.rst b/models/tts/NeMo/docs/source/nemo_text_processing/inverse_text_normalization.rst new file mode 100644 index 0000000000..0e828c78bf --- /dev/null +++ b/models/tts/NeMo/docs/source/nemo_text_processing/inverse_text_normalization.rst @@ -0,0 +1,90 @@ +Inverse Text Normalization +========================== + +Inverse text normalization (ITN) is a part of the Automatic Speech Recognition (ASR) post-processing pipeline. +ITN is the task of converting the raw spoken output of the ASR model into its written form to improve text readability. + +For example, +`"in nineteen seventy"` -> `"in 1975"` +and `"it costs one hundred and twenty three dollars"` -> `"it costs $123"`. + +NeMo ITN :cite:`textprocessing-itn-zhang2021nemo` is based on WFST-grammars :cite:`textprocessing-itn-Mohri2009`. We also provide a deployment route to C++ using `Sparrowhawk `_ :cite:`textprocessing-itn-sparrowhawk` -- an open-source version of Google Kestrel :cite:`textprocessing-itn-ebden2015kestrel`. +See :doc:`Text Procesing Deployment <../tools/text_processing_deployment>` for details. + +.. note:: + + For more details, see the tutorial `NeMo/tutorials/text_processing/Inverse_Text_Normalization.ipynb `__ in `Google's Colab `_. + + + + + + +Classes +---------------------------------- + + +The base class for every grammar is :class:`GraphFst`. +This tool is designed as a two-stage application: 1. `classification` of the input into semiotic tokens and 2. `verbalization` into written form. +For every stage and every semiotic token class there is a corresponding grammar, e.g. :class:`taggers.CardinalFst` +and :class:`verbalizers.CardinalFst`. +Together, they compose the final grammars :class:`ClassifyFst` and +:class:`VerbalizeFinalFst` that are compiled into WFST and used for inference. + + + + + + +.. autoclass:: nemo_text_processing.inverse_text_normalization.ClassifyFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.inverse_text_normalization.VerbalizeFinalFst + :show-inheritance: + :members: + + +Prediction +---------------------------------- + +Example prediction run: + +.. code:: + + python run_prediction.py --input= --output= [--verbose] + + +The input is expected to be lower-cased. Punctuation are outputted with separating spaces after semiotic tokens, e.g. `"i see, it is ten o'clock..."` -> `"I see, it is 10:00 . . ."`. +Inner-sentence white-space characters in the input are not maintained. + + + +Data Cleaning for Evaluation +---------------------------------- + +.. code:: + + python clean_eval_data.py --input= + + + +Evaluation +---------------------------------- + +Example evaluation run on (cleaned) `Google's text normalization dataset `__ :cite:`textprocessing-itn-sproat2016rnn`: + +.. code:: + + python run_evaluation.py --input=./en_with_types/output-00001-of-00100 [--cat CLASS_CATEGORY] [--filter] + + + + +References +---------- + +.. bibliography:: textprocessing_all.bib + :style: plain + :labelprefix: TEXTPROCESSING-ITN + :keyprefix: textprocessing-itn- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nemo_text_processing/text_normalization.rst b/models/tts/NeMo/docs/source/nemo_text_processing/text_normalization.rst new file mode 100644 index 0000000000..3f225cacd7 --- /dev/null +++ b/models/tts/NeMo/docs/source/nemo_text_processing/text_normalization.rst @@ -0,0 +1,78 @@ +Text Normalization +================== + +NeMo Text Normalization converts text from written form into its verbalized form. It is used as a preprocessing step before Text to Speech (TTS). It could also be used for preprocessing Automatic Speech Recognition (ASR) training transcripts. + + +For example, +`"at 10:00"` -> `"at ten o'clock"` +and `"it weighs 10kg."` -> `"it weights ten kilograms ."`. + + +NeMo Text Normalization :cite:`textprocessing-norm-zhang2021nemo` is based on WFST-grammars :cite:`textprocessing-norm-Mohri2009`. We also provide a deployment route to C++ using `Sparrowhawk `_ :cite:`textprocessing-norm-sparrowhawk` -- an open-source version of Google Kestrel :cite:`textprocessing-norm-ebden2015kestrel`. +See :doc:`Text Procesing Deployment <../tools/text_processing_deployment>` for details. + + +.. note:: + + For more details, see the tutorial `NeMo/tutorials/text_processing/Text_Normalization.ipynb `__ in `Google's Colab `_. + + + + + + +Classes +---------------------------------- + + +The base class for every grammar is :class:`GraphFst`. +This tool is designed as a two-stage application: 1. `classification` of the input into semiotic tokens and 2. `verbalization` into written form. +For every stage and every semiotic token class there is a corresponding grammar, e.g. :class:`taggers.CardinalFst` +and :class:`verbalizers.CardinalFst`. +Together, they compose the final grammars :class:`ClassifyFst` and +:class:`VerbalizeFinalFst` that are compiled into WFST and used for inference. + + + + +.. autoclass:: nemo_text_processing.text_normalization.ClassifyFst + :show-inheritance: + :members: + +.. autoclass:: nemo_text_processing.text_normalization.VerbalizeFinalFst + :show-inheritance: + :members: + + +Prediction +---------------------------------- + +Example prediction run: + +.. code:: + + python run_prediction.py <--input INPUT_TEXT_FILE> <--output OUTPUT_PATH> [--input_case INPUT_CASE] + +``INPUT_CASE`` specifies whether to treat the input as lower-cased or case sensitive. By default treat the input as cased since this is more informative, especially for abbreviations. Punctuation are outputted with separating spaces after semiotic tokens, e.g. `"I see, it is 10:00..."` -> `"I see, it is ten o'clock . . ."`. +Inner-sentence white-space characters in the input are not maintained. + + +Evaluation +---------------------------------- + +Example evaluation run on `Google's text normalization dataset `__ :cite:`textprocessing-norm-sproat2016rnn`: + +.. code:: + + python run_evaluation.py --input=./en_with_types/output-00001-of-00100 [--cat CLASS_CATEGORY] [--input_case INPUT_CASE] + + + +References +---------- + +.. bibliography:: textprocessing_all.bib + :style: plain + :labelprefix: TEXTPROCESSING-NORM + :keyprefix: textprocessing-norm- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nemo_text_processing/textprocessing_all.bib b/models/tts/NeMo/docs/source/nemo_text_processing/textprocessing_all.bib new file mode 100644 index 0000000000..63eb101cae --- /dev/null +++ b/models/tts/NeMo/docs/source/nemo_text_processing/textprocessing_all.bib @@ -0,0 +1,58 @@ +@article{ebden2015kestrel, + title={The Kestrel TTS text normalization system}, + author={Ebden, Peter and Sproat, Richard}, + journal={Natural Language Engineering}, + volume={21}, + number={3}, + pages={333}, + year={2015}, + publisher={Cambridge University Press} +} + + +@article{sproat2016rnn, + title={RNN approaches to text normalization: A challenge}, + author={Sproat, Richard and Jaitly, Navdeep}, + journal={arXiv preprint arXiv:1611.00068}, + year={2016} +} + +@book{taylor2009text, + title={Text-to-speech synthesis}, + author={Taylor, Paul}, + year={2009}, + publisher={Cambridge university press} +} + +@misc{zhang2021nemo, + title={NeMo Inverse Text Normalization: From Development To Production}, + author={Yang Zhang and Evelina Bakhturina and Kyle Gorman and Boris Ginsburg}, + year={2021}, + eprint={2104.05055}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} + +@Inbook{Mohri2009, +author="Mohri, Mehryar", +editor="Droste, Manfred +and Kuich, Werner +and Vogler, Heiko", +title="Weighted Automata Algorithms", +bookTitle="Handbook of Weighted Automata", +year="2009", +publisher="Springer Berlin Heidelberg", +address="Berlin, Heidelberg", +pages="213--254", +abstract="Weighted automata and transducers are widely used in modern applications in bioinformatics and text, speech, and image processing. This chapter describes several fundamental weighted automata and shortest-distance algorithms including composition, determinization, minimization, and synchronization, as well as single-source and all-pairs shortest distance algorithms over general semirings. It presents the pseudocode of these algorithms, gives an analysis of their running time complexity, and illustrates their use in some simple cases. Many other complex weighted automata and transducer algorithms used in practice can be obtained by combining these core algorithms.", +isbn="978-3-642-01492-5", +doi="10.1007/978-3-642-01492-5_6", +url="https://doi.org/10.1007/978-3-642-01492-5_6" +} + +@inproceedings{sparrowhawk, + title = {TTS for Low Resource Languages: A Bangla Synthesizer}, + author = {Alexander Gutkin and Linne Ha and Martin Jansche and Knot Pipatsrisawat and Richard Sproat}, + booktitle = {10th Language Resources and Evaluation Conference}, + year = {2016}, +} \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/api.rst b/models/tts/NeMo/docs/source/nlp/api.rst new file mode 100755 index 0000000000..18a81d76e3 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/api.rst @@ -0,0 +1,74 @@ +NeMo NLP collection API +======================= + +Model Classes +------------- + +.. autoclass:: nemo.collections.nlp.models.TextClassificationModel + :show-inheritance: + :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact, classifytext + +.. autoclass:: nemo.collections.nlp.models.GLUEModel + :show-inheritance: + :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + +.. autoclass:: nemo.collections.nlp.models.PunctuationCapitalizationModel + :show-inheritance: + :members: add_punctuation_capitalization, setup_training_data, setup_optimization, setup_validation_data, setup_test_data, multi_validation_epoch_end, register_artifact + +.. autoclass:: nemo.collections.nlp.models.TokenClassificationModel + :show-inheritance: + :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, register_artifact + +.. autoclass:: nemo.collections.nlp.models.QAModel + :show-inheritance: + :members: setup_training_data, setup_optimization, setup_validation_data, setup_test_data, inference, validation_epoch_end, test_epoch_end + +.. autoclass:: nemo.collections.nlp.models.BERTLMModel + :show-inheritance: + :members: setup_training_data, setup_optimization + +Modules +------- + +.. autoclass:: nemo.collections.nlp.modules.BertModule + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.MegatronBertEncoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.AlbertEncoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.BertEncoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.DistilBertEncoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.RobertaEncoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.SequenceClassifier + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.SequenceRegression + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.nlp.modules.SequenceTokenClassifier + :show-inheritance: + :members: + +.. autofunction:: nemo.collections.nlp.modules.get_lm_model + +.. autofunction:: nemo.collections.nlp.modules.get_pretrained_lm_models_list + +.. autofunction:: nemo.collections.nlp.modules.get_megatron_lm_models_list \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/bert_pretraining.rst b/models/tts/NeMo/docs/source/nlp/bert_pretraining.rst new file mode 100755 index 0000000000..eee23bd550 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/bert_pretraining.rst @@ -0,0 +1,133 @@ +.. _bert_pretraining: + +BERT +==== + +BERT is an autoencoding language model with a final loss composed of: + +- masked language model loss +- next sentence prediction + +The model architecture is published in `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding `__ :cite:`nlp-bert-devlin2018bert`. +The model is originally trained on English Wikipedia and BookCorpus. BERT is often used as a language model encoder for downstream tasks, for example, :ref:`token_classification`, :ref:`text_classification`, :ref:`question_answering`, etc. +Domain-specific BERT models can be advantageous for a wide range of applications. One notable application is the domain-specific BERT in a biomedical setting, +e.g. BioBERT :cite:`nlp-bert-lee2019biobert` or its improved derivative BioMegatron :cite:`nlp-bert-shin2020biomegatron`. For the latter, refer to :ref:`megatron_finetuning`. + +Quick Start Guide +----------------- + +.. code-block:: python + + from nemo.collections.nlp.models import BERTLMModel + + # to get the list of pre-trained models + BERTLMModel.list_available_models() + + # Download and load the pre-trained BERT-based model + model = BERTLMModel.from_pretrained("bertbaseuncased") + +Available Models +^^^^^^^^^^^^^^^^ + +.. list-table:: *Pretrained Models* + :widths: 5 10 + :header-rows: 1 + + * - Model + - Pretrained Checkpoint + * - BERT-base uncased + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:bertbaseuncased + * - BERT-large uncased + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:bertlargeuncased + +.. _dataset_bert_pretraining: + +Data Input for the BERT model +----------------------------- + +Data preprocessing can be either done on-the-fly during training or offline before training. The latter is optimized and recommended +for large text corpora. This was also used in the original paper to train the model on Wikipedia and BookCorpus. For on-the-fly data +processing, provide text files with sentences for training and validation, where words are separated by spaces, i.e.: ``[WORD] [SPACE] [WORD] [SPACE] [WORD]``. +To use this pipeline in training, use the dedicated configuration file `NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml`. + +To process data offline in advance, refer to the `BERT Quick Start Guide `__. +To recreate the original Wikipedia and BookCorpus datasets, follow steps 1-5 in the Quick Start Guide and run the script ``./data/create_datasets_from_start.sh`` inside the Docker container. +The ``downloaded`` folder should include two sub folders ``lower_case_[0,1]_seq_len_128_max_pred_20_masked_lm_prob_0.15_random_seed_12345_dupe_factor_5`` +and ``lower_case_[0,1]_seq_len_512_max_pred_80_masked_lm_prob_0.15_random_seed_12345_dupe_factor_5``, containing sequences of length 128 with a maximum of 20 masked tokens +and sequences of length 512 with a maximum of 80 masked tokens respectively. To use this pipeline in training, use the dedicated configuration file ``NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml`` +and specify the path to the created hd5f files. + + +Training the BERT model +----------------------- + +Example of model configuration for on-the-fly data preprocessing: `NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml `__. +Example of model configuration for offline data preprocessing: `NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml `__. + +The specification can be grouped into three categories: + +- Parameters that describe the training process: **trainer** +- Parameters that describe the datasets: **model.train_ds**, **model.validation_ds** +- Parameters that describe the model: **model**, **model.tokenizer**, **model.language_model** + +More details about parameters in the config file can be found below: + ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.only_mlm_loss** | bool | Only uses masked language model without next sentence prediction. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **train_ds.data_file** | string | Name of the text file or hdf5 data directory. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **train_ds.num_samples** | integer | Number of samples to use from the training dataset, ``-1`` - to use all. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +More details about parameters for offline data preprocessing can be found below: + ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **train_ds.max_predictions_per_seq** | integer | Maximum number of masked tokens in a sequence in the preprocessed data. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +More details about parameters for online data preprocessing can be found below: + ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.max_seq_length** | integer | The maximum total input sequence length after tokenization. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.mask_prob** | float | Probability of masking a token in the input text during data processing. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.short_seq_prob** | float | Probability of having a sequence shorter than the maximum sequence length. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +.. note:: + + For offline data preprocessing, **model.tokenizer** is null. For downstream task, use the same tokenizer that was used for + offline preprocessing. For online data preprocessing, **model.tokenizer** needs to be specified. See also :ref:`nlp_model` for + details. + +Example of the command for training the model: + +.. code:: + + python bert_pretraining.py \ + model.train_ds.data_file= \ + trainer.max_epochs= \ + trainer.gpus=[] + + +Fine-tuning on Downstream Tasks +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +To use a trained BERT model checkpoint on a NeMo NLP downstream task, e.g. :ref:`question_answering`, specify +:code:`model.language_model.lm_checkpoint=`. + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-BERT + :keyprefix: nlp-bert- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/glue_benchmark.rst b/models/tts/NeMo/docs/source/nlp/glue_benchmark.rst new file mode 100644 index 0000000000..37e0912520 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/glue_benchmark.rst @@ -0,0 +1,10 @@ +.. _glue_benchmark: + +GLUE Benchmark +============== + +We recommend you try the GLUE Benchmark model in a Jupyter notebook (can run on `Google's Colab `_): `NeMo/tutorials/nlp/GLUE_Benchmark.ipynb `__. + +Connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for the hardware accelerator). + +An example script on how to train the model can be found here: `NeMo/examples/nlp/glue_benchmark/glue_benchmark.py `__. diff --git a/models/tts/NeMo/docs/source/nlp/information_retrieval.rst b/models/tts/NeMo/docs/source/nlp/information_retrieval.rst new file mode 100644 index 0000000000..e4c047a40b --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/information_retrieval.rst @@ -0,0 +1,10 @@ +.. _information_retrieval: + +Information Retrieval +===================== + +We recommend you try the Information Retrieval model in a Jupyter notebook (can run on `Google's Colab `_): `NeMo/tutorials/nlp/Information_Retrieval_MSMARCO.ipynb `__. + +Connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for hardware the accelerator), + +An example script on how to train the model can be found here: `NeMo/examples/nlp/information_retrieval `__. diff --git a/models/tts/NeMo/docs/source/nlp/intro.rst b/models/tts/NeMo/docs/source/nlp/intro.rst new file mode 100755 index 0000000000..7703dd06b3 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/intro.rst @@ -0,0 +1,15 @@ +Natural Language Processing (NLP) +================================= + +.. toctree:: + :maxdepth: 2 + + models + megatron_finetuning + api + + +Example scripts for the NLP collection can be found under `NeMo/example/nlp/ `__. + +NLP notebook-tutorials are located under `NeMo/tutorials/nlp/ `__. Most NeMo +tutorials can be run on `Google's Colab `_.), with more details in the :ref:`tutorials` section. diff --git a/models/tts/NeMo/docs/source/nlp/joint_intent_slot.rst b/models/tts/NeMo/docs/source/nlp/joint_intent_slot.rst new file mode 100644 index 0000000000..5dc00430a2 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/joint_intent_slot.rst @@ -0,0 +1,244 @@ +.. _intent_slot: + +Joint Intent and Slot Classification +==================================== + +Joint Intent and Slot classification is a NLU task for classifying an intent and detecting all +relevant slots (Entities) for the intent in a query. For example, in the query ``What is the weather in Santa Clara tomorrow morning?``, +we would like to classify the query as a ``weather intent``, detect ``Santa Clara`` as a `location slot`, +and ``tomorrow morning`` as a ``date_time slot``. Intent and Slot names are usually task-specific and +defined as labels in the training data. This is a fundamental step that is executed in any +task-driven conversational assistant. + +Our BERT-based model implementation allows you to train and detect both of these tasks together. + +.. note:: + + We recommend you try the Joint Intent and Slot Classification model in a Jupyter notebook (can run on `Google's Colab `_.): `NeMo/tutorials/nlp/Joint_Intent_and_Slot_Classification.ipynb `__. + + Connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for the hardware accelerator). + + An example script on how to train the model can be found here: `NeMo/examples/nlp/intent_slot_classification `__. + + +NeMo Data Format +---------------- + +When training the model, the dataset should be first converted to the required data format, which requires the following files: + +- :code:`dict.intents.csv` - A list of all intent names in the data. One line per an intent name. The index of the intent line + (starting from ``0``) is used to identify the appropriate intent in ``train.tsv`` and ``test.tsv`` files. + +.. code:: + + weather + alarm + meeting + ... + +- :code:`dict.slots.csv` - A list of all slot names in the data. One line per slot name. The index of the slot line + (starting from ``0``) is used to identify the appropriate slots in the queries in ``train_slot.tsv`` and ``test_slot.tsv`` files. + In the last line of this dictionary ``O`` slot name is used to identify all ``out of scope`` slots, which are usually the majority of the tokens + in the queries. + +.. code:: + + date + time + city + ... + O + +- :code:`train.tsv/test.tsv` - A list of original queries, one per line, with the intent number + separated by a tab (e.g. "what alarms do i have set right now 0"). Intent numbers are + set according to the intent line in the intent dictionary file (:code:`dict.intents.csv`), + starting from ``0``. The first line in these files should contain the header line ``sentence + label``. + +- :code:`train_slot.tvs/test_slot.tsv` - A list that contains one line per query, when each word from the original text queries + is replaced by a token number from the slots dictionary file (``dict.slots.csv``), counted starting from ``0``. All the words + which do not contain a relevant slot are replaced by ``out-of scope`` token number, which is also a part of the slot dictionary file, + usually as the last entry there. For example a line from these files should look similar to: "54 0 0 54 54 12 12" (the numbers are + separated by a space). These files do not contain a header line. + + +Dataset Conversion +------------------ + +To convert to the format of the model data, use the ``import_datasets`` utility, which implements +the conversion for the Assistant dataset. Download the dataset `here `_ or you can +write your own converter for the format that you are using for data annotation. + +For a dataset that follows your own annotation format, we recommend using one text file for all +samples of the same intent, with the name of the file as the name of the intent. Use one line per +query, with brackets to define slot names. This is very similar to the assistant format, and you can +adapt this converter utility or your own format with small changes: + +:: + + did i set an alarm to [alarm_type : wake up] in the [timeofday : morning] + +Run the ``dataset_converter`` command: + +.. code:: + + python examples/nlp/intent_slot_classification/data/import_datasets.py + --source_data_dir=`source_data_dir` \ + --target_data_dir=`target_data_dir` \ + --dataset_name=['assistant'|'snips'|'atis'] + +- :code:`source_data_dir`: the directory location of the your dataset +- :code:`target_data_dir`: the directory location where the converted dataset should be saved +- :code:`dataset_name`: one of the implemented dataset names + +After conversion, ``target_data_dir`` should contain the following files: + +.. code:: + + . + |--target_data_dir + |-- dict.intents.csv + |-- dict.slots.csv + |-- train.tsv + |-- train_slots.tsv + |-- test.tsv + |-- test_slots.tsv + +Model Training +-------------- + +This is a pretrained BERT based model with 2 linear classifier heads on the top of it, one for classifying an intent of the query and +another for classifying slots for each token of the query. This model is trained with the combined loss function on the Intent and Slot +classification task on the given dataset. The model architecture is based on the paper `BERT for Joint Intent Classification and Slot Filling `__:cite:`nlp-jis-chen2019bert`. + +For each query, the model classifies it as one the intents from the intent dictionary and for each word of the query it will classify +it as one of the slots from the slot dictionary, including out of scope slot for all the remaining words in the query which does not +fall in another slot category. Out of scope slot (``O``) is a part of slot dictionary that the model is trained on. + +Example of model configuration file for training the model can be found at: `NeMo/examples/nlp/intent_slot_classification/conf/intent_slot_classification.yaml `__. +In the configuration file, define the parameters of the training and the model, although most of the default values will work well. + +The specification can be roughly grouped into three categories: + +- Parameters that describe the training process: **trainer** +- Parameters that describe the model: **model** +- Parameters that describe the datasets: **model.train_ds**, **model.validation_ds**, **model.test_ds**, + +More details about parameters in the spec file can be found below: + ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Default** | **Description** | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.data_dir** | string | -- | The path of the data converted to the specified format. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.class_balancing** | string | ``null`` | Choose from ``[null, weighted_loss]``. The ``weighted_los``s enables weighted class balancing of the loss. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.intent_loss_weight** | float | ``0.6`` | The elation of intent-to-slot loss in the total loss. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.pad_label** | integer | ``-1`` | A value to pad the inputs. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.ignore_extra_tokens** | boolean | ``false`` | A flag that specifies whether to ignore extra tokens. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.ignore_start_end** | boolean | ``true`` | A flag that specifies whether to not use the first and last token for slot training. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.head.num_output_layers** | integer | ``2`` | The number of fully connected layers of the classifier on top of the BERT model. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **model.head.fc_dropout** | float | ``0.1`` | The dropout ratio of the fully connected layers. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.prefix** | string | ``train`` | A prefix for the training file names. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.prefix** | string | ``dev`` | A prefix for the validation file names. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ +| **test_ds.prefix** | string | ``test`` | A prefix for the test file names. | ++-------------------------------------------+-----------------+----------------------------------------------------------------------------------+--------------------------------------------------------------------------------------------------------------+ + +For additional config parameters common to all NLP models, refer to the `nlp_model doc `__. + +The following is an example of the command for training the model: + +.. code:: + + python examples/nlp/intent_slot_classification/intent_slot_classification.py + model.data_dir= \ + trainer.max_epochs= \ + trainer.gpus=[] + + +Required Arguments for Training +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`model.data_dir`: the dataset directory + + +Optional Arguments +^^^^^^^^^^^^^^^^^^ + +Most of the default parameters in the existing configuration file are already set appropriately, however, there are some parameters +you may want to experiment with. + +- ``trainer.max_epochs``: the number of training epochs (reasonable to be between 10 to 100) +- ``model.class_balancing`` - value ``weighted_loss`` may help to train the model when there is unbalanced set of classes +- ``model.intent_loss_weight`` - a number between 0 to 1 that defines a weight of the intent lost versus a slot loss during training. A +default value 0.6 gives a slight preference for the intent lose optimization. + +Training Procedure +^^^^^^^^^^^^^^^^^^ + +At the start of evaluation, NeMo will print out a log of the experiment specification, a summary of the training dataset, and the +model architecture. + +As the model starts training, you should see a progress bar per epoch. During training, after each epoch, NeMo will display accuracy +metrics on the validation dataset for every intent and slot separately, as well as the total accuracy. You can expect these numbers +to grow up to 50-100 epochs, depending on the size of the trained data. Since this is a joint iIntent and slot training, usually +intent's accuracy will grow first for the initial 10-20 epochs, and after that, slot's accuracy will start improving as well. + +At the end of training, NeMo saves the best checkpoint on the validation dataset at the path specified by the experiment spec file +before finishing. + +.. code:: + + GPU available: True, used: True + TPU available: None, using: 0 TPU cores + LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0,1,2] + [NeMo W 2021-01-28 14:52:19 exp_manager:299] There was no checkpoint folder at checkpoint_dir :results/checkpoints. Training from scratch. + [NeMo I 2021-01-28 14:52:19 exp_manager:186] Experiments will be logged at results + ... + label precision recall f1 support + weather.weather (label_id: 0) 0.00 0.00 0.00 128 + weather.temperature (label_id: 1) 0.00 0.00 0.00 0 + weather.temperature_yes_no (label_id: 2) 0.00 0.00 0.00 0 + weather.rainfall (label_id: 3) 0.00 0.00 0.00 0 + weather.rainfall_yes_no (label_id: 4) 0.00 0.00 0.00 0 + weather.snow (label_id: 5) 0.00 0.00 0.00 0 + weather.snow_yes_no (label_id: 6) 0.00 0.00 0.00 0 + weather.humidity (label_id: 7) 0.00 0.00 0.00 0 + weather.humidity_yes_no (label_id: 8) 0.00 0.00 0.00 0 + weather.windspeed (label_id: 9) 0.00 0.00 0.00 0 + weather.sunny (label_id: 10) 0.00 0.00 0.00 0 + weather.cloudy (label_id: 11) 0.00 0.00 0.00 0 + weather.alert (label_id: 12) 0.00 0.00 0.00 0 + context.weather (label_id: 13) 0.00 0.00 0.00 0 + context.continue (label_id: 14) 0.00 0.00 0.00 0 + context.navigation (label_id: 15) 0.00 0.00 0.00 0 + context.rating (label_id: 16) 0.00 0.00 0.00 0 + context.distance (label_id: 17) 0.00 0.00 0.00 0 + ------------------- + micro avg 0.00 0.00 0.00 128 + macro avg 0.00 0.00 0.00 128 + weighted avg 0.00 0.00 0.00 128 + +Model Evaluation and Inference +------------------------------ + +There is no separate script for the evaluation and inference of this model in NeMo, however, inside of the example file `examples/nlp/intent_slot_classification/intent_slot_classification.py` +after the training part is finished, you can see the code that evaluates the trained model on an evaluation test set and then an example of doing inference using a list of given queries. + +For the deployment in the production environment, refer to `NVIDIA Jarvis `__ and `NVIDIA TLT documentation `__. + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-JIS + :keyprefix: nlp-jis- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/language_modeling.rst b/models/tts/NeMo/docs/source/nlp/language_modeling.rst new file mode 100644 index 0000000000..5565e5d191 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/language_modeling.rst @@ -0,0 +1,6 @@ +.. _language_modeling: + +Language Modeling +================= + +An example script on how to train the model can be found here: `NeMo/examples/nlp/language_modeling `__. diff --git a/models/tts/NeMo/docs/source/nlp/machine_translation.rst b/models/tts/NeMo/docs/source/nlp/machine_translation.rst new file mode 100644 index 0000000000..8f3ed7d53c --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/machine_translation.rst @@ -0,0 +1,466 @@ +.. _machine_translation: + +Machine Translation Models +========================== +Machine translation is the task of translating text from one language to another. For example, from English to Spanish. Models are +based on the Transformer sequence-to-sequence architecture :cite:`nlp-machine_translation-vaswani2017attention`. + +An example script on how to train the model can be found here: `NeMo/examples/nlp/machine_translation/enc_dec_nmt.py `__. +The default configuration file for the model can be found at: `NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml `__. + +Quick Start Guide +----------------- + +.. code-block:: python + + from nemo.collections.nlp.models import MTEncDecModel + + # To get the list of pre-trained models + MTEncDecModel.list_available_models() + + # Download and load the a pre-trained to translate from English to Spanish + model = MTEncDecModel.from_pretrained("nmt_en_es_transformer12x2") + + # Translate a sentence or list of sentences + translations = model.translate(["Hello!"], source_lang="en", target_lang="es") + +Available Models +^^^^^^^^^^^^^^^^ + +.. list-table:: *Pretrained Models* + :widths: 5 10 + :header-rows: 1 + + * - Model + - Pretrained Checkpoint + * - English -> German + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_de_transformer12x2 + * - German -> English + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_de_en_transformer12x2 + * - English -> Spanish + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_es_transformer12x2 + * - Spanish -> English + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_es_en_transformer12x2 + * - English -> French + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_fr_transformer12x2 + * - French -> English + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_fr_en_transformer12x2 + * - English -> Russian + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_ru_transformer6x6 + * - Russian -> English + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_ru_en_transformer6x6 + * - English -> Chinese + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_en_zh_transformer6x6 + * - Chinese -> English + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:nmt_zh_en_transformer6x6 + +Data Format +----------- + +Supervised machine translation models require parallel corpora which comprises many examples of sentences in a source language and +their corresponding translation in a target language. We use parallel data formatted as seperate text files for source and target +languages where sentences in corresponding files are aligned like in the table below. + +.. list-table:: *Parallel Coprus* + :widths: 10 10 + :header-rows: 1 + + * - train.english.txt + - train.spanish.txt + * - Hello . + - Hola . + * - Thank you . + - Gracias . + * - You can now translate from English to Spanish in NeMo . + - Ahora puedes traducir del inglés al español en NeMo . + +It is common practice to apply data cleaning, normalization, and tokenization to the data prior to training a translation model and +NeMo expects already cleaned, normalized, and tokenized data. The only data pre-processing NeMo does is subword tokenization with BPE +:cite:`nlp-machine_translation-sennrich2015neural`. + +Data Cleaning, Normalization & Tokenization +------------------------------------------- + +We recommend applying the following steps to clean, normalize, and tokenize your data. All pre-trained models released, apply these data pre-processing steps. + +#. Language ID filtering - This step filters out examples from your training dataset that aren't in the correct language. For example, + many datasets contain examples where source and target sentences are in the same language. You can use a pre-trained language ID + classifier from `fastText `__. You can then run our script using the + ``lid.176.bin`` model downloaded from the fastText website. + + .. code :: + + python NeMo/scripts/neural_machine_translation/filter_by_language.py \ + --input-src train.en \ + --input-tgt train.es \ + --output-src train_lang_filtered.en \ + --output-tgt train_lang_filtered.es \ + --source-lang en \ + --target-lang es \ + --removed-src train_noise.en \ + --removed-tgt train_noise.de \ + --fasttext-model lid.176.bin + +#. Length filtering - We filter out sentences from the data that are below a minimum length (1) or exceed a maximum length (250). We + also filter out sentences where the ratio between source and target lengths exceeds 1.3 except for English <-> Chinese models. + `Moses `__ is a statistical machine translation toolkit that contains many useful + pre-processing scripts. + + .. code :: + + perl mosesdecoder/scripts/training/clean-corpus-n.perl -ratio 1.3 train en es train.filter 1 250 + +#. Data cleaning - While language ID filtering can sometimes help with filtering out noisy sentences that contain too many puncutations, + it does not help in cases where the translations are potentially incorrect, disfluent, or incomplete. We use `bicleaner `__ + a tool to identify such sentences. It trains a classifier based on many features included pre-trained language model fluency, word + alignment scores from a word-alignment model like `Giza++ `__ etc. We use their available + pre-trained models wherever possible and train models ourselves using their framework for remaining languages. The following script + applies a pre-trained bicleaner model to the data and pick sentences that are clean with probability > 0.5. + + .. code :: + + awk '{print "-\t-"}' train.en \ + | paste -d "\t" - train.filter.en train.filter.es \ + | bicleaner-classify - - > train.en-es.bicleaner.score + +#. Data deduplication - We use `bifixer `__ (which uses xxHash) to hash the source and target + sentences based on which we remove duplicate entries from the file. You may want to do something similar to remove training examples + that are in the test dataset. + + .. code :: + + cat train.en-es.bicleaner.score \ + | parallel -j 25 --pipe -k -l 30000 python bifixer.py --ignore-segmentation -q - - en es \ + > train.en-es.bifixer.score + + awk -F awk -F "\t" '!seen[$6]++' train.en-es.bifixer.score > train.en-es.bifixer.dedup.score + +#. Filter out data that bifixer assigns probability < 0.5 to. + + .. code :: + + awk -F "\t" '{ if ($5>0.5) {print $3}}' train.en-es.bifixer.dedup.score > train.cleaned.en + awk -F "\t" '{ if ($5>0.5) {print $4}}' train.en-es.bifixer.dedup.score > train.cleaned.es + +#. Punctuation Normalization - Punctuation, especially things like quotes can be written in different ways. + It's often useful to normalize the way they appear in text. We use the moses punctuation normalizer on all languages except Chinese. + + .. code :: + + perl mosesdecoder/scripts/tokenizer/normalize-punctuation.perl -l es < train.cleaned.es > train.normalized.es + perl mosesdecoder/scripts/tokenizer/normalize-punctuation.perl -l en < train.cleaned.en > train.normalized.en + + For example: + + .. code :: + + Before - Aquí se encuentran joyerías como Tiffany`s entre negocios tradicionales suizos como la confitería Sprüngli. + After - Aquí se encuentran joyerías como Tiffany's entre negocios tradicionales suizos como la confitería Sprüngli. + +#. Tokenization and word segmentation for Chinese - Naturally written text often contains punctuation markers like commas, full-stops + and apostrophes that are attached to words. Tokenization by just splitting a string on spaces will result in separate token IDs for + very similar items like ``NeMo`` and ``NeMo.``. Tokenization splits punctuation from the word to create two separate tokens. In the + previous example ``NeMo.`` becomes ``NeMo .`` which when split by space, results in two tokens and adressess the earlier problem. + + For example: + + .. code :: + + Before - Especialmente porque se enfrentará "a Mathieu (Debuchy), Yohan (Cabaye) y Adil (Rami) ", recuerda. + After - Especialmente porque se enfrentará " a Mathieu ( Debuchy ) , Yohan ( Cabaye ) y Adil ( Rami ) " , recuerda . + + We use the Moses tokenizer for all languages except Chinese. + + .. code :: + + perl mosesdecoder/scripts/tokenizer/tokenizer.perl -l es -no-escape < train.normalized.es > train.tokenized.es + perl mosesdecoder/scripts/tokenizer/tokenizer.perl -l en -no-escape < train.normalized.en > train.tokenized.en + + For languages like Chinese where there is no explicit marker like spaces that separate words, we use `Jieba `__ to segment a string into words that are space separated. + + For example: + + .. code :: + + Before - 同时，卫生局认为有必要接种的其他人员，包括公共部门，卫生局将主动联络有关机构取得名单后由卫生中心安排接种。 + After - 同时 ， 卫生局 认为 有 必要 接种 的 其他 人员 ， 包括 公共部门 ， 卫生局 将 主动 联络 有关 机构 取得 名单 后 由 卫生 中心 安排 接种 。 + +Training a BPE Tokenization +--------------------------- + +Byte-pair encoding (BPE) :cite:`nlp-machine_translation-sennrich2015neural` is a sub-word tokenization algorithm that is commonly used +to reduce the large vocabulary size of datasets by splitting words into frequently occuring sub-words. Currently, mMachine translation +only supports the `YouTokenToMe `__ BPE tokenizer. One can set the tokenization configuration +as follows: + ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Default** | **Description** | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **model.{encoder_tokenizer,decoder_tokenizer}.tokenizer_name** | str | ``yttm`` | BPE library name. Only supports ``yttm`` for now. | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **model.{encoder_tokenizer,decoder_tokenizer}.tokenizer_model** | str | ``null`` | Path to an existing YTTM BPE model. If ``null``, will train one from scratch on the provided data. | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **model.{encoder_tokenizer,decoder_tokenizer}.vocab_size** | int | ``null`` | Desired vocabulary size after BPE tokenization. | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **model.{encoder_tokenizer,decoder_tokenizer}.bpe_dropout** | float | ``null`` | BPE dropout probability. :cite:`nlp-machine_translation-provilkov2019bpe`. | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **model.{encoder_tokenizer,decoder_tokenizer}.vocab_file** | str | ``null`` | Path to pre-computed vocab file if exists. | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ +| **model.shared_tokenizer** | bool | ``True`` | Whether to share the tokenizer between the encoder and decoder. | ++-----------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------+ + + +Applying BPE Tokenization, Batching, Bucketing and Padding +---------------------------------------------------------- + +Given BPE tokenizers, and a cleaned parallel corpus, the following steps are applied to create a `TranslationDataset `__ object. + +#. Text to IDs - This performs subword tokenization with the BPE model on an input string and maps it to a sequence of tokens for the + source and target text. + +#. Bucketing - Sentences vary in length and when creating minibatches, we'd like sentences in them to have roughly the same length to + minimize the number of ```` tokens and to maximize computational efficiency. This step groups sentences roughly the same length + into buckets. + +#. Batching and padding - Creates minibatches with a maximum number of tokens specified by ``model.{train_ds,validation_ds,test_ds}.tokens_in_batch`` + from buckets and pads, so they can be packed into a tensor. + +Datasets can be configured as follows: + ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Default** | **Description** | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.src_file_name** | str | ``null`` | Path to the source language file. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.tgt_file_name** | str | ``null`` | Path to the target language file. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.tokens_in_batch** | int | ``512`` | Maximum number of tokens per minibatch. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.clean** | bool | ``true`` | Whether to clean the dataset by discarding examples that are greater than ``max_seq_length``. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.max_seq_length** | int | ``512`` | Maximum sequence to be used with the ``clean`` argument above. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.shuffle** | bool | ``true`` | Whether to shuffle minibatches in the PyTorch DataLoader. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.num_samples** | int | ``-1`` | Number of samples to use. ``-1`` for the entire dataset. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.drop_last** | bool | ``false`` | Drop last minibatch if it is not of equal size to the others. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.pin_memory** | bool | ``false`` | Whether to pin memory in the PyTorch DataLoader. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.num_workers** | int | ``8`` | Number of workers for the PyTorch DataLoader. | ++-------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------------+ + + +Tarred Datasets for Large Corpora +--------------------------------- + +When training with ``DistributedDataParallel``, each process has its own copy of the dataset. For large datasets, this may not always +fit in CPU memory. `Webdatasets `__ circumvents this problem by efficiently iterating over +tar files stored on disk. Each tar file can contain hundreds to thousands of pickle files, each containing a single minibatch. + +We recommend using this method when working with datasets with > 1 million sentence pairs. + +Tarred datasets can be configured as follows: + ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Default** | **Description** | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.use_tarred_dataset** | bool | ``false`` | Whether to use tarred datasets. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.tar_files** | str | ``null`` | String specifying path to all tar files. Example with 100 tarfiles ``/path/to/tarfiles._OP_1..100_CL_.tar``. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.metadata_file** | str | ``null`` | Path to JSON metadata file that contains only a single entry for the total number of batches in the dataset. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.lines_per_dataset_fragment** | int | ``1000000`` | | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.num_batches_per_tarfile** | int | ``100`` | Maximum sequence to be used with the ``clean`` argument above. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.tar_shuffle_n** | int | ``100`` | Whether to cache IDs to avoid re-tokenizing data. This will be deprecated in favor of tarred datasets. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.{train_ds,validation_ds,test_ds}.shard_strategy** | str | ``scatter`` | Whether to cache IDs in each of the nodes in multi-node training. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ +| **model.preproc_out_dir** | str | ``null`` | Path to folder that contains processed tar files or directory where new tar files are written. | ++-----------------------------------------------------------------------+-----------------+----------------+----------------------------------------------------------------------------------------------------------------+ + +Tarred datasets can be created in two ways: + +#. Using the Hydra config and `training script `__. + + For example: + + .. code :: + + python examples/nlp/machine_translation/enc_dec_nmt.py \ + -cn aayn_base \ + do_training=false \ + model.preproc_out_dir=/path/to/preproc_dir \ + model.train_ds.use_tarred_dataset=true \ + model.train_ds.lines_per_dataset_fragment=1000000 \ + model.train_ds.num_batches_per_tarfile=200 \ + model.train_ds.src_file_name=train.tokenized.en \ + model.train_ds.tgt_file_name=train.tokenized.es \ + model.validation_ds.src_file_name=validation.tokenized.en \ + model.validation_ds.tgt_file_name=validation.tokenized.es \ + model.encoder_tokenizer.vocab_size=32000 \ + model.decoder_tokenizer.vocab_size=32000 \ + ~model.test_ds \ + trainer.gpus=[0,1,2,3] \ + +trainer.fast_dev_run=true \ + exp_manager=null \ + + The above script processes the parallel tokenized text files into tarred datasets that are written to ``/path/to/preproc_dir``. Since + ``do_training`` is set to ``False``, the above script only creates tarred datasets and then exits. If ``do_training`` is set ``True``, + then one of two things happen: + + (a) If no tar files are present in ``model.preproc_out_dir``, the script first creates those files and then commences training. + (b) If tar files are already present in ``model.preproc_out_dir``, the script starts training from the provided tar files. + +#. Using a separate script without Hydra. + + Tarred datasets for parallel corpora can also be created with a script that doesn't require specifying a configs via Hydra and + just uses Python argparse. + + For example: + + .. code :: + + python examples/nlp/machine_translation/create_tarred_parallel_dataset.py \ + --shared_tokenizer \ + --clean \ + --bpe_dropout 0.1 \ + --src_fname train.tokenized.en \ + --tgt_fname train.tokenized.es \ + --out_dir /path/to/preproc_dir \ + --vocab_size 32000 \ + --max_seq_length 512 \ + --min_seq_length 1 \ + --tokens_in_batch 8192 \ + --lines_per_dataset_fragment 1000000 \ + --num_batches_per_tarfile 200 + + You can then set `model.preproc_out_dir=/path/to/preproc_dir` and `model.train_ds.use_tarred_dataset=true` to train with this data. + +Model Configuration and Training +-------------------------------- + +The overall model consists of an encoder, decoder, and classification head. Encoders and decoders have the following configuration +options: + ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Default** | **Description** | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.max_sequence_length** | int | ``512`` | Maximum sequence length of positional encodings. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.embedding_dropout** | float | ``0.1`` | Path to JSON metadata file that contains only a single entry for the total number of batches in the dataset. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.learn_positional_encodings** | bool | ``false`` | If ``True``, this is a regular learnable embedding layer. If ``False``, fixes position encodings to sinusoidal. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.hidden_size** | int | ``512`` | Size of the transformer hidden states. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.num_layers** | int | ``6`` | Number of transformer layers. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.inner_size** | int | ``2048`` | Size of the hidden states within the feedforward layers. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.num_attention_heads** | int | ``8`` | Number of attention heads. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.ffn_dropout** | float | ``0.1`` | Dropout probability within the feedforward layers. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.attn_score_dropout** | float | ``0.1`` | Dropout probability of the attention scores before softmax normalization. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.attn_layer_dropout** | float | ``0.1`` | Dropout probability of the attention query, key, and value projection activations. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.hidden_act** | str | ``relu`` | Activation function throughout the network. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.mask_future** | bool | ``false``, ``true`` | Whether to mask future timesteps for attention. Defaults to ``True`` for decoder and ``False`` for encoder. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ +| **model.{encoder,decoder}.pre_ln** | bool | ``false`` | Whether to apply layer-normalization before (``true``) or after (``false``) a sub-layer. | ++-------------------------------------------------------------------+-----------------+-----------------------+-----------------------------------------------------------------------------------------------------------------+ + +Our pre-trained models are optimized with Adam, with a maximum learning of 0.0004, beta of (0.9, 0.98), an inverse square root learning +rate schedule from :cite:`nlp-machine_translation-vaswani2017attention`. The **model.optim** section sets the optimization parameters. + +The following script creates tarred datasets based on the provided parallel corpus and train a model based on the ``base`` configuration +from :cite:`nlp-machine_translation-vaswani2017attention`. + +.. code :: + + python examples/nlp/machine_translation/enc_dec_nmt.py \ + -cn aayn_base \ + do_training=true \ + trainer.gpus=8 \ + ~trainer.max_epochs \ + +trainer.max_steps=100000 \ + +trainer.val_check_interval=1000 \ + +exp_manager.exp_dir=/path/to/store/results \ + +exp_manager.create_checkpoint_callback=True \ + +exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \ + +exp_manager.checkpoint_callback_params.mode=max \ + +exp_manager.checkpoint_callback_params.save_top_k=5 \ + model.preproc_out_dir=/path/to/preproc_dir \ + model.train_ds.use_tarred_dataset=true \ + model.train_ds.lines_per_dataset_fragment=1000000 \ + model.train_ds.num_batches_per_tarfile=200 \ + model.train_ds.src_file_name=train.tokenized.en \ + model.train_ds.tgt_file_name=train.tokenized.es \ + model.validation_ds.src_file_name=validation.tokenized.en \ + model.validation_ds.tgt_file_name=validation.tokenized.es \ + model.encoder_tokenizer.vocab_size=32000 \ + model.decoder_tokenizer.vocab_size=32000 \ + ~model.test_ds \ + +The trainer keeps track of the sacreBLEU score :cite:`nlp-machine_translation-post2018call` on the provided validation set and saves +the checkpoints that had the top 5 (by default) sacreBLEU scores. + +At the end of training, a ``.nemo`` file is written to the result directory which allows to run inference on a test set. + +Multi-Validation +---------------- + +To run validation on multiple datasets, specify ``validation_ds.src_file_name`` and ``validation_ds.tgt_file_name`` with a list of file paths: + +.. code-block:: bash + + model.validation_ds.src_file_name=[/data/wmt13-en-de.src,/data/wmt14-en-de.src] \ + model.validation_ds.tgt_file_name=[/data/wmt13-en-de.ref,/data/wmt14-en-de.ref] \ + +When using ``val_loss`` or ``val_sacreBLEU`` for the ``exp_manager.checkpoint_callback_params.monitor`` +then the 0th indexed dataset will be used as the monitor. + +To use other indexes, append the index: + +.. code-block:: bash + + exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU_dl_index_1 + +Multiple test datasets work exactly the same way as validation datasets, simply replace ``validation_ds`` by ``test_ds`` in the above examples. + +Model Inference +--------------- + +To generate translations on a test set and compute sacreBLEU scores, run the inference script: + +.. code :: + + python examples/nlp/machine_translation/nmt_transformer_infer.py \ + --model /path/to/model.nemo \ + --srctext test.en \ + --tgtout test.en-es.translations \ + --batch_size 128 \ + --source_lang en \ + --target_lang es + +The ``--srctext`` file must be provided before tokenization and normalization. The resulting ``--tgtout`` file is detokenized and +can be used to compute sacreBLEU scores. + +.. code :: + + cat test.en-es.translations | sacrebleu test.es + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: nlp-machine_translation + :keyprefix: nlp-machine_translation- diff --git a/models/tts/NeMo/docs/source/nlp/megatron_finetuning.rst b/models/tts/NeMo/docs/source/nlp/megatron_finetuning.rst new file mode 100644 index 0000000000..fec82f2d1d --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/megatron_finetuning.rst @@ -0,0 +1,216 @@ +.. _megatron_finetuning: + +Megatron-LM for Downstream Tasks +================================ + +Megatron-LM :cite:`nlp-megatron-shoeybi2019megatron` is a large, powerful transformer developed by the Applied Deep Learning Research +team at NVIDIA. Unlike BERT, the position of the layer normalization and the residual connection in the model architecture (similar to +GPT-2 architucture) are swapped, which allows the models to continue to improve as they were scaled up. This model reaches higher +scores compared to BERT on a range of Natural Language Processing (NLP) tasks. More details on efficient, model-parallel and multi-node +pre-training of GPT and BERT using mixed precision can be found in the `Megatron-LM GitHub repo `_. + + +Fine-tuning +----------- + +Pre-trained Megatron-LM (BERT) can be used for most of the NLP downstream tasks from `NeMo/examples/nlp `_. +Specify the ``model.language_model.pretrained_model_name`` parameter: + +.. code-block:: bash + + model.language_model.pretrained_model_name=megatron-bert-345m-uncased + +Available pre-trained Megatron-LM models: + +- ``megatron-bert-345m-cased`` +- ``megatron-bert-345m-uncased`` +- ``biomegatron-bert-345m-uncased`` +- ``biomegatron-bert-345m-cased`` + +For example, to fine-tune SQuAD v1.1 with Megatron-LM, run: + +.. code-block:: bash + + python question_answering_squad.py \ + model.train_ds.file= \ + model.validation_ds= \ + model.language_model.pretrained_model_name=megatron-bert-345m-uncased + +If you have a different checkpoint or model configuration (pre-trained with `Megatron-LM GitHub repo `_), +use ``model.language_model.pretrained_model_name=megatron-bert-uncased`` or ``model.language_model.pretrained_model_name=megatron-bert-cased`` +and specify ``--bert_config`` and ``--bert_checkpoint`` for your model. + +.. note:: + Megatron-LM has its own set of training arguments (including tokenizer) that are ignored during fine-tuning in NeMo. Use downstream + task training scripts for all NeMo supported arguments. + +BioMegatron +----------- + +BioMegatron has the same network architecture as the Megatron-LM, but is pretrained on a different dataset - `PubMed `_, +a large biomedical text corpus, which achieves better performance in biomedical downstream tasks than the original Megatron-LM. + +Examples of using BioMegatron on biomedical downstream tasks can be found at (can be executed with `Google's Colab `_): +`NeMo/tutorials/nlp/Relation_Extraction-BioMegatron.ipynb `__ and `NeMo/tutorials/nlp/Token_Classification-BioMegatron.ipynb `__. + +Model Parallelism +----------------- + +`Megatron-LM `_ is a highly optimized and efficient library for training large language models. +With Megatron model parallelism, language models can be trained with billions of weights and then used in NeMo for downstream tasks. + +NeMo handles pretrained model parallel checkpoints from Megatron-LM automatically and model parallel models in NeMo have the all +the same features as other NeMo Models. + +.. note:: + + Currently, NeMo only supports tensor model parallelism. + +Training +^^^^^^^^ + +All of the necessary logic to train model parallel models in NeMo with PyTorch Lightning is contained in the ``NLPDDPPlugin``. +The ``NLPDDPPlugin`` subclasses the PyTorch Lightning training type plugin ``DDPPlugin``. +See `plugins `_ for more information on PyTorch Lightning Plugins. + +To enable model parallel training in NeMo: + +.. code-block:: python + + trainer = Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer) + +Megatron-LM checkpoints have a specific format. One checkpoint is saved for each model parallel rank: + +.. code-block:: bash + + iter_0080000/ + ├── mp_rank_00 + │ └── model_optim_rng.pt + └── mp_rank_01 + └── model_optim_rng.pt + + +To start fine-tuning from a Megatron-LM checkpoint, simply pass the path to the Megatron-LM checkpoint +via the language model config: + +.. code-block:: bash + + model.language_model.lm_checkpoint=/raid/megatron/bert/iter_0080000 \ + +We also need to input the model configuration. This can be done via json: + +.. code-block:: json + + { + "hidden-size": 1024, + "num-attention-heads": 16, + "num-layers": 24, + "max-seq-length": 512 + } + +And input via command line: + +.. code-block:: bash + + model.language_model.config_file=/raid/data/megatron/bert/config.json \ + +Or the model configuration can be input via YAML: + +.. code-block:: YAML + + model: + language_model: + config: + hidden_size: 1024 + num_attention_heads: 16 + num_layers: 24 + max_position_embeddings: 512 + +Additionally, Megatron-LM requires a vocab file: + +.. code-block:: bash + + model.tokenizer.vocab_file=/path/to/vocab.txt + +If using the Megatron-LM default tokenizer for training BERT the vocab file can be omitted: + +.. code-block:: bash + + # uncased model + model.tokenizer.tokenizer_name=megatron-bert-uncased + +.. code-block:: bash + + # cased model + model.tokenizer.tokenizer_name=megatron-bert-uncased + +Auto-Resume +^^^^^^^^^^^ + +Resuming training with NeMo experiment manager and PyTorch Lightning works exactly the same as other NeMo models. +While training with PTL, model parallel checkpoint will be saved and loaded properly. + +.. code-block:: bash + + checkpoints/ + ├── mp_rank_00 + │ ├── mp_autoresume-last.ckpt + │ ├── mp_autoresume---val_loss=0.35-epoch=0.ckpt + │ ├── mp_autoresume---val_loss=0.38-epoch=1.ckpt + │ └── mp_autoresume---val_loss=0.39-epoch=2.ckpt + └── mp_rank_01 + ├── mp_autoresume-last.ckpt + ├── mp_autoresume---val_loss=0.35-epoch=0.ckpt + ├── mp_autoresume---val_loss=0.38-epoch=1.ckpt + └── mp_autoresume---val_loss=0.39-epoch=2.ckpt + +Save and Restore +^^^^^^^^^^^^^^^^ + +Model parallel .nemo files behave the same as all other .nemo files. Calling ``.save_to`` will save +a checkpoint for each model parallel rank inside the .nemo file: + +.. code-block:: bash + + text_class_350m + ├── megatron-bert-uncased_encoder_config.json + ├── megatron_checkpoint_version.json + ├── model_config.yaml + ├── mp_rank_00 + │ └── model_weights.ckpt + ├── mp_rank_01 + │ └── model_weights.ckpt + ├── tokenizer_vocab_dict.json + └── tokenizer.vocab_file + +When restoring a model parallel .nemo file, we must pass in the ``Trainer`` as model parallel requires DDP: + +.. code-block:: python + + model = TokenClassificationModel.restore_from(cfg.pretrained_model, trainer=trainer) + +Evaluation +^^^^^^^^^^ + +Since model parallel models always require more than one GPU, the ``Trainer`` is needed for evaluation: + +.. code-block:: python + + trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer) + + model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer) + model.setup_test_data(test_data_config=cfg.model.test_ds) + + trainer.test(model=model, ckpt_path=None) + + + + + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-MEGATRON + :keyprefix: nlp-megatron- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/models.rst b/models/tts/NeMo/docs/source/nlp/models.rst new file mode 100755 index 0000000000..c6b9adbad0 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/models.rst @@ -0,0 +1,22 @@ +.. _nlp_models: + +Models +====== + +NeMo's NLP collection supports the following models: + +.. toctree:: + :maxdepth: 1 + + punctuation_and_capitalization + token_classification + joint_intent_slot + text_classification + bert_pretraining + language_modeling + question_answering + sgd_qa + glue_benchmark + information_retrieval + nlp_model + machine_translation diff --git a/models/tts/NeMo/docs/source/nlp/nlp_all.bib b/models/tts/NeMo/docs/source/nlp/nlp_all.bib new file mode 100644 index 0000000000..f291310941 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/nlp_all.bib @@ -0,0 +1,109 @@ +@article{devlin2018bert, + title={Bert: Pre-training of deep bidirectional transformers for language understanding}, + author={Devlin, Jacob and Chang, Ming-Wei and Lee, Kenton and Toutanova, Kristina}, + journal={arXiv preprint arXiv:1810.04805}, + year={2018} +} + +@article{shoeybi2019megatron, + title={Megatron-lm: Training multi-billion parameter language models using model parallelism}, + author={Shoeybi, Mohammad and Patwary, Mostofa and Puri, Raul and LeGresley, Patrick and Casper, Jared and Catanzaro, Bryan}, + journal={arXiv preprint arXiv:1909.08053}, + year={2019} +} + +@InProceedings{maas2011, + author = {Maas, Andrew L. and Daly, Raymond E. and Pham, Peter T. and Huang, Dan and Ng, Andrew Y. and Potts, Christopher}, + title = {Learning Word Vectors for Sentiment Analysis}, + booktitle = {Proceedings of the 49th Annual Meeting of the Association for Computational Linguistics: Human Language Technologies}, + month = {June}, + year = {2011}, + address = {Portland, Oregon, USA}, + publisher = {Association for Computational Linguistics}, + pages = {142--150}, + url = {http://www.aclweb.org/anthology/P11-1015} +} + +@inproceedings{socher2013, + title = "Recursive Deep Models for Semantic Compositionality Over a Sentiment Treebank", + author = "Socher, Richard and Perelygin, Alex and Wu, Jean and Chuang, Jason and Manning, Christopher D. and Ng, Andrew and Potts, Christopher", + booktitle = "Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing", + month = oct, + year = "2013", + address = "Seattle, Washington, USA", + publisher = "Association for Computational Linguistics", + url = "https://www.aclweb.org/anthology/D13-1170", + pages = "1631--1642", +} + +@article{lim2018chemical, + title={Chemical--gene relation extraction using recursive neural network}, + author={Lim, Sangrak and Kang, Jaewoo}, + journal={Database}, + volume={2018}, + year={2018}, + publisher={Oxford Academic} +} + +@inproceedings{li2007scalable, + title={Scalable term selection for text categorization}, + author={Li, Jingyang and Sun, Maosong}, + booktitle={Proceedings of the 2007 Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning (EMNLP-CoNLL)}, + pages={774--782}, + year={2007} +} + +@misc{lee2019biobert, + title={BioBERT: a pre-trained biomedical language representation model for biomedical text mining}, + author={Jinhyuk Lee and Wonjin Yoon and Sungdong Kim and Donghyeon Kim and Sunkyu Kim and Chan Ho So and Jaewoo Kang}, + year={2019}, + eprint={1901.08746}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} + +@misc{shin2020biomegatron, + title={BioMegatron: Larger Biomedical Domain Language Model}, + author={Hoo-Chang Shin and Yang Zhang and Evelina Bakhturina and Raul Puri and Mostofa Patwary and Mohammad Shoeybi and Raghav Mani}, + year={2020}, + eprint={2010.06060}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} + +@article{vaswani2017attention, + title={Attention is all you need}, + author={Vaswani, Ashish and Shazeer, Noam and Parmar, Niki and Uszkoreit, Jakob and Jones, Llion and Gomez, Aidan N and Kaiser, Lukasz and Polosukhin, Illia}, + journal={arXiv preprint arXiv:1706.03762}, + year={2017} +} + +@article{sennrich2015neural, + title={Neural machine translation of rare words with subword units}, + author={Sennrich, Rico and Haddow, Barry and Birch, Alexandra}, + journal={arXiv preprint arXiv:1508.07909}, + year={2015} +} + +@article{provilkov2019bpe, + title={Bpe-dropout: Simple and effective subword regularization}, + author={Provilkov, Ivan and Emelianenko, Dmitrii and Voita, Elena}, + journal={arXiv preprint arXiv:1910.13267}, + year={2019} +} + +@article{post2018call, + title={A call for clarity in reporting BLEU scores}, + author={Post, Matt}, + journal={arXiv preprint arXiv:1804.08771}, + year={2018} +} + +@misc{zhang2021sgdqa, + title={SGD-QA: Fast Schema-Guided Dialogue State Tracking for Unseen Services}, + author={Yang Zhang and Vahid Noroozi and Evelina Bakhturina and Boris Ginsburg}, + year={2021}, + eprint={2105.08049}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/nlp_model.rst b/models/tts/NeMo/docs/source/nlp/nlp_model.rst new file mode 100755 index 0000000000..8c47518c40 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/nlp_model.rst @@ -0,0 +1,43 @@ +.. _nlp_model: + +Model NLP +========= + +The config file for NLP models contain three main sections: + + - ``trainer``: contains the configs for PTL training. For more information, refer to :doc:`../core/core` and `PTL Trainer class API '. + - ``exp_manager``: the configs of the experiment manager. For more information, refer to :doc:`../core/core`. + - ``model``: contains the configs of the datasets, model architecture, tokenizer, optimizer, scheduler, etc. + +The following sub-sections of the model section are shared among most of the NLP models. + + - ``tokenizer``: specifies the tokenizer + - ``language_model``: specifies the underlying model to be used as the encoder + - ``optim``: the configs of the optimizer and scheduler :doc:`../core/core` + +The ``tokenizer`` and ``language_model`` sections have the following parameters: + ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.tokenizer.tokenizer_name** | string | Tokenizer name will be filled automatically based on ``model.language_model.pretrained_model_name``. | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.tokenizer.vocab_file** | string | Path to tokenizer vocabulary. | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.tokenizer.tokenizer_model** | string | Path to tokenizer model (only for sentencepiece tokenizer). | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.language_model.pretrained_model_name** | string | Pre-trained language model name, for example: ``bert-base-cased`` or ``bert-base-uncased``. | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.language_model.lm_checkpoint** | string | Path to the pre-trained language model checkpoint. | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.language_model.config_file** | string | Path to the pre-trained language model config file. | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.language_model.config** | dictionary | Config of the pre-trained language model. | ++------------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +The parameter **model.language_model.pretrained_model_name** can be one of the following: + + - ``megatron-bert-345m-uncased``, ``megatron-bert-345m-cased``, ``biomegatron-bert-345m-uncased``, ``biomegatron-bert-345m-cased``, ``bert-base-uncased``, ``bert-large-uncased``, ``bert-base-cased``, ``bert-large-cased`` + - ``distilbert-base-uncased``, ``distilbert-base-cased`` + - ``roberta-base``, ``roberta-large``, ``distilroberta-base`` + - ``albert-base-v1``, ``albert-large-v1``, ``albert-xlarge-v1``, ``albert-xxlarge-v1``, ``albert-base-v2``, ``albert-large-v2``, ``albert-xlarge-v2``, ``albert-xxlarge-v2`` diff --git a/models/tts/NeMo/docs/source/nlp/punctuation_and_capitalization.rst b/models/tts/NeMo/docs/source/nlp/punctuation_and_capitalization.rst new file mode 100755 index 0000000000..4cce7b34ce --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/punctuation_and_capitalization.rst @@ -0,0 +1,300 @@ +.. _punctuation_and_capitalization: + +Punctuation and Capitalization Model +==================================== + +Automatic Speech Recognition (ASR) systems typically generate text with no punctuation and capitalization of the words. +There are two issues with non-punctuated ASR output: + +- it could be difficult to read and understand +- models for some downstream tasks, such as named entity recognition, machine translation, or text-to-speech, are usually trained +on punctuated datasets and using raw ASR output as the input to these models could deteriorate their performance + +Quick Start Guide +----------------- + +.. code-block:: python + + from nemo.collections.nlp.models import PunctuationCapitalizationModel + + # to get the list of pre-trained models + PunctuationCapitalizationModel.list_available_models() + + # Download and load the pre-trained BERT-based model + model = PunctuationCapitalizationModel.from_pretrained("punctuation_en_bert") + + # try the model on a few examples + model.add_punctuation_capitalization(['how are you', 'great how about you']) + +Model Description +----------------- + +For each word in the input text, the Punctuation and Capitalization model: + +- predicts a punctuation mark that should follow the word (if any). By default, the model supports commas, periods, and question marks. +- predicts if the word should be capitalized or not + +In the Punctuation and Capitalization model, we are jointly training two token-level classifiers on top of a pre-trained +language model, such as `BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding `__ :cite:`nlp-punct-devlin2018bert`. + +.. note:: + + We recommend you try this model in a Jupyter notebook (run on `Google's Colab `_.): `NeMo/tutorials/nlp/Punctuation_and_Capitalization.ipynb `__. + + Connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for the hardware accelerator). + + An example script on how to train the model can be found at: `NeMo/examples/nlp/token_classification/punctuation_capitalization_train.py `__. + + An example script on how to run evaluation and inference can be found at: `NeMo/examples/nlp/token_classification/punctuation_capitalization_evaluate.py `__. + + The default configuration file for the model can be found at: `NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml `__. + +.. _raw_data_format_punct: + +Raw Data Format +--------------- + +The Punctuation and Capitalization model can work with any text dataset, although it is recommended to balance the data, especially +for the punctuation task. Before pre-processing the data to the format expected by the model, the data should be split into ``train.txt`` +and ``dev.txt`` (and optionally ``test.txt``). Each line in the ``train.txt/dev.txt/test.txt`` should represent one or more full +and/or truncated sentences. + +Example of the ``train.txt``/``dev.txt`` file: + +.. code:: + + When is the next flight to New York? + The next flight is ... + .... + + +The ``source_data_dir`` structure should look similar to the following: + +.. code:: + + . + |--sourced_data_dir + |-- dev.txt + |-- train.txt + +NeMo Data Format +---------------- + +The Punctuation and Capitalization model expects the data in the following format: + +The training and evaluation data is divided into 2 files: +- ``text.txt`` +- ``labels.txt`` + +Each line of the ``text.txt`` file contains text sequences, where words are separated with spaces. + +[WORD] [SPACE] [WORD] [SPACE] [WORD], for example: + + :: + + when is the next flight to new york + the next flight is ... + ... + +The ``labels.txt`` file contains corresponding labels for each word in ``text.txt``, the labels are separated with spaces. +Each label in ``labels.txt`` file consists of 2 symbols: + +- the first symbol of the label indicates what punctuation mark should follow the word (where ``O`` means no punctuation needed) +- the second symbol determines if a word needs to be capitalized or not (where ``U`` indicates that the word should be upper cased, +and ``O`` - no capitalization needed) + +By default, the following punctuation marks are considered: commas, periods, and question marks; the remaining punctuation marks were +removed from the data. This can be changed by introducing new labels in the ``labels.txt`` files. + +Each line of the ``labels.txt`` should follow the format: ``[LABEL] [SPACE] [LABEL] [SPACE] [LABEL]`` (for ``labels.txt``). For example, +labels for the above ``text.txt`` file should be: + + :: + + OU OO OO OO OO OO OU ?U + OU OO OO OO ... + ... + +The complete list of all possible labels used in this tutorial are: + +- ``OO`` +- ``O`` +- ``.O`` +- ``?O`` +- ``OU`` +- +- ``U`` +- ``.U`` +- ``?U`` + +Converting Raw Data to NeMo Format +---------------------------------- + +To pre-process the raw text data, stored under :code:`sourced_data_dir` (see the :ref:`raw_data_format_punct` +section), run the following command: + +.. code:: + + python examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \ + -s + -o + + +Required Argument for Dataset Conversion +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`-s` or :code:`--source_file`: path to the raw file +- :code:`-o` or :code:`--output_dir` - path to the directory to store the converted files + +After the conversion, the :code:`output_dir` should contain :code:`labels_*.txt` and :code:`text_*.txt` files. The default names +for the training and evaluation in the :code:`conf/punctuation_capitalization_config.yaml` are the following: + +.. code:: + + . + |--output_dir + |-- labels_dev.txt + |-- labels_train.txt + |-- text_dev.txt + |-- text_train.txt + +Training Punctuation and Capitalization Model +--------------------------------------------- + +The language model is initialized with the pre-trained model from `HuggingFace Transformers `__, +unless the user provides a pre-trained checkpoint for the language model. Example of model configuration file for training the model can be found at: `NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml `__. + +The specification is roughly grouped into the following categories: + +- Parameters that describe the training process: **trainer** +- Parameters that describe the datasets: **model.dataset**, **model.train_ds**, **model.validation_ds** +- Parameters that describe the model: **model** + +More details about parameters in the config file can be found below and in the `model's config file `__: + ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **pretrained_model** | string | Path to the pre-trained model ``.nemo`` file or pre-trained model name. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.dataset.data_dir** | string | Path to the data converted to the specified above format. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.punct_head.punct_num_fc_layers** | integer | Number of fully connected layers. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.punct_head.fc_dropout** | float | Activation to use between fully connected layers. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.punct_head.activation** | string | Dropout to apply to the input hidden states. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.punct_head.use_transrormer_init** | bool | Whether to initialize the weights of the classifier head with the same approach used in Transformer. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.capit_head.punct_num_fc_layers** | integer | Number of fully connected layers. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.capit_head.fc_dropout** | float | Dropout to apply to the input hidden states. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.capit_head.activation** | string | Activation function to use between fully connected layers. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.capit_head.use_transrormer_init** | bool | Whether to initialize the weights of the classifier head with the same approach used in Transformer. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.text_file** | string | Name of the text training file located at ``data_dir``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.labels_file** | string | Name of the labels training file located at ``data_dir``, such as ``labels_train.txt``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.num_samples** | integer | Number of samples to use from the training dataset, ``-1`` - to use all. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.text_file** | string | Name of the text file for evaluation, located at ``data_dir``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.labels_file** | string | Name of the labels dev file located at ``data_dir``, such as ``labels_dev.txt``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.num_samples** | integer | Number of samples to use from the dev set, ``-1`` - to use all. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +For more information, refer to the :ref:`nlp_model` section. + +To train the model from scratch, run: + +.. code:: + + python examples/nlp/token_classification/punctuation_and_capitalization_train.py \ + model.dataset.data_dir= \ + trainer.gpus=[0,1] \ + optim.name=adam \ + optim.lr=0.0001 + +The above command will start model training on GPUs 0 and 1 with Adam optimizer and learning rate of 0.0001; and the trained model is +stored in the ``nemo_experiments/Punctuation_and_Capitalization`` folder. + +To train from the pre-trained model, run: + +.. code:: + + python examples/nlp/token_classification/punctuation_and_capitalization_train.py \ + model.dataset.data_dir= \ + pretrained_model= + + +Required Arguments for Training +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +* :code:`model.dataset.data_dir`: Path to the `data_dir` with the pre-processed data files. + + +.. note:: + + All parameters defined in the configuration file can be changed with command arguments. For example, the sample config file + mentioned above has :code:`validation_ds.batch_size` set to ``64``. However, if you see that the GPU utilization can be + optimized further by using a larger batch size, you may override to the desired value by adding the field :code:`validation_ds.batch_size=128` + over the command-line. You can repeat this with any of the parameters defined in the sample configuration file. + +Inference +--------- + +An example script on how to run inference on a few examples, can be found at `examples/nlp/token_classification/punctuation_capitalization_evaluate.py `_. + +To start inference with a pre-trained model on a few examples, run: + +.. code:: + + python punctuation_capitalization_evaluate.py \ + pretrained_model= + + +Model Evaluation +---------------- + +An example script on how to evaluate the pre-trained model, can be found at `examples/nlp/token_classification/punctuation_capitalization_evaluate.py `_. + +To start evaluation of the pre-trained model, run: + +.. code:: + + python punctuation_capitalization_evaluate.py \ + model.dataset.data_dir= \ + pretrained_model=punctuation_en_bert \ + model.test_ds.text_file= \ + model.test_ds.labels_file= + + +Required Arguments +^^^^^^^^^^^^^^^^^^ + +- :code:`pretrained_model`: pretrained Punctuation and Capitalization model from ``list_available_models()`` or path to a ``.nemo`` +file. For example: ``punctuation_en_bert`` or ``your_model.nemo``. +- :code:`model.dataset.data_dir`: path to the directory that containes :code:`model.test_ds.text_file` and :code:`model.test_ds.labels_file` + +During evaluation of the :code:`test_ds`, the script generates two classification reports: one for capitalization task and another +one for punctuation task. This classification reports include the following metrics: + +- :code:`Precision` +- :code:`Recall` +- :code:`F1` + +More details about these metrics can be found `here `__. + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-PUNCT + :keyprefix: nlp-punct- + diff --git a/models/tts/NeMo/docs/source/nlp/question_answering.rst b/models/tts/NeMo/docs/source/nlp/question_answering.rst new file mode 100644 index 0000000000..567be7dc62 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/question_answering.rst @@ -0,0 +1,300 @@ +.. _question_answering: + +Question Answering model +======================== + +With Question Answering, or Reading Comprehension, given a question and a passage of content (context) that may contain an answer for +the question, the model predicts the span within the text with a start and end position indicating the answer to the question. For +datasets like SQuAD 2.0, this model supports cases when the answer is not contained in the content. + +For every word in the context of a given question, the model is trained to predict: + +- The likelihood this word is the start of the span. +- The likelihood this word is the end of the span. + +The model chooses the start and end words with maximal probabilities. When the content does not contain the answer, we would like the +start and end span to be set for the first token. + +A pretrained BERT encoder with two span prediction heads is used for the prediction start and the end position of the answer. The span +predictions are token classifiers consisting of a single linear layer. + +Quick Start Guide +----------------- + +.. code-block:: python + + from nemo.collections.nlp.models import QAModel + + # to get the list of pre-trained models + QAModel.list_available_models() + + # Download and load the pre-trained BERT-based model + model = QAModel.from_pretrained("qa_squadv1.1_bertbase") + + # try the model on a few examples + model.inference(test_file) + + +.. note:: + + We recommend you try Question Answering model in a Jupyter notebook (can run on `Google's Colab `_.): + `NeMo/tutorials/nlp/Question_Answering_Squad.ipynb `__. + + Connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for the hardware accelerator). + + An example script on how to train and evaluate the model can be found here: `NeMo/examples/nlp/question_answering/question_answering_squad.py `__. + + The default configuration file for the model can be found at: `NeMo/examples/nlp/question_answering/conf/question_answering_squad.yaml `__. + + + +Available models +^^^^^^^^^^^^^^^^ + +.. list-table:: *Pretrained Models* + :widths: 5 10 + :header-rows: 1 + + * - Model + - Pretrained Checkpoint + * - qa_squadv1.1_bertbase + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv1_1_bertbase + * - qa_squadv2.0_bertbase + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv2_0_bertbase + * - qa_squadv1.1_bertlarge + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv1_1_bertlarge + * - qa_squadv2.0_bertlarge + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv2_0_bertlarge + * - qa_squadv1.1_megatron_cased + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv1_1_megatron_cased + * - qa_squadv2.0_megatron_cased + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv2_0_megatron_cased + * - qa_squadv1.1_megatron_uncased + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv1_1_megatron_uncased + * - qa_squadv2.0_megatron_uncased + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:qa_squadv2_0_megatron_uncased + + +.. _dataset_question_answering: + +Data Format +----------- + +This model expects the dataset in `SQuAD format`_ (i.e., a JSON file for each dataset split). The code snippet below shows an example +of the training file. Each title has one or multiple paragraph entries, each consisting of the "context" and question-answer entries. +Each question-answer entry has: + +- A question +- A globally unique id +- The Boolean flag ``is_impossible``, which shows whether a question is answerable or not: + - if the question is answerable, one answer entry contains the text span and its starting character index in the context + - if the question is not answerable, an empty ``answers`` list is provided + +.. _SQuAD format: https://rajpurkar.github.io/SQuAD-explorer/ + +The evaluation files (for validation and testing) follow the above format, except that it can provide more than one answer to the +same question. The inference file also follows the above format, except that it does not require the ``answers`` and ``is_impossible`` +keywords. + +The following is an example of the data format (JSON file): + +.. code:: + + { + "data": [ + { + "title": "Super_Bowl_50", + "paragraphs": [ + { + "context": "Super Bowl 50 was an American football game to determine the champion of the National Football League (NFL) for the 2015 season. The American Football Conference (AFC) champion Denver Broncos defeated the National Football Conference (NFC) champion Carolina Panthers 24\u201310 to earn their third Super Bowl title. The game was played on February 7, 2016, at Levi's Stadium in the San Francisco Bay Area at Santa Clara, California. As this was the 50th Super Bowl, the league emphasized the \"golden anniversary\" with various gold-themed initiatives, as well as temporarily suspending the tradition of naming each Super Bowl game with Roman numerals (under which the game would have been known as \"Super Bowl L\"), so that the logo could prominently feature the Arabic numerals 50.", + "qas": [ + { + "question": "Where did Super Bowl 50 take place?", + "is_impossible": "false", + "id": "56be4db0acb8001400a502ee", + "answers": [ + { + "answer_start": "403", + "text": "Santa Clara, California" + } + ] + }, + { + "question": "What was the winning score of the Super Bowl 50?", + "is_impossible": "true", + "id": "56be4db0acb8001400a502ez", + "answers": [ + ] + } + ] + } + ] + } + ] + } + + +Dataset Download +---------------- + +To perform training of the Question Answering model on the SQuAD dataset, you must first download it from `here +`_ or run: + +.. code:: + + python get_squad.py + +There are two versions: + +- SQuAD version 1.1, which does not contain questions without the answer and has 100,000+ question-answer pairs on 500+ articles. +- SQuAD version 2.0, which combines the 100,000 questions from SQuAD 1.1 with over 50,000 unanswerable questions. To do well with + SQuAD 2.0, a system must not only answer questions when possible, but also determine when no answer is supported by the paragraph + and abstain from answering. + +After downloading the files, you should have a :code:`squad` data folder that contains the following four files for training and +evaluation: + +.. code:: + + . + |--squad + |-- v1.1/train-v1.1.json + |-- v1.1/dev-v1.1.json + |-- v2.0/train-v2.0.json + |-- v2.0/dev-v2.0.json + + +.. _model_training_question_answering: + +Model Training +-------------- + +In the Question Answering Model, we are training a span prediction head on top of a pre-trained language model, such as +`BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding `__ :cite:`nlp-qa-devlin2018bert`. +Unless the user provides a pre-trained checkpoint for the language model, the language model is initialized with the pre-trained model +from `HuggingFace Transformers `__. + +Example of model configuration file for training the model can be found at: `NeMo/examples/nlp/question_answering/conf/question_answering_squad_config.yaml `__. + +The specification can be grouped into three categories: + +- Parameters that describe the training process: **trainer** +- Parameters that describe the datasets: **model.dataset**, **model.train_ds**, **model.validation_ds**, **model.test_ds** +- Parameters that describe the model: **model** + +More details about parameters in the spec file can be found below: + ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **pretrained_model** | string | Pretrained QA model model from ``list_available_models()`` or path to a ``.nemo`` file. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **do_training** | bool | If ``true``, starts training, otherwise, skips training and continues with evaluation/inference. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.dataset.version_2_with_negative** | bool | Set to ``true`` to allow examples without an answer, e.g. for SQuAD v2.0. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.dataset.do_lower_case** | bool | If ``true``, converts text to lower case, only import for inference/evaluation. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.dataset.use_cache** | bool | If ``true``, either loads all preprocessed data from cache or saves preprocessed data for future use. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.file** | string | The training file path. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.num_samples** | integer | The number of samples to use from the training dataset (use ``-1`` to specify all samples). | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.file** | string | The validation file path. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.num_samples** | integer | The number of samples to use from the validation dataset (use ``-1`` to specify all samples). | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **test_ds.file** | string | The test file path (optional). | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **test_ds.num_samples** | integer | The number of samples to use from the test dataset (use ``-1`` to specify all samples). | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +Example of the command for training the model: + +.. code:: + + python question_answering_squad.py \ + model.train_ds.file= \ + model.validation_ds.file= \ + model.dataset.version_2_with_negative= \ + model.dataset.do_lower_case= \ + trainer.max_epochs= \ + trainer.gpus=[] + +.. Note:: + + The first time you train, it will take an extra 5-10 minutes to process the dataset. For future training runs, it will use the + processed dataset if :code:`model.dataset.use_cache=true`, which is automatically cached in the files in the same directory as + the data. + +Required Arguments for Training +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`model.train_ds.file`: path to the training file in JSON format +- :code:`model.validation_ds.file`: path to the validation file in JSON format + +Fine-tuning Procedure +^^^^^^^^^^^^^^^^^^^^^ + +Fine-tuning procedure and logs look similar to what's described in the Model Training section, with the addition of the model +that is initially loaded from a previously trained checkpoint, e.g. by specifying :code:`pretrained_model=`. + +Inference +--------- + +An example script on how to run inference can be found at `examples/nlp/question_answering/question_answering_squad.py `_. + +To run inference with the pre-trained model, run: + +.. code:: + + python question_answering_squad.py \ + pretrained_model= \ + model.dataset.version_2_with_negative= \ + model.dataset.do_lower_case= \ + do_training=false \ + model.validation_ds.file= + +Required Arguments for inference: +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`pretrained_model`: pretrained QA Model model from ``list_available_models()`` or path to a ``.nemo`` file + +Model Evaluation +---------------- + +An example script on how to evaluate the pre-trained model, can be found at `examples/nlp/question_answering/question_answering_squad.py `_. + +To run evaluation of the pre-trained model, run: + +.. code:: + + python question_answering_squad.py \ + pretrained_model= \ + model.dataset.version_2_with_negative= \ + model.dataset.do_lower_case= \ + do_training=false \ + model.test_ds.file= + + +Required Arguments: +^^^^^^^^^^^^^^^^^^^ + +- :code:`pretrained_model`: pretrained QA model from `list_available_models()`` or path to a ``.nemo`` file +- :code:`model.test_ds.file`: path to test file + +During evaluation of the :code:`test_ds`, the script generates the following metrics: + +- :code:`Exact Match (EM)` +- :code:`F1` + +More details about these metrics can be found `here `__. + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-QA + :keyprefix: nlp-qa- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/sgd_qa.rst b/models/tts/NeMo/docs/source/nlp/sgd_qa.rst new file mode 100644 index 0000000000..351ffe520e --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/sgd_qa.rst @@ -0,0 +1,17 @@ +.. _sgd_qa: + +Dialogue State Tracking - SGD-QA Model +====================================== + +More details can be found in the paper +`SGD-QA: Fast Schema-Guided Dialogue State Tracking for Unseen Services `__ :cite:`nlp-sgdqa-zhang2021sgdqa`. +An example script on how to train the model can be found here: `NeMo/examples/nlp/dialogue_state_tracking/sgd_qa.py `__. + + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-SGDQA + :keyprefix: nlp-sgdqa- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nlp/text_classification.rst b/models/tts/NeMo/docs/source/nlp/text_classification.rst new file mode 100644 index 0000000000..ea90c8503a --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/text_classification.rst @@ -0,0 +1,169 @@ +.. _text_classification: + +Text Classification model +========================= + +Text Classification is a sequence classification model based on BERT-based encoders. It can be used for a +variety of tasks like text classification, sentiment analysis, domain/intent detection for dialogue systems, etc. +The model takes a text input and predicts a label/class for the whole sequence. Megatron-LM and most of the BERT-based encoders +supported by HuggingFace including BERT, RoBERTa, and DistilBERT. + +An example script on how to train the model can be found here: `NeMo/examples/nlp/text_classification/text_classification_with_bert.py `__. +The default configuration file for the model can be found at: `NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml `__. + +There is also a Jupyter notebook which explains how to work with this model. We recommend you try this model in the Jupyter notebook (can run on `Google's Colab `_.): +`NeMo/tutorials/nlp/Text_Classification_Sentiment_Analysis.ipynb `__. +This tutorial shows an example of how run the Text Classification model on a sentiment analysis task. You may connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for the hardware accelerator) to run the notebook. + +Data Format +----------- + +The Text Classification model uses a simple text format as the dataset. It requires the data to be stored in TAB separated files +(``.tsv``) with two columns: sentence and label. Each line of the data file contains text sequences, where words are separated with spaces and the label is separated with ``[TAB]``, i.e.: + +.. code:: + + [WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL] + +Labels need to be integers starting from ``0``. Some examples taken from the SST2 dataset, which is a two-class dataset for sentiment analysis: + +.. code:: + + saw how bad this movie was 0 + lend some dignity to a dumb story 0 + the greatest musicians 1 + +You may need separate files for train, validation, and test with this format. + +Dataset Conversion +------------------ + +If your dataset is stored in another format, you need to convert it to NeMo's format to use this model. There are some conversion +scripts available for the following datasets: + +- SST2 :cite:`nlp-textclassify-socher2013` +- IMDB :cite:`nlp-textclassify-maas2011` +- ChemProt :cite:`nlp-textclassify-lim2018chemical` +- THUCnews :cite:`nlp-textclassify-li2007scalable` + +You can convert them from their original format to NeMo's format. To convert the original datasets to NeMo's format, use the ``examples/text_classification/data/import_datasets.py`` script: + +.. code:: python + + python import_datasets.py \ + --dataset_name DATASET_NAME \ + --target_data_dir TARGET_PATH \ + --source_data_dir SOURCE_PATH + +It reads the dataset specified by ``DATASET_NAME`` from ``SOURCE_PATH`` and converts it to NeMo's format. It then saves the new +dataset at ``TARGET_PATH``. + +Arguments: + +- ``dataset_name``: name of the dataset to convert (``sst-2``, ``chemprot``, ``imdb``, and ``thucnews`` are currently supported) +- ``source_data_dir``: directory of your dataset +- ``target_data_dir``: directory to save the converted dataset + +After the conversion, the ``TARGET_PATH`` should contain the following files: + +.. code:: + + . + |--TARGET_PATH + |-- train.tsv + |-- dev.tsv + |-- test.tsv + +Some datasets do not have the test set or their test set does not have any labels, therefore, the corresponding file may be missing. + +Model Training +-------------- + +You may find an example of a config file to be used for training of the Text Classification model at `NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml `__. +You can change any of these parameters directly from the config file or update them with the command-line arguments. + +The config file of the Text Classification model contains three main sections of ``trainer``, ``exp_manager``, and ``model``. You can +find more details about the ``trainer`` and ``exp_manager`` at :doc:`./nlp_model`. Some sub-sections of the model section including +``tokenizer``, ``language_model``, and ``optim`` are shared among most of the NLP models. The details of these sections can be found +at :doc:`./nlp_model`. + +Example of a command for training a Text Classification model on two GPUs for 50 epochs: + +.. code:: + + python examples/nlp/text_classification/text_classification_with_bert.py \ + model.training_ds.file_path= \ + model.validation_ds.file_path= \ + trainer.max_epochs=50 \ + trainer.gpus=[0,1] \ + optim.name=adam \ + optim.lr=0.0001 \ + model.nemo_path= + +At the start of each training experiment, there is a printed log of the experiment specification including any parameters added or +overridden via the command-line. It also shows additional information, such as which GPUs are available, where logs are saved, +and some samples from the datasets with their corresponding inputs to the model. It also provides some stats on the lengths of +sequences in the dataset. + +After each epoch, you should see a summary table of metrics on the validation set which include the following metrics: + +- :code:`Precision` +- :code:`Recall` +- :code:`F1` + +At the end of training, NeMo saves the last checkpoint at the path specified by ``NEMO_FILE_PATH`` in ``.nemo`` format. + +Model Arguments +^^^^^^^^^^^^^^^ + +The following table lists some of the model's parameters you can use in the config files or set them from the command-line when +training a model: + ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Default** | **Description** | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.class_labels.class_labels_file** | string | ``null`` | Path to an optional file containing the labels; each line is the string label corresponding to a label. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.dataset.num_classes** | int | ``Required`` | Number of the categories or classes, ``0`` < ``Label`` < ``num_classes``. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.dataset.do_lower_case** | boolean | ``true`` for uncased models, ``false`` for cased | Specifies if inputs should be made lower case, would be set automatically if pre-trained model is used. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.dataset.max_seq_length** | int | ``256`` | Maximum length of the input sequences. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.dataset.class_balancing** | string | ``null`` | ``null`` or ``weighted_loss``. ``weighted_loss`` enables the weighted class balancing to handle unbalanced classes. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.dataset.use_cache** | boolean | ``false`` | Uses a cache to store the processed dataset, you can use it for large datasets for speed up. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.classifier_head.num_output_layers** | integer | ``2`` | Number of fully connected layers of the classifier on top of the BERT model. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **model.classifier_head.fc_dropout** | float | ``0.1`` | Dropout ratio of the fully connected layers. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.file_path** | string | ``Required`` | Path of the training ``.tsv`` file. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.batch_size** | integer | ``32`` | Data loader's batch size. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.num_workers** | integer | ``2`` | Number of worker threads for data loader. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.shuffle** | boolean | ``true`` (training), ``false`` (test and validation) | Shuffles data for each epoch. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.drop_last** | boolean | ``false`` | Specifies if last batch of data needs to get dropped if it is smaller than batch size. | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.pin_memory** | boolean | ``false`` | Enables pin_memory of PyTorch's data loader to enhance speed | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ +| **{training,validation,test}_ds.num_samples** | integer | ``-1`` | Number of samples to be used from the dataset; -1 means all samples | ++-----------------------------------------------+-----------------+--------------------------------------------------------+-----------------------------------------------------------------------------------------------------------------------+ + + +Model Evaluation and Inference +------------------------------ + +After saving the model in ``.nemo`` format, you can load the model and perform evaluation or inference on the model. You can find +some examples in the example script: `NeMo/examples/nlp/text_classification/text_classification_with_bert.py `__. + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-TEXTCLASSIFY + :keyprefix: nlp-textclassify- diff --git a/models/tts/NeMo/docs/source/nlp/token_classification.rst b/models/tts/NeMo/docs/source/nlp/token_classification.rst new file mode 100755 index 0000000000..cab12eda12 --- /dev/null +++ b/models/tts/NeMo/docs/source/nlp/token_classification.rst @@ -0,0 +1,219 @@ +.. _token_classification: + +Token Classification (Named Entity Recognition) Model +===================================================== + +Token Classification model supports named entity recognition (NER) and other token level classification tasks, as long as the data +follows the format specified below. + +We're going to use NER task throughout this section. NER, also referred to as entity chunking, identification or extraction, is the +task of detecting and classifying key information (entities) in text. In other words, a NER model takes a piece of text as input and +for each word in the text, the model identifies a category the word belongs to. For example, in a sentence: ``Mary lives in Santa Clara +and works at NVIDIA``, the model should detect that ``Mary`` is a person, ``Santa Clara`` is a location and ``NVIDIA`` is a company. + +Quick Start Guide +----------------- + +.. code-block:: python + + from nemo.collections.nlp.models import TokenClassificationModel + + # to get the list of pre-trained models + TokenClassificationModel.list_available_models() + + # Download and load the pre-trained BERT-based model + model = TokenClassificationModel.from_pretrained("ner_en_bert") + + # try the model on a few examples + model.add_predictions(['we bought four shirts from the nvidia gear store in santa clara.', 'NVIDIA is a company.']) + +.. note:: + + We recommend you try this model in a Jupyter notebook (run on `Google's Colab `_.): + `NeMo/tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb `__. + + Connect to an instance with a GPU (**Runtime** -> **Change runtime type** -> select **GPU** for the hardware accelerator). + + An example script on how to train the model can be found here: `NeMo/examples/nlp/token_classification/token_classification_train.py `__. + + An example script on how to run evaluation and inference can be found here: `NeMo/examples/nlp/token_classification/token_classification_evaluate.py `__. + + The default configuration file for the model can be found here: `NeMo/examples/nlp/token_classification/conf/token_classification_config.yaml `__. + +.. _dataset_token_classification: + +Data Input for Token Classification Model +----------------------------------------- + +For pre-training or fine-tuning of the model, the data should be split into 2 files: + +- ``text.txt`` +- ``labels.txt`` + +Each line of the ``text.txt`` file contains text sequences, where words are separated with spaces, i.e.: ``[WORD] [SPACE] [WORD] [SPACE] [WORD]``. +The ``labels.txt`` file contains corresponding labels for each word in ``text.txt``, the labels are separated with spaces, i.e.: ``[LABEL] [SPACE] [LABEL] [SPACE] [LABEL]``. +Example of a ``text.txt`` file: + + Jennifer is from New York City . + She likes ... + ... + +Corresponding ``labels.txt`` file: + + B-PER O O B-LOC I-LOC I-LOC O + O O ... + ... + +Dataset Conversion +------------------ + +To convert an `IOB format `__ (short for inside, outside, beginning) data to the format required for training, use +`examples/nlp/token_classification/data/import_from_iob_format.py `_. + +.. code:: + + # For conversion from IOB format, for example, for CoNLL-2003 dataset: + python import_from_iob_format.py --data_file= + +Convert Dataset Required Arguments +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`--data_file`: path to the file to convert from IOB to NeMo format + +After running the above command, the data directory, where the :code:`--data_file` is stored, should contain :code:`text_*.txt` and :code:`labels_*.txt` files. +The default names for the training and evaluation in the :code:`conf/token_classification_config.yaml` are the following: + +.. code:: + + . + |--data_dir + |-- labels_dev.txt + |-- labels_train.txt + |-- text_dev.txt + |-- text_train.txt + + +Training The Token Classification model +--------------------------------------- + +In the Token Classification model, we are jointly training a classifier on top of a pre-trained language model, such as +`BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding `__ :cite:`nlp-ner-devlin2018bert`. +Unless the user provides a pre-trained checkpoint for the language model, the language model is initialized with the pre-trained model +from `HuggingFace Transformers `__. + +Example of model configuration file for training the model can be found at: `NeMo/examples/nlp/token_classification/conf/token_classification_config.yaml `__. + +The specification can be roughly grouped into three categories: + +- Parameters that describe the training process: **trainer** +- Parameters that describe the datasets: **model.dataset**, **model.train_ds**, **model.validation_ds** +- Parameters that describe the model: **model** + +More details about parameters in the spec file can be found below: + ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **Parameter** | **Data Type** | **Description** | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.dataset.data_dir** | string | Path to the data converted to the specified above format. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.head.num_fc_layers** | integer | Number of fully connected layers. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.head.fc_dropout** | float | Dropout to apply to the input hidden states. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.head.activation** | string | Activation to use between fully connected layers. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **model.punct_head.use_transrormer_init** | bool | Whether to initialize the weights of the classifier head with the same approach used in Transformer. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.text_file** | string | Name of the text training file located at ``data_dir``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.labels_file** | string | Name of the labels training file located at ``data_dir``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **training_ds.num_samples** | integer | Number of samples to use from the training dataset, ``-1`` - to use all. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.text_file** | string | Name of the text file for evaluation, located at ``data_dir``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.labels_file** | string | Name of the labels dev file located at ``data_dir``. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ +| **validation_ds.num_samples** | integer | Number of samples to use from the dev set, ``-1`` - to use all. | ++-------------------------------------------+-----------------+--------------------------------------------------------------------------------------------------------------+ + +For more information, see :ref:`nlp_model`. + +Example of the command for training the model: + +.. code:: + + python token_classification_train.py \ + model.dataset.data_dir= \ + trainer.max_epochs= \ + trainer.gpus=[] + + +Required Arguments for Training +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`model.dataset.data_dir`: path to the directory with pre-processed data. + +.. note:: + + While the arguments are defined in the spec file, if you want to override these parameter definitions in the spec file and + experiment with them, use the command-line to define the parameter. For example, the sample spec file mentioned above has + :code:`validation_ds.batch_size` set to ``64``. However, if the GPU utilization can be optimized further by + using a larger batch size, override it to the desired value by adding the field :code:`validation_ds.batch_size=128` from + the command-line. You can repeat this with any of the parameters defined in the sample spec file. + +Inference +--------- + +An example script on how to run inference can be found at `examples/nlp/token_classification/token_classification_evaluate.py `_. + +To run inference with the pre-trained model, run: + +.. code:: + + python token_classification_evaluate.py \ + pretrained_model= + +Required Arguments for Inference +^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ + +- :code:`pretrained_model`: pretrained Token Classification model from ``list_available_models()`` or path to a ``.nemo`` file. For example, ``ner_en_bert`` or ``your_model.nemo`` + +Model Evaluation +---------------- + +An example script on how to evaluate the pre-trained model can be found at `examples/nlp/token_classification/token_classification_evaluate.py `_. + +To start evaluation of the pre-trained model, run: + +.. code:: + + python token_classification_evaluate.py \ + model.dataset.data_dir= \ + pretrained_model=ner_en_bert \ + model.test_ds.text_file= \ + model.test_ds.labels_file= \ + model.dataset.max_seq_length=512 + + +Required Arguments +^^^^^^^^^^^^^^^^^^ + +- :code:`pretrained_model`: pretrained Token Classification model from ``list_available_models()`` or path to a ``.nemo`` file. For example, ``ner_en_bert`` or ``your_model.nemo`` +- :code:`model.dataset.data_dir`: path to the directory that containes :code:`model.test_ds.text_file` and :code:`model.test_ds.labels_file` + +During evaluation of the :code:`test_ds`, the script generates a classification report that includes the following metrics: + +- :code:`Precision` +- :code:`Recall` +- :code:`F1` + +For more information, see `here `__. + +References +---------- + +.. bibliography:: nlp_all.bib + :style: plain + :labelprefix: NLP-NER + :keyprefix: nlp-ner- \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/nv_logo.png b/models/tts/NeMo/docs/source/nv_logo.png new file mode 100644 index 0000000000..e764a71852 Binary files /dev/null and b/models/tts/NeMo/docs/source/nv_logo.png differ diff --git a/models/tts/NeMo/docs/source/starthere/best-practices.rst b/models/tts/NeMo/docs/source/starthere/best-practices.rst new file mode 100644 index 0000000000..2c0d1238ed --- /dev/null +++ b/models/tts/NeMo/docs/source/starthere/best-practices.rst @@ -0,0 +1,312 @@ +.. _best-practices: + +Best Practices +============== + +The NVIDIA NeMo Toolkit is available on GitHub as `open source `_ as well as +a `Docker container on NGC `_. It's assumed the user has +already installed NeMo by following the :ref:`quick-start-guide` instructions. + +The conversational AI pipeline consists of three major stages: + +- Automatic Speech Recognition (ASR) +- Natural Language Processing (NLP) or Natural Language Understanding (NLU) +- Text-to-Speech (TTS) or voice synthesis + +As you talk to a computer, the ASR phase converts the audio signal into text, the NLP stage interprets the question +and generates a smart response, and finally the TTS phase converts the text into speech signals to generate audio for +the user. The toolkit enables development and training of deep learning models involved in conversational AI and easily +chain them together. + +Why NeMo? +--------- + +Deep learning model development for conversational AI is complex. It involves defining, building, and training several +models in specific domains; experimenting several times to get high accuracy, fine tuning on multiple tasks and domain +specific data, ensuring training performance and making sure the models are ready for deployment to inference applications. +Neural modules are logical blocks of AI applications which take some typed inputs and produce certain typed outputs. By +separating a model into its essential components in a building block manner, NeMo helps researchers develop state-of-the-art +accuracy models for domain specific data faster and easier. + +Collections of modules for core tasks as well as specific to speech recognition, natural language, speech synthesis help +develop modular, flexible, and reusable pipelines. + +A neural module’s inputs/outputs have a neural type, that describes the semantics, the axis order and meaning, and the dimensions +of the input/output tensors. This typing allows neural modules to be safely chained together to build models for applications. + +NeMo can be used to train new models or perform transfer learning on existing pre-trained models. Pre-trained weights per module +(such as encoder, decoder) help accelerate model training for domain specific data. + +ASR, NLP and TTS pre-trained models are trained on multiple datasets (including some languages such as Mandarin) and optimized +for high accuracy. They can be used for transfer learning as well. + +NeMo supports developing models that work with Mandarin Chinese data. Tutorials help users train or fine tune models for +conversational AI with the Mandarin Chinese language. The export method provided in NeMo makes it easy to transform a trained +model into inference ready format for deployment. + +A key area of development in the toolkit is interoperability with other tools used by speech researchers. Data layer for Kaldi +compatibility is one such example. + +NeMo, PyTorch Lightning, And Hydra +---------------------------------- + +Conversational AI architectures are typically very large and require a lot of data and compute for training. NeMo uses +`Pytorch Lightning `_ for easy and performant multi-GPU/multi-node +mixed precision training. + +Pytorch Lightning is a high-performance PyTorch wrapper that organizes PyTorch code, scales model training, and reduces +boilerplate. PyTorch Lightning has two main components, the ``LightningModule`` and the Trainer. The ``LightningModule`` is +used to organize PyTorch code so that deep learning experiments can be easily understood and reproduced. The Pytorch Lightning +Trainer is then able to take the ``LightningModule`` and automate everything needed for deep learning training. + +NeMo models are LightningModules that come equipped with all supporting infrastructure for training and reproducibility. This +includes the deep learning model architecture, data preprocessing, optimizer, check-pointing and experiment logging. NeMo +models, like LightningModules, are also PyTorch modules and are fully compatible with the broader PyTorch ecosystem. Any NeMo +model can be taken and plugged into any PyTorch workflow. + +Configuring conversational AI applications is difficult due to the need to bring together many different Python libraries into +one end-to-end system. NeMo uses Hydra for configuring both NeMo models and the PyTorch Lightning Trainer. `Hydra `_ +is a flexible solution that makes it easy to configure all of these libraries from a configuration file or from the command-line. + +Every NeMo model has an example configuration file and a corresponding script that contains all configurations needed for training +to state-of-the-art accuracy. NeMo models have the same look and feel so that it is easy to do conversational AI research across +multiple domains. + +Using Optimized Pretrained Models With NeMo +------------------------------------------- + +`NVIDIA GPU Cloud (NGC) `_ is a software repository that has containers and models optimized +for deep learning. NGC hosts many conversational AI models developed with NeMo that have been trained to state-of-the-art accuracy +on large datasets. NeMo models on NGC can be automatically downloaded and used for transfer learning tasks. Pretrained models +are the quickest way to get started with conversational AI on your own data. NeMo has many `example scripts `_ +and `Jupyter Notebook tutorials `_ showing step-by-step how to fine-tune pretrained NeMo +models on your own domain-specific datasets. + +For BERT based models, the model weights provided are ready for +downstream NLU tasks. For speech models, it can be helpful to start with a pretrained model and then continue pretraining on your +own domain-specific data. Jasper and QuartzNet base model pretrained weights have been known to be very efficient when used as +base models. For an easy to follow guide on transfer learning and building domain specific ASR models, you can follow this `blog `_. +All pre-trained NeMo models can be found on the `NGC NeMo Collection `_. Everything needed to quickly get started +with NeMo ASR, NLP, and TTS models is there. + +Pre-trained models are packaged as a ``.nemo`` file and contain the PyTorch checkpoint along with everything needed to use the model. +NeMo models are trained to state-of-the-art accuracy and trained on multiple datasets so that they are robust to small differences +in data. NeMo contains a large variety of models such as speaker identification and Megatron BERT and the best models in speech and +language are constantly being added as they become available. NeMo is the premier toolkit for conversational AI model building and +training. + +For a list of supported models, refer to the :ref:`tutorials` section. + +ASR Guidance +------------ + +This section is to help guide your decision making by answering our most asked ASR questions. + +**Q: Is there a way to add domain specific vocabulary in NeMo? If so, how do I do that?** +A: QuartzNet and Jasper models are character-based. So pretrained models we provide for these two output lowercase English +letters and ‘. Users can re-retrain them on vocabulary with upper case letters and punctuation symbols. + +**Q: When training, there are “Reference” lines and “Decoded” lines that are printed out. It seems like the reference line should +be the “truth” line and the decoded line should be what the ASR is transcribing. Why do I see that even the reference lines do not +appear to be correct?** +A: Because our pre-trained models can only output lowercase letters and apostrophe, everything else is dropped. So the model will +transcribe 10 as ten. The best way forward is to prepare the training data first by transforming everything to lowercase and convert +the numbers from digit representation to word representation using a simple library such as `inflect `_. Then, add the uppercase letters +and punctuation back using the NLP punctuation model. Here is an example of how this is incorporated: `NeMo voice swap demo `_. + +**Q: What languages are supported in NeMo currently?** +A: Along with English, we provide pre-trained models for Zh, Es, Fr, De, Ru, It, Ca and Pl languages. +For more information, see `NeMo Speech Models `_. + +Data Augmentation +----------------- + +Data augmentation in ASR is invaluable. It comes at the cost of increased training time if samples are augmented during training +time. To save training time, it is recommended to pre-process the dataset offline for a one time preprocessing cost and then train +the dataset on this augmented training set. + +For example, processing a single sample involves: + +- Speed perturbation +- Time stretch perturbation (sample level) +- Noise perturbation +- Impulse perturbation +- Time stretch augmentation (batch level, neural module) + +A simple tutorial guides users on how to use these utilities provided in `GitHub: NeMo `_. + +Speech Data Explorer +-------------------- + +Speech data explorer is a `Dash-based tool `_ for interactive exploration of ASR/TTS datasets. + +Speech data explorer collects: + +- dataset statistics (alphabet, vocabulary, and duration-based histograms) +- navigation across datasets (sorting and filtering) +- inspections of individual utterances (waveform, spectrogram, and audio player) +- errors analysis (word error rate, character error rate, word match rate, mean word accuracy, and diff) + +In order to use the tool, it needs to be installed separately. Perform the steps `here `_ to install speech data explorer. + +Using Kaldi Formatted Data +-------------------------- + +The `Kaldi Speech Recognition Toolkit `_ project began in 2009 at `Johns Hopkins University `. It is a toolkit written in C++. If +researchers have used Kaldi and have datasets that are formatted to be used with the toolkit; they can use NeMo to develop models +based on that data. + +To load Kaldi-formatted data, you can simply use ``KaldiFeatureDataLayer`` instead of ``AudioToTextDataLayer``. The ``KaldiFeatureDataLayer`` +takes in the argument ``kaldi_dir`` instead of a ``manifest_filepath``. The ``manifest_filepath`` argument should be set to the directory +that contains the files ``feats.scp`` and ``text``. + +Using Speech Command Recognition Task For ASR Models +---------------------------------------------------- + +Speech Command Recognition is the task of classifying an input audio pattern into a set of discrete classes. It is a subset of ASR, +sometimes referred to as Key Word Spotting, in which a model is constantly analyzing speech patterns to detect certain ``action`` classes. + +Upon detection of these commands, a specific action can be taken. An example Jupyter notebook provided in NeMo shows how to train a +QuartzNet model with a modified decoder head trained on a speech commands dataset. + +.. note:: It is preferred that you use absolute paths to ``data_dir`` when preprocessing the dataset. + +NLP Fine-Tuning BERT +-------------------- + +BERT, or Bidirectional Encoder Representations from Transformers, is a neural approach to pre-train language representations which +obtains near state-of-the-art results on a wide array of Natural Language Processing (NLP) tasks, including the GLUE benchmark and +SQuAD Question & Answering dataset. + +BERT model checkpoints (`BERT-large-uncased `_ and `BERT-base-uncased `_) are provided can be used for either fine tuning BERT on your custom +dataset, or fine tuning downstream tasks, including GLUE benchmark tasks, Question & Answering tasks, Joint Intent & Slot detection, +Punctuation and Capitalization, Named Entity Recognition, and Speech Recognition post processing model to correct mistakes. + +.. note:: Almost all NLP examples also support RoBERTa and ALBERT models for downstream fine-tuning tasks (see the list of all +supported models by calling ``nemo.collections.nlp.modules.common.lm_utils.get_pretrained_lm_models_list()``. The user needs to specify +the name of the model desired while running the example scripts. + +BioMegatron Medical BERT +------------------------ + +BioMegatron is a large language model (Megatron-LM) trained on larger domain text corpus (PubMed abstract + full-text-commercial). +It achieves state-of-the-art results for certain tasks such as Relationship Extraction, Named Entity Recognition and Question & +Answering. Follow these tutorials to learn how to train and fine tune BioMegatron; pretrained models are provided on NGC: + +- `Relation Extraction BioMegatron `_ +- `Token Classification BioMegatron `_ + +Efficient Training With NeMo +---------------------------- + +Using Mixed Precision +^^^^^^^^^^^^^^^^^^^^^ + +Mixed precision accelerates training speed while protecting against noticeable loss. Tensor Cores is a specific hardware unit that +comes starting with the Volta and Turing architectures to accelerate large matrix to matrix multiply-add operations by operating them +on half precision inputs and returning the result in full precision. + +Neural networks which usually use massive matrix multiplications can be significantly sped up with mixed precision and Tensor Cores. +However, some neural network layers are numerically more sensitive than others. Apex AMP is an NVIDIA library that maximizes the +benefit of mixed precision and Tensor Cores usage for a given network. + +Multi-GPU Training +^^^^^^^^^^^^^^^^^^ + +This section is to help guide your decision making by answering our most asked multi-GPU training questions. + +**Q: Why is multi-GPU training preferred over other types of training?** +A: Multi-GPU training can reduce the total training time by distributing the workload onto multiple compute instances. This is +particularly important for large neural networks which would otherwise take weeks to train until convergence. Since NeMo supports +multi-GPU training, no code change is needed to move from single to multi-GPU training, only a slight change in your launch command +is required. + +**Q: What are the advantages of mixed precision training?** +A: Mixed precision accelerates training speed while protecting against noticeable loss in precision. Tensor Cores is a specific +hardware unit that comes starting with the Volta and Turing architectures to accelerate large matrix multiply-add operations by +operating on half precision inputs and returning the result in full precision in order to prevent loss in precision. Neural +networks which usually use massive matrix multiplications can be significantly sped up with mixed precision and Tensor Cores. +However, some neural network layers are numerically more sensitive than others. Apex AMP is a NVIDIA library that maximizes the +benefit of mixed precision and Tensor Core usage for a given network. + +**Q: What is the difference between multi-GPU and multi-node training?** +A: Multi-node is an abstraction of multi-GPU training, which requires a distributed compute cluster, where each node can have multiple +GPUs. Multi-node training is needed to scale training beyond a single node to large amounts of GPUs. + +From the framework perspective, nothing changes from moving to multi-node training. However, a master address and port needs to be set +up for inter-node communication. Multi-GPU training will then be launched on each node with passed information. You might also consider +the underlying inter-node network topology and type to achieve full performance, such as HPC-style hardware such as NVLink, InfiniBand +networking, or Ethernet. + + +Recommendations For Optimization And FAQs +----------------------------------------- + +This section is to help guide your decision making by answering our most asked NeMo questions. + +**Q: Are there areas where performance can be increased?** +A: You should try using mixed precision for improved performance. Note that typically when using mixed precision, memory consumption +is decreased and larger batch sizes could be used to further improve the performance. + +When fine-tuning ASR models on your data, it is almost always possible to take advantage of NeMo's pre-trained modules. Even if you +have a different target vocabulary, or even a different language; you can still try starting with pre-trained weights from Jasper or +QuartzNet ``encoder`` and only adjust the ``decoder`` for your needs. + +**Q: What is the recommended sampling rate for ASR?** +A: The released models are based on 16 KHz audio, therefore, ensure you use models with 16 KHz audio. Reduced performance should be +expected for any audio that is up-sampled from a sampling frequency less than 16 KHz data. + +**Q: How do we use this toolkit for audio with different types of compression and frequency than the training domain for ASR?** +A: You have to match the compression and frequency. + +**Q: How do you replace the 6-gram out of the ASR model with a custom language model? What is the language format supported in NeMo?** +A: NeMo’s Beam Search decoder with Levenberg-Marquardt (LM) neural module supports the KenLM language model. + +- You should retrain the KenLM language model on your own dataset. Refer to `KenLM’s documentation `_. +- If you want to use a different language model, other than KenLM, you will need to implement a corresponding decoder module. +- Transformer-XL example is present in OS2S. It would need to be updated to work with NeMo. `Here is the code `_. + +**Q: How do I use text-to-speech (TTS)?** +A: +- Obtain speech data ideally at 22050 Hz or alternatively at a higher sample rate and then down sample to 22050 Hz. + - If less than 22050 Hz and above 16000 Hz: + - Retrain WaveGlow on your own dataset. + - Tweak the spectrogram generation parameters, namely the ``window_size`` and the ``window_stride`` for their fourier + transforms. + - For below 16000 Hz, look into obtaining new data. +- In terms of bitrate/quantization, the general advice is the higher the better. We have not experimented enough to state how much + this impacts quality. +- For the amount of data, again the more the better, and the more diverse in terms of phonemes the better. Aim for around 20 hours + of speech after filtering for silences and non-speech audio. +- Most open speech datasets are in ~10 second format so training Tacotron 2 on audio on the order of 10s - 20s per sample is known + to work. Additionally, the longer the speech samples, the more difficult it will be to train Tacotron 2. +- Audio files should be clean. There should be little background noise or music. Data recorded from a studio mic is likely to be easier + to train compared to data captured using a phone. +- To ensure pronunciation of words are accurate; the technical challenge is related to the dataset, text to phonetic spelling is + required, use phonetic alphabet (notation) that has the name correctly pronounced. +- Here are some example parameters you can use to train Tacotron 2: + - Tacotron 2 requires a single speaker dataset + - Use AMP level O0 + - Trim long silences in the beginning and end + - ``optimizer="adam"`` + - ``beta1 = 0.9`` + - ``beta2 = 0.999`` + - ``lr=0.001 (constant)`` + - ``amp_opt_level="O0"`` + - ``weight_decay=1e-6`` + - ``batch_size=48 (per GPU)`` + - ``trim_silence=True`` + +Resources +--------- + +Ensure you are familiar with the following resources for NeMo. + +- Developer blogs + - `How to Build Domain Specific Automatic Speech Recognition Models on GPUs `_ + - `Develop Smaller Speech Recognition Models with NVIDIA’s NeMo Framework `_ + - `Neural Modules for Fast Development of Speech and Language Models `_ + +- Domain specific, transfer learning, Docker container with Jupyter Notebooks + - `Domain Specific NeMo ASR Application `_ + diff --git a/models/tts/NeMo/docs/source/starthere/intro.rst b/models/tts/NeMo/docs/source/starthere/intro.rst new file mode 100644 index 0000000000..d80a3dfac9 --- /dev/null +++ b/models/tts/NeMo/docs/source/starthere/intro.rst @@ -0,0 +1,157 @@ +Introduction +============ + +.. # define a hard line break for html +.. |br| raw:: html + +
+ +.. _dummy_header: + +`NVIDIA NeMo `_ is a toolkit for building new state-of-the-art +conversational AI models. NeMo has separate collections for Automatic Speech Recognition (ASR), +Natural Language Processing (NLP), and Text-to-Speech (TTS) models. Each collection consists of +prebuilt modules that include everything needed to train on your data. +Every module can easily be customized, extended, and composed to create new conversational AI +model architectures. + +Conversational AI architectures are typically large and require a lot of data and compute +for training. NeMo uses `PyTorch Lightning `_ for easy and performant multi-GPU/multi-node +mixed-precision training. + +.. raw:: html + +
+ +
+ +For more information and questions, visit the `NVIDIA NeMo Discussion Board `_. + +Prerequisites +------------- + +Before you begin using NeMo, it's assumed you meet the following prerequisites. + +#. You have Python version 3.6, 3.7 or 3.8. + +#. You have Pytorch version 1.8.1. + +#. You have access to an NVIDIA GPU for training. + +.. _quick_start_guide: + +Quick Start Guide +----------------- + +This NeMo Quick Start Guide is a starting point for users who want to try out NeMo; specifically, this guide enables users to quickly get started with the NeMo fundamentals by walking you through an example audio translator and voice swap. + +If you're new to NeMo, the best way to get started is to take a look at the following tutorials: + +* `Text Classification (Sentiment Analysis) `__ - demonstrates the Text Classification model using the NeMo NLP collection. +* `NeMo Primer `__ - introduces NeMo, PyTorch Lightning, and OmegaConf, and shows how to use, modify, save, and restore NeMo models. +* `NeMo Models `__ - explains the fundamental concepts of the NeMo model. +* `NeMo voice swap demo `__ - demonstrates how to swap a voice in the audio fragment with a computer generated one using NeMo. + +Below we is the code snippet of Audio Translator application. + +.. code-block:: python + + # Import NeMo and it's ASR, NLP and TTS collections + import nemo + # Import Speech Recognition collection + import nemo.collections.asr as nemo_asr + # Import Natural Language Processing colleciton + import nemo.collections.nlp as nemo_nlp + # Import Speech Synthesis collection + import nemo.collections.tts as nemo_tts + + # Next, we instantiate all the necessary models directly from NVIDIA NGC + # Speech Recognition model - QuartzNet trained on Russian part of MCV 6.0 + quartznet = nemo_asr.models.EncDecCTCModel.from_pretrained(model_name="stt_ru_quartznet15x5").cuda() + # Neural Machine Translation model + nmt_model = nemo_nlp.models.MTEncDecModel.from_pretrained(model_name='nmt_ru_en_transformer6x6').cuda() + # Spectrogram generator which takes text as an input and produces spectrogram + spectrogram_generator = nemo_tts.models.Tacotron2Model.from_pretrained(model_name="tts_en_tacotron2").cuda() + # Vocoder model which takes spectrogram and produces actual audio + vocoder = nemo_tts.models.WaveGlowModel.from_pretrained(model_name="tts_waveglow_88m").cuda() + # Transcribe an audio file + # IMPORTANT: The audio must be mono with 16Khz sampling rate + # Get example from: https://nemo-public.s3.us-east-2.amazonaws.com/mcv-samples-ru/common_voice_ru_19034087.wav + russian_text = quartznet.transcribe(['Path_to_audio_file']) + print(russian_text) + # You should see russian text here. Let's translate it to English + english_text = nmt_model.translate(russian_text) + print(english_text) + # After this you should see English translation + # Let's convert it into audio + # A helper function which combines Tacotron2 and WaveGlow to go directly from + # text to audio + def text_to_audio(text): + parsed = spectrogram_generator.parse(text) + spectrogram = spectrogram_generator.generate_spectrogram(tokens=parsed) + audio = vocoder.convert_spectrogram_to_audio(spec=spectrogram) + return audio.to('cpu').numpy() + audio = text_to_audio(english_text[0]) + + +Installation +------------ + +Pip +~~~ +Use this installation mode if you want the latest released version. + +.. code-block:: bash + + apt-get update && apt-get install -y libsndfile1 ffmpeg + pip install Cython + pip install nemo_toolkit[all] + +Pip from source +~~~~~~~~~~~~~~~ +Use this installation mode if you want the version from a particular GitHub branch (for example, ``main``). + +.. code-block:: bash + + apt-get update && apt-get install -y libsndfile1 ffmpeg + pip install Cython + python -m pip install git+https://github.com/NVIDIA/NeMo.git@{BRANCH}#egg=nemo_toolkit[all] + # For v1.0.2, replace {BRANCH} with v1.0.2 like so: + # python -m pip install git+https://github.com/NVIDIA/NeMo.git@v1.0.2#egg=nemo_toolkit[all] + +From source +~~~~~~~~~~~ +Use this installation mode if you are contributing to NeMo. + +.. code-block:: bash + + apt-get update && apt-get install -y libsndfile1 ffmpeg + git clone https://github.com/NVIDIA/NeMo + cd NeMo + ./reinstall.sh + +Docker containers +~~~~~~~~~~~~~~~~~ +If you chose to work with the ``main`` branch, we recommend using `NVIDIA's PyTorch container version 21.05-py3 `_, then install from GitHub. + +.. code-block:: bash + + docker run --gpus all -it --rm -v :/NeMo --shm-size=8g \ + -p 8888:8888 -p 6006:6006 --ulimit memlock=-1 --ulimit \ + stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:21.05-py3 + + +FAQ +--- +Have a look at our `discussions board `_ and feel free to post a question or start a discussion. + + +Contributing +------------ + +We welcome community contributions! Refer to the `CONTRIBUTING.md `_ file for the process. + +License +------- + +NeMo is under `Apache 2.0 license `_. \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/starthere/tutorials.rst b/models/tts/NeMo/docs/source/starthere/tutorials.rst new file mode 100644 index 0000000000..76b4dd744e --- /dev/null +++ b/models/tts/NeMo/docs/source/starthere/tutorials.rst @@ -0,0 +1,117 @@ +.. _tutorials: + +Tutorials +========= + +The best way to get started with NeMo is to start wtih one of our tutorials. + +Most NeMo tutorials can be run on `Google's Colab `_. + +To run a tutorial: + +#. Click the **Colab** link (see table below). +#. Connect to an instance with a GPU. For example, click **Runtime** > **Change runtime type** and select **GPU** for the hardware accelerator. + +.. list-table:: **Tutorials** + :widths: 15 25 25 + :header-rows: 1 + + * - Domain + - Title + - GitHub URL + * - General + - Getting Started: Exploring Nemo Fundamentals + - `NeMo Fundamentals `_ + * - General + - Getting Started: Sample Conversational AI application + - `Audio translator example `_ + * - General + - Getting Started: Voice swap application + - `Voice swap example `_ + * - General + - Exploring NeMo Model Construction + - `NeMo Models `_ + * - ASR + - ASR with NeMo + - `ASR with NeMo `_ + * - ASR + - ASR with Subword Tokenization + - `ASR with Subword Tokenization `_ + * - ASR + - Offline ASR Inference with Beam Search and External Language Model Rescoring + - `Offline ASR `_ + * - ASR + - Online ASR inference with Microphone + - `Online ASR Microphone `_ + * - ASR + - Fine-tuning CTC Models on New Languages + - `ASR CTC Language Fine-Tuning `_ + * - ASR + - Speech Commands + - `Speech Commands `_ + * - ASR + - Online and Offline Speech Commands Inference + - `Online Offline Microphone Speech Commands `_ + * - ASR + - Voice Activiy Detection (VAD) + - `Voice Activiy Detection `_ + * - ASR + - Online and Offline VAD Inference + - `Online Offline Microphone VAD `_ + * - ASR + - Speaker Recognition and Verification + - `Speaker Recognition and Verification `_ + * - ASR + - Speaker Diarization Inference + - `Speaker Diarization Inference `_ + * - ASR + - ASR with Speaker Diarization + - `ASR with Speaker Diarization `_ + * - ASR + - Online Noise Augmentation + - `Online Noise Augmentation `_ + * - NLP + - Using Pretrained Language Models for Downstream Tasks + - `Pretrained Language Models for Downstream Tasks `_ + * - NLP + - Exploring NeMo NLP Tokenizers + - `NLP Tokenizers `_ + * - NLP + - Text Classification (Sentiment Analysis) with BERT + - `Text Classification (Sentiment Analysis) `_ + * - NLP + - Question Answering with SQuAD + - `Question Answering Squad `_ + * - NLP + - Token Classification (Named Entity Recognition) + - `Token Classification: Named Entity Recognition `_ + * - NLP + - Joint Intent Classification and Slot Filling + - `Joint Intent and Slot Classification `_ + * - NLP + - GLUE Benchmark + - `GLUE Benchmark `_ + * - NLP + - Punctuation and Capitalization + - `Punctuation and Capitalization `_ + * - NLP + - Named Entity Recognition - BioMegatron + - `Named Entity Recognition - BioMegatron `_ + * - NLP + - Relation Extraction - BioMegatron + - `Relation Extraction - BioMegatron `_ + * - TTS + - Speech Synthesis + - `TTS Inference `_ + * - TTS + - Speech Synthesis + - `Tacotron2 Training `_ + * - Tools + - CTC Segmentation + - `CTC Segmentation `_ + * - Text Processing + - Text Normalization for TTS + - `Text Normalization `_ + * - Text Processing + - Inverse Text Normalization for ASR + - `Inverse Text Normalization `_ \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/tools/ctc_segmentation.rst b/models/tts/NeMo/docs/source/tools/ctc_segmentation.rst new file mode 100644 index 0000000000..80e4f6cf85 --- /dev/null +++ b/models/tts/NeMo/docs/source/tools/ctc_segmentation.rst @@ -0,0 +1,20 @@ +Dataset Creation Tool Based on CTC-Segmentation +=============================================== + +This tool provides functionality to align long audio files and the corresponding transcripts into shorter fragments +that are suitable for an Automatic Speech Recognition (ASR) model training. + +More details could be found in `NeMo/tutorials/tools/CTC_Segmentation_Tutorial.ipynb `__ (can be executed with `Google's Colab `_). + +The tool is based on the `CTC-Segmentation `__ package and +`CTC-Segmentation of Large Corpora for German End-to-end Speech Recognition +`__ :cite:`tools-kurzinger2020ctc` + + +References +---------- + +.. bibliography:: tools_all.bib + :style: plain + :labelprefix: TOOLS + :keyprefix: tools- diff --git a/models/tts/NeMo/docs/source/tools/intro.rst b/models/tts/NeMo/docs/source/tools/intro.rst new file mode 100644 index 0000000000..ff7a61fcb1 --- /dev/null +++ b/models/tts/NeMo/docs/source/tools/intro.rst @@ -0,0 +1,14 @@ +Tools +===== + +NeMo provides a set of tools useful for developing Automatic Speech Recognitions (ASR) and Text-to-Speech (TTS) models: \ +`https://github.com/NVIDIA/NeMo/tree/main/tools `__ . + +.. toctree:: + :maxdepth: 1 + + ctc_segmentation + speech_data_explorer + text_processing_deployment + + diff --git a/models/tts/NeMo/docs/source/tools/speech_data_explorer.rst b/models/tts/NeMo/docs/source/tools/speech_data_explorer.rst new file mode 100644 index 0000000000..92d1d0752a --- /dev/null +++ b/models/tts/NeMo/docs/source/tools/speech_data_explorer.rst @@ -0,0 +1,33 @@ +Speech Data Explorer +==================== + +.. note:: + + The tool could be found under `NeMo/tools/speech_data_explorer `__. + +`Dash `__-based tool for interactive exploration of ASR/TTS datasets. + +Features: + +* dataset's statistics (alphabet, vocabulary, duration-based histograms) +* navigation across dataset (sorting, filtering) +* inspection of individual utterances (waveform, spectrogram, audio player) +* errors' analysis (Word Error Rate, Character Error Rate, Word Match Rate, Mean Word Accuracy, diff) + +Please make sure that requirements are installed. Then run: + +.. code:: + + python data_explorer.py path_to_manifest.json + + +JSON manifest file should contain the following fields: + +* `audio_filepath` (path to audio file) +* `duration` (duration of the audio file in seconds) +* `text` (reference transcript) + +Errors' analysis requires "pred_text" (ASR transcript) for all utterances. + +Any additional field will be parsed and displayed in 'Samples' tab. + diff --git a/models/tts/NeMo/docs/source/tools/text_processing_deployment.rst b/models/tts/NeMo/docs/source/tools/text_processing_deployment.rst new file mode 100644 index 0000000000..3b277bdfcf --- /dev/null +++ b/models/tts/NeMo/docs/source/tools/text_processing_deployment.rst @@ -0,0 +1,78 @@ +NeMo Text Processing Deployment +=============================================== + +This tool deploys :doc:`NeMo Inverse Text Normalization (ITN) <../nemo_text_processing/inverse_text_normalization>` and :doc:`NeMo Text Normalization (TN) <../nemo_text_processing/text_normalization>` for production :cite:`tools-itn_deploy-zhang2021nemo`. +It uses `Sparrowhawk `_ :cite:`tools-itn_deploy-sparrowhawk` -- an open-source version of Google Kestrel :cite:`tools-itn_deploy-ebden2015kestrel`. + +Requirements +------------------------ + +:doc:`nemo_text_processing <../nemo_text_processing/intro>` package + + +Usage +------------ + +Starts docker container with production backend with plugged in grammars. This is entry point script. + +Arguments: +^^^^^^^^^ +* ``GRAMMARS`` - ``tn_grammars`` or ``itn_grammars`` to export either TN or ITN grammars from :doc:`nemo_text_processing <../nemo_text_processing/intro>`. +* ``INPUT_CASE`` - ``cased`` or ``lower_cased`` (lower_cased is supported only in TN grammars). +* ``MODE`` - choose ``test`` to run test on the grammars inside the container. + +For example: + + +.. code-block:: bash + + # to export ITN grammars + bash export_grammar.sh --GRAMMARS=itn_grammars + + # to export and test TN grammars + bash export_grammar.sh --GRAMMARS=itn_grammars --INPUT_CASE=cased --MODE=test + +This script runs the following steps in sequence: + +Exports grammar `ClassifyFst` and `VerbalizeFst` from :doc:`nemo_text_processing <../nemo_text_processing/intro>` to `OUTPUT_DIR/classify/tokenize_and_classify.far` and `OUTPUT_DIR/verbalize/verbalize.far` respectively. + +.. code-block:: bash + + python pynini_export.py <--output_dir OUTPUT_DIR> <--grammars GRAMMARS> <--input_case INPUT_CASE> + +Builds C++ production backend docker + +.. code-block:: bash + + bash docker/build.sh + + +Plugs in grammars into production backend by mounting grammar directory `classify/` and `verbalize/` with sparrowhawk grammar directory inside docker. Returns docker prompt + +.. code-block:: bash + + # to launch container with the exported grammars + bash docker/launch.sh + + # to launch container with the exported grammars and run tests on TN grammars + bash docker/launch.sh test_tn_grammars + + # to launch container with the exported grammars and run tests on ITN grammars + bash docker/launch.sh test_itn_grammars + + +Runs TN or ITN in docker container: + +.. code-block:: bash + + echo "two dollars fifty" | ../../src/bin/normalizer_main --config=sparrowhawk_configuration.ascii_proto + +This returns $2.50 for ITN. + +References +---------- + +.. bibliography:: tools_all.bib + :style: plain + :labelprefix: TOOLS-ITN_DEPLOY + :keyprefix: tools-itn_deploy- diff --git a/models/tts/NeMo/docs/source/tools/tools_all.bib b/models/tts/NeMo/docs/source/tools/tools_all.bib new file mode 100644 index 0000000000..4b86ee2b74 --- /dev/null +++ b/models/tts/NeMo/docs/source/tools/tools_all.bib @@ -0,0 +1,49 @@ +@inproceedings{kurzinger2020ctc, + title={CTC-segmentation of large corpora for german end-to-end speech recognition}, + author={K{\"u}rzinger, Ludwig and Winkelbauer, Dominik and Li, Lujun and Watzel, Tobias and Rigoll, Gerhard}, + booktitle={International Conference on Speech and Computer}, + pages={267--278}, + year={2020}, + organization={Springer} +} + +@article{ebden2015kestrel, + title={The Kestrel TTS text normalization system}, + author={Ebden, Peter and Sproat, Richard}, + journal={Natural Language Engineering}, + volume={21}, + number={3}, + pages={333}, + year={2015}, + publisher={Cambridge University Press} +} + +@article{sproat2016rnn, + title={RNN approaches to text normalization: A challenge}, + author={Sproat, Richard and Jaitly, Navdeep}, + journal={arXiv preprint arXiv:1611.00068}, + year={2016} +} + +@book{taylor2009text, + title={Text-to-speech synthesis}, + author={Taylor, Paul}, + year={2009}, + publisher={Cambridge university press} +} + +@misc{zhang2021nemo, + title={NeMo Inverse Text Normalization: From Development To Production}, + author={Yang Zhang and Evelina Bakhturina and Kyle Gorman and Boris Ginsburg}, + year={2021}, + eprint={2104.05055}, + archivePrefix={arXiv}, + primaryClass={cs.CL} +} + +@inproceedings{sparrowhawk, + title = {TTS for Low Resource Languages: A Bangla Synthesizer}, + author = {Alexander Gutkin and Linne Ha and Martin Jansche and Knot Pipatsrisawat and Richard Sproat}, + booktitle = {10th Language Resources and Evaluation Conference}, + year = {2016}, +} \ No newline at end of file diff --git a/models/tts/NeMo/docs/source/tts/api.rst b/models/tts/NeMo/docs/source/tts/api.rst new file mode 100644 index 0000000000..2255af4110 --- /dev/null +++ b/models/tts/NeMo/docs/source/tts/api.rst @@ -0,0 +1,47 @@ +NeMo TTS Collection API +======================= + +TTS Base Classes +---------------- + +The classes below are the base of the TTS pipeline. +To read more about them, see the `Base Classes <./intro.html#Base Classes>`__ section of the intro page. + +.. autoclass:: nemo.collections.tts.models.base.SpectrogramGenerator + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.tts.models.base.Vocoder + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.tts.models.base.TextToWaveform + :show-inheritance: + :members: + +TTS Datasets +------------ + +.. autoclass:: nemo.collections.tts.data.datalayers.AudioDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.tts.data.datalayers.MelAudioDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.tts.data.datalayers.SplicedAudioDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.tts.data.datalayers.NoisySpecsDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.asr.data.audio_to_text.FastPitchDataset + :show-inheritance: + :members: + +.. autoclass:: nemo.collections.tts.data.datalayers.FastSpeech2Dataset + :show-inheritance: + :members: diff --git a/models/tts/NeMo/docs/source/tts/datasets.rst b/models/tts/NeMo/docs/source/tts/datasets.rst new file mode 100644 index 0000000000..ef3236e703 --- /dev/null +++ b/models/tts/NeMo/docs/source/tts/datasets.rst @@ -0,0 +1,187 @@ +TTS Datasets +============ + +Most TTS scripts work out-of-the-box with the LJ Speech dataset. +However, some models may require supplementary data for training and validation. +The following sections contain details of required data formats and instructions for running preprocessing scripts for +such models. + + +FastSpeech 2 +------------ + +The FastSpeech 2 model converts from phonemes to Mel Spectrograms and predicts phoneme durations, pitches, and +energies. +Therefore, in addition to the manifest and audio samples, it requires some supplementary files and data: + +* A **mappings file** for converting from words to phonemes and phonemes to indices +* **Phoneme durations** (for each audio sample) +* **Pitch per frame** (for each training audio sample) +* **Energy per frame** (for each training audio sample) +* (Optional) **Ignore list** for filtering out samples with OOV words + +The ``FastSpeech2Dataset`` uses the manifest format shared with all other NeMo speech tasks. +Each line of the manifest should describe one sample, and should be in the following format: + +.. code:: + + {"audio_filepath": "/data/wavs/audio.wav", "text": "an unused transcription", "duration": 23.147} + +See the documentation on :ref:`Preparing Custom ASR Data` for more details on the NeMo speech manifest format. + +.. note:: + The ``FastSpeech2Dataset`` ignores the ``"text"`` field of the manifest, since the model reads phoneme indices from + the supplementary duration files instead. + +The **mappings file** should be a JSON-formatted file that contains dictionaries ``word2phones`` and ``phone2idx``. +``word2phones`` is a mapping that converts from your vocabulary to phonemes. +For example, one entry in ``word2phones`` might be ``"nemo": ["N", "EH1", "M", "OW0"]``. +``phone2idx`` is an (arbitrary) mapping that converts each phoneme to an index for prediction. +It is also used to determine the ``pad_id``, which is set to the length of ``phone2idx``. + +Together, these two dictionaries are also used for inference, as an input sentence is first normalized, then converted +to phonemes using ``word2phones``, and then converted from phonemes to indices using ``phone2idx``. + +For each audio sample, there should be a corresponding **phoneme durations** file that contains ``token_duration`` +(duration per phoneme in frames) and ``text_encoded`` (index of each phoneme corresponding to ``phone2idx``). + +For each sample, there should also be a **pitch file** that contains the pitches (F0) per frame, and an **energy file** +that contains the energies (L2-norm of STFT frame amplitudes) per frame. + +The **ignore list** is a pickled list of file base names (with no extension) that tells the ``FastSpeech2Dataset`` +which audio samples to discard. +You should use this to exclude samples from your manifest that contain OOV words. + + +Directory Structure +^^^^^^^^^^^^^^^^^^^ + +While the mappings file and ignore list paths are passed in to the ``FastSpeech2Dataset`` directly, the +``FastSpeech2Dataset`` infers the path to the supplementary files corresponding to each audio sample based on the +audio samples' paths found in the manifest. +For each sample, the paths to the corresponding duration, pitch, and energy files are inferred by replacing ``wavs/`` +with ``phoneme_durations/``, ``pitches/``, and ``energies``, and swapping out the file extension (``.wav``) with +``.pt``, ``.npy``, and ``.npy`` respectively. + +For example, given manifest audio path ``/data/LJSpeech/wavs/LJ001-0001.wav``, the inferred duration and phonemes file +path would be ``/data/LJSpeech/phoneme_durations/LJ001-0001.pt``. + +Your directory structure should therefore look something like this: + +.. code:: + + data/ + ├── manifest.json + ├── mappings.json + ├── ignore_file.pkl + ├── energies/ + │   ├── basename_1.npy + │   ├── basename_2.npy + │   ├── ... + │   └── basename_n.npy + ├── phoneme_durations/ + │   ├── basename_1.pt + │   ├── basename_2.pt + │   ├── ... + │   └── basename_n.pt + ├── pitches/ + │   ├── basename_1.npy + │   ├── basename_2.npy + │   ├── ... + │   └── basename_n.npy + └── wavs/ +    ├── basename_1.wav +    ├── basename_2.wav +    ├── ... +    └── basename_n.wav + + +LJ Speech Dataset +^^^^^^^^^^^^^^^^^ + +NeMo comes with scripts for preprocessing the LJ Speech dataset for use with FastSpeech 2, which can be found in the +``/scripts/dataset_processing/ljspeech/`` directory. +These scripts assume that you have downloaded the LJ Speech dataset and extracted it to ````. + +They perform the following dataset preparation steps: + +* Create manifest files and ``.txt`` files with normalized transcripts +* Extract pitches +* Extract energies +* Downloads the `CMU Pronouncing Dictionary `_ to convert normalized + transcripts to phonemes +* Uses the `Montreal Forced Aligner `_ (MFA) to + perform alignment +* Calculates durations based on the alignment +* Generates the mappings file and ignore file for the dataset + +To run the scripts, follow the steps below. + +#. Download and extract the `LJ Speech dataset `_ to ````. + +#. Create the manifest files and normalized text files (for MFA to discover later) by running: + + .. code-block:: bash + + python create_manifests_and_textfiles.py --ljspeech_base= + +#. Extract pitches and energies from the audio files: + + .. code-block:: bash + + python extract_ljspeech_energy_pitch.py --ljspeech_base= + + This will write the extracted pitches and energies to ``/pitches/`` and + ``/energies/``. + +#. Run the phoneme extraction and duration calculation script. + This script will set up a Conda environment to download and install MFA and its dependencies, so make sure that you + have Anaconda or Miniconda before running the script. + + Also note that MFA does sometimes have trouble finding OpenBlas, so you may have to manually install it with the + command ``sudo apt-get install libopenblas-dev``. + + .. code-block:: bash + + ./extract_ljspeech_phonemes_and_durs.sh + + This script takes additional options ``--skip_env_setup`` if you have already set up the ``aligner`` environment, + and ``--g2p_dict=`` if you have already downloaded CMU dict or another G2P dictionary you prefer. + + The alignment step will take a while to run, so be prepared to wait upwards of an hour. + + In addition to alignments and durations, this script will create the mappings and ignore files as well. + It will also generate some intermediate files that you are free to delete, as you only need the files listed above. + +When the following steps are finished, your ```` directory will look like this. +(Starred files and directories indicate files from intermediate steps that are safe to remove, but that you may want +to keep for bookkeeping purposes.) + +.. code:: + + / + ├── alignments/ * + ├── cmudict.dict * + ├── energies/ + │   ├── LJ001-0001.npy + │   └── ... + ├── ljs_audio_text_test_filelist.txt * + ├── ljs_audio_text_train_filelist.txt * + ├── ljs_audio_text_val_filelist.txt * + ├── ljspeech_test.json + ├── ljspeech_train.json + ├── ljspeech_val.json + ├── mappings.json + ├── metadata.csv + ├── phoneme_durations/ + │   ├── LJ001-0001.pt + │   └── ... + ├── pitches/ + │   ├── LJ001-0001.npy + │   └── ... + ├── README + ├── uncommented_cmudict.dict * + ├── wavs/ + │   ├── LJ001-0001.wav + │   └── ... + └── wavs_to_ignore.pkl diff --git a/models/tts/NeMo/docs/source/tts/intro.rst b/models/tts/NeMo/docs/source/tts/intro.rst new file mode 100644 index 0000000000..7aeb5e1e38 --- /dev/null +++ b/models/tts/NeMo/docs/source/tts/intro.rst @@ -0,0 +1,153 @@ +Speech Synthesis (TTS) +====================== + +.. toctree:: + :maxdepth: 4 + :hidden: + + datasets + api + +Speech Synthesis or Text-to-Speech (TTS) involves turning text into human speech. The NeMo TTS collection currently +supports two pipelines for TTS: + +1) The two stage pipeline. First, a model is used to generate a mel spectrogram from text. Second, a model is used +to generate audio from a mel spectrogram. +2) The "end-to-end" approach that uses one model to generate audio straight from text. + +Quick Start: + +.. code-block:: python + + import soundfile as sf + from nemo.collections.tts.models.base import SpectrogramGenerator, Vocoder + + # Download and load the pretrained tacotron2 model + spec_generator = SpectrogramGenerator.from_pretrained("tts_en_tacotron2") + # Download and load the pretrained waveglow model + vocoder = Vocoder.from_pretrained("tts_waveglow_88m") + + # All spectrogram generators start by parsing raw strings to a tokenized version of the string + parsed = spec_generator.parse("You can type your sentence here to get nemo to produce speech.") + # They then take the tokenized string and produce a spectrogram + spectrogram = spec_generator.generate_spectrogram(tokens=parsed) + # Finally, a vocoder converts the spectrogram to audio + audio = vocoder.convert_spectrogram_to_audio(spec=spectrogram) + + # Save the audio to disk in a file called speech.wav + # Note vocoder return a batch of audio. In this example, we just take the first and only sample. + sf.write("speech.wav", audio.to('cpu').detach().numpy()[0], 22050) + +.. note:: + + For an interactive version of the quick start above, refer to the TTS inference notebook that can be found on the + github readme. + +Available Models +################ + +NeMo supports a variety of models that can be used for TTS. For beginners, we recommend starting with the Tacotron2 + +WaveGlow combination. Then we suggest trying using (TalkNet or FastPitch) + HiFiGAN if you want to continue exploring +the two stage pipeline, or the FastPitch_HifiGan_E2E model for the end-to-end pipeline. + +.. list-table:: *TTS Models* + :widths: 5 5 10 25 + :header-rows: 1 + + * - Model + - Base Class + - Pretrained Checkpoint + - Description + * - Tacotron2 + - :class:`SpectrogramGenerator` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_en_tacotron2 + - LSTM encoder decoder based model that generates spectrograms + * - GlowTTS + - :class:`SpectrogramGenerator` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_en_glowtts + - Glow-based spectrogram generator + * - FastSpeech 2 + - :class:`SpectrogramGenerator` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_en_fastspeech_2 + - Non-autoregressive transformer-based spectrogram generator that predicts duration, energy, and pitch + * - FastPitch + - :class:`SpectrogramGenerator` + - (Coming soon) + - Non-autoregressive transformer-based spectrogram generator that predicts duration and pitch + * - TalkNet + - :class:`SpectrogramGenerator` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_en_talknet + - Non-autoregressive convolution-based spectrogram generator + * - WaveGlow + - :class:`Vocoder` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_waveglow_88m + - Glow-based vocoder + * - SqueezeWave + - :class:`Vocoder` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_squeezewave + - Glow-based vocoder based on WaveGlow but with fewer parameters + * - UniGlow + - :class:`Vocoder` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_uniglow + - Glow-based vocoder based on WaveGlow but shares 1 set of parameters across all flow steps + * - MelGAN + - :class:`Vocoder` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_melgan + - GAN-based vocoder + * - HiFiGAN + - :class:`Vocoder` + - https://ngc.nvidia.com/catalog/models/nvidia:nemo:tts_hifigan + - GAN-based vocoder + * - FastPitch_HifiGan_E2E + - :class:`TextToWaveform` + - (Coming soon) + - End-to-end model based on composing FastPitch and HiFiGAN + * - FastSpeech2_HifiGan_E2E + - :class:`TextToWaveform` + - (Coming soon) + - End-to-end model based on composing FastSpeech2 and HiFiGAN + +Base Classes +############ + +Two Stage Pipeline +****************** + +NeMo TTS has two base classes corresponding to the two stage pipeline: + + - :class:`SpectrogramGenerator` + - :class:`Vocoder` + +The :class:`SpectrogramGenerator` class has two important +functions: :py:meth:`parse` which +accepts raw python strings and returns a torch.tensor that represents tokenized text ready to pass to, and +:py:meth:`generate_spectrogram` which +accepts a batch of tokenized text and returns a torch.tensor that represents a batch of spectrograms + +The :class:`Vocoder` class has one important functions +:py:meth:`convert_spectrogram_to_audio` which +accepts a batch of spectrograms and returns a torch.tensor that represents a batch of raw audio. + +End-to-End Pipeline +******************* + +Correspondingly, NeMo TTS has one base class for the end-to-end pipeline: + + - :class:`TextToWaveform` + +Similarly to the SpectrogramGenerator, :class:`TextToWaveform` +implements two functions: :py:meth:`parse` which +accepts raw python strings and returns a torch.tensor that represents tokenized text ready to pass to, and +:py:meth:`generate_spectrogram` which +accepts a batch of tokenized text and returns a torch.tensor that represents a batch of audio. + +TTS Training +############ + +Training of TTS models can be done using the scripts inside the NeMo ``examples/tts folders``. The majority of the TTS +YAML configurations should work out of the box with the LJSpeech dataset. If you want to train on other data, it is +recommended that you walk through the Tacotron 2 Training notebook. Please pay special attention to the sample rate and +FFT parameters for new data. + +Some models, like FastSpeech 2, require supplementary data such as phoneme durations, pitches, and energies. +For more information about how to preprocess datasets for such models, see the `TTS Datasets <./datasets.html>`__ page. diff --git a/models/tts/NeMo/docs/update_docs_docker.sh b/models/tts/NeMo/docs/update_docs_docker.sh new file mode 100755 index 0000000000..c7a45f8705 --- /dev/null +++ b/models/tts/NeMo/docs/update_docs_docker.sh @@ -0,0 +1,3 @@ +cd ../ +docker run --rm -v $PWD:/workspace python:3.7 /bin/bash -c "cd /workspace && \ +pip install -r requirements/requirements_docs.txt && cd docs/ && rm -rf build && make clean && make html && make html" diff --git a/models/tts/NeMo/examples/asr/asr_webapp/Dockerfile b/models/tts/NeMo/examples/asr/asr_webapp/Dockerfile new file mode 100644 index 0000000000..8e5ee21a7f --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/Dockerfile @@ -0,0 +1,41 @@ +# syntax=docker/dockerfile:experimental + +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +ARG BASE_IMAGE=nvcr.io/nvidia/nemo:1.0.0rc1 + +# build an image that includes only the nemo dependencies, ensures that dependencies +# are included first for optimal caching, and useful for building a development +# image (by specifying build target as `nemo-deps`) +FROM ${BASE_IMAGE} as nemo-deps + +# copy asr service source into a temp directory +WORKDIR /tmp/nemo-service +COPY . . + +# override nemo installation with dependency from requirements.txt +RUN /bin/echo "export BASE_IMAGE=${BASE_IMAGE}" >> /root/.bashrc +RUN cd /tmp/nemo-service && pip install -r "requirements.txt" + +# copy webapp into container for end user +WORKDIR /workspace/nemo-service +COPY . /workspace/nemo-service/ + +# pre-import some asr dependencies which would take time on each worker +RUN python -c "import nemo.collections.asr as nemo_asr" + +# expose port 8000 and launch gunicorn +EXPOSE 8000 +CMD ["gunicorn", "wsgi:app"] diff --git a/models/tts/NeMo/examples/asr/asr_webapp/README.rst b/models/tts/NeMo/examples/asr/asr_webapp/README.rst new file mode 100644 index 0000000000..1f05b5f27e --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/README.rst @@ -0,0 +1,69 @@ +**ASR EVALUATION SERVICE ** +=========================== + +Usage +----- + +A simple service that will load a pre-trained model from NeMo's ASR collection and after uploading some files, will transcribe them. Any pre-trained model (or model which has been uploaded to the service) can be selected and if the local / docker environment supports GPUs then it will be possible to transcribe the audio segments with GPU. + +All uploaded files are immediately deleted after transcriptions are obtained. + + .. note:: + + When using `Gunicorn `_, you might notice that each pretrained checkpoint takes a long time for the first transcription. This is because that worker is instantiating the model into memory and then moving the model onto the GPU. For large models, this step might take a significant amount of time. However, this step is cached, and subsequent transcription requests should be much faster (especially if a GPU is utilized). + +From inside the ``asr_webapp`` directory, there are three options to run the ASR Service to transcribe audio - + +Local Setup +----------- + +1) Install all the dependencies (``pip install -r requirements.txt``). Note : It will utilize the latest branch set in the requirements and will override current nemo installation. + +2) Simply run ``python asr_service.py``. This will launch a single service worker that can be run locally. + +3) The app will run at ``127.0.0.1:5000`` by default. + +Gunicorn Setup +-------------- + +1) Follow Step (1) of the ``Local Setup`` section. + +2) Edit the configuration of ``gunicorn`` inside ``gunicorn.conf.py`` if you wish to change the port or number of workers (though this can be achieved via command line overrides as well). + +3) Simply run ``gunicorn wsgi:app``. This will launch two workers (default) in a gunicorn environment. + +4) The app will run at ``0.0.0.0:8000`` by default. + +Docker Setup +------------ + +The cleanest approach of the three, and requires simply building and running a docker container. + +1) Build the docker by executing ``bash docker_container_build.sh``. This will build a container using the latest branch of nemo. + +2) Run the container by executing ``bash docker_container_run.sh``. This will run a detached container that can be used by visiting ``0.0.0.0:8000`` on a modern browser. + +Note About Uploading Models +--------------------------- + +Uploading models is a useful method to evaluate models quickly without having to write scripts. +By design, models uploaded via the browser are visible to **all** users who have access to the web app - so that there only needs to be one user who uploads the model and everyone else who can access this service will be able to evaluate with that model. + +This also means that models are uploaded indefinitely - and cannot be removed easily without explicitly deleting the ``models`` directory or by shutting down the container (if a container is used). + +There are several reasons this approach was taken - + +* It allows for easy evaluation amongst peers without duplicating model uploads (which may be very big). +* When using ``gunicorn``, every worker has access to a central cache of models which can be loaded into memory. +* Model upload mode is meant to be used in a local, protected environment. + +If you wish to disable model uploading, please open ``templates/main.html`` and comment out the following section: + +.. code-block:: html + + + +
+ NeMo File : + + diff --git a/models/tts/NeMo/examples/asr/asr_webapp/asr_service.py b/models/tts/NeMo/examples/asr/asr_webapp/asr_service.py new file mode 100644 index 0000000000..1c53d1593c --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/asr_service.py @@ -0,0 +1,374 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 atexit +import glob +import os +import shutil +import time +from html import unescape +from uuid import uuid4 + +import model_api +import torch +import werkzeug +from flask import Flask, make_response, render_template, request, url_for +from flask_cors import CORS +from werkzeug.utils import secure_filename + +from nemo.utils import logging + +app = Flask(__name__) +CORS(app) + +# Upload folder for audio files; models are stored in permanent cache +# which gets deleted once the container shuts down +app.config[f'UPLOAD_FOLDER'] = f"tmp/" + + +@app.route('/initialize_model', methods=['POST']) +def initialize_model(): + """ + API Endpoint to instantiate a model + + Loads ASR model by its pretrained checkpoint name or upload ASR model that is provided by the user, + then load that checkpoint into the cache. + + Loading of the model into cache is done once per worker. Number of workers should be limited + so as not to exhaust the GPU memory available on device (if GPU is being used). + """ + logging.info("Starting ASR service") + if torch.cuda.is_available(): + logging.info("CUDA is available. Running on GPU") + else: + logging.info("CUDA is not available. Defaulting to CPUs") + + # get form fields + model_name = request.form['model_names_select'] + use_gpu_if_available = request.form.get('use_gpu_ckbx', "off") + + # get nemo model from user (if not none) + nemo_model_file = request.files.get('nemo_model', '') + + # if nemo model is not None, upload it to model cache + if nemo_model_file != '': + model_name = _store_model(nemo_model_file) + + # Alert user that model has been uploaded into the model cache, + # and they should refresh the page to access the model + result = render_template( + 'toast_msg.html', toast_message=f"Model {model_name} has been uploaded. " f"Refresh page !", timeout=5000 + ) + + else: + # Alert user that model has been loaded onto a workers memory + result = render_template( + 'toast_msg.html', toast_message=f"Model {model_name} has been initialized !", timeout=2000 + ) + + # Load model into memory cache + model_api.initialize_model(model_name=model_name) + + # reset file banner + reset_nemo_model_file_script = """ + + """ + + result = result + reset_nemo_model_file_script + result = make_response(result) + + # set cookies + result.set_cookie("model_name", model_name) + result.set_cookie("use_gpu", use_gpu_if_available) + return result + + +def _store_model(nemo_model_file): + """ + Preserve the model supplied by user into permanent cache + This cache needs to be manually deleted (if run locally), or gets deleted automatically + (when the container gets shutdown / killed). + + Args: + nemo_model_file: User path to .nemo checkpoint. + + Returns: + A file name (with a .nemo) at the end - to signify this is an uploaded checkpoint. + """ + filename = secure_filename(nemo_model_file.filename) + file_basename = os.path.basename(filename) + model_dir = os.path.splitext(file_basename)[0] + + model_store = os.path.join('models', model_dir) + if not os.path.exists(model_store): + os.makedirs(model_store) + + # upload model + model_path = os.path.join(model_store, filename) + nemo_model_file.save(model_path) + return file_basename + + +@app.route('/upload_audio_files', methods=['POST']) +def upload_audio_files(): + """ + API Endpoint to upload audio files for inference. + + The uploaded files must be wav files, 16 KHz sample rate, mono-channel audio samples. + """ + # Try to get one or more files from form + try: + f = request.files.getlist('file') + except werkzeug.exceptions.BadRequestKeyError: + f = None + + # If user did not select any file to upload, notify them. + if f is None or len(f) == 0: + toast = render_template('toast_msg.html', toast_message="No file has been selected to upload !", timeout=2000) + result = render_template('updates/upload_files_failed.html', pre_exec=toast, url=url_for('upload_audio_files')) + result = unescape(result) + return result + + # temporary id to store data + uuid = str(uuid4()) + data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], uuid) + + # If the user attempt to upload another set of files without first transcribing them, + # delete the old cache of files and create a new cache entirely + _remove_older_files_if_exists() + + # Save each file into this unique cache + for fn in f: + filename = secure_filename(fn.filename) + if not os.path.exists(data_store): + os.makedirs(data_store) + + fn.save(os.path.join(data_store, filename)) + logging.info(f"Saving file : {fn.filename}") + + # Update user that N files were uploaded. + msg = f"{len(f)} file(s) uploaded. Click to upload more !" + toast = render_template('toast_msg.html', toast_message=f"{len(f)} file(s) uploaded !", timeout=2000) + result = render_template( + 'updates/upload_files_successful.html', pre_exec=toast, msg=msg, url=url_for('upload_audio_files') + ) + result = unescape(result) + + result = make_response(result) + result.set_cookie("uuid", uuid) + return result + + +def _remove_older_files_if_exists(): + """ + Helper method to prevent cache leakage when user attempts to upload another set of files + without first transcribing the files already uploaded. + """ + # remove old data store (if exists) + old_uuid = request.cookies.get('uuid', '') + if old_uuid is not None and old_uuid != '': + # delete old data store + old_data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], old_uuid) + + logging.info("Tried uploading more data without using old uploaded data. Purging data cache.") + shutil.rmtree(old_data_store, ignore_errors=True) + + +@app.route('/remove_audio_files', methods=['POST']) +def remove_audio_files(): + """ + API Endpoint for removing audio files + + # Note: Sometimes data may persist due to set of circumstances: + + - User uploads audio then closes app without transcribing anything + + In such a case, the files will be deleted when gunicorn shutsdown, or container is stopped. + However the data may not be automatically deleted if the flast server is used as is. + """ + # Get the unique cache id from cookie + uuid = request.cookies.get("uuid", "") + data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], uuid) + + # If the data does not exist (cache is empty), notify user + if not os.path.exists(data_store) or uuid == "": + files_dont_exist = render_template( + 'toast_msg.html', toast_message="No files have been uploaded !", timeout=2000 + ) + result = render_template( + 'updates/remove_files.html', pre_exec=files_dont_exist, url=url_for('remove_audio_files') + ) + result = unescape(result) + return result + + else: + # delete data that exists in cache + shutil.rmtree(data_store, ignore_errors=True) + + logging.info("Removed all data") + + # Notify user that cache was deleted. + toast = render_template('toast_msg.html', toast_message="All files removed !", timeout=2000) + result = render_template('updates/remove_files.html', pre_exec=toast, url=url_for('remove_audio_files')) + result = unescape(result) + + result = make_response(result) + result.set_cookie("uuid", '', expires=0) + return result + + +@app.route('/transcribe', methods=['POST']) +def transcribe(): + """ + API Endpoint to transcribe a set of audio files. + + The files are sorted according to their name, so order may not be same as upload order. + + Utilizing the cached info inside the cookies, a model with selected name will be loaded into memory, + and maybe onto a GPU (if it is supported on the device). + + Then the transcription api will be called from the model_api. If all is successful, a template is updated + with results. If some issue occurs (memory ran out, file is invalid format), notify the user. + """ + # load model name from cookie + model_name = request.cookies.get('model_name') + logging.info(f"Model name : {model_name}") + + # If model name is not selected via Load Model, notify user. + if model_name is None or model_name == '': + result = render_template('toast_msg.html', toast_message="Model has not been initialized !", timeout=2000) + return result + + # load whether gpu should be used + use_gpu_if_available = request.cookies.get('use_gpu') == 'on' + gpu_used = torch.cuda.is_available() and use_gpu_if_available + + # Load audio from paths + uuid = request.cookies.get("uuid", "") + data_store = os.path.join(app.config[f'UPLOAD_FOLDER'], uuid) + + files = list(glob.glob(os.path.join(data_store, "*.wav"))) + + # If no files found in cache, notify user + if len(files) == 0: + result = render_template('toast_msg.html', toast_message="No audio files were found !", timeout=2000) + return result + + # transcribe file via model api + t1 = time.time() + transcriptions = model_api.transcribe_all(files, model_name, use_gpu_if_available=use_gpu_if_available) + t2 = time.time() + + # delete all transcribed files immediately + for fp in files: + try: + os.remove(fp) + except FileNotFoundError: + logging.info(f"Failed to delete transcribed file : {os.path.basename(fp)}") + + # delete temporary transcription directory + shutil.rmtree(data_store, ignore_errors=True) + + # If something happened during transcription, and it failed, notify user. + if type(transcriptions) == str and transcriptions == model_api.TAG_ERROR_DURING_TRANSCRIPTION: + toast = render_template( + 'toast_msg.html', + toast_message=f"Failed to transcribe files due to unknown reason. " + f"Please provide 16 KHz Monochannel wav files onle.", + timeout=5000, + ) + transcriptions = ["" for _ in range(len(files))] + + else: + # Transcriptions obtained successfully, notify user. + toast = render_template( + 'toast_msg.html', + toast_message=f"Transcribed {len(files)} files using {model_name} (gpu={gpu_used}), " + f"in {(t2 - t1): 0.2f} s", + timeout=5000, + ) + + # Write results to data table + results = [] + for filename, transcript in zip(files, transcriptions): + results.append(dict(filename=os.path.basename(filename), transcription=transcript)) + + result = render_template('transcripts.html', transcripts=results) + result = toast + result + result = unescape(result) + + result = make_response(result) + result.set_cookie("uuid", "", expires=0) + return result + + +def remove_tmp_dir_at_exit(): + """ + Helper method to attempt a deletion of audio file cache on flask api exit. + Gunicorn and Docker container (based on gunicorn) will delete any remaining files on + shutdown of the gunicorn server or the docker container. + + This is a patch that might not always work for Flask server, but in general should ensure + that local audio file cache is deleted. + + This does *not* impact the model cache. Flask and Gunicorn servers will *never* delete uploaded models. + Docker container will delete models *only* when the container is killed (since models are uploaded to + local storage path inside container). + """ + try: + uuid = request.cookies.get("uuid", "") + + if uuid is not None or uuid != "": + cache_dir = os.path.join(os.path.join(app.config[f'UPLOAD_FOLDER'], uuid)) + logging.info(f"Removing cache file for worker : {os.getpid()}") + + if os.path.exists(cache_dir): + shutil.rmtree(cache_dir, ignore_errors=True) + logging.info(f"Deleted tmp folder : {cache_dir}") + + except RuntimeError: + # Working outside of request context (probably shutdown) + # simply delete entire tmp folder + shutil.rmtree(app.config[f'UPLOAD_FOLDER'], ignore_errors=True) + + +@app.route('/') +def main(): + """ + API Endpoint for ASR Service. + """ + nemo_model_names, local_model_names = model_api.get_model_names() + model_names = [] + model_names.extend(local_model_names) # prioritize local models + model_names.extend(nemo_model_names) # attach all other pretrained models + + # page initializations + result = render_template('main.html', model_names=model_names) + result = make_response(result) + + # Reset cookies + result.set_cookie("model_name", '', expires=0) # model name from pretrained model list + result.set_cookie("use_gpu", '', expires=0) # flag to use gpu (if available) + result.set_cookie("uuid", '', expires=0) # session id + return result + + +# Register hook to delete file cache (for flask server only) +atexit.register(remove_tmp_dir_at_exit) + + +if __name__ == '__main__': + app.run(debug=True) diff --git a/models/tts/NeMo/examples/asr/asr_webapp/docker_container_build.sh b/models/tts/NeMo/examples/asr/asr_webapp/docker_container_build.sh new file mode 100644 index 0000000000..b94987e710 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/docker_container_build.sh @@ -0,0 +1,19 @@ +#!/bin/bash + +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# build container +DOCKER_BUILDKIT=1 docker build -f Dockerfile --progress=plain -t nemo-asr-service:latest . + diff --git a/models/tts/NeMo/examples/asr/asr_webapp/docker_container_run.sh b/models/tts/NeMo/examples/asr/asr_webapp/docker_container_run.sh new file mode 100644 index 0000000000..6ea348c47f --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/docker_container_run.sh @@ -0,0 +1,20 @@ +#!/bin/bash + +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# run container +docker run -d --gpus=all -p 8000:8000 --shm-size=8g --ulimit memlock=-1 --ulimit \ + stack=67108864 nemo-asr-service:latest + diff --git a/models/tts/NeMo/examples/asr/asr_webapp/gunicorn.conf.py b/models/tts/NeMo/examples/asr/asr_webapp/gunicorn.conf.py new file mode 100644 index 0000000000..1fc0e361d4 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/gunicorn.conf.py @@ -0,0 +1,36 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# https://docs.gunicorn.org/en/stable/settings.html + +# NOTE: Do not import nemo / torch code here +# Gunicorn creates forked processes - and CUDA cannot be used in forked multiprocessing environment. + +import shutil + +# General config +bind = "0.0.0.0:8000" +workers = 2 + +# Worker specific config +worker_connections = 1000 +timeout = 180 # 3 minutes of timeout + + +def on_exit(server): + # delete tmp dir + print("Shutting down server ...") + print("Deleteing tmp directory ...") + shutil.rmtree('tmp/', ignore_errors=True) + print("Tmp directory deleted !") diff --git a/models/tts/NeMo/examples/asr/asr_webapp/model_api.py b/models/tts/NeMo/examples/asr/asr_webapp/model_api.py new file mode 100644 index 0000000000..e14bcb3a8a --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/model_api.py @@ -0,0 +1,131 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 contextlib +import glob +import os + +import torch + +import nemo.collections.asr as nemo_asr +from nemo.utils import logging, model_utils + +# setup AMP (optional) +if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'): + logging.info("AMP enabled!\n") + autocast = torch.cuda.amp.autocast +else: + + @contextlib.contextmanager + def autocast(): + yield + + +MODEL_CACHE = {} + +# Special tags for fallbacks / user notifications +TAG_ERROR_DURING_TRANSCRIPTION = "" + + +def get_model_names(): + # Populate local copy of models + local_model_paths = glob.glob(os.path.join('models', "**", "*.nemo"), recursive=True) + local_model_names = list(sorted([os.path.basename(path) for path in local_model_paths])) + + # Populate with pretrained nemo checkpoint list + nemo_model_names = set() + for model_info in nemo_asr.models.ASRModel.list_available_models(): + for superclass in model_info.class_.mro(): + if 'CTC' in superclass.__name__ or 'RNNT' in superclass.__name__: + nemo_model_names.add(model_info.pretrained_model_name) + break + nemo_model_names = list(sorted(nemo_model_names)) + return nemo_model_names, local_model_names + + +def initialize_model(model_name): + # load model + if model_name not in MODEL_CACHE: + if '.nemo' in model_name: + # use local model + model_name_no_ext = os.path.splitext(model_name)[0] + model_path = os.path.join('models', model_name_no_ext, model_name) + + # Extract config + model_cfg = nemo_asr.models.ASRModel.restore_from(restore_path=model_path, return_config=True) + classpath = model_cfg.target # original class path + imported_class = model_utils.import_class_by_path(classpath) # type: ASRModel + logging.info(f"Restoring local model : {imported_class.__name__}") + + # load model from checkpoint + model = imported_class.restore_from(restore_path=model_path, map_location='cpu') # type: ASRModel + + else: + # use pretrained model + model = nemo_asr.models.ASRModel.from_pretrained(model_name, map_location='cpu') + + model.freeze() + + # cache model + MODEL_CACHE[model_name] = model + + model = MODEL_CACHE[model_name] + return model + + +def transcribe_all(filepaths, model_name, use_gpu_if_available=True): + # instantiate model + if model_name in MODEL_CACHE: + model = MODEL_CACHE[model_name] + else: + model = initialize_model(model_name) + + if torch.cuda.is_available() and use_gpu_if_available: + model = model.cuda() + + # transcribe audio + logging.info("Begin transcribing audio...") + try: + with autocast(): + with torch.no_grad(): + transcriptions = model.transcribe(filepaths, batch_size=32) + + except RuntimeError: + # Purge the cache to clear some memory + MODEL_CACHE.clear() + + logging.info("Ran out of memory on device - performing inference on CPU for now") + + try: + model = model.cpu() + + with torch.no_grad(): + transcriptions = model.transcribe(filepaths, batch_size=32) + + except Exception as e: + logging.info(f"Exception {e} occured while attemting to transcribe audio. Returning error message") + return TAG_ERROR_DURING_TRANSCRIPTION + + logging.info(f"Finished transcribing {len(filepaths)} files !") + + # If RNNT models transcribe, they return a tuple (greedy, beam_scores) + if type(transcriptions[0]) == list and len(transcriptions) == 2: + # get greedy transcriptions only + transcriptions = transcriptions[0] + + # Force onto CPU + model = model.cpu() + MODEL_CACHE[model_name] = model + + return transcriptions diff --git a/models/tts/NeMo/examples/asr/asr_webapp/requirements.txt b/models/tts/NeMo/examples/asr/asr_webapp/requirements.txt new file mode 100644 index 0000000000..b08f825d02 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/requirements.txt @@ -0,0 +1,4 @@ +-e git+git://github.com/NVIDIA/NeMo.git@v1.0.2#egg=nemo_toolkit[all] +flask +flask-cors +gunicorn diff --git a/models/tts/NeMo/examples/asr/asr_webapp/templates/main.html b/models/tts/NeMo/examples/asr/asr_webapp/templates/main.html new file mode 100644 index 0000000000..b6bafcb2da --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/templates/main.html @@ -0,0 +1,150 @@ + + + + + + + + Automatic Speech Recognition + + + + + + + + + + + + + + + +
+
+
+ + NeMo Automatic Speech Recognition Demo + +
+ + +
+
+
+ ASR Transcription + +
+
+
+ + +
+ +
+ +
Model Name : + + +
+ +
+ + + + + + + +  
+ NeMo File : + + + + +
+ + + + +
+ +
+
+ + +
+ + + + + +
+
+
+ + + + +
+ + + + +
+
+ +
+ +
+
+
+ + + \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/asr_webapp/templates/toast_msg.html b/models/tts/NeMo/examples/asr/asr_webapp/templates/toast_msg.html new file mode 100644 index 0000000000..794d0e7a4a --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/templates/toast_msg.html @@ -0,0 +1,28 @@ + + + \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/asr_webapp/templates/transcripts.html b/models/tts/NeMo/examples/asr/asr_webapp/templates/transcripts.html new file mode 100644 index 0000000000..5da42c7a22 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/templates/transcripts.html @@ -0,0 +1,43 @@ + + + + + + + + + + + {% for transcript in transcripts %} + + + + + {% endfor %} + +
FilenameTranscript
+
+ {{ transcript.filename }} +
+ 		+
+ +
+
\ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/remove_files.html b/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/remove_files.html new file mode 100644 index 0000000000..185f5d8b84 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/remove_files.html @@ -0,0 +1,22 @@ + + +{{ pre_exec }} + \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_failed.html b/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_failed.html new file mode 100644 index 0000000000..61ad36ba33 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_failed.html @@ -0,0 +1,22 @@ + + +{{ pre_exec }} + \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_successful.html b/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_successful.html new file mode 100644 index 0000000000..c3d15da2d4 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/templates/updates/upload_files_successful.html @@ -0,0 +1,22 @@ + + +{{ pre_exec }} + \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/asr_webapp/wsgi.py b/models/tts/NeMo/examples/asr/asr_webapp/wsgi.py new file mode 100644 index 0000000000..64410091e4 --- /dev/null +++ b/models/tts/NeMo/examples/asr/asr_webapp/wsgi.py @@ -0,0 +1,18 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from asr_service import app + +if __name__ == '__main__': + app.run() diff --git a/models/tts/NeMo/examples/asr/conf/citrinet/citrinet_384.yaml b/models/tts/NeMo/examples/asr/conf/citrinet/citrinet_384.yaml new file mode 100644 index 0000000000..1ba7085208 --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/citrinet/citrinet_384.yaml @@ -0,0 +1,428 @@ +# This config contains the default values for training a Citrinet model with CTC loss and BPE-based vocabulary. +# Default learning parameters in this config are set for effective batch size of 1k on 32 GPUs. +# To train it with smaller batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches. +# If training for a short time, you can also reduce weight decay to 0. + +# Training Recipe +# This model can be trained using the default settings in this config with FP32 precision. +# When training under AMP, increase `warmup_steps` to 5000 for stable training. +# In order to create Citrinet-C, find-replace `filters: 384` with `filters: C`. +# When reducing the receptive field of these models, it is advised to reduce the amount of augmentation +# for larger models from 10x time masking to 5x or 2x time masking. +# For further details regarding Citrinet, visit - https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#citrinet + +name: &name "Citrinet-384-8x-Stride" + +model: + sample_rate: &sample_rate 16000 + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + trim_silence: false + max_duration: 16.7 + shuffle: true + is_tarred: false + tarred_audio_filepaths: null + use_start_end_token: false + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: false + use_start_end_token: false + + test_ds: + manifest_filepath: null + sample_rate: 16000 + batch_size: 32 + shuffle: false + use_start_end_token: false + + model_defaults: + repeat: 5 + dropout: 0.0 + separable: true + se: true + se_context_size: -1 + kernel_size_factor: 1.0 + + tokenizer: + dir: ??? # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe) + type: ??? # Can be either bpe or wpe + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + sample_rate: *sample_rate + normalize: "per_feature" + window_size: 0.025 + window_stride: 0.01 + window: "hann" + features: &n_mels 80 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + pad_to: 16 + stft_conv: false + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 + time_masks: 2 + freq_width: 27 + time_width: 0.05 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 384 + repeat: 1 + kernel: [5] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [11] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [19] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [21] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [13] + stride: [2] # *stride + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [19] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [21] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [23] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [25] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [25] + stride: [2] # stride + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [27] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [29] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [31] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [33] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [35] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [37] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 384 + repeat: ${model.model_defaults.repeat} + kernel: [39] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: &enc_final 640 + repeat: 1 + kernel: [41] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_final + num_classes: -1 # filled with vocabulary size from tokenizer at runtime + vocabulary: [] # filled with vocabulary from tokenizer at runtime + + optim: + name: novograd + lr: 0.05 + + # optimizer arguments + betas: [0.8, 0.25] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # scheduler config override + warmup_steps: 1000 + warmup_ratio: null + min_lr: 1e-5 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 100 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + log_every_n_steps: 100 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + check_val_every_n_epoch: 1 + precision: 32 + sync_batchnorm: false + benchmark: false + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: true + create_checkpoint_callback: true + checkpoint_callback_params: + monitor: "val_wer" + mode: "min" + save_top_k: 3 + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null + entity: null + resume_if_exists: false + resume_ignore_no_checkpoint: false + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/citrinet/citrinet_512.yaml b/models/tts/NeMo/examples/asr/conf/citrinet/citrinet_512.yaml new file mode 100644 index 0000000000..2aab672008 --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/citrinet/citrinet_512.yaml @@ -0,0 +1,428 @@ +# This config contains the default values for training a Citrinet model with CTC loss and BPE-based vocabulary. +# Default learning parameters in this config are set for effective batch size of 1k on 32 GPUs. +# To train it with smaller batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches. +# If training for a short time, you can also reduce weight decay to 0. + +# Training Recipe +# This model can be trained using the default settings in this config with FP32 precision. +# When training under AMP, increase `warmup_steps` to 5000 for stable training. +# In order to create Citrinet-C, find-replace `filters: 384` with `filters: C`. +# When reducing the receptive field of these models, it is advised to reduce the amount of augmentation +# for larger models from 10x time masking to 5x or 2x time masking. +# For further details regarding Citrinet, visit - https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#citrinet + +name: &name "Citrinet-512-8x-Stride" + +model: + sample_rate: &sample_rate 16000 + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + trim_silence: false + max_duration: 16.7 + shuffle: true + is_tarred: false + tarred_audio_filepaths: null + use_start_end_token: false + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: false + use_start_end_token: false + + test_ds: + manifest_filepath: null + sample_rate: 16000 + batch_size: 32 + shuffle: false + use_start_end_token: false + + model_defaults: + repeat: 5 + dropout: 0.1 + separable: true + se: true + se_context_size: -1 + kernel_size_factor: 1.0 + + tokenizer: + dir: ??? # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe) + type: ??? # Can be either bpe or wpe + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + sample_rate: *sample_rate + normalize: "per_feature" + window_size: 0.025 + window_stride: 0.01 + window: "hann" + features: &n_mels 80 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + pad_to: 16 + stft_conv: false + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 + time_masks: 10 + freq_width: 27 + time_width: 0.05 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 512 + repeat: 1 + kernel: [5] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [11] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [19] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [21] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [13] + stride: [2] # *stride + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [19] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [21] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [23] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [25] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [25] + stride: [2] # stride + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [27] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [29] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [31] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [33] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [35] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [37] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: 512 + repeat: ${model.model_defaults.repeat} + kernel: [39] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + - filters: &enc_final 640 + repeat: 1 + kernel: [41] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + kernel_size_factor: ${model.model_defaults.kernel_size_factor} + + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_final + num_classes: -1 # filled with vocabulary size from tokenizer at runtime + vocabulary: [] # filled with vocabulary from tokenizer at runtime + + optim: + name: novograd + lr: 0.05 + + # optimizer arguments + betas: [0.8, 0.25] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # scheduler config override + warmup_steps: 1000 + warmup_ratio: null + min_lr: 1e-5 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 100 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + log_every_n_steps: 100 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + check_val_every_n_epoch: 1 + precision: 32 + sync_batchnorm: false + benchmark: false + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: true + create_checkpoint_callback: true + checkpoint_callback_params: + monitor: "val_wer" + mode: "min" + save_top_k: 3 + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null + entity: null + resume_if_exists: false + resume_ignore_no_checkpoint: false + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/citrinet/config_bpe.yaml b/models/tts/NeMo/examples/asr/conf/citrinet/config_bpe.yaml new file mode 100644 index 0000000000..4b7a7bfb8e --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/citrinet/config_bpe.yaml @@ -0,0 +1,183 @@ +name: &name "ContextNet5x1" +sample_rate: &sample_rate 16000 +repeat: &repeat 1 +dropout: &dropout 0.0 +separable: &separable true + +model: + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + trim_silence: True + max_duration: 16.7 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: False + + tokenizer: + dir: ??? # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe) + type: ??? # Can be either bpe or wpe + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [13] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [15] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [17] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [19] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: 1 + kernel: [21] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: *separable + se: true + se_context_size: -1 + + - filters: &enc_feat_out 1024 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: *separable + se: true + se_context_size: -1 + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: 1024 + num_classes: -1 # filled with vocabulary size from tokenizer at runtime + vocabulary: [] # filled with vocabulary from tokenizer at runtime + + optim: + name: adam + # _target_: nemo.core.optim.optimizers.Adam + lr: .1 + + # optimizer arguments + betas: [0.9, 0.999] + weight_decay: 0.0001 + + # scheduler setup + sched: + name: CosineAnnealing + + # scheduler config override + warmup_steps: null + warmup_ratio: 0.05 + min_lr: 1e-6 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/config.yaml b/models/tts/NeMo/examples/asr/conf/config.yaml new file mode 100644 index 0000000000..f18a23fc38 --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/config.yaml @@ -0,0 +1,183 @@ +name: &name "QuartzNet15x5" +sample_rate: &sample_rate 16000 +repeat: &repeat 1 +dropout: &dropout 0.0 +separable: &separable true +labels: &labels [" ", "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", "'"] + +model: + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + trim_silence: True + max_duration: 16.7 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + shuffle: False + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + stft_conv: false + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [13] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [15] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [17] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: *repeat + kernel: [19] + stride: [1] + dilation: [1] + dropout: *dropout + residual: true + separable: *separable + se: true + se_context_size: -1 + + - filters: 256 + repeat: 1 + kernel: [21] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: *separable + se: true + se_context_size: -1 + + - filters: &enc_feat_out 1024 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: *separable + se: true + se_context_size: -1 + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: 1024 + num_classes: 28 + vocabulary: *labels + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: .01 + # optimizer arguments + betas: [0.8, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + + # Scheduler params + warmup_steps: null + warmup_ratio: null + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/conformer/conformer_ctc_bpe.yaml b/models/tts/NeMo/examples/asr/conf/conformer/conformer_ctc_bpe.yaml new file mode 100644 index 0000000000..1560fc36ed --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/conformer/conformer_ctc_bpe.yaml @@ -0,0 +1,189 @@ +# It contains the default values for training a Conformer-CTC ASR model, large size (~120M) with CTC loss and sub-word encoding. + +# Architecture and training config: +# Default learning parameters in this config are set for effective batch size of 2K. To train it with smaller effective +# batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches. +# Here are the recommended configs for different variants of Conformer-CTC, other parameters are the same as in this config file. +# One extra layer (compared to original paper) is added to the medium and large variants to compensate for replacing the LSTM decoder with a linear one. +# +# +-------------+---------+---------+----------+------------+-----+ +# | Model | d_model | n_heads | n_layers | time_masks | lr | +# +=============+=========+========+===========+============+=====+ +# | Small (13M)| 176 | 4 | 16 | 5 | 5.0 | +# +-------------+---------+--------+-----------+------------+-----+ +# | Medium (30M)| 256 | 4 | 18 | 5 | 5.0 | +# +-------------+---------+--------+-----------+------------+-----+ +# | Large (121M)| 512 | 8 | 18 | 10 | 2.0 | +# +---------------------------------------------------------------+ +# +# If you do not want to train with AMP, you may use weight decay of 0.0 or reduce the number of time maskings to 2 +# with time_width=100. It may help when you want to train for fewer epochs and need faster convergence. +# With weight_decay=0.0, learning rate may need to get reduced to 2.0. + +# You may find more info about Conformer-CTC here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc +# Pre-trained models of Conformer-CTC can be found here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/results.html +# The checkpoint of the large model trained on LibriSpeech with this recipe can be found here: https://ngc.nvidia.com/catalog/models/nvidia:nemo:stt_en_conformer_ctc_large_ls + +name: "Conformer-CTC-BPE" + +model: + sample_rate: 16000 + log_prediction: true # enables logging sample predictions in the output during training + ctc_reduction: 'mean_batch' + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: true + num_workers: 8 + pin_memory: true + use_start_end_token: false + trim_silence: false + max_duration: 16.7 # it is set for LibriSpeech, you may need to update it for your dataset + min_duration: 0.1 + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: false + num_workers: 8 + pin_memory: true + use_start_end_token: false + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: false + num_workers: 8 + pin_memory: true + use_start_end_token: false + + # recommend small vocab size of 128 or 256 when using 4x sub-sampling + # you may find more detail on how to train a tokenizer at: /scripts/tokenizers/process_asr_text_tokenizer.py + tokenizer: + dir: ??? # path to directory which contains either tokenizer.model (bpe) or vocab.txt (wpe) + type: wpe # Can be either bpe (SentencePiece tokenizer) or wpe (WordPiece tokenizer) + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + sample_rate: ${model.sample_rate} + normalize: "per_feature" + window_size: 0.025 + window_stride: 0.01 + window: "hann" + features: 80 + n_fft: 512 + log: true + frame_splicing: 1 + dither: 0.00001 + pad_to: 0 + pad_value: 0.0 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 # set to zero to disable it + # you may use lower time_masks for smaller models to have a faster convergence + time_masks: 10 # set to zero to disable it + freq_width: 27 + time_width: 0.05 + + encoder: + _target_: nemo.collections.asr.modules.ConformerEncoder + feat_in: ${model.preprocessor.features} + feat_out: -1 # you may set it if you need different output size other than the default d_model + n_layers: 18 + d_model: 512 + + # Sub-sampling params + subsampling: striding # vggnet or striding, vggnet may give better results but needs more memory + subsampling_factor: 4 # must be power of 2 + subsampling_conv_channels: -1 # -1 sets it to d_model + + # Feed forward module's params + ff_expansion_factor: 4 + + # Multi-headed Attention Module's params + self_attention_model: rel_pos # rel_pos or abs_pos + n_heads: 8 # may need to be lower for smaller d_models + # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention + att_context_size: [-1, -1] # -1 means unlimited context + xscaling: true # scales up the input embeddings by sqrt(d_model) + untie_biases: true # unties the biases of the TransformerXL layers + pos_emb_max_len: 5000 + + # Convolution module's params + conv_kernel_size: 31 + + ### regularization + dropout: 0.1 # The dropout used in most of the Conformer Modules + dropout_emb: 0.0 # The dropout used for embeddings + dropout_att: 0.1 # The dropout for multi-headed attention modules + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: null + num_classes: -1 + vocabulary: [] + + optim: + name: adamw + lr: 2.0 + # optimizer arguments + betas: [0.9, 0.98] + # less necessity for weight_decay as we already have large augmentations with SpecAug + # you may need weight_decay for large models, stable AMP training, small datasets, or when lower augmentations are used + # weight decay of 0.0 with lr of 2.0 also works fine + weight_decay: 1e-3 + + # scheduler setup + sched: + name: NoamAnnealing + d_model: ${model.encoder.d_model} + # scheduler config override + warmup_steps: 10000 + warmup_ratio: null + min_lr: 1e-6 + +trainer: + gpus: -1 # number of GPUs, -1 would use all available GPUs + num_nodes: 1 + max_epochs: 1000 + max_steps: null # computed at runtime if not set + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + accelerator: ddp + accumulate_grad_batches: 1 + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP. + log_every_n_steps: 10 # Interval of logging. + progress_bar_refresh_rate: 10 + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs + sync_batchnorm: true + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: true + create_checkpoint_callback: true + checkpoint_callback_params: + # in case of multiple validation sets, first one is used + monitor: "val_wer" + mode: "min" + save_top_k: 3 + + # you need to set these two to True to continue the training + resume_if_exists: false + resume_ignore_no_checkpoint: false + + # You may use this section to create a W&B logger + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null diff --git a/models/tts/NeMo/examples/asr/conf/conformer/conformer_ctc_char.yaml b/models/tts/NeMo/examples/asr/conf/conformer/conformer_ctc_char.yaml new file mode 100644 index 0000000000..735b075d6b --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/conformer/conformer_ctc_char.yaml @@ -0,0 +1,165 @@ +# It contains the default values for training a Conformer-CTC ASR model, large size (~120M) with CTC loss and char-based vocabulary. +# Char-based encoding may give lower accuracy than sub-word encoding (conformer_ctc_bpe.yaml). + +# You may find more detail on Conformer-CTC at `examples/asr/conf/conformer/conformer_ctc_bpe.yaml` + +name: "Conformer-CTC-Char" + +model: + sample_rate: 16000 + labels: [" ", "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", "'"] + log_prediction: true # enables logging sample predictions in the output during training + ctc_reduction: 'mean_batch' + + train_ds: + manifest_filepath: ??? + labels: ${model.labels} + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: true + num_workers: 8 + pin_memory: true + trim_silence: false + max_duration: 16.7 # it is set for LibriSpeech, you may need to update it for your dataset + min_duration: 0.1 + + validation_ds: + manifest_filepath: ??? + labels: ${model.labels} + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: false + num_workers: 8 + pin_memory: true + + test_ds: + manifest_filepath: null + labels: ${model.labels} + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: false + num_workers: 8 + pin_memory: true + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + sample_rate: ${model.sample_rate} + normalize: "per_feature" + window_size: 0.025 + window_stride: 0.01 + window: "hann" + features: 80 + n_fft: 512 + log: true + frame_splicing: 1 + dither: 0.00001 + pad_to: 0 + pad_value: 0.0 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 # set to zero to disable it + # you may use lower time_masks for smaller models to have a faster convergence + time_masks: 10 # set to zero to disable it + freq_width: 27 + time_width: 0.05 + + encoder: + _target_: nemo.collections.asr.modules.ConformerEncoder + feat_in: ${model.preprocessor.features} + feat_out: -1 # you may set it if you need different output size other than the default d_model + n_layers: 16 + d_model: 256 + + # Sub-sampling params + subsampling: striding # vggnet or striding, vggnet may give better results but needs more memory + subsampling_factor: 4 # must be power of 2 + subsampling_conv_channels: -1 # set to -1 to make it equal to the d_model + + # Feed forward module's params + ff_expansion_factor: 4 + + # Multi-headed Attention Module's params + self_attention_model: rel_pos # rel_pos or abs_pos + n_heads: 8 # may need to be lower for smaller d_models + # [left, right] specifies the number of steps to be seen from left and right of each step in self-attention + att_context_size: [-1, -1] # -1 means unlimited context + xscaling: true # scales up the input embeddings by sqrt(d_model) + untie_biases: true # unties the biases of the TransformerXL layers + pos_emb_max_len: 5000 + + # Convolution module's params + conv_kernel_size: 31 + + ### regularization + dropout: 0.1 # The dropout used in most of the Conformer Modules + dropout_emb: 0.0 # The dropout used for embeddings + dropout_att: 0.1 # The dropout for multi-headed attention modules + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: null + num_classes: -1 + vocabulary: ${model.labels} + + optim: + name: adamw + lr: 2.0 + # optimizer arguments + betas: [0.9, 0.98] + # less necessity for weight_decay as we already have large augmentations with SpecAug + # you may need weight_decay for large models, stable AMP training, small datasets, or when lower augmentations are used + # weight decay of 0.0 with lr of 2.0 also works fine + weight_decay: 1e-3 + + # scheduler setup + sched: + name: NoamAnnealing + d_model: ${model.encoder.d_model} + # scheduler config override + warmup_steps: 10000 + warmup_ratio: null + min_lr: 1e-6 + +trainer: + gpus: -1 # number of GPUs, -1 would use all available GPUs + num_nodes: 1 + max_epochs: 1000 + max_steps: null # computed at runtime if not set + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + accelerator: ddp + accumulate_grad_batches: 1 + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP. + log_every_n_steps: 10 # Interval of logging. + progress_bar_refresh_rate: 10 + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs + sync_batchnorm: true + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: true + create_checkpoint_callback: true + checkpoint_callback_params: + # in case of multiple validation sets, first one is used + monitor: "val_wer" + mode: "min" + save_top_k: 3 + + # you need to set these two to True to continue the training + resume_if_exists: false + resume_ignore_no_checkpoint: false + + # You may use this section to create a W&B logger + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null diff --git a/models/tts/NeMo/examples/asr/conf/jasper/jasper_10x5dr.yaml b/models/tts/NeMo/examples/asr/conf/jasper/jasper_10x5dr.yaml new file mode 100644 index 0000000000..2f5309aeed --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/jasper/jasper_10x5dr.yaml @@ -0,0 +1,208 @@ +name: &name "Jasper10x5" + +model: + sample_rate: &sample_rate 16000 + labels: &labels [" ", "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", "'"] + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + trim_silence: True + max_duration: 16.7 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + shuffle: False + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + jasper: + - dilation: [1] + dropout: 0.2 + filters: 256 + kernel: [11] + repeat: 1 + residual: false + stride: [2] + - dilation: [1] + dropout: 0.2 + filters: 256 + kernel: [11] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.2 + filters: 256 + kernel: [11] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.2 + filters: 384 + kernel: [13] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.2 + filters: 384 + kernel: [13] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.2 + filters: 512 + kernel: [17] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.2 + filters: 512 + kernel: [17] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.3 + filters: 640 + kernel: [21] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.3 + filters: 640 + kernel: [21] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.3 + filters: 768 + kernel: [25] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [1] + dropout: 0.3 + filters: 768 + kernel: [25] + repeat: 5 + residual: true + residual_dense: true + stride: [1] + - dilation: [2] + dropout: 0.4 + filters: 896 + kernel: [29] + repeat: 1 + residual: false + stride: [1] + - dilation: [1] + dropout: 0.4 + filters: 1024 + kernel: [1] + repeat: 1 + residual: false + stride: [1] + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: 1024 + num_classes: 28 + vocabulary: *labels + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: .01 + # optimizer arguments + betas: [0.8, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # pytorch lightning args + # monitor: val_loss + # reduce_on_plateau: false + + # Scheduler params + warmup_steps: null + warmup_ratio: null + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/marblenet/marblenet_3x2x64.yaml b/models/tts/NeMo/examples/asr/conf/marblenet/marblenet_3x2x64.yaml new file mode 100644 index 0000000000..e3d34cb1f3 --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/marblenet/marblenet_3x2x64.yaml @@ -0,0 +1,180 @@ +name: &name "MarbleNet-3x2x64" + +model: + sample_rate: 16000 + repeat: 2 + dropout: 0.0 + kernel_size_factor: 1.0 + + labels: ['background', 'speech'] + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + augmentor: + shift: + prob: 1.0 + min_shift_ms: -5.0 + max_shift_ms: 5.0 + white_noise: + prob: 1.0 + min_level: -90 + max_level: -46 + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: False + val_loss_idx: 0 + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: False + test_loss_idx: 0 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMFCCPreprocessor + window_size: 0.025 + window_stride: 0.01 + window: "hann" + n_mels: &n_mels 64 + n_mfcc: *n_mels + n_fft: 512 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 + time_masks: 2 + freq_width: 15 + time_width: 25 + rect_masks: 5 + rect_time: 25 + rect_freq: 15 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: false + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 128 + repeat: 1 + kernel: [29] + stride: [1] + dilation: [2] + dropout: ${model.dropout} + residual: false + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: &enc_final_filters 128 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: false + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoderClassification + feat_in: *enc_final_filters + return_logits: true + pooling_type: 'avg' + + optim: + name: sgd + lr: 0.01 + # optimizer arguments + weight_decay: 0.001 + momentum: 0.9 + + # scheduler setup + sched: + name: PolynomialHoldDecayAnnealing + # Scheduler params + power: 2.0 + warmup_ratio: 0.05 + hold_ratio: 0.45 + min_lr: 0.001 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 150 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v1.yaml b/models/tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v1.yaml new file mode 100644 index 0000000000..3016539bc1 --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v1.yaml @@ -0,0 +1,193 @@ +name: &name "MatchboxNet-3x1x64-v1" + +model: + sample_rate: 16000 + timesteps: 128 + repeat: 1 + dropout: 0.0 + kernel_size_factor: 1.0 + + labels_full: ['bed', 'bird', 'cat', 'dog', 'down', 'eight', 'five', 'four', 'go', 'happy', 'house', 'left', 'marvin', + 'nine', 'no', 'off', 'on', 'one', 'right', 'seven', 'sheila', 'six', 'stop', 'three', 'tree', 'two', 'up', + 'wow', 'yes', 'zero'] + + labels_subset: ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go", "unknown", "silence"] + + labels: ${model.labels_full} + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + augmentor: + shift: + prob: 1.0 + min_shift_ms: -5.0 + max_shift_ms: 5.0 + white_noise: + prob: 1.0 + min_level: -90 + max_level: -46 + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: False + val_loss_idx: 0 + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: False + test_loss_idx: 0 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMFCCPreprocessor + window_size: 0.025 + window_stride: 0.01 + window: "hann" + n_mels: &n_mels 64 + n_mfcc: *n_mels + n_fft: 512 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 + time_masks: 2 + freq_width: 15 + time_width: 25 + rect_masks: 5 + rect_time: 25 + rect_freq: 15 + + crop_or_pad_augment: + _target_: nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation + audio_length: ${model.timesteps} + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: false + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 128 + repeat: 1 + kernel: [29] + stride: [1] + dilation: [2] + dropout: ${model.dropout} + residual: false + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: &enc_final_filters 128 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: false + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoderClassification + feat_in: *enc_final_filters + return_logits: true + pooling_type: 'avg' + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: 0.05 + # optimizer arguments + betas: [0.95, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: PolynomialHoldDecayAnnealing + + # Scheduler params + power: 2.0 + warmup_ratio: 0.05 + hold_ratio: 0.45 + min_lr: 0.001 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 200 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v2.yaml b/models/tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v2.yaml new file mode 100644 index 0000000000..5e088b065b --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/matchboxnet/matchboxnet_3x1x64_v2.yaml @@ -0,0 +1,193 @@ +name: &name "MatchboxNet-3x1x64-v2" + +model: + sample_rate: 16000 + timesteps: 128 + repeat: 1 + dropout: 0.0 + kernel_size_factor: 1.0 + + labels_full: ['visual', 'wow', 'learn', 'backward', 'dog', 'two', 'left', 'happy', 'nine', 'go', 'up', 'bed', 'stop', + 'one', 'zero', 'tree', 'seven', 'on', 'four', 'bird', 'right', 'eight', 'no', 'six', 'forward', 'house', + 'marvin', 'sheila', 'five', 'off', 'three', 'down', 'cat', 'follow', 'yes'] + + labels_subset: ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go", "unknown", "silence"] + + labels: ${model.labels_full} + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + augmentor: + shift: + prob: 1.0 + min_shift_ms: -5.0 + max_shift_ms: 5.0 + white_noise: + prob: 1.0 + min_level: -90 + max_level: -46 + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: False + val_loss_idx: 0 + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + labels: ${model.labels} + batch_size: 128 + shuffle: False + test_loss_idx: 0 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMFCCPreprocessor + window_size: 0.025 + window_stride: 0.01 + window: "hann" + n_mels: &n_mels 64 + n_mfcc: *n_mels + n_fft: 512 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 + time_masks: 2 + freq_width: 15 + time_width: 25 + rect_masks: 5 + rect_time: 25 + rect_freq: 15 + + crop_or_pad_augment: + _target_: nemo.collections.asr.modules.CropOrPadSpectrogramAugmentation + audio_length: ${model.timesteps} + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: false + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 64 + repeat: ${model.repeat} + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: 128 + repeat: 1 + kernel: [29] + stride: [1] + dilation: [2] + dropout: ${model.dropout} + residual: false + separable: true + kernel_size_factor: ${model.kernel_size_factor} + + - filters: &enc_final_filters 128 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: false + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoderClassification + feat_in: *enc_final_filters + return_logits: true + pooling_type: 'avg' + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: 0.05 + # optimizer arguments + betas: [0.95, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: PolynomialHoldDecayAnnealing + + # Scheduler params + power: 2.0 + warmup_ratio: 0.05 + hold_ratio: 0.45 + min_lr: 0.001 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 200 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5.yaml b/models/tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5.yaml new file mode 100644 index 0000000000..ad44ed16f0 --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5.yaml @@ -0,0 +1,281 @@ +name: &name "QuartzNet15x5" + +model: + sample_rate: &sample_rate 16000 + repeat: &repeat 5 + dropout: &dropout 0.0 + separable: &separable true + labels: &labels [" ", "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", "'"] + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + trim_silence: True + max_duration: 16.7 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + shuffle: False + + test_ds: + manifest_filepath: null + sample_rate: 16000 + labels: *labels + batch_size: 32 + shuffle: False + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: 1 + residual: false + separable: *separable + stride: [2] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [2] + dropout: *dropout + filters: 512 + kernel: [87] + repeat: 1 + residual: false + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: &enc_filters 1024 + kernel: [1] + repeat: 1 + residual: false + stride: [1] + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_filters + num_classes: 28 + vocabulary: *labels + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: .01 + # optimizer arguments + betas: [0.8, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # pytorch lightning args + # monitor: val_loss + # reduce_on_plateau: false + + # Scheduler params + warmup_steps: null + warmup_ratio: null + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_wer" + mode: "min" + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5_zh.yaml b/models/tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5_zh.yaml new file mode 100644 index 0000000000..576ef81b6d --- /dev/null +++ b/models/tts/NeMo/examples/asr/conf/quartznet/quartznet_15x5_zh.yaml @@ -0,0 +1,483 @@ +name: &name "QuartzNet15x5" + +model: + sample_rate: &sample_rate 16000 + repeat: &repeat 5 + dropout: &dropout 0.0 + separable: &separable true + labels: &labels [' ', '''', 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, 㶧, 䶮, 一, 丁, 七, 万, 丈, 三, 上, 下, 不, 与, 丐, 丑, 专, 且, 丕, + 世, 丘, 丙, 业, 丛, 东, 丝, 丞, 丢, 两, 严, 丧, 个, 丫, 中, 丰, 串, 临, 丸, 丹, 为, 主, 丽, 举, 乃, + 久, 么, 义, 之, 乌, 乍, 乎, 乏, 乐, 乒, 乓, 乔, 乖, 乘, 乙, 九, 乞, 也, 习, 乡, 书, 买, 乱, 乳, 乾, + 了, 予, 争, 事, 二, 于, 亏, 云, 互, 五, 井, 亘, 亚, 些, 亟, 亡, 亢, 交, 亥, 亦, 产, 亨, 亩, 享, 京, + 亭, 亮, 亲, 亳, 亵, 人, 亿, 什, 仁, 仄, 仅, 仆, 仇, 今, 介, 仍, 从, 仑, 仓, 仔, 仕, 他, 仗, 付, 仙, + 仞, 仟, 仡, 代, 令, 以, 仨, 仪, 们, 仰, 仲, 件, 价, 任, 份, 仿, 企, 伉, 伊, 伍, 伎, 伏, 伐, 休, 众, + 优, 伙, 会, 伞, 伟, 传, 伢, 伤, 伦, 伪, 伫, 伯, 估, 伴, 伶, 伸, 伺, 似, 伽, 佃, 但, 位, 低, 住, 佐, + 佑, 体, 何, 佗, 佘, 余, 佚, 佛, 作, 佝, 佟, 你, 佣, 佩, 佬, 佯, 佰, 佳, 佶, 佻, 佼, 使, 侃, 侄, 侈, + 例, 侍, 侏, 侑, 侗, 供, 依, 侠, 侣, 侥, 侦, 侧, 侨, 侬, 侮, 侯, 侵, 便, 促, 俄, 俊, 俎, 俏, 俐, 俑, + 俗, 俘, 俚, 保, 俞, 俟, 信, 俨, 俩, 俪, 俭, 修, 俯, 俱, 俸, 俺, 俾, 倌, 倍, 倒, 倔, 倘, 候, 倚, 倜, + 借, 倡, 倦, 倩, 倪, 倭, 债, 值, 倾, 偃, 假, 偈, 偌, 偎, 偏, 偓, 偕, 做, 停, 健, 偶, 偷, 偻, 偿, 傀, + 傅, 傍, 傣, 傥, 储, 催, 傲, 傻, 像, 僚, 僧, 僮, 僵, 僻, 儋, 儒, 儡, 儿, 兀, 允, 元, 兄, 充, 兆, 先, + 光, 克, 免, 兑, 兔, 兖, 党, 兜, 兢, 入, 全, 八, 公, 六, 兮, 兰, 共, 关, 兴, 兵, 其, 具, 典, 兹, 养, + 兼, 兽, 冀, 内, 冈, 冉, 册, 再, 冒, 冕, 冗, 写, 军, 农, 冠, 冢, 冤, 冥, 冬, 冯, 冰, 冲, 决, 况, 冶, + 冷, 冻, 冼, 冽, 净, 凄, 准, 凇, 凉, 凋, 凌, 减, 凑, 凛, 凝, 几, 凡, 凤, 凭, 凯, 凰, 凳, 凶, 凸, 凹, + 出, 击, 函, 凿, 刀, 刁, 刃, 分, 切, 刊, 刍, 刎, 刑, 划, 列, 刘, 则, 刚, 创, 初, 删, 判, 刨, 利, 别, + 刮, 到, 制, 刷, 券, 刹, 刺, 刻, 刽, 剁, 剂, 剃, 削, 剌, 前, 剐, 剑, 剔, 剖, 剜, 剥, 剧, 剩, 剪, 副, + 割, 剽, 剿, 劈, 力, 劝, 办, 功, 加, 务, 劣, 动, 助, 努, 劫, 劭, 励, 劲, 劳, 劵, 劾, 势, 勃, 勇, 勉, + 勋, 勐, 勒, 勘, 募, 勤, 勺, 勾, 勿, 匀, 包, 匆, 匈, 匏, 匕, 化, 北, 匙, 匝, 匠, 匡, 匣, 匪, 匮, 匹, + 区, 医, 匾, 匿, 十, 千, 升, 午, 卉, 半, 华, 协, 卑, 卒, 卓, 单, 卖, 南, 博, 卜, 卞, 占, 卡, 卢, 卤, + 卦, 卧, 卫, 卯, 印, 危, 卲, 即, 却, 卵, 卷, 卸, 卿, 厂, 厄, 厅, 历, 厉, 压, 厌, 厕, 厘, 厚, 厝, 原, + 厢, 厥, 厦, 厨, 厩, 厮, 去, 县, 叁, 参, 又, 叉, 及, 友, 双, 反, 发, 叔, 取, 受, 变, 叙, 叛, 叠, 口, + 古, 句, 另, 叨, 叩, 只, 叫, 召, 叭, 叮, 可, 台, 叱, 史, 右, 叵, 叶, 号, 司, 叹, 叼, 叽, 吁, 吃, 各, + 吆, 合, 吉, 吊, 吋, 同, 名, 后, 吏, 吐, 向, 吒, 吓, 吕, 吖, 吗, 君, 吝, 吞, 吟, 吠, 否, 吧, 吨, 吩, + 含, 听, 吭, 吮, 启, 吱, 吴, 吵, 吸, 吹, 吻, 吼, 吾, 呀, 呃, 呆, 呈, 告, 呐, 呕, 呗, 员, 呛, 呜, 呢, + 呦, 周, 呱, 呲, 味, 呵, 呷, 呸, 呻, 呼, 命, 咀, 咂, 咄, 咆, 咋, 和, 咎, 咏, 咐, 咒, 咔, 咕, 咖, 咘, + 咙, 咚, 咝, 咣, 咤, 咦, 咧, 咨, 咩, 咪, 咫, 咬, 咭, 咯, 咱, 咳, 咸, 咻, 咽, 哀, 品, 哂, 哄, 哆, 哇, + 哈, 哉, 响, 哎, 哐, 哑, 哒, 哔, 哕, 哗, 哟, 哥, 哦, 哨, 哩, 哪, 哭, 哮, 哲, 哺, 哼, 哽, 唁, 唆, 唇, + 唉, 唏, 唐, 唑, 唛, 唠, 唢, 唤, 唧, 唬, 售, 唯, 唰, 唱, 唳, 唷, 唾, 啃, 啄, 商, 啊, 啕, 啖, 啜, 啡, + 啤, 啥, 啦, 啧, 啪, 啬, 啰, 啲, 啵, 啶, 啸, 啼, 啾, 喀, 喁, 喂, 喃, 善, 喆, 喇, 喉, 喊, 喋, 喔, 喘, + 喜, 喝, 喟, 喧, 喱, 喳, 喵, 喷, 喻, 喽, 嗄, 嗅, 嗑, 嗒, 嗓, 嗔, 嗖, 嗜, 嗝, 嗡, 嗣, 嗤, 嗦, 嗨, 嗪, + 嗫, 嗬, 嗯, 嗲, 嗷, 嗽, 嘀, 嘈, 嘉, 嘎, 嘏, 嘘, 嘛, 嘞, 嘟, 嘣, 嘭, 嘱, 嘲, 嘴, 嘶, 嘹, 嘻, 嘿, 噌, + 噎, 噗, 噘, 噙, 噜, 噢, 噤, 器, 噩, 噪, 噬, 噱, 噶, 噻, 噼, 嚎, 嚏, 嚓, 嚣, 嚷, 嚼, 囊, 囍, 囔, 囗, + 囚, 四, 回, 因, 团, 囤, 囧, 囫, 园, 囯, 困, 囱, 围, 囵, 囹, 固, 国, 图, 圃, 圄, 圆, 圈, 土, 圣, 在, + 圩, 圪, 圭, 地, 圳, 圹, 场, 圻, 圾, 址, 坂, 均, 坊, 坍, 坎, 坏, 坐, 坑, 块, 坚, 坛, 坝, 坞, 坟, 坠, + 坡, 坤, 坦, 坨, 坩, 坪, 坭, 坯, 坳, 坷, 坻, 垂, 垃, 垄, 垅, 型, 垌, 垒, 垚, 垛, 垡, 垢, 垣, 垤, 垦, + 垩, 垫, 垭, 垮, 埂, 埃, 埇, 埋, 城, 埔, 埕, 埚, 埝, 域, 埠, 埭, 埸, 培, 基, 堀, 堂, 堃, 堆, 堇, 堕, + 堡, 堤, 堪, 堰, 堵, 堺, 塌, 塍, 塑, 塔, 塘, 塞, 填, 塬, 塾, 境, 墅, 墉, 墓, 増, 墙, 增, 墟, 墨, 墩, + 壁, 壑, 壕, 壤, 士, 壬, 壮, 声, 壳, 壶, 壹, 处, 备, 复, 夏, 夔, 夕, 外, 夙, 多, 夜, 够, 大, 天, 太, + 夫, 夭, 央, 夯, 失, 头, 夷, 夸, 夹, 夺, 奁, 奂, 奄, 奇, 奈, 奉, 奋, 奎, 奏, 契, 奔, 奕, 奖, 套, 奘, + 奚, 奠, 奢, 奥, 女, 奴, 奶, 奸, 她, 好, 如, 妃, 妄, 妆, 妇, 妈, 妊, 妍, 妒, 妓, 妖, 妙, 妞, 妤, 妥, + 妨, 妩, 妪, 妫, 妮, 妯, 妲, 妹, 妻, 妾, 姆, 姊, 始, 姐, 姑, 姓, 委, 姗, 姚, 姜, 姝, 姣, 姥, 姨, 姬, + 姻, 姿, 威, 娃, 娄, 娅, 娆, 娇, 娈, 娉, 娌, 娓, 娘, 娜, 娟, 娠, 娣, 娥, 娩, 娱, 娲, 娴, 娶, 娼, 婀, + 婆, 婉, 婊, 婕, 婚, 婢, 婧, 婪, 婴, 婵, 婶, 婷, 婺, 婿, 媒, 媚, 媛, 媞, 媲, 媳, 嫁, 嫂, 嫉, 嫌, 嫒, + 嫔, 嫖, 嫚, 嫡, 嫣, 嫦, 嫩, 嫫, 嬅, 嬉, 嬗, 嬛, 嬴, 嬷, 孀, 子, 孑, 孔, 孕, 字, 存, 孙, 孚, 孛, 孜, + 孝, 孟, 孢, 季, 孤, 学, 孩, 孪, 孬, 孰, 孱, 孳, 孵, 孺, 孽, 宁, 它, 宅, 宇, 守, 安, 宋, 完, 宏, 宓, + 宕, 宗, 官, 宙, 定, 宛, 宜, 宝, 实, 宠, 审, 客, 宣, 室, 宥, 宦, 宪, 宫, 宰, 害, 宴, 宵, 家, 宸, 容, + 宽, 宾, 宿, 寂, 寄, 寅, 密, 寇, 富, 寐, 寒, 寓, 寝, 寞, 察, 寡, 寥, 寨, 寮, 寰, 寸, 对, 寺, 寻, 导, + 寿, 封, 射, 尅, 将, 尉, 尊, 小, 少, 尔, 尕, 尖, 尘, 尚, 尝, 尤, 尧, 尬, 就, 尴, 尸, 尹, 尺, 尼, 尽, + 尾, 尿, 局, 屁, 层, 居, 屈, 屉, 届, 屋, 屌, 屎, 屏, 屐, 屑, 展, 属, 屠, 屡, 履, 屯, 山, 屹, 屿, 岁, + 岂, 岌, 岐, 岑, 岔, 岖, 岗, 岚, 岛, 岩, 岬, 岭, 岱, 岳, 岷, 岸, 峁, 峋, 峒, 峙, 峡, 峥, 峦, 峨, 峪, + 峭, 峰, 峻, 崂, 崃, 崆, 崇, 崎, 崔, 崖, 崛, 崧, 崩, 崭, 崮, 崴, 崽, 嵇, 嵊, 嵋, 嵌, 嵘, 嵛, 嵩, 嵬, + 嶂, 嶙, 嶝, 巅, 巍, 川, 州, 巡, 巢, 工, 左, 巧, 巨, 巩, 巫, 差, 己, 已, 巳, 巴, 巷, 巾, 币, 市, 布, + 帅, 帆, 师, 希, 帐, 帕, 帖, 帘, 帚, 帛, 帜, 帝, 带, 帧, 席, 帮, 帷, 常, 帼, 帽, 幂, 幄, 幅, 幌, 幔, + 幕, 幡, 幢, 干, 平, 年, 并, 幸, 幺, 幻, 幼, 幽, 广, 庄, 庆, 庇, 床, 序, 庐, 库, 应, 底, 庖, 店, 庙, + 庚, 府, 庞, 废, 度, 座, 庭, 庵, 庶, 康, 庸, 庹, 庾, 廉, 廊, 廓, 廖, 延, 廷, 建, 开, 异, 弃, 弄, 弈, + 弊, 弋, 式, 弑, 弓, 引, 弗, 弘, 弛, 弟, 张, 弥, 弦, 弧, 弩, 弭, 弯, 弱, 弹, 强, 弼, 归, 当, 录, 彗, + 彝, 形, 彤, 彦, 彩, 彪, 彬, 彭, 彰, 影, 彷, 役, 彻, 彼, 往, 征, 径, 待, 徇, 很, 徉, 徊, 律, 徐, 徒, + 得, 徘, 徙, 徜, 御, 徨, 循, 微, 德, 徽, 心, 必, 忆, 忌, 忍, 忏, 忐, 忑, 忒, 忖, 志, 忘, 忙, 忠, 忡, + 忤, 忧, 忪, 快, 忱, 念, 忻, 忽, 忿, 怀, 态, 怂, 怄, 怅, 怆, 怎, 怒, 怕, 怖, 怜, 思, 怠, 怡, 急, 怦, + 性, 怨, 怪, 怫, 怯, 怵, 总, 怼, 怿, 恁, 恃, 恋, 恍, 恐, 恒, 恕, 恙, 恢, 恣, 恤, 恨, 恩, 恪, 恬, 恭, + 息, 恰, 恳, 恶, 恸, 恺, 恻, 恼, 恿, 悄, 悉, 悌, 悍, 悔, 悖, 悚, 悟, 悠, 患, 悦, 您, 悬, 悭, 悯, 悱, + 悲, 悴, 悸, 悻, 悼, 情, 惆, 惊, 惋, 惑, 惕, 惚, 惜, 惟, 惠, 惦, 惧, 惨, 惩, 惫, 惬, 惭, 惮, 惯, 惰, + 想, 惶, 惹, 惺, 愁, 愈, 愉, 意, 愕, 愚, 感, 愣, 愤, 愧, 愫, 愿, 慈, 慌, 慎, 慑, 慕, 慢, 慧, 慨, 慰, + 慵, 慷, 憋, 憎, 憔, 憧, 憨, 憩, 憬, 憷, 憾, 懂, 懈, 懊, 懋, 懑, 懒, 懦, 懵, 懿, 戈, 戊, 戌, 戍, 戎, + 戏, 成, 我, 戒, 或, 戗, 战, 戚, 戛, 戟, 截, 戬, 戮, 戳, 戴, 户, 戾, 房, 所, 扁, 扇, 扈, 扉, 手, 才, + 扎, 扑, 扒, 打, 扔, 托, 扛, 扞, 扣, 扦, 执, 扩, 扪, 扫, 扬, 扭, 扮, 扯, 扰, 扳, 扶, 批, 扼, 找, 承, + 技, 抄, 抉, 把, 抑, 抒, 抓, 投, 抖, 抗, 折, 抚, 抛, 抠, 抡, 抢, 护, 报, 抨, 披, 抬, 抱, 抵, 抹, 押, + 抽, 抿, 拂, 拄, 担, 拆, 拇, 拈, 拉, 拌, 拍, 拎, 拐, 拒, 拓, 拔, 拖, 拗, 拘, 拙, 拚, 招, 拜, 拟, 拢, + 拣, 拥, 拦, 拧, 拨, 择, 括, 拭, 拮, 拯, 拱, 拳, 拴, 拷, 拼, 拽, 拾, 拿, 持, 挂, 指, 按, 挎, 挑, 挖, + 挚, 挛, 挝, 挞, 挟, 挠, 挡, 挣, 挤, 挥, 挨, 挪, 挫, 振, 挺, 挽, 捂, 捅, 捆, 捉, 捋, 捌, 捍, 捎, 捏, + 捐, 捕, 捞, 损, 捡, 换, 捣, 捧, 据, 捶, 捷, 捺, 捻, 掀, 掂, 掇, 授, 掉, 掌, 掏, 掐, 排, 掖, 掘, 掠, + 探, 掣, 接, 控, 推, 掩, 措, 掬, 掮, 掰, 掳, 掴, 掷, 掸, 掺, 揄, 揉, 揍, 描, 提, 插, 握, 揣, 揩, 揪, + 揭, 援, 揶, 揽, 搀, 搁, 搂, 搅, 搏, 搐, 搓, 搔, 搜, 搞, 搡, 搧, 搪, 搬, 搭, 携, 搽, 摁, 摄, 摆, 摇, + 摈, 摊, 摒, 摔, 摘, 摞, 摧, 摩, 摸, 摹, 撂, 撅, 撇, 撑, 撒, 撕, 撞, 撤, 撩, 撬, 播, 撮, 撰, 撵, 撸, + 撺, 撼, 擀, 擂, 擅, 操, 擎, 擒, 擘, 擞, 擢, 擦, 攀, 攒, 攘, 攥, 攫, 支, 收, 攸, 改, 攻, 放, 政, 故, + 效, 敌, 敏, 救, 敕, 敖, 教, 敛, 敝, 敞, 敢, 散, 敦, 敬, 数, 敲, 整, 敷, 文, 斋, 斌, 斐, 斑, 斓, 斗, + 料, 斛, 斜, 斟, 斡, 斤, 斥, 斧, 斩, 断, 斯, 新, 方, 施, 旁, 旅, 旋, 旌, 族, 旖, 旗, 无, 既, 日, 旦, + 旧, 旨, 早, 旬, 旭, 旮, 旯, 旱, 时, 旷, 旺, 旻, 昀, 昂, 昆, 昊, 昌, 明, 昏, 易, 昔, 昕, 昙, 昝, 星, + 映, 春, 昧, 昨, 昭, 是, 昱, 昴, 昵, 昶, 昼, 显, 晃, 晋, 晌, 晏, 晒, 晓, 晔, 晕, 晖, 晗, 晚, 晞, 晟, + 晤, 晦, 晨, 普, 景, 晰, 晴, 晶, 晷, 智, 晾, 暂, 暄, 暇, 暌, 暑, 暖, 暗, 暧, 暨, 暮, 暴, 暹, 暾, 曈, + 曙, 曜, 曝, 曦, 曰, 曲, 曳, 更, 曹, 曼, 曾, 替, 最, 月, 有, 朋, 服, 朐, 朔, 朕, 朗, 望, 朝, 期, 朦, + 木, 未, 末, 本, 札, 术, 朱, 朴, 朵, 机, 朽, 杀, 杂, 权, 杆, 杈, 杉, 李, 杏, 材, 村, 杓, 杖, 杜, 杞, + 束, 杠, 条, 来, 杨, 杭, 杯, 杰, 杳, 杵, 杷, 松, 板, 极, 构, 枇, 枉, 枋, 析, 枕, 林, 枚, 果, 枝, 枞, + 枢, 枣, 枥, 枪, 枫, 枭, 枯, 枰, 枳, 架, 枷, 枸, 柃, 柄, 柏, 某, 柑, 柒, 染, 柔, 柘, 柚, 柜, 柞, 柠, + 查, 柩, 柬, 柯, 柱, 柳, 柴, 柿, 栀, 栅, 标, 栈, 栋, 栌, 栎, 栏, 树, 栓, 栖, 栗, 校, 栩, 株, 样, 核, + 根, 格, 栽, 栾, 桁, 桂, 桃, 框, 案, 桉, 桌, 桎, 桐, 桑, 桓, 桔, 桠, 桢, 档, 桥, 桦, 桨, 桩, 桴, 桶, + 桷, 梁, 梅, 梆, 梏, 梓, 梗, 梢, 梦, 梧, 梨, 梭, 梯, 械, 梳, 梵, 检, 棂, 棉, 棋, 棍, 棒, 棕, 棘, 棚, + 棠, 棣, 森, 棱, 棵, 棺, 椁, 椅, 椋, 植, 椎, 椒, 椟, 椤, 椭, 椰, 椴, 椹, 椿, 楂, 楔, 楚, 楞, 楠, 楣, + 楷, 楸, 楼, 概, 榄, 榆, 榈, 榉, 榔, 榕, 榛, 榜, 榨, 榫, 榭, 榴, 榷, 榻, 槃, 槌, 槎, 槐, 槛, 槟, 槭, + 槽, 槿, 樊, 樟, 模, 樨, 横, 樯, 樱, 樵, 樽, 樾, 橄, 橇, 橐, 橘, 橙, 橡, 橱, 檀, 檐, 檗, 檬, 欠, 次, + 欢, 欣, 欧, 欲, 欸, 欺, 款, 歆, 歇, 歉, 歌, 歙, 止, 正, 此, 步, 武, 歧, 歩, 歪, 歹, 死, 歼, 殁, 殃, + 殆, 殇, 殉, 殊, 残, 殒, 殓, 殖, 殚, 殡, 殴, 段, 殷, 殿, 毁, 毂, 毅, 毋, 母, 每, 毒, 毓, 比, 毕, 毗, + 毙, 毛, 毡, 毫, 毯, 毽, 氏, 民, 氓, 气, 氚, 氛, 氟, 氢, 氤, 氦, 氧, 氨, 氪, 氮, 氯, 氰, 氲, 水, 永, + 汀, 汁, 求, 汇, 汉, 汊, 汐, 汕, 汗, 汛, 汝, 汞, 江, 池, 污, 汤, 汨, 汩, 汪, 汰, 汲, 汴, 汶, 汹, 汽, + 汾, 沁, 沂, 沃, 沅, 沈, 沉, 沌, 沏, 沐, 沓, 沙, 沛, 沟, 没, 沢, 沣, 沥, 沦, 沧, 沪, 沫, 沭, 沮, 沱, + 河, 沸, 油, 治, 沼, 沽, 沾, 沿, 泄, 泉, 泊, 泌, 泓, 泔, 法, 泖, 泗, 泛, 泞, 泠, 泡, 波, 泣, 泥, 注, + 泪, 泫, 泮, 泯, 泰, 泱, 泳, 泵, 泷, 泸, 泺, 泻, 泼, 泽, 泾, 洁, 洋, 洒, 洗, 洙, 洛, 洞, 津, 洪, 洮, + 洱, 洲, 洵, 洹, 洺, 活, 洼, 洽, 派, 流, 浃, 浅, 浆, 浇, 浈, 浊, 测, 济, 浏, 浐, 浑, 浒, 浓, 浔, 浙, + 浚, 浜, 浠, 浣, 浦, 浩, 浪, 浮, 浴, 海, 浸, 涂, 涅, 消, 涉, 涌, 涎, 涑, 涓, 涕, 涛, 涝, 涞, 涟, 涠, + 涡, 涣, 涤, 润, 涧, 涨, 涩, 涪, 涮, 涯, 液, 涵, 涸, 涿, 淀, 淄, 淅, 淆, 淇, 淋, 淌, 淑, 淖, 淘, 淝, + 淞, 淡, 淤, 淦, 淫, 淬, 淮, 深, 淳, 混, 淹, 添, 淼, 清, 渊, 渌, 渍, 渎, 渐, 渑, 渔, 渗, 渚, 渝, 渠, + 渡, 渣, 渤, 渥, 温, 渭, 港, 渲, 渴, 游, 渺, 湃, 湄, 湉, 湍, 湎, 湖, 湘, 湛, 湫, 湾, 湿, 溃, 溅, 溆, + 溉, 溏, 源, 溜, 溟, 溢, 溥, 溧, 溪, 溯, 溶, 溺, 滁, 滇, 滋, 滑, 滔, 滕, 滘, 滚, 滞, 满, 滢, 滤, 滥, + 滦, 滨, 滩, 滴, 滹, 漂, 漆, 漉, 漏, 漓, 演, 漕, 漠, 漩, 漪, 漫, 漭, 漯, 漱, 漳, 漾, 潆, 潇, 潋, 潍, + 潘, 潜, 潞, 潢, 潦, 潭, 潮, 潸, 潺, 潼, 澄, 澈, 澍, 澎, 澜, 澡, 澧, 澳, 澶, 激, 濂, 濑, 濒, 濠, 濡, + 濮, 濯, 瀑, 瀚, 瀛, 灌, 灏, 灞, 火, 灭, 灯, 灰, 灵, 灶, 灸, 灼, 灾, 灿, 炀, 炅, 炉, 炊, 炎, 炒, 炔, + 炕, 炖, 炙, 炜, 炫, 炬, 炭, 炮, 炯, 炳, 炷, 炸, 点, 炼, 炽, 烀, 烁, 烂, 烃, 烈, 烊, 烘, 烙, 烛, 烟, + 烤, 烦, 烧, 烨, 烩, 烫, 烬, 热, 烯, 烷, 烹, 烽, 焉, 焊, 焓, 焕, 焖, 焗, 焘, 焙, 焚, 焦, 焯, 焰, 焱, + 然, 煊, 煌, 煎, 煜, 煞, 煤, 煦, 照, 煨, 煮, 煲, 煳, 煽, 熄, 熊, 熏, 熔, 熙, 熟, 熠, 熨, 熬, 熵, 熹, + 燃, 燊, 燎, 燕, 燥, 燮, 爆, 爪, 爬, 爱, 爵, 父, 爷, 爸, 爹, 爽, 片, 版, 牌, 牍, 牒, 牙, 牛, 牟, 牠, + 牡, 牢, 牧, 物, 牲, 牵, 特, 牺, 牾, 犀, 犁, 犄, 犇, 犊, 犒, 犟, 犬, 犯, 状, 犷, 犸, 犹, 狂, 狄, 狈, + 狐, 狒, 狗, 狙, 狞, 狠, 狡, 狩, 独, 狭, 狮, 狰, 狱, 狸, 狼, 猁, 猎, 猖, 猛, 猜, 猝, 猥, 猩, 猪, 猫, + 猬, 献, 猴, 猷, 猹, 猾, 猿, 獒, 獗, 獭, 獾, 玄, 率, 玉, 王, 玑, 玖, 玛, 玟, 玥, 玩, 玫, 玮, 环, 现, + 玲, 玳, 玷, 玹, 玺, 玻, 珀, 珂, 珈, 珉, 珊, 珍, 珏, 珑, 珙, 珞, 珠, 珥, 班, 珮, 珲, 珺, 球, 琅, 理, + 琉, 琊, 琏, 琐, 琛, 琢, 琤, 琥, 琦, 琨, 琪, 琬, 琮, 琰, 琳, 琴, 琵, 琶, 琼, 瑁, 瑄, 瑕, 瑙, 瑚, 瑛, + 瑜, 瑞, 瑟, 瑠, 瑭, 瑰, 瑶, 瑷, 瑾, 璀, 璃, 璇, 璋, 璐, 璞, 璟, 璧, 璨, 瓜, 瓢, 瓣, 瓦, 瓮, 瓯, 瓶, + 瓷, 甄, 甘, 甚, 甜, 生, 甥, 用, 甩, 甫, 甬, 甭, 田, 由, 甲, 申, 电, 男, 甸, 町, 画, 畅, 畈, 畊, 界, + 畏, 畔, 留, 畜, 略, 番, 畴, 畸, 畿, 疃, 疆, 疏, 疑, 疖, 疗, 疙, 疚, 疝, 疟, 疡, 疣, 疤, 疫, 疮, 疯, + 疱, 疲, 疴, 疵, 疸, 疹, 疼, 疽, 疾, 病, 症, 痉, 痊, 痍, 痒, 痔, 痕, 痘, 痛, 痞, 痢, 痣, 痧, 痨, 痪, + 痫, 痰, 痱, 痴, 痹, 痼, 瘀, 瘁, 瘙, 瘟, 瘠, 瘢, 瘤, 瘦, 瘩, 瘪, 瘫, 瘳, 瘴, 瘸, 瘾, 癌, 癖, 癜, 癞, + 癣, 癫, 登, 白, 百, 皂, 的, 皆, 皇, 皋, 皎, 皑, 皓, 皖, 皙, 皮, 皱, 皿, 盂, 盅, 盆, 盈, 益, 盎, 盏, + 盐, 监, 盒, 盔, 盖, 盗, 盘, 盛, 盟, 目, 盯, 盱, 盲, 直, 相, 盹, 盼, 盾, 省, 眈, 眉, 看, 眙, 真, 眠, + 眨, 眩, 眬, 眯, 眶, 眷, 眸, 眺, 眼, 着, 睁, 睇, 睐, 睑, 睛, 睡, 睢, 督, 睦, 睫, 睬, 睹, 睽, 睾, 睿, + 瞄, 瞅, 瞌, 瞎, 瞑, 瞒, 瞟, 瞠, 瞥, 瞧, 瞩, 瞪, 瞬, 瞭, 瞰, 瞳, 瞻, 瞿, 矍, 矗, 矛, 矜, 矢, 矣, 知, + 矩, 矫, 矬, 短, 矮, 石, 矶, 矸, 矽, 矾, 矿, 砀, 码, 砂, 砌, 砍, 砒, 研, 砖, 砚, 砝, 砣, 砥, 砭, 砰, + 破, 砷, 砸, 砺, 砼, 砾, 础, 硅, 硌, 硒, 硕, 硖, 硚, 硝, 硫, 硬, 确, 硼, 碉, 碌, 碍, 碎, 碑, 碓, 碗, + 碘, 碚, 碜, 碟, 碣, 碧, 碰, 碱, 碳, 碴, 碾, 磁, 磅, 磊, 磋, 磐, 磕, 磨, 磴, 磷, 磺, 礁, 示, 礼, 社, + 祀, 祁, 祈, 祉, 祎, 祐, 祖, 祚, 祛, 祝, 神, 祟, 祠, 祢, 祥, 票, 祭, 祯, 祷, 祸, 祺, 禀, 禁, 禄, 禅, + 福, 禧, 禹, 禺, 离, 禽, 禾, 秀, 私, 秃, 秆, 秉, 秋, 种, 科, 秒, 秘, 租, 秣, 秤, 秦, 秧, 秩, 积, 称, + 秸, 移, 秽, 稀, 程, 稍, 税, 稔, 稚, 稞, 稠, 稣, 稳, 稷, 稹, 稻, 稼, 稽, 稿, 穆, 穗, 穴, 究, 穷, 穹, + 空, 穿, 突, 窃, 窄, 窈, 窍, 窑, 窒, 窕, 窖, 窗, 窘, 窜, 窝, 窟, 窠, 窥, 窦, 窨, 窿, 立, 竖, 站, 竞, + 竟, 章, 竣, 童, 竭, 端, 竹, 竺, 竽, 竿, 笃, 笆, 笈, 笋, 笑, 笔, 笙, 笛, 笠, 符, 笨, 第, 笳, 笸, 笼, + 等, 筋, 筏, 筐, 筑, 筒, 答, 策, 筛, 筝, 筠, 筱, 筵, 筷, 筹, 签, 简, 箍, 箔, 箕, 算, 管, 箩, 箫, 箭, + 箱, 箴, 篁, 篆, 篇, 篑, 篓, 篝, 篡, 篦, 篪, 篮, 篱, 篷, 篼, 簇, 簋, 簧, 簪, 簸, 簿, 籁, 籍, 米, 类, + 籼, 籽, 粉, 粑, 粒, 粕, 粗, 粘, 粟, 粤, 粥, 粪, 粮, 粱, 粲, 粳, 粹, 粼, 粽, 精, 糊, 糕, 糖, 糗, 糙, + 糟, 糠, 糯, 系, 紊, 素, 索, 紧, 紫, 累, 絮, 綦, 繁, 纂, 纠, 纡, 红, 纣, 纤, 约, 级, 纨, 纪, 纫, 纬, + 纭, 纯, 纰, 纱, 纲, 纳, 纵, 纶, 纷, 纸, 纹, 纺, 纽, 纾, 线, 绀, 练, 组, 绅, 细, 织, 终, 绉, 绊, 绋, + 绌, 绍, 绎, 经, 绑, 绒, 结, 绔, 绕, 绘, 给, 绚, 绛, 络, 绝, 绞, 统, 绢, 绣, 绥, 继, 绩, 绪, 绫, 续, + 绮, 绯, 绰, 绳, 维, 绵, 绷, 绸, 绻, 综, 绽, 绿, 缀, 缄, 缅, 缆, 缇, 缉, 缎, 缓, 缔, 缕, 编, 缘, 缙, + 缚, 缛, 缜, 缝, 缠, 缢, 缤, 缨, 缩, 缪, 缬, 缭, 缮, 缰, 缱, 缴, 缸, 缺, 罂, 罄, 罐, 网, 罔, 罕, 罗, + 罚, 罡, 罢, 罩, 罪, 置, 署, 罹, 羁, 羊, 羌, 美, 羔, 羚, 羞, 羡, 群, 羧, 羯, 羲, 羸, 羹, 羽, 羿, 翁, + 翅, 翊, 翌, 翎, 翔, 翘, 翟, 翠, 翡, 翩, 翰, 翱, 翻, 翼, 耀, 老, 考, 耄, 者, 耆, 耋, 而, 耍, 耐, 耒, + 耕, 耗, 耘, 耙, 耜, 耪, 耳, 耶, 耷, 耸, 耻, 耽, 耿, 聂, 聆, 聊, 聋, 职, 联, 聘, 聚, 聪, 肃, 肆, 肇, + 肉, 肋, 肌, 肖, 肘, 肚, 肛, 肝, 肠, 股, 肢, 肤, 肥, 肩, 肪, 肮, 肯, 肱, 育, 肴, 肺, 肾, 肿, 胀, 胁, + 胃, 胆, 背, 胎, 胖, 胗, 胚, 胛, 胜, 胞, 胡, 胤, 胥, 胧, 胫, 胭, 胯, 胰, 胱, 胳, 胶, 胸, 胺, 能, 脂, + 脆, 脉, 脊, 脍, 脏, 脐, 脑, 脓, 脖, 脚, 脯, 脱, 脸, 脾, 腆, 腈, 腊, 腋, 腌, 腐, 腑, 腓, 腔, 腕, 腥, + 腩, 腭, 腮, 腰, 腱, 腴, 腹, 腺, 腻, 腼, 腾, 腿, 膀, 膈, 膊, 膏, 膑, 膛, 膜, 膝, 膨, 膳, 膺, 臀, 臂, + 臃, 臆, 臊, 臣, 臧, 自, 臬, 臭, 至, 致, 臻, 臼, 舀, 舅, 舆, 舌, 舍, 舐, 舒, 舔, 舛, 舜, 舞, 舟, 航, + 舫, 般, 舰, 舱, 舵, 舶, 舷, 舸, 船, 艇, 艋, 艘, 艮, 良, 艰, 色, 艳, 艺, 艾, 艿, 节, 芊, 芋, 芍, 芒, + 芗, 芙, 芜, 芝, 芥, 芦, 芩, 芪, 芬, 芭, 芮, 芯, 花, 芳, 芷, 芸, 芹, 芽, 芾, 苇, 苋, 苍, 苏, 苑, 苓, + 苔, 苗, 苛, 苞, 苟, 苡, 苣, 若, 苦, 苫, 苯, 英, 苷, 苹, 茁, 茂, 范, 茄, 茅, 茆, 茉, 茌, 茎, 茗, 茛, + 茜, 茧, 茨, 茫, 茬, 茯, 茱, 茳, 茴, 茵, 茶, 茸, 茹, 茼, 荀, 荃, 荆, 荇, 草, 荏, 荐, 荒, 荔, 荚, 荛, + 荞, 荟, 荠, 荡, 荣, 荤, 荧, 荨, 荫, 药, 荷, 荸, 荻, 荼, 莅, 莆, 莉, 莎, 莒, 莓, 莘, 莜, 莞, 莠, 莪, + 莫, 莱, 莲, 莴, 获, 莹, 莺, 莽, 菀, 菁, 菅, 菇, 菊, 菌, 菏, 菖, 菘, 菜, 菠, 菡, 菩, 菱, 菲, 萃, 萄, + 萋, 萌, 萍, 萎, 萝, 萤, 营, 萦, 萧, 萨, 萱, 萸, 落, 葆, 著, 葚, 葛, 葡, 董, 葩, 葫, 葬, 葱, 葳, 葵, + 葺, 蒂, 蒋, 蒙, 蒜, 蒯, 蒲, 蒸, 蒿, 蓁, 蓄, 蓉, 蓓, 蓝, 蓟, 蓥, 蓦, 蓬, 蓼, 蔑, 蔓, 蔗, 蔚, 蔡, 蔫, + 蔬, 蔷, 蔺, 蔻, 蔼, 蔽, 蕃, 蕉, 蕊, 蕙, 蕨, 蕲, 蕴, 蕾, 薄, 薇, 薏, 薛, 薪, 薯, 薰, 薷, 藁, 藉, 藏, + 藐, 藓, 藕, 藜, 藠, 藤, 藩, 藻, 藿, 蘑, 蘸, 虎, 虏, 虐, 虑, 虔, 虚, 虞, 虫, 虱, 虹, 虻, 虽, 虾, 蚀, + 蚁, 蚂, 蚊, 蚌, 蚓, 蚕, 蚝, 蚣, 蚤, 蚪, 蚬, 蚯, 蚱, 蚴, 蛀, 蛆, 蛇, 蛉, 蛊, 蛋, 蛎, 蛐, 蛔, 蛙, 蛛, + 蛟, 蛤, 蛮, 蛰, 蛳, 蛹, 蛾, 蜀, 蜂, 蜃, 蜇, 蜈, 蜊, 蜍, 蜒, 蜓, 蜕, 蜗, 蜘, 蜚, 蜜, 蜡, 蜢, 蜥, 蜱, + 蜴, 蜷, 蜻, 蜿, 蝇, 蝈, 蝉, 蝌, 蝎, 蝗, 蝙, 蝠, 蝮, 蝴, 蝶, 蝽, 螂, 螃, 螈, 融, 螨, 螳, 螺, 蟀, 蟆, + 蟊, 蟋, 蟑, 蟒, 蟠, 蟹, 蟾, 蠊, 蠕, 蠡, 蠢, 血, 衅, 行, 衍, 衔, 街, 衙, 衡, 衢, 衣, 补, 表, 衩, 衫, + 衬, 衮, 衰, 衲, 衷, 袁, 袂, 袄, 袅, 袈, 袋, 袍, 袒, 袖, 袜, 被, 袭, 袱, 裁, 裂, 装, 裆, 裔, 裕, 裘, + 裙, 裟, 裤, 裨, 裱, 裳, 裴, 裸, 裹, 褂, 褐, 褒, 褓, 褔, 褚, 褛, 褥, 褪, 褴, 褶, 襁, 襄, 襟, 西, 要, + 覃, 覆, 见, 观, 规, 觅, 视, 览, 觉, 觊, 觎, 觐, 觑, 角, 觞, 解, 觥, 触, 言, 訾, 詹, 誉, 誓, 警, 譬, + 计, 订, 讣, 认, 讥, 讧, 讨, 让, 讪, 训, 议, 讯, 记, 讲, 讳, 讴, 讶, 讷, 许, 讹, 论, 讼, 讽, 设, 访, + 诀, 证, 诃, 评, 诅, 识, 诈, 诉, 诊, 诋, 词, 诏, 译, 诓, 试, 诗, 诘, 诙, 诚, 诛, 话, 诞, 诟, 诠, 诡, + 询, 诣, 诤, 该, 详, 诧, 诩, 诫, 诬, 语, 误, 诱, 诲, 说, 诵, 诶, 请, 诸, 诹, 诺, 读, 诽, 课, 诿, 谀, + 谁, 调, 谅, 谆, 谈, 谊, 谋, 谌, 谍, 谎, 谏, 谐, 谑, 谓, 谕, 谖, 谘, 谙, 谚, 谛, 谜, 谟, 谢, 谣, 谤, + 谦, 谧, 谨, 谩, 谬, 谭, 谮, 谯, 谱, 谴, 谶, 谷, 豁, 豆, 豇, 豉, 豌, 豚, 象, 豢, 豪, 豫, 豹, 豺, 貂, + 貅, 貉, 貌, 貔, 贝, 贞, 负, 贡, 财, 责, 贤, 败, 账, 货, 质, 贩, 贪, 贫, 贬, 购, 贮, 贯, 贰, 贱, 贲, + 贴, 贵, 贷, 贸, 费, 贺, 贻, 贼, 贾, 贿, 赁, 赂, 赃, 资, 赅, 赈, 赉, 赊, 赋, 赌, 赎, 赏, 赐, 赓, 赔, + 赖, 赘, 赚, 赛, 赝, 赞, 赠, 赡, 赢, 赣, 赤, 赦, 赫, 走, 赳, 赴, 赵, 赶, 起, 趁, 超, 越, 趋, 趟, 趣, + 足, 趴, 趵, 趸, 趺, 趾, 跃, 跄, 跆, 跋, 跌, 跎, 跑, 跚, 跛, 距, 跟, 跤, 跨, 跪, 跬, 路, 跳, 践, 跶, + 跷, 跹, 跺, 跻, 踉, 踊, 踌, 踏, 踝, 踞, 踢, 踩, 踪, 踮, 踯, 踱, 踵, 踹, 踺, 蹁, 蹂, 蹄, 蹈, 蹉, 蹊, + 蹋, 蹒, 蹚, 蹦, 蹩, 蹬, 蹭, 蹲, 蹴, 蹶, 蹼, 蹿, 躁, 躅, 躇, 躏, 身, 躬, 躯, 躲, 躺, 车, 轧, 轨, 轩, + 轫, 转, 轮, 软, 轰, 轱, 轲, 轳, 轴, 轶, 轸, 轻, 轼, 载, 轿, 较, 辄, 辅, 辆, 辈, 辉, 辊, 辍, 辐, 辑, + 输, 辕, 辖, 辗, 辘, 辙, 辛, 辜, 辞, 辟, 辣, 辨, 辩, 辫, 辰, 辱, 边, 辽, 达, 迁, 迂, 迄, 迅, 过, 迈, + 迎, 运, 近, 返, 还, 这, 进, 远, 违, 连, 迟, 迢, 迥, 迦, 迩, 迪, 迫, 迭, 述, 迷, 迸, 迹, 追, 退, 送, + 适, 逃, 逅, 逆, 选, 逊, 逋, 逍, 透, 逐, 逑, 递, 途, 逗, 通, 逛, 逝, 逞, 速, 造, 逡, 逢, 逮, 逯, 逵, + 逸, 逻, 逼, 逾, 遁, 遂, 遇, 遍, 遏, 遐, 遑, 道, 遗, 遛, 遢, 遣, 遥, 遨, 遭, 遮, 遴, 遵, 避, 邀, 邂, + 邃, 邋, 邑, 邓, 邕, 邙, 邛, 邝, 邡, 邢, 那, 邦, 邪, 邬, 邮, 邯, 邰, 邱, 邳, 邵, 邸, 邹, 邺, 邻, 郁, + 郅, 郇, 郊, 郎, 郑, 郓, 郜, 郝, 郡, 郧, 部, 郫, 郭, 郯, 郴, 郸, 都, 鄂, 鄙, 鄞, 鄢, 鄱, 酉, 酊, 酋, + 酌, 配, 酐, 酒, 酗, 酚, 酝, 酞, 酣, 酥, 酩, 酪, 酬, 酮, 酯, 酰, 酱, 酵, 酶, 酷, 酸, 酿, 醇, 醉, 醋, + 醍, 醐, 醒, 醛, 醺, 采, 釉, 释, 里, 重, 野, 量, 金, 釜, 鉴, 銮, 鏖, 鑫, 钇, 针, 钉, 钊, 钎, 钏, 钐, + 钒, 钓, 钗, 钙, 钛, 钜, 钝, 钞, 钟, 钠, 钢, 钣, 钥, 钦, 钧, 钨, 钩, 钮, 钯, 钰, 钱, 钲, 钳, 钴, 钵, + 钻, 钼, 钾, 钿, 铀, 铁, 铂, 铃, 铄, 铅, 铆, 铉, 铋, 铍, 铎, 铐, 铑, 铖, 铛, 铜, 铝, 铟, 铠, 铡, 铣, + 铤, 铧, 铨, 铩, 铬, 铭, 铮, 铰, 铲, 银, 铷, 铸, 铺, 链, 铿, 销, 锁, 锂, 锄, 锅, 锆, 锈, 锉, 锋, 锌, + 锏, 锐, 锑, 锒, 错, 锚, 锟, 锡, 锢, 锣, 锤, 锥, 锦, 锨, 锭, 键, 锯, 锰, 锲, 锴, 锵, 锷, 锹, 锻, 镀, + 镁, 镂, 镇, 镉, 镊, 镌, 镍, 镏, 镐, 镑, 镔, 镕, 镖, 镜, 镣, 镭, 镯, 镰, 镳, 镶, 长, 门, 闩, 闪, 闫, + 闭, 问, 闯, 闰, 闲, 闳, 间, 闵, 闷, 闸, 闹, 闺, 闻, 闽, 闾, 阀, 阁, 阂, 阄, 阅, 阆, 阉, 阎, 阐, 阑, + 阔, 阕, 阖, 阙, 阚, 阜, 队, 阡, 阪, 阮, 阱, 防, 阳, 阴, 阵, 阶, 阻, 阿, 陀, 陂, 附, 际, 陆, 陇, 陈, + 陉, 陋, 陌, 降, 限, 陕, 陛, 陡, 院, 除, 陨, 险, 陪, 陬, 陵, 陶, 陷, 隅, 隆, 隋, 隍, 随, 隐, 隔, 隗, + 隘, 隙, 障, 隧, 隶, 隼, 隽, 难, 雀, 雁, 雄, 雅, 集, 雇, 雉, 雌, 雍, 雏, 雒, 雕, 雨, 雪, 雯, 雳, 零, + 雷, 雹, 雾, 需, 霁, 霄, 霆, 震, 霈, 霉, 霍, 霎, 霏, 霓, 霖, 霜, 霞, 霪, 露, 霸, 霹, 霾, 靑, 青, 靓, + 靖, 静, 靛, 非, 靠, 靡, 面, 革, 靳, 靴, 靶, 鞅, 鞋, 鞍, 鞑, 鞘, 鞠, 鞭, 韦, 韧, 韩, 韫, 韬, 韭, 音, + 韵, 韶, 页, 顶, 顷, 项, 顺, 须, 顽, 顾, 顿, 颀, 颁, 颂, 预, 颅, 领, 颇, 颈, 颊, 颌, 颍, 颐, 频, 颓, + 颖, 颗, 题, 颚, 颜, 额, 颠, 颢, 颤, 颦, 颧, 风, 飒, 飓, 飘, 飙, 飚, 飞, 食, 飧, 餍, 餐, 餮, 饕, 饥, + 饨, 饪, 饭, 饮, 饯, 饰, 饱, 饲, 饴, 饵, 饶, 饷, 饺, 饼, 饽, 饿, 馀, 馁, 馄, 馅, 馆, 馈, 馊, 馋, 馍, + 馏, 馑, 馒, 馕, 首, 馗, 香, 馥, 馨, 马, 驭, 驮, 驯, 驰, 驱, 驳, 驴, 驶, 驷, 驸, 驹, 驻, 驼, 驾, 驿, + 骁, 骂, 骄, 骅, 骆, 骇, 骈, 骊, 骋, 验, 骏, 骐, 骑, 骓, 骗, 骚, 骛, 骜, 骝, 骞, 骠, 骡, 骤, 骥, 骨, + 骰, 骷, 骸, 骺, 骼, 髂, 髅, 髋, 髌, 髓, 高, 髦, 髯, 鬃, 鬓, 鬟, 鬼, 魁, 魂, 魄, 魅, 魇, 魉, 魍, 魏, + 魔, 魟, 鱼, 鱿, 鲁, 鲅, 鲈, 鲍, 鲑, 鲜, 鲟, 鲠, 鲢, 鲤, 鲨, 鲫, 鲭, 鲳, 鲶, 鲷, 鲸, 鲼, 鳃, 鳄, 鳅, + 鳌, 鳍, 鳕, 鳖, 鳗, 鳝, 鳞, 鳟, 鸟, 鸠, 鸡, 鸢, 鸣, 鸥, 鸦, 鸩, 鸪, 鸫, 鸭, 鸯, 鸳, 鸵, 鸽, 鸾, 鸿, + 鹁, 鹂, 鹃, 鹅, 鹉, 鹊, 鹌, 鹏, 鹑, 鹜, 鹞, 鹤, 鹦, 鹧, 鹫, 鹭, 鹰, 鹳, 鹿, 麂, 麋, 麒, 麓, 麝, 麟, + 麦, 麸, 麻, 麾, 黄, 黍, 黎, 黏, 黑, 黔, 默, 黛, 黝, 黟, 黯, 鼎, 鼓, 鼠, 鼬, 鼹, 鼻, 鼾, 齐, 齿, 龃, + 龄, 龅, 龈, 龉, 龊, 龌, 龙, 龚, 龟, "\U0002B5AF", "\U0002B689"] + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + trim_silence: True + normalize: False + max_duration: 16.7 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + normalize: False + labels: *labels + batch_size: 32 + shuffle: False + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: 1 + residual: false + separable: *separable + stride: [2] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [2] + dropout: *dropout + filters: 512 + kernel: [87] + repeat: 1 + residual: false + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: &enc_filters 1024 + kernel: [1] + repeat: 1 + residual: false + stride: [1] + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_filters + num_classes: 5206 + vocabulary: *labels + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: .01 + # optimizer arguments + betas: [0.8, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # pytorch lightning args + # monitor: val_loss + # reduce_on_plateau: false + + # Scheduler params + warmup_steps: null + warmup_ratio: null + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_wer" + mode: "min" + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/experimental/conformer/conformer_rnnt_bpe.yaml b/models/tts/NeMo/examples/asr/experimental/conformer/conformer_rnnt_bpe.yaml new file mode 100644 index 0000000000..e7c2beac47 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/conformer/conformer_rnnt_bpe.yaml @@ -0,0 +1,201 @@ +# It contains the default values for training a Conformer ASR model medium size with RNNT loss and BPE-based vocabulary. +# Default learning parameters in this config are set for effective batch size of 2K. +# To train it with smaller batch sizes, you may need to re-tune the learning parameters or use higher accumulate_grad_batches. + +# You may find more info about Conformer-CTC here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/models.html#conformer-ctc +# Pre-trained models of Conformer-CTC can be found here: https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/stable/asr/results.html + +name: "Conformer-RNNT-BPE" + +model: + sample_rate: &sample_rate 16000 + compute_eval_loss: false # disables loss evaluation for validation and test data + log_prediction: true # enables logging sample predictions in the output during training + ctc_reduction: 'mean_batch' + + model_defaults: + se: true + se_context_size: -1 + # encoder / decoder / joint values + enc_hidden: ${model.encoder.d_model} + pred_hidden: 640 + joint_hidden: 640 + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 16 # you may increase batch_size if your memory allows + shuffle: true + num_workers: 8 + pin_memory: false + use_start_end_token: true + trim_silence: false + max_duration: 16.7 # it is set for LibriSpeech, you may need to update it for your dataset + min_duration: 0.1 + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 16 + shuffle: false + num_workers: 8 + pin_memory: false + use_start_end_token: true + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + batch_size: 16 + shuffle: false + num_workers: 8 + pin_memory: false + use_start_end_token: true + + # you may find more detail on how to train a tokenizer at: /scripts/tokenizers/process_asr_text_tokenizer.py + tokenizer: + dir: ??? # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe) + type: wpe # Can be either bpe (SentencePiece tokenizer) or wpe (WordPiece tokenizer) + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + sample_rate: *sample_rate + normalize: "per_feature" + window_size: 0.025 + window_stride: 0.01 + window: "hann" + features: 80 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + pad_to: 0 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + freq_masks: 2 # set to zero to disable it + time_masks: 10 # set to zero to disable it + freq_width: 27 + time_width: 0.05 + + encoder: + _target_: nemo.collections.asr.modules.ConformerEncoder + feat_in: ${model.preprocessor.features} + feat_out: -1 # you may set it if you need different output size other than the default d_model + n_layers: 16 + d_model: 256 + + # Sub-sampling params + subsampling: striding # vggnet or striding + subsampling_factor: 4 # must be power of 2 + subsampling_conv_channels: -1 # set to -1 to make it equal to the d_model + + # Feed forward module's params + ff_expansion_factor: 4 + + # Multi-headed Attention Module's params + self_attention_model: rel_pos # rel_pos or abs_pos + n_heads: 4 + xscaling: true # scales up the inputs by sqrt(d_model) + untie_biases: true # unties the biases of the TransformerXL layers + pos_emb_max_len: 5000 + + # Convolution module's params + conv_kernel_size: 31 + + ### regularization + dropout: 0.1 # The dropout used in most of the Conformer Modules + dropout_emb: 0.0 # The dropout used for embeddings + dropout_att: 0.1 # The dropout for multi-headed attention modules + + decoder: + _target_: nemo.collections.asr.modules.RNNTDecoder + normalization_mode: null + random_state_sampling: false + blank_as_pad: true + + prednet: + pred_hidden: ${model.model_defaults.pred_hidden} + pred_rnn_layers: 1 + t_max: null + dropout: 0.1 + + joint: + _target_: nemo.collections.asr.modules.RNNTJoint + log_softmax: null # 'null' would set it automatically according to CPU/GPU device + + # fused mode + experimental_fuse_loss_wer: true + fused_batch_size: 16 + + jointnet: + joint_hidden: ${model.model_defaults.joint_hidden} + activation: "relu" + dropout: 0.1 + + decoding: + strategy: "greedy_batch" + + # greedy strategy config + greedy: + max_symbols: 30 + + # beam strategy config + beam: + beam_size: 2 + score_norm: true + tsd_max_sym_exp: 50 # for Time Synchronous Decoding + alsd_max_target_len: 2.0 # for Alignment-Length Synchronous Decoding + + optim: + name: adam + lr: 5.0 + # optimizer arguments + betas: [0.9, 0.98] + weight_decay: 1e-6 + + # scheduler setup + sched: + name: NoamAnnealing + d_model: ${model.encoder.d_model} + # scheduler config override + warmup_steps: 10000 + warmup_ratio: null + min_lr: 1e-6 + +trainer: + gpus: -1 # number of GPUs, -1 would use all available GPUs + num_nodes: 1 + max_epochs: 1000 + max_steps: null # computed at runtime if not set + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + accelerator: ddp + accumulate_grad_batches: 1 + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP. + log_every_n_steps: 10 # Interval of logging. + progress_bar_refresh_rate: 10 + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs + sync_batchnorm: true + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: true + create_checkpoint_callback: true + checkpoint_callback_params: + # in case of multiple validation sets, first one is used + monitor: "val_wer" + mode: "min" + save_top_k: 3 + resume_if_exists: false + resume_ignore_no_checkpoint: false + + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null diff --git a/models/tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt.yaml b/models/tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt.yaml new file mode 100644 index 0000000000..306ffe094c --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt.yaml @@ -0,0 +1,237 @@ +name: &name "ConvRNNT5x1" + +model: + sample_rate: 16000 + compute_eval_loss: true + + labels: [ " ", "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", "'" ] + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 32 + trim_silence: true + max_duration: 16.7 + labels: ${model.labels} + shuffle: true + is_tarred: false + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 32 + shuffle: false + labels: ${model.labels} + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + batch_size: 32 + shuffle: false + labels: ${model.labels} + + model_defaults: + repeat: 5 + dropout: 0.0 + separable: true + se: true + se_context_size: -1 + # encoder / decoder / joint values + enc_hidden: 1024 + pred_hidden: 320 + joint_hidden: 320 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: ${model.sample_rate} + window_stride: 0.01 + window: "hann" + features: &n_mels 80 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [13] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [17] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [19] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 256 + repeat: 1 + kernel: [21] + stride: [2] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + - filters: ${model.model_defaults.enc_hidden} + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + decoder: + _target_: nemo.collections.asr.modules.RNNTDecoder + normalization_mode: null + random_state_sampling: false + blank_as_pad: true + + prednet: + pred_hidden: ${model.model_defaults.pred_hidden} + pred_rnn_layers: 1 + t_max: null + dropout: 0.0 + + joint: + _target_: nemo.collections.asr.modules.RNNTJoint + log_softmax: null # sets it according to cpu/gpu device + + # fused mode + experimental_fuse_loss_wer: false + fused_batch_size: 1 + + jointnet: + joint_hidden: ${model.model_defaults.joint_hidden} + activation: "relu" + dropout: 0.0 + + decoding: + strategy: "greedy_batch" + + # greedy strategy config + greedy: + max_symbols: 30 + + # beam strategy config + beam: + beam_size: 2 + score_norm: true + tsd_max_sym_exp: 50 # for Time Synchronous Decoding + alsd_max_target_len: 2.0 # for Alignment-Length Synchronous Decoding + + optim: + name: adam + # _target_: nemo.core.optim.optimizers.Adam + lr: .1 + + # optimizer arguments + betas: [0.9, 0.999] + weight_decay: 0.0001 + + # scheduler setup + sched: + name: CosineAnnealing + + # scheduler config override + warmup_steps: null + warmup_ratio: 0.05 + min_lr: 1e-6 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + precision: 32 + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_wer" + mode: "min" + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null diff --git a/models/tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt_bpe.yaml b/models/tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt_bpe.yaml new file mode 100644 index 0000000000..c259ba581c --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/contextnet_rnnt/config_rnnt_bpe.yaml @@ -0,0 +1,238 @@ +name: &name "ConvRNNTBPE5x1" + +model: + sample_rate: 16000 + compute_eval_loss: true + + tokenizer: + dir: ??? # path to directory which contains either tokenizer.model (bpe) or vocab.txt (for wpe) + type: ??? # Can be either bpe or wpe + + train_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 32 + trim_silence: true + max_duration: 16.7 + labels: [] + shuffle: true + is_tarred: false + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: ${model.sample_rate} + batch_size: 32 + shuffle: false + labels: [] + + test_ds: + manifest_filepath: null + sample_rate: ${model.sample_rate} + batch_size: 32 + shuffle: false + labels: [] + + model_defaults: + repeat: 5 + dropout: 0.0 + separable: true + se: true + se_context_size: -1 + # encoder / decoder / joint values + enc_hidden: 1024 + pred_hidden: 320 + joint_hidden: 320 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: ${model.sample_rate} + window_stride: 0.01 + window: "hann" + features: &n_mels 80 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: 128 + repeat: 1 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [13] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [17] + stride: [2] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + - filters: 256 + repeat: ${model.model_defaults.repeat} + kernel: [19] + stride: [1] + dilation: [1] + dropout: ${model.model_defaults.dropout} + residual: true + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + - filters: 256 + repeat: 1 + kernel: [21] + stride: [2] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + stride_last: true + residual_mode: "stride_add" + + - filters: ${model.model_defaults.enc_hidden} + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: ${model.model_defaults.separable} + se: ${model.model_defaults.se} + se_context_size: ${model.model_defaults.se_context_size} + + decoder: + _target_: nemo.collections.asr.modules.RNNTDecoder + normalization_mode: null + random_state_sampling: false + blank_as_pad: true + + prednet: + pred_hidden: ${model.model_defaults.pred_hidden} + pred_rnn_layers: 1 + t_max: null + dropout: 0.0 + + joint: + _target_: nemo.collections.asr.modules.RNNTJoint + log_softmax: null # sets it according to cpu/gpu device + + # fused mode + experimental_fuse_loss_wer: false + fused_batch_size: 1 + + jointnet: + joint_hidden: ${model.model_defaults.joint_hidden} + activation: "relu" + dropout: 0.0 + + decoding: + strategy: "greedy_batch" + + # greedy strategy config + greedy: + max_symbols: 30 + + # beam strategy config + beam: + beam_size: 2 + score_norm: true + tsd_max_sym_exp: 50 # for Time Synchronous Decoding + alsd_max_target_len: 2.0 # for Alignment-Length Synchronous Decoding + + optim: + name: adam + # _target_: nemo.core.optim.optimizers.Adam + lr: .1 + + # optimizer arguments + betas: [0.9, 0.999] + weight_decay: 0.0001 + + # scheduler setup + sched: + name: CosineAnnealing + + # scheduler config override + warmup_steps: null + warmup_ratio: 0.05 + min_lr: 1e-6 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + precision: 32 + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_wer" + mode: "min" + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null diff --git a/models/tts/NeMo/examples/asr/experimental/quartznet/quartznet_15x5_aug.yaml b/models/tts/NeMo/examples/asr/experimental/quartznet/quartznet_15x5_aug.yaml new file mode 100644 index 0000000000..8ebc2f477c --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/quartznet/quartznet_15x5_aug.yaml @@ -0,0 +1,290 @@ +name: &name "QuartzNet15x5" + +model: + sample_rate: &sample_rate 16000 + repeat: &repeat 5 + dropout: &dropout 0.0 + separable: &separable true + labels: &labels [" ", "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", "'"] + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + trim_silence: True + max_duration: 16.7 + shuffle: True + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + augmentor: + rir_noise_aug: + prob: 0.5 + rir_manifest_path: ??? + rir_tar_filepaths: ??? + rir_prob: 0.5 + noise_manifest_paths: ??? + noise_tar_filepaths: ??? + min_snr_db: [0,0] + max_snr_db: [30,30] + orig_sample_rate: ??? + bg_noise_manifest_paths: ??? + bg_noise_tar_filepaths: ??? + bg_min_snr_db: [10,10] + bg_max_snr_db: [40,40] + bg_orig_sample_rate: ??? + transcode_aug: + prob: 0.1 + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: *labels + batch_size: 32 + shuffle: False + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + stft_conv: false + + spec_augment: + _target_: nemo.collections.asr.modules.SpectrogramAugmentation + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: 1 + residual: false + separable: *separable + stride: [2] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [2] + dropout: *dropout + filters: 512 + kernel: [87] + repeat: 1 + residual: false + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: &enc_filters 1024 + kernel: [1] + repeat: 1 + residual: false + stride: [1] + + decoder: + _target_: nemo.collections.asr.modules.ConvASRDecoder + feat_in: *enc_filters + num_classes: 28 + vocabulary: *labels + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: .01 + # optimizer arguments + betas: [0.8, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # pytorch lightning args + # monitor: val_loss + # reduce_on_plateau: false + + # Scheduler params + warmup_steps: null + warmup_ratio: null + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True + create_wandb_logger: False + wandb_logger_kwargs: + name: null + project: null + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/asr/experimental/sclite/speech_to_text_sclite.py b/models/tts/NeMo/examples/asr/experimental/sclite/speech_to_text_sclite.py new file mode 100644 index 0000000000..b90258212d --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/sclite/speech_to_text_sclite.py @@ -0,0 +1,158 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script is based on speech_to_text_infer.py and allows you to score the hypotheses +with sclite. A local installation from https://github.com/usnistgov/SCTK is required. +Hypotheses and references are first saved in trn format and are scored after applying a glm +file (if provided). +""" + +import errno +import json +import os +import subprocess +from argparse import ArgumentParser + +import torch + +from nemo.collections.asr.metrics.wer import WER +from nemo.collections.asr.models import EncDecCTCModel +from nemo.utils import logging + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +def score_with_sctk(sctk_dir, ref_fname, hyp_fname, out_dir, glm=""): + sclite_path = os.path.join(sctk_dir, "bin", "sclite") + if not os.path.exists(sclite_path): + raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), sclite_path) + # apply glm + if os.path.exists(glm): + rfilter_path = os.path.join(sctk_dir, "bin", "rfilter1") + if not os.path.exists(rfilter_path): + raise FileNotFoundError(errno.ENOENT, os.strerror(errno.ENOENT), rfilter_path) + hypglm = os.path.join(out_dir, os.path.basename(hyp_fname)) + ".glm" + rfilt_cmd = [rfilter_path] + [glm] + with open(hypglm, "w") as hypf, open(hyp_fname, "r") as hyp_in: + subprocess.run(rfilt_cmd, stdin=hyp_in, stdout=hypf) + refglm = os.path.join(out_dir, os.path.basename(ref_fname)) + ".glm" + with open(refglm, "w") as reff, open(ref_fname, "r") as ref_in: + subprocess.run(rfilt_cmd, stdin=ref_in, stdout=reff) + else: + refglm = ref_fname + hypglm = hyp_fname + + _ = subprocess.check_output(f"{sclite_path} -h {hypglm} -r {refglm} -i wsj -o all", shell=True) + + +can_gpu = torch.cuda.is_available() + + +def get_utt_info(manifest_path): + info_list = [] + with open(manifest_path, "r") as utt_f: + for line in utt_f: + utt = json.loads(line) + info_list.append(utt) + + return info_list + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--asr_model", type=str, default="QuartzNet15x5Base-En", required=False, help="Pass: 'QuartzNet15x5Base-En'", + ) + parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data") + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument( + "--dont_normalize_text", + default=False, + action='store_true', + help="Turn off trasnscript normalization. Recommended for non-English.", + ) + parser.add_argument("--out_dir", type=str, required=True, help="Destination dir for output files") + parser.add_argument("--sctk_dir", type=str, required=False, default="", help="Path to sctk root dir") + parser.add_argument("--glm", type=str, required=False, default="", help="Path to glm file") + args = parser.parse_args() + torch.set_grad_enabled(False) + + if not os.path.exists(args.out_dir): + os.makedirs(args.out_dir) + + use_sctk = os.path.exists(args.sctk_dir) + + if args.asr_model.endswith('.nemo'): + logging.info(f"Using local ASR model from {args.asr_model}") + asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model) + else: + logging.info(f"Using NGC cloud ASR model {args.asr_model}") + asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model) + asr_model.setup_test_data( + test_data_config={ + 'sample_rate': 16000, + 'manifest_filepath': args.dataset, + 'labels': asr_model.decoder.vocabulary, + 'batch_size': args.batch_size, + 'normalize_transcripts': not args.dont_normalize_text, + } + ) + if can_gpu: + asr_model = asr_model.cuda() + asr_model.eval() + labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))]) + + wer = WER(vocabulary=asr_model.decoder.vocabulary) + hypotheses = [] + references = [] + all_log_probs = [] + for test_batch in asr_model.test_dataloader(): + if can_gpu: + test_batch = [x.cuda() for x in test_batch] + with autocast(): + log_probs, encoded_len, greedy_predictions = asr_model( + input_signal=test_batch[0], input_signal_length=test_batch[1] + ) + for r in log_probs.cpu().numpy(): + all_log_probs.append(r) + hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions) + for batch_ind in range(greedy_predictions.shape[0]): + reference = ''.join([labels_map[c] for c in test_batch[2][batch_ind].cpu().detach().numpy()]) + references.append(reference) + del test_batch + + info_list = get_utt_info(args.dataset) + hypfile = os.path.join(args.out_dir, "hyp.trn") + reffile = os.path.join(args.out_dir, "ref.trn") + with open(hypfile, "w") as hyp_f, open(reffile, "w") as ref_f: + for i in range(len(hypotheses)): + utt_id = os.path.splitext(os.path.basename(info_list[i]['audio_filepath']))[0] + # rfilter in sctk likes each transcript to have a space at the beginning + hyp_f.write(" " + hypotheses[i] + " (" + utt_id + ")" + "\n") + ref_f.write(" " + references[i] + " (" + utt_id + ")" + "\n") + + if use_sctk: + score_with_sctk(args.sctk_dir, reffile, hypfile, args.out_dir, glm=args.glm) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/experimental/structured/conf/quartznet_15x5.yaml b/models/tts/NeMo/examples/asr/experimental/structured/conf/quartznet_15x5.yaml new file mode 100644 index 0000000000..5bb9925c29 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/structured/conf/quartznet_15x5.yaml @@ -0,0 +1,235 @@ +name: &name "QuartzNet15x5" + +model: + sample_rate: &sample_rate 16000 + repeat: &repeat 5 + dropout: &dropout 0.0 + separable: &separable true + labels: &labels [" ", "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", "'"] + + train_ds: + manifest_filepath: ??? + sample_rate: *sample_rate + labels: *labels + batch_size: 32 + max_duration: 16.7 + is_tarred: False + tarred_audio_filepaths: null + + validation_ds: + labels: *labels + manifest_filepath: ??? + sample_rate: *sample_rate + batch_size: 32 + + preprocessor: + sample_rate: *sample_rate + window_size: 0.02 + window_stride: 0.01 + features: &n_mels 64 + + spec_augment: + rect_freq: 50 + rect_masks: 5 + rect_time: 120 + + encoder: + feat_in: *n_mels + activation: relu + + jasper: + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: 1 + residual: false + separable: *separable + stride: [2] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [33] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 256 + kernel: [39] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [51] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [63] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: 512 + kernel: [75] + repeat: *repeat + residual: true + separable: *separable + stride: [1] + + - dilation: [2] + dropout: *dropout + filters: 512 + kernel: [87] + repeat: 1 + residual: false + separable: *separable + stride: [1] + + - dilation: [1] + dropout: *dropout + filters: &enc_filters 1024 + kernel: [1] + repeat: 1 + residual: false + stride: [1] + + decoder: + feat_in: *enc_filters + num_classes: 28 + vocabulary: *labels + + optim: + name: novograd + lr: .01 + + # optimizer arguments + betas: [0.8, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + + # Scheduler params + warmup_steps: null + warmup_ratio: null + min_lr: 0.0 + +trainer: + gpus: 0 # number of gpus + max_epochs: 5 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + name: *name diff --git a/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_hybrid.py b/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_hybrid.py new file mode 100644 index 0000000000..e396c447e4 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_hybrid.py @@ -0,0 +1,56 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.asr.models import EncDecCTCModel, configs +from nemo.core.config import hydra_runner +from nemo.utils.config_utils import update_model_config +from nemo.utils.exp_manager import exp_manager + + +""" +python speech_to_text_hybrid.py \ + --config-path="conf/quartznet" \ + --config-name="quartznet_15x5" \ + model.train_ds.manifest_filepath="/home/smajumdar/PycharmProjects/NeMo-som/examples/asr/an4/train_manifest.json" \ + model.validation_ds.manifest_filepath="/home/smajumdar/PycharmProjects/NeMo-som/examples/asr/an4/test_manifest.json" \ + trainer.gpus=1 +""" + + +@hydra_runner(config_path="conf/quartznet", config_name="quartznet_15x5") +def main(cfg): + # Generate default asr model config + asr_model_config = configs.EncDecCTCModelConfig() + + # Merge hydra updates with model config + # `drop_missing_subconfig=True` is necessary here. Without it, while the data class will instantiate and be added + # to the config, it contains test_ds.sample_rate = MISSING and test_ds.labels = MISSING. + # This will raise a OmegaConf MissingMandatoryValue error when processing the dataloaders inside + # model_utils.resolve_test_dataloaders(model=self) (used for multi data loader support). + # In general, any operation that tries to use a DictConfig with MISSING in it will fail, + # other than explicit update operations to change MISSING to some actual value. + asr_model_config = update_model_config(asr_model_config, cfg, drop_missing_subconfigs=True) + + # From here on out, its a general OmegaConf DictConfig, directly usable by our code. + trainer = pl.Trainer(**asr_model_config.trainer) + exp_manager(trainer, asr_model_config.get("exp_manager", None)) + asr_model = EncDecCTCModel(cfg=asr_model_config.model, trainer=trainer) + + trainer.fit(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured.py b/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured.py new file mode 100644 index 0000000000..23f33f04f3 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured.py @@ -0,0 +1,146 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from dataclasses import asdict + +import pytorch_lightning as pl + +import nemo.collections.asr as nemo_asr +from nemo.collections.asr.models import EncDecCTCModel, configs +from nemo.utils.exp_manager import exp_manager + +""" +python speech_to_text_structured.py +""" + +# Generate default asr model config +cfg = configs.EncDecCTCModelConfig() + +# set global values +cfg.model.repeat = 5 +cfg.model.separable = True + +# fmt: off +LABELS = [ + " ", "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", "'", +] +# fmt: on + +qn_15x5 = [ + nemo_asr.modules.conv_asr.JasperEncoderConfig( + filters=256, + repeat=1, + kernel=[33], + stride=[2], + separable=cfg.model.separable, + dilation=[1], + dropout=cfg.model.dropout, + residual=False, + ), + nemo_asr.modules.conv_asr.JasperEncoderConfig( + filters=256, + repeat=1, + kernel=[33], + stride=[1], + separable=cfg.model.separable, + dilation=[1], + dropout=cfg.model.dropout, + residual=True, + ), + # ... repeat 14 more times + nemo_asr.modules.conv_asr.JasperEncoderConfig( + filters=1024, repeat=1, kernel=[1], stride=[1], dilation=[1], dropout=cfg.model.dropout, residual=False, + ), +] + + +def main(): + # Update values + # MODEL UPDATES + cfg.name = "Mini QuartzNet" + cfg.model.labels = LABELS + + # train ds + cfg.model.train_ds.manifest_filepath = "" + cfg.model.train_ds.labels = LABELS + cfg.model.train_ds.sample_rate = cfg.model.sample_rate + + # validation ds + cfg.model.validation_ds.manifest_filepath = "" + cfg.model.validation_ds.labels = LABELS + cfg.model.validation_ds.sample_rate = cfg.model.sample_rate + + # del `test_ds` does not work! + # Refer - https://stackoverflow.com/questions/58119758/how-to-remove-dataclass-attributes + # Hydra/OmegaConf dont allow custom .asdict() methods either + # For now, explicitly set parameters + cfg.model.test_ds.sample_rate = cfg.model.sample_rate + cfg.model.test_ds.labels = cfg.model.labels + + # preprocessor + cfg.model.preprocessor.sample_rate = cfg.model.sample_rate + + # spec aug + cfg.model.spec_augment.rect_masks = 5 + cfg.model.spec_augment.rect_freq = 50 + cfg.model.spec_augment.rect_time = 120 + + # encoder + cfg.model.encoder.feat_in = cfg.model.preprocessor.features + cfg.model.encoder.activation = 'relu' + cfg.model.encoder.jasper = qn_15x5 + + # decoder + cfg.model.decoder.feat_in = qn_15x5[-1].filters + cfg.model.decoder.num_classes = len(LABELS) + cfg.model.decoder.vocabulary = LABELS + + # optim + cfg.model.optim.name = 'novograd' + cfg.model.optim.lr = 0.01 + + # `betas` dont exist inside the base config, + # so they cannot be added as such! + # Same for `weight_decay`. + cfg.model.optim.betas = [0.8, 0.5] + cfg.model.optim.weight_decay = 0.001 + + # sched + # As parameters such as warmup_steps and warmup_ratio + # dont exist inside the shell config, these values are not added! + cfg.model.optim.sched.name = "CosineAnnealing" + cfg.model.optim.sched.warmup_steps = None + cfg.model.optim.sched.warmup_ratio = 0.01 + + # Trainer config + cfg.trainer.gpus = 1 + cfg.trainer.max_epochs = 5 + + # Exp Manager config + cfg.exp_manager.name = cfg.name + + # Note usage of asdict + trainer = pl.Trainer(**asdict(cfg.trainer)) + exp_manager(trainer, asdict(cfg.exp_manager)) + asr_model = EncDecCTCModel(cfg=cfg.model, trainer=trainer) + + trainer.fit(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured_v2.py b/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured_v2.py new file mode 100644 index 0000000000..1fc9cc8341 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/structured/speech_to_text_structured_v2.py @@ -0,0 +1,90 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from dataclasses import asdict + +import pytorch_lightning as pl + +from nemo.collections.asr.models import EncDecCTCModel, configs +from nemo.core.config import modelPT, optimizers, schedulers +from nemo.utils.exp_manager import exp_manager + +""" +python speech_to_text_structured_v2.py +""" + +# fmt: off +LABELS = [ + " ", "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", "'", +] + +optim_cfg = optimizers.NovogradParams( + lr=0.01, + betas=(0.8, 0.5), + weight_decay=0.001 +) + +sched_cfg = schedulers.CosineAnnealingParams( + warmup_steps=None, + warmup_ratio=None, + min_lr=0.0, +) +# fmt: on + + +def main(): + # NeMo Model config + cfg = modelPT.NemoConfig(name='Custom QuartzNet') + + # Generate default asr model config + builder = configs.EncDecCTCModelConfigBuilder(name='quartznet_15x5') + + # set model global values + builder.set_labels(LABELS) + builder.set_optim(cfg=optim_cfg, sched_cfg=sched_cfg) + + model_cfg = builder.build() + + # set the model config to the NeMo Model + cfg.model = model_cfg + + # Update values + # MODEL UPDATES + # train ds + model_cfg.train_ds.manifest_filepath = "" + + # validation ds + model_cfg.validation_ds.manifest_filepath = "" + + # Trainer config + cfg.trainer.gpus = 1 + cfg.trainer.max_epochs = 5 + + # Exp Manager config + cfg.exp_manager.name = cfg.name + + # Note usage of asdict + trainer = pl.Trainer(**asdict(cfg.trainer)) + exp_manager(trainer, asdict(cfg.exp_manager)) + asr_model = EncDecCTCModel(cfg=cfg.model, trainer=trainer) + + trainer.fit(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain.yaml b/models/tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain.yaml new file mode 100644 index 0000000000..00268fcefe --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain.yaml @@ -0,0 +1,123 @@ +name: &name "Wav2vec_Pretrain" +# For more details on the model parameters, see nemo/collections/asr/models/wav2vec/wav2vec_config.py +# These parameters are based off the FairSeq implementation. +# See here: https://github.com/pytorch/fairseq/blob/master/examples/wav2vec/config/pretraining/wav2vec2_base_librispeech.yaml +model: + transformer_encoder: + dropout: 0.1 # Dropout probability for the transformer + conv: # Config for convolutional model that generates positional embeddings required for attention layer + conv_pos: 128 # Number of filters for convolutional positional embeddings + conv_pos_groups: 16 # Number of groups for convolutional positional embeddings + encoder: + encoder_layers: 12 # Number of encoder layers in transformer model + encoder_layerdrop: 0.05 # Probability of dropping transformer layers + embedding_dim: 768 # Encoder embedding dim + ffn_embedding_dim: 3072 # Encoder embedding dim for feed forward + num_attention_heads: 8 # Number of encoder attention heads + dropout: 0.1 # Dropout probability for transformer encoder + activation_fn: gelu # Activation for transformer + layer_norm_first: False # Apply layer norm first within the transformer + loss: + prob_ppl_weight: 0.1 # Weight applied to quantized prob perplexity loss + feature_loss_weight: 0 # Weight applied to feature L2 Norm + quantizer: + quantize_targets: True # Quantize targets + quantize_input: False # Quantize inputs + same_quantizer: False # Use the same quantizer for inputs and targets (If both enabled) + latent_vars: 320 # Number of latent variables in each group of the codebook + latent_groups: 2 # Number of groups within the codebook + latent_dim: 0 # If greater than 0, use the dim for latest variables else inferred by final dim / latent groups + latent_temp: [2, 0.5, 0.999995] # Quantize temperature (start, stop, decay factor) + conv_feature_encoder: # Config for feature extractor taking raw audio waveform and converting to features + extractor_mode: default # Mode for feature extractor. [Default, layer_norm] + conv_bias: False # Include bias in convolution feature extractor model + masking: + mask_prob: 0.65 # Probability of replacing token with mask + mask_type: static # Used to select configuration to compute mask lengths [static, uniform, normal, poisson] + mask_other: 0 # Secondary mask used for complex distributions (see help in compute_mask_indices) + mask_length: 10 # Length of mask when masking time steps + no_mask_overlap: False # Whether to allow masks to overlap + mask_min_space: 1 # Minimum space beetween spans (if no overlap is enabled) + mask_channel_prob: 0 # Probability of replacing a feature with 0 + mask_channel_type: static # Used to select configuration to compute mask channel lengths [static, uniform, normal, poisson] + mask_channel_other: 0 # Secondary mask argument (used for more complex distributions (see help in compute_mask_indices) + mask_channel_length: 10 # Length of masks for features (channels) + no_mask_channel_overlap: False # Whether to allow channel masks to overlap + mask_channel_min_space: 1 # Minimum space between spans (if no overlap is enabled) + dropout_input: 0.1 # Dropout applied to input raw features + dropout_features: 0.1 # Dropout applied to the features generator by convolutions + final_dim: 0 # Project final representations and targets to this dimension + n_negatives: 100 # Number of negatives to sample from the same audio sample + cross_sample_negatives: 0 # Number of negatives to sample from any sample in the batch + codebook_negatives: 0 # Number of negative examples in codebook + negatives_from_everywhere: False # Sample negatives from everywhere, not just masked states + logit_temp: 0.1 # Temperature to divide logits by + target_glu: False # Adds project and applies GLU to targets + feature_grad_mult: 0.1 # Multiply extracted feature gradients + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: True + num_workers: 4 + labels: [ ] # No vocab needed for pre-training + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: False + num_workers: 1 + labels: [ ] + test_ds: + manifest_filepath: null + sample_rate: 16000 + batch_size: 32 + shuffle: False + num_workers: 4 + labels: [ ] + optim: + name: adamw + lr: 0.0005 + eps: 1e-06 + # optimizer arguments + betas: [ 0.9, 0.98 ] + weight_decay: 0.01 + + # scheduler setup + sched: + name: PolynomialDecayAnnealing + min_lr: 0.000 + # Scheduler params + warmup_steps: null + warmup_ratio: 0.1 +trainer: + gpus: 0 # number of gpus + num_nodes: 1 + max_epochs: 100 + max_steps: null # computed at runtime if not set + val_check_interval: 0.5 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + accelerator: ddp + accumulate_grad_batches: 1 + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP. + log_every_n_steps: 10 # Interval of logging. + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs + sync_batchnorm: true + + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: true + create_checkpoint_callback: true + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null + resume_if_exists: false + resume_ignore_no_checkpoint: false \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain_large.yaml b/models/tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain_large.yaml new file mode 100644 index 0000000000..4a92a281c4 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/wav2vec/configs/wav2vec_pretrain_large.yaml @@ -0,0 +1,123 @@ +name: &name "Wav2vec_large_Pretrain" +# For more details on the model parameters, see nemo/collections/asr/models/wav2vec/wav2vec_config.py +# These parameters are based off the FairSeq implementation. +# See here: https://github.com/pytorch/fairseq/blob/master/examples/wav2vec/config/pretraining/wav2vec2_large_librivox.yaml +model: + transformer_encoder: + dropout: 0.0 # Dropout probability for the transformer + conv: # Config for convolutional model that generates positional embeddings required for attention layer + conv_pos: 128 # Number of filters for convolutional positional embeddings + conv_pos_groups: 16 # Number of groups for convolutional positional embeddings + encoder: + encoder_layers: 24 # Number of encoder layers in transformer model + encoder_layerdrop: 0.0 # Probability of dropping transformer layers + embedding_dim: 1024 # Encoder embedding dim + ffn_embedding_dim: 4096 # Encoder embedding dim for feed forward + num_attention_heads: 16 # Number of encoder attention heads + dropout: 0.0 # Dropout probability for transformer encoder + activation_fn: gelu # Activation for transformer + layer_norm_first: True # Apply layer norm first within the transformer + loss: + prob_ppl_weight: 0.1 # Weight applied to quantized prob perplexity loss + feature_loss_weight: 0 # Weight applied to feature L2 Norm + quantizer: + quantize_targets: True # Quantize targets + quantize_input: False # Quantize inputs + same_quantizer: False # Use the same quantizer for inputs and targets (If both enabled) + latent_vars: 320 # Number of latent variables in each group of the codebook + latent_groups: 2 # Number of groups within the codebook + latent_dim: 0 # If greater than 0, use the dim for latest variables else inferred by final dim / latent groups + latent_temp: [2, 0.1, 0.999995] # Quantize temperature (start, stop, decay factor) + conv_feature_encoder: # Config for feature extractor taking raw audio waveform and converting to features + extractor_mode: layer_norm # Mode for feature extractor. [default, layer_norm] + conv_bias: True # Include bias in convolution feature extractor model + masking: + mask_prob: 0.65 # Probability of replacing token with mask + mask_type: static # Used to select configuration to compute mask lengths [static, uniform, normal, poisson] + mask_other: 0 # Secondary mask used for complex distributions (see help in compute_mask_indices) + mask_length: 10 # Length of mask when masking time steps + no_mask_overlap: False # Whether to allow masks to overlap + mask_min_space: 1 # Minimum space beetween spans (if no overlap is enabled) + mask_channel_prob: 0 # Probability of replacing a feature with 0 + mask_channel_type: static # Used to select configuration to compute mask channel lengths [static, uniform, normal, poisson] + mask_channel_other: 0 # Secondary mask argument (used for more complex distributions (see help in compute_mask_indices) + mask_channel_length: 10 # Length of masks for features (channels) + no_mask_channel_overlap: False # Whether to allow channel masks to overlap + mask_channel_min_space: 1 # Minimum space between spans (if no overlap is enabled) + dropout_input: 0.0 # Dropout applied to input raw features + dropout_features: 0.0 # Dropout applied to the features generator by convolutions + final_dim: 768 # Project final representations and targets to this dimension + n_negatives: 100 # Number of negatives to sample from the same audio sample + cross_sample_negatives: 0 # Number of negatives to sample from any sample in the batch + codebook_negatives: 0 # Number of negative examples in codebook + negatives_from_everywhere: False # Sample negatives from everywhere, not just masked states + logit_temp: 0.1 # Temperature to divide logits by + target_glu: False # Adds project and applies GLU to targets + feature_grad_mult: 1.0 # Multiply extracted feature gradients + + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: True + num_workers: 4 + labels: [ ] # No vocab needed for pre-training + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + batch_size: 32 + shuffle: False + num_workers: 1 + labels: [ ] + test_ds: + manifest_filepath: null + sample_rate: 16000 + batch_size: 32 + shuffle: False + num_workers: 4 + labels: [ ] + optim: + name: adamw + lr: 0.0005 + eps: 1e-06 + # optimizer arguments + betas: [ 0.9, 0.98 ] + weight_decay: 0.01 + + # scheduler setup + sched: + name: PolynomialDecayAnnealing + min_lr: 0.000 + # Scheduler params + warmup_steps: null + warmup_ratio: 0.1 +trainer: + gpus: 0 # number of gpus + num_nodes: 1 + max_epochs: 100 + max_steps: null # computed at runtime if not set + val_check_interval: 0.5 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + accelerator: ddp + accumulate_grad_batches: 1 + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP. + log_every_n_steps: 10 # Interval of logging. + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + check_val_every_n_epoch: 1 # number of evaluations on validation every n epochs + sync_batchnorm: true + + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: true + create_checkpoint_callback: true + create_wandb_logger: false + wandb_logger_kwargs: + name: null + project: null + resume_if_exists: false + resume_ignore_no_checkpoint: false \ No newline at end of file diff --git a/models/tts/NeMo/examples/asr/experimental/wav2vec/wav2vec_pretrain.py b/models/tts/NeMo/examples/asr/experimental/wav2vec/wav2vec_pretrain.py new file mode 100644 index 0000000000..2e5c1db320 --- /dev/null +++ b/models/tts/NeMo/examples/asr/experimental/wav2vec/wav2vec_pretrain.py @@ -0,0 +1,97 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.asr.models.wav2vec.wav2vec_model import Wav2VecEncoderModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +""" +Pre-train a wav2vec 2.0 transformer model on audio. Uses a contrastive loss function to pre-train on unlabelled audio, +using a task similar to masked language modeling in NLP. In wav2vec, we mask portions of the audio +and the model is trained by minimising the distance of the ground truth for the masked section, +using the ground truth quantized codebook representation. Distractors are obtained from other time steps. +See :class:`Wav2VecCriterion` for more information. + +Reference: https://arxiv.org/abs/2006.11477 + + python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \ + model.train_ds.manifest_filepath="./examples/asr/train.tsv" \ + model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \ + trainer.gpus=1 \ + trainer.max_epochs=100 + +Basic run (on CPU for 50 epochs): + python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \ + model.train_ds.manifest_filepath="./examples/asr/train.tsv" \ + model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \ + trainer.gpus=1 \ + trainer.max_epochs=50 + +Using wav2vec-large with mixed precision: + python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \ + --config-name=wav2vec_pretrain_large \ + model.train_ds.manifest_filepath="./examples/asr/train.tsv" \ + model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \ + trainer.gpus=1 \ + trainer.max_epochs=100 \ + trainer.precision=16 + +Add PyTorch Lightning Trainer arguments from CLI: + python wav2vec.py \ + ... \ + +trainer.fast_dev_run=true + +Hydra logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/.hydra)" +PTL logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/lightning_logs)" + +Override some args of optimizer: + python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \ + model.train_ds.manifest_filepath="./examples/asr/train.tsv" \ + model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \ + trainer.gpus=2 \ + trainer.max_epochs=2 \ + model.optim.args.params.betas=[0.8,0.5] \ + model.optim.args.params.weight_decay=0.0001 + +Override optimizer entirely + python examples/asr/experimental/wav2vec/wav2vec_pretrain.py \ + model.train_ds.manifest_filepath="./examples/asr/train.tsv" \ + model.validation_ds.manifest_filepath="./examples/asr/valid.tsv" \ + trainer.gpus=2 \ + trainer.max_epochs=2 \ + ~model.optim.args \ + +model.optim.args.betas=[0.8,0.5]\ + +model.optim.args.weight_decay=0.0005 + +""" + + +@hydra_runner(config_path="configs", config_name="wav2vec_pretrain") +def main(cfg: DictConfig): + logging.info("Application config\n" + OmegaConf.to_yaml(cfg)) + + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + wav2vec_encoder_model = Wav2VecEncoderModel(cfg=cfg.model, trainer=trainer) + trainer.fit(wav2vec_encoder_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/quantization/speech_to_text_calibrate.py b/models/tts/NeMo/examples/asr/quantization/speech_to_text_calibrate.py new file mode 100644 index 0000000000..264806c7b1 --- /dev/null +++ b/models/tts/NeMo/examples/asr/quantization/speech_to_text_calibrate.py @@ -0,0 +1,160 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Script for calibrating a pretrained ASR model for quantization +""" + +from argparse import ArgumentParser + +import torch +from omegaconf import open_dict + +from nemo.collections.asr.models import EncDecCTCModel +from nemo.utils import logging + +try: + from pytorch_quantization import calib + from pytorch_quantization import nn as quant_nn + from pytorch_quantization import quant_modules + from pytorch_quantization.tensor_quant import QuantDescriptor +except ImportError: + raise ImportError( + "pytorch-quantization is not installed. Install from " + "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization." + ) + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +can_gpu = torch.cuda.is_available() + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'", + ) + parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data") + parser.add_argument("--batch_size", type=int, default=256) + parser.add_argument( + "--dont_normalize_text", + default=False, + action='store_false', + help="Turn off trasnscript normalization. Recommended for non-English.", + ) + parser.add_argument('--num_calib_batch', default=1, type=int, help="Number of batches for calibration.") + parser.add_argument('--calibrator', type=str, choices=["max", "histogram"], default="max") + parser.add_argument('--percentile', nargs='+', type=float, default=[99.9, 99.99, 99.999, 99.9999]) + parser.add_argument("--amp", action="store_true", help="Use AMP in calibration.") + parser.set_defaults(amp=False) + + args = parser.parse_args() + torch.set_grad_enabled(False) + + # Initialize quantization + quant_desc_input = QuantDescriptor(calib_method=args.calibrator) + quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input) + quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input) + quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input) + + if args.asr_model.endswith('.nemo'): + logging.info(f"Using local ASR model from {args.asr_model}") + asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True) + with open_dict(asr_model_cfg): + asr_model_cfg.encoder.quantize = True + asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg) + + else: + logging.info(f"Using NGC cloud ASR model {args.asr_model}") + asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True) + with open_dict(asr_model_cfg): + asr_model_cfg.encoder.quantize = True + asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg) + + asr_model.setup_test_data( + test_data_config={ + 'sample_rate': 16000, + 'manifest_filepath': args.dataset, + 'labels': asr_model.decoder.vocabulary, + 'batch_size': args.batch_size, + 'normalize_transcripts': args.dont_normalize_text, + 'shuffle': True, + } + ) + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + if can_gpu: + asr_model = asr_model.cuda() + asr_model.eval() + + # Enable calibrators + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer): + if module._calibrator is not None: + module.disable_quant() + module.enable_calib() + else: + module.disable() + + for i, test_batch in enumerate(asr_model.test_dataloader()): + if can_gpu: + test_batch = [x.cuda() for x in test_batch] + if args.amp: + with autocast(): + _ = asr_model(input_signal=test_batch[0], input_signal_length=test_batch[1]) + else: + _ = asr_model(input_signal=test_batch[0], input_signal_length=test_batch[1]) + if i >= args.num_calib_batch: + break + + # Save calibrated model(s) + model_name = args.asr_model.replace(".nemo", "") if args.asr_model.endswith(".nemo") else args.asr_model + if not args.calibrator == "histogram": + compute_amax(asr_model, method="max") + asr_model.save_to(F"{model_name}-max-{args.num_calib_batch*args.batch_size}.nemo") + else: + for percentile in args.percentile: + print(F"{percentile} percentile calibration") + compute_amax(asr_model, method="percentile") + asr_model.save_to(F"{model_name}-percentile-{percentile}-{args.num_calib_batch*args.batch_size}.nemo") + + for method in ["mse", "entropy"]: + print(F"{method} calibration") + compute_amax(asr_model, method=method) + asr_model.save_to(F"{model_name}-{method}-{args.num_calib_batch*args.batch_size}.nemo") + + +def compute_amax(model, **kwargs): + for name, module in model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer): + if module._calibrator is not None: + if isinstance(module._calibrator, calib.MaxCalibrator): + module.load_calib_amax() + else: + module.load_calib_amax(**kwargs) + print(F"{name:40}: {module}") + if can_gpu: + model.cuda() + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer.py b/models/tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer.py new file mode 100644 index 0000000000..b072fe8e90 --- /dev/null +++ b/models/tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer.py @@ -0,0 +1,216 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Script for post training quantization of ASR models +""" + +import collections +from argparse import ArgumentParser +from pprint import pprint + +import torch +from omegaconf import open_dict + +from nemo.collections.asr.metrics.wer import WER, word_error_rate +from nemo.collections.asr.models import EncDecCTCModel +from nemo.utils import logging + +try: + from pytorch_quantization import nn as quant_nn + from pytorch_quantization import quant_modules +except ImportError: + raise ImportError( + "pytorch-quantization is not installed. Install from " + "https://github.com/NVIDIA/TensorRT/tree/master/tools/pytorch-quantization." + ) + + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +can_gpu = torch.cuda.is_available() + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'", + ) + parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data") + parser.add_argument("--wer_target", type=float, default=None, help="used by test") + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--wer_tolerance", type=float, default=1.0, help="used by test") + parser.add_argument( + "--dont_normalize_text", + default=False, + action='store_false', + help="Turn off trasnscript normalization. Recommended for non-English.", + ) + parser.add_argument( + "--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric" + ) + parser.add_argument('--sensitivity', action="store_true", help="Perform sensitivity analysis") + parser.add_argument('--onnx', action="store_true", help="Export to ONNX") + parser.add_argument('--quant-disable-keyword', type=str, nargs='+', help='disable quantizers by keyword') + args = parser.parse_args() + torch.set_grad_enabled(False) + + quant_modules.initialize() + + if args.asr_model.endswith('.nemo'): + logging.info(f"Using local ASR model from {args.asr_model}") + asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True) + with open_dict(asr_model_cfg): + asr_model_cfg.encoder.quantize = True + asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg) + + else: + logging.info(f"Using NGC cloud ASR model {args.asr_model}") + asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True) + with open_dict(asr_model_cfg): + asr_model_cfg.encoder.quantize = True + asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg) + asr_model.setup_test_data( + test_data_config={ + 'sample_rate': 16000, + 'manifest_filepath': args.dataset, + 'labels': asr_model.decoder.vocabulary, + 'batch_size': args.batch_size, + 'normalize_transcripts': args.dont_normalize_text, + } + ) + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + if can_gpu: + asr_model = asr_model.cuda() + asr_model.eval() + + if args.quant_disable_keyword: + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer): + for keyword in args.quant_disable_keyword: + if keyword in name: + logging.warning(F"Disable {name}") + module.disable() + + labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))]) + wer = WER(vocabulary=asr_model.decoder.vocabulary, use_cer=args.use_cer) + wer_quant = evaluate(asr_model, labels_map, wer) + logging.info(f'Got WER of {wer_quant}. Tolerance was {args.wer_tolerance}') + + if args.sensitivity: + if wer_quant < args.wer_tolerance: + logging.info("Tolerance is already met. Skip sensitivity analyasis.") + return + quant_layer_names = [] + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer): + module.disable() + layer_name = name.replace("._input_quantizer", "").replace("._weight_quantizer", "") + if layer_name not in quant_layer_names: + quant_layer_names.append(layer_name) + logging.info(F"{len(quant_layer_names)} quantized layers found.") + + # Build sensitivity profile + quant_layer_sensitivity = {} + for i, quant_layer in enumerate(quant_layer_names): + logging.info(F"Enable {quant_layer}") + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer) and quant_layer in name: + module.enable() + logging.info(F"{name:40}: {module}") + + # Eval the model + wer_value = evaluate(asr_model, labels_map, wer) + logging.info(F"WER: {wer_value}") + quant_layer_sensitivity[quant_layer] = args.wer_tolerance - wer_value + + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer) and quant_layer in name: + module.disable() + logging.info(F"{name:40}: {module}") + + # Skip most sensitive layers until WER target is met + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer): + module.enable() + quant_layer_sensitivity = collections.OrderedDict(sorted(quant_layer_sensitivity.items(), key=lambda x: x[1])) + pprint(quant_layer_sensitivity) + skipped_layers = [] + for quant_layer, _ in quant_layer_sensitivity.items(): + for name, module in asr_model.named_modules(): + if isinstance(module, quant_nn.TensorQuantizer): + if quant_layer in name: + logging.info(F"Disable {name}") + if not quant_layer in skipped_layers: + skipped_layers.append(quant_layer) + module.disable() + wer_value = evaluate(asr_model, labels_map, wer) + if wer_value <= args.wer_tolerance: + logging.info( + F"WER tolerance {args.wer_tolerance} is met by skipping {len(skipped_layers)} sensitive layers." + ) + print(skipped_layers) + export_onnx(args, asr_model) + return + raise ValueError(f"WER tolerance {args.wer_tolerance} can not be met with any layer quantized!") + + export_onnx(args, asr_model) + + +def export_onnx(args, asr_model): + if args.onnx: + if args.asr_model.endswith("nemo"): + onnx_name = args.asr_model.replace(".nemo", ".onnx") + else: + onnx_name = args.asr_model + logging.info(F"Export to {onnx_name}") + quant_nn.TensorQuantizer.use_fb_fake_quant = True + asr_model.export(onnx_name, onnx_opset_version=13) + quant_nn.TensorQuantizer.use_fb_fake_quant = False + + +def evaluate(asr_model, labels_map, wer): + # Eval the model + hypotheses = [] + references = [] + for test_batch in asr_model.test_dataloader(): + if can_gpu: + test_batch = [x.cuda() for x in test_batch] + with autocast(): + log_probs, encoded_len, greedy_predictions = asr_model( + input_signal=test_batch[0], input_signal_length=test_batch[1] + ) + hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions) + for batch_ind in range(greedy_predictions.shape[0]): + seq_len = test_batch[3][batch_ind].cpu().detach().numpy() + seq_ids = test_batch[2][batch_ind].cpu().detach().numpy() + reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]]) + references.append(reference) + del test_batch + wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=wer.use_cer) + + return wer_value + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer_trt.py b/models/tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer_trt.py new file mode 100644 index 0000000000..747bd5bd43 --- /dev/null +++ b/models/tts/NeMo/examples/asr/quantization/speech_to_text_quant_infer_trt.py @@ -0,0 +1,226 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Script for inference ASR models using TensorRT +""" + +import collections +import os +import time +from argparse import ArgumentParser +from pprint import pprint + +import numpy as np +import pycuda.autoinit +import pycuda.driver as cuda +import tensorrt as trt +import torch +from omegaconf import open_dict + +from nemo.collections.asr.metrics.wer import WER, word_error_rate +from nemo.collections.asr.models import EncDecCTCModel +from nemo.utils import logging + +TRT_LOGGER = trt.Logger() + + +can_gpu = torch.cuda.is_available() + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'", + ) + parser.add_argument( + "--asr_onnx", + type=str, + default="./QuartzNet15x5Base-En-max-32.onnx", + help="Pass: 'QuartzNet15x5Base-En-max-32.onnx'", + ) + parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data") + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument( + "--dont_normalize_text", + default=False, + action='store_false', + help="Turn off trasnscript normalization. Recommended for non-English.", + ) + parser.add_argument( + "--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric" + ) + parser.add_argument('--qat', action="store_true", help="Use onnx file exported from QAT tools") + args = parser.parse_args() + torch.set_grad_enabled(False) + + if args.asr_model.endswith('.nemo'): + logging.info(f"Using local ASR model from {args.asr_model}") + asr_model_cfg = EncDecCTCModel.restore_from(restore_path=args.asr_model, return_config=True) + with open_dict(asr_model_cfg): + asr_model_cfg.encoder.quantize = True + asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model, override_config_path=asr_model_cfg) + + else: + logging.info(f"Using NGC cloud ASR model {args.asr_model}") + asr_model_cfg = EncDecCTCModel.from_pretrained(model_name=args.asr_model, return_config=True) + with open_dict(asr_model_cfg): + asr_model_cfg.encoder.quantize = True + asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model, override_config_path=asr_model_cfg) + asr_model.setup_test_data( + test_data_config={ + 'sample_rate': 16000, + 'manifest_filepath': args.dataset, + 'labels': asr_model.decoder.vocabulary, + 'batch_size': args.batch_size, + 'normalize_transcripts': args.dont_normalize_text, + } + ) + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + if can_gpu: + asr_model = asr_model.cuda() + asr_model.eval() + labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))]) + wer = WER(vocabulary=asr_model.decoder.vocabulary, use_cer=args.use_cer) + wer_result = evaluate(asr_model, args.asr_onnx, labels_map, wer, args.qat) + logging.info(f'Got WER of {wer_result}.') + + +def get_min_max_input_shape(asr_model): + max_shape = (1, 64, 1) + min_shape = (64, 64, 99999) + for test_batch in asr_model.test_dataloader(): + test_batch = [x.cuda() for x in test_batch] + processed_signal, processed_signal_length = asr_model.preprocessor( + input_signal=test_batch[0], length=test_batch[1] + ) + shape = processed_signal.cpu().numpy().shape + if shape[0] > max_shape[0]: + max_shape = (shape[0], *max_shape[1:]) + if shape[0] < min_shape[0]: + min_shape = (shape[0], *min_shape[1:]) + if shape[2] > max_shape[2]: + max_shape = (*max_shape[0:2], shape[2]) + if shape[2] < min_shape[2]: + min_shape = (*min_shape[0:2], shape[2]) + return min_shape, max_shape + + +def build_trt_engine(asr_model, onnx_path, qat): + trt_engine_path = "{}.trt".format(onnx_path) + if os.path.exists(trt_engine_path): + return trt_engine_path + + min_input_shape, max_input_shape = get_min_max_input_shape(asr_model) + workspace_size = 512 + with trt.Builder(TRT_LOGGER) as builder: + network_flags = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + if qat: + network_flags |= 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_PRECISION) + with builder.create_network(flags=network_flags) as network, trt.OnnxParser( + network, TRT_LOGGER + ) as parser, builder.create_builder_config() as builder_config: + parser.parse_from_file(onnx_path) + builder_config.max_workspace_size = workspace_size * (1024 * 1024) + if qat: + builder_config.set_flag(trt.BuilderFlag.INT8) + + profile = builder.create_optimization_profile() + profile.set_shape("audio_signal", min=min_input_shape, opt=max_input_shape, max=max_input_shape) + builder_config.add_optimization_profile(profile) + + engine = builder.build_engine(network, builder_config) + serialized_engine = engine.serialize() + with open(trt_engine_path, "wb") as fout: + fout.write(serialized_engine) + return trt_engine_path + + +def trt_inference(stream, trt_ctx, d_input, d_output, input_signal, input_signal_length): + print("infer with shape: {}".format(input_signal.shape)) + + trt_ctx.set_binding_shape(0, input_signal.shape) + assert trt_ctx.all_binding_shapes_specified + + h_output = cuda.pagelocked_empty(tuple(trt_ctx.get_binding_shape(1)), dtype=np.float32) + + h_input_signal = cuda.register_host_memory(np.ascontiguousarray(input_signal.cpu().numpy().ravel())) + cuda.memcpy_htod_async(d_input, h_input_signal, stream) + trt_ctx.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle) + cuda.memcpy_dtoh_async(h_output, d_output, stream) + stream.synchronize() + + greedy_predictions = torch.tensor(h_output).argmax(dim=-1, keepdim=False) + return greedy_predictions + + +def evaluate(asr_model, asr_onnx, labels_map, wer, qat): + # Eval the model + hypotheses = [] + references = [] + stream = cuda.Stream() + vocabulary_size = len(labels_map) + 1 + engine_file_path = build_trt_engine(asr_model, asr_onnx, qat) + with open(engine_file_path, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime: + trt_engine = runtime.deserialize_cuda_engine(f.read()) + trt_ctx = trt_engine.create_execution_context() + + profile_shape = trt_engine.get_profile_shape(profile_index=0, binding=0) + print("profile shape min:{}, opt:{}, max:{}".format(profile_shape[0], profile_shape[1], profile_shape[2])) + max_input_shape = profile_shape[2] + input_nbytes = trt.volume(max_input_shape) * trt.float32.itemsize + d_input = cuda.mem_alloc(input_nbytes) + max_output_shape = [max_input_shape[0], vocabulary_size, (max_input_shape[-1] + 1) // 2] + output_nbytes = trt.volume(max_output_shape) * trt.float32.itemsize + d_output = cuda.mem_alloc(output_nbytes) + + for test_batch in asr_model.test_dataloader(): + if can_gpu: + test_batch = [x.cuda() for x in test_batch] + processed_signal, processed_signal_length = asr_model.preprocessor( + input_signal=test_batch[0], length=test_batch[1] + ) + + greedy_predictions = trt_inference( + stream, + trt_ctx, + d_input, + d_output, + input_signal=processed_signal, + input_signal_length=processed_signal_length, + ) + hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions) + for batch_ind in range(greedy_predictions.shape[0]): + seq_len = test_batch[3][batch_ind].cpu().detach().numpy() + seq_ids = test_batch[2][batch_ind].cpu().detach().numpy() + reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]]) + references.append(reference) + del test_batch + wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=wer.use_cer) + + return wer_value + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/speech_to_label.py b/models/tts/NeMo/examples/asr/speech_to_label.py new file mode 100644 index 0000000000..369b0f026b --- /dev/null +++ b/models/tts/NeMo/examples/asr/speech_to_label.py @@ -0,0 +1,146 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +# Task 1: Speech Command + +## Preparing the dataset +Use the `process_speech_commands_data.py` script under /scripts/dataset_processing in order to prepare the dataset. + +```sh +python /scripts/dataset_processing/process_speech_commands_data.py \ + --data_root= \ + --data_version= \ + --class_split= \ + --rebalance \ + --log +``` + +## Train to convergence +```sh +python speech_to_label.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath=["",""] \ + trainer.gpus=2 \ + trainer.accelerator="ddp" \ + trainer.max_epochs=200 \ + exp_manager.create_wandb_logger=True \ + exp_manager.wandb_logger_kwargs.name="MatchboxNet-3x1x64-v1" \ + exp_manager.wandb_logger_kwargs.project="MatchboxNet-v1" \ + +trainer.precision=16 \ + +trainer.amp_level=O1 # needed if using PyTorch < 1.6 +``` + + +# Task 2: Voice Activity Detection + +## Preparing the dataset +Use the `process_vad_data.py` script under /scripts/dataset_processing in order to prepare the dataset. + +```sh +python process_vad_data.py \ + --out_dir= \ + --speech_data_root= \ + --background_data_root= \ + --rebalance_method=<'under' or 'over' of 'fixed'> \ + --log + (Optional --demo (for demonstration in tutorial). If you want to use your own background noise data, make sure to delete --demo) +``` + +## Train to convergence +```sh +python speech_to_label.py \ + --config-path= + --config-name= \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath=["",""] \ + trainer.gpus=2 \ + trainer.accelerator="ddp" \ + trainer.max_epochs=200 \ + exp_manager.create_wandb_logger=True \ + exp_manager.wandb_logger_kwargs.name="MatchboxNet-3x1x64-vad" \ + exp_manager.wandb_logger_kwargs.project="MatchboxNet-vad" \ + +trainer.precision=16 \ + +trainer.amp_level=O1 # needed if using PyTorch < 1.6 +``` + +# Optional: Use tarred dataset to speed up data loading. Apply to both tasks. +## Prepare tarred dataset. + Prepare ONE manifest that contains all training data you would like to include. Validation should use non-tarred dataset. + Note that it's possible that tarred datasets impacts validation scores because it drop values in order to have same amount of files per tarfile; + Scores might be off since some data is missing. + + Use the `convert_to_tarred_audio_dataset.py` script under /scripts/speech_recognition in order to prepare tarred audio dataset. + For details, please see TarredAudioToClassificationLabelDataset in /nemo/collections/asr/data/audio_to_label.py + +python speech_to_label.py \ + --config-path= + --config-name= \ + model.train_ds.manifest_filepath= \ + model.train_ds.is_tarred=True \ + model.train_ds.tarred_audio_filepaths= \ + +model.train_ds.num_worker= \ + model.validation_ds.manifest_filepath=\ + trainer.gpus=2 \ + trainer.accelerator="ddp" \ + trainer.max_epochs=200 \ + exp_manager.create_wandb_logger=True \ + exp_manager.wandb_logger_kwargs.name="MatchboxNet-3x1x64-vad" \ + exp_manager.wandb_logger_kwargs.project="MatchboxNet-vad" \ + +trainer.precision=16 \ + +trainer.amp_level=O1 # needed if using PyTorch < 1.6 + +# Fine-tune a model + +For documentation on fine-tuning this model, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations + +# Pretrained Models + +For documentation on existing pretrained models, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speech_classification/results.html# + +""" +import pytorch_lightning as pl +from omegaconf import OmegaConf + +from nemo.collections.asr.models import EncDecClassificationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf/matchboxnet", config_name="matchboxnet_3x1x64_v1") +def main(cfg): + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + asr_model = EncDecClassificationModel(cfg=cfg.model, trainer=trainer) + + # Initialize the weights of the model from another model, if provided via config + asr_model.maybe_init_from_pretrained_checkpoint(cfg) + + trainer.fit(asr_model) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + trainer = pl.Trainer(gpus=gpu) + if asr_model.prepare_test(trainer): + trainer.test(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/speech_to_text.py b/models/tts/NeMo/examples/asr/speech_to_text.py new file mode 100644 index 0000000000..ced4cb8e3a --- /dev/null +++ b/models/tts/NeMo/examples/asr/speech_to_text.py @@ -0,0 +1,105 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +# Training the model + +Basic run (on CPU for 50 epochs): + python examples/asr/speech_to_text.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath="" \ + trainer.gpus=0 \ + trainer.max_epochs=50 + + +Add PyTorch Lightning Trainer arguments from CLI: + python speech_to_text.py \ + ... \ + +trainer.fast_dev_run=true + +Hydra logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/.hydra)" +PTL logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/lightning_logs)" + +Override some args of optimizer: + python speech_to_text.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath="./an4/train_manifest.json" \ + model.validation_ds.manifest_filepath="./an4/test_manifest.json" \ + trainer.gpus=2 \ + trainer.max_epochs=2 \ + model.optim.args.betas=[0.8,0.5] \ + model.optim.args.weight_decay=0.0001 + +Override optimizer entirely + python speech_to_text.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath="./an4/train_manifest.json" \ + model.validation_ds.manifest_filepath="./an4/test_manifest.json" \ + trainer.gpus=2 \ + trainer.max_epochs=2 \ + model.optim.name=adamw \ + model.optim.lr=0.001 \ + ~model.optim.args \ + +model.optim.args.betas=[0.8,0.5]\ + +model.optim.args.weight_decay=0.0005 + +# Fine-tune a model + +For documentation on fine-tuning this model, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations + +# Pretrained Models + +For documentation on existing pretrained models, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/results.html + +""" + +import pytorch_lightning as pl +from omegaconf import OmegaConf + +from nemo.collections.asr.models import EncDecCTCModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="config") +def main(cfg): + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + asr_model = EncDecCTCModel(cfg=cfg.model, trainer=trainer) + + # Initialize the weights of the model from another model, if provided via config + asr_model.maybe_init_from_pretrained_checkpoint(cfg) + + trainer.fit(asr_model) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + test_trainer = pl.Trainer( + gpus=gpu, + precision=trainer.precision, + amp_level=trainer.accelerator_connector.amp_level, + amp_backend=cfg.trainer.get("amp_backend", "native"), + ) + if asr_model.prepare_test(test_trainer): + test_trainer.test(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/speech_to_text_bpe.py b/models/tts/NeMo/examples/asr/speech_to_text_bpe.py new file mode 100644 index 0000000000..090535f3ac --- /dev/null +++ b/models/tts/NeMo/examples/asr/speech_to_text_bpe.py @@ -0,0 +1,101 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +# Preparing the Tokenizer for the dataset +Use the `process_asr_text_tokenizer.py` script under /scripts/tokenizers/ in order to prepare the tokenizer. + +```sh +python /scripts/tokenizers/process_asr_text_tokenizer.py \ + --manifest= + OR + --data_file= \ + --data_root="" \ + --vocab_size= \ + --tokenizer=<"spe" or "wpe"> \ + --no_lower_case \ + --spe_type=<"unigram", "bpe", "char" or "word"> \ + --spe_character_coverage=1.0 \ + --log +``` + +# Training the model +```sh +python speech_to_text_bpe.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath= \ + model.validation_ds.manifest_filepath= \ + model.tokenizer.dir= \ + model.tokenizer.type= \ + trainer.gpus=-1 \ + trainer.accelerator="ddp" \ + trainer.max_epochs=100 \ + model.optim.name="adamw" \ + model.optim.lr=0.001 \ + model.optim.betas=[0.9,0.999] \ + model.optim.weight_decay=0.0001 \ + model.optim.sched.warmup_steps=2000 + exp_manager.create_wandb_logger=True \ + exp_manager.wandb_logger_kwargs.name="" \ + exp_manager.wandb_logger_kwargs.project="" +``` + +# Fine-tune a model + +For documentation on fine-tuning this model, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations + +# Pretrained Models + +For documentation on existing pretrained models, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/results.html + +""" + +import pytorch_lightning as pl +from omegaconf import OmegaConf + +from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf/citrinet/", config_name="config_bpe") +def main(cfg): + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + asr_model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer) + + # Initialize the weights of the model from another model, if provided via config + asr_model.maybe_init_from_pretrained_checkpoint(cfg) + + trainer.fit(asr_model) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + test_trainer = pl.Trainer( + gpus=gpu, + precision=trainer.precision, + amp_level=trainer.accelerator_connector.amp_level, + amp_backend=cfg.trainer.get("amp_backend", "native"), + ) + if asr_model.prepare_test(test_trainer): + test_trainer.test(asr_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/asr/speech_to_text_infer.py b/models/tts/NeMo/examples/asr/speech_to_text_infer.py new file mode 100644 index 0000000000..bbba1a6119 --- /dev/null +++ b/models/tts/NeMo/examples/asr/speech_to_text_infer.py @@ -0,0 +1,114 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script serves three goals: + (1) Demonstrate how to use NeMo Models outside of PytorchLightning + (2) Shows example of batch ASR inference + (3) Serves as CI test for pre-trained checkpoint +""" + +from argparse import ArgumentParser + +import torch + +from nemo.collections.asr.metrics.wer import WER, word_error_rate +from nemo.collections.asr.models import EncDecCTCModel +from nemo.utils import logging + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +can_gpu = torch.cuda.is_available() + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--asr_model", type=str, default="QuartzNet15x5Base-En", required=True, help="Pass: 'QuartzNet15x5Base-En'", + ) + parser.add_argument("--dataset", type=str, required=True, help="path to evaluation data") + parser.add_argument("--batch_size", type=int, default=4) + parser.add_argument("--wer_tolerance", type=float, default=1.0, help="used by test") + parser.add_argument( + "--dont_normalize_text", + default=False, + action='store_true', + help="Turn off trasnscript normalization. Recommended for non-English.", + ) + parser.add_argument( + "--use_cer", default=False, action='store_true', help="Use Character Error Rate as the evaluation metric" + ) + args = parser.parse_args() + torch.set_grad_enabled(False) + + if args.asr_model.endswith('.nemo'): + logging.info(f"Using local ASR model from {args.asr_model}") + asr_model = EncDecCTCModel.restore_from(restore_path=args.asr_model) + else: + logging.info(f"Using NGC cloud ASR model {args.asr_model}") + asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model) + + asr_model.preprocessor.featurizer.pad_to = 0 + asr_model.preprocessor.featurizer.dither = 0.0 + + asr_model.setup_test_data( + test_data_config={ + 'sample_rate': 16000, + 'manifest_filepath': args.dataset, + 'labels': asr_model.decoder.vocabulary, + 'batch_size': args.batch_size, + 'normalize_transcripts': not args.dont_normalize_text, + } + ) + if can_gpu: + asr_model = asr_model.cuda() + asr_model.eval() + labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))]) + wer = WER(vocabulary=asr_model.decoder.vocabulary) + hypotheses = [] + references = [] + for test_batch in asr_model.test_dataloader(): + if can_gpu: + test_batch = [x.cuda() for x in test_batch] + with autocast(): + log_probs, encoded_len, greedy_predictions = asr_model( + input_signal=test_batch[0], input_signal_length=test_batch[1] + ) + hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions) + for batch_ind in range(greedy_predictions.shape[0]): + seq_len = test_batch[3][batch_ind].cpu().detach().numpy() + seq_ids = test_batch[2][batch_ind].cpu().detach().numpy() + reference = ''.join([labels_map[c] for c in seq_ids[0:seq_len]]) + references.append(reference) + del test_batch + + wer_value = word_error_rate(hypotheses=hypotheses, references=references, use_cer=args.use_cer) + if not args.use_cer: + if wer_value > args.wer_tolerance: + raise ValueError(f"got wer of {wer_value}. it was higher than {args.wer_tolerance}") + logging.info(f'Got WER of {wer_value}. Tolerance was {args.wer_tolerance}') + else: + logging.info(f'Got CER of {wer_value}') + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/speech_to_text_rnnt.py b/models/tts/NeMo/examples/asr/speech_to_text_rnnt.py new file mode 100644 index 0000000000..39d3be9ca7 --- /dev/null +++ b/models/tts/NeMo/examples/asr/speech_to_text_rnnt.py @@ -0,0 +1,104 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +# Training the model + +Basic run (on CPU for 50 epochs): + python examples/asr/speech_to_text_rnnt.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath="" \ + model.validation_ds.manifest_filepath="" \ + trainer.gpus=0 \ + trainer.max_epochs=50 + + +Add PyTorch Lightning Trainer arguments from CLI: + python speech_to_text_rnnt.py \ + ... \ + +trainer.fast_dev_run=true + +Hydra logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/.hydra)" +PTL logs will be found in "$(./outputs/$(date +"%y-%m-%d")/$(date +"%H-%M-%S")/lightning_logs)" + +Override some args of optimizer: + python speech_to_text_rnnt.py \ + --config-path="experimental/contextnet_rnnt" \ + --config-name="config_rnnt" \ + model.train_ds.manifest_filepath="./an4/train_manifest.json" \ + model.validation_ds.manifest_filepath="./an4/test_manifest.json" \ + trainer.gpus=2 \ + trainer.precision=16 \ + trainer.max_epochs=2 \ + model.optim.betas=[0.8,0.5] \ + model.optim.weight_decay=0.0001 + +Override optimizer entirely + python speech_to_text_rnnt.py \ + --config-path="experimental/contextnet_rnnt" \ + --config-name="config_rnnt" \ + model.train_ds.manifest_filepath="./an4/train_manifest.json" \ + model.validation_ds.manifest_filepath="./an4/test_manifest.json" \ + trainer.gpus=2 \ + trainer.precision=16 \ + trainer.max_epochs=2 \ + model.optim.name=adamw \ + model.optim.lr=0.001 \ + ~model.optim.args \ + +model.optim.args.betas=[0.8,0.5]\ + +model.optim.args.weight_decay=0.0005 + +# Fine-tune a model + +For documentation on fine-tuning this model, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations + +""" + +import pytorch_lightning as pl +from omegaconf import OmegaConf + +from nemo.collections.asr.models import EncDecRNNTModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="experimental/contextnet_rnnt", config_name="config_rnnt") +def main(cfg): + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + asr_model = EncDecRNNTModel(cfg=cfg.model, trainer=trainer) + + # Initialize the weights of the model from another model, if provided via config + asr_model.maybe_init_from_pretrained_checkpoint(cfg) + + trainer.fit(asr_model) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + test_trainer = pl.Trainer( + gpus=gpu, + precision=trainer.precision, + amp_level=trainer.accelerator_connector.amp_level, + amp_backend=cfg.trainer.get("amp_backend", "native"), + ) + if asr_model.prepare_test(test_trainer): + test_trainer.test(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/speech_to_text_rnnt_bpe.py b/models/tts/NeMo/examples/asr/speech_to_text_rnnt_bpe.py new file mode 100644 index 0000000000..91e41deaf0 --- /dev/null +++ b/models/tts/NeMo/examples/asr/speech_to_text_rnnt_bpe.py @@ -0,0 +1,96 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +# Preparing the Tokenizer for the dataset +Use the `process_asr_text_tokenizer.py` script under /scripts/tokenizers/ in order to prepare the tokenizer. + +```sh +python /scripts/tokenizers/process_asr_text_tokenizer.py \ + --manifest= + OR + --data_file= \ + --data_root="" \ + --vocab_size= \ + --tokenizer=<"spe" or "wpe"> \ + --no_lower_case \ + --spe_type=<"unigram", "bpe", "char" or "word"> \ + --spe_character_coverage=1.0 \ + --log +``` + +# Training the model +```sh +python speech_to_text_rnnt_bpe.py \ + # (Optional: --config-path= --config-name=) \ + model.train_ds.manifest_filepath= \ + model.validation_ds.manifest_filepath= \ + model.tokenizer.dir= \ + model.tokenizer.type= \ + trainer.gpus=-1 \ + trainer.accelerator="ddp" \ + trainer.max_epochs=100 \ + model.optim.name="adamw" \ + model.optim.lr=0.001 \ + model.optim.betas=[0.9,0.999] \ + model.optim.weight_decay=0.0001 \ + model.optim.sched.warmup_steps=2000 + exp_manager.create_wandb_logger=True \ + exp_manager.wandb_logger_kwargs.name="" \ + exp_manager.wandb_logger_kwargs.project="" +``` + +# Fine-tune a model + +For documentation on fine-tuning this model, please visit - +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/configs.html#fine-tuning-configurations + +""" + +import pytorch_lightning as pl +from omegaconf import OmegaConf + +from nemo.collections.asr.models import EncDecRNNTBPEModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="experimental/contextnet_rnnt", config_name="config_rnnt_bpe") +def main(cfg): + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + asr_model = EncDecRNNTBPEModel(cfg=cfg.model, trainer=trainer) + + # Initialize the weights of the model from another model, if provided via config + asr_model.maybe_init_from_pretrained_checkpoint(cfg) + + trainer.fit(asr_model) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + test_trainer = pl.Trainer( + gpus=gpu, + precision=trainer.precision, + amp_level=trainer.accelerator_connector.amp_level, + amp_backend=cfg.trainer.get("amp_backend", "native"), + ) + if asr_model.prepare_test(test_trainer): + test_trainer.test(asr_model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/transcribe_speech.py b/models/tts/NeMo/examples/asr/transcribe_speech.py new file mode 100644 index 0000000000..b186c22757 --- /dev/null +++ b/models/tts/NeMo/examples/asr/transcribe_speech.py @@ -0,0 +1,167 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 contextlib +import glob +import json +import os +from dataclasses import dataclass +from typing import Optional + +import pytorch_lightning as pl +import torch +from omegaconf import OmegaConf + +from nemo.collections.asr.metrics.rnnt_wer import RNNTDecodingConfig +from nemo.collections.asr.models import ASRModel +from nemo.core.config import hydra_runner +from nemo.utils import logging, model_utils + + +""" +# Transcribe audio +# Arguments +# model_path: path to .nemo ASR checkpoint +# pretrained_name: name of pretrained ASR model (from NGC registry) +# audio_dir: path to directory with audio files +# dataset_manifest: path to dataset JSON manifest file (in NeMo format) +# +# ASR model can be specified by either "model_path" or "pretrained_name". +# Data for transcription can be defined with either "audio_dir" or "dataset_manifest". +# Results are returned in a JSON manifest file. + +python transcribe_speech.py \ + model_path=null \ + pretrained_name=null \ + audio_dir="" \ + dataset_manifest="" \ + output_filename="" +""" + + +@dataclass +class TranscriptionConfig: + # Required configs + model_path: Optional[str] = None # Path to a .nemo file + pretrained_name: Optional[str] = None # Name of a pretrained model + audio_dir: Optional[str] = None # Path to a directory which contains audio files + dataset_manifest: Optional[str] = None # Path to dataset's JSON manifest + + # General configs + output_filename: Optional[str] = None + batch_size: int = 32 + cuda: Optional[bool] = None # will switch to cuda if available, defaults to CPU otherwise + amp: bool = False + audio_type: str = "wav" + + # decoding strategy for RNNT models + rnnt_decoding: RNNTDecodingConfig = RNNTDecodingConfig() + + +@hydra_runner(config_name="TranscriptionConfig", schema=TranscriptionConfig) +def main(cfg: TranscriptionConfig): + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + + if cfg.model_path is None and cfg.pretrained_name is None: + raise ValueError("Both cfg.model_path and cfg.pretrained_name cannot be None!") + if cfg.audio_dir is None and cfg.dataset_manifest is None: + raise ValueError("Both cfg.audio_dir and cfg.dataset_manifest cannot be None!") + + # setup GPU + if cfg.cuda is None: + cfg.cuda = torch.cuda.is_available() + + if type(cfg.cuda) == int: + device_id = int(cfg.cuda) + else: + device_id = 0 + + device = torch.device(f'cuda:{device_id}' if cfg.cuda else 'cpu') + + # setup model + if cfg.model_path is not None: + # restore model from .nemo file path + model_cfg = ASRModel.restore_from(restore_path=cfg.model_path, return_config=True) + classpath = model_cfg.target # original class path + imported_class = model_utils.import_class_by_path(classpath) # type: ASRModel + logging.info(f"Restoring model : {imported_class.__name__}") + asr_model = imported_class.restore_from(restore_path=cfg.model_path, map_location=device) # type: ASRModel + model_name = os.path.splitext(os.path.basename(cfg.model_path))[0] + else: + # restore model by name + asr_model = ASRModel.from_pretrained(model_name=cfg.pretrained_name, map_location=device) # type: ASRModel + model_name = cfg.pretrained_name + + trainer = pl.Trainer(gpus=int(cfg.cuda)) + asr_model.set_trainer(trainer) + asr_model = asr_model.eval() + + # Setup decoding strategy + if hasattr(asr_model, 'change_decoding_strategy'): + asr_model.change_decoding_strategy(cfg.rnnt_decoding) + + # get audio filenames + if cfg.audio_dir is not None: + filepaths = list(glob.glob(os.path.join(cfg.audio_dir, f"*.{cfg.audio_type}"))) + else: + # get filenames from manifest + filepaths = [] + with open(cfg.dataset_manifest, 'r') as f: + for line in f: + item = json.loads(line) + filepaths.append(item['audio_filepath']) + logging.info(f"\nTranscribing {len(filepaths)} files...\n") + + # setup AMP (optional) + if cfg.amp and torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'): + logging.info("AMP enabled!\n") + autocast = torch.cuda.amp.autocast + else: + + @contextlib.contextmanager + def autocast(): + yield + + # transcribe audio + with autocast(): + with torch.no_grad(): + transcriptions = asr_model.transcribe(filepaths, batch_size=cfg.batch_size) + logging.info(f"Finished transcribing {len(filepaths)} files !") + + if cfg.output_filename is None: + # create default output filename + if cfg.audio_dir is not None: + cfg.output_filename = os.path.dirname(os.path.join(cfg.audio_dir, '.')) + '.json' + else: + cfg.output_filename = cfg.dataset_manifest.replace('.json', f'_{model_name}.json') + + logging.info(f"Writing transcriptions into file: {cfg.output_filename}") + + with open(cfg.output_filename, 'w', encoding='utf-8') as f: + if cfg.audio_dir is not None: + for idx, text in enumerate(transcriptions): + item = {'audio_filepath': filepaths[idx], 'pred_text': text} + f.write(json.dumps(item) + "\n") + else: + with open(cfg.dataset_manifest, 'r') as fr: + for idx, line in enumerate(fr): + item = json.loads(line) + item['pred_text'] = transcriptions[idx] + f.write(json.dumps(item) + "\n") + + logging.info("Finished writing predictions !") + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/asr/vad_infer.py b/models/tts/NeMo/examples/asr/vad_infer.py new file mode 100644 index 0000000000..ba036029f2 --- /dev/null +++ b/models/tts/NeMo/examples/asr/vad_infer.py @@ -0,0 +1,171 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +During inference, we perform frame-level prediction by two approaches: + 1) shift the window of length time_length (e.g. 0.63s) by shift_length (e.g. 10ms) to generate the frame and use the prediction of the window to represent the label for the frame; + [this script demonstrate how to do this approach] + 2) generate predictions with overlapping input segments. Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments. + [get frame level prediction by this script and use vad_overlap_posterior.py in NeMo/scripts/voice_activity_detection + One can also find posterior about converting frame level prediction + to speech/no-speech segment in start and end times format in that script.] + + Image https://raw.githubusercontent.com/NVIDIA/NeMo/main/tutorials/asr/images/vad_post_overlap_diagram.png + will help you understand this method. + +Usage: +python vad_infer.py --vad_model="vad_marblenet" --dataset= --out_dir='frame/demo' --time_length=0.63 + +""" + + +import json +import logging +import os +from argparse import ArgumentParser + +import torch + +from nemo.collections.asr.models import EncDecClassificationModel +from nemo.collections.asr.parts.utils.vad_utils import get_vad_stream_status, prepare_manifest +from nemo.utils import logging + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--vad_model", type=str, default="MatchboxNet-VAD-3x2", required=False, help="Pass: 'MatchboxNet-VAD-3x2'" + ) + parser.add_argument( + "--dataset", + type=str, + required=True, + help="Path of json file of evaluation data. Audio files should have unique names.", + ) + parser.add_argument("--out_dir", type=str, default="vad_frame", help="Dir of your vad outputs") + parser.add_argument("--time_length", type=float, default=0.63) + parser.add_argument("--shift_length", type=float, default=0.01) + parser.add_argument("--normalize_audio", type=bool, default=False) + parser.add_argument("--num_workers", type=float, default=20) + parser.add_argument("--split_duration", type=float, default=400) + parser.add_argument( + "--dont_auto_split", + default=False, + action='store_true', + help="Whether to automatically split manifest entry by split_duration to avoid potential CUDA out of memory issue.", + ) + + args = parser.parse_args() + + torch.set_grad_enabled(False) + + if args.vad_model.endswith('.nemo'): + logging.info(f"Using local VAD model from {args.vad_model}") + vad_model = EncDecClassificationModel.restore_from(restore_path=args.vad_model) + else: + logging.info(f"Using NGC cloud VAD model {args.vad_model}") + vad_model = EncDecClassificationModel.from_pretrained(model_name=args.vad_model) + + if not os.path.exists(args.out_dir): + os.mkdir(args.out_dir) + + # Prepare manifest for streaming VAD + manifest_vad_input = args.dataset + if not args.dont_auto_split: + logging.info("Split long audio file to avoid CUDA memory issue") + logging.debug("Try smaller split_duration if you still have CUDA memory issue") + config = { + 'manifest_filepath': manifest_vad_input, + 'time_length': args.time_length, + 'split_duration': args.split_duration, + 'num_workers': args.num_workers, + } + manifest_vad_input = prepare_manifest(config) + else: + logging.warning( + "If you encounter CUDA memory issue, try splitting manifest entry by split_duration to avoid it." + ) + + # setup_test_data + vad_model.setup_test_data( + test_data_config={ + 'vad_stream': True, + 'sample_rate': 16000, + 'manifest_filepath': manifest_vad_input, + 'labels': ['infer',], + 'num_workers': args.num_workers, + 'shuffle': False, + 'time_length': args.time_length, + 'shift_length': args.shift_length, + 'trim_silence': False, + 'normalize_audio': args.normalize_audio, + } + ) + + vad_model = vad_model.to(device) + vad_model.eval() + + time_unit = int(args.time_length / args.shift_length) + trunc = int(time_unit / 2) + trunc_l = time_unit - trunc + all_len = 0 + + data = [] + for line in open(args.dataset, 'r'): + file = json.loads(line)['audio_filepath'].split("/")[-1] + data.append(file.split(".wav")[0]) + logging.info(f"Inference on {len(data)} audio files/json lines!") + + status = get_vad_stream_status(data) + for i, test_batch in enumerate(vad_model.test_dataloader()): + test_batch = [x.to(device) for x in test_batch] + with autocast(): + log_probs = vad_model(input_signal=test_batch[0], input_signal_length=test_batch[1]) + probs = torch.softmax(log_probs, dim=-1) + pred = probs[:, 1] + + if status[i] == 'start': + to_save = pred[:-trunc] + elif status[i] == 'next': + to_save = pred[trunc:-trunc_l] + elif status[i] == 'end': + to_save = pred[trunc_l:] + else: + to_save = pred + + all_len += len(to_save) + outpath = os.path.join(args.out_dir, data[i] + ".frame") + with open(outpath, "a") as fout: + for f in range(len(to_save)): + fout.write('{0:0.4f}\n'.format(to_save[f])) + del test_batch + if status[i] == 'end' or status[i] == 'single': + logging.debug(f"Overall length of prediction of {data[i]} is {all_len}!") + all_len = 0 + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/cv/mnist_lenet5_image_classification_pure_lightning.py b/models/tts/NeMo/examples/cv/mnist_lenet5_image_classification_pure_lightning.py new file mode 100644 index 0000000000..ff03a9e146 --- /dev/null +++ b/models/tts/NeMo/examples/cv/mnist_lenet5_image_classification_pure_lightning.py @@ -0,0 +1,65 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from dataclasses import dataclass + +# from hydra.utils import instantiate +import pytorch_lightning as ptl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.cv.models import MNISTLeNet5, MNISTLeNet5Config +from nemo.core.config import Config, TrainerConfig, hydra_runner +from nemo.utils import logging + + +@dataclass +class AppConfig(Config): + """ + This is structured config for this application. + + Args: + name: Description of the application. + trainer: configuration of the trainer. + model: configuation of the model. + """ + + name: str = "Training of a LeNet-5 Model using a pure PyTorchLightning approach - using DDP on 2 GPUs." + trainer: TrainerConfig = TrainerConfig(gpus=2, accelerator="ddp") + model: MNISTLeNet5Config = MNISTLeNet5Config() + + +@hydra_runner(schema=AppConfig, config_name="AppConfig") +def main(cfg: DictConfig): + # Show configuration - user can modify every parameter from command line! + logging.info("Application config\n" + OmegaConf.to_yaml(cfg)) + + # The "model" - with dataloader/dataset inside of it. + lenet5 = MNISTLeNet5(cfg.model) + + # Setup train data loader and optimizer + lenet5.setup_training_data() + + # Setup optimizer and scheduler + lenet5.setup_optimization() + + # Create trainer. + trainer = ptl.Trainer(**(cfg.trainer)) + # trainer = instantiate(cfg.trainer) + + # Train. + trainer.fit(model=lenet5) + + +if __name__ == "__main__": + main() # TODO: No cfg in function call, and no hydra runner diff --git a/models/tts/NeMo/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml b/models/tts/NeMo/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml new file mode 100644 index 0000000000..7ecc5652a6 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml @@ -0,0 +1,115 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + + +pretrained_model: null # pretrained QAModel model from list_available_models() +do_training: true # training mode, for testing change to false +trainer: + gpus: 8 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1] + num_nodes: 1 + max_epochs: 3 + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + gradient_clip_val: 1.0 + amp_level: O1 # O1/O2 for mixed precision + precision: 16 # Should be set to 16 for O1 and O2 to enable the AMP. + accelerator: ddp + log_every_n_steps: 5 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + +model: + nemo_path: null # filename to save the model and associated artifacts to .nemo file + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null + + language_model: + pretrained_model_name: bert-base-cased + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null + + encoder: + dropout: 0.1 + + dataset: + max_seq_length: 128 # the maximum length BERT supports is 512 + max_num_cat_slot: 6 # maximum number of different categorical slots per service in dataset + max_num_noncat_slot: 12 # maximum number of different non-categorical slots per service in dataset + max_value_per_cat_slot: 12 # maximum number of different categorical slot values per service in dataset + max_num_intent: 4 # maximum number of different intents per service in dataset + task_name: sgd_single_domain # or from [sgd_all, sgd_all_single, sgd_multi_domain, debug_sample] + state_tracker: nemotracker # or baseline + data_dir: ??? + dialogues_example_dir: ??? # cache location to store or load preprocessed data from. To use cache, specify `use_cache` + subsample: false # balances negative and positive training examples for improved performance + use_cache: false # uses a cache to store the processed dataset, you may use it for large datasets for speed up + use_fuzzy_match: true # Whether to use fuzzy string matching when comparing non-categorical slot values. Should be set to False when conducting multiwoz style evaluation. + joint_acc_across_turn: false # Whether to compute joint goal accuracy across turn instead of across service. Should be set to True when conducting multiwoz style evaluation. + + train_ds: + ds_item: "train" + batch_size: 128 + shuffle: true + num_workers: 3 + drop_last: false + pin_memory: false + + validation_ds: + ds_item: ["dev"] # ["dev", "test"] + batch_size: 8 + shuffle: false + num_workers: 3 + drop_last: false + pin_memory: false + + test_ds: + ds_item: ["test"] + batch_size: 8 + shuffle: false + num_workers: 3 + drop_last: false + pin_memory: false + + optim: + name: adamw + lr: 1e-4 + # optimizer arguments + betas: [0.9, 0.999] + weight_decay: 0.01 + + # scheduler setup + sched: + name: PolynomialDecayAnnealing + # Scheduler params + warmup_steps: null + warmup_ratio: 0.02 + last_epoch: -1 + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + +exp_manager: + exp_dir: null # exp_dir for your experiment, if None, defaults to "./nemo_experiments" + name: "SGDQA" # The name of your model + create_tensorboard_logger: True # Whether you want exp_manger to create a tb logger + create_checkpoint_callback: True # Whether you want exp_manager to create a modelcheckpoint callback + resume_if_exists: false + resume_ignore_no_checkpoint: false diff --git a/models/tts/NeMo/examples/nlp/dialogue_state_tracking/sgd_qa.py b/models/tts/NeMo/examples/nlp/dialogue_state_tracking/sgd_qa.py new file mode 100644 index 0000000000..85dbae19af --- /dev/null +++ b/models/tts/NeMo/examples/nlp/dialogue_state_tracking/sgd_qa.py @@ -0,0 +1,145 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved. +# Copyright 2019 The Google Research Authors. +# +# 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. + +""" +This script contains an example of how to train and test the NeMo SGD-QA Model (https://arxiv.org/abs/2105.08049). +The SGD-QA model is a fast multi-pass schema-guided state-tracking model, that is trained on the Google schema-guided state tracking dataset (https://arxiv.org/abs/1909.05855). +The model takes dialogue as input and outputs the dialogue state, which includes slot-value pairs. +The model consists of two components: a neural natural language understanding model (NLU), and a rule-based state tracker. +The NLU takes in a dialogue turn and different schema (entity) information options and outputs their match score. The state tracker takes the highest rated entities and composes +the dialogue state across turns. + +***Data format*** +The SGD-QA requires a JSON schema and dialogue files for each dataset split. +In the following we will show an example for a service entry in the schema file. +* service_name +* description +* slots + * name + * description + * is_categorical + * possible values +* intents + * name + * description + * required_slots (not used) + * is_transactional (not used) + * optional_slots (not used) + * result_slots (not used) + + +In the following we will show an example for a dialogue. More information at https://arxiv.org/abs/1909.05855 +* dialogue_id +* services +* turns + * frames + * actions + * act + * slot + * values + * service + * slots + * exclusive_end + * slot + * start + * state + * active_intent + * requeste_slots + * slot_values + * speaker - [USER, SYSTEM] + * utterance + + +***Downloading the dataset*** +# git clone https://github.com/google-research-datasets/dstc8-schema-guided-dialogue.git + +***Setting the configs*** +The model and the PT trainer are defined in a config file that declares multiple important sections. +The most important ones are: + model: All arguments that are related to the Model - language model, SGD-QA encoder and decoder, loss, optimizer, + schedulers, and datasets/data loaders. + trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs, + precision level, etc. + +This script uses the `/examples/nlp/dialogue_state_tracking/conf/sgdqa_config.yaml` config file +by default. You may update the config file from the file directly. The other option is to set another config file via command-line arguments by `--config-name=CONFIG_FILE_PATH'. + + +***Model Training*** +# python sgd_qa.py + do_training=True + model.dataset.data_dir= + model.dataset.dialogues_example_dir= + model.validation_ds.ds_item= + trainer.max_epochs= + trainer.gpus=[] + + +***Model Evaluation*** +# python sgd_qa.py + do_training=False + model.test_ds.ds_item= + +To load a pretrained checkpoint from the cloud prior to training (e.g. for fine-tuning) or evaluation you can set cfg.from_pretrained=. You can find all pretrained model names by using +SGDQAModel.list_available_models(). To load a local checkpoint use model.restore_from() + +""" + +import os + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.dialogue_state_tracking.sgdqa_model import SGDQAModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="sgdqa_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + if cfg.pretrained_model or (cfg.model.nemo_path and os.path.exists(cfg.model.nemo_path)): + if cfg.pretrained_model: + logging.info(f'Loading pretrained model {cfg.pretrained_model}') + model = SGDQAModel.from_pretrained(cfg.pretrained_model) + else: + logging.info(f'Restoring model from {cfg.model.nemo_path}') + model = SGDQAModel.restore_from(cfg.model.nemo_path) + if cfg.do_training: + model.setup_training_data(train_data_config=cfg.model.train_ds) + model.setup_multiple_validation_data(val_data_config=cfg.model.validation_ds) + else: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + model = SGDQAModel(cfg.model, trainer=trainer) + + if cfg.do_training: + trainer.fit(model) + if cfg.model.nemo_path: + model.save_to(cfg.model.nemo_path) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.ds_item is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + trainer = pl.Trainer(gpus=gpu) + model.setup_multiple_test_data(test_data_config=cfg.model.test_ds) + if model.prepare_test(trainer): + trainer.test(model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/entity_linking/build_index.py b/models/tts/NeMo/examples/nlp/entity_linking/build_index.py new file mode 100644 index 0000000000..eeba5c8313 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/entity_linking/build_index.py @@ -0,0 +1,201 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os +import pickle as pkl +import random +from argparse import ArgumentParser + +import h5py +import numpy as np +import torch +from omegaconf import DictConfig, OmegaConf +from sklearn.decomposition import PCA +from tqdm import tqdm + +from nemo.collections.nlp.models import EntityLinkingModel +from nemo.utils import logging + +try: + import faiss +except ModuleNotFoundError: + logging.warning("Faiss is required for building the index. Please install faiss-gpu") + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + +def build_index(cfg: DictConfig, model: object): + """ + Builds faiss index from index dataset specified in the config. + + Args: + cfg (DictConfig): Config file specifying index parameters + model (object): Encoder model + """ + + # Get index dataset embeddings + # PCA model exists and index embeddings have already been PCAed, no need to re-extract/PCA them + if cfg.apply_pca and os.path.isfile(cfg.pca.pca_save_name) and os.path.isfile(cfg.pca_embeddings_save_name): + logging.info("Loading reduced dimensionality embeddings") + embeddings = h5py.File(cfg.pca_embeddings_save_name, "r") + embeddings = embeddings[cfg.index_ds.name][:] + + elif os.path.isfile(cfg.embedding_save_name): + logging.info("Loading previously extracted index dataset embeddings") + embeddings = h5py.File(cfg.embedding_save_name, "r") + embeddings = embeddings[cfg.index_ds.name][:] + + else: + logging.info("Encoding index dataset, this may take a while") + index_dataloader = model.setup_dataloader(cfg.index_ds, is_index_data=True) + embeddings, concept_ids = get_index_embeddings(cfg, index_dataloader, model) + + # Create pca model to reduce dimensionality of index dataset and decrease memory footprint + if cfg.apply_pca: + + # Need to train PCA model and apply PCA transformation with newly trained model + if not os.path.isfile(cfg.pca.pca_save_name): + logging.info("Fitting PCA model for embedding dimensionality reduction") + pca_train_set = random.sample(list(embeddings), k=int(len(embeddings) * cfg.pca.sample_fraction)) + pca = PCA(n_components=cfg.pca.output_dim) + pca.fit(pca_train_set) + pkl.dump(pca, open(cfg.pca.pca_save_name, "wb")) + embeddings = reduce_embedding_dim(pca, embeddings, cfg) + + # PCA model already trained, just need to reduce dimensionality of all embeddings + elif not os.path.isfile(cfg.pca_embeddings_save_name): + pca = pkl.load(open(cfg.pca.pca_save_name, "rb")) + embeddings = reduce_embedding_dim(pca, embeddings, cfg) + + # Build faiss index from embeddings + logging.info(f"Training index with embedding dim size {cfg.dims} using {faiss.get_num_gpus()} gpus") + quantizer = faiss.IndexFlatL2(cfg.dims) + index = faiss.IndexIVFFlat(quantizer, cfg.dims, cfg.nlist) + index = faiss.index_cpu_to_all_gpus(index) + index.train(embeddings) + + logging.info("Adding dataset embeddings to index") + for i in tqdm(range(0, embeddings.shape[0], cfg.index_batch_size)): + index.add(embeddings[i : i + cfg.index_batch_size]) + + logging.info("Saving index") + faiss.write_index(faiss.index_gpu_to_cpu(index), cfg.index_save_name) + logging.info("Index built and saved") + + +def reduce_embedding_dim(pca, embeddings, cfg): + """Apply PCA transformation to index dataset embeddings""" + + logging.info("Applying PCA transformation to entire index dataset") + embeddings = np.array(pca.transform(embeddings), dtype=np.float32) + emb_file = h5py.File(cfg.pca_embeddings_save_name, "w") + emb_file.create_dataset(cfg.index_ds.name, data=embeddings) + emb_file.close() + + return embeddings + + +def get_index_embeddings(cfg: DictConfig, dataloader: object, model: object): + """Use entity linking encoder to get embeddings for full index dataset""" + embeddings = [] + concept_ids = [] + + with torch.no_grad(): + for batch in tqdm(dataloader): + input_ids, token_type_ids, input_mask, batch_concept_ids = batch + input_ids = input_ids.to(device) + token_type_ids = token_type_ids.to(device) + input_mask = input_mask.to(device) + batch_embeddings = model.forward( + input_ids=input_ids, token_type_ids=token_type_ids, attention_mask=input_mask + ) + + embeddings.extend(batch_embeddings.detach().cpu().numpy()) + concept_ids.extend(batch_concept_ids.numpy()) + + emb_file = h5py.File(cfg.embedding_save_name, "w") + emb_file.create_dataset(cfg.index_ds.name, data=embeddings) + emb_file.close() + + pkl.dump(concept_ids, open(cfg.concept_id_save_name, "wb")) + + return embeddings, concept_ids + + +def load_model(cfg: DictConfig, restore: bool): + """ + Loads encoder model. + + Args: + cfg: Config file specifying model parameters + restore: Whether to restore model weights trained + by the user. Otherwise will load weights + used before self alignment pretraining. + """ + + if restore: + model = EntityLinkingModel.restore_from(cfg.nemo_path) + else: + cfg.train_ds = None + cfg.validation_ds = None + cfg.test_ds = None + model = EntityLinkingModel(cfg) + + model = model.to(device) + + return model + + +def main(cfg: DictConfig, restore: bool): + """ + Builds new index if one hasn't been built yet. + + Args: + cfg: Config file specifying index parameters + restore: Whether to restore model weights trained + by the user. Otherwise will load weights + used before self alignment pretraining. + """ + + logging.info("Loading entity linking encoder model") + model = load_model(cfg.model, restore) + + if not os.path.isfile(cfg.index.index_save_name) or ( + cfg.apply_pca and not os.path.isfile(cfg.index.pca.pca_save_name) + ): + logging.info("Building index") + build_index(cfg.index, model) + else: + logging.info("Index and pca model (if required) already exists. Skipping build index step.") + + if not os.path.isfile(cfg.index.idx_to_id): + logging.info("Mapping entity index postions to ids") + map_idx_to_ids(cfg.index) + else: + logging.info("Map from concept index to id already exists. Skipping mapping step.") + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument( + "--restore", action="store_true", help="Whether to restore encoder model weights from nemo path" + ) + parser.add_argument("--project_dir", required=False, type=str, default=".") + parser.add_argument("--cfg", required=False, type=str, default="./conf/umls_medical_entity_linking_config.yaml") + args = parser.parse_args() + + cfg = OmegaConf.load(args.cfg) + cfg.project_dir = args.project_dir + + main(cfg, args.restore) diff --git a/models/tts/NeMo/examples/nlp/entity_linking/conf/tiny_example_entity_linking_config.yaml b/models/tts/NeMo/examples/nlp/entity_linking/conf/tiny_example_entity_linking_config.yaml new file mode 100644 index 0000000000..36a0273060 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/entity_linking/conf/tiny_example_entity_linking_config.yaml @@ -0,0 +1,90 @@ +project_dir: null +name: SelfAlignmentPretrainingForMedicalEntityLinking +trainer: + gpus: 1 + num_nodes: 1 + max_epochs: 2 + max_steps: null + accumulate_grad_batches: 1 + precision: 16 + amp_level: O1 + accelerator: ddp + gradient_clip_val: 0.0 + log_every_n_steps: 1 + val_check_interval: 2 + checkpoint_callback: false + logger: false +model: + nemo_path: ??? + max_seq_length: 128 + language_model: + pretrained_model_name: bert-base-uncased + config_file: null + config: null + lm_checkpoint: null + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} + vocab_file: null + tokenizer_model: null + do_lower_case: true + loss_params: null + train_ds: + data_file: ??? + max_seq_length: ${model.max_seq_length} + batch_size: 8 + shuffle: true + num_workers: 2 + pin_memory: false + drop_last: false + validation_ds: + data_file: ??? + max_seq_length: ${model.max_seq_length} + batch_size: 8 + shuffle: false + num_workers: 2 + pin_memory: false + drop_last: false + optim: + name: adam + lr: 3.0e-05 + weight_decay: 0.0 + sched: + name: CosineAnnealing + warmup_steps: null + warmup_ratio: 0.1 + min_lr: 0.0 + last_epoch: -1 +index: + dims: 768 + nlist: 2 + top_n: 3 + query_num_factor: 20 + index_save_name: ??? + index_batch_size: 10 + index_ds: + name: tiny_example + data_file: ??? + max_seq_length: ${model.max_seq_length} + batch_size: 100 + shuffle: false + num_workers: 2 + pin_memory: false + drop_last: false + idx_to_id: ${project_dir}/idx_to_id.pkl + id_to_string: ${project_dir}/id_to_string.pkl + concept_id_save_name: ${project_dir}/tiny_example_concept_ids.pkl + embedding_save_name: ${project_dir}/tiny_example_concept_embeddings.hdf5 + pca_embeddings_save_name: null + apply_pca: false + pca: null +exp_manager: + exp_dir: . + name: ${project_dir}/SelfAlignmentPretrainingTinyExample + create_tensorboard_logger: true + create_checkpoint_callback: true +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/nlp/entity_linking/conf/umls_medical_entity_linking_config.yaml b/models/tts/NeMo/examples/nlp/entity_linking/conf/umls_medical_entity_linking_config.yaml new file mode 100644 index 0000000000..5b6709e722 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/entity_linking/conf/umls_medical_entity_linking_config.yaml @@ -0,0 +1,95 @@ +project_dir: ??? +name: SelfAlignmentPretrainingForMedicalEntityLinking +trainer: + gpus: 1 + num_nodes: 1 + max_epochs: 2 + max_steps: null + accumulate_grad_batches: 1 + precision: 16 + amp_level: O1 + accelerator: ddp + gradient_clip_val: 0.0 + log_every_n_steps: 1 + val_check_interval: 1000 + checkpoint_callback: false + logger: false +model: + nemo_path: ${project_dir}/sap_bert_umls.nemo + raw_data: ${project_dir}/data/MRCONSO.RRF + max_seq_length: 128 + language_model: + pretrained_model_name: bert-base-uncased + config_file: null + config: null + lm_checkpoint: null + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} + vocab_file: null + tokenizer_model: null + do_lower_case: true + train_ds: + data_file: ${project_dir}/data/umls_train_pairs.tsv + max_seq_length: ${model.max_seq_length} + batch_size: 128 + shuffle: true + num_workers: 2 + pin_memory: false + drop_last: false + validation_ds: + data_file: ${project_dir}/data/umls_validation_pairs.tsv + max_seq_length: ${model.max_seq_length} + batch_size: 128 + shuffle: false + num_workers: 2 + pin_memory: false + drop_last: false + optim: + name: adam + lr: 3.0e-05 + weight_decay: 0.0 + sched: + name: CosineAnnealing + warmup_steps: null + warmup_ratio: 0.1 + min_lr: 0.0 + last_epoch: -1 +index: + dims: 256 + nlist: 300 + top_n: 5 + query_num_factor: 20 + index_save_name: ${project_dir}/medical_entity_linking_index + index_batch_size: 1000 + raw_data: ${model.raw_data} + index_ds: + name: umls + data_file: ${project_dir}/data/umls_index_concepts.tsv + max_seq_length: ${model.max_seq_length} + batch_size: 128 + shuffle: false + num_workers: 2 + pin_memory: false + drop_last: false + idx_to_id: ${project_dir}/data/idx_to_id.pkl + id_to_string: ${project_dir}/data/id_to_string.pkl + concept_id_save_name: ${project_dir}/data/concept_ids.pkl + embedding_save_name: ${project_dir}/data/medical_concept_embeddings.hdf5 + pca_embeddings_save_name: ${project_dir}/data/medical_concept_reduced_${index.dims}dim_embeddings.hdf5 + apply_pca: true + pca: + input_dim: 756 + output_dim: ${index.dims} + sample_fraction: 0.5 + pca_save_name: ${project_dir}/${index.pca.input_dim}_to_${index.pca.output_dim}_pca_model.pkl +exp_manager: + exp_dir: ${project_dir}/medical_entity_linking_experiments + name: sap_bert_umls + create_tensorboard_logger: true + create_checkpoint_callback: true +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/nlp/entity_linking/data/umls_dataset_processing.py b/models/tts/NeMo/examples/nlp/entity_linking/data/umls_dataset_processing.py new file mode 100644 index 0000000000..03a17da3c0 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/entity_linking/data/umls_dataset_processing.py @@ -0,0 +1,189 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 itertools +import pickle as pkl +import random +from argparse import ArgumentParser + +import pandas as pd +from omegaconf import OmegaConf +from tqdm import tqdm + +# Info on these headers can be found here on the UMLS website https://www.ncbi.nlm.nih.gov/books/NBK9685/ +# section 3.3.4 Table 1 +HEADERS = [ + 'CUI', + 'LAT', + 'TS', + 'LUI', + 'STT', + 'SUI', + 'ISPREF', + 'AUI', + 'SAUI', + 'SCUI', + 'SDUI', + 'SAB', + 'TTY', + 'CODE', + 'STR', + 'SRL', + 'SUPPRESS', + 'CVF', +] + + +def process_umls_training_dataset(data_path, train_save_name, val_save_name, max_pairs, train_split, headers): + """ + Generates and saves UMLS self alignment pretraining train and validation data. Takes the raw .RRF UMLS + data file and creates different pair combinations for entities with the same CUI. Each row in the output + will be formatted as 'CUI EntitySynonym1 EntitySynonym2' with each item in a row separated by tabs. + Saves two .tsv output files, one for the train split and one for the validation split. + Only data marked as English is added to the train and val splits. + + Arguments: + data_path (str): path to MRCONSO.RRF UMLS data file + train_save_name (str): path to where training data will be saved + val_save_name (str): path to where validation data will be saved + max_pairs (int): max number of pairs for any one CUI added to the train + or validation splits + train_split (float): precentage of raw data to be added to train set split + headers (list): column lables within MRCONSO.RRF + """ + + print("Loading training data file...") + df = pd.read_table(data_path, names=headers, index_col=False, delimiter='|') + train_file = open(train_save_name, 'w') + val_file = open(val_save_name, 'w') + + cui = df["CUI"].iloc[0] + names = [] + random.seed(2021) + + for idx in tqdm(range(len(df))): + # Address incorrectly formatted data + if type(df["STR"].iloc[idx]) != str or "|" in df["STR"].iloc[idx]: + continue + + # Collect all english concept strings matching the current CUI + if df["CUI"].iloc[idx] == cui and df["LAT"].iloc[idx] == "ENG": + concept_string = df["STR"].iloc[idx] + names.append(concept_string) + + else: + # Pair off concept synonyms to make training and val sets + pairs = list(itertools.combinations(names, 2)) + + if len(pairs) == 0: + # Not enough concepts gathered to make a pair + cui = df["CUI"].iloc[idx] + names = [df["STR"].iloc[idx]] + continue + + # Removing leading C to convert label string to int + cui = int(cui[1:]) + random.shuffle(pairs) + + # Keep up to max pairs number pairs for any one concept + for pair in pairs[:max_pairs]: + + # Want concepts in train and val splits to be randomly selected and mutually exclusive + add_to_train = random.random() + + if add_to_train <= train_split: + train_file.write(f'{cui}\t{pair[0]}\t{pair[1]}\n') + else: + val_file.write(f'{cui}\t{pair[0]}\t{pair[1]}\n') + + # Switch to next concept + cui = df["CUI"].iloc[idx] + names = [df["STR"].iloc[idx]] + + train_file.close() + val_file.close() + print("Finished making training and validation data") + + +def process_umls_index_dataset(data_path, data_savename, id2string_savename, headers): + """ + Generates data file needed to build a UMLS index and a hash table mapping each + CUI to one canonical concept string. Takes the raw .RRF data file and saves + a .tsv indec concept file as well as the a .pkl file of cui to concept string + mappings. Only data marked as English is added to the index data file. + + Arguments: + data_path (str): path to MRCONSO.RRF UMLS data file + data_savename (str): path to where .tsv index data will be saved + id2string_savename (str): path to where .pkl cui to string mapping will + be saved + headers (list): column lables within MRCONSO.RRF + """ + + print("Loading index data file...") + df = pd.read_table(data_path, names=headers, index_col=False, delimiter='|') + id2string = {} + + with open(data_savename, "w") as outfile: + for idx, row in tqdm(df.iterrows(), total=df.shape[0]): + # Address incorrectly formatted data + if type(row["STR"]) != str or "|" in row["STR"]: + continue + + cui = row["CUI"] + sent = row["STR"] + + # Removing leading C to convert label string to int + cui = int(cui[1:]) + + # Only keeping english concepts + if row["LAT"] == "ENG": + outfile.write(f'{cui}\t{sent}\n') + + # Matching each cui to one canonical string represention + if cui not in id2string and ":" not in sent: + id2string[cui] = sent + + outfile.close() + pkl.dump(id2string, open(id2string_savename, "wb")) + print("Finished saving index data and id to concept mapping") + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument("--index", action="store_true", help="Whether to process data for building an index") + parser.add_argument("--project_dir", required=False, type=str, default=".") + parser.add_argument("--cfg", required=False, type=str, default="conf/umls_medical_entity_linking_config.yaml") + parser.add_argument( + "--max_pairs", required=False, type=int, default=50, help="Max number of train pairs for a single concepts" + ) + parser.add_argument( + "--train_split", required=False, type=float, default=0.99, help="Precentage of data to add to train set" + ) + + args = parser.parse_args() + cfg = OmegaConf.load(args.cfg) + cfg.project_dir = args.project_dir + + if args.index: + process_umls_index_dataset(cfg.index.raw_data, cfg.index.index_ds.data_file, cfg.index.id_to_string, HEADERS) + else: + process_umls_training_dataset( + cfg.model.raw_data, + cfg.model.train_ds.data_file, + cfg.model.validation_ds.data_file, + args.max_pairs, + args.train_split, + HEADERS, + ) diff --git a/models/tts/NeMo/examples/nlp/entity_linking/query_index.py b/models/tts/NeMo/examples/nlp/entity_linking/query_index.py new file mode 100644 index 0000000000..6cb51a7de1 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/entity_linking/query_index.py @@ -0,0 +1,166 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os +import pickle as pkl +from argparse import ArgumentParser +from collections import OrderedDict +from typing import Dict + +import numpy as np +import torch +from build_index import load_model +from omegaconf import DictConfig, OmegaConf + +from nemo.utils import logging + +try: + import faiss +except ModuleNotFoundError: + logging.warning("Faiss is required for building the index. Please install faiss-gpu") + +device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') + + +def get_query_embedding(query, model): + """Use entity linking encoder to get embedding for index query""" + model_input = model.tokenizer( + query, + add_special_tokens=True, + padding=True, + truncation=True, + max_length=512, + return_token_type_ids=True, + return_attention_mask=True, + ) + + query_emb = model.forward( + input_ids=torch.LongTensor([model_input["input_ids"]]).to(device), + token_type_ids=torch.LongTensor([model_input["token_type_ids"]]).to(device), + attention_mask=torch.LongTensor([model_input["attention_mask"]]).to(device), + ) + + return query_emb + + +def query_index( + query: str, cfg: DictConfig, model: object, index: object, pca: object, idx2id: dict, id2string: dict, +) -> Dict: + + """ + Query the nearest neighbor index of entities to find the + concepts in the index dataset that are most similar to the + query. + + Args: + query (str): entity to look up in the index + cfg (DictConfig): config object to specifiy query parameters + model (EntityLinkingModel): entity linking encoder model + index (object): faiss index + pca (object): sklearn pca transformation to be applied to queries + idx2id (dict): dictionary mapping unique concept dataset index to + its CUI + id2string (dict): dictionary mapping each unqiue CUI to a + representative english description of + the concept + Returns: + A dictionary with the concept ids of the index's most similar + entities as the keys and a tuple containing the string + representation of that concept and its cosine similarity to + the query as the values. + """ + query_emb = get_query_embedding(query, model).detach().cpu().numpy() + + if cfg.apply_pca: + query_emb = pca.transform(query_emb) + + dist, neighbors = index.search(query_emb.astype(np.float32), cfg.query_num_factor * cfg.top_n) + dist, neighbors = dist[0], neighbors[0] + unique_ids = OrderedDict() + neighbor_idx = 0 + + # Many of nearest neighbors could map to the same concept id, their idx is their unique identifier + while len(unique_ids) < cfg.top_n and neighbor_idx < len(neighbors): + concept_id_idx = neighbors[neighbor_idx] + concept_id = idx2id[concept_id_idx] + + # Only want one instance of each unique concept + if concept_id not in unique_ids: + concept = id2string[concept_id] + unique_ids[concept_id] = (concept, 1 - dist[neighbor_idx]) + + neighbor_idx += 1 + + unique_ids = dict(unique_ids) + + return unique_ids + + +def main(cfg: DictConfig, restore: bool): + """ + Loads faiss index and allows commandline queries + to the index. Builds new index if one hasn't been built yet. + + Args: + cfg: Config file specifying index parameters + restore: Whether to restore model weights trained + by the user. Otherwise will load weights + used before self alignment pretraining. + """ + + if not os.path.isfile(cfg.index.index_save_name) or ( + cfg.apply_pca and not os.path.isfile(cfg.index.pca.pca_save_name) or not os.path.isfile(cfg.index.idx_to_id) + ): + logging.warning("Either no index and/or no mapping from entity idx to ids exists. Please run `build_index.py`") + return + + logging.info("Loading entity linking encoder model") + model = load_model(cfg.model, restore) + + logging.info("Loading index and associated files") + index = faiss.read_index(cfg.index.index_save_name) + idx2id = pkl.load(open(cfg.index.idx_to_id, "rb")) + id2string = pkl.load(open(cfg.index.id_to_string, "rb")) # Should be created during dataset prep + + if cfg.index.apply_pca: + pca = pkl.load(open(cfg.index.pca.pca_save_name, "rb")) + + while True: + query = input("enter index query: ") + output = query_index(query, cfg.top_n, cfg.index, model, index, pca, idx2id, id2string) + + if query == "exit": + break + + for concept_id in output: + concept_details = output[concept_id] + concept_id = "C" + str(concept_id).zfill(7) + print(concept_id, concept_details) + + print("----------------\n") + + +if __name__ == '__main__': + parser = ArgumentParser() + parser.add_argument( + "--restore", action="store_true", help="Whether to restore encoder model weights from nemo path" + ) + parser.add_argument("--project_dir", required=False, type=str, default=".") + parser.add_argument("--cfg", required=False, type=str, default="./conf/umls_medical_entity_linking_config.yaml") + args = parser.parse_args() + + cfg = OmegaConf.load(args.cfg) + cfg.project_dir = args.project_dir + + main(cfg, args.restore) diff --git a/models/tts/NeMo/examples/nlp/entity_linking/self_alignment_pretraining.py b/models/tts/NeMo/examples/nlp/entity_linking/self_alignment_pretraining.py new file mode 100644 index 0000000000..0d0e0d3683 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/entity_linking/self_alignment_pretraining.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + + +# Please see tutorial at Nemo/tutorials/nlp/Entity_Linking_Medical.ipynb for +# more information on entity linking and self alignment pretraining. + +from omegaconf import DictConfig, OmegaConf +from pytorch_lightning import Trainer + +from nemo.collections.nlp.models import EntityLinkingModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="umls_medical_entity_linking_config.yaml") +def main(cfg: DictConfig) -> None: + logging.info(f"\nConfig Params:\n{OmegaConf.to_yaml(cfg)}") + trainer = Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + logging.info(f"Loading weights from pretrained model {cfg.model.language_model.pretrained_model_name}") + model = EntityLinkingModel(cfg=cfg.model, trainer=trainer) + logging.info("===========================================================================================") + logging.info('Starting training...') + trainer.fit(model) + logging.info('Training finished!') + logging.info("===========================================================================================") + + if cfg.model.nemo_path: + # '.nemo' file contains the last checkpoint and the params to initialize the model + model.save_to(cfg.model.nemo_path) + logging.info(f'Model is saved into `.nemo` file: {cfg.model.nemo_path}') + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark.py b/models/tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark.py new file mode 100644 index 0000000000..3242986749 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark.py @@ -0,0 +1,73 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +## Tasks +This script works with all GLUE Benchmark tasks, more details about the GLUE Benchmark could be found at +https://gluebenchmark.com/ + +More details on how to use this script could be found in tutorials/nlp/GLUE_Benchmark.ipynb + +## Model Training + +To train GLUEModel with the default config file, run: + python glue_benchmark.py \ + model.dataset.data_dir= \ + model.task_name=TASK_NAME \ + trainer.max_epochs= \ + trainer.gpus="[] + +Supported task names: +["cola", "sst-2", "mrpc", "sts-b", "qqp", "mnli", "qnli", "rte", "wnli"] +Note, MNLI task includes both matched and mismatched dev sets +""" + +import os + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import GLUEModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_name="glue_benchmark_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_manager_cfg = cfg.get("exp_manager", None) + + if exp_manager_cfg: + exp_manager_cfg.name = cfg.model.task_name + logging.info(f'Setting task_name to {exp_manager_cfg.name} in exp_manager') + exp_manager(trainer, exp_manager_cfg) + + if cfg.model.nemo_path and os.path.exists(cfg.model.nemo_path): + model = GLUEModel.restore_from(cfg.model.nemo_path) + logging.info(f'Restoring model from {cfg.model.nemo_path}') + model.update_data_dir(data_dir=cfg.model.dataset.data_dir) + model.setup_training_data() + model.setup_multiple_validation_data() + trainer.fit(model) + else: + model = GLUEModel(cfg.model, trainer=trainer) + trainer.fit(model) + if cfg.model.nemo_path: + model.save_to(cfg.model.nemo_path) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark_config.yaml b/models/tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark_config.yaml new file mode 100644 index 0000000000..23b5522ae3 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/glue_benchmark/glue_benchmark_config.yaml @@ -0,0 +1,82 @@ +# GLUE Benchmark with pre-trained BERT models +supported_tasks: &supported_tasks ['cola', 'sst-2', 'mrpc', 'sts-b', 'qqp', 'mnli', 'qnli', 'rte', 'wnli'] + +trainer: + gpus: 1 # the number of gpus, 0 for CPU + num_nodes: 1 + max_epochs: 3 + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + amp_level: O0 # O1/O2 for mixed precision + precision: 16 + accelerator: ddp + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + +model: + task_name: &task_name mrpc # choose from: ["cola", "sst-2", "mrpc", "sts-b", "qqp", "mnli", "qnli", "rte", "wnli"] GLUE task name, MNLI includes both matched and mismatched dev sets + supported_tasks: *supported_tasks + output_dir: null # dir to write write predictions + nemo_path: null # filename to save the model and associated artifacts to .nemo file + dataset: + data_dir: ??? # /path/to/data + max_seq_length: 128 + use_cache: true + + # shared across dataloaders: + num_workers: 2 + pin_memory: false + drop_last: false + + train_ds: + ds_item: 'train.tsv' + shuffle: true + num_samples: -1 + batch_size: 32 + + validation_ds: + ds_item: 'dev.tsv' # for MNLI 'dev_matched.tsv' and 'dev_mismatched.tsv' will de used + shuffle: false + num_samples: -1 + batch_size: 32 + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently. + + language_model: + pretrained_model_name: bert-base-uncased + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null + + optim: + name: adam + lr: 5e-5 + weight_decay: 0.00 + + sched: + name: WarmupAnnealing + # Scheduler params + warmup_steps: null + warmup_ratio: 0.1 + last_epoch: -1 + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + +exp_manager: + exp_dir: null # exp_dir for your experiment, if None, defaults to "./NeMo_experiments" + name: *task_name # The name of your model + create_tensorboard_logger: True # Whether you want exp_manger to create a tb logger + create_checkpoint_callback: True # Whether you want exp_manager to create a modelcheckpoint callback + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/nlp/information_retrieval/bert_dpr.py b/models/tts/NeMo/examples/nlp/information_retrieval/bert_dpr.py new file mode 100644 index 0000000000..2d9cd962ff --- /dev/null +++ b/models/tts/NeMo/examples/nlp/information_retrieval/bert_dpr.py @@ -0,0 +1,35 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import BertDPRModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="bert_ir_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + bert_dpr_model = BertDPRModel(cfg.model, trainer=trainer) + trainer.fit(bert_dpr_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/information_retrieval/bert_joint_ir.py b/models/tts/NeMo/examples/nlp/information_retrieval/bert_joint_ir.py new file mode 100644 index 0000000000..1bb164e580 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/information_retrieval/bert_joint_ir.py @@ -0,0 +1,35 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import BertJointIRModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="bert_ir_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + bert_joint_ir_model = BertJointIRModel(cfg.model, trainer=trainer) + trainer.fit(bert_joint_ir_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/information_retrieval/conf/bert_ir_config.yaml b/models/tts/NeMo/examples/nlp/information_retrieval/conf/bert_ir_config.yaml new file mode 100644 index 0000000000..e4fe72b349 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/information_retrieval/conf/bert_ir_config.yaml @@ -0,0 +1,99 @@ +# Fine-tuning BERT model for information retrieval +name: &name BertIR +trainer: + gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices + num_nodes: 1 + max_epochs: 2 # the number of training epochs + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + precision: 16 # 16 to use AMP + amp_level: O1 # O1 or O2 if using AMP + accelerator: ddp + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 0.05 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + checkpoint_callback: false # provided by exp_manager + logger: false # provided by exp_manager + +model: + nemo_path: null # exported .nemo path + + language_model: + pretrained_model_name: bert-base-uncased + sim_score_dropout: 0.1 + lm_checkpoint: null + config: + attention_probs_dropout_prob: 0.1 + hidden_act: gelu + hidden_dropout_prob: 0.1 + hidden_size: 768 + initializer_range: 0.02 + intermediate_size: 3072 + max_position_embeddings: 512 + num_attention_heads: 12 + num_hidden_layers: 12 + type_vocab_size: 2 + vocab_size: 30522 + config_file: null # json file, precedence over config + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec + vocab_file: null # path to vocab file + tokenizer_model: null # tokenizer model for sentencepiece + special_tokens: null + + train_ds: + passages: null # path to file with passages and their indices + queries: null # path to file with training queries and their indices + query_to_passages: null + # path to file with training examples which have the form of + # (query_id, relevant_passage_id, irrelevant_passage_1_id, ..., irrelevant_passage_n_id) + num_negatives: 10 + batch_size: 6 + psg_cache_format: npz + shuffle: true + num_samples: -1 # number of samples to be considered, -1 means all the dataset + num_workers: 1 + drop_last: false + pin_memory: false + + validation_ds: + passages: null # path to file with passages and their indices + queries: null # path to file with validation queries and their indices + query_to_passages: null # path to file with passages to re-rank for each validation query + num_negatives: 10 + batch_size: 6 + psg_cache_format: pkl + shuffle: false + num_samples: -1 # number of samples to be considered, -1 means all the dataset + num_workers: 1 + drop_last: false + pin_memory: false + + optim: + name: adam + lr: 1e-5 + betas: [0.9, 0.999] + weight_decay: 0 + + sched: + name: WarmupAnnealing + warmup_steps: null + warmup_ratio: 0.05 + last_epoch: -1 + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + +exp_manager: + exp_dir: null # where to store logs and checkpoints + name: *name # name of experiment + create_tensorboard_logger: True + create_checkpoint_callback: True + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/nlp/information_retrieval/construct_random_negatives.py b/models/tts/NeMo/examples/nlp/information_retrieval/construct_random_negatives.py new file mode 100644 index 0000000000..36da96f503 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/information_retrieval/construct_random_negatives.py @@ -0,0 +1,49 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse + +import numpy as np +import pandas as pd + + +def construct_negatives(input_file, output_file, num_passages, num_negatives): + qrels = pd.read_csv(input_file, delimiter="\t", header=None) + with open(output_file, "w") as f: + for i in range(len(qrels)): + query_id, rel_passage_id = qrels[0][i], qrels[2][i] + negatives = np.random.randint(num_passages, size=num_negatives) + output_ids = [query_id, rel_passage_id] + negatives.tolist() + output_str = [str(id_) for id_ in output_ids] + print("\t".join(output_str), file=f) + + +def main(): + parser = argparse.ArgumentParser(description="Negative passages construction") + parser.add_argument("--data", type=str, default="msmarco_dataset", help="path to folder with data") + parser.add_argument("--num_passages", type=int, default=8841823, help="total number of passages") + parser.add_argument("--num_negatives", type=int, default=10, help="number of negatives per positive") + args = parser.parse_args() + + for mode in ["train", "dev"]: + construct_negatives( + input_file=f"{args.data}/qrels.{mode}.tsv", + output_file=f"{args.data}/query2passages.{mode}.tsv", + num_passages=args.num_passages, + num_negatives=args.num_negatives, + ) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/information_retrieval/get_msmarco.sh b/models/tts/NeMo/examples/nlp/information_retrieval/get_msmarco.sh new file mode 100644 index 0000000000..0bf286c3f7 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/information_retrieval/get_msmarco.sh @@ -0,0 +1,22 @@ +echo "=== Acquiring MSMARCO dataset ===" +echo "---" + +mkdir -p msmarco_dataset +cd msmarco_dataset + +echo "- Downloading passages" +wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/collection.tar.gz +tar -xzvf collection.tar.gz +rm collection.tar.gz + +echo "- Downloading queries" +wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/queries.tar.gz +tar -xzvf queries.tar.gz +rm queries.tar.gz +rm queries.eval.tsv + +echo "- Downloading relevance labels" +wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.train.tsv +wget --quiet --continue https://msmarco.blob.core.windows.net/msmarcoranking/qrels.dev.tsv + +echo "---" \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml b/models/tts/NeMo/examples/nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml new file mode 100644 index 0000000000..502458c1d3 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml @@ -0,0 +1,111 @@ +# Intent and Slot classification with pretrained BERT models + +trainer: + gpus: 1 # the number of gpus, 0 for CPU + num_nodes: 1 + max_epochs: 50 + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 amp_level to enable the AMP. + accelerator: ddp + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + +model: + nemo_path: null # filename to save the model and associated artifacts to .nemo file + data_dir: ??? # /path/to/data + class_labels: + intent_labels_file: intent_labels.csv + slot_labels_file: slot_labels.csv + class_balancing: null # or weighted_loss + intent_loss_weight: 0.6 # relation of intent to slot loss in total loss (between 0 to 1) + pad_label: -1 # if -1 not slot token will be used + ignore_extra_tokens: false + ignore_start_end: true # do not use first and last token for slot training + + train_ds: + prefix: train + batch_size: 32 + shuffle: true + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + validation_ds: + prefix: test + batch_size: 32 + shuffle: false + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + test_ds: + prefix: test + batch_size: 32 + shuffle: false + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null + + language_model: + max_seq_length: 50 + pretrained_model_name: bert-base-uncased + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null + + head: + num_output_layers: 2 + fc_dropout: 0.1 + + optim: + name: adam + lr: 2e-5 + args: + name: auto + params: + weight_decay: 0.01 + + sched: + name: WarmupAnnealing + iters_per_batch: null # computed at runtime + max_steps: null # computed at runtime or explicitly set here + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + + # scheduler config override + args: + name: auto + params: + warmup_steps: null + warmup_ratio: 0.1 + last_epoch: -1 + +exp_manager: + exp_dir: null # exp_dir for your experiment, if None, defaults to "./nemo_experiments" + name: "IntentSlot" # The name of your model + create_tensorboard_logger: true # Whether you want exp_manger to create a tb logger + create_checkpoint_callback: true # Whether you want exp_manager to create a modelcheckpoint callback + +hydra: + run: + dir: . + job_logging: + root: + handlers: null \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/intent_slot_classification/data/assistant_utils.py b/models/tts/NeMo/examples/nlp/intent_slot_classification/data/assistant_utils.py new file mode 100644 index 0000000000..f47094e276 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/intent_slot_classification/data/assistant_utils.py @@ -0,0 +1,151 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import re +import shutil + +from nemo.collections.nlp.data.data_utils.data_preprocessing import DATABASE_EXISTS_TMP, if_exist, write_files +from nemo.utils import logging + + +def copy_input_files(infold): + """ Put training files in convenient place for conversion to our format. """ + our_infold = infold + "/dataset" + + if os.path.exists(our_infold + "/trainset") and os.path.exists(our_infold + "/testset"): + logging.info("Input folders exists") + return + + logging.info(f"Copying files to input folder: {our_infold}") + os.makedirs(infold, exist_ok=True) + + old_infold = ( + infold + '/CrossValidation/autoGeneFromRealAnno/autoGene_2018_03_22-13_01_25_169/CrossValidation/KFold_1' + ) + if not os.path.exists(our_infold + "/trainset"): + shutil.copytree(old_infold + '/trainset', our_infold + '/trainset') + + if not os.path.exists(our_infold + "/testset"): + shutil.copytree(old_infold + '/testset/csv', our_infold + '/testset') + + +def get_intents(infold): + """ Get list of intents from file names. """ + intents = [f[:-4] for f in os.listdir(infold)] + intents.sort() + print(f'Found {len(intents)} intents') + return intents + + +def get_intent_queries(infold, intent_names, mode): + """ Get list of queries with their corresponding intent number. """ + intent_queries = ['sentence\tlabel\n'] + + for index, intent in enumerate(intent_names): + queries = open(f'{infold}/{mode}set/{intent}.csv', 'r').readlines() + for query in queries[1:]: + phrases = query.split(";") + intent_query = phrases[4][1:-1] + "\t" + str(index) + intent_queries.append(intent_query) + + return intent_queries + + +def get_slots(infold, modes): + """ + Find a lost of unique slot types in training and testing data. + We use a single slot type name both for starting and continuation tokes (not using B-, I- notation). + """ + slots = set() + + for mode in modes: + path = f'{infold}/{mode}set' + for filename in os.listdir(path): + lines = open(f'{path}/{filename}', 'r').readlines() + for line in lines[1:]: + query = line.split(";")[3] + slot_phrases = re.findall('\[.*?\]', query) + for slot_phrase in slot_phrases: + slot = slot_phrase.split(" : ")[0][1:] + slots.add(slot) + + slots = sorted(slots) + slots.append("O") + print(f'Found {len(slots)} slot types') + return slots + + +def get_slot_queries(infold, slot_dict, mode, intent_names): + """ Convert each word in a query to corresponding slot number. """ + slot_queries = [] + outside_slot = len(slot_dict) - 1 + + # keep the same order of files/queries as for intents + for intent in intent_names: + lines = open(f'{infold}/{mode}set/{intent}.csv', 'r').readlines() + for line in lines[1:]: + slot_query = "" + query = line.split(";")[3] + words = query.split(" ") + current_slot = outside_slot + for word in words: + if word[0] == "[": + current_slot = slot_dict[word[1:]] + elif word[0] == ":": + continue + else: + slot_query += str(current_slot) + " " + if word[-1] == ']': + current_slot = outside_slot + + slot_queries.append(slot_query.strip()) + + return slot_queries + + +def process_assistant(infold, outfold, modes=['train', 'test']): + """ + https://github.com/xliuhw/NLU-Evaluation-Data - this dataset includes + about 25 thousand examples with 66 various multi-domain intents and 57 entity types. + """ + if if_exist(outfold, [f'{mode}_slots.tsv' for mode in modes]): + logging.info(DATABASE_EXISTS_TMP.format('robot', outfold)) + return outfold + + logging.info(f'Processing assistant commands dataset and store at {outfold}') + os.makedirs(outfold, exist_ok=True) + + # copy train/test files to the convenient directory to work with + copy_input_files(infold) + infold += "/dataset" + + # get list of intents from train folder (test folder supposed to be the same) + intent_names = get_intents(infold + "/trainset") + write_files(intent_names, f'{outfold}/dict.intents.csv') + + # get all train and test queries with their intent + for mode in modes: + intent_queries = get_intent_queries(infold, intent_names, mode) + write_files(intent_queries, f'{outfold}/{mode}.tsv') + + # get list of all unique slots in training and testing files + slot_types = get_slots(infold, modes) + write_files(slot_types, f'{outfold}/dict.slots.csv') + + # create files of slot queries + slot_dict = {k: v for v, k in enumerate(slot_types)} + for mode in modes: + slot_queries = get_slot_queries(infold, slot_dict, mode, intent_names) + write_files(slot_queries, f'{outfold}/{mode}_slots.tsv') diff --git a/models/tts/NeMo/examples/nlp/intent_slot_classification/data/import_datasets.py b/models/tts/NeMo/examples/nlp/intent_slot_classification/data/import_datasets.py new file mode 100755 index 0000000000..9c58ded287 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/intent_slot_classification/data/import_datasets.py @@ -0,0 +1,265 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os +import shutil +from os.path import exists + +from assistant_utils import process_assistant + +from nemo.collections.nlp.data.data_utils.data_preprocessing import ( + DATABASE_EXISTS_TMP, + MODE_EXISTS_TMP, + create_dataset, + get_dataset, + get_vocab, + if_exist, +) +from nemo.utils import logging + + +def ids2text(ids, vocab): + return ' '.join([vocab[int(id_)] for id_ in ids]) + + +def process_atis(infold, outfold, modes=['train', 'test'], do_lower_case=False): + """ MSFT's dataset, processed by Kaggle + https://www.kaggle.com/siddhadev/atis-dataset-from-ms-cntk + """ + vocab = get_vocab(f'{infold}/atis.dict.vocab.csv') + + if if_exist(outfold, [f'{mode}.tsv' for mode in modes]): + logging.info(DATABASE_EXISTS_TMP.format('ATIS', outfold)) + return outfold + logging.info(f'Processing ATIS dataset and storing at {outfold}.') + + os.makedirs(outfold, exist_ok=True) + + outfiles = {} + for mode in modes: + outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w') + outfiles[mode].write('sentence\tlabel\n') + outfiles[mode + '_slots'] = open(f'{outfold}/{mode}_slots.tsv', 'w') + + queries = open(f'{infold}/atis.{mode}.query.csv', 'r').readlines() + intents = open(f'{infold}/atis.{mode}.intent.csv', 'r').readlines() + slots = open(f'{infold}/atis.{mode}.slots.csv', 'r').readlines() + + for i, query in enumerate(queries): + sentence = ids2text(query.strip().split()[1:-1], vocab) + if do_lower_case: + sentence = sentence.lower() + outfiles[mode].write(f'{sentence}\t{intents[i].strip()}\n') + slot = ' '.join(slots[i].strip().split()[1:-1]) + outfiles[mode + '_slots'].write(slot + '\n') + + shutil.copyfile(f'{infold}/atis.dict.intent.csv', f'{outfold}/dict.intents.csv') + shutil.copyfile(f'{infold}/atis.dict.slots.csv', f'{outfold}/dict.slots.csv') + for mode in modes: + outfiles[mode].close() + + +def process_snips(infold, outfold, do_lower_case, modes=['train', 'test'], dev_split=0.1): + if not os.path.exists(infold): + link = 'https://github.com/snipsco/spoken-language-understanding-research-datasets' + raise ValueError(f'Data not found at {infold}. ' f'You may request to download the SNIPS dataset from {link}.') + + exist = True + for dataset in ['light', 'speak', 'all']: + if if_exist(f'{outfold}/{dataset}', [f'{mode}.tsv' for mode in modes]): + logging.info(DATABASE_EXISTS_TMP.format('SNIPS-' + dataset, outfold)) + else: + exist = False + if exist: + return outfold + + logging.info(f'Processing SNIPS dataset and storing at folders "speak", "light" and "all" under {outfold}.') + logging.info( + f'Processing and importing "smart-speaker-en-close-field" -> "speak" and "smart-speaker-en-close-field" -> "light".' + ) + + os.makedirs(outfold, exist_ok=True) + + speak_dir = 'smart-speaker-en-close-field' + light_dir = 'smart-lights-en-close-field' + + light_files = [f'{infold}/{light_dir}/dataset.json'] + speak_files = [f'{infold}/{speak_dir}/training_dataset.json'] + speak_files.append(f'{infold}/{speak_dir}/test_dataset.json') + + light_train, light_dev, light_slots, light_intents = get_dataset(light_files, dev_split) + speak_train, speak_dev, speak_slots, speak_intents = get_dataset(speak_files) + + create_dataset(light_train, light_dev, light_slots, light_intents, do_lower_case, f'{outfold}/light') + create_dataset(speak_train, speak_dev, speak_slots, speak_intents, do_lower_case, f'{outfold}/speak') + create_dataset( + light_train + speak_train, + light_dev + speak_dev, + light_slots | speak_slots, + light_intents | speak_intents, + do_lower_case, + f'{outfold}/all', + ) + + +def process_jarvis_datasets( + infold, outfold, modes=['train', 'test', 'dev'], do_lower_case=False, ignore_prev_intent=False +): + """ process and convert Jarvis datasets into NeMo's BIO format + """ + dataset_name = "jarvis" + if if_exist(outfold, ['dict.intents.csv', 'dict.slots.csv']): + logging.info(DATABASE_EXISTS_TMP.format(dataset_name, outfold)) + return outfold + + logging.info(f'Processing {dataset_name} dataset and storing at {outfold}') + + os.makedirs(outfold, exist_ok=True) + + outfiles = {} + intents_list = {} + slots_list = {} + slots_list_all = {} + + outfiles['dict_intents'] = open(f'{outfold}/dict.intents.csv', 'w') + outfiles['dict_slots'] = open(f'{outfold}/dict.slots.csv', 'w') + + outfiles['dict_slots'].write('O\n') + slots_list["O"] = 0 + slots_list_all["O"] = 0 + + for mode in modes: + if if_exist(outfold, [f'{mode}.tsv']): + logging.info(MODE_EXISTS_TMP.format(mode, dataset_name, outfold, mode)) + continue + + if not if_exist(infold, [f'{mode}.tsv']): + logging.info(f'{mode} mode of {dataset_name}' f' is skipped as it was not found.') + continue + + outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'w') + outfiles[mode].write('sentence\tlabel\n') + outfiles[mode + '_slots'] = open(f'{outfold}/{mode}_slots.tsv', 'w') + + queries = open(f'{infold}/{mode}.tsv', 'r').readlines() + + for i, query in enumerate(queries): + line_splits = query.strip().split("\t") + if len(line_splits) == 3: + intent_str, slot_tags_str, sentence = line_splits + else: + intent_str, sentence = line_splits + slot_tags_str = "" + + if intent_str not in intents_list: + intents_list[intent_str] = len(intents_list) + outfiles['dict_intents'].write(f'{intent_str}\n') + + if ignore_prev_intent: + start_token = 2 + else: + start_token = 1 + + if do_lower_case: + sentence = sentence.lower() + sentence_cld = " ".join(sentence.strip().split()[start_token:-1]) + outfiles[mode].write(f'{sentence_cld}\t' f'{str(intents_list[intent_str])}\n') + + slot_tags_list = [] + if slot_tags_str.strip(): + slot_tags = slot_tags_str.strip().split(",") + for st in slot_tags: + if not st.strip(): + continue + [start_i, end_i, slot_name] = st.strip().split(":") + slot_tags_list.append([int(start_i), int(end_i), slot_name]) + if slot_name not in slots_list: + slots_list[slot_name] = len(slots_list) + slots_list_all[f'B-{slot_name}'] = len(slots_list_all) + slots_list_all[f'I-{slot_name}'] = len(slots_list_all) + outfiles['dict_slots'].write(f'B-{slot_name}\n') + outfiles['dict_slots'].write(f'I-{slot_name}\n') + + slot_tags_list.sort(key=lambda x: x[0]) + slots = [] + processed_index = 0 + for tag_start, tag_end, tag_str in slot_tags_list: + if tag_start > processed_index: + words_list = sentence[processed_index:tag_start].strip().split() + slots.extend([str(slots_list_all['O'])] * len(words_list)) + words_list = sentence[tag_start:tag_end].strip().split() + slots.append(str(slots_list_all[f'B-{tag_str}'])) + slots.extend([str(slots_list_all[f'I-{tag_str}'])] * (len(words_list) - 1)) + processed_index = tag_end + + if processed_index < len(sentence): + words_list = sentence[processed_index:].strip().split() + slots.extend([str(slots_list_all['O'])] * len(words_list)) + + slots = slots[1:-1] + slot = ' '.join(slots) + outfiles[mode + '_slots'].write(slot + '\n') + + outfiles[mode + '_slots'].close() + outfiles[mode].close() + + outfiles['dict_slots'].close() + outfiles['dict_intents'].close() + + return outfold + + +if __name__ == "__main__": + # Parse the command-line arguments. + parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.") + parser.add_argument( + "--dataset_name", required=True, type=str, choices=['atis', 'snips', 'jarvis', 'assistant'], + ) + parser.add_argument( + "--source_data_dir", required=True, type=str, help='path to the folder containing the dataset files' + ) + parser.add_argument("--target_data_dir", required=True, type=str, help='path to save the processed dataset') + parser.add_argument("--do_lower_case", action='store_true') + parser.add_argument( + "--ignore_prev_intent", + action='store_true', + help='ignores previous intent while importing datasets in jarvis\'s format', + ) + + args = parser.parse_args() + + dataset_name = args.dataset_name + source_dir = args.source_data_dir + target_dir = args.target_data_dir + + if not exists(source_dir): + raise FileNotFoundError(f"{source_dir} does not exist.") + + if dataset_name == 'atis': + process_atis(infold=source_dir, outfold=target_dir, do_lower_case=args.do_lower_case) + elif dataset_name == 'snips': + process_snips(infold=source_dir, outfold=target_dir, do_lower_case=args.do_lower_case) + elif dataset_name == 'jarvis': + process_jarvis_datasets( + infold=source_dir, + outfold=target_dir, + modes=["train", "test", "dev"], + do_lower_case=args.do_lower_case, + ignore_prev_intent=args.ignore_prev_intent, + ) + elif dataset_name == 'assistant': + process_assistant(infold=source_dir, outfold=target_dir) + else: + raise ValueError(f'Dataset {dataset_name} is not supported.') diff --git a/models/tts/NeMo/examples/nlp/intent_slot_classification/intent_slot_classification.py b/models/tts/NeMo/examples/nlp/intent_slot_classification/intent_slot_classification.py new file mode 100644 index 0000000000..39b4b7c5dd --- /dev/null +++ b/models/tts/NeMo/examples/nlp/intent_slot_classification/intent_slot_classification.py @@ -0,0 +1,85 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import IntentSlotClassificationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="intent_slot_classification_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config Params:\n {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + # initialize the model using the config file + model = IntentSlotClassificationModel(cfg.model, trainer=trainer) + + # training + logging.info("================================================================================================") + logging.info('Starting training...') + trainer.fit(model) + logging.info('Training finished!') + + # Stop further testing as fast_dev_run does not save checkpoints + if trainer.fast_dev_run: + return + + # after model training is done, you can load the model from the saved checkpoint + # and evaluate it on a data file or on given queries. + logging.info("================================================================================================") + logging.info("Starting the testing of the trained model on test set...") + logging.info("We will load the latest model saved checkpoint from the training...") + + # for evaluation and inference you can load the previously trained model saved in .nemo file + # like this in your code, but we will just reuse the trained model here + # eval_model = IntentSlotClassificationModel.restore_from(restore_path=checkpoint_path) + eval_model = model + + # we will setup testing data reusing the same config (test section) + eval_model.update_data_dir_for_testing(data_dir=cfg.model.data_dir) + eval_model.setup_test_data(test_data_config=cfg.model.test_ds) + + trainer.test(model=eval_model, ckpt_path=None, verbose=False) + logging.info("Testing finished!") + + # run an inference on a few examples + logging.info("======================================================================================") + logging.info("Evaluate the model on the given queries...") + + # this will work well if you train the model on Assistant dataset + # for your own dataset change the examples appropriately + queries = [ + 'set alarm for seven thirty am', + 'lower volume by fifty percent', + 'what is my schedule for tomorrow', + ] + + pred_intents, pred_slots = eval_model.predict_from_examples(queries, cfg.model.test_ds) + + logging.info('The prediction results of some sample queries with the trained model:') + for query, intent, slots in zip(queries, pred_intents, pred_slots): + logging.info(f'Query : {query}') + logging.info(f'Predicted Intent: {intent}') + logging.info(f'Predicted Slots: {slots}') + + logging.info("Inference finished!") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/language_modeling/bert_pretraining.py b/models/tts/NeMo/examples/nlp/language_modeling/bert_pretraining.py new file mode 100644 index 0000000000..7760b73709 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/bert_pretraining.py @@ -0,0 +1,38 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf +from pytorch_lightning.plugins import DDPPlugin + +from nemo.collections.nlp.models.language_modeling import BERTLMModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="bert_pretraining_from_text_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config:\n {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(plugins=[DDPPlugin(find_unused_parameters=True)], **cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + bert_model = BERTLMModel(cfg.model, trainer=trainer) + trainer.fit(bert_model) + if cfg.model.nemo_path: + bert_model.save_to(cfg.model.nemo_path) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml b/models/tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml new file mode 100644 index 0000000000..cc516c1e85 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_preprocessed_config.yaml @@ -0,0 +1,80 @@ +# BERT Pretraining from Preprocessed (tokenized) data +name: &name PretrainingBERTFromPreprocessed +trainer: + gpus: 8 # the number of gpus, 0 for CPU, or list with gpu indices + num_nodes: 1 + max_steps: 2285714 # precedence over max_epochs + num_sanity_val_steps: 0 # needed for bert pretraining from preproc + replace_sampler_ddp: false # needed for bert pretraining from preproc + accumulate_grad_batches: 1 # accumulates grads every k batches + precision: 16 # 16 to use AMP + amp_level: O1 # O1 or O2 if using AMP + accelerator: ddp + gradient_clip_val: 1.0 + log_every_n_steps: 1 + val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch + checkpoint_callback: false # provided by exp_manager + logger: false # provided by exp_manager + +model: + nemo_path: null # exported .nemo path + only_mlm_loss: true # only use masked language model without next sentence prediction + num_tok_classification_layers: 1 # number of token classification head output layers + num_seq_classification_layers: 2 # number of sequence classification head output layers + + + language_model: + pretrained_model_name: bert-base-uncased # huggingface model name + lm_checkpoint: null + config: + attention_probs_dropout_prob: 0.1 + hidden_act: gelu + hidden_dropout_prob: 0.1 + hidden_size: 768 + initializer_range: 0.02 + intermediate_size: 3072 + max_position_embeddings: 512 + num_attention_heads: 12 + num_hidden_layers: 12 + type_vocab_size: 2 + vocab_size: 30522 + config_file: null # json file, precedence over config + + tokenizer: null + + train_ds: + data_file: null # path to hdf5 file (or directory) + max_predictions_per_seq: 80 + batch_size: 16 + shuffle: true + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + optim: + name: adamw + lr: 0.4375e-4 + weight_decay: 0.01 + + sched: + name: SquareRootAnnealing + warmup_steps: null + warmup_ratio: 0.01 + min_lr: 0.0 + last_epoch: -1 + + +exp_manager: + exp_dir: null # where to store logs and checkpoints + name: *name # name of experiment + create_tensorboard_logger: True + create_checkpoint_callback: True + + +hydra: + run: + dir: . + job_logging: + root: + handlers: null \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml b/models/tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml new file mode 100644 index 0000000000..02df49edf5 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/conf/bert_pretraining_from_text_config.yaml @@ -0,0 +1,109 @@ +# BERT Pretraining from Text +name: &name PretrainingBERTFromText +trainer: + gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices + num_nodes: 1 + max_epochs: 2 # the number of training epochs + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + precision: 16 # 16 to use AMP + amp_level: O1 # O1 or O2 if using AMP + accelerator: ddp + gradient_clip_val: 0.0 + log_every_n_steps: 1 + val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch + checkpoint_callback: false # provided by exp_manager + logger: false # provided by exp_manager + +model: + nemo_path: null # exported .nemo path + only_mlm_loss: false # only use masked language model without next sentence prediction + num_tok_classification_layers: 1 # number of token classification head output layers + num_seq_classification_layers: 2 # number of sequence classification head output layers + max_seq_length: 128 + # The maximum total input sequence length after tokenization. Sequences longer than this + # will be truncated, and sequences shorter than this will be padded. + mask_prob: 0.15 + # Probability of masking a token in the input text during data processing. + short_seq_prob: 0.1 + # Probability of having a sequence shorter than the maximum sequence length `max_seq_length` in data processing.", + + language_model: + pretrained_model_name: bert-base-uncased + lm_checkpoint: null + config: + attention_probs_dropout_prob: 0.1 + hidden_act: gelu + hidden_dropout_prob: 0.1 + hidden_size: 768 + initializer_range: 0.02 + intermediate_size: 3072 + max_position_embeddings: 512 + num_attention_heads: 12 + num_hidden_layers: 12 + type_vocab_size: 2 + vocab_size: 30522 + config_file: null # json file, precedence over config + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec + vocab_file: null # path to vocab file + tokenizer_model: null # tokenizer model for sentencepiece + special_tokens: # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently. + unk_token: '[UNK]' + sep_token: '[SEP]' + pad_token: '[PAD]' + bos_token: '[CLS]' + mask_token: '[MASK]' + eos_token: '[SEP]' + cls_token: '[CLS]' + + train_ds: + data_file: null # path to data file + max_seq_length: ${model.max_seq_length} + mask_prob: ${model.mask_prob} + short_seq_prob: ${model.short_seq_prob} + batch_size: 16 # per GPU + shuffle: true + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + validation_ds: + data_file: null # path to data file + max_seq_length: ${model.max_seq_length} + mask_prob: ${model.mask_prob} + short_seq_prob: ${model.short_seq_prob} + batch_size: 16 # per GPU + shuffle: false + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + optim: + name: adamw + lr: 3e-5 + weight_decay: 0.0 + + sched: + name: CosineAnnealing + warmup_steps: null + warmup_ratio: 0.1 + min_lr: 0.0 + last_epoch: -1 + + +exp_manager: + exp_dir: null # where to store logs and checkpoints + name: *name # name of experiment + create_tensorboard_logger: True + create_checkpoint_callback: True + +hydra: + run: + dir: . + job_logging: + root: + handlers: null \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/language_modeling/conf/transformer_lm_config.yaml b/models/tts/NeMo/examples/nlp/language_modeling/conf/transformer_lm_config.yaml new file mode 100644 index 0000000000..cb968bdc06 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/conf/transformer_lm_config.yaml @@ -0,0 +1,96 @@ +name: TransformerLanguageModel +do_training: True # set to False if only preprocessing data + +model: + label_smoothing: 0.0 + preproc_out_dir: null # path to store data preprocessing outputs + + train_ds: + use_tarred_dataset: true + tar_files: ??? # path to tarred files + metadata_file: ??? # metadata for tarred dataset + shard_strategy: scatter + tokens_in_batch: 4096 + tar_shuffle_n: 256 + + validation_ds: + file_name: ??? # path to file with validation data + tokens_in_batch: 512 + clean: false + shuffle: false + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + test_ds: + file_name: ??? # path to file with test data + tokens_in_batch: 512 + clean: false + shuffle: false + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + optim: + name: adam + lr: 0.001 + betas: + - 0.9 + - 0.98 + weight_decay: 0.0 + sched: + name: InverseSquareRootAnnealing + min_lr: 0.0 + last_epoch: -1 + warmup_ratio: 0.1 + + tokenizer: + tokenizer_name: yttm + tokenizer_model: ??? + vocab_file: null + special_tokens: null + training_sample_size: null # valid for sentencepiece tokenizer + + encoder: + library: nemo + model_name: null + pretrained: false + max_sequence_length: 512 + num_token_types: 0 + embedding_dropout: 0.1 + learn_positional_encodings: false + hidden_size: 512 + num_layers: 6 + inner_size: 2048 + num_attention_heads: 8 + ffn_dropout: 0.1 + attn_score_dropout: 0.1 + attn_layer_dropout: 0.1 + hidden_act: relu + mask_future: true + pre_ln: false + + head: + num_layers: 1 + activation: relu + log_softmax: true + dropout: 0.0 + use_transformer_init: true + +trainer: + gpus: 4 + num_nodes: 1 + max_epochs: 200 + amp_level: O2 # O1/O2 for mixed precision + precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0 + accelerator: ddp + checkpoint_callback: False + logger: False + log_every_n_steps: 50 # Interval of logging. + check_val_every_n_epoch: 1 + +exp_manager: + name: TransformerLM + files_to_copy: [] \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/language_modeling/convert_weights_to_nemo1.0.py b/models/tts/NeMo/examples/nlp/language_modeling/convert_weights_to_nemo1.0.py new file mode 100644 index 0000000000..b58b33942e --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/convert_weights_to_nemo1.0.py @@ -0,0 +1,61 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Converts BERT NeMo0.* checkpoints to NeMo1.0 format. +""" + +from argparse import ArgumentParser + +import torch + +parser = ArgumentParser() +parser.add_argument("--bert_encoder", required=True, help="path to BERT encoder, e.g. /../BERT-STEP-2285714.pt") +parser.add_argument( + "--bert_token_classifier", + required=True, + help="path to BERT token classifier, e.g. /../BertTokenClassifier-STEP-2285714.pt", +) +parser.add_argument( + "--bert_sequence_classifier", + required=False, + default=None, + help="path to BERT sequence classifier, e.g /../SequenceClassifier-STEP-2285714.pt", +) +parser.add_argument( + "--output_path", required=False, default="converted_model.pt", help="output path to newly converted model" +) +args = parser.parse_args() + +bert_in = torch.load(args.bert_encoder) +tok_in = torch.load(args.bert_token_classifier) +if args.bert_sequence_classifier: + seq_in = torch.load(args.bert_sequence_classifier) + +new_dict = {} +new_model = {"state_dict": new_dict} +for k in bert_in: + new_name = k.replace("bert.", "bert_model.") + new_dict[new_name] = bert_in[k] + +for k in tok_in: + new_name = "mlm_classifier." + k + new_dict[new_name] = tok_in[k] + +if args.bert_sequence_classifier: + for k in seq_in: + new_name = "nsp_classifier." + k + new_dict[new_name] = seq_in[k] + +torch.save(new_model, args.output_path) diff --git a/models/tts/NeMo/examples/nlp/language_modeling/get_wkt2.sh b/models/tts/NeMo/examples/nlp/language_modeling/get_wkt2.sh new file mode 100644 index 0000000000..33b55b2a73 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/get_wkt2.sh @@ -0,0 +1,23 @@ +#!/bin/bash + +""" +This file is adapted from +https://github.com/salesforce/awd-lstm-lm/blob/master/getdata.sh +Copyright by the AWD LSTM authors. +""" +DATA_DIR=$1 +echo "- Downloading WikiText-2" + +wget --continue -P $DATA_DIR https://s3.amazonaws.com/research.metamind.io/wikitext/wikitext-2-v1.zip +unzip -q $DATA_DIR/wikitext-2-v1.zip -d $DATA_DIR +cd $DATA_DIR/wikitext-2 +mv wiki.train.tokens train.txt +sed -i -e "s//[UNK]/g" train.txt +mv wiki.valid.tokens valid.txt +sed -i -e "s//[UNK]/g" valid.txt +mv wiki.test.tokens test.txt +sed -i -e "s//[UNK]/g" test.txt +cd .. +rm wikitext-2-v1.zip + +echo "- WikiText-2 saved at $DATA_DIR/wikitext-2" diff --git a/models/tts/NeMo/examples/nlp/language_modeling/transformer_lm.py b/models/tts/NeMo/examples/nlp/language_modeling/transformer_lm.py new file mode 100644 index 0000000000..caaa0e0d29 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/language_modeling/transformer_lm.py @@ -0,0 +1,35 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.language_modeling import TransformerLMModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="transformer_lm_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + transformer_lm = TransformerLMModel(cfg.model, trainer=trainer) + trainer.fit(transformer_lm) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml b/models/tts/NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml new file mode 100644 index 0000000000..a932f63152 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/conf/aayn_base.yaml @@ -0,0 +1,156 @@ +name: AttentionIsAllYouNeed +do_training: True # set to False if only preprocessing data +do_testing: False # set to True to run evaluation on test data after training + +model: + beam_size: 4 + len_pen: 0.6 + multilingual: False + max_generation_delta: 5 + label_smoothing: 0.1 + shared_tokenizer: True # train tokenizer model across src and tgt train data + preproc_out_dir: null # path to store data preprocessing outputs + src_language: 'en' + tgt_language: 'de' + + train_ds: + src_file_name: null + tgt_file_name: null + use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically) + # config for preprocessing training data and creating a tarred datset automatically + tar_file_prefix: parallel # prefix for tar file names + tar_files: null # if data has already been preprocessed (rest of config ignored) + metadata_file: null # metadata for tarred dataset + lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding + num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile + tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled + shard_strategy: scatter # tarred dataset shard distribution strategy + n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2) + tokens_in_batch: 512 + clean: true + max_seq_length: 512 + shuffle: true + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + concat_sampling_technique: temperature # only used with ConcatTranslationDataset + concat_sampling_temperature: 5 # only used with ConcatTranslationDataset + concat_sampling_probabilities: null # only used with ConcatTranslationDataset + + validation_ds: + src_file_name: ??? + tgt_file_name: ??? + tokens_in_batch: 512 + clean: false + max_seq_length: 512 + shuffle: false + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + test_ds: + src_file_name: ??? + tgt_file_name: ??? + tokens_in_batch: 512 + clean: false + max_seq_length: 512 + shuffle: false + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + optim: + name: adam + lr: 0.001 + betas: + - 0.9 + - 0.98 + weight_decay: 0.0 + sched: + name: InverseSquareRootAnnealing + min_lr: 0.0 + last_epoch: -1 + warmup_ratio: 0.1 + + encoder_tokenizer: + library: yttm + tokenizer_model: null + vocab_size: null # vocab size for training bpe + bpe_dropout: null + vocab_file: null + special_tokens: null + training_sample_size: null # valid for sentencepiece tokenizer + + decoder_tokenizer: + library: yttm + tokenizer_model: null + vocab_size: null # vocab size for training bpe + bpe_dropout: null + vocab_file: null + special_tokens: null + training_sample_size: null # valid for sentencepiece tokenizer + + encoder: + library: nemo + model_name: null + pretrained: false + max_sequence_length: 512 + num_token_types: 0 + embedding_dropout: 0.1 + learn_positional_encodings: false + hidden_size: 512 + num_layers: 6 + inner_size: 2048 + num_attention_heads: 8 + ffn_dropout: 0.1 + attn_score_dropout: 0.1 + attn_layer_dropout: 0.1 + hidden_act: relu + mask_future: false + pre_ln: false + pre_ln_final_layer_norm: true + + decoder: + library: nemo + model_name: null + pretrained: false + max_sequence_length: 512 + num_token_types: 0 + embedding_dropout: 0.1 + learn_positional_encodings: false + hidden_size: 512 + inner_size: 2048 + num_layers: 6 + num_attention_heads: 8 + ffn_dropout: 0.1 + attn_score_dropout: 0.1 + attn_layer_dropout: 0.1 + hidden_act: relu + pre_ln: false + pre_ln_final_layer_norm: true + + head: + num_layers: 1 + activation: relu + log_softmax: true + dropout: 0.0 + use_transformer_init: true + +trainer: + gpus: 4 + num_nodes: 1 + max_epochs: 200 + amp_level: O2 # O1/O2 for mixed precision + precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0 + accelerator: ddp + checkpoint_callback: False + logger: False + log_every_n_steps: 50 # Interval of logging. + check_val_every_n_epoch: 1 + +exp_manager: + name: AAYNBase + files_to_copy: [] \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/machine_translation/conf/huggingface.yaml b/models/tts/NeMo/examples/nlp/machine_translation/conf/huggingface.yaml new file mode 100644 index 0000000000..abb5e9a844 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/conf/huggingface.yaml @@ -0,0 +1,132 @@ +name: HuggingFaceEncoder +do_training: True # set to False if only preprocessing data +do_testing: False # set to True to run evaluation on test data after training + +model: + beam_size: 4 + len_pen: 0.6 + max_generation_delta: 5 + label_smoothing: 0.1 + shared_tokenizer: false + preproc_out_dir: null + src_language: 'en' + tgt_language: 'de' + + train_ds: + src_file_name: null + tgt_file_name: null + use_tarred_dataset: False # if true tar_file_name and meta_file_name will be used (or created automatically) + # config for preprocessing training data and creating a tarred datset automatically + tar_file_prefix: parallel # prefix for tar file names + tar_files: null # if data has already been preprocessed (rest of config ignored) + metadata_file: null # metadata for tarred dataset + lines_per_dataset_fragment: 1000000 # Number of lines to consider for bucketing and padding + num_batches_per_tarfile: 100 # Number of batches (pickle files) within each tarfile + tar_shuffle_n: 100 # How many samples to look ahead and load to be shuffled + shard_strategy: scatter # tarred dataset shard distribution strategy + n_preproc_jobs: -2 # number of processes to use for data preprocessing (-2 means all but 2) + tokens_in_batch: 512 + clean: true + max_seq_length: 512 + shuffle: true + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + validation_ds: + src_file_name: null + tgt_file_name: null + tokens_in_batch: 512 + clean: false + max_seq_length: 512 + shuffle: false + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + test_ds: + src_file_name: null + tgt_file_name: null + tokens_in_batch: 512 + clean: false + max_seq_length: 512 + shuffle: false + num_samples: -1 + drop_last: false + pin_memory: false + num_workers: 8 + + optim: + name: adam + lr: 0.001 + betas: + - 0.9 + - 0.98 + weight_decay: 0.0 + sched: + name: InverseSquareRootAnnealing + min_lr: 0.0 + last_epoch: -1 + warmup_ratio: 0.1 + + encoder_tokenizer: + library: huggingface + tokenizer_model: null + vocab_file: null + special_tokens: null + vocab_size: null + + decoder_tokenizer: + library: yttm + tokenizer_model: null + vocab_file: null + special_tokens: null + vocab_size: null + + encoder: + library: huggingface + model_name: bert-base-uncased + pretrained: false + + decoder: + library: nemo + model_name: null + pretrained: false + max_sequence_length: 512 + num_token_types: 2 + embedding_dropout: 0.1 + learn_positional_encodings: false + hidden_size: 512 + inner_size: 2048 + num_layers: 6 + num_attention_heads: 8 + ffn_dropout: 0.1 + attn_score_dropout: 0.1 + attn_layer_dropout: 0.1 + hidden_act: relu + pre_ln: false + + head: + num_layers: 1 + activation: relu + log_softmax: true + dropout: 0.0 + use_transformer_init: true + +trainer: + gpus: 4 + num_nodes: 1 + max_epochs: 200 + amp_level: O2 # O1/O2 for mixed precision + precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0 + accelerator: ddp + checkpoint_callback: False + logger: False + log_every_n_steps: 50 # Interval of logging. + check_val_every_n_epoch: 1 + +exp_manager: + name: HuggingFaceEncoder + files_to_copy: [] \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/machine_translation/create_tarred_monolingual_dataset.py b/models/tts/NeMo/examples/nlp/machine_translation/create_tarred_monolingual_dataset.py new file mode 100644 index 0000000000..5ad5f789d7 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/create_tarred_monolingual_dataset.py @@ -0,0 +1,68 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os + +from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='NMT dataset pre-processing') + parser.add_argument('--tokenizer_model', type=str, required=True, help='Path to tokenizer model') + parser.add_argument('--tokenizer_name', type=str, default='yttm', help='BPE Tokenizer Name, Options: [yttm]') + parser.add_argument('--bpe_droput', type=float, default=0.0, help='BPE dropout to use') + parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff') + parser.add_argument('--pkl_file_prefix', type=str, default='parallel', help='Prefix for tar and pickle files') + parser.add_argument('--fname', type=str, required=True, help='Path to monolingual data file') + parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models') + parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length') + parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length') + parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU') + parser.add_argument( + '--lines_per_dataset_fragment', + type=int, + default=1000000, + help='Number of lines to consider for bucketing and padding', + ) + parser.add_argument( + '--num_batches_per_tarfile', + type=int, + default=1000, + help='Number of batches (pickle files) within each tarfile', + ) + + args = parser.parse_args() + if not os.path.exists(args.out_dir): + os.mkdir(args.out_dir) + if not os.path.exists(args.tokenizer_model): + assert FileNotFoundError("Could not find tokenizer model %s" % (args.tokenizer)) + + tokenizer_model = MTDataPreproc.get_monolingual_tokenizer( + tokenizer_name=args.tokenizer_name, tokenizer_model=args.tokenizer_model, bpe_dropout=args.bpe_droput + ) + + MTDataPreproc.preprocess_monolingual_dataset( + clean=args.clean, + fname=args.fname, + out_dir=args.out_dir, + tokenizer=tokenizer_model, + max_seq_length=args.max_seq_length, + min_seq_length=args.min_seq_length, + tokens_in_batch=args.tokens_in_batch, + lines_per_dataset_fragment=args.lines_per_dataset_fragment, + num_batches_per_tarfile=args.num_batches_per_tarfile, + pkl_file_prefix=args.pkl_file_prefix, + global_rank=0, + world_size=1, + ) diff --git a/models/tts/NeMo/examples/nlp/machine_translation/create_tarred_parallel_dataset.py b/models/tts/NeMo/examples/nlp/machine_translation/create_tarred_parallel_dataset.py new file mode 100644 index 0000000000..d2cd59b37a --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/create_tarred_parallel_dataset.py @@ -0,0 +1,143 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os + +from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='NMT dataset pre-processing') + parser.add_argument('--shared_tokenizer', action="store_true", help='Whether to share encoder/decoder tokenizers') + parser.add_argument('--clean', action="store_true", help='Whether to clean dataset based on length diff') + parser.add_argument('--tar_file_prefix', type=str, default='parallel', help='Prefix for tar files') + parser.add_argument('--bpe_dropout', type=float, default=0.1, help='Whether to share encoder/decoder tokenizers') + parser.add_argument('--src_fname', type=str, required=True, help='Path to the source file') + parser.add_argument('--tgt_fname', type=str, required=True, help='Path to the target file') + parser.add_argument('--out_dir', type=str, required=True, help='Path to store dataloader and tokenizer models') + parser.add_argument('--encoder_model_name', type=str, default=None, help='For use with pretrained encoders') + parser.add_argument( + '--decoder_model_name', type=str, default=None, help='For use with pretrained decoders (not yet supported)' + ) + parser.add_argument( + '--encoder_tokenizer_model', type=str, default='None', help='Path to pre-trained encoder tokenizer model' + ) + parser.add_argument( + '--encoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]' + ) + parser.add_argument('--encoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE') + parser.add_argument( + '--encoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE' + ) + parser.add_argument('--encoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob') + parser.add_argument( + '--decoder_tokenizer_model', type=str, default='None', help='Path to pre-trained decoder tokenizer model' + ) + parser.add_argument( + '--decoder_tokenizer_name', type=str, default='yttm', help='Encoder BPE Tokenizer Name, Options: [yttm]' + ) + parser.add_argument('--decoder_tokenizer_vocab_size', type=int, default=32000, help='Encoder Vocab size after BPE') + parser.add_argument( + '--decoder_tokenizer_coverage', type=float, default=0.999, help='Encoder Character coverage for BPE' + ) + parser.add_argument('--decoder_tokenizer_bpe_dropout', type=float, default=0.1, help='Encoder BPE dropout prob') + parser.add_argument('--max_seq_length', type=int, default=512, help='Max Sequence Length') + parser.add_argument('--min_seq_length', type=int, default=1, help='Min Sequence Length') + parser.add_argument('--tokens_in_batch', type=int, default=16000, help='# Tokens per batch per GPU') + parser.add_argument('--coverage', type=float, default=0.999, help='BPE character coverage [0-1]') + parser.add_argument( + '--lines_per_dataset_fragment', + type=int, + default=1000000, + help='Number of lines to consider for bucketing and padding', + ) + parser.add_argument( + '--num_batches_per_tarfile', + type=int, + default=1000, + help='Number of batches (pickle files) within each tarfile', + ) + parser.add_argument( + '--n_preproc_jobs', type=int, default=-2, help='Number of processes to use for creating the tarred dataset.', + ) + + args = parser.parse_args() + if not os.path.exists(args.out_dir): + os.mkdir(args.out_dir) + + if ( + args.encoder_tokenizer_model != 'None' + and args.decoder_tokenizer_model == 'None' + or args.decoder_tokenizer_model != 'None' + and args.encoder_tokenizer_model == 'None' + ): + if args.shared_tokenizer: + raise ValueError( + ''' + If using a pre-trained shared tokenizer, + both encoder and decoder tokenizers must be the same + ''' + ) + else: + raise ValueError('Both encoder and decoder pre-trained tokenizer models must be specified') + + if args.encoder_tokenizer_model == 'None' and args.decoder_tokenizer_model == 'None': + encoder_tokenizer_model, decoder_tokenizer_model = MTDataPreproc.train_tokenizers( + out_dir=args.out_dir, + src_fname=args.src_fname, + tgt_fname=args.tgt_fname, + shared_tokenizer=args.shared_tokenizer, + encoder_tokenizer_name=args.encoder_tokenizer_name, + encoder_tokenizer_vocab_size=args.encoder_tokenizer_vocab_size, + encoder_tokenizer_coverage=args.encoder_tokenizer_coverage, + decoder_tokenizer_name=args.decoder_tokenizer_name, + decoder_tokenizer_vocab_size=args.decoder_tokenizer_vocab_size, + decoder_tokenizer_coverage=args.decoder_tokenizer_coverage, + global_rank=0, + ) + else: + encoder_tokenizer_model, decoder_tokenizer_model = args.encoder_tokenizer_model, args.decoder_tokenizer_model + + encoder_tokenizer, decoder_tokenizer = MTDataPreproc.get_enc_dec_tokenizers( + encoder_tokenizer_name=args.encoder_tokenizer_name, + encoder_tokenizer_model=encoder_tokenizer_model, + encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout, + decoder_tokenizer_name=args.decoder_tokenizer_name, + decoder_tokenizer_model=decoder_tokenizer_model, + decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout, + ) + + _, _ = MTDataPreproc.preprocess_parallel_dataset( + clean=args.clean, + src_fname=args.src_fname, + tgt_fname=args.tgt_fname, + out_dir=args.out_dir, + encoder_tokenizer_name=args.encoder_tokenizer_name, + encoder_model_name=args.encoder_model_name, + encoder_tokenizer_model=encoder_tokenizer_model, + encoder_bpe_dropout=args.encoder_tokenizer_bpe_dropout, + decoder_tokenizer_name=args.decoder_tokenizer_name, + decoder_model_name=args.decoder_model_name, + decoder_tokenizer_model=decoder_tokenizer_model, + decoder_bpe_dropout=args.decoder_tokenizer_bpe_dropout, + max_seq_length=args.max_seq_length, + min_seq_length=args.min_seq_length, + tokens_in_batch=args.tokens_in_batch, + lines_per_dataset_fragment=args.lines_per_dataset_fragment, + num_batches_per_tarfile=args.num_batches_per_tarfile, + tar_file_prefix=args.tar_file_prefix, + global_rank=0, + world_size=1, + n_jobs=args.n_preproc_jobs, + ) diff --git a/models/tts/NeMo/examples/nlp/machine_translation/enc_dec_nmt.py b/models/tts/NeMo/examples/nlp/machine_translation/enc_dec_nmt.py new file mode 100644 index 0000000000..1cfb8bcfe5 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/enc_dec_nmt.py @@ -0,0 +1,137 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from dataclasses import dataclass +from typing import Optional + +from omegaconf import OmegaConf +from pytorch_lightning import Trainer + +from nemo.collections.nlp.data.machine_translation.preproc_mt_data import MTDataPreproc +from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTEncDecModelConfig +from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel +from nemo.core.config import hydra_runner +from nemo.core.config.modelPT import NemoConfig +from nemo.core.config.pytorch_lightning import TrainerConfig +from nemo.utils import logging +from nemo.utils.config_utils import update_model_config +from nemo.utils.exp_manager import ExpManagerConfig, exp_manager + + +""" +Usage: + 1. If you need to start docker and install NeMo, otherwise skip this step: + + a. ```docker run --gpus all -it --rm -v /home/okuchaiev/repos/NeMo/:/NeMo -p 6006:6006 -v /mnt:/mnt --shm-size=16g --ulimit memlock=-1 --ulimit stack=67108864 --device=/dev/snd nvcr.io/nvidia/pytorch:20.11-py3``` + b. ```cd /NeMo``` + c. ```./reinstall.sh``` + + 2. Train a new tokenizer (or use pre-trained one): + ```yttm bpe --data /mnt/D1/Data/NMT/wmt16_de_en/train.clean.en-de.shuffled.common --model tokenizer.BPE.8192.model --vocab_size 8192``` + +(To use WANDB, optionally, do login first) +``wandb login [YOUR WANDB login]`` + + 3. Start training: + + + (This example for "base" model on 2 GPUs for 150000 steps with batch size of 12500 tokens per GPU) + + python enc_dec_nmt.py \ + --config-path=conf \ + --config-name=aayn_base \ + trainer.gpus=[0,1] \ + ~trainer.max_epochs \ + +trainer.max_steps=150000 \ + model.beam_size=4 \ + model.max_generation_delta=5 \ + model.label_smoothing=0.1 \ + model.encoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model \ + model.decoder_tokenizer.tokenizer_model=tokenizer.BPE.8192.model \ + model.encoder.num_layers=6 \ + model.encoder.hidden_size=512 \ + model.encoder.inner_size=2048 \ + model.encoder.num_attention_heads=8 \ + model.encoder.ffn_dropout=0.1 \ + model.decoder.num_layers=6 \ + model.decoder.hidden_size=512 \ + model.decoder.inner_size=2048 \ + model.decoder.num_attention_heads=8 \ + model.decoder.ffn_dropout=0.1 \ + model.train_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.de.shuffled \ + model.train_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/train.clean.en.shuffled \ + model.train_ds.tokens_in_batch=12500 \ + model.validation_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \ + model.validation_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \ + model.validation_ds.tokens_in_batch=8192 \ + model.test_ds.src_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.ref \ + model.test_ds.tgt_file_name=/mnt/D1/Data/NMT/wmt16_de_en/wmt14-en-de.src \ + model.optim.lr=0.001 \ + model.optim.sched.warmup_ratio=0.05 \ + +exp_manager.create_wandb_logger=True \ + +exp_manager.wandb_logger_kwargs.name=TEST-nmt-base \ + +exp_manager.wandb_logger_kwargs.project=nmt-de-en \ + +exp_manager.create_checkpoint_callback=True \ + +exp_manager.checkpoint_callback_params.monitor=val_sacreBLEU \ + +exp_manager.exp_dir=nmt_base \ + +exp_manager.checkpoint_callback_params.mode=max +""" + + +@dataclass +class MTEncDecConfig(NemoConfig): + name: Optional[str] = 'MTEncDec' + do_training: bool = True + do_testing: bool = False + model: MTEncDecModelConfig = MTEncDecModelConfig() + trainer: Optional[TrainerConfig] = TrainerConfig() + exp_manager: Optional[ExpManagerConfig] = ExpManagerConfig(name='MTEncDec', files_to_copy=[]) + + +@hydra_runner(config_path="conf", config_name="aayn_base") +def main(cfg: MTEncDecConfig) -> None: + # merge default config with user specified config + default_cfg = MTEncDecConfig() + cfg = update_model_config(default_cfg, cfg) + logging.info("\n\n************** Experiment configuration ***********") + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + + # training is managed by PyTorch Lightning + trainer = Trainer(**cfg.trainer) + + # tokenizers will be trained and and tarred training data will be created if needed + # model config is then updated + if cfg.model.preproc_out_dir is not None: + MTDataPreproc(cfg=cfg.model, trainer=trainer) + + # experiment logs, checkpoints, and auto-resume are managed by exp_manager and PyTorch Lightning + exp_manager(trainer, cfg.exp_manager) + + # everything needed to train translation models is encapsulated in the NeMo MTEncdDecModel + mt_model = MTEncDecModel(cfg.model, trainer=trainer) + + logging.info("\n\n************** Model parameters and their sizes ***********") + for name, param in mt_model.named_parameters(): + print(name, param.size()) + logging.info("***********************************************************\n\n") + + if cfg.do_training: + trainer.fit(mt_model) + + if cfg.do_testing: + trainer.test(mt_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py b/models/tts/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py new file mode 100644 index 0000000000..91f2fe9e94 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_transformer_infer.py @@ -0,0 +1,122 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Given NMT model's .nemo file, this script can be used to translate text. +USAGE Example: +1. Obtain text file in src language. You can use sacrebleu to obtain standard test sets like so: + sacrebleu -t wmt14 -l de-en --echo src > wmt14-de-en.src +2. Translate: + python nmt_transformer_infer.py --model=[Path to .nemo file] --srctext=wmt14-de-en.src --tgtout=wmt14-de-en.pre +""" + + +from argparse import ArgumentParser + +import torch + +import nemo.collections.nlp as nemo_nlp +from nemo.collections.nlp.modules.common.transformer import ( + BeamSearchSequenceGenerator, + BeamSearchSequenceGeneratorWithLanguageModel, +) +from nemo.utils import logging + + +def main(): + parser = ArgumentParser() + parser.add_argument("--model", type=str, required=True, help="") + parser.add_argument("--srctext", type=str, required=True, help="") + parser.add_argument("--tgtout", type=str, required=True, help="") + parser.add_argument("--batch_size", type=int, default=256, help="") + parser.add_argument("--beam_size", type=int, default=4, help="") + parser.add_argument("--len_pen", type=float, default=0.6, help="") + parser.add_argument("--max_delta_length", type=int, default=5, help="") + parser.add_argument("--target_lang", type=str, default=None, help="") + parser.add_argument("--source_lang", type=str, default=None, help="") + # shallow fusion specific parameters + parser.add_argument("--lm_model", type=str, default=None, help="") + parser.add_argument("--fusion_coef", type=float, default=0.0, help="") + + args = parser.parse_args() + torch.set_grad_enabled(False) + if args.model.endswith(".nemo"): + logging.info("Attempting to initialize from .nemo file") + model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=args.model) + src_text = [] + tgt_text = [] + else: + raise NotImplemented(f"Only support .nemo files, but got: {args.model}") + + if torch.cuda.is_available(): + model = model.cuda() + + if args.lm_model is not None: + lm_model = nemo_nlp.models.language_modeling.TransformerLMModel.restore_from(restore_path=args.lm_model).eval() + model.beam_search = BeamSearchSequenceGeneratorWithLanguageModel( + embedding=model.decoder.embedding, + decoder=model.decoder.decoder, + log_softmax=model.log_softmax, + bos=model.decoder_tokenizer.bos_id, + pad=model.decoder_tokenizer.pad_id, + eos=model.decoder_tokenizer.eos_id, + language_model=lm_model, + fusion_coef=args.fusion_coef, + max_sequence_length=model.decoder.max_sequence_length, + beam_size=args.beam_size, + len_pen=args.len_pen, + max_delta_length=args.max_delta_length, + ) + else: + model.beam_search = BeamSearchSequenceGenerator( + embedding=model.decoder.embedding, + decoder=model.decoder.decoder, + log_softmax=model.log_softmax, + bos=model.decoder_tokenizer.bos_id, + pad=model.decoder_tokenizer.pad_id, + eos=model.decoder_tokenizer.eos_id, + max_sequence_length=model.decoder.max_sequence_length, + beam_size=args.beam_size, + len_pen=args.len_pen, + max_delta_length=args.max_delta_length, + ) + + logging.info(f"Translating: {args.srctext}") + + count = 0 + with open(args.srctext, 'r') as src_f: + for line in src_f: + src_text.append(line.strip()) + if len(src_text) == args.batch_size: + res = model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang) + if len(res) != len(src_text): + print(len(res)) + print(len(src_text)) + print(res) + print(src_text) + tgt_text += res + src_text = [] + count += 1 + # if count % 300 == 0: + # print(f"Translated {count} sentences") + if len(src_text) > 0: + tgt_text += model.translate(text=src_text, source_lang=args.source_lang, target_lang=args.target_lang) + + with open(args.tgtout, 'w') as tgt_f: + for line in tgt_text: + tgt_f.write(line + "\n") + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/README.rst b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/README.rst new file mode 100644 index 0000000000..dd453d165b --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/README.rst @@ -0,0 +1,7 @@ +**STEPS 2 RUN** +=============== +0. (Make sure you have Flask installed - ``pip install -r requirements.txt``) +1. Edit "config.json" file to only contain models you need. If model's location starts with "NGC/" - it will load this model from NVIDIA's NGC. Otherwise, specify full path to .nemo file. +4. To run: ``python nmt_service.py`` +5. To translate: ``http://127.0.0.1:5000/translate?text=Frohe%20Weihnachten`` (here %20 means space) +6. Run web UI: ``python -m http.server`` \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/config.json b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/config.json new file mode 100644 index 0000000000..938d91a0f3 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/config.json @@ -0,0 +1,12 @@ +{ + "en-de": "NGC/nmt_en_de_transformer12x2", + "de-en": "NGC/nmt_de_en_transformer12x2", + "en-es": "NGC/nmt_en_es_transformer12x2", + "es-en": "NGC/nmt_es_en_transformer12x2", + "en-ru": "NGC/nmt_en_ru_transformer6x6", + "ru-en": "NGC/nmt_ru_en_transformer6x6", + "en-fr": "NGC/nmt_en_fr_transformer12x2", + "fr-en": "NGC/nmt_fr_en_transformer12x2", + "en-zh": "NGC/nmt_en_zh_transformer6x6", + "zh-en": "NGC/nmt_zh_en_transformer6x6" +} diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/index.html b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/index.html new file mode 100644 index 0000000000..2f639d5860 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/index.html @@ -0,0 +1,140 @@ + + + + + + Neural Machine Translation Demo App + + + + + + + + + +
+

Neural Machine Translation Demo

+

+
+

Running NeMo commit id: v1.0.2

+
+

+

+

+ Select translation direction: + + + + + + + + + + + + + + + + + + + + + + + +
+ + +             + + +             + + +             + + +
+ + +             + + +             + + +             + + +             + + +             + + +
+
+
+ + + + + + + + + + +
+ +
+ +
+ +
+
+

+
+ + + + diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/nmt_service.py b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/nmt_service.py new file mode 100644 index 0000000000..3661fcdb4a --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/nmt_service.py @@ -0,0 +1,100 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 json +import time + +import flask +import torch +from flask import Flask, json, request +from flask_cors import CORS + +import nemo.collections.nlp as nemo_nlp +from nemo.utils import logging + +MODELS_DICT = {} + +model = None +api = Flask(__name__) +CORS(api) + + +def initialize(config_file_path: str): + """ + Loads 'language-pair to NMT model mapping' + """ + __MODELS_DICT = None + + logging.info("Starting NMT service") + logging.info(f"I will attempt to load all the models listed in {config_file_path}.") + logging.info(f"Edit {config_file_path} to disable models you don't need.") + if torch.cuda.is_available(): + logging.info("CUDA is available. Running on GPU") + else: + logging.info("CUDA is not available. Defaulting to CPUs") + + # read config + with open(config_file_path) as f: + __MODELS_DICT = json.load(f) + + if __MODELS_DICT is not None: + for key, value in __MODELS_DICT.items(): + logging.info(f"Loading model for {key} from file: {value}") + if value.startswith("NGC/"): + model = nemo_nlp.models.machine_translation.MTEncDecModel.from_pretrained(model_name=value[4:]) + else: + model = nemo_nlp.models.machine_translation.MTEncDecModel.restore_from(restore_path=value) + if torch.cuda.is_available(): + model = model.cuda() + MODELS_DICT[key] = model + else: + raise ValueError("Did not find the config.json or it was empty") + logging.info("NMT service started") + + +@api.route('/translate', methods=['GET', 'POST', 'OPTIONS']) +def get_translation(): + try: + time_s = time.time() + langpair = request.args["langpair"] + src = request.args["text"] + do_moses = request.args.get('do_moses', False) + if langpair in MODELS_DICT: + if do_moses: + result = MODELS_DICT[langpair].translate( + [src], source_lang=langpair.split('-')[0], target_lang=langpair.split('-')[1] + ) + else: + result = MODELS_DICT[langpair].translate([src]) + + duration = time.time() - time_s + logging.info( + f"Translated in {duration}. Input was: {request.args['text']} <############> Translation was: {result[0]}" + ) + res = {'translation': result[0]} + response = flask.jsonify(res) + response.headers.add('Access-Control-Allow-Origin', '*') + return response + + else: + logging.error(f"Got the following langpair: {langpair} which was not found") + except Exception as ex: + res = {'translation': str(ex)} + response = flask.jsonify(res) + response.headers.add('Access-Control-Allow-Origin', '*') + return res + + +if __name__ == '__main__': + initialize('config.json') + api.run(host='0.0.0.0') diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/requirements.txt b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/requirements.txt new file mode 100644 index 0000000000..cd6e0e7486 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/requirements.txt @@ -0,0 +1,3 @@ +nemo_toolkit[all]>=1.0.0rc1 +flask +flask_cors \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/style.css b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/style.css new file mode 100644 index 0000000000..58f23342f5 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/nmt_webapp/style.css @@ -0,0 +1,111 @@ +@import url("https://fonts.googleapis.com/css?family=Spartan&display=swap"); +* { + margin: 0; + padding: 0; +} + +body { + /*background: url("https://img.wallpapersafari.com/desktop/1680/1050/65/18/8AYkSt.jpg") center center;*/ + background: url("https://images5.alphacoders.com/550/550948.jpg") center center; + background-size: cover; + width: 100vw; + height: 100vh; + display: grid; + align-items: center; + justify-items: center; + font-size: 12pt; + font-family: "Spartan"; + color: #ffffff; +} + +.contact-us { + /*background: #1f1d1d;*/ + padding: 50px 100px; + border-top: 10px solid #229955; +} + +label, input, textarea { + display: block; + width: 100%; + font-size: 12pt; + line-height: 24pt; + font-family: "Spartan"; +} + +input { + margin-bottom: 24pt; +} + +h1 { + color: #ffffff; + font-weight: bold; + font-size: 22pt; + line-height: 24pt; + font-style: italic; + margin: 0 0 0.5em 0; +} + +span { + font-size: 8pt; +} + +em { + color: #229955; + font-weight: bold; +} + +input, textarea { + border: none; + border: 1px solid rgba(0, 0, 0, 0.1); + border-radius: 2px; + background: #f8f4e5; + padding-left: 5px; + outline: none; +} + +input:focus, textarea:focus { + border: 1px solid #6bd4b1; +} + +textarea { + resize: none; +} + +fieldset { + overflow: hidden + } + +button { + display: block; + float: left; + line-height: 24pt; + padding: 0 20px; + border: none; + background: #f45702; + color: white; + letter-spacing: 2px; + transition: 0.2s all ease-in-out; + border-bottom: 2px solid transparent; + outline: none; +} +button:hover { + background: inherit; + color: #ea6153; + border-bottom: 2px solid #f45702; +} + +::selection { + background: #ffc7b8; +} + +input:-webkit-autofill, +input:-webkit-autofill:hover, +input:-webkit-autofill:focus, +textarea:-webkit-autofill, +textarea:-webkit-autofill:hover, +textarea:-webkit-autofill:focus { + border: 1px solid #6bd4b1; + -webkit-text-fill-color: #2A293E; + -webkit-box-shadow: 0 0 0px 1000px #f8f4e5 inset; + transition: background-color 5000s ease-in-out 0s; +} \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/machine_translation/translate_ddp.py b/models/tts/NeMo/examples/nlp/machine_translation/translate_ddp.py new file mode 100644 index 0000000000..bf057df9a1 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/machine_translation/translate_ddp.py @@ -0,0 +1,122 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +from argparse import ArgumentParser + +import torch +import torch.multiprocessing as mp +from torch.utils.data import DataLoader + +from nemo.collections.nlp.data.language_modeling import TarredSentenceDataset +from nemo.collections.nlp.data.machine_translation import TarredTranslationDataset +from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel +from nemo.utils import logging + + +def get_args(): + parser = ArgumentParser(description='Batch translation of sentences from a pre-trained model on multiple GPUs') + parser.add_argument("--model", type=str, required=True, help="Path to the .nemo translation model file") + parser.add_argument( + "--text2translate", type=str, required=True, help="Path to the pre-processed tarfiles for translation" + ) + parser.add_argument("--result_dir", type=str, required=True, help="Folder to write translation results") + parser.add_argument( + "--twoside", action="store_true", help="Set flag when translating the source side of a parallel dataset" + ) + parser.add_argument( + '--metadata_path', type=str, required=True, help="Path to the JSON file that contains dataset info" + ) + parser.add_argument('--topk', type=int, default=500, help="Value of k for topk sampling") + parser.add_argument('--src_language', type=str, required=True, help="Source lang ID for detokenization") + parser.add_argument('--tgt_language', type=str, required=True, help="Target lang ID for detokenization") + parser.add_argument( + '--reverse_lang_direction', + action="store_true", + help="Reverse source and target language direction for parallel dataset", + ) + parser.add_argument('--n_gpus', type=int, default=-1, help="Number of GPUs to use") + args = parser.parse_args() + return args + + +def translate(rank, world_size, args): + if args.model.endswith(".nemo"): + logging.info("Attempting to initialize from .nemo file") + model = MTEncDecModel.restore_from(restore_path=args.model, map_location=f"cuda:{rank}") + elif args.model.endswith(".ckpt"): + logging.info("Attempting to initialize from .ckpt file") + model = MTEncDecModel.load_from_checkpoint(checkpoint_path=args.model, map_location=f"cuda:{rank}") + model.replace_beam_with_sampling(topk=args.topk) + model.eval() + if args.twoside: + dataset = TarredTranslationDataset( + text_tar_filepaths=args.text2translate, + metadata_path=args.metadata_path, + encoder_tokenizer=model.encoder_tokenizer, + decoder_tokenizer=model.decoder_tokenizer, + shuffle_n=100, + shard_strategy="scatter", + world_size=world_size, + global_rank=rank, + reverse_lang_direction=args.reverse_lang_direction, + ) + else: + dataset = TarredSentenceDataset( + text_tar_filepaths=args.text2translate, + metadata_path=args.metadata_path, + tokenizer=model.encoder_tokenizer, + shuffle_n=100, + shard_strategy="scatter", + world_size=world_size, + global_rank=rank, + ) + loader = DataLoader(dataset, batch_size=1) + result_dir = os.path.join(args.result_dir, f'rank{rank}') + os.makedirs(result_dir, exist_ok=True) + originals_file_name = os.path.join(result_dir, 'originals.txt') + translations_file_name = os.path.join(result_dir, 'translations.txt') + num_translated_sentences = 0 + + with open(originals_file_name, 'w') as of, open(translations_file_name, 'w') as tf: + for batch_idx, batch in enumerate(loader): + for i in range(len(batch)): + if batch[i].ndim == 3: + batch[i] = batch[i].squeeze(dim=0) + batch[i] = batch[i].to(rank) + if args.twoside: + src_ids, src_mask, _, _, _ = batch + else: + src_ids, src_mask = batch + if batch_idx % 100 == 0: + logging.info( + f"{batch_idx} batches ({num_translated_sentences} sentences) were translated by process with " + f"rank {rank}" + ) + num_translated_sentences += len(src_ids) + inputs, translations = model.batch_translate(src=src_ids, src_mask=src_mask) + for src, translation in zip(inputs, translations): + of.write(src + '\n') + tf.write(translation + '\n') + + +def main() -> None: + args = get_args() + world_size = torch.cuda.device_count() if args.n_gpus == -1 else args.n_gpus + mp.spawn(translate, args=(world_size, args), nprocs=world_size, join=True) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/question_answering/conf/question_answering_squad_config.yaml b/models/tts/NeMo/examples/nlp/question_answering/conf/question_answering_squad_config.yaml new file mode 100644 index 0000000000..929409d53d --- /dev/null +++ b/models/tts/NeMo/examples/nlp/question_answering/conf/question_answering_squad_config.yaml @@ -0,0 +1,135 @@ +# Question Answering with SQUAD +name: &name QA + +pretrained_model: null # pretrained QAModel model from list_available_models() +do_training: true # training mode, for testing change to false +trainer: + gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices + num_nodes: 1 + max_epochs: 2 # the number of training epochs + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + precision: 16 # 16 to use AMP + amp_level: O1 # O1 or O2 if using AMP + accelerator: ddp + gradient_clip_val: 0.0 + val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch + checkpoint_callback: false # provided by exp_manager + logger: false # provided by exp_manager + num_sanity_val_steps: 0 + log_every_n_steps: 1 # Interval of logging. + +model: + nemo_path: null # exported .nemo path + dataset: + version_2_with_negative: false + # If true, the examples contain some that do not have an answer. + doc_stride: 128 + # When splitting up a long document into chunks, + # how much stride to take between chunks. + max_query_length: 64 + # The maximum number of tokens for the question. + # Questions longer than this will be truncated to + # this length. + max_seq_length: 384 + # The maximum total input sequence length after + # WordPiece tokenization. Sequences longer than this + # will be truncated, and sequences shorter than this + # will be padded. + max_answer_length: 30 + # The maximum length of an answer that can be + # generated. This is needed because the start + # and end predictions are not conditioned + # on one another. + null_score_diff_threshold: 0.0 + # If null_score - best_non_null is greater than the threshold predict null. + n_best_size: 20 + # The total number of n-best predictions to generate at testing. + use_cache: true + do_lower_case: true + # if true does lower case + + num_workers: 2 + pin_memory: false + drop_last: false + + train_ds: + file: null # .json file + batch_size: 24 # per GPU + shuffle: true + num_samples: -1 + # Default values for the following params are retrieved from dataset config section, but you may override them + num_workers: ${model.dataset.num_workers} + drop_last: ${model.dataset.drop_last} + pin_memory: ${model.dataset.pin_memory} + + validation_ds: + file: null # .json file + batch_size: 24 # per GPU + shuffle: false + num_samples: -1 + # Default values for the following params are retrieved from dataset config section, but you may override them + num_workers: ${model.dataset.num_workers} + drop_last: ${model.dataset.drop_last} + pin_memory: ${model.dataset.pin_memory} + + test_ds: + file: null # .json file + batch_size: 24 # per GPU + shuffle: false + num_samples: -1 + # Default values for the following params are retrieved from dataset config section, but you may override them + num_workers: ${model.dataset.num_workers} + drop_last: ${model.dataset.drop_last} + pin_memory: ${model.dataset.pin_memory} + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null # expand the following to a dictionary if special tokens need to be added. + # only necessary for adding transformer/bert-specific special tokens to tokenizer if the tokenizer does not already have these inherently. + + language_model: + pretrained_model_name: bert-base-uncased # BERT-like model name + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null # if specified initializes model from scratch + + token_classifier: + num_layers: 1 + dropout: 0.0 + num_classes: 2 + activation: relu + log_softmax: false + use_transformer_init: true + + + optim: + name: adamw + lr: 3e-5 + weight_decay: 0.0 + sched: + name: SquareRootAnnealing + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + + # scheduler config override + warmup_steps: null + warmup_ratio: 0.0 + last_epoch: -1 + +exp_manager: + exp_dir: null # where to store logs and checkpoints + name: *name # name of experiment + create_tensorboard_logger: True + create_checkpoint_callback: True + +hydra: + run: + dir: . + job_logging: + root: + handlers: null \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/question_answering/get_squad.py b/models/tts/NeMo/examples/nlp/question_answering/get_squad.py new file mode 100755 index 0000000000..255b040c36 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/question_answering/get_squad.py @@ -0,0 +1,68 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os +import urllib.request + +from nemo.utils import logging + + +class SquadDownloader: + def __init__(self, save_path): + self.save_path = save_path + '/squad' + + if not os.path.exists(self.save_path): + os.makedirs(self.save_path) + + if not os.path.exists(self.save_path + '/v1.1'): + os.makedirs(self.save_path + '/v1.1') + + if not os.path.exists(self.save_path + '/v2.0'): + os.makedirs(self.save_path + '/v2.0') + + self.download_urls = { + 'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/train-v1.1.json': 'v1.1/train-v1.1.json', + 'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/dev-v1.1.json': 'v1.1/dev-v1.1.json', + 'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/train-v2.0.json': 'v2.0/train-v2.0.json', + 'https://rajpurkar.github.io/SQuAD-explorer' '/dataset/dev-v2.0.json': 'v2.0/dev-v2.0.json', + } + + def download(self): + for item in self.download_urls: + url = item + file = self.download_urls[item] + + logging.info('Downloading: %s', url) + if os.path.isfile(self.save_path + '/' + file): + logging.info('** Download file already exists, skipping download') + else: + response = urllib.request.urlopen(url) + with open(self.save_path + '/' + file, "wb") as handle: + handle.write(response.read()) + + +if __name__ == '__main__': + parser = argparse.ArgumentParser(description='Download Squad') + parser.add_argument( + '--destDir', + type=str, + required=False, + help='directory to store data', + default=os.path.split(os.path.abspath(__file__))[0], + ) + args = parser.parse_args() + logging.info(args.destDir) + squad_dl = SquadDownloader(args.destDir) + squad_dl.download() diff --git a/models/tts/NeMo/examples/nlp/question_answering/question_answering_squad.py b/models/tts/NeMo/examples/nlp/question_answering/question_answering_squad.py new file mode 100644 index 0000000000..269b4bfe48 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/question_answering/question_answering_squad.py @@ -0,0 +1,134 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + + +""" +This script contains an example on how to train, evaluate and perform inference with the question answering model. +The QAModel in NeMo supports extractive question answering problems for data in the SQuAD (https://rajpurkar.github.io/SQuAD-explorer/) format. + +***Data format*** +The QAModel requires a JSON file for each dataset split. +In the following we will show example for a training file. Each title has one or multiple paragraph entries, each consisting of the text - "context", and question-answer entries. Each question-answer entry has: +* a question +* a globally unique id +* a boolean flag "is_impossible" which shows if the question is answerable or not +* in case the question is answerable one answer entry, which contains the text span and its starting character index in the context. If not answerable, the "answers" list is empty + +The evaluation file follows the above format except for it can provide more than one answers to the same question. +The inference file follows the above format except for it does not require the "answers" and "is_impossible" keywords. + + +***Downloading the dataset*** +Run ./NeMo/examples/nlp/question_answering/get_squad.py to download the SQuAD dataset: + +# python get_squad.py --destDir= + +***Setting the configs*** +The model and the PT trainer are defined in a config file which declares multiple important sections. +The most important ones are: + model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer, + schedulers, and datasets/data loaders. + trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs, + precision level, etc. + +This script uses the `/examples/nlp/question_answering/conf/question_answering_squad_config.yaml` config file +by default. You may update the config file from the file directly. The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + + +***Model Training*** +# python question_answering_squad.py + model.train_ds.file= + model.validation_ds= + trainer.max_epochs= + trainer.gpus=[] + + +***Model Evaluation*** +Set `do_training=False` in the script and run: + +# python question_answering_squad.py + model.test_file= + +To load a pretrained checkpoint from cloud prior to training (e.g. for fine-tuning) or evaluation you can set cfg.from_pretrained=, +e.g. MODEL_NAME='BERTBaseUncasedSQuADv1.1'. You can find all pretrained model names by using +QAModel.list_available_models(). To load a local checkpoint use qa_model.restore_from() + + +***Model Inference*** +For inference use + qa_model.inference( + file=, + batch_size=, + output_nbest_file=, + output_prediction_file= + ) + +More details on how to use this script can be found in +./NeMo/tutorials/nlp/Question_Answering_Squad.ipynb +""" + +import os + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.question_answering.qa_model import QAModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="question_answering_squad_config") +def main(cfg: DictConfig) -> None: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + exp_dir = exp_manager(trainer, cfg.get("exp_manager", None)) + + if not cfg.pretrained_model: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + model = QAModel(cfg.model, trainer=trainer) + else: + logging.info(f'Loading pretrained model {cfg.pretrained_model}') + model = QAModel.from_pretrained(cfg.pretrained_model) + if cfg.do_training: + model.setup_training_data(train_data_config=cfg.model.train_ds) + model.setup_validation_data(val_data_config=cfg.model.validation_ds) + + if cfg.do_training: + trainer.fit(model) + if cfg.model.nemo_path: + model.save_to(cfg.model.nemo_path) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.file is not None: + trainer.test(model) + + # change to path if you want results to be written to file e.g. os.path.join(exp_dir, "output_nbest_file.txt") + output_nbest_file = None + # change to path if you want results to be written to file e.g. os.path.join(exp_dir, "output_prediction_file.txt") + output_prediction_file = None + inference_samples = 5 # for test purposes. To use entire inference dataset set to -1 + all_preds, all_nbests = model.inference( + file=cfg.model.validation_ds.file, + batch_size=1, + num_samples=inference_samples, + output_nbest_file=output_nbest_file, + output_prediction_file=output_prediction_file, + ) + + for _, item in all_preds.items(): + print(f"question: {item[0]} answer: {item[1]}") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/text2sparql/conf/text2sparql_config.yaml b/models/tts/NeMo/examples/nlp/text2sparql/conf/text2sparql_config.yaml new file mode 100644 index 0000000000..1603279ce8 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text2sparql/conf/text2sparql_config.yaml @@ -0,0 +1,105 @@ +# Text2Sparql with BART + +name: &name Text2Sparql + +trainer: + gpus: 1 # the number of gpus, 0 for CPU, or list with gpu indices + num_nodes: 1 + max_epochs: 2 # the number of training epochs + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + accelerator: ddp + gradient_clip_val: 0.0 + log_every_n_steps: 1 + val_check_interval: 1.0 # check once per epoch .25 for 4 times per epoch + checkpoint_callback: false # provided by exp_manager + logger: false # provided by exp_manager + +model: + nemo_path: null # exported .nemo path + max_seq_length: 150 + batch_size: 16 + convert_labels: true # true if Bart, false otherwise (converts pad_id to -100 for masked loss) + data_dir: null + + language_model: + pretrained_model_name: facebook/bart-base # huggingface end-to-end model name + pretrained_encoder_model_name: null # huggingface encoder model name + pretrained_decoder_model_name: null # huggingface decoder model name + lm_checkpoint: null + config: null + config_file: null # json file, precedence over config + + encoder_tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec + vocab_file: null # path to vocab file + tokenizer_model: null # tokenizer model for sentencepiece + special_tokens: null + add_special_tokens: true + + decoder_tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # tokenizer that inherits from TokenizerSpec + vocab_file: null # path to vocab file + tokenizer_model: null # tokenizer model for sentencepiece + special_tokens: null + add_special_tokens: true + + train_ds: + filepath: ${model.data_dir}/train.tsv # path to data file + shuffle: true + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + validation_ds: + filepath: ${model.data_dir}/test_easy.tsv # path to data file + shuffle: false + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + test_ds: + filepath: ${model.data_dir}/test_hard.tsv # path to data file + shuffle: false + num_samples: -1 + num_workers: 2 + drop_last: false + pin_memory: false + + optim: + name: adamw + lr: 4e-5 + weight_decay: 0.0 + + sched: + name: CosineAnnealing + warmup_steps: null + warmup_ratio: 0.06 + min_lr: 0.0 + last_epoch: -1 + + generate: + max_length: ${model.max_seq_length} + num_beams: 1 + length_penalty: 2.0 + early_stopping: true + repetition_penalty: 1.0 + do_sample: false + top_k: null + top_p: null + num_return_sequences: 1 + +exp_manager: + exp_dir: null # where to store logs and checkpoints + name: *name # name of experiment + create_tensorboard_logger: True + create_checkpoint_callback: True + +hydra: + run: + dir: . + job_logging: + root: + handlers: null \ No newline at end of file diff --git a/models/tts/NeMo/examples/nlp/text2sparql/data/import_datasets.py b/models/tts/NeMo/examples/nlp/text2sparql/data/import_datasets.py new file mode 100644 index 0000000000..ce5f86393b --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text2sparql/data/import_datasets.py @@ -0,0 +1,134 @@ +# Copyright (c) 2020, MeetKai Inc. All rights reserved. +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script downloads Text2Sparql data and processes it into NeMo's neural machine translation dataset format. + +Text2Sparql data consists of 3 files which are saved to source_data_dir: + - train_queries_v3.tsv + - test_easy_queries_v3.tsv + - test_hard_queries_v3.tsv + +After processing, the script saves them to the target_data_dir as: + - train.tsv + - test_easy.tsv + - test_hard.tsv + + +You may run it with: + +python import_datasets \ + --source_data_dir ./text2sparql_src \ + --target_data_dir ./text2sparql_tgt +""" + +import argparse +import csv +import os +from urllib.request import Request, urlopen + +from nemo.collections.nlp.data.data_utils.data_preprocessing import MODE_EXISTS_TMP, if_exist +from nemo.utils import logging + +base_url = "https://m.meetkai.com/public_datasets/knowledge/" +prefix_map = { + "train_queries_v3.tsv": "train.tsv", + "test_easy_queries_v3.tsv": "test_easy.tsv", + "test_hard_queries_v3.tsv": "test_hard.tsv", +} + + +def download_text2sparql(infold: str): + """Downloads text2sparql train, test_easy, and test_hard data + + Args: + infold: save directory path + """ + os.makedirs(infold, exist_ok=True) + + for prefix in prefix_map: + url = base_url + prefix + + logging.info(f"Downloading: {url}") + if if_exist(infold, [prefix]): + logging.info("** Download file already exists, skipping download") + else: + req = Request(url, headers={"User-Agent": "Mozilla/5.0"}) + with open(os.path.join(infold, prefix), "wb") as handle: + handle.write(urlopen(req, timeout=20).read()) + + +def process_text2sparql(infold: str, outfold: str, do_lower_case: bool): + """ Process and convert MeetKai's text2sparql datasets to NeMo's neural machine translation format. + + Args: + infold: directory path to raw text2sparql data containing + train.tsv, test_easy.tsv, test_hard.tsv + outfold: output directory path to save formatted data for NeuralMachineTranslationDataset + the first line is header (sentence [tab] label) + each line should be [sentence][tab][label] + do_lower_case: if true, convert all sentences and labels to lower + """ + logging.info(f"Processing Text2Sparql dataset and storing at: {outfold}") + + os.makedirs(outfold, exist_ok=True) + + dataset_name = "Text2Sparql" + for prefix in prefix_map: + input_file = os.path.join(infold, prefix) + output_file = os.path.join(outfold, prefix_map[prefix]) + + if if_exist(outfold, [prefix_map[prefix]]): + logging.info(f"** {MODE_EXISTS_TMP.format(prefix_map[prefix], dataset_name, output_file)}") + continue + + if not if_exist(infold, [prefix]): + logging.info(f"** {prefix} of {dataset_name}" f" is skipped as it was not found") + continue + + assert input_file != output_file, "input file cannot equal output file" + with open(input_file, "r") as in_file: + with open(output_file, "w") as out_file: + reader = csv.reader(in_file, delimiter="\t") + + # replace headers + out_file.write("sentence\tlabel\n") + next(reader) + + for line in reader: + sentence = line[0] + label = line[1] + if do_lower_case: + sentence = sentence.lower() + label = label.lower() + out_file.write(f"{sentence}\t{label}\n") + + +if __name__ == "__main__": + # Parse the command-line arguments. + parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo's format") + parser.add_argument( + "--source_data_dir", required=True, type=str, help="Path to the folder containing the dataset files" + ) + parser.add_argument("--target_data_dir", required=True, type=str, help="Path to save the processed dataset") + parser.add_argument("--do_lower_case", action="store_true") + args = parser.parse_args() + + source_dir = args.source_data_dir + target_dir = args.target_data_dir + do_lower_case = args.do_lower_case + + download_text2sparql(infold=source_dir) + process_text2sparql(infold=source_dir, outfold=target_dir, do_lower_case=do_lower_case) diff --git a/models/tts/NeMo/examples/nlp/text2sparql/evaluate_text2sparql.py b/models/tts/NeMo/examples/nlp/text2sparql/evaluate_text2sparql.py new file mode 100644 index 0000000000..050d142e99 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text2sparql/evaluate_text2sparql.py @@ -0,0 +1,72 @@ +# Copyright (c) 2020, MeetKai Inc. All rights reserved. +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script contains an example on how to evaluate a NeuralMachineTranslationModel. +To load the example Text2Sparql dataset, please refer to ./data/import_datasets.py. +To train a model, please refer to text2sparql.py. + + +***Setting the configs*** +This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default. +You may update the config file from the file directly or by using the command line arguments. +Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +Please refer to text2sparql.py for detailed instructions on setting the configuration. + + +***How to run the script?*** +- To reload and evaluate the model, run: + +python evaluate_text2sparql.py \ + model.test_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \ + model.batch_size=16 \ + model.nemo_path=./NeMo_logs/bart.nemo \ + exp_manager.exp_dir=./NeMo_logs +""" + +import os + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging + + +@hydra_runner(config_path="conf", config_name="text2sparql_config") +def main(cfg: DictConfig) -> None: + logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}") + trainer = pl.Trainer(gpus=cfg.trainer.gpus) + nmt_model = NeuralMachineTranslationModel.restore_from(restore_path=cfg.model.nemo_path) + nmt_model.setup_test_data(cfg.model.test_ds) + results = trainer.test(nmt_model) + + with open(cfg.model.test_ds.filepath, "r") as f: + lines = f.readlines() + + lines[0] = lines[0].strip() + f"\tpredictions\n" + for i, res in enumerate(results[0]["texts"]): + lines[i + 1] = lines[i + 1].strip() + f"\t{res}\n" + + savepath = os.path.join(cfg.exp_manager.exp_dir, os.path.basename(cfg.model.test_ds.filepath)) + with open(savepath, "w") as f: + f.writelines(lines) + logging.info(f"Predictions saved to {savepath}") + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/examples/nlp/text2sparql/text2sparql.py b/models/tts/NeMo/examples/nlp/text2sparql/text2sparql.py new file mode 100644 index 0000000000..062094ccd3 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text2sparql/text2sparql.py @@ -0,0 +1,112 @@ +# Copyright (c) 2020, MeetKai Inc. All rights reserved. +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script contains an example on how to train and save a NeuralMachineTranslationModel. +NeuralMachineTranslationModel in NeMo supports sequence to sequence problems such as language translation +and text summarization, provided the data follows the format specified below. + + +***Data format*** +NeuralMachineTranslationModel requires the data to be stored in TAB separated files (.tsv) with two columns of +sentence and label, where the first line is a header of format: + sentence[TAB]label +And each line is of the format: + [SENTENCE][TAB][LABEL] + +If your dataset is stored in another format, you need to convert it to this format to use a +NeuralMachineTranslationModel. + + +***Setting the configs*** +This script uses the `/examples/nlp/text2sparql/conf/text2sparql_config.yaml` config file by default. +You may update the config file from the file directly or by using the command line arguments. +Another other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +A NeuralMachineTranslationModel's config file declares multiple import sections. They are: + - trainer: Arguments to be passed to PyTorch Lightning. + - model: All arguments that relate to the Model - language_model, tokenizers, datasets, optimizer, generate. + - exp_manager: Arguments to be passed to NeMo's experiment manager. + - hydra: Arguments to be passed to Hydra. + +If using text2sparql_config.yaml, you must first update the following fields in the config: + - model.nemo_path: Model save path. Eg. [PATH]/bart.nemo + - model.data_dir: Path to data directory. Alternatively, you can adjust the file paths directly: + - model.train_ds.filepath + - model.validation_ds.filepath + - model.test_ds.filepath + - exp_manager.exp_dir: Directory to log results from the experiment. + +It is highly recommended to also adjust these parameters as necessary: + - trainer.gpus: Set to 0 to use CPU. Otherwise the number denotes the number of GPUs. + - trainer.max_epochs: Maximum number of epochs to train for. + - model.batch_size: 8 is sufficient to train a decent Bart model for Text2Sparql. + - model.max_seq_length: Maximum (tokenized) sequence length. 150 is sufficient for Text2Sparql. + - model.language_model.pretrained_model_name: End2end pretrained model name from huggingface. + - model.encoder_tokenizer.tokenizer_name: Pretrained tokenizer name from huggingface. + - model.decoder_tokenizer.tokenizer_name: The same as above, as the tokenizer will handle encoding and decoding. + - model.optim.lr: Learning rate. + +You can also specify an encoder and decoder rather than using an end2end model like Bart by defining these parameters: + - model.language_model.pretrained_encoder_model_name: Pretrained huggingface encoder model name. + - model.encoder_tokenizer.tokenizer_name: Pretrained huggingface encoder tokenizer name. + - model.language_model.pretrained_decoder_model_name: Pretrained huggingface decoder model name. + - model.decoder_tokenizer.tokenizer_name: Pretrained huggingface decoder tokenizer name. + - model.language_model.pretrained_model_name: Set this to null. + + +***How to run the script?*** +- First, download the data to TGT_DATA_DIR (see ./data/import_datasets.py): + +SRC_DATA_DIR=./data/text2sparql_src +TGT_DATA_DIR=./data/text2sparql_tgt + +python ./data/import_datasets.py \ + --source_data_dir $SRC_DATA_DIR \ + --target_data_dir $TGT_DATA_DIR + +- And run the following to train and save the model: + +python text2sparql.py \ + model.train_ds.filepath="$TGT_DATA_DIR"/train.tsv \ + model.validation_ds.filepath="$TGT_DATA_DIR"/test_easy.tsv \ + model.test_ds.filepath="$TGT_DATA_DIR"/test_hard.tsv \ + model.batch_size=16 \ + model.nemo_path=./NeMo_logs/bart.nemo \ + exp_manager.exp_dir=./NeMo_logs +""" + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.neural_machine_translation import NeuralMachineTranslationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="text2sparql_config") +def main(cfg: DictConfig) -> None: + logging.info(f"Config:\n {OmegaConf.to_yaml(cfg)}") + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + nmt_model = NeuralMachineTranslationModel(cfg.model, trainer=trainer) + trainer.fit(nmt_model) + if cfg.model.nemo_path: + nmt_model.save_to(cfg.model.nemo_path) + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml b/models/tts/NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml new file mode 100644 index 0000000000..8470991ea8 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text_classification/conf/text_classification_config.yaml @@ -0,0 +1,119 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# Config file for text classification with pre-trained BERT models + +trainer: + gpus: 1 # number of GPUs (0 for CPU), or list of the GPUs to use e.g. [0, 1] + num_nodes: 1 + max_epochs: 100 + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 32 # Should be set to 16 for O1 and O2 to enable the AMP. + accelerator: ddp + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + num_sanity_val_steps: 0 # number of steps to perform validation steps for sanity check the validation process before starting the training, setting to 0 disables it + + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + +model: + nemo_path: text_classification_model.nemo # filename to save the model and associated artifacts to .nemo file + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null + + language_model: + pretrained_model_name: bert-base-uncased + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null + + classifier_head: + num_output_layers: 2 + fc_dropout: 0.1 + + class_labels: + class_labels_file : null # optional to specify a file containing the list of the labels + + dataset: + num_classes: ??? # The number of classes. 0 < Label ", "") + outfiles[mode].write(f'{review}\t{label}\n') + + for mode in modes: + outfiles[mode].close() + + class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w') + class_labels_file.write('negative\npositive\n') + class_labels_file.close() + + +def process_sst2(infold, outfold, uncased, splits=['train', 'dev']): + """Process sst2 dataset.""" + # "test" split doesn't have labels, so it is skipped + if not os.path.exists(infold): + link = 'https://dl.fbaipublicfiles.com/glue/data/SST-2.zip' + raise ValueError( + f'Data not found at {infold}. Please download SST-2 dataset from `{link}` and ' + f'extract it into the folder specified by `source_data_dir` argument.' + ) + + logging.info(f'Processing SST-2 dataset') + os.makedirs(outfold, exist_ok=True) + + def _read_tsv(input_file, quotechar=None): + """Read a tab separated value file.""" + with open(input_file, "r") as f: + reader = csv.reader(f, delimiter="\t", quotechar=quotechar) + lines = [] + for line in reader: + lines.append(line) + return lines + + for split in splits: + # Load input file. + input_file = os.path.join(infold, split + '.tsv') + lines = _read_tsv(input_file) + # Create output. + outfile = open(os.path.join(outfold, split + '.tsv'), 'w') + + # Copy lines, skip the header (line 0). + for line in lines[1:]: + text = line[0] + label = line[1] + # Lowercase when required. + if uncased: + text = text.lower() + # Write output. + outfile.write(f'{text}\t{label}\n') + # Close file. + outfile.close() + + class_labels_file = open(os.path.join(outfold, 'label_ids.tsv'), 'w') + class_labels_file.write('negative\npositive\n') + class_labels_file.close() + + logging.info(f'Result stored at {outfold}') + + +def process_chemprot(source_dir, target_dir, uncased, modes=['train', 'test', 'dev']): + if not os.path.exists(source_dir): + link = 'https://github.com/arwhirang/recursive_chemprot/tree/master/Demo/tree_LSTM/data' + raise ValueError(f'Data not found at {source_dir}. ' f'Please download ChemProt from {link}.') + + logging.info(f'Processing Chemprot dataset and store at {target_dir}') + os.makedirs(target_dir, exist_ok=True) + + naming_map = {'train': 'trainingPosit_chem', 'test': 'testPosit_chem', 'dev': 'developPosit_chem'} + + def _read_tsv(input_file, quotechar=None): + """Reads a tab separated value file.""" + with open(input_file, "r") as f: + reader = csv.reader(f, delimiter="\t", quotechar=quotechar) + lines = [] + for line in reader: + lines.append(line) + return lines + + outfiles = {} + label_mapping = {} + out_label_mapping = open(os.path.join(target_dir, 'label_mapping.tsv'), 'w') + for mode in modes: + outfiles[mode] = open(os.path.join(target_dir, mode + '.tsv'), 'w') + input_file = os.path.join(source_dir, naming_map[mode]) + lines = _read_tsv(input_file) + for line in lines: + text = line[1] + label = line[2] + if label == "True": + label = line[3] + if uncased: + text = text.lower() + if label not in label_mapping: + out_label_mapping.write(f'{label}\t{len(label_mapping)}\n') + label_mapping[label] = len(label_mapping) + label = label_mapping[label] + outfiles[mode].write(f'{text}\t{label}\n') + for mode in modes: + outfiles[mode].close() + out_label_mapping.close() + + +def process_thucnews(infold, outfold): + modes = ['train', 'test'] + train_size = 0.8 + if not os.path.exists(infold): + link = 'thuctc.thunlp.org/' + raise ValueError(f'Data not found at {infold}. ' f'Please download THUCNews from {link}.') + + logging.info(f'Processing THUCNews dataset and store at {outfold}') + os.makedirs(outfold, exist_ok=True) + + outfiles = {} + for mode in modes: + outfiles[mode] = open(os.path.join(outfold, mode + '.tsv'), 'a+', encoding='utf-8') + categories = ['体育', '娱乐', '家居', '彩票', '房产', '教育', '时尚', '时政', '星座', '游戏', '社会', '科技', '股票', '财经'] + for category in categories: + label = categories.index(category) + category_files = glob.glob(f'{infold}/{category}/*.txt') + test_num = int(len(category_files) * (1 - train_size)) + test_files = category_files[:test_num] + train_files = category_files[test_num:] + + for mode in modes: + logging.info(f'Processing {mode} data of the category {category}') + if mode == 'test': + files = test_files + else: + files = train_files + + if len(files) == 0: + logging.info(f'Skipping category {category} for {mode} mode') + continue + + for file in tqdm.tqdm(files): + with open(file, 'r', encoding='utf-8') as f: + news = f.read().strip().replace('\r', '') + news = news.replace('\n', '').replace('\t', ' ') + outfiles[mode].write(f'{news}\t{label}\n') + for mode in modes: + outfiles[mode].close() + + +if __name__ == "__main__": + # Parse the command-line arguments. + parser = argparse.ArgumentParser(description="Process and convert datasets into NeMo\'s format.") + parser.add_argument("--dataset_name", required=True, type=str, choices=['imdb', 'thucnews', 'chemprot']) + parser.add_argument( + "--source_data_dir", required=True, type=str, help='The path to the folder containing the dataset files.' + ) + parser.add_argument("--target_data_dir", required=True, type=str) + parser.add_argument("--do_lower_case", action='store_true') + args = parser.parse_args() + + dataset_name = args.dataset_name + do_lower_case = args.do_lower_case + source_dir = args.source_data_dir + target_dir = args.target_data_dir + + if not exists(source_dir): + raise FileNotFoundError(f"{source_dir} does not exist.") + + if dataset_name == 'imdb': + process_imdb(source_dir, target_dir, do_lower_case) + elif dataset_name == 'thucnews': + process_thucnews(source_dir, target_dir) + elif dataset_name == "chemprot": + process_chemprot(source_dir, target_dir, do_lower_case) + elif dataset_name == "sst-2": + process_sst2(source_dir, target_dir, do_lower_case) + else: + raise ValueError( + f'Dataset {dataset_name} is not supported.' + + "Please make sure that you build the preprocessing process for it. " + + "NeMo's format assumes that a data file has a header and each line of the file follows " + + "the format: text [TAB] label. Label is assumed to be an integer." + ) diff --git a/models/tts/NeMo/examples/nlp/text_classification/model_parallel_text_classification_evaluation.py b/models/tts/NeMo/examples/nlp/text_classification/model_parallel_text_classification_evaluation.py new file mode 100644 index 0000000000..255a27ee79 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text_classification/model_parallel_text_classification_evaluation.py @@ -0,0 +1,41 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +""" +This script runs model parallel text classification evaluation. +""" +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.text_classification import TextClassificationModel +from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="text_classification_config") +def main(cfg: DictConfig) -> None: + logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + # TODO: can we drop strict=False + model = TextClassificationModel.restore_from(cfg.model.nemo_path, trainer=trainer, strict=False) + model.setup_test_data(test_data_config=cfg.model.test_ds) + + trainer.test(model=model, ckpt_path=None) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/text_classification/text_classification_with_bert.py b/models/tts/NeMo/examples/nlp/text_classification/text_classification_with_bert.py new file mode 100644 index 0000000000..1cad923381 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/text_classification/text_classification_with_bert.py @@ -0,0 +1,154 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script contains an example on how to train, evaluate and perform inference with the TextClassificationModel. +TextClassificationModel in NeMo supports text classification problems such as sentiment analysis or +domain/intent detection for dialogue systems, as long as the data follows the format specified below. + +***Data format*** +TextClassificationModel requires the data to be stored in TAB separated files (.tsv) with two columns of sentence and +label. Each line of the data file contains text sequences, where words are separated with spaces and label separated +with [TAB], i.e.: + +[WORD][SPACE][WORD][SPACE][WORD][TAB][LABEL] + +For example: + +hide new secretions from the parental units[TAB]0 +that loves its characters and communicates something rather beautiful about human nature[TAB]1 +... + +If your dataset is stored in another format, you need to convert it to this format to use the TextClassificationModel. + + +***Setting the configs*** +The model and the PT trainer are defined in a config file which declares multiple important sections. +The most important ones are: + model: All arguments that are related to the Model - language model, tokenizer, head classifier, optimizer, + schedulers, and datasets/data loaders. + trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs, + precision level, etc. + +This script uses the `/examples/nlp/text_classification/conf/text_classification_config.yaml` default config file +by default. You may update the config file from the file directly or by using the command line arguments. +Other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +You first need to set the num_classes in the config file which specifies the number of classes in the dataset. +Notice that some config lines, including `model.dataset.classes_num`, have `???` as their value, this means that values +for these fields are required to be specified by the user. We need to specify and set the `model.train_ds.file_name`, +`model.validation_ds.file_name`, and `model.test_ds.file_name` in the config file to the paths of the train, validation, + and test files if they exist. We may do it by updating the config file or by setting them from the command line. + + +***How to run the script?*** +For example the following would train a model for 50 epochs in 2 GPUs on a classification task with 2 classes: + +# python text_classification_with_bert.py + model.dataset.num_classes=2 + model.train_ds=PATH_TO_TRAIN_FILE + model.validation_ds=PATH_TO_VAL_FILE + trainer.max_epochs=50 + trainer.gpus=2 + +This script would also reload the last checkpoint after the training is done and does evaluation on the dev set, +then performs inference on some sample queries. + +By default, this script uses examples/nlp/text_classification/conf/text_classifciation_config.py config file, and +you may update all the params in the config file from the command line. You may also use another config file like this: + +# python text_classification_with_bert.py --config-name==PATH_TO_CONFIG_FILE + model.dataset.num_classes=2 + model.train_ds=PATH_TO_TRAIN_FILE + model.validation_ds=PATH_TO_VAL_FILE + trainer.max_epochs=50 + trainer.gpus=2 + +***Load a saved model*** +This script would save the model after training into '.nemo' checkpoint file specified by nemo_path of the model config. +You may restore the saved model like this: + model = TextClassificationModel.restore_from(restore_path=NEMO_FILE_PATH) + +***Evaluation a saved model on another dataset*** +# If you wanted to evaluate the saved model on another dataset, you may restore the model and create a new data loader: + eval_model = TextClassificationModel.restore_from(restore_path=checkpoint_path) + +# Then, you may create a dataloader config for evaluation: + eval_config = OmegaConf.create( + {'file_path': cfg.model.test_ds.file_path, 'batch_size': 64, 'shuffle': False, 'num_workers': 3} + ) + eval_model.setup_test_data(test_data_config=eval_config) + +# You need to create a new trainer: + eval_trainer = pl.Trainer(gpus=1) + eval_model.set_trainer(eval_trainer) + eval_trainer.test(model=eval_model, verbose=False) +""" +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models.text_classification import TextClassificationModel +from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="text_classification_config") +def main(cfg: DictConfig) -> None: + logging.info(f'\nConfig Params:\n{OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + if not cfg.model.train_ds.file_path: + raise ValueError("'train_ds.file_path' need to be set for the training!") + + model = TextClassificationModel(cfg.model, trainer=trainer) + logging.info("===========================================================================================") + logging.info('Starting training...') + trainer.fit(model) + logging.info('Training finished!') + logging.info("===========================================================================================") + + if cfg.model.nemo_path: + # '.nemo' file contains the last checkpoint and the params to initialize the model + model.save_to(cfg.model.nemo_path) + logging.info(f'Model is saved into `.nemo` file: {cfg.model.nemo_path}') + + # We evaluate the trained model on the test set if test_ds is set in the config file + if cfg.model.test_ds.file_path: + logging.info("===========================================================================================") + logging.info("Starting the testing of the trained model on test set...") + trainer.test(model=model, ckpt_path=None, verbose=False) + logging.info("Testing finished!") + logging.info("===========================================================================================") + + # perform inference on a list of queries. + if "infer_samples" in cfg.model and cfg.model.infer_samples: + logging.info("===========================================================================================") + logging.info("Starting the inference on some sample queries...") + + # max_seq_length=512 is the maximum length BERT supports. + results = model.classifytext(queries=cfg.model.infer_samples, batch_size=16, max_seq_length=512) + logging.info('The prediction results of some sample queries with the trained model:') + for query, result in zip(cfg.model.infer_samples, results): + logging.info(f'Query : {query}') + logging.info(f'Predicted label: {result}') + + logging.info("Inference finished!") + logging.info("===========================================================================================") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml b/models/tts/NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml new file mode 100644 index 0000000000..0d27907751 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml @@ -0,0 +1,133 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# Punctuation and capitalization model with pretrained BERT-like models + +pretrained_model: null # pretrained Punctuation and Capitalization model from list_available_models(), for example: +# punctuation_en_bert or punctuation_en_distilbert +# or your_model.nemo +trainer: + gpus: 1 # the number of gpus, 0 for CPU + num_nodes: 1 + max_epochs: 3 + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0 + accelerator: ddp + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + +exp_manager: + exp_dir: null # exp_dir for your experiment, if None, defaults to "./nemo_experiments" + name: Punctuation_and_Capitalization # The name of your model + create_tensorboard_logger: true # Whether you want exp_manger to create a tb logger + create_checkpoint_callback: true # Whether you want exp_manager to create a model checkpoint callback + +model: + punct_label_ids: null # will be filled during training + capit_label_ids: null # will be filled during training + class_labels: + punct_labels_file: punct_label_ids.csv + capit_labels_file: capit_label_ids.csv + dataset: + data_dir: ??? # /path/to/training/data + # parameters for dataset preprocessing + max_seq_length: 128 + pad_label: 'O' + ignore_extra_tokens: false + ignore_start_end: false + use_cache: true + # shared among dataloaders + num_workers: 2 + pin_memory: false + drop_last: false + + train_ds: + text_file: text_train.txt + labels_file: labels_train.txt + shuffle: true + num_samples: -1 + batch_size: 64 + + validation_ds: + # if evaluation data is not in the model.dataset.data_dir as the training data or multiple datasets are used for + # evaluation is needed, specify ds_item, otherwise by default data_dir is used + ds_item: null # expected format: [PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets) + text_file: text_dev.txt + labels_file: labels_dev.txt + shuffle: false + num_samples: -1 + batch_size: 64 + + test_ds: + # if evaluation data is not in the model.dataset.data_dir as the training data or multiple datasets are used for + # evaluation is needed, specify ds_item, otherwise by default data_dir is used + ds_item: null # expected format: [PATH_TO_DEV1,PATH_TO_DEV2] (Note no space between the paths and square brackets) + text_file: text_dev.txt + labels_file: labels_dev.txt + shuffle: false + num_samples: -1 + batch_size: 64 + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null + + language_model: + pretrained_model_name: bert-base-uncased + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null + + punct_head: + punct_num_fc_layers: 1 + fc_dropout: 0.1 + activation: 'relu' + use_transformer_init: True + + capit_head: + capit_num_fc_layers: 1 + fc_dropout: 0.1 + activation: 'relu' + use_transformer_init: true + + optim: + name: adam + lr: 1e-4 + weight_decay: 0.00 + + sched: + name: WarmupAnnealing + # Scheduler params + warmup_steps: null + warmup_ratio: 0.1 + last_epoch: -1 + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/nlp/token_classification/conf/token_classification_config.yaml b/models/tts/NeMo/examples/nlp/token_classification/conf/token_classification_config.yaml new file mode 100644 index 0000000000..8cb162a60d --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/conf/token_classification_config.yaml @@ -0,0 +1,119 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# Token Classification tasks (for example, Named Entity Recognition) with pretrained BERT-like models + +pretrained_model: null # pretrained TokenClassification model from list_available_models() or path to a .nemo file, +# for example: ner_en_bert or your_model.nemo +trainer: + gpus: 1 # the number of gpus, 0 for CPU + num_nodes: 1 + max_epochs: 5 + max_steps: null # precedence over max_epochs + accumulate_grad_batches: 1 # accumulates grads every k batches + gradient_clip_val: 0.0 + amp_level: O0 # O1/O2 for mixed precision + precision: 16 # Should be set to 16 for O1 and O2, default is 16 as PT ignores it when am_level is O0 + accelerator: ddp + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + resume_from_checkpoint: null # The path to a checkpoint file to continue the training, restores the whole state including the epoch, step, LR schedulers, apex, etc. + +exp_manager: + exp_dir: null # exp_dir for your experiment, if None, defaults to "./nemo_experiments" + name: token_classification_model # The name of your model + create_tensorboard_logger: true # Whether you want exp_manager to create a tb logger + create_checkpoint_callback: true # Whether you want exp_manager to create a model checkpoint callback + +model: + label_ids: null # will be filled during training + class_labels: + class_labels_file: label_ids.csv # will be generated during training and saved in .nemo file + dataset: + data_dir: ??? # /path/to/data + class_balancing: null # choose from [null, weighted_loss]. Weighted_loss enables the weighted class balancing of the loss, may be used for handling unbalanced classes + max_seq_length: 128 + pad_label: 'O' + ignore_extra_tokens: false + ignore_start_end: false + use_cache: true + # shared among dataloaders + num_workers: 2 + pin_memory: false + drop_last: false + + train_ds: + text_file: text_train.txt + labels_file: labels_train.txt + shuffle: true + num_samples: -1 + batch_size: 64 + + validation_ds: + text_file: text_dev.txt + labels_file: labels_dev.txt + shuffle: false + num_samples: -1 + batch_size: 64 + + test_ds: + text_file: text_dev.txt + labels_file: labels_dev.txt + shuffle: false + num_samples: -1 + batch_size: 64 + + tokenizer: + tokenizer_name: ${model.language_model.pretrained_model_name} # or sentencepiece + vocab_file: null # path to vocab file + tokenizer_model: null # only used if tokenizer is sentencepiece + special_tokens: null + + language_model: + pretrained_model_name: bert-base-uncased + lm_checkpoint: null + config_file: null # json file, precedence over config + config: null + + + head: + num_fc_layers: 2 + fc_dropout: 0.5 + activation: 'relu' + use_transformer_init: True + + optim: + name: adam + lr: 5e-5 + weight_decay: 0.00 + + sched: + name: WarmupAnnealing + # Scheduler params + warmup_steps: null + warmup_ratio: 0.1 + last_epoch: -1 + + # pytorch lightning args + monitor: val_loss + reduce_on_plateau: false + +hydra: + run: + dir: . + job_logging: + root: + handlers: null diff --git a/models/tts/NeMo/examples/nlp/token_classification/data/get_tatoeba_data.py b/models/tts/NeMo/examples/nlp/token_classification/data/get_tatoeba_data.py new file mode 100644 index 0000000000..d895e81543 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/data/get_tatoeba_data.py @@ -0,0 +1,242 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 argparse +import logging +import os +import random +import re +import string +import subprocess + +from nemo.utils import logging + +URL = {'tatoeba': 'https://downloads.tatoeba.org/exports/sentences.csv'} + + +def __maybe_download_file(destination: str, source: str): + """ + Downloads source to destination if not exists. + If exists, skips download + Args: + destination: local filepath + source: url of resource + """ + source = URL[source] + if not os.path.exists(destination): + logging.info(f'Downloading {source} to {destination}') + subprocess.run(['wget', '-O', destination, source]) + else: + logging.info(f'{destination} found. Skipping download') + + +def __process_english_sentences( + in_file: str, out_file: str, percent_to_cut: float = 0, num_to_combine: int = 1, num_samples: int = -1 +): + """ + Extract English sentences from the Tatoeba dataset. + + Expected in_file format + that + contrain letters and punctuation marks (,.?). + Chop and combine sentences. + Args: + in_file: local filepath to the tatoeba dataset. + Format: id [TAB] region_name [TAB] sentence, + for example: "1276\teng\tLet's try something.\n" + out_file: local filepath to the clean dataset + percent_to_cut: Percent of sentences to cut in the middle + to get examples of incomplete sentences. + This could be useful since ASR output not always + represents a complete sentence + num_to_combine: Number of sentences to combine into + a single example + num_samples: Number of samples in the final dataset + """ + if not os.path.exists(in_file): + raise FileNotFoundError(f'{in_file} not found.') + + in_file = open(in_file, 'r') + out_file = open(out_file, 'w') + lines_to_combine = [] + samples_count = 0 + + for line in in_file: + line = line.split('\t') + # use only English sentences + if line[1] == 'eng': + line = line[2].strip() + if re.match("^[A-Z][A-Za-z.,'?\s]+$", line): # nopep8 + # chop some sentences in the middle + if percent_to_cut > 0: + line = line.split() + if random.random() < percent_to_cut: + line = line[: len(line) // 2] + line = ' '.join(line) + + # combine multiple sentences into a single example + # to make it harder for the model to learn eos punctuation + if len(lines_to_combine) >= num_to_combine: + if samples_count == num_samples: + return + out_file.write(' '.join(lines_to_combine) + '\n') + lines_to_combine = [] + samples_count += 1 + lines_to_combine.append(line) + + if len(lines_to_combine) > 0 and (samples_count < num_samples or num_samples < 0): + out_file.write(' '.join(lines_to_combine) + '\n') + + +def __split_into_train_dev(in_file: str, train_file: str, dev_file: str, percent_dev: float): + """ + Create train and dev split of the dataset. + Args: + in_file: local filepath to the dataset + train_file: local filepath to the train dataset + dev_file: local filepath to the dev dataset + percent_dev: Percent of the sentences in the dev set + """ + if not os.path.exists(in_file): + raise FileNotFoundError(f'{in_file} not found.') + + lines = open(in_file, 'r').readlines() + train_file = open(train_file, 'w') + dev_file = open(dev_file, 'w') + + dev_size = int(len(lines) * percent_dev) + train_file.write(' '.join(lines[:-dev_size])) + dev_file.write(' '.join(lines[-dev_size:])) + + +def remove_punctuation(word: str): + """ + Removes all punctuation marks from a word except for ' + that is often a part of word: don't, it's, and so on + """ + all_punct_marks = string.punctuation.replace("'", '') + return re.sub('[' + all_punct_marks + ']', '', word) + + +def create_text_and_labels(output_dir: str, file_path: str, punct_marks: str = ',.?'): + """ + Create datasets for training and evaluation. + + Args: + output_dir: path to the output data directory + file_path: path to file name + punct_marks: supported punctuation marks + + The data will be split into 2 files: text.txt and labels.txt. \ + Each line of the text.txt file contains text sequences, where words\ + are separated with spaces. The labels.txt file contains \ + corresponding labels for each word in text.txt, the labels are \ + separated with spaces. Each line of the files should follow the \ + format: \ + [WORD] [SPACE] [WORD] [SPACE] [WORD] (for text.txt) and \ + [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt).' + """ + if not os.path.exists(file_path): + raise ValueError(f'{file_path} not found') + + os.makedirs(output_dir, exist_ok=True) + + base_name = os.path.basename(file_path) + labels_file = os.path.join(output_dir, 'labels_' + base_name) + text_file = os.path.join(output_dir, 'text_' + base_name) + + with open(file_path, 'r') as f: + with open(text_file, 'w') as text_f: + with open(labels_file, 'w') as labels_f: + for line in f: + line = line.split() + text = '' + labels = '' + for word in line: + label = word[-1] if word[-1] in punct_marks else 'O' + word = remove_punctuation(word) + if len(word) > 0: + if word[0].isupper(): + label += 'U' + else: + label += 'O' + + word = word.lower() + text += word + ' ' + labels += label + ' ' + + text_f.write(text.strip() + '\n') + labels_f.write(labels.strip() + '\n') + + print(f'{text_file} and {labels_file} created from {file_path}.') + + +def __delete_file(file_to_del: str): + """ + Deletes the file + Args: + file_to_del: local filepath to the file to delete + """ + if os.path.exists(file_to_del): + os.remove(file_to_del) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Prepare tatoeba dataset') + parser.add_argument("--data_dir", required=True, type=str) + parser.add_argument("--dataset", default='tatoeba', type=str) + parser.add_argument("--num_samples", default=-1, type=int, help='-1 to use the whole dataset') + parser.add_argument("--percent_to_cut", default=0, type=float, help='Percent of sentences to cut in the middle') + parser.add_argument( + "--num_lines_to_combine", default=1, type=int, help='Number of lines to combine into single example' + ) + parser.add_argument("--percent_dev", default=0.2, type=float, help='Size of the dev set, float') + parser.add_argument("--clean_dir", action='store_true') + args = parser.parse_args() + + if not os.path.exists(args.data_dir): + os.makedirs(args.data_dir) + + if args.dataset != 'tatoeba': + raise ValueError("Unsupported dataset.") + + logging.info(f'Downloading tatoeba dataset') + tatoeba_dataset = os.path.join(args.data_dir, 'sentences.csv') + __maybe_download_file(tatoeba_dataset, args.dataset) + + logging.info(f'Processing English sentences...') + clean_eng_sentences = os.path.join(args.data_dir, 'clean_eng_sentences.txt') + __process_english_sentences( + tatoeba_dataset, clean_eng_sentences, args.percent_to_cut, args.num_lines_to_combine, args.num_samples + ) + + train_file = os.path.join(args.data_dir, 'train.txt') + dev_file = os.path.join(args.data_dir, 'dev.txt') + + logging.info( + f'Splitting the {args.dataset} dataset into train and dev sets' + ' and creating labels and text files' + ) + __split_into_train_dev(clean_eng_sentences, train_file, dev_file, args.percent_dev) + + logging.info(f'Creating text and label files for training') + create_text_and_labels(args.data_dir, os.path.join(args.data_dir, 'train.txt')) + create_text_and_labels(args.data_dir, os.path.join(args.data_dir, 'dev.txt')) + + if args.clean_dir: + logging.info(f'Cleaning up {args.data_dir}') + __delete_file(clean_eng_sentences) + __delete_file(tatoeba_dataset) + __delete_file(train_file) + __delete_file(dev_file) + logging.info(f'Processing of the {args.dataset} is complete') diff --git a/models/tts/NeMo/examples/nlp/token_classification/data/import_from_iob_format.py b/models/tts/NeMo/examples/nlp/token_classification/data/import_from_iob_format.py new file mode 100644 index 0000000000..4a6f15442b --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/data/import_from_iob_format.py @@ -0,0 +1,124 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 argparse +import os + +from nemo.utils import logging + + +def __convert_data(in_file: str, out_text_f: str, out_labels_f: str, max_length: int): + """ + Convert data from the IOB format to NeMo accepted format described below. + in_file should be in the IOB format, see example here: + https://www.clips.uantwerpen.be/conll2003/ner/. + + Args: + in_file: input file name + out_text_f: output file with text + out_labels_f: output file with labels + max_length: use -1 to leave the examples' length as is, otherwise long examples will be split into multiple + examples + After the conversion, the dataset is split into 2 files: text.txt + and labels.txt. + Each line of the text.txt file contains text sequences, where words + are separated with spaces. The labels.txt file contains corresponding + labels for each word in text.txt, the labels are separated with spaces. + Each line of the files should follow the format: + [WORD] [SPACE] [WORD] [SPACE] [WORD] (for text.txt) and + [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt). + + """ + in_file = open(in_file, 'r') + + if max_length == -1: + with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels: + for line in in_file: + if line == '\n': + out_text.write(line) + out_labels.write(line) + else: + line = line.split() + out_text.write(line[0] + ' ') + out_labels.write(line[-1] + ' ') + + else: + words = [] + labels = [] + with open(out_text_f, 'w') as out_text, open(out_labels_f, 'w') as out_labels: + lines = in_file.readlines() + for line_id, line in enumerate(lines): + logging.info(f"{line_id} {len(lines)}") + contends = line.strip() + if len(contends) == 0: + assert len(words) == len(labels) + if len(words) > max_length: + # split if the sentence is longer than max_length + while len(words) > max_length: + tmplabel = labels[:max_length] + for iidx in range(len(tmplabel)): + if tmplabel.pop() == 'O': + break + l = ' '.join([label for label in labels[: len(tmplabel) + 1] if len(label) > 0]) + w = ' '.join([word for word in words[: len(tmplabel) + 1] if len(word) > 0]) + + out_text.write(w + "\n") + out_labels.write(l + "\n") + words = words[len(tmplabel) + 1 :] + labels = labels[len(tmplabel) + 1 :] + + if len(words) == 0: + continue + l = ' '.join([label for label in labels if len(label) > 0]) + w = ' '.join([word for word in words if len(word) > 0]) + + out_text.write(w + "\n") + out_labels.write(l + "\n") + words = [] + labels = [] + continue + + word = line.strip().split()[0] + label = line.strip().split()[-1] + words.append(word) + labels.append(label) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description='Convert data from IOB format to the format compatible with \ + nlp/examples/token_classification/scripts/token_classification_train.py and \ + token_classification_evaluate.py' + ) + parser.add_argument("--data_file", required=True, type=str, help='path to a file in IOB format') + parser.add_argument( + "--max_length", + default=-1, + type=int, + help='use -1 to leave the examples\'s length as is, ' + 'otherwise long examples will be split into multiple examples', + ) + args = parser.parse_args() + + data_dir, basename = os.path.split(args.data_file) + prefix = os.path.splitext(basename)[0] + if not os.path.exists(args.data_file): + raise FileNotFoundError(f"{args.data_file} not found") + + logging.info(f'Processing {args.data_file}') + out_text = os.path.join(data_dir, 'text_' + prefix + '.txt') + out_labels = os.path.join(data_dir, 'labels_' + prefix + '.txt') + + __convert_data(args.data_file, out_text, out_labels, args.max_length) + logging.info(f'Processing of the {args.data_file} is complete') diff --git a/models/tts/NeMo/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py b/models/tts/NeMo/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py new file mode 100644 index 0000000000..f31e6fc391 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py @@ -0,0 +1,99 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +""" +The script converts raw text to the NeMo format for punctuation and capitalization task. + +Raw Data Format +--------------- + +The Punctuation and Capitalization model can work with any text dataset, although it is recommended to balance the data, especially for the punctuation task. +Before pre-processing the data to the format expected by the model, the data should be split into train.txt and dev.txt (and optionally test.txt). +Each line in the **train.txt/dev.txt/test.txt** should represent one or more full and/or truncated sentences. + +Example of the train.txt/dev.txt file: + When is the next flight to New York? + The next flight is ... + .... + + +The `source_data_dir` structure should look like this: + . + |--sourced_data_dir + |-- dev.txt + |-- train.txt + + + +NeMo Data Format for training the model +--------------------------------------- + +The punctuation and capitalization model expects the data in the following format: + +The training and evaluation data is divided into 2 files: text.txt and labels.txt. \ +Each line of the **text.txt** file contains text sequences, where words are separated with spaces, i.e. + +[WORD] [SPACE] [WORD] [SPACE] [WORD], for example: + when is the next flight to new york + the next flight is ... + ... + +The **labels.txt** file contains corresponding labels for each word in text.txt, the labels are separated with spaces. \ +Each label in labels.txt file consists of 2 symbols: + +* the first symbol of the label indicates what punctuation mark should follow the word (where O means no punctuation needed); +* the second symbol determines if a word needs to be capitalized or not (where U indicates that the word should be upper cased, and O - no capitalization needed.) + +By default the following punctuation marks are considered: commas, periods, and question marks; the rest punctuation marks were removed from the data. +This can be changed by introducing new labels in the labels.txt files + +Each line of the labels.txt should follow the format: [LABEL] [SPACE] [LABEL] [SPACE] [LABEL] (for labels.txt). \ +For example, labels for the above text.txt file should be: + + OU OO OO OO OO OO OU ?U + OU OO OO OO ... + ... + +The complete list of all possible labels for this task used in this tutorial is: OO, ,O, .O, ?O, OU, ,U, .U, ?U. + +Converting Raw data to NeMo format +---------------------------------- + +To pre-process the raw text data, stored under :code:`sourced_data_dir` (see the :ref:`raw_data_format_punct` +section), run the following command: + + python examples/nlp/token_classification/data/prepare_data_for_punctuation_capitalization.py \ + -s + -o + +""" + +import argparse +import os + +from get_tatoeba_data import create_text_and_labels + +if __name__ == "__main__": + parser = argparse.ArgumentParser(description='Prepare data for punctuation and capitalization tasks') + parser.add_argument("-s", "--source_file", required=True, type=str, help="Path to the source file") + parser.add_argument("-o", "--output_dir", required=True, type=str, help="Path to the output directory") + args = parser.parse_args() + + if not os.path.exists(args.source_file): + raise ValueError(f'{args.source_file} was not found') + + os.makedirs(args.output_dir, exist_ok=True) + create_text_and_labels(args.output_dir, args.source_file) + + print(f'Processing of the {args.source_file} is complete') diff --git a/models/tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_evaluate.py b/models/tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_evaluate.py new file mode 100644 index 0000000000..b91b332a87 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_evaluate.py @@ -0,0 +1,127 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +import pytorch_lightning as pl +from omegaconf import DictConfig + +from nemo.collections.nlp.models import PunctuationCapitalizationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +""" +This script shows how to perform evaluation and runs inference of a few examples. + +More details on the task and data format could be found in tutorials/nlp/Punctuation_and_Capitalization.ipynb + +*** Setting the configs *** + +This script uses the `/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml` config file +by default. You may update the config file from the file directly. +The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb + + +*** Model Evaluation *** + + python punctuation_capitalization_evaluate.py \ + model.dataset.data_dir= \ + pretrained_model=punctuation_en_distilbert + + - a directory that contains test_ds.text_file and test_ds.labels_file (see the config) +pretrained_model - pretrained PunctuationCapitalizationModel model from list_available_models() or + path to a .nemo file, for example: punctuation_en_bert or your_model.nemo + +""" + + +@hydra_runner(config_path="conf", config_name="punctuation_capitalization_config") +def main(cfg: DictConfig) -> None: + logging.info( + 'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU and \ + no DDP to obtain accurate results' + ) + + if not hasattr(cfg.model, 'test_ds'): + raise ValueError(f'model.test_ds was not found in the config, skipping evaluation') + else: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + + trainer = pl.Trainer( + gpus=gpu, + precision=cfg.trainer.precision, + amp_level=cfg.trainer.amp_level, + logger=False, + checkpoint_callback=False, + ) + exp_dir = exp_manager(trainer, cfg.exp_manager) + + if not cfg.pretrained_model: + raise ValueError( + 'To run evaluation and inference script a pre-trained model or .nemo file must be provided.' + f'Choose from {PunctuationCapitalizationModel.list_available_models()} or "pretrained_model"="your_model.nemo"' + ) + + if os.path.exists(cfg.pretrained_model): + model = PunctuationCapitalizationModel.restore_from(cfg.pretrained_model) + elif cfg.pretrained_model in PunctuationCapitalizationModel.get_available_model_names(): + model = PunctuationCapitalizationModel.from_pretrained(cfg.pretrained_model) + else: + raise ValueError( + f'Provide path to the pre-trained .nemo file or choose from {PunctuationCapitalizationModel.list_available_models()}' + ) + + data_dir = cfg.model.dataset.get('data_dir', None) + + if data_dir is None: + logging.error( + 'No dataset directory provided. Skipping evaluation. ' + 'To run evaluation on a file, specify path to the directory that contains test_ds.text_file and test_ds.labels_file with "model.dataset.data_dir" argument.' + ) + elif not os.path.exists(data_dir): + logging.error(f'{data_dir} is not found, skipping evaluation on the test set.') + else: + model.update_data_dir(data_dir=data_dir) + model._cfg.dataset = cfg.model.dataset + + if not hasattr(cfg.model, 'test_ds'): + logging.error(f'model.test_ds was not found in the config, skipping evaluation') + elif model.prepare_test(trainer): + model.setup_test_data(cfg.model.test_ds) + trainer.test(model) + else: + logging.error('Skipping the evaluation. The trainer is not setup properly.') + + # run an inference on a few examples + queries = [ + 'we bought four shirts one pen and a mug from the nvidia gear store in santa clara', + 'what can i do for you today', + 'how are you', + ] + + inference_results = model.add_punctuation_capitalization(queries, batch_size=len(queries), max_seq_length=512) + + for query, result in zip(queries, inference_results): + logging.info(f'Query : {query}') + logging.info(f'Result: {result.strip()}\n') + + logging.info(f'Results are saved at {exp_dir}') + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_train.py b/models/tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_train.py new file mode 100644 index 0000000000..ae55448d3d --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/punctuation_capitalization_train.py @@ -0,0 +1,101 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import PunctuationCapitalizationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +""" +This script show how to train a Punctuation and Capitalization Model. +More details on the task and data format could be found in tutorials/nlp/Punctuation_and_Capitalization.ipynb + +*** Setting the configs *** + +The model and the PT trainer are defined in a config file which declares multiple important sections. +The most important ones are: + model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer, + schedulers, and datasets/data loaders. + trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs, + precision level, etc. +This script uses the `/examples/nlp/token_classification/conf/punctuation_capitalization_config.yaml` config file +by default. You may update the config file from the file directly. +The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb + +*** Model training *** + +To run this script and train the model from scratch, use: + python punctuation_and_capitalization_train.py \ + model.dataset.data_dir= + +To use one of the pretrained versions of the model and finetune it, run: + python punctuation_and_capitalization.py \ + pretrained_model=punctuation_en_bert \ + model.dataset.data_dir= + + - a directory that contains test_ds.text_file and test_ds.labels_file (see the config) + pretrained_model - pretrained PunctuationCapitalization model from list_available_models() or + path to a .nemo file, for example: punctuation_en_bert or model.nemo + +""" + + +@hydra_runner(config_path="conf", config_name="punctuation_capitalization_config") +def main(cfg: DictConfig) -> None: + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + if not cfg.pretrained_model: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + model = PunctuationCapitalizationModel(cfg.model, trainer=trainer) + else: + if os.path.exists(cfg.pretrained_model): + model = PunctuationCapitalizationModel.restore_from(cfg.pretrained_model) + elif cfg.pretrained_model in PunctuationCapitalizationModel.get_available_model_names(): + model = PunctuationCapitalizationModel.from_pretrained(cfg.pretrained_model) + else: + raise ValueError( + f'Provide path to the pre-trained .nemo file or choose from {PunctuationCapitalizationModel.list_available_models()}' + ) + + data_dir = cfg.model.dataset.get('data_dir', None) + if data_dir: + if not os.path.exists(data_dir): + raise ValueError(f'{data_dir} is not found at') + + # we can also do finetuning of the pretrained model but we would need to update the data dir + model.update_data_dir(data_dir) + # setup train and validation Pytorch DataLoaders + model.setup_training_data() + model.setup_validation_data() + logging.info(f'Using config file of the pretrained model') + else: + raise ValueError( + 'Specify a valid dataset directory that contains test_ds.text_file and test_ds.labels_file \ + with "model.dataset.data_dir" argument' + ) + + trainer.fit(model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/token_classification/token_classification_evaluate.py b/models/tts/NeMo/examples/nlp/token_classification/token_classification_evaluate.py new file mode 100644 index 0000000000..4a888f50d7 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/token_classification_evaluate.py @@ -0,0 +1,137 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +import pytorch_lightning as pl +from omegaconf import DictConfig + +from nemo.collections.nlp.models import TokenClassificationModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +""" +This script shows how to perform evaluation and runs inference of a few examples. + +More details on Token Classification model could be found in tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb + +*** Setting the configs *** + +This script uses the `/examples/nlp/token_classification/conf/token_classification_config.yaml` config file +by default. You may update the config file from the file directly. +The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb + +*** Model Evaluation *** + +The script runs two types of evaluation: + * model.test() - this eval will use the config setting for evaluation such as model.dataset.max_seq_length + * model.evaluate_from_file(): + * disregards model.dataset.max_seq_length and evaluates all the tokens, BERT max seq length - 512 tokens after tokenization + * creates confusion matrix + * saves predictions and labels (if provided) + +To run the script: + + python token_classification_evaluate.py \ + model.dataset.data_dir= \ + pretrained_model=ner_en_bert + + - a directory that contains test_ds.text_file and test_ds.labels_file (see the config) +pretrained_model - pretrained TokenClassification model from list_available_models() or + path to a .nemo file, for example: ner_en_bert or your_model.nemo + +""" + + +@hydra_runner(config_path="conf", config_name="token_classification_config") +def main(cfg: DictConfig) -> None: + logging.info( + 'During evaluation/testing, it is currently advisable to construct a new Trainer with single GPU and \ + no DDP to obtain accurate results' + ) + + if not hasattr(cfg.model, 'test_ds'): + raise ValueError(f'model.test_ds was not found in the config, skipping evaluation') + else: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + + trainer = pl.Trainer( + gpus=gpu, + precision=cfg.trainer.precision, + amp_level=cfg.trainer.amp_level, + logger=False, + checkpoint_callback=False, + ) + exp_dir = exp_manager(trainer, cfg.exp_manager) + + if not cfg.pretrained_model: + raise ValueError( + 'To run evaluation and inference script a pre-trained model or .nemo file must be provided.' + f'Choose from {TokenClassificationModel.list_available_models()} or "pretrained_model"="your_model.nemo"' + ) + + if os.path.exists(cfg.pretrained_model): + model = TokenClassificationModel.restore_from(cfg.pretrained_model) + elif cfg.pretrained_model in TokenClassificationModel.get_available_model_names(): + model = TokenClassificationModel.from_pretrained(cfg.pretrained_model) + else: + raise ValueError( + f'Provide path to the pre-trained .nemo checkpoint or choose from {TokenClassificationModel.list_available_models()}' + ) + + data_dir = cfg.model.dataset.get('data_dir', None) + if data_dir is None: + logging.error( + 'No dataset directory provided. Skipping evaluation. ' + 'To run evaluation on a file, specify path to the directory that contains test_ds.text_file and test_ds.labels_file with "model.dataset.data_dir" argument.' + ) + elif not os.path.exists(data_dir): + logging.error(f'{data_dir} is not found, skipping evaluation on the test set.') + else: + model.update_data_dir(data_dir=data_dir) + model._cfg.dataset = cfg.model.dataset + + if not hasattr(cfg.model, 'test_ds'): + logging.error(f'model.test_ds was not found in the config, skipping evaluation') + elif model.prepare_test(trainer): + model.setup_test_data(cfg.model.test_ds) + trainer.test(model) + + model.evaluate_from_file( + text_file=os.path.join(data_dir, cfg.model.test_ds.text_file), + labels_file=os.path.join(data_dir, cfg.model.test_ds.labels_file), + output_dir=exp_dir, + add_confusion_matrix=True, + normalize_confusion_matrix=True, + ) + else: + logging.error('Skipping the evaluation. The trainer is not setup properly.') + + # run an inference on a few examples + queries = ['we bought four shirts from the nvidia gear store in santa clara.', 'Nvidia is a company.'] + results = model.add_predictions(queries, output_file='predictions.txt') + + for query, result in zip(queries, results): + logging.info(f'Query : {query}') + logging.info(f'Result: {result.strip()}\n') + + logging.info(f'Results are saved at {exp_dir}') + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/nlp/token_classification/token_classification_train.py b/models/tts/NeMo/examples/nlp/token_classification/token_classification_train.py new file mode 100644 index 0000000000..f4ffb61c55 --- /dev/null +++ b/models/tts/NeMo/examples/nlp/token_classification/token_classification_train.py @@ -0,0 +1,147 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +import pytorch_lightning as pl +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import TokenClassificationModel +from nemo.collections.nlp.parts.nlp_overrides import NLPDDPPlugin +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +""" +This scripts shows how to train a Token Classification model. + +The Token Classification model supports Named Entity Recognition task and other token level classification tasks, +as long as the data follows the format specified below. + +More details on how to use this script could be found in +tutorials/nlp/Token_Classification_Named_Entity_Recognition.ipynb + +*** Data Format *** +Token Classification Model requires the data to be split into 2 files: text.txt and labels.txt. +Each line of the text.txt file contains text sequences, where words are separated with spaces, i.e.: +[WORD] [SPACE] [WORD] [SPACE] [WORD]. +The labels.txt file contains corresponding labels for each word in text.txt, the labels are separated with spaces, i.e.: +[LABEL] [SPACE] [LABEL] [SPACE] [LABEL]. + +Example of a text.txt file: +Jennifer is from New York City . +She likes ... +... + +Corresponding labels.txt file: +B-PER O O B-LOC I-LOC I-LOC O +O O ... +... + +*** Preparing the dataset *** + +To convert an IOB format data to the format required for training, run +examples/nlp/token_classification/data/import_from_iob_format.py on your train and dev files, as follows: + +python examples/nlp/token_classification/data/import_from_iob_format.py --data_file PATH_TO_IOB_FORMAT_DATAFILE + +*** Setting the configs *** + +The model and the PT trainer are defined in a config file which declares multiple important sections. +The most important ones are: + model: All arguments that are related to the Model - language model, tokenizer, token classifier, optimizer, + schedulers, and datasets/data loaders. + trainer: Any argument to be passed to PyTorch Lightning including number of epochs, number of GPUs, + precision level, etc. +This script uses the `/examples/nlp/token_classification/conf/token_classification_config.yaml` config file +by default. You may update the config file from the file directly. +The other option is to set another config file via command line arguments by `--config-name=CONFIG_FILE_PATH'. + +For more details about the config files and different ways of model restoration, see tutorials/00_NeMo_Primer.ipynb + +*** Model Training *** + +To train TokenClassification model from scratch with the default config file, run: + + python token_classification_train.py \ + model.dataset.data_dir= \ + trainer.max_epochs= \ + trainer.gpus=[] + +To use one of the pretrained versions of the model specify a `pretrained_model` arg with either +TokenClassification model from list_available_models() or path to a .nemo file, for example: +ner_en_bert or model.nemo, run: + + python token_classification_train.py pretrained_model=ner_en_bert + +To use one of the pretrained versions of the model and fine-tune it, run: + + python token_classification_train.py \ + model.dataset.data_dir= \ + pretrained_model=ner_en_bert + + - a directory that contains test_ds.text_file and test_ds.labels_file (see the config) +pretrained_model - pretrained TokenClassification model from list_available_models() or + path to a .nemo file, for example: ner_en_bert or model.nemo + +For more ways of restoring a pre-trained model, see tutorials/00_NeMo_Primer.ipynb +""" + + +@hydra_runner(config_path="conf", config_name="token_classification_config") +def main(cfg: DictConfig) -> None: + trainer = pl.Trainer(plugins=[NLPDDPPlugin()], **cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + + if not cfg.pretrained_model: + logging.info(f'Config: {OmegaConf.to_yaml(cfg)}') + model = TokenClassificationModel(cfg.model, trainer=trainer) + else: + if os.path.exists(cfg.pretrained_model): + # TODO: can we drop strict=False? + model = TokenClassificationModel.restore_from(cfg.pretrained_model, trainer=trainer, strict=False) + elif cfg.pretrained_model in TokenClassificationModel.get_available_model_names(): + model = TokenClassificationModel.from_pretrained(cfg.pretrained_model) + else: + raise ValueError( + f'Provide path to the pre-trained .nemo file or choose from {TokenClassificationModel.list_available_models()}' + ) + + data_dir = cfg.model.dataset.get('data_dir', None) + if data_dir: + if not os.path.exists(data_dir): + raise ValueError(f'{data_dir} is not found at') + + # we can also do finetuning of the pretrained model but it will require + # setup the data dir to get class weights statistics + model.update_data_dir(data_dir=data_dir) + # finally, setup train and validation Pytorch DataLoaders + model.setup_training_data() + model.setup_validation_data() + # then we're setting up loss, use model.dataset.class_balancing, + # if you want to add class weights to the CrossEntropyLoss + model.setup_loss(class_balancing=cfg.model.dataset.class_balancing) + logging.info(f'Using config file of the pretrained model') + else: + raise ValueError( + 'Specify a valid dataset directory that contains test_ds.text_file and test_ds.labels_file \ + with "model.dataset.data_dir" argument' + ) + + trainer.fit(model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/speaker_recognition/README.md b/models/tts/NeMo/examples/speaker_recognition/README.md new file mode 100644 index 0000000000..4e321b3553 --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/README.md @@ -0,0 +1,21 @@ +# Speaker Recognition + +Documentation section for speaker related tasks can be found at: + - [Speaker Recognition and Verification](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speaker_recognition/intro.html) + - [Speaker Diarization](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/speaker_diarization/intro.html) + +This folder contains scripts for performing various speaker related tasks as breiefly explained below: + - **speaker_reco.py** +: Training script for speaker verification and recognition + - **speaker_reco_infer.py** +: Inference for speaker recognition. Can be used to map labels using trained model on test set with known speaker labels from train set. (speaker id) + - **extract_speaker_embeddings.py** +: Extract speaker embeddings using trained model for speaker verification purposes. + - **speaker_reco_finetune.py** +: Sample script to demonstrate on how to finetune a already trained model for specific sub tasks. + - **voxceleb_eval.py** +: Helper script to evaluate speaker embeddings on voxceleb trail files. + - **speaker_diarize.py** +: Inference script for speaker diarization - knowing who spoke when. + - **conf/** +: Folder containing configuration files for training speaker tasks. We currently only support inference for speaker diarization. diff --git a/models/tts/NeMo/examples/speaker_recognition/conf/SpeakerNet_recognition_3x2x512.yaml b/models/tts/NeMo/examples/speaker_recognition/conf/SpeakerNet_recognition_3x2x512.yaml new file mode 100644 index 0000000000..4bba3a780d --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/conf/SpeakerNet_recognition_3x2x512.yaml @@ -0,0 +1,157 @@ +name: &name "SpeakerNet" +sample_rate: &sample_rate 16000 +repeat: &rep 2 +dropout: &drop 0.5 +separable: &separable True +n_filters: &n_filters 512 + +model: + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: null + batch_size: 64 + shuffle: True + time_length: 8 + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: null + batch_size: 128 + shuffle: False + time_length: 8 + + test_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: null + batch_size: 1 + shuffle: False + time_length: 8 + embedding_dir: '.' + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: *n_filters + repeat: 1 + kernel: [3] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: *n_filters + repeat: *rep + kernel: [7] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: *n_filters + repeat: *rep + kernel: [11] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: *n_filters + repeat: *rep + kernel: [15] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: &enc_feat_out 1500 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: *separable + + decoder: + _target_: nemo.collections.asr.modules.SpeakerDecoder + feat_in: *enc_feat_out + num_classes: 7205 + pool_mode: 'xvector' + emb_sizes: 512,512 + angular: False + + loss: + scale: 30 + margin: 0.2 + + optim: + name: novograd + # _target_: nemo.core.optim.optimizers.Novograd + lr: .008 + # optimizer arguments + args: + name: auto + # _target_: nemo.core.config.optimizers.NovogradParams + betas: [0.95, 0.5] + weight_decay: 0.001 + + # scheduler setup + sched: + name: CosineAnnealing + iters_per_batch: 1 # computed at runtime + max_steps: null # computed at runtime or explicitly set here + + # scheduler config override + args: + name: auto + # _target_: nemo.core.config.schedulers.CosineAnnealingParams + warmup_steps: null + warmup_ratio: 0.1 + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 1 # number of gpus + max_epochs: 200 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + amp_level: O0 + deterministic: True + checkpoint_callback: False + logger: False + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/speaker_recognition/conf/SpeakerNet_verification_3x2x512.yaml b/models/tts/NeMo/examples/speaker_recognition/conf/SpeakerNet_verification_3x2x512.yaml new file mode 100644 index 0000000000..9ce23f1f8d --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/conf/SpeakerNet_verification_3x2x512.yaml @@ -0,0 +1,145 @@ +name: &name "SpeakerNet" +sample_rate: &sample_rate 16000 +repeat: &rep 2 +dropout: &drop 0.5 +separable: &separable True +n_filters: &n_filters 512 + +model: + train_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: null + batch_size: 64 + shuffle: True + time_length: 8 + is_tarred: False + tarred_audio_filepaths: null + tarred_shard_strategy: "scatter" + + validation_ds: + manifest_filepath: ??? + sample_rate: 16000 + labels: null + batch_size: 128 + shuffle: False + time_length: 8 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + normalize: "per_feature" + window_size: 0.02 + sample_rate: *sample_rate + window_stride: 0.01 + window: "hann" + features: &n_mels 64 + n_fft: 512 + frame_splicing: 1 + dither: 0.00001 + + encoder: + _target_: nemo.collections.asr.modules.ConvASREncoder + feat_in: *n_mels + activation: relu + conv_mask: true + + jasper: + - filters: *n_filters + repeat: 1 + kernel: [3] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: *n_filters + repeat: *rep + kernel: [7] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: *n_filters + repeat: *rep + kernel: [11] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: *n_filters + repeat: *rep + kernel: [15] + stride: [1] + dilation: [1] + dropout: *drop + residual: true + separable: *separable + + - filters: &enc_feat_out 1500 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: 0.0 + residual: false + separable: *separable + + decoder: + _target_: nemo.collections.asr.modules.SpeakerDecoder + feat_in: *enc_feat_out + num_classes: 7205 + pool_mode: 'xvector' + emb_sizes: 256 + angular: True + + loss: + scale: 30 + margin: 0.2 + + optim: + name: sgd + lr: .006 + # optimizer arguments + args: + name: auto + weight_decay: 0.001 + momentum: 0.9 + + # scheduler setup + sched: + name: CosineAnnealing + iters_per_batch: 1 # computed at runtime + max_steps: null # computed at runtime or explicitly set here + + # scheduler config override + args: + name: auto + warmup_steps: null + warmup_ratio: 0.1 + min_lr: 0.0 + last_epoch: -1 + +trainer: + gpus: 1 # number of gpus + max_epochs: 200 + max_steps: null # computed at runtime if not set + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + amp_level: O0 + deterministic: True + checkpoint_callback: False + logger: False + log_every_n_steps: 1 # Interval of logging. + val_check_interval: 1.0 # Set to 0.25 to check 4 times per epoch, or an int for number of iterations + +exp_manager: + exp_dir: null + name: *name + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/speaker_recognition/conf/speaker_diarization.yaml b/models/tts/NeMo/examples/speaker_recognition/conf/speaker_diarization.yaml new file mode 100644 index 0000000000..482bf4682a --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/conf/speaker_diarization.yaml @@ -0,0 +1,27 @@ +name: &name "ClusterDiarizer" + +num_workers: 4 +sample_rate: &sample_rate 16000 + +diarizer: + oracle_num_speakers: ??? # oracle number of speakers or pass null if not known. [int or path to file containing uniq-id with num of speakers of that session] + max_num_speakers: 20 # max number of speakers for each recording. If oracle num speakers is passed this value is ignored. + out_dir: ??? + paths2audio_files: ??? #file containing paths to audio files for inference + path2groundtruth_rttm_files: null #file containing paths to ground truth rttm files for score computation or null if no reference rttms are present + + vad: + model_path: null #.nemo local model path or pretrained model name or none + window_length_in_sec: 0.15 + shift_length_in_sec: 0.01 + threshold: 0.7 # tune threshold on dev set. Check /scripts/voice_activity_detection/vad_tune_threshold.py + vad_decision_smoothing: True + smoothing_params: + method: "median" + overlap: 0.875 + + speaker_embeddings: + oracle_vad_manifest: null + model_path: ??? #.nemo local model path or pretrained model name (speakerverification_speakernet or speakerdiarization_speakernet) + window_length_in_sec: 1.5 # window length in sec for speaker embedding extraction + shift_length_in_sec: 0.75 # shift length in sec for speaker embedding extraction diff --git a/models/tts/NeMo/examples/speaker_recognition/extract_speaker_embeddings.py b/models/tts/NeMo/examples/speaker_recognition/extract_speaker_embeddings.py new file mode 100644 index 0000000000..ce7ce0ac5f --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/extract_speaker_embeddings.py @@ -0,0 +1,101 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This is a helper script to extract speaker embeddings based on manifest file +Usage: +python extract_speaker_embeddings.py --manifest=/path/to/manifest/file' +--model_path='/path/to/.nemo/file'(optional) +--embedding_dir='/path/to/embedding/directory' + +Args: +--manifest: path to manifest file containing audio_file paths for which embeddings need to be extracted +--model_path(optional): path to .nemo speaker verification model file to extract embeddings, if not passed SpeakerNet-M model would + be downloaded from NGC and used to extract embeddings +--embeddings_dir(optional): path to directory where embeddings need to stored default:'./' + + +""" + +from argparse import ArgumentParser + +import pytorch_lightning as pl +import torch +from omegaconf import OmegaConf + +from nemo.collections.asr.models.label_models import ExtractSpeakerEmbeddingsModel +from nemo.utils import logging + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--manifest", type=str, required=True, help="Path to manifest file", + ) + parser.add_argument( + "--model_path", + type=str, + default='speakerverification_speakernet', + required=False, + help="path to .nemo speaker verification model file to extract embeddings, if not passed SpeakerNet-M model would be downloaded from NGC and used to extract embeddings", + ) + parser.add_argument( + "--embedding_dir", + type=str, + default='./', + required=False, + help="path to directory where embeddings need to stored default:'./'", + ) + args = parser.parse_args() + torch.set_grad_enabled(False) + + if args.model_path.endswith('.nemo'): + logging.info(f"Using local speaker model from {args.model_path}") + speaker_model = ExtractSpeakerEmbeddingsModel.restore_from(restore_path=args.model_path) + else: + speaker_model = ExtractSpeakerEmbeddingsModel.from_pretrained(model_name="speakerverification_speakernet") + logging.info(f"using pretrained speaker verification model from NGC") + + num_gpus = 1 if torch.cuda.is_available() else 0 + if not num_gpus: + logging.warning("Running model on CPU, for faster performance it is adviced to use atleast one NVIDIA GPUs") + + trainer = pl.Trainer(gpus=num_gpus, accelerator=None) + + test_config = OmegaConf.create( + dict( + manifest_filepath=args.manifest, + sample_rate=16000, + labels=None, + batch_size=1, + shuffle=False, + time_length=8, + embedding_dir=args.embedding_dir, + ) + ) + speaker_model.setup_test_data(test_config) + trainer.test(speaker_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/speaker_recognition/speaker_diarize.py b/models/tts/NeMo/examples/speaker_recognition/speaker_diarize.py new file mode 100644 index 0000000000..a9013134c7 --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/speaker_diarize.py @@ -0,0 +1,49 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from omegaconf import OmegaConf +from pytorch_lightning import seed_everything + +from nemo.collections.asr.models import ClusteringDiarizer +from nemo.core.config import hydra_runner +from nemo.utils import logging + +""" +This script demonstrates how to use run speaker diarization. +Usage: + python speaker_diarize.py \ + diarizer.paths2audio_files= \ + diarizer.path2groundtruth_rttm_files=<(Optional) either list of rttm file paths or file containing paths to rttm files> \ + diarizer.vad.model_path="vad_telephony_marblenet" \ + diarizer.vad.threshold=0.7 \ + diarizer.speaker_embeddings.model_path="speakerverification_speakernet" \ + diarizer.out_dir="demo" + +Check out whole parameters in ./conf/speaker_diarization.yaml and their meanings. +For details, have a look at /tutorials/speaker_recognition/Speaker_Diarization_Inference.ipynb +""" + +seed_everything(42) + + +@hydra_runner(config_path="conf", config_name="speaker_diarization.yaml") +def main(cfg): + + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + sd_model = ClusteringDiarizer(cfg=cfg) + sd_model.diarize() + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/speaker_recognition/speaker_reco.py b/models/tts/NeMo/examples/speaker_recognition/speaker_reco.py new file mode 100644 index 0000000000..90bb75388b --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/speaker_reco.py @@ -0,0 +1,73 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os + +import pytorch_lightning as pl +from omegaconf import ListConfig, OmegaConf +from pytorch_lightning import seed_everything + +from nemo.collections.asr.models import EncDecSpeakerLabelModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + +""" +Basic run (on GPU for 10 epochs for 2 class training): +EXP_NAME=sample_run +python ./speaker_reco.py --config-path='conf' --config-name='SpeakerNet_recognition_3x2x512.yaml' \ + trainer.max_epochs=10 \ + model.train_ds.batch_size=64 model.validation_ds.batch_size=64 \ + model.train_ds.manifest_filepath="" model.validation_ds.manifest_filepath="" \ + model.test_ds.manifest_filepath="" \ + trainer.gpus=1 \ + model.decoder.params.num_classes=2 \ + exp_manager.name=$EXP_NAME +exp_manager.use_datetime_version=False \ + exp_manager.exp_dir='./speaker_exps' + +See https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_recognition/Speaker_Recognition_Verification.ipynb for notebook tutorial + +Optional: Use tarred dataset to speech up data loading. + Prepare ONE manifest that contains all training data you would like to include. Validation should use non-tarred dataset. + Note that it's possible that tarred datasets impacts validation scores because it drop values in order to have same amount of files per tarfile; + Scores might be off since some data is missing. + + Use the `convert_to_tarred_audio_dataset.py` script under /speech_recognition/scripts in order to prepare tarred audio dataset. + For details, please see TarredAudioToClassificationLabelDataset in /nemo/collections/asr/data/audio_to_label.py +""" + +seed_everything(42) + + +@hydra_runner(config_path="conf", config_name="SpeakerNet_recognition_3x2x512.yaml") +def main(cfg): + + logging.info(f'Hydra config: {OmegaConf.to_yaml(cfg)}') + trainer = pl.Trainer(**cfg.trainer) + log_dir = exp_manager(trainer, cfg.get("exp_manager", None)) + speaker_model = EncDecSpeakerLabelModel(cfg=cfg.model, trainer=trainer) + trainer.fit(speaker_model) + if not trainer.fast_dev_run: + model_path = os.path.join(log_dir, '..', 'spkr.nemo') + speaker_model.save_to(model_path) + + if hasattr(cfg.model, 'test_ds') and cfg.model.test_ds.manifest_filepath is not None: + gpu = 1 if cfg.trainer.gpus != 0 else 0 + trainer = pl.Trainer(gpus=gpu) + if speaker_model.prepare_test(trainer): + trainer.test(speaker_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/speaker_recognition/speaker_reco_finetune.py b/models/tts/NeMo/examples/speaker_recognition/speaker_reco_finetune.py new file mode 100644 index 0000000000..c8ed4d1835 --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/speaker_reco_finetune.py @@ -0,0 +1,98 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script demonstrates how to use pretrained .nemo speaker recognition model to finetune on your domain set +Usage: +python speaker_reco_finetune.py --pretrained_model='/path/to/.nemo_or_.ckpt/file' --finetune_config_file=/path/to/finetune/config/yaml/file' + +in finetune config file make sure to change +- train manifest +- validation manifest +- decoder num classes +- learning rate +- num of epochs + +for finetuning tips see: https://github.com/NVIDIA/NeMo/blob/main/tutorials/speaker_recognition/Speaker_Recognition_Verification.ipynb +""" + + +from argparse import ArgumentParser + +import pytorch_lightning as pl +from omegaconf import OmegaConf + +from nemo.collections.asr.models import EncDecSpeakerLabelModel +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--pretrained_model", + type=str, + default="speakerverification_speakernet", + required=False, + help="Pass your trained .nemo model", + ) + parser.add_argument( + "--finetune_config_file", + type=str, + required=True, + help="path to speakernet config yaml file to load train, validation dataset and also for trainer parameters", + ) + + parser.add_argument( + "--freeze_encoder", + type=bool, + required=False, + default=True, + help="True if speakernet encoder paramteres needs to be frozen while finetuning", + ) + + args = parser.parse_args() + + if args.pretrained_model.endswith('.nemo'): + logging.info(f"Using local speaker model from {args.pretrained_model}") + speaker_model = EncDecSpeakerLabelModel.restore_from(restore_path=args.pretrained_model) + elif args.pretrained_model.endswith('.ckpt'): + logging.info(f"Using local speaker model from checkpoint {args.pretrained_model}") + speaker_model = EncDecSpeakerLabelModel.load_from_checkpoint(checkpoint_path=args.pretrained_model) + else: + logging.info("Using pretrained speaker recognition model from NGC") + speaker_model = EncDecSpeakerLabelModel.from_pretrained(model_name=args.pretrained_model) + + finetune_config = OmegaConf.load(args.finetune_config_file) + + if 'test_ds' in finetune_config.model and finetune_config.model.test_ds is not None: + finetune_config.model.test_ds = None + logging.warning("Removing test ds") + + speaker_model.setup_finetune_model(finetune_config.model) + finetune_trainer = pl.Trainer(**finetune_config.trainer) + speaker_model.set_trainer(finetune_trainer) + + _ = exp_manager(finetune_trainer, finetune_config.get('exp_manager', None)) + speaker_model.setup_optimization(finetune_config.optim) + + if args.freeze_encoder: + for param in speaker_model.encoder.parameters(): + param.requires_grad = False + + finetune_trainer.fit(speaker_model) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/speaker_recognition/speaker_reco_infer.py b/models/tts/NeMo/examples/speaker_recognition/speaker_reco_infer.py new file mode 100644 index 0000000000..bc3ceaf642 --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/speaker_reco_infer.py @@ -0,0 +1,139 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script serves two goals: + (1) Demonstrate how to use NeMo Models outside of PytorchLightning + (2) Shows example of batch speaker recognition inference +Usage: +python speaker_reco_infer.py --spkr_model='/path/to/.nemo/file' --train_manifest=/path/to/train/manifest/file' +--test_manifest=/path/to/train/manifest/file' --batch_size=32 +train_manifest file is used to map the labels from which model was trained so it is mandatory to +pass the train manifest file + +for finetuning tips see: https://colab.research.google.com/github/NVIDIA/NeMo/blob/main/tutorials/speaker_recognition/Speaker_Recognition_Verification.ipynb +""" + + +import json +import os +import sys +from argparse import ArgumentParser + +import numpy as np +import torch +from tqdm import tqdm + +from nemo.collections.asr.models import EncDecSpeakerLabelModel +from nemo.utils import logging + +try: + from torch.cuda.amp import autocast +except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + +can_gpu = torch.cuda.is_available() + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--spkr_model", + type=str, + default="speakerrecognition_speakernet", + required=True, + help="Pass your trained .nemo model", + ) + parser.add_argument( + "--train_manifest", type=str, required=True, help="path to train manifest file to match labels" + ) + parser.add_argument( + "--test_manifest", type=str, required=True, help="path to test manifest file to perform inference" + ) + parser.add_argument("--batch_size", type=int, default=32) + args = parser.parse_args() + torch.set_grad_enabled(False) + + if args.spkr_model.endswith('.nemo'): + logging.info(f"Using local speaker model from {args.spkr_model}") + speaker_model = EncDecSpeakerLabelModel.restore_from(restore_path=args.spkr_model) + else: + logging.error(f"Please pass a trained .nemo file") + sys.exit() + + labels = [] + with open(args.train_manifest, 'rb') as f: + lines = f.readlines() + for line in lines: + line = line.strip() + item = json.loads(line) + labels.append(item['label']) + + labels_map = sorted(set(labels)) + label2id, id2label = {}, {} + for label_id, label in enumerate(labels_map): + label2id[label] = label_id + id2label[label_id] = label + + speaker_model.setup_test_data( + test_data_layer_params={ + 'sample_rate': 16000, + 'manifest_filepath': args.test_manifest, + 'labels': labels_map, + 'batch_size': args.batch_size, + 'trim_silence': False, + 'shuffle': False, + } + ) + if can_gpu: + speaker_model = speaker_model.cuda() + speaker_model.eval() + + speaker_model.test_dataloader() + all_labels = [] + all_logits = [] + for test_batch in tqdm(speaker_model.test_dataloader()): + if can_gpu: + test_batch = [x.cuda() for x in test_batch] + with autocast(): + audio_signal, audio_signal_len, labels, _ = test_batch + logits, _ = speaker_model.forward(input_signal=audio_signal, input_signal_length=audio_signal_len) + + all_logits.extend(logits.cpu().numpy()) + all_labels.extend(labels.cpu().numpy()) + + all_logits, true_labels = np.asarray(all_logits), np.asarray(all_labels) + infer_labels = all_logits.argmax(axis=1) + + out_manifest = os.path.basename(args.test_manifest).split('.')[0] + '_infer.json' + out_manifest = os.path.join(os.path.dirname(args.test_manifest), out_manifest) + with open(args.test_manifest, 'rb') as f1, open(out_manifest, 'w') as f2: + lines = f1.readlines() + for idx, line in enumerate(lines): + line = line.strip() + item = json.loads(line) + item['infer'] = id2label[infer_labels[idx]] + json.dump(item, f2) + f2.write('\n') + + logging.info("Inference labels have been written to {} manifest file".format(out_manifest)) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/examples/speaker_recognition/voxceleb_eval.py b/models/tts/NeMo/examples/speaker_recognition/voxceleb_eval.py new file mode 100644 index 0000000000..cdfd1db033 --- /dev/null +++ b/models/tts/NeMo/examples/speaker_recognition/voxceleb_eval.py @@ -0,0 +1,110 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os +import pickle as pkl +import sys + +import numpy as np +from scipy.interpolate import interp1d +from scipy.optimize import brentq +from sklearn.metrics import roc_curve +from tqdm import tqdm + + +""" +This script faciliates to get EER % based on cosine-smilarity +for Voxceleb dataset. + +Args: + trial_file str: path to voxceleb trial file + emb : path to pickle file of embeddings dictionary (generated from spkr_get_emb.py) + save_kaldi_emb: if required pass this argument to save kaldi embeddings for KALDI PLDA training later + Note: order of audio files in manifest file should match the embeddings +""" + + +def get_acc(trial_file='', emb='', save_kaldi_emb=False): + + trial_score = open('trial_score.txt', 'w') + dirname = os.path.dirname(trial_file) + emb = pkl.load(open(emb, 'rb')) + trial_embs = [] + keys = [] + all_scores = [] + all_keys = [] + + # for each trials in trial file + with open(trial_file, 'r') as f: + tmp_file = f.readlines() + for line in tqdm(tmp_file): + line = line.strip() + truth, x_speaker, y_speaker = line.split() + + x_speaker = x_speaker.split('/') + x_speaker = '@'.join(x_speaker) + + y_speaker = y_speaker.split('/') + y_speaker = '@'.join(y_speaker) + + X = emb[x_speaker] + Y = emb[y_speaker] + + if save_kaldi_emb and x_speaker not in keys: + keys.append(x_speaker) + trial_embs.extend([X]) + + if save_kaldi_emb and y_speaker not in keys: + keys.append(y_speaker) + trial_embs.extend([Y]) + + score = (X @ Y.T) / (((X @ X.T) * (Y @ Y.T)) ** 0.5) + score = (score + 1) / 2 + + all_scores.append(score) + trial_score.write(str(score) + "\t" + truth) + truth = int(truth) + all_keys.append(truth) + + trial_score.write('\n') + trial_score.close() + + if save_kaldi_emb: + np.save(dirname + '/all_embs_voxceleb.npy', np.asarray(trial_embs)) + np.save(dirname + '/all_ids_voxceleb.npy', np.asarray(keys)) + print("Saved KALDI PLDA related embeddings to {}".format(dirname)) + + return np.asarray(all_scores), np.asarray(all_keys) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--trial_file", help="path to voxceleb trial file", type=str, required=True) + parser.add_argument("--emb", help="path to numpy file of embeddings", type=str, required=True) + parser.add_argument( + "--save_kaldi_emb", + help=":save kaldi embeddings for KALDI PLDA training later", + required=False, + action='store_true', + ) + + args = parser.parse_args() + trial_file, emb, save_kaldi_emb = args.trial_file, args.emb, args.save_kaldi_emb + + y_score, y = get_acc(trial_file=trial_file, emb=emb, save_kaldi_emb=save_kaldi_emb) + fpr, tpr, thresholds = roc_curve(y, y_score, pos_label=1) + + eer = brentq(lambda x: 1.0 - x - interp1d(fpr, tpr)(x), 0.0, 1.0) + sys.stdout.write("{0:.2f}\n".format(eer * 100)) diff --git a/models/tts/NeMo/examples/tts/aligner.py b/models/tts/NeMo/examples/tts/aligner.py new file mode 100644 index 0000000000..e32c0444ca --- /dev/null +++ b/models/tts/NeMo/examples/tts/aligner.py @@ -0,0 +1,33 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import AlignerModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path='conf', config_name='aligner') +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get('exp_manager', None)) + model = AlignerModel(cfg=cfg.model, trainer=trainer) + trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/conf/aligner.yaml b/models/tts/NeMo/examples/tts/conf/aligner.yaml new file mode 100644 index 0000000000..801b5770eb --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/aligner.yaml @@ -0,0 +1,123 @@ +name: Aligner +sample_rate: 22050 + +train_dataset: ??? +validation_datasets: ??? +attn_prior_folder: ??? + +model: + d_char: 512 + n_mels: 80 + bin_start_ratio: 0.3 + + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithPriorDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + int_values: false + normalize: false + sample_rate: ${sample_rate} + trim: false + attn_prior_folder: ${attn_prior_folder} + n_window_stride: 256 + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: None + pad_with_space: True + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithPriorDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + int_values: false + normalize: false + sample_rate: ${sample_rate} + trim: false + attn_prior_folder: ${attn_prior_folder} + n_window_stride: 256 + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: None + pad_with_space: True + dataloader_params: + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 1 + + preprocessor: + _target_: "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor" + dither: 0.0 + features: ${model.n_mels} + frame_splicing: 1 + highfreq: 8000 + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + window_size: False + n_window_stride: 256 + window_stride: False + normalize: null + pad_to: 1 + pad_value: -11.52 + preemph: null + sample_rate: ${sample_rate} + stft_conv: false + window: hann + + alignment_encoder: + _target_: "nemo.collections.tts.modules.aligner.AlignmentEncoder" + n_mel_channels: ${model.n_mels} + n_text_channels: ${model.d_char} + n_att_channels: ${model.n_mels} + + optim: + name: adam + lr: 1e-3 + weight_decay: 1e-6 + + sched: + name: CosineAnnealing + min_lr: 5e-5 + warmup_ratio: 0.35 + +trainer: + gpus: 1 + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False + logger: False + gradient_clip_val: 1.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: False + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_forward_sum_loss" + mode: "min" \ No newline at end of file diff --git a/models/tts/NeMo/examples/tts/conf/degli.yaml b/models/tts/NeMo/examples/tts/conf/degli.yaml new file mode 100644 index 0000000000..5327694142 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/degli.yaml @@ -0,0 +1,107 @@ +name: "DeepGriffinLim" +sample_rate: 22050 +n_fft: 1024 +hop_length: 256 +destination: "/tmp" + +model: + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.NoisySpecsDataset" + destination: ${destination} + subdir: "/train" + n_fft: ${n_fft} + hop_length: ${hop_length} + num_snr: 1 + dataloader_params: + drop_last: false + shuffle: true + batch_size: 16 + num_workers: 1 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.NoisySpecsDataset" + destination: ${destination} + subdir: "/valid" + n_fft: ${n_fft} + hop_length: ${hop_length} + num_snr: 1 + + dataloader_params: + drop_last: false + shuffle: false + batch_size: 16 + num_workers: 1 + + preprocessor: + valid_filelist: none + train_filelist: none + n_fft: ${n_fft} + hop_length: ${hop_length} + num_snr: 1 + destination: ${destination} + + degli: + _target_: nemo.collections.tts.modules.degli.DegliModule + n_fft: ${n_fft} + hop_length: ${hop_length} + depth: 1 + out_all_block: True + tiny: False # overrides all below configurations if true. + layers: 6 + k_x: 3 + k_y: 3 + s_x: 1 + s_y: 2 + widening: 16 + use_bn: False + linear_finalizer: True + convGlu: False + act: relu + act2: selu + glu_bn: True + use_weight_norm: True + + + optim: + name: novograd + lr: 0.01 + weight_decay: 0.001 + + sched: + name: ReduceLROnPlateau + mode: min + factor: 0.6 + patience: 5 + verbose: True + threshold: 0.01 + threshold_mode: 'rel' + cooldown: 0 + min_lr: 0.00001 + eps: 1.e-08 + + train_params: + repeat_training: 2 + repeat_validation: 32 + sampling_rate: ${sample_rate} + +trainer: + gpus: 1 # number of gpus + max_epochs: 200 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + flush_logs_every_n_steps: 2000 + log_every_n_steps: 200 + check_val_every_n_epoch: 5 + precision: 32 + gradient_clip_val: 4.0 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/ed_mel2spec.yaml b/models/tts/NeMo/examples/tts/conf/ed_mel2spec.yaml new file mode 100755 index 0000000000..d98464e1b9 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/ed_mel2spec.yaml @@ -0,0 +1,103 @@ +name: "EncoderDecoderMelToSpec" +sample_rate: 22050 +n_fft: 1024 +hop_length: 256 +destination: "/tmp" + +model: + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.NoisySpecsDataset" + destination: ${destination} + subdir: "/train" + n_fft: ${n_fft} + hop_length: ${hop_length} + num_snr: 1 + dataloader_params: + drop_last: false + shuffle: true + batch_size: 16 + num_workers: 1 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.NoisySpecsDataset" + destination: ${destination} + subdir: "/valid" + n_fft: ${n_fft} + hop_length: ${hop_length} + num_snr: 1 + + dataloader_params: + drop_last: false + shuffle: false + batch_size: 16 + num_workers: 1 + + preprocessor: + valid_filelist: none + train_filelist: none + n_fft: ${n_fft} + hop_length: ${hop_length} + num_snr: 1 + destination: ${destination} + + mel2spec: + _target_: nemo.collections.tts.modules.ed_mel2spec.EDMel2SpecModule + n_fft: ${n_fft} + hop_length: ${hop_length} + mel_fmin: 0 + mel_fmax: null + mel_freq: 80 + layers: 8 + sampling_rate: ${sample_rate} + widening: 16 + use_batchnorm: True + droprate: 0.2 + num_dropout: 3 + pre_final_lin: True + act1: lrelu + act2: relu + use_weight_norm: False + + optim: + name: adam + lr: 5.e-4 + weight_decay: 1.e-3 + + sched: + name: ReduceLROnPlateau + mode: min + factor: 0.8 + patience: 5 + verbose: True + threshold: 0.01 + threshold_mode: 'rel' + cooldown: 0 + min_lr: 5.e-5 + eps: 1.e-08 + + train_params: + lreg_factor: 1.0 + loss_mode: 3 + validate_scores: True # Slow! + +trainer: + gpus: 1 # number of gpus + max_epochs: 200 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + flush_logs_every_n_steps: 2000 + log_every_n_steps: 200 + check_val_every_n_epoch: 10 + precision: 32 + gradient_clip_val: 4.0 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/fastpitch.yaml b/models/tts/NeMo/examples/tts/conf/fastpitch.yaml new file mode 100644 index 0000000000..8336d7e502 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/fastpitch.yaml @@ -0,0 +1,144 @@ +name: "FastPitch" +sample_rate: 22050 + +labels: [' ', '!', "'", '(', ')', ',', '-', '.', ':', ';', '?', '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'] +train_dataset: ??? +validation_datasets: ??? +test_datasets: null + +model: + labels: ${labels} + n_speakers: 1 + symbols_embedding_dim: 384 + max_token_duration: 75 + n_mel_channels: 80 + pitch_embedding_kernel_size: 3 + + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.FastPitchDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + trim: false + int_values: false + sample_rate: ${sample_rate} + dataloader_params: + drop_last: false + shuffle: true + batch_size: 48 + num_workers: 12 + + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.FastPitchDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: null + int_values: false + sample_rate: ${sample_rate} + trim: false + dataloader_params: + drop_last: false + shuffle: false + batch_size: 48 + num_workers: 8 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + dither: 0.0 + features: ${model.n_mel_channels} + frame_splicing: 1 + highfreq: 8000 + log: true + log_zero_guard_type: add + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + n_window_stride: 256 + normalize: null + pad_to: 1 + pad_value: 0 + preemph: null + sample_rate: ${sample_rate} + window: hann + window_size: null + window_stride: null + + input_fft: + _target_: nemo.collections.tts.modules.transformer.FFTransformerEncoder + n_layer: 6 + n_head: 1 + d_model: ${model.symbols_embedding_dim} + d_head: 64 + d_inner: 1536 + kernel_size: 3 + dropout: 0.1 + dropatt: 0.1 + dropemb: 0.0 + n_embed: 148 # NOTE Should match # of tokens in `symbol_set` + d_embed: ${model.symbols_embedding_dim} + padding_idx: 0 + + output_fft: + _target_: nemo.collections.tts.modules.transformer.FFTransformerDecoder + n_layer: 6 + n_head: 1 + d_model: ${model.symbols_embedding_dim} + d_head: 64 + d_inner: 1536 + kernel_size: 3 + dropout: 0.1 + dropatt: 0.1 + dropemb: 0.0 + + duration_predictor: + _target_: nemo.collections.tts.modules.fastpitch.TemporalPredictor + input_size: ${model.symbols_embedding_dim} + kernel_size: 3 + filter_size: 256 + dropout: 0.1 + n_layers: 2 + + pitch_predictor: + _target_: nemo.collections.tts.modules.fastpitch.TemporalPredictor + input_size: ${model.symbols_embedding_dim} + kernel_size: 3 + filter_size: 256 + dropout: 0.1 + n_layers: 2 + + optim: + name: lamb + lr: 1e-1 + betas: [0.9, 0.98] + weight_decay: 1e-6 + + sched: + name: NoamAnnealing + warmup_steps: 1000 + last_epoch: -1 + d_model: 1 # Disable scaling based on model dim + +trainer: + gpus: -1 # number of gpus + max_epochs: 1500 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + gradient_clip_val: 1000.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 100 + check_val_every_n_epoch: 5 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/fastpitch_align.yaml b/models/tts/NeMo/examples/tts/conf/fastpitch_align.yaml new file mode 100644 index 0000000000..4b26618b9b --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/fastpitch_align.yaml @@ -0,0 +1,185 @@ +name: "FastPitch" +sample_rate: 22050 +train_dataset: ??? +validation_datasets: ??? +prior_folder: ??? + +model: + learn_alignment: true + n_speakers: 1 + symbols_embedding_dim: 384 + max_token_duration: 75 + n_mel_channels: 80 + pitch_embedding_kernel_size: 3 + n_window_size: 1024 + n_window_stride: 256 + pitch_fmin: 80 + pitch_fmax: 640 + pitch_avg: 211.27540199742586 + pitch_std: 52.1851002822779 + + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithPriorAndPitchDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + int_values: false + normalize: true + sample_rate: ${sample_rate} + trim: false + sup_data_path: ${prior_folder} + n_window_stride: ${model.n_window_stride} + n_window_size: ${model.n_window_size} + pitch_fmin: ${model.pitch_fmin} + pitch_fmax: ${model.pitch_fmax} + pitch_avg: ${model.pitch_avg} + pitch_std: ${model.pitch_std} + vocab: + notation: phonemes + punct: true + spaces: true + stresses: true + add_blank_at: None + pad_with_space: True + chars: true + improved_version_g2p: true + dataloader_params: + drop_last: false + shuffle: true + batch_size: 32 + num_workers: 12 + + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithPriorAndPitchDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: null + int_values: false + normalize: true + sample_rate: ${sample_rate} + trim: false + sup_data_path: ${prior_folder} + n_window_stride: ${model.n_window_stride} + n_window_size: ${model.n_window_size} + pitch_fmin: ${model.pitch_fmin} + pitch_fmax: ${model.pitch_fmax} + pitch_avg: ${model.pitch_avg} + pitch_std: ${model.pitch_std} + vocab: + notation: phonemes + punct: true + spaces: true + stresses: true + add_blank_at: None + pad_with_space: True + chars: true + improved_version_g2p: true + dataloader_params: + drop_last: false + shuffle: false + batch_size: 32 + num_workers: 8 + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + dither: 0.0 + features: ${model.n_mel_channels} + frame_splicing: 1 + highfreq: null + log: true + log_zero_guard_type: add + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${model.n_window_size} + n_window_size: ${model.n_window_size} + n_window_stride: ${model.n_window_stride} + normalize: null + pad_to: 1 + pad_value: 0 + preemph: null + sample_rate: ${sample_rate} + window: hann + window_size: null + window_stride: null + + input_fft: #n_embed and padding_idx are added by the model + _target_: nemo.collections.tts.modules.transformer.FFTransformerEncoder + n_layer: 6 + n_head: 1 + d_model: ${model.symbols_embedding_dim} + d_head: 64 + d_inner: 1536 + kernel_size: 3 + dropout: 0.1 + dropatt: 0.1 + dropemb: 0.0 + d_embed: ${model.symbols_embedding_dim} + + output_fft: + _target_: nemo.collections.tts.modules.transformer.FFTransformerDecoder + n_layer: 6 + n_head: 1 + d_model: ${model.symbols_embedding_dim} + d_head: 64 + d_inner: 1536 + kernel_size: 3 + dropout: 0.1 + dropatt: 0.1 + dropemb: 0.0 + + alignment_module: + _target_: nemo.collections.tts.modules.aligner.AlignmentEncoder + n_text_channels: ${model.symbols_embedding_dim} + + duration_predictor: + _target_: nemo.collections.tts.modules.fastpitch.TemporalPredictor + input_size: ${model.symbols_embedding_dim} + kernel_size: 3 + filter_size: 256 + dropout: 0.1 + n_layers: 2 + + pitch_predictor: + _target_: nemo.collections.tts.modules.fastpitch.TemporalPredictor + input_size: ${model.symbols_embedding_dim} + kernel_size: 3 + filter_size: 256 + dropout: 0.1 + n_layers: 2 + + optim: + name: lamb + lr: 1e-1 + betas: [0.9, 0.98] + weight_decay: 1e-6 + + sched: + name: NoamAnnealing + warmup_steps: 1000 + last_epoch: -1 + d_model: 1 # Disable scaling based on model dim + +trainer: + gpus: -1 # number of gpus + max_epochs: 1500 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + gradient_clip_val: 1000.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 100 + check_val_every_n_epoch: 5 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: v_loss diff --git a/models/tts/NeMo/examples/tts/conf/fastpitch_hifigan_e2e.yaml b/models/tts/NeMo/examples/tts/conf/fastpitch_hifigan_e2e.yaml new file mode 100644 index 0000000000..c0570ee031 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/fastpitch_hifigan_e2e.yaml @@ -0,0 +1,138 @@ +name: "FastPitch_HiFiGAN_E2E" + +labels: [' ', '!', "'", '(', ')', ',', '-', '.', ':', ';', '?', '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'] +train_dataset: ??? +validation_datasets: ??? +test_datasets: null + +model: + sample_rate: 22050 + splice_length: 64 + lr: 3e-2 + labels: ${labels} + n_speakers: 1 + symbols_embedding_dim: 384 + max_token_duration: 75 + n_mel_channels: 80 + pitch_embedding_kernel_size: 3 + mel_loss_coeff: 40 + hop_size: 256 + + train_ds: + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + sample_rate: ${model.sample_rate} + trim: false + parser: null + drop_last: true + shuffle: true + batch_size: 64 + num_workers: 12 + + validation_ds: + manifest_filepath: ${validation_datasets} + sample_rate: ${model.sample_rate} + trim: false + parser: null + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 8 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${model.n_mel_channels} + frame_splicing: 1 + highfreq: 8000 + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + n_window_stride: ${model.hop_size} + normalize: null + pad_to: 1 + pad_value: 0 + preemph: null + sample_rate: ${model.sample_rate} + window: hann + exact_pad: true + use_grads: false + + input_fft: + _target_: nemo.collections.tts.modules.transformer.FFTransformerEncoder + n_layer: 6 + n_head: 1 + d_model: ${model.symbols_embedding_dim} + d_head: 64 + d_inner: 1536 + kernel_size: 3 + dropout: 0.1 + dropatt: 0.1 + dropemb: 0.0 + n_embed: 148 # NOTE Should match # of tokens in `symbol_set` + d_embed: ${model.symbols_embedding_dim} + padding_idx: 0 + + output_fft: + _target_: nemo.collections.tts.modules.transformer.FFTransformerDecoder + n_layer: 6 + n_head: 1 + d_model: ${model.symbols_embedding_dim} + d_head: 64 + d_inner: 1536 + kernel_size: 3 + dropout: 0.1 + dropatt: 0.1 + dropemb: 0.0 + + duration_predictor: + _target_: nemo.collections.tts.modules.fastpitch.TemporalPredictor + input_size: ${model.symbols_embedding_dim} + kernel_size: 3 + filter_size: 256 + dropout: 0.1 + n_layers: 2 + + pitch_predictor: + _target_: nemo.collections.tts.modules.fastpitch.TemporalPredictor + input_size: ${model.symbols_embedding_dim} + kernel_size: 3 + filter_size: 256 + dropout: 0.1 + n_layers: 2 + + generator: + _target_: nemo.collections.tts.modules.hifigan_modules.Generator + upsample_kernel_sizes: [16,16,4,4] + upsample_rates: [8,8,2,2] + upsample_initial_channel: 512 + resblock_kernel_sizes: [3,7,11] + resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]] + resblock: 1 + initial_input_size: 384 + +trainer: + gpus: -1 # number of gpus + max_epochs: 1500 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + gradient_clip_val: 1000.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 100 + check_val_every_n_epoch: 5 + precision: 16 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/fastspeech2.yaml b/models/tts/NeMo/examples/tts/conf/fastspeech2.yaml new file mode 100644 index 0000000000..ed1dea2373 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/fastspeech2.yaml @@ -0,0 +1,144 @@ +name: FastSpeech2 +sample_rate: 22050 +n_mels: 80 + +# Dataset info +train_dataset: ??? +validation_datasets: ??? +mappings_file: ??? # JSON file with word->phone and phone->idx mappings +ignore_file: ??? + +model: + add_pitch_predictor: True + add_energy_predictor: True + d_model: 256 + duration_coeff: 0.25 + mappings_filepath: ${mappings_file} + + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.FastSpeech2Dataset" + manifest_filepath: ${train_dataset} + mappings_filepath: ${mappings_file} + sample_rate: ${sample_rate} + max_duration: null + min_duration: 0.1 + ignore_file: ${ignore_file} + trim: false + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.FastSpeech2Dataset" + manifest_filepath: ${validation_datasets} + mappings_filepath: ${mappings_file} + sample_rate: ${sample_rate} + ignore_file: ${ignore_file} + trim: false + load_supplementary_values: false + dataloader_params: + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 4 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${n_mels} + frame_splicing: 1 + highfreq: 8000 + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + n_window_stride: 256 + normalize: null + pad_to: 16 + pad_value: -11.52 + preemph: null + sample_rate: ${sample_rate} + window: hann + + # FFTransformer encoder + encoder: + _target_: nemo.collections.tts.modules.fastspeech2.FastSpeech2Encoder + d_model: ${model.d_model} + n_layers: 4 + n_attn_heads: 2 + d_attn_head: 256 + d_inner: 1024 + kernel_size: 9 + dropout: 0.1 + attn_dropout: 0.1 + n_embed: 74 # Should match number of tokens in symbol set +1 (pad_idx) + padding_idx: 73 + + # FFTransformer mel spec decoder + decoder: + _target_: nemo.collections.tts.modules.fastspeech2.MelSpecDecoder + d_model: ${model.d_model} + d_out: ${n_mels} + n_layers: 4 + n_attn_heads: 2 + d_attn_head: 256 + d_inner: 1024 + kernel_size: 9 + dropout: 0.1 + attn_dropout: 0.1 + + # VarianceAdaptor + variance_adaptor: + _target_: nemo.collections.tts.modules.fastspeech2.VarianceAdaptor + d_model: ${model.d_model} + dropout: 0.2 + dur_d_hidden: 256 + dur_kernel_size: 3 + pitch: ${model.add_pitch_predictor} + log_pitch: True + n_f0_bins: 256 + pitch_kernel_size: 3 + pitch_min: 80.0 + pitch_max: 800.0 + energy: ${model.add_energy_predictor} + n_energy_bins: 256 + energy_kernel_size: 3 + energy_min: 0.0 + energy_max: 600.0 + + optim: + name: adam + betas: [0.9,0.98] + lr: 1e-3 + weight_decay: 1e-6 + # scheduler setup + sched: + name: NoamAnnealing + warmup_steps: 4000 + min_lr: 1e-5 + d_model: ${model.d_model} + +trainer: + gpus: 1 # number of gpus + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/fastspeech2_hifigan_e2e.yaml b/models/tts/NeMo/examples/tts/conf/fastspeech2_hifigan_e2e.yaml new file mode 100644 index 0000000000..95acf3d769 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/fastspeech2_hifigan_e2e.yaml @@ -0,0 +1,136 @@ +name: FastSpeech2_HiFiGAN_E2E + +# Dataset info +train_dataset: ??? +validation_datasets: ??? +mappings_file: ??? # JSON file with word->phone and phone->idx mappings +ignore_file: ??? + + +model: + lr: 0.5 + splice_length: 64 # Length in frames + mel_loss_coeff: 45. + pitch_loss_coeff: 0.003 + energy_loss_coeff: 0.03 + add_energy_predictor: True + add_pitch_predictor: True + mappings_filepath: ${mappings_file} + sample_rate: 22050 + n_mels: 80 + hop_size: 256 + + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.FastSpeech2Dataset" + manifest_filepath: ${train_dataset} + mappings_filepath: ${mappings_file} + sample_rate: ${model.sample_rate} + max_duration: null + min_duration: 0.75 + ignore_file: ${ignore_file} + trim: false + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.FastSpeech2Dataset" + manifest_filepath: ${validation_datasets} + mappings_filepath: ${mappings_file} + sample_rate: ${model.sample_rate} + ignore_file: ${ignore_file} + trim: false + load_supplementary_values: false + dataloader_params: + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 4 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${n_mels} + frame_splicing: 1 + highfreq: 8000 + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + n_window_stride: ${model.hop_size} + normalize: null + pad_to: 0 + preemph: null + sample_rate: ${model.sample_rate} + window: hann + + # FFTransformer encoder + encoder: + _target_: nemo.collections.tts.modules.fastspeech2.FastSpeech2Encoder + d_model: 256 + n_layers: 4 + n_attn_heads: 2 + d_attn_head: 256 + d_inner: 1024 + kernel_size: 9 + dropout: 0.1 + attn_dropout: 0.1 + n_embed: 74 # Should match number of tokens in symbol set +1 (pad_idx) + padding_idx: 73 + + # VarianceAdaptor + variance_adaptor: + _target_: nemo.collections.tts.modules.fastspeech2.VarianceAdaptor + d_model: 256 + dropout: 0.2 + dur_d_hidden: 256 + dur_kernel_size: 3 + + pitch: ${model.add_pitch_predictor} + log_pitch: True + n_f0_bins: 256 + pitch_kernel_size: 3 + pitch_min: 80.0 + pitch_max: 800.0 + energy: ${model.add_energy_predictor} + n_energy_bins: 256 + energy_kernel_size: 3 + energy_min: 0.0 + energy_max: 600.0 + + generator: + _target_: nemo.collections.tts.modules.hifigan_modules.Generator + upsample_kernel_sizes: [16,16,4,4] + upsample_rates: [8,8,2,2] + upsample_initial_channel: 512 + resblock_kernel_sizes: [3,7,11] + resblock_dilation_sizes: [[1,3,5], [1,3,5], [1,3,5]] + resblock: 1 + initial_input_size: 256 + apply_weight_init_conv_pre: true + +trainer: + gpus: 1 # number of gpus + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + move_metrics_to_cpu: False + precision: 16 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True \ No newline at end of file diff --git a/models/tts/NeMo/examples/tts/conf/glow_tts.yaml b/models/tts/NeMo/examples/tts/conf/glow_tts.yaml new file mode 100644 index 0000000000..f082d9f1f4 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/glow_tts.yaml @@ -0,0 +1,143 @@ +name: "GlowTTS" +sample_rate: 22050 +n_fft: 1024 +n_mels: 80 +fmax: 8000 +pad_value: -11.52 +gin_channels: 0 + +train_dataset: ??? +validation_datasets: ??? +test_datasets: null +cmu_dict_path: null + +model: + + n_speakers: 1 + gin_channels: ${gin_channels} + + train_ds: + manifest_filepath: ${train_dataset} + sample_rate: ${sample_rate} + batch_size: 32 + max_duration: null + min_duration: 0.1 + trim: false + shuffle: True + num_workers: 8 + cmu_dict_path: ${cmu_dict_path} + + + validation_ds: + manifest_filepath: ${validation_datasets} + sample_rate: ${sample_rate} + batch_size: 8 + max_duration: null + min_duration: 0.1 + trim: false + shuffle: False + num_workers: 8 + cmu_dict_path: ${cmu_dict_path} + + test_ds: + manifest_filepath: ${test_datasets} + sample_rate: ${sample_rate} + batch_size: 1 + max_duration: null + min_duration: 0.1 + trim: false + shuffle: False + num_workers: 8 + cmu_dict_path: ${cmu_dict_path} + + parser: + _target_: nemo.collections.tts.modules.glow_tts_parser.GlowTTSParser + cmu_dict_path: ${cmu_dict_path} + + preprocessor: + _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor + dither: 0.0 + features: ${n_mels} + frame_splicing: 1 + highfreq: ${fmax} + log: true + log_zero_guard_type: add + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${n_fft} + n_window_size: 1024 + n_window_stride: 256 + normalize: null + pad_to: 16 + pad_value: ${pad_value} + preemph: null + sample_rate: ${sample_rate} + window: hann + window_size: null + window_stride: null + + encoder: + _target_: nemo.collections.tts.modules.glow_tts.TextEncoder + n_vocab: 148 + out_channels: ${n_mels} + hidden_channels: 192 + filter_channels: 768 + filter_channels_dp: 256 + kernel_size: 3 + p_dropout: 0.1 + n_layers: 6 + n_heads: 2 + window_size: 4 + prenet: true + mean_only: true + gin_channels: ${gin_channels} + + decoder: + _target_: nemo.collections.tts.modules.glow_tts.FlowSpecDecoder + in_channels: ${n_mels} + hidden_channels: 192 + kernel_size: 5 + n_blocks: 12 + n_layers: 4 + n_sqz: 2 + n_split: 4 + sigmoid_scale: false + p_dropout: 0.05 + dilation_rate: 1 + gin_channels: ${gin_channels} + + optim: + name: adam + lr: 1e-3 + # optimizer arguments + betas: [0.9, 0.98] + weight_decay: 0 + + # scheduler setup + sched: + name: CosineAnnealing + + # Scheduler params + warmup_steps: 3000 + min_lr: 1e-4 + last_epoch: -1 + +trainer: + gpus: -1 # number of gpus + max_epochs: 350 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + gradient_clip_val: 1.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 50 + check_val_every_n_epoch: 5 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/hifigan.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/hifigan.yaml new file mode 100644 index 0000000000..9412ced669 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/hifigan.yaml @@ -0,0 +1,67 @@ +name: "HifiGan" +train_dataset: ??? +validation_datasets: ??? + +defaults: + - model/generator: v1 + - model/train_ds: train_ds + - model/validation_ds: val_ds + +model: + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + frame_splicing: 1 + nfilt: 80 + highfreq: 8000 + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + n_window_stride: 256 + normalize: null + pad_to: 0 + pad_value: -11.52 + preemph: null + sample_rate: 22050 + window: hann + use_grads: false + exact_pad: true + + optim: + _target_: torch.optim.AdamW + lr: 0.0002 + betas: [0.8, 0.99] + + sched: + name: CosineAnnealing + min_lr: 1e-5 + warmup_ratio: 0.02 + + max_steps: 25000000 + l1_loss_factor: 45 + denoise_strength: 0.0025 + +trainer: + gpus: -1 # number of gpus + max_steps: ${model.max_steps} + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + flush_logs_every_n_steps: 200 + log_every_n_steps: 100 + check_val_every_n_epoch: 10 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_loss" + mode: "min" diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v1.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v1.yaml new file mode 100644 index 0000000000..30b819e979 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v1.yaml @@ -0,0 +1,8 @@ +# @package _group_ +_target_: nemo.collections.tts.modules.hifigan_modules.Generator +resblock: 1 +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]] diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v2.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v2.yaml new file mode 100644 index 0000000000..ef33801d5c --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v2.yaml @@ -0,0 +1,8 @@ +# @package _group_ +_target_: nemo.collections.tts.modules.hifigan_modules.Generator +resblock: 1 +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]] diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v3.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v3.yaml new file mode 100644 index 0000000000..734f0a330b --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/generator/v3.yaml @@ -0,0 +1,8 @@ +# @package _group_ +_target_: nemo.collections.tts.modules.hifigan_modules.Generator +resblock: 2 +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]] diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/train_ds/train_ds.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/train_ds/train_ds.yaml new file mode 100644 index 0000000000..d7bfce88c2 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/train_ds/train_ds.yaml @@ -0,0 +1,13 @@ +# @package _group_ +dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.75 + n_segments: 8192 + trim: false +dataloader_params: + drop_last: false + shuffle: true + batch_size: 16 + num_workers: 4 diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/train_ds/train_ds_finetune.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/train_ds/train_ds_finetune.yaml new file mode 100644 index 0000000000..195f3b2324 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/train_ds/train_ds_finetune.yaml @@ -0,0 +1,11 @@ +# @package _group_ +dataset: + _target_: "nemo.collections.tts.data.datalayers.MelAudioDataset" + manifest_filepath: ${train_dataset} + min_duration: 0.75 + n_segments: 8192 +dataloader_params: + drop_last: false + shuffle: true + batch_size: 16 + num_workers: 4 diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/validation_ds/val_ds.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/validation_ds/val_ds.yaml new file mode 100644 index 0000000000..cfe6003837 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/validation_ds/val_ds.yaml @@ -0,0 +1,13 @@ +# @package _group_ +dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: null + n_segments: -1 + trim: false +dataloader_params: + drop_last: false + shuffle: false + batch_size: 16 + num_workers: 1 diff --git a/models/tts/NeMo/examples/tts/conf/hifigan/model/validation_ds/val_ds_finetune.yaml b/models/tts/NeMo/examples/tts/conf/hifigan/model/validation_ds/val_ds_finetune.yaml new file mode 100644 index 0000000000..ad9e2b9714 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/hifigan/model/validation_ds/val_ds_finetune.yaml @@ -0,0 +1,11 @@ +# @package _group_ +dataset: + _target_: "nemo.collections.tts.data.datalayers.MelAudioDataset" + manifest_filepath: ${validation_datasets} + min_duration: 3 + n_segments: 66048 +dataloader_params: + drop_last: false + shuffle: false + batch_size: 16 + num_workers: 4 diff --git a/models/tts/NeMo/examples/tts/conf/melgan.yaml b/models/tts/NeMo/examples/tts/conf/melgan.yaml new file mode 100644 index 0000000000..92599889d2 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/melgan.yaml @@ -0,0 +1,108 @@ +name: "MelGan" +sample_rate: 22050 +n_fft: 1024 +fmax: 8000 +pad_value: -11.52 +train_dataset: ??? +validation_datasets: ??? +n_mels: 80 + +model: + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.75 + n_segments: 16384 + trim: false + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: null + n_segments: -1 + trim: false + dataloader_params: + drop_last: false + shuffle: false + batch_size: 13 + num_workers: 4 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + frame_splicing: 1 + nfilt: ${n_mels} + highfreq: ${fmax} + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${n_fft} + + n_window_size: 1024 + n_window_stride: 256 + + normalize: null + pad_to: 0 + pad_value: ${pad_value} + preemph: null + sample_rate: ${sample_rate} + window: hann + exact_pad: true + + generator: + _target_: nemo.collections.tts.modules.melgan_modules.MelGANGenerator + in_channels: ${n_mels} # Number of input channels. + out_channels: 1 # Number of output channels. + kernel_size: 7 # Kernel size of initial and final conv layers. + channels: 512 # Initial number of channels for conv layers. + upsample_scales: [8, 4, 4, 2] # List of Upsampling scales. + stack_kernel_size: 3 # Kernel size of dilated conv layers in residual stack. + stacks: 4 # Number of stacks in a single residual stack module. + use_weight_norm: True # Whether to use weight normalization. + + discriminator: + _target_: nemo.collections.tts.modules.melgan_modules.MelGANMultiScaleDiscriminator + scales: 3 + out_channels: 1 # Number of output channels. + kernel_sizes: [5, 3] # List of kernel size. + channels: 16 # Number of channels of the initial conv layer. + max_downsample_channels: 512 # Maximum number of channels of downsampling layers. + downsample_scales: [4, 4, 4] # List of downsampling scales. + use_weight_norm: True # Whether to use weight norm. + use_spectral_norm: False # Whether to use spectral norm. + + loss: + _target_: nemo.collections.tts.losses.stftlosses.MultiResolutionSTFTLoss + fft_sizes: [1024, 2048, 512] + hop_sizes: [120, 240, 50] + win_lengths: [600, 1200, 240] + window: "hann_window" + +trainer: + gpus: 1 # number of gpus + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + log_every_n_steps: 1000 + check_val_every_n_epoch: 150 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/squeezewave.yaml b/models/tts/NeMo/examples/tts/conf/squeezewave.yaml new file mode 100644 index 0000000000..9fc620f401 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/squeezewave.yaml @@ -0,0 +1,104 @@ +name: "SqueezeWave" +sample_rate: 22050 +n_fft: 1024 +n_mels: 80 +fmax: 8000 +pad_value: -11.52 +train_dataset: ??? +validation_datasets: ??? + +model: + sigma: 1.0 + train_ds: + dataset: + # Use SplicedAudioDataset to use all of training data + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + n_segments: 16384 + trim: false + dataloader_params: + drop_last: true + pin_memory: true + shuffle: true + batch_size: 96 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + n_segments: -1 + trim: false + # This value is hard-coded to reflect n_window_stride * lcm(2, pad_to) + truncate_to: 512 + dataloader_params: + drop_last: false + shuffle: false + batch_size: 1 + num_workers: 4 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${n_mels} + frame_splicing: 1 + highfreq: ${fmax} + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${n_fft} + + # SqueezeWave is currently hardcoded to these values for window size and stride + # Changing these parameters are not recommended + n_window_size: 1024 + n_window_stride: 256 + + normalize: null + pad_to: 1 + pad_value: ${pad_value} + preemph: null + sample_rate: ${sample_rate} + window: hann + exact_pad: true + + squeezewave: + _target_: nemo.collections.tts.modules.squeezewave.SqueezeWaveModule + n_early_every: 2 + n_early_size: 16 + n_flows: 12 + n_group: 128 + n_mel_channels: ${n_mels} + n_wn_channels: 256 + n_wn_layers: 8 + wn_kernel_size: 3 + + optim: + name: adam + lr: 4e-4 + +trainer: + gpus: 1 # number of gpus + max_epochs: 10000 + max_steps: 600000 + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: false # Provided by exp_manager + logger: false # Provided by exp_manager + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + precision: 16 + amp_level: O1 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: true + create_checkpoint_callback: true diff --git a/models/tts/NeMo/examples/tts/conf/tacotron2.yaml b/models/tts/NeMo/examples/tts/conf/tacotron2.yaml new file mode 100644 index 0000000000..11b37cf4d1 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/tacotron2.yaml @@ -0,0 +1,140 @@ +name: Tacotron2 +sample_rate: 22050 +# , , will be added by the tacotron2.py script +labels: [' ', '!', '"', "'", '(', ')', ',', '-', '.', ':', ';', '?', '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', '[', ']', + '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'] +n_fft: 1024 +n_mels: 80 +fmax: 8000 +n_stride: 256 +pad_value: -11.52 +train_dataset: ??? +validation_datasets: ??? + +model: + labels: ${labels} + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + trim: false + int_values: false + normalize: true + sample_rate: ${sample_rate} + # bos_id: 66 + # eos_id: 67 + # pad_id: 68 These parameters are added automatically in Tacotron2 + dataloader_params: + drop_last: false + shuffle: true + batch_size: 48 + num_workers: 4 + + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + int_values: false + normalize: true + sample_rate: ${sample_rate} + trim: false + # bos_id: 66 + # eos_id: 67 + # pad_id: 68 These parameters are added automatically in Tacotron2 + dataloader_params: + drop_last: false + shuffle: false + batch_size: 48 + num_workers: 8 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${n_mels} + frame_splicing: 1 + highfreq: ${fmax} + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${n_fft} + n_window_size: 1024 + n_window_stride: ${n_stride} + normalize: null + pad_to: 16 + pad_value: ${pad_value} + preemph: null + sample_rate: ${sample_rate} + window: hann + + encoder: + _target_: nemo.collections.tts.modules.tacotron2.Encoder + encoder_kernel_size: 5 + encoder_n_convolutions: 3 + encoder_embedding_dim: 512 + + decoder: + _target_: nemo.collections.tts.modules.tacotron2.Decoder + decoder_rnn_dim: 1024 + encoder_embedding_dim: ${model.encoder.encoder_embedding_dim} + gate_threshold: 0.5 + max_decoder_steps: 1000 + n_frames_per_step: 1 # currently only 1 is supported + n_mel_channels: ${n_mels} + p_attention_dropout: 0.1 + p_decoder_dropout: 0.1 + prenet_dim: 256 + prenet_p_dropout: 0.5 + # Attention parameters + attention_dim: 128 + attention_rnn_dim: 1024 + # AttentionLocation Layer parameters + attention_location_kernel_size: 31 + attention_location_n_filters: 32 + early_stopping: true + + postnet: + _target_: nemo.collections.tts.modules.tacotron2.Postnet + n_mel_channels: ${n_mels} + p_dropout: 0.5 + postnet_embedding_dim: 512 + postnet_kernel_size: 5 + postnet_n_convolutions: 5 + + optim: + name: adam + lr: 1e-3 + weight_decay: 1e-6 + + # scheduler setup + sched: + name: CosineAnnealing + min_lr: 1e-5 + + +trainer: + gpus: 1 # number of gpus + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + gradient_clip_val: 1.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/talknet-durs.yaml b/models/tts/NeMo/examples/tts/conf/talknet-durs.yaml new file mode 100644 index 0000000000..aaa977d376 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/talknet-durs.yaml @@ -0,0 +1,175 @@ +name: TalkNetDurs +sample_rate: 22050 + +train_dataset: ??? +validation_datasets: ??? +durs_file: ??? +f0_file: ??? + +model: + d_char: 128 + dropout: 0.1 + separable: true + + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + int_values: false + load_audio: false + normalize: false + sample_rate: ${sample_rate} + trim: false + durs_file: ${durs_file} + f0_file: ${f0_file} + blanking: true + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: last + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + int_values: false + load_audio: false + normalize: false + sample_rate: ${sample_rate} + trim: false + durs_file: ${durs_file} + f0_file: ${f0_file} + blanking: true + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: last + dataloader_params: + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 1 + + model: + _target_: "nemo.collections.asr.modules.ConvASREncoder" + feat_in: ${model.d_char} + activation: relu + conv_mask: true + + jasper: + - filters: 256 + repeat: 1 + kernel: [3] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + + - filters: 256 + repeat: 5 + kernel: [5] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [7] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [9] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 512 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + + optim: + name: adam + lr: 1e-3 + weight_decay: 1e-6 + + sched: + name: CosineAnnealing + min_lr: 1e-5 + warmup_ratio: 0.02 + +trainer: + gpus: 1 + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False + logger: False + gradient_clip_val: 1.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_acc" + mode: "max" \ No newline at end of file diff --git a/models/tts/NeMo/examples/tts/conf/talknet-pitch.yaml b/models/tts/NeMo/examples/tts/conf/talknet-pitch.yaml new file mode 100644 index 0000000000..722c429045 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/talknet-pitch.yaml @@ -0,0 +1,179 @@ +name: TalkNetPitch +sample_rate: 22050 + +train_dataset: ??? +validation_datasets: ??? +durs_file: ??? +f0_file: ??? + +model: + d_char: 128 + dropout: 0.1 + separable: true + + # F0 stats (depends on dataset, current for LJSpeech) + f0_mean: 178.7767791748047 + f0_std: 33.51659393310547 + + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + int_values: false + load_audio: false + normalize: false + sample_rate: ${sample_rate} + trim: false + durs_file: ${durs_file} + f0_file: ${f0_file} + blanking: true + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: last + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + int_values: false + load_audio: false + normalize: false + sample_rate: ${sample_rate} + trim: false + durs_file: ${durs_file} + f0_file: ${f0_file} + blanking: true + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: last + dataloader_params: + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 1 + + model: + _target_: "nemo.collections.asr.modules.ConvASREncoder" + feat_in: ${model.d_char} + activation: relu + conv_mask: true + + jasper: + - filters: 256 + repeat: 1 + kernel: [3] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + + - filters: 256 + repeat: 5 + kernel: [5] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [7] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [9] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [11] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: True + se_reduction_ratio: 8 + + - filters: 512 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + + optim: + name: adam + lr: 1e-3 + weight_decay: 1e-6 + + sched: + name: CosineAnnealing + min_lr: 1e-5 + warmup_ratio: 0.02 + +trainer: + gpus: 1 + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False + logger: False + gradient_clip_val: 1.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_body_mae" + mode: "min" \ No newline at end of file diff --git a/models/tts/NeMo/examples/tts/conf/talknet-spect.yaml b/models/tts/NeMo/examples/tts/conf/talknet-spect.yaml new file mode 100644 index 0000000000..e476b4d62b --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/talknet-spect.yaml @@ -0,0 +1,244 @@ +name: TalkNetSpect +sample_rate: 22050 + +train_dataset: ??? +validation_datasets: ??? +durs_file: ??? +f0_file: ??? + +model: + d_char: 128 + n_mels: 80 + dropout: 0.1 + separable: true + + train_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + int_values: false + load_audio: true + normalize: false + sample_rate: ${sample_rate} + trim: false + durs_file: ${durs_file} + f0_file: ${f0_file} + blanking: true + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: last + dataloader_params: + drop_last: false + shuffle: true + batch_size: 64 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.asr.data.audio_to_text.AudioToCharWithDursF0Dataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + int_values: false + load_audio: true + normalize: false + sample_rate: ${sample_rate} + trim: false + durs_file: ${durs_file} + f0_file: ${f0_file} + blanking: true + vocab: + notation: phonemes + punct: true + spaces: true + stresses: false + add_blank_at: last + dataloader_params: + drop_last: false + shuffle: false + batch_size: 64 + num_workers: 1 + + preprocessor: + _target_: "nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor" + dither: 0.0 + features: ${model.n_mels} + frame_splicing: 1 + highfreq: 8000 + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: 1024 + n_window_size: 1024 + window_size: False + n_window_stride: 256 + window_stride: False + normalize: null + pad_to: 1 + pad_value: -11.52 + preemph: null + sample_rate: ${sample_rate} + stft_conv: false + window: hann + + model: + _target_: "nemo.collections.asr.modules.ConvASREncoder" + feat_in: ${model.d_char} + activation: relu + conv_mask: true + + jasper: + - filters: 256 + repeat: 3 + kernel: [3] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + + - filters: 256 + repeat: 5 + kernel: [5] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [7] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [9] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [13] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [15] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 256 + repeat: 5 + kernel: [17] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 512 + repeat: 5 + kernel: [21] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 512 + repeat: 5 + kernel: [23] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 512 + repeat: 5 + kernel: [25] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + separable: ${model.separable} + se: true + se_reduction_ratio: 8 + + - filters: 1024 + repeat: 1 + kernel: [1] + stride: [1] + dilation: [1] + dropout: ${model.dropout} + residual: true + + optim: + name: adam + lr: 1e-3 + weight_decay: 1e-6 + + sched: + name: CosineAnnealing + min_lr: 1e-5 + warmup_ratio: 0.02 + +trainer: + gpus: 1 + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False + logger: False + gradient_clip_val: 1.0 + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True + checkpoint_callback_params: + monitor: "val_loss" + mode: "min" diff --git a/models/tts/NeMo/examples/tts/conf/uniglow.yaml b/models/tts/NeMo/examples/tts/conf/uniglow.yaml new file mode 100644 index 0000000000..b74584d799 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/uniglow.yaml @@ -0,0 +1,96 @@ +name: "UniGlow" +sample_rate: 22050 +n_fft: 1024 +n_mels: 80 +fmax: 8000 +pad_value: -11.52 +train_dataset: ??? +validation_datasets: ??? +resume_from_ckpt: null + +model: + sigma: 1.0 + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + n_segments: 16384 + trim: false + dataloader_params: + drop_last: false + shuffle: true + batch_size: 12 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + n_segments: 49152 # 3 secs + trim: false + dataloader_params: + drop_last: false + shuffle: false + batch_size: 12 + num_workers: 4 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${n_mels} + frame_splicing: 1 + highfreq: ${fmax} + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${n_fft} + + # UniGlow is currently hardcoded to these values for window size and stride + # Changing these parameters are not recommended + n_window_size: 1024 + n_window_stride: 256 + + normalize: null + pad_to: 0 #16 + pad_value: ${pad_value} + preemph: null + sample_rate: ${sample_rate} + window: hann + + uniglow: + n_flows: 16 + n_group: 8 + n_mel_channels: ${n_mels} + n_wn_channels: 256 + n_wn_layers: 8 + wn_kernel_size: 3 + stft_loss_coef: 0.1 + + optim: + name: adam + lr: 1e-4 + +trainer: + gpus: 1 # number of gpus + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + precision: 16 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/conf/waveglow.yaml b/models/tts/NeMo/examples/tts/conf/waveglow.yaml new file mode 100644 index 0000000000..83f97a0474 --- /dev/null +++ b/models/tts/NeMo/examples/tts/conf/waveglow.yaml @@ -0,0 +1,97 @@ +name: "WaveGlow" +sample_rate: 22050 +n_fft: 1024 +n_mels: 80 +fmax: 8000 +pad_value: -11.52 +train_dataset: ??? +validation_datasets: ??? + +model: + sigma: 1.0 + train_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${train_dataset} + max_duration: null + min_duration: 0.1 + n_segments: 16000 + trim: false + dataloader_params: + drop_last: false + shuffle: true + batch_size: 12 + num_workers: 4 + + validation_ds: + dataset: + _target_: "nemo.collections.tts.data.datalayers.AudioDataset" + manifest_filepath: ${validation_datasets} + max_duration: null + min_duration: 0.1 + n_segments: -1 + trim: false + dataloader_params: + drop_last: false + shuffle: false + batch_size: 8 + num_workers: 4 + + preprocessor: + _target_: nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures + dither: 0.0 + nfilt: ${n_mels} + frame_splicing: 1 + highfreq: ${fmax} + log: true + log_zero_guard_type: clamp + log_zero_guard_value: 1e-05 + lowfreq: 0 + mag_power: 1.0 + n_fft: ${n_fft} + + # Waveglow is currently hardcoded to these values for window size and stride + # Changing these parameters are not recommended + n_window_size: 1024 + n_window_stride: 256 + + normalize: null + pad_to: 16 + pad_value: ${pad_value} + preemph: null + sample_rate: ${sample_rate} + window: hann + + waveglow: + _target_: nemo.collections.tts.modules.waveglow.WaveGlowModule + n_early_every: 4 + n_early_size: 2 + n_flows: 12 + n_group: 8 + n_mel_channels: ${n_mels} + n_wn_channels: 256 + n_wn_layers: 8 + wn_kernel_size: 3 + + optim: + name: adam + lr: 1e-4 + +trainer: + gpus: 1 # number of gpus + max_epochs: ??? + num_nodes: 1 + accelerator: ddp + accumulate_grad_batches: 1 + checkpoint_callback: False # Provided by exp_manager + logger: False # Provided by exp_manager + flush_logs_every_n_steps: 1000 + log_every_n_steps: 200 + check_val_every_n_epoch: 25 + precision: 16 + +exp_manager: + exp_dir: null + name: ${name} + create_tensorboard_logger: True + create_checkpoint_callback: True diff --git a/models/tts/NeMo/examples/tts/degli.py b/models/tts/NeMo/examples/tts/degli.py new file mode 100755 index 0000000000..b38b4eb159 --- /dev/null +++ b/models/tts/NeMo/examples/tts/degli.py @@ -0,0 +1,36 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.data.datalayers import preprocess_linear_specs_dataset +from nemo.collections.tts.models import DegliModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="degli") +def main(cfg): + preprocess_linear_specs_dataset(**cfg.model.preprocessor) + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = DegliModel(cfg=cfg.model, trainer=trainer) + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/ed_mel2spec.py b/models/tts/NeMo/examples/tts/ed_mel2spec.py new file mode 100755 index 0000000000..8815018407 --- /dev/null +++ b/models/tts/NeMo/examples/tts/ed_mel2spec.py @@ -0,0 +1,36 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.data.datalayers import preprocess_linear_specs_dataset +from nemo.collections.tts.models import EDMel2SpecModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="ed_mel2spec") +def main(cfg): + preprocess_linear_specs_dataset(**cfg.model.preprocessor) + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = EDMel2SpecModel(cfg=cfg.model, trainer=trainer) + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/fastpitch.py b/models/tts/NeMo/examples/tts/fastpitch.py new file mode 100644 index 0000000000..3afb305a99 --- /dev/null +++ b/models/tts/NeMo/examples/tts/fastpitch.py @@ -0,0 +1,35 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import FastPitchModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="fastpitch") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = FastPitchModel(cfg=cfg.model, trainer=trainer) + lr_logger = pl.callbacks.LearningRateMonitor() + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([lr_logger, epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/fastpitch_hifigan_e2e.py b/models/tts/NeMo/examples/tts/fastpitch_hifigan_e2e.py new file mode 100644 index 0000000000..f8b52537be --- /dev/null +++ b/models/tts/NeMo/examples/tts/fastpitch_hifigan_e2e.py @@ -0,0 +1,36 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from pytorch_lightning.plugins import DDPPlugin + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models.fastpitch_hifigan_e2e import FastPitchHifiGanE2EModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="fastpitch_hifigan_e2e") +def main(cfg): + trainer = pl.Trainer(plugins=[DDPPlugin(find_unused_parameters=True)], **cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = FastPitchHifiGanE2EModel(cfg=cfg.model, trainer=trainer) + lr_logger = pl.callbacks.LearningRateMonitor() + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([lr_logger, epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/fastspeech2.py b/models/tts/NeMo/examples/tts/fastspeech2.py new file mode 100644 index 0000000000..0b95dff3bb --- /dev/null +++ b/models/tts/NeMo/examples/tts/fastspeech2.py @@ -0,0 +1,35 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import FastSpeech2Model +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="fastspeech2") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = FastSpeech2Model(cfg=cfg.model, trainer=trainer) + lr_logger = pl.callbacks.LearningRateMonitor() + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([lr_logger, epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/fastspeech2_hifigan_e2e.py b/models/tts/NeMo/examples/tts/fastspeech2_hifigan_e2e.py new file mode 100644 index 0000000000..4122eefcc4 --- /dev/null +++ b/models/tts/NeMo/examples/tts/fastspeech2_hifigan_e2e.py @@ -0,0 +1,35 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models.fastspeech2_hifigan_e2e import FastSpeech2HifiGanE2EModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="fastspeech2_hifigan_e2e") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = FastSpeech2HifiGanE2EModel(cfg=cfg.model, trainer=trainer) + lr_logger = pl.callbacks.LearningRateMonitor() + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([lr_logger, epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/glow_tts.py b/models/tts/NeMo/examples/tts/glow_tts.py new file mode 100644 index 0000000000..7d797ff2fe --- /dev/null +++ b/models/tts/NeMo/examples/tts/glow_tts.py @@ -0,0 +1,35 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import GlowTTSModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="glow_tts") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = GlowTTSModel(cfg=cfg.model, trainer=trainer) + lr_logger = pl.callbacks.LearningRateMonitor() + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([lr_logger, epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/hifigan.py b/models/tts/NeMo/examples/tts/hifigan.py new file mode 100644 index 0000000000..5c3406a2f2 --- /dev/null +++ b/models/tts/NeMo/examples/tts/hifigan.py @@ -0,0 +1,31 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.tts.models import HifiGanModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf/hifigan", config_name="hifigan") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = HifiGanModel(cfg=cfg.model, trainer=trainer) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/melgan.py b/models/tts/NeMo/examples/tts/melgan.py new file mode 100644 index 0000000000..946f2b76e8 --- /dev/null +++ b/models/tts/NeMo/examples/tts/melgan.py @@ -0,0 +1,35 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl +from pytorch_lightning.plugins import DDPPlugin + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import MelGanModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="melgan") +def main(cfg): + trainer = pl.Trainer(plugins=[DDPPlugin(find_unused_parameters=True)], **cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = MelGanModel(cfg=cfg.model, trainer=trainer) + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/squeezewave.py b/models/tts/NeMo/examples/tts/squeezewave.py new file mode 100755 index 0000000000..d7a5ad8a1d --- /dev/null +++ b/models/tts/NeMo/examples/tts/squeezewave.py @@ -0,0 +1,34 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import SqueezeWaveModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="squeezewave") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = SqueezeWaveModel(cfg=cfg.model, trainer=trainer) + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/tacotron2.py b/models/tts/NeMo/examples/tts/tacotron2.py new file mode 100755 index 0000000000..a5446c35f7 --- /dev/null +++ b/models/tts/NeMo/examples/tts/tacotron2.py @@ -0,0 +1,44 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import Tacotron2Model +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +# hydra_runner is a thin NeMo wrapper around Hydra +# It looks for a config named tacotron2.yaml inside the conf folder +# Hydra parses the yaml and returns it as a Omegaconf DictConfig +@hydra_runner(config_path="conf", config_name="tacotron2") +def main(cfg): + # Define the Lightning trainer + trainer = pl.Trainer(**cfg.trainer) + # exp_manager is a NeMo construct that helps with logging and checkpointing + exp_manager(trainer, cfg.get("exp_manager", None)) + # Define the Tacotron 2 model, this will construct the model as well as + # define the training and validation dataloaders + model = Tacotron2Model(cfg=cfg.model, trainer=trainer) + # Let's add a few more callbacks + lr_logger = pl.callbacks.LearningRateMonitor() + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([lr_logger, epoch_time_logger]) + # Call lightning trainer's fit() to train the model + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/talknet_durs.py b/models/tts/NeMo/examples/tts/talknet_durs.py new file mode 100644 index 0000000000..f3543f69a2 --- /dev/null +++ b/models/tts/NeMo/examples/tts/talknet_durs.py @@ -0,0 +1,33 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import TalkNetDursModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path='conf', config_name='talknet-durs') +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get('exp_manager', None)) + model = TalkNetDursModel(cfg=cfg.model, trainer=trainer) + trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/talknet_pitch.py b/models/tts/NeMo/examples/tts/talknet_pitch.py new file mode 100644 index 0000000000..6836d63dc9 --- /dev/null +++ b/models/tts/NeMo/examples/tts/talknet_pitch.py @@ -0,0 +1,33 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import TalkNetPitchModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path='conf', config_name='talknet-pitch') +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get('exp_manager', None)) + model = TalkNetPitchModel(cfg=cfg.model, trainer=trainer) + trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/talknet_spect.py b/models/tts/NeMo/examples/tts/talknet_spect.py new file mode 100644 index 0000000000..619403c30f --- /dev/null +++ b/models/tts/NeMo/examples/tts/talknet_spect.py @@ -0,0 +1,33 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import TalkNetSpectModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path='conf', config_name='talknet-spect') +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get('exp_manager', None)) + model = TalkNetSpectModel(cfg=cfg.model, trainer=trainer) + trainer.callbacks.extend([pl.callbacks.LearningRateMonitor(), LogEpochTimeCallback()]) # noqa + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/test_tts_infer.py b/models/tts/NeMo/examples/tts/test_tts_infer.py new file mode 100644 index 0000000000..ce98fef1f3 --- /dev/null +++ b/models/tts/NeMo/examples/tts/test_tts_infer.py @@ -0,0 +1,153 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +This script is used as a CI test and shows how to chain TTS and ASR models +""" + +from argparse import ArgumentParser +from math import ceil +from pathlib import Path + +import librosa +import soundfile +import torch + +from nemo.collections.asr.metrics.wer import word_error_rate +from nemo.collections.asr.models import EncDecCTCModel +from nemo.collections.common.parts.preprocessing import parsers +from nemo.collections.tts.models.base import SpectrogramGenerator, Vocoder +from nemo.utils import logging + +LIST_OF_TEST_STRINGS = [ + "Hey, this is a test of the speech synthesis system.", + "roupell received the announcement with a cheerful countenance.", + "with thirteen dollars, eighty-seven cents when considerably greater resources were available to him.", + "Two other witnesses were able to offer partial descriptions of a man they saw in the southeast corner window.", + "'just to steady their legs a little' in other words, to add his weight to that of the hanging bodies.", + "The discussion above has already set forth examples of his expression of hatred for the United States.", + "At two:thirty-eight p.m., Eastern Standard Time, Lyndon Baines Johnson took the oath of office as the thirty-sixth President of the United States.", + "or, quote, other high government officials in the nature of a complaint coupled with an expressed or implied determination to use a means.", + "As for my return entrance visa please consider it separately. End quote.", + "it appears that Marina Oswald also complained that her husband was not able to provide more material things for her.", + "appeared in The Dallas Times Herald on November fifteen, nineteen sixty-three.", + "The only exit from the office in the direction Oswald was moving was through the door to the front stairway.", +] + + +def main(): + parser = ArgumentParser() + parser.add_argument( + "--asr_model", + type=str, + default="QuartzNet15x5Base-En", + choices=[x.pretrained_model_name for x in EncDecCTCModel.list_available_models()], + ) + parser.add_argument( + "--tts_model_spec", + type=str, + default="tts_en_tacotron2", + choices=[x.pretrained_model_name for x in SpectrogramGenerator.list_available_models()], + ) + parser.add_argument( + "--tts_model_vocoder", + type=str, + default="tts_waveglow_88m", + choices=[x.pretrained_model_name for x in Vocoder.list_available_models()], + ) + parser.add_argument("--wer_tolerance", type=float, default=1.0, help="used by test") + parser.add_argument("--trim", action="store_true") + parser.add_argument("--debug", action="store_true") + args = parser.parse_args() + torch.set_grad_enabled(False) + + if args.debug: + logging.set_verbosity(logging.DEBUG) + + logging.info(f"Using NGC cloud ASR model {args.asr_model}") + asr_model = EncDecCTCModel.from_pretrained(model_name=args.asr_model) + logging.info(f"Using NGC cloud TTS Spectrogram Generator model {args.tts_model_spec}") + tts_model_spec = SpectrogramGenerator.from_pretrained(model_name=args.tts_model_spec) + logging.info(f"Using NGC cloud TTS Vocoder model {args.tts_model_vocoder}") + tts_model_vocoder = Vocoder.from_pretrained(model_name=args.tts_model_vocoder) + models = [asr_model, tts_model_spec, tts_model_vocoder] + + if torch.cuda.is_available(): + for i, m in enumerate(models): + models[i] = m.cuda() + for m in models: + m.eval() + + asr_model, tts_model_spec, tts_model_vocoder = models + + parser = parsers.make_parser( + labels=asr_model.decoder.vocabulary, name="en", unk_id=-1, blank_id=-1, do_normalize=True, + ) + labels_map = dict([(i, asr_model.decoder.vocabulary[i]) for i in range(len(asr_model.decoder.vocabulary))]) + + tts_input = [] + asr_references = [] + longest_tts_input = 0 + for test_str in LIST_OF_TEST_STRINGS: + tts_parsed_input = tts_model_spec.parse(test_str) + if len(tts_parsed_input[0]) > longest_tts_input: + longest_tts_input = len(tts_parsed_input[0]) + tts_input.append(tts_parsed_input.squeeze()) + + asr_parsed = parser(test_str) + asr_parsed = ''.join([labels_map[c] for c in asr_parsed]) + asr_references.append(asr_parsed) + + # Pad TTS Inputs + for i, text in enumerate(tts_input): + pad = (0, longest_tts_input - len(text)) + tts_input[i] = torch.nn.functional.pad(text, pad, value=68) + + logging.debug(tts_input) + + # Do TTS + tts_input = torch.stack(tts_input) + if torch.cuda.is_available(): + tts_input = tts_input.cuda() + specs = tts_model_spec.generate_spectrogram(tokens=tts_input) + audio = [] + step = ceil(len(specs) / 4) + for i in range(4): + audio.append(tts_model_vocoder.convert_spectrogram_to_audio(spec=specs[i * step : i * step + step])) + + audio = [item for sublist in audio for item in sublist] + audio_file_paths = [] + # Save audio + logging.debug(f"args.trim: {args.trim}") + for i, aud in enumerate(audio): + aud = aud.cpu().numpy() + if args.trim: + aud = librosa.effects.trim(aud, top_db=40)[0] + soundfile.write(f"{i}.wav", aud, samplerate=22050) + audio_file_paths.append(str(Path(f"{i}.wav"))) + + # Do ASR + hypotheses = asr_model.transcribe(audio_file_paths) + for i, _ in enumerate(hypotheses): + logging.debug(f"{i}") + logging.debug(f"ref:'{asr_references[i]}'") + logging.debug(f"hyp:'{hypotheses[i]}'") + wer_value = word_error_rate(hypotheses=hypotheses, references=asr_references) + if wer_value > args.wer_tolerance: + raise ValueError(f"Got WER of {wer_value}. It was higher than {args.wer_tolerance}") + logging.info(f'Got WER of {wer_value}. Tolerance was {args.wer_tolerance}') + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/uniglow.py b/models/tts/NeMo/examples/tts/uniglow.py new file mode 100755 index 0000000000..39a4be5987 --- /dev/null +++ b/models/tts/NeMo/examples/tts/uniglow.py @@ -0,0 +1,44 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import UniGlowModel +from nemo.core.config import hydra_runner +from nemo.utils import logging +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="uniglow") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([epoch_time_logger]) + exp_manager(trainer, cfg.get("exp_manager", None)) + + if cfg.resume_from_ckpt is None: + logging.info("Training UniGlow from scratch") + model = UniGlowModel(cfg=cfg.model, trainer=trainer) + else: + logging.info("Fine-tuning UniGlow from {cfg.resume_from_ckpt}") + model = UniGlowModel.restore_from(cfg.resume_from_ckpt) + model.setup_training_data(cfg.model.train_ds) + model.setup_validation_data(cfg.model.validation_ds) + + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/examples/tts/waveglow.py b/models/tts/NeMo/examples/tts/waveglow.py new file mode 100755 index 0000000000..66b13491ab --- /dev/null +++ b/models/tts/NeMo/examples/tts/waveglow.py @@ -0,0 +1,34 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytorch_lightning as pl + +from nemo.collections.common.callbacks import LogEpochTimeCallback +from nemo.collections.tts.models import WaveGlowModel +from nemo.core.config import hydra_runner +from nemo.utils.exp_manager import exp_manager + + +@hydra_runner(config_path="conf", config_name="waveglow") +def main(cfg): + trainer = pl.Trainer(**cfg.trainer) + exp_manager(trainer, cfg.get("exp_manager", None)) + model = WaveGlowModel(cfg=cfg.model, trainer=trainer) + epoch_time_logger = LogEpochTimeCallback() + trainer.callbacks.extend([epoch_time_logger]) + trainer.fit(model) + + +if __name__ == '__main__': + main() # noqa pylint: disable=no-value-for-parameter diff --git a/models/tts/NeMo/external/get_collections.py b/models/tts/NeMo/external/get_collections.py new file mode 100644 index 0000000000..803db2afd2 --- /dev/null +++ b/models/tts/NeMo/external/get_collections.py @@ -0,0 +1,90 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" Script responsible for generation of a JSON file with list of NeMo collections. """ + +import argparse +import importlib +import json +import os + +import nemo +from nemo.utils import logging + + +def process_collection(id, col): + """ Helper function processing the collection. + + Args: + id: (short) name of the collection. + col: a collection (python module). + """ + return { + "id": id, + "name": col.__name__, + "description": col.__description__, + "version": col.__version__, + "author": col.__author__, + } + + +def main(): + """ Main function generating a JSON file with list of NeMo collections. """ + # Parse filename. + parser = argparse.ArgumentParser() + parser.add_argument('--filename', help='Name of the output JSON file', type=str, default="collections.json") + args = parser.parse_args() + + # Get collections directory. + colletions_dir = os.path.dirname(nemo.collections.__file__) + logging.info('Analysing collections in `{}`'.format(colletions_dir)) + + # Generate list of NeMo collections - from the list of collection subfolders. + collections = {} + for sub_dir in os.listdir(colletions_dir): + # Skip cache. + if sub_dir == "__pycache__": + continue + # Check if it is a directory. + if os.path.isdir(os.path.join(colletions_dir, sub_dir)): + collections[sub_dir] = "nemo.collections." + sub_dir + + output_list = [] + # Iterate over all collections. + for key, val in collections.items(): + # Try to get module specification. + module_spec = importlib.util.find_spec(val) + if module_spec is None: + logging.warning(" * Failed to process `{}`".format(val)) + else: + try: + # Import the module from the module specification. + module = importlib.util.module_from_spec(module_spec) + module_spec.loader.exec_module(module) + # Add to list. + output_list.append(process_collection(key, module)) + logging.info(" * Processed `{}`".format(val)) + except AttributeError: + logging.warning(" * Failed to process `{}`".format(val)) + + # Export to JSON. + with open(args.filename, 'w') as outfile: + json.dump(output_list, outfile) + + logging.info('Finshed the analysis, results exported to `{}`.'.format(args.filename)) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/external/get_modules.py b/models/tts/NeMo/external/get_modules.py new file mode 100644 index 0000000000..ba8f2f60d8 --- /dev/null +++ b/models/tts/NeMo/external/get_modules.py @@ -0,0 +1,159 @@ +# -*- coding: utf-8 -*- +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" Script responsible for generation of a JSON file containing list of modules of a given collection. """ + +import argparse +import importlib +import inspect +import json +import os + +import nemo +from nemo.utils import logging + + +def process_member(name, obj, module_list): + """ Helper function processing the passed object and, if ok, adding a record to the module list. + + Args: + name: name of the member + obj: member (class/function etc.) + module_list: list of modules that (probably) will be expanded. + """ + # It is not a class - skip it. + if not inspect.isclass(obj): + return + + # Check inheritance - we know that all our datasets/modules/losses inherit from Serialization, + # Btw. Serialization is also required by this script. + if not issubclass(obj, nemo.core.Serialization): + return + + logging.info(" * Processing `{}`".format(str(obj))) + + module_list.append( + { + "name": name, + "cls": str(obj), + # Temporary solution: mockup arguments. + "arguments": [ + "jasper", + "activation", + "feat_in", + "normalization_mode", + "residual_mode", + "norm_groups", + "conv_mask", + "frame_splicing", + "init_mode", + ], + # Temporary solution: mockup input types. + "input_types": { + "audio_signal": "axes: (batch, dimension, time); elements_type: MelSpectrogramType", + "length": "axes: (batch,); elements_type: LengthType", + }, + # Temporary solution: mockup output types. + "output_types": { + "encoder_output": "axes: (batch, dimension, time); elements_type: AcousticEncodedRepresentation" + }, + } + ) + + +def main(): + """ Main function analysing the indicated NeMo collection and generating a JSON file with module descriptions. """ + # Parse filename. + parser = argparse.ArgumentParser() + parser.add_argument('--collection', help='ID of the collection', type=str) + parser.add_argument('--filename', help='Name of the output JSON file', type=str, default="modules.json") + args = parser.parse_args() + + # Get collections directory. + colletions_dir = os.path.dirname(nemo.collections.__file__) + logging.info('Analysing collections in `{}`'.format(colletions_dir)) + + # Generate list of NeMo collections - from the list of collection subfolders. + collections = {} + for sub_dir in os.listdir(colletions_dir): + # Skip cache. + if sub_dir == "__pycache__": + continue + # Check if it is a directory. + if os.path.isdir(os.path.join(colletions_dir, sub_dir)): + collections[sub_dir] = "nemo.collections." + sub_dir + + # Check the collection. + if args.collection not in collections.keys(): + logging.error("Coudn't process the incidated `{}` collection".format(args.collection)) + logging.info( + "Please select one of the existing collections using `--collection [{}]`".format("|".join(collections)) + ) + exit(-1) + + # Load the collection specification. + collection_spec = importlib.util.find_spec(collections[args.collection]) + if collection_spec is None: + logging.error("Failed to load the `{}` collection".format(val)) + + # Import the module from the module specification. + collection = importlib.util.module_from_spec(collection_spec) + collection_spec.loader.exec_module(collection) + + module_list = [] + # Iterate over the packages in the indicated collection. + logging.info("Analysing the `{}` collection".format(args.collection)) + + try: # Datasets in dataset folder + logging.info("Analysing the 'data' package") + for name, obj in inspect.getmembers(collection.data): + process_member(name, obj, module_list) + except AttributeError as e: + logging.info(" * No datasets found") + + try: # Datasets in dataset folder + logging.info("Analysing the 'datasets' package") + for name, obj in inspect.getmembers(collection.datasets): + process_member(name, obj, module_list) + except AttributeError as e: + logging.info(" * No datasets found") + + try: # Modules + logging.info("Analysing the 'modules' package") + for name, obj in inspect.getmembers(collection.modules): + process_member(name, obj, module_list) + except AttributeError as e: + logging.info(" * No modules found") + + try: # Losses + logging.info("Analysing the 'losses' package") + for name, obj in inspect.getmembers(collection.losses): + process_member(name, obj, module_list) + except AttributeError as e: + logging.info(" * No losses found") + + # Add prefix - only for default name. + filename = args.filename if args.filename != "modules.json" else args.collection + "_" + args.filename + # Export to JSON. + with open(filename, 'w') as outfile: + json.dump(module_list, outfile) + + logging.info( + 'Finished analysis of the `{}` collection, results exported to `{}`.'.format(args.collection, filename) + ) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/nemo_text_processing/README.md b/models/tts/NeMo/nemo_text_processing/README.md new file mode 100644 index 0000000000..fd0f0197c0 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/README.md @@ -0,0 +1,8 @@ +**nemo_text_processing** +========================== + +Introduction +------------ + +NeMo's `nemo_text_processing` is a Python package that is installed with the `nemo_toolkit`. +See [documentation](https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/) for details. \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/__init__.py b/models/tts/NeMo/nemo_text_processing/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/README.md b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/README.md new file mode 100644 index 0000000000..30d2f22309 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/README.md @@ -0,0 +1,12 @@ +Inverse Text Normalization system for english, e.g. `one hundred twenty three kilograms` -> `123 kg` +Offers prediction and evaluation on text normalization data, e.g. [Google text normalization dataset](https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish). + + +Install dependencies: +bash setup.sh + +Example prediction run: +python run_predict.py --input=`INPUT_FILE` --output=`OUTPUT_FILE` [--verbose] + +Example evaluation run: +python run_evaluate.py --input=./en_with_types/output-00001-of-00100 [--cat CATEGORY] [--filter] diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/__init__.py new file mode 100644 index 0000000000..51e4804c62 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/__init__.py @@ -0,0 +1,32 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from nemo_text_processing.inverse_text_normalization.taggers.tokenize_and_classify import ClassifyFst +from nemo_text_processing.inverse_text_normalization.verbalizers.verbalize import VerbalizeFst +from nemo_text_processing.inverse_text_normalization.verbalizers.verbalize_final import VerbalizeFinalFst + +from nemo.utils import logging + +try: + import pynini + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + logging.warning( + "`pynini` is not installed ! \n" + "Please run the `nemo_text_processing/setup.sh` script" + "prior to usage of this toolkit." + ) + + PYNINI_AVAILABLE = False diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/clean_eval_data.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/clean_eval_data.py new file mode 100644 index 0000000000..d9bc2fccb6 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/clean_eval_data.py @@ -0,0 +1,342 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser +from typing import List + +import regex as re +from nemo_text_processing.text_normalization.data_loader_utils import ( + EOS_TYPE, + Instance, + load_files, + training_data_to_sentences, +) + + +""" +This file is for evaluation purposes. +filter_loaded_data() cleans data (list of instances) for inverse text normalization. Filters and cleaners can be specified for each semiotic class individually. +For example, normalized text should only include characters and whitespace characters but no punctuation. + Cardinal unnormalized instances should contain at least one integer and all other characters are removed. +""" + + +class Filter: + """ + Filter class + + Args: + class_type: semiotic class used in dataset + process_func: function to transform text + filter_func: function to filter text + + """ + + def __init__(self, class_type: str, process_func: object, filter_func: object): + self.class_type = class_type + self.process_func = process_func + self.filter_func = filter_func + + def filter(self, instance: Instance) -> bool: + """ + filter function + + Args: + filters given instance with filter function + + Returns: True if given instance fulfills criteria or does not belong to class type + """ + if instance.token_type != self.class_type: + return True + return self.filter_func(instance) + + def process(self, instance: Instance) -> Instance: + """ + process function + + Args: + processes given instance with process function + + Returns: processed instance if instance belongs to expected class type or original instance + """ + if instance.token_type != self.class_type: + return instance + return self.process_func(instance) + + +def filter_cardinal_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_cardinal_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r"[^0-9]", "", un_normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_ordinal_1(instance: Instance) -> bool: + ok = re.search(r"(st|nd|rd|th)\s*$", instance.un_normalized) + return ok + + +def process_ordinal_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r"[,\s]", "", un_normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_decimal_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_decimal_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + un_normalized = re.sub(r",", "", un_normalized) + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_measure_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_measure_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r",", "", un_normalized) + un_normalized = re.sub(r"m2", "m²", un_normalized) + un_normalized = re.sub(r"(\d)([^\d.\s])", r"\1 \2", un_normalized) + normalized = re.sub(r"[^a-z\s]", "", normalized) + normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_money_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_money_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r",", "", un_normalized) + un_normalized = re.sub(r"a\$", r"$", un_normalized) + un_normalized = re.sub(r"us\$", r"$", un_normalized) + un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) + un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_time_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_time_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + un_normalized = re.sub(r": ", ":", un_normalized) + un_normalized = re.sub(r"(\d)\s?a\s?m\s?", r"\1 a.m.", un_normalized) + un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_plain_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_plain_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_punct_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_punct_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_date_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_date_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + un_normalized = re.sub(r",", "", un_normalized) + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_letters_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_letters_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_verbatim_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_verbatim_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_digit_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_digit_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_telephone_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_telephone_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_electronic_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_electronic_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_fraction_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_fraction_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_address_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_address_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +filters = [] +filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) +filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) +filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) +filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) +filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) + +filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) +filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) +filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) +filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) +filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) +filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) +filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) +filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) +filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) + + +def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Instance]: + """ + Filters list of instances + + Args: + data: list of instances + + Returns: filtered and transformed list of instances + """ + updates_instances = [] + for instance in data: + updated_instance = False + for fil in filters: + if fil.class_type == instance.token_type and fil.filter(instance): + instance = fil.process(instance) + updated_instance = True + if updated_instance: + if verbose: + print(instance) + updates_instances.append(instance) + return updates_instances + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') + parser.add_argument("--verbose", help="print filtered instances", action='store_true') + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + file_path = args.input + + print("Loading training data: " + file_path) + instance_list = load_files([file_path]) # List of instances + filtered_instance_list = filter_loaded_data(instance_list, args.verbose) + training_data_to_sentences(filtered_instance_list) diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/currency.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/currency.tsv new file mode 100644 index 0000000000..1c300bf774 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/currency.tsv @@ -0,0 +1,35 @@ +$ dollar +$ us dollar +$ united states dollar +£ pound +€ euro +₩ won +nzd new zealand dollar +rs rupee +chf swiss franc +dkk danish kroner +fim finnish markka +aed arab emirates dirham +¥ yen +czk czech koruna +mro mauritanian ouguiya +pkr pakistani rupee +crc costa rican colon +hk$ hong kong dollar +npr nepalese rupee +awg aruban florin +nok norwegian kroner +tzs tanzanian shilling +sek swedish kronor +cyp cypriot pound +r real +sar saudi riyal +cve cape verde escudo +rsd serbian dinar +dm german mark +shp saint helena pounds +php philippine peso +cad canadian dollar +ssp south sudanese pound +scr seychelles rupee +mvr maldivian rufiyaa \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/domain.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/domain.tsv new file mode 100644 index 0000000000..cb4bc3f8f4 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/domain.tsv @@ -0,0 +1 @@ +com com \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/server_name.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/server_name.tsv new file mode 100644 index 0000000000..c447cf4682 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/electronic/server_name.tsv @@ -0,0 +1,2 @@ +gmail g mail +nvidia n vidia \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/magnitudes.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/magnitudes.tsv new file mode 100644 index 0000000000..3bfa4849e8 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/magnitudes.tsv @@ -0,0 +1,4 @@ +k thousand +m million +b billion +t trillion \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/measurements.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/measurements.tsv new file mode 100644 index 0000000000..db22e2aee2 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/measurements.tsv @@ -0,0 +1,110 @@ +f fahrenheit +c celsius +km kilometer +m meter +cm centimeter +mm millimeter +ha hectare +mi mile +m² square meter +km² square kilometer +ft foot +% percent +hz hertz +kw kilowatt +hp horsepower +mg milligram +kg kilogram +ghz gigahertz +khz kilohertz +mhz megahertz +lb pound +v volt +h hour +mc mega coulomb +s second +nm nanometer +rpm revolution per minute +min minute +mA milli ampere +% per cent +kwh kilo watt hour +m³ cubic meter +mph mile per hour +tw tera watt +mv milli volt +mw megawatt +μm micrometer +" inch +tb terabyte +cc c c +g gram +da dalton +atm atmosphere +ω ohm +db decibel +ps peta second +oz ounce +hl hecto liter +μg microgram +pg petagram +gb gigabyte +kb kilobit +ev electron volt +mb megabyte +kb kilobyte +kbps kilobit per second +mbps megabit per second +st stone +kl kilo liter +tj tera joule +kv kilo volt +mv mega volt +kn kilonewton +mm megameter +au astronomical unit +yd yard +rad radian +lm lumen +hs hecto second +mol mole +gpa giga pascal +ml milliliter +gw gigawatt +ma mega ampere +kt knot +kgf kilogram force +ng nano gram +ns nanosecond +ms mega siemens +bar bar +gl giga liter +μs microsecond +da deci ampere +pa pascal +ds deci second +ms milli second +dm deci meter +dm³ cubic deci meter +amu atomic mass unit +mb megabit +mf mega farad +bq becquerel +pb petabit +mm² square millimeter +cm² square centimeter +sq mi square mile +sq ft square foot +kpa kilopascal +cd candela +tl tera liter +ms mega second +mpa megapascal +pm peta meter +pb peta byte +gwh giga watt hour +kcal kilo calory +gy gray +sv sievert +cwt hundredweight +cc c c diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/months.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/months.tsv new file mode 100644 index 0000000000..6e7903c0ba --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/months.tsv @@ -0,0 +1,12 @@ +january +february +march +april +may +june +july +august +september +october +november +december \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/digit.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/digit.tsv new file mode 100644 index 0000000000..fa329799db --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/digit.tsv @@ -0,0 +1,9 @@ +one 1 +two 2 +three 3 +four 4 +five 5 +six 6 +seven 7 +eight 8 +nine 9 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/hundred.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/hundred.tsv new file mode 100644 index 0000000000..10e2b9645a --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/hundred.tsv @@ -0,0 +1 @@ +hundred \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/teen.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/teen.tsv new file mode 100644 index 0000000000..8e60fa1d3c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/teen.tsv @@ -0,0 +1,10 @@ +ten 10 +eleven 11 +twelve 12 +thirteen 13 +fourteen 14 +fifteen 15 +sixteen 16 +seventeen 17 +eighteen 18 +nineteen 19 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/thousands.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/thousands.tsv new file mode 100644 index 0000000000..bf30542bb6 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/thousands.tsv @@ -0,0 +1,22 @@ +thousand +million +billion +trillion +quadrillion +quintillion +sextillion +septillion +octillion +nonillion +decillion +undecillion +duodecillion +tredecillion +quattuordecillion +quindecillion +sexdecillion +septendecillion +octodecillion +novemdecillion +vigintillion +centillion \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/ties.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/ties.tsv new file mode 100644 index 0000000000..d7ec82e114 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/ties.tsv @@ -0,0 +1,9 @@ +twenty 2 +thirty 3 +forty 4 +fourty 4 +fifty 5 +sixty 6 +seventy 7 +eighty 8 +ninety 9 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/zero.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/zero.tsv new file mode 100644 index 0000000000..c479272d40 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/numbers/zero.tsv @@ -0,0 +1 @@ +zero 0 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/digit.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/digit.tsv new file mode 100644 index 0000000000..4b7fc24f16 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/digit.tsv @@ -0,0 +1,9 @@ +first one +second two +third three +fourth four +fifth five +sixth sixth +seventh seven +eighth eight +ninth nine \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/teen.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/teen.tsv new file mode 100644 index 0000000000..496fefc086 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/ordinals/teen.tsv @@ -0,0 +1 @@ +twelfth twelve \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_suffix.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_suffix.tsv new file mode 100644 index 0000000000..b39a9ef9a9 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_suffix.tsv @@ -0,0 +1,8 @@ +p m p.m. +pm p.m. +p.m. +p.m p.m. +am a.m. +a.m. +a.m a.m. +a m a.m. \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_to.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_to.tsv new file mode 100644 index 0000000000..1b006ea3ea --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_to.tsv @@ -0,0 +1,12 @@ +one 12 +two 1 +three 2 +four 3 +five 4 +six 5 +seven 6 +eigh 7 +nine 8 +ten 9 +eleven 10 +twelve 11 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_zone.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_zone.tsv new file mode 100644 index 0000000000..3e0ade467a --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/time/time_zone.tsv @@ -0,0 +1,7 @@ +cst c s t +cet c e t +pst p s t +est e s t +pt p t +et e t +gmt g m t diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/whitelist.tsv b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/whitelist.tsv new file mode 100644 index 0000000000..92f3877d7e --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/data/whitelist.tsv @@ -0,0 +1,10 @@ +e.g. for example +i.e. that is +dr. doctor +mr. mister +mrs. misses +ms. miss +st. saint +7-eleven seven eleven +es3 e s three +s&p s and p diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/inverse_normalize.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/inverse_normalize.py new file mode 100644 index 0000000000..44f4d5ece9 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/inverse_normalize.py @@ -0,0 +1,71 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser +from typing import List + +from nemo_text_processing.inverse_text_normalization.taggers.tokenize_and_classify import ClassifyFst +from nemo_text_processing.inverse_text_normalization.verbalizers.verbalize_final import VerbalizeFinalFst +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.token_parser import TokenParser + + +class InverseNormalizer(Normalizer): + """ + Inverse normalizer that converts text from spoken to written form. Useful for ASR postprocessing. + Input is expected to have no punctuation and be lower cased. + """ + + def __init__(self): + self.tagger = ClassifyFst() + self.verbalizer = VerbalizeFinalFst() + self.parser = TokenParser() + + def inverse_normalize_list(self, texts: List[str], verbose=False) -> List[str]: + """ + NeMo inverse text normalizer + + Args: + texts: list of input strings + verbose: whether to print intermediate meta information + + Returns converted list of input strings + """ + return self.normalize_list(texts=texts, verbose=verbose) + + def inverse_normalize(self, text: str, verbose: bool) -> str: + """ + Main function. Inverse normalizes tokens from spoken to written form + e.g. twelve kilograms -> 12 kg + + Args: + text: string that may include semiotic classes + verbose: whether to print intermediate meta information + + Returns: written form + """ + return self.normalize(text=text, verbose=verbose) + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("input_string", help="input string", type=str) + parser.add_argument("--verbose", help="print info for debugging", action='store_true') + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + inverse_normalizer = InverseNormalizer() + print(inverse_normalizer.inverse_normalize(args.input_string, verbose=args.verbose)) diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/run_evaluate.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/run_evaluate.py new file mode 100644 index 0000000000..845486f322 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/run_evaluate.py @@ -0,0 +1,109 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser + +from nemo_text_processing.inverse_text_normalization.clean_eval_data import filter_loaded_data +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.data_loader_utils import ( + evaluate, + known_types, + load_files, + training_data_to_sentences, + training_data_to_tokens, +) + + +''' +Runs Evaluation on data in the format of : \t\t<`self` if trivial class or normalized text> +like the Google text normalization data https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish +''' + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--input", help="input file path", type=str) + parser.add_argument( + "--cat", + dest="category", + help="focus on class only (" + ", ".join(known_types) + ")", + type=str, + default=None, + choices=known_types, + ) + parser.add_argument("--filter", action='store_true', help="clean data for inverse normalization purposes") + return parser.parse_args() + + +if __name__ == "__main__": + # Example usage: + # python run_evaluate.py --input= --cat= --filter + args = parse_args() + file_path = args.input + inverse_normalizer = InverseNormalizer() + + print("Loading training data: " + file_path) + training_data = load_files([file_path]) + + if args.filter: + training_data = filter_loaded_data(training_data) + + if args.category is None: + print("Sentence level evaluation...") + sentences_un_normalized, sentences_normalized, _ = training_data_to_sentences(training_data) + print("- Data: " + str(len(sentences_normalized)) + " sentences") + sentences_prediction = inverse_normalizer.inverse_normalize_list(sentences_normalized) + print("- Denormalized. Evaluating...") + sentences_accuracy = evaluate( + preds=sentences_prediction, labels=sentences_un_normalized, input=sentences_normalized + ) + print("- Accuracy: " + str(sentences_accuracy)) + + print("Token level evaluation...") + tokens_per_type = training_data_to_tokens(training_data, category=args.category) + token_accuracy = {} + for token_type in tokens_per_type: + print("- Token type: " + token_type) + tokens_un_normalized, tokens_normalized = tokens_per_type[token_type] + print(" - Data: " + str(len(tokens_normalized)) + " tokens") + tokens_prediction = inverse_normalizer.inverse_normalize_list(tokens_normalized) + print(" - Denormalized. Evaluating...") + token_accuracy[token_type] = evaluate(tokens_prediction, tokens_un_normalized, input=tokens_normalized) + print(" - Accuracy: " + str(token_accuracy[token_type])) + token_count_per_type = {token_type: len(tokens_per_type[token_type][0]) for token_type in tokens_per_type} + token_weighted_accuracy = [ + token_count_per_type[token_type] * accuracy for token_type, accuracy in token_accuracy.items() + ] + print("- Accuracy: " + str(sum(token_weighted_accuracy) / sum(token_count_per_type.values()))) + + print(" - Total: " + str(sum(token_count_per_type.values())), '\n') + + for token_type in token_accuracy: + if token_type not in known_types: + raise ValueError("Unexpected token type: " + token_type) + + if args.category is None: + c1 = ['Class', 'sent level'] + known_types + c2 = ['Num Tokens', len(sentences_normalized)] + [ + token_count_per_type[known_type] if known_type in tokens_per_type else '0' for known_type in known_types + ] + c3 = ["Denormalization", sentences_accuracy] + [ + token_accuracy[known_type] if known_type in token_accuracy else '0' for known_type in known_types + ] + + for i in range(len(c1)): + print(f'{str(c1[i]):10s} | {str(c2[i]):10s} | {str(c3[i]):5s}') + else: + print(f'numbers\t{token_count_per_type[args.category]}') + print(f'Denormalization\t{token_accuracy[args.category]}') diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/run_predict.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/run_predict.py new file mode 100644 index 0000000000..39ce2ca40d --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/run_predict.py @@ -0,0 +1,75 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser +from typing import List + +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer + + +''' +Runs inverse text normalization on text data +''' + + +def load_file(file_path: str) -> List[str]: + """ + Load given text file into list of string. + + Args: + file_path: file path + + Returns: flat list of string + """ + res = [] + with open(file_path, 'r') as fp: + for line in fp: + res.append(line) + return res + + +def write_file(file_path: str, data: List[str]): + """ + Writes out list of string to file. + + Args: + file_path: file path + data: list of string + + """ + with open(file_path, 'w') as fp: + for line in data: + fp.write(line + '\n') + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--input", help="input file path", required=True, type=str) + parser.add_argument("--output", help="output file path", required=True, type=str) + parser.add_argument("--verbose", help="print denormalization info. For debugging", action='store_true') + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + file_path = args.input + inverse_normalizer = InverseNormalizer() + + print("Loading data: " + file_path) + data = load_file(file_path) + + print("- Data: " + str(len(data)) + " sentences") + inverse_normalizer_prediction = inverse_normalizer.inverse_normalize_list(data, verbose=args.verbose) + write_file(args.output, inverse_normalizer_prediction) + print(f"- Denormalized. Writing out to {args.output}") diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/cardinal.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/cardinal.py new file mode 100644 index 0000000000..20c78af2d8 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/cardinal.py @@ -0,0 +1,134 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path, num_to_word +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_DIGIT, + NEMO_SIGMA, + NEMO_SPACE, + GraphFst, + delete_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class CardinalFst(GraphFst): + """ + Finite state transducer for classifying cardinals + e.g. minus twenty three -> cardinal { integer: "23" negative: "-" } } + Numbers below thirteen are not converted. + """ + + def __init__(self): + super().__init__(name="cardinal", kind="classify") + graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) + graph_ties = pynini.string_file(get_abs_path("data/numbers/ties.tsv")) + graph_teen = pynini.string_file(get_abs_path("data/numbers/teen.tsv")) + + graph_hundred = pynini.cross("hundred", "") + + graph_hundred_component = pynini.union(graph_digit + delete_space + graph_hundred, pynutil.insert("0")) + graph_hundred_component += delete_space + graph_hundred_component += pynini.union( + graph_teen | pynutil.insert("00"), + (graph_ties | pynutil.insert("0")) + delete_space + (graph_digit | pynutil.insert("0")), + ) + + graph_hundred_component_at_least_one_none_zero_digit = graph_hundred_component @ ( + pynini.closure(NEMO_DIGIT) + (NEMO_DIGIT - "0") + pynini.closure(NEMO_DIGIT) + ) + self.graph_hundred_component_at_least_one_none_zero_digit = ( + graph_hundred_component_at_least_one_none_zero_digit + ) + + graph_thousands = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("thousand"), + pynutil.insert("000", weight=0.1), + ) + + graph_million = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("million"), + pynutil.insert("000", weight=0.1), + ) + graph_billion = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("billion"), + pynutil.insert("000", weight=0.1), + ) + graph_trillion = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("trillion"), + pynutil.insert("000", weight=0.1), + ) + graph_quadrillion = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("quadrillion"), + pynutil.insert("000", weight=0.1), + ) + graph_quintillion = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("quintillion"), + pynutil.insert("000", weight=0.1), + ) + graph_sextillion = pynini.union( + graph_hundred_component_at_least_one_none_zero_digit + delete_space + pynutil.delete("sextillion"), + pynutil.insert("000", weight=0.1), + ) + + graph = pynini.union( + graph_sextillion + + delete_space + + graph_quintillion + + delete_space + + graph_quadrillion + + delete_space + + graph_trillion + + delete_space + + graph_billion + + delete_space + + graph_million + + delete_space + + graph_thousands + + delete_space + + graph_hundred_component, + graph_zero, + ) + + graph = graph @ pynini.union( + pynutil.delete(pynini.closure("0")) + pynini.difference(NEMO_DIGIT, "0") + pynini.closure(NEMO_DIGIT), "0" + ) + + labels_exception = [num_to_word(x) for x in range(0, 13)] + graph_exception = pynini.union(*labels_exception) + + graph = pynini.cdrewrite(pynutil.delete("and"), NEMO_SPACE, NEMO_SPACE, NEMO_SIGMA) @ graph + + self.graph_no_exception = graph + + self.graph = (pynini.project(graph, "input") - graph_exception.arcsort()) @ graph + + optional_minus_graph = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("minus", "\"-\"") + NEMO_SPACE, 0, 1 + ) + + final_graph = optional_minus_graph + pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/date.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/date.py new file mode 100644 index 0000000000..60088db97d --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/date.py @@ -0,0 +1,163 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_ALPHA, + NEMO_DIGIT, + GraphFst, + delete_extra_space, + delete_space, +) + +try: + import pynini + from pynini.lib import pynutil + + graph_teen = pynini.string_file(get_abs_path("data/numbers/teen.tsv")).optimize() + graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")).optimize() + ties_graph = pynini.string_file(get_abs_path("data/numbers/ties.tsv")).optimize() + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + graph_teen = None + graph_digit = None + ties_graph = None + + PYNINI_AVAILABLE = True + + +def _get_month_graph(): + """ + Transducer for month, e.g. march -> march + """ + month_graph = pynini.string_file(get_abs_path("data/months.tsv")) + return month_graph + + +def _get_ties_graph(): + """ + Transducer for 20-99 e.g + twenty three -> 23 + """ + graph = ties_graph + (delete_space + graph_digit | pynutil.insert("0")) + return graph + + +def _get_range_graph(): + """ + Transducer for decades (1**0s, 2**0s), centuries (2*00s, 1*00s), millennia (2000s) + """ + graph_ties = _get_ties_graph() + graph = (graph_ties | graph_teen) + delete_space + pynini.cross("hundreds", "00s") + graph |= pynini.cross("two", "2") + delete_space + pynini.cross("thousands", "000s") + graph |= ( + (graph_ties | graph_teen) + + delete_space + + (pynini.closure(NEMO_ALPHA, 1) + (pynini.cross("ies", "y") | pynutil.delete("s"))) + @ (graph_ties | pynini.cross("ten", "10")) + + pynutil.insert("s") + ) + graph @= pynini.union("1", "2") + NEMO_DIGIT + NEMO_DIGIT + NEMO_DIGIT + "s" + return graph + + +def _get_year_graph(): + """ + Transducer for year, e.g. twenty twenty -> 2020 + """ + + def _get_digits_graph(): + zero = pynini.cross((pynini.accep("oh") | pynini.accep("o")), "0") + graph = zero + delete_space + graph_digit + graph.optimize() + return graph + + def _get_thousands_graph(): + graph_ties = _get_ties_graph() + graph_hundred_component = (graph_digit + delete_space + pynutil.delete("hundred")) | pynutil.insert("0") + graph = ( + graph_digit + + delete_space + + pynutil.delete("thousand") + + delete_space + + graph_hundred_component + + delete_space + + (graph_teen | graph_ties) + ) + return graph + + graph_ties = _get_ties_graph() + graph_digits = _get_digits_graph() + graph_thousands = _get_thousands_graph() + year_graph = ( + # 20 19, 40 12, 2012 - assuming no limit on the year + (graph_teen + delete_space + (graph_ties | graph_digits | graph_teen)) + | (graph_ties + delete_space + (graph_ties | graph_digits | graph_teen)) + | graph_thousands + ) + year_graph.optimize() + return year_graph + + +class DateFst(GraphFst): + """ + Finite state transducer for classifying date, + e.g. january fifth twenty twelve -> date { month: "january" day: "5" year: "2012" preserve_order: true } + e.g. the fifth of january twenty twelve -> date { day: "5" month: "january" year: "2012" preserve_order: true } + e.g. twenty twenty -> date { year: "2012" preserve_order: true } + + Args: + ordinal: OrdinalFst + """ + + def __init__(self, ordinal: GraphFst): + super().__init__(name="date", kind="classify") + + ordinal_graph = ordinal.graph + year_graph = _get_year_graph() + YEAR_WEIGHT = 0.001 + year_graph = pynutil.add_weight(year_graph, YEAR_WEIGHT) + month_graph = _get_month_graph() + + month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + + day_graph = pynutil.insert("day: \"") + pynutil.add_weight(ordinal_graph, -0.7) + pynutil.insert("\"") + optional_day_graph = pynini.closure(delete_extra_space + day_graph, 0, 1) + optional_graph_year = pynini.closure( + delete_extra_space + + pynutil.insert("year: \"") + + pynutil.add_weight(year_graph, -YEAR_WEIGHT) + + pynutil.insert("\""), + 0, + 1, + ) + graph_mdy = month_graph + optional_day_graph + optional_graph_year + graph_dmy = ( + pynutil.delete("the") + + delete_space + + day_graph + + delete_space + + pynutil.delete("of") + + delete_extra_space + + month_graph + + optional_graph_year + ) + graph_year = pynutil.insert("year: \"") + (year_graph | _get_range_graph()) + pynutil.insert("\"") + + final_graph = graph_mdy | graph_dmy | graph_year + final_graph += pynutil.insert(" preserve_order: true") + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/decimal.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/decimal.py new file mode 100644 index 0000000000..47017f4d36 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/decimal.py @@ -0,0 +1,95 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import NEMO_DIGIT, GraphFst, delete_extra_space, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': + """ + Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, + e.g. one million -> integer_part: "1" quantity: "million" + e.g. one point five million -> integer_part: "1" fractional_part: "5" quantity: "million" + + Args: + decimal: decimal FST + cardinal_up_to_hundred: cardinal FST + """ + numbers = cardinal_up_to_hundred @ ( + pynutil.delete(pynini.closure("0")) + pynini.difference(NEMO_DIGIT, "0") + pynini.closure(NEMO_DIGIT) + ) + suffix = pynini.union("million", "billion", "trillion", "quadrillion", "quintillion", "sextillion") + res = ( + pynutil.insert("integer_part: \"") + + numbers + + pynutil.insert("\"") + + delete_extra_space + + pynutil.insert("quantity: \"") + + suffix + + pynutil.insert("\"") + ) + res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "thousand") + pynutil.insert("\"") + return res + + +class DecimalFst(GraphFst): + """ + Finite state transducer for classifying decimal + e.g. minus twelve point five o o six billion -> decimal { negative: "true" integer_part: "12" fractional_part: "5006" quantity: "billion" } + e.g. one billion -> decimal { integer_part: "1" quantity: "billion" } + Args: + cardinal: CardinalFst + """ + + def __init__(self, cardinal: GraphFst): + super().__init__(name="decimal", kind="classify") + + cardinal_graph = cardinal.graph_no_exception + + graph_decimal = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) + graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) | pynini.cross("o", "0") + + graph_decimal = pynini.closure(graph_decimal + delete_space) + graph_decimal + self.graph = graph_decimal + + point = pynutil.delete("point") + + optional_graph_negative = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + ) + + graph_fractional = pynutil.insert("fractional_part: \"") + graph_decimal + pynutil.insert("\"") + graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + final_graph_wo_sign = ( + pynini.closure(graph_integer + delete_extra_space, 0, 1) + point + delete_extra_space + graph_fractional + ) + final_graph = optional_graph_negative + final_graph_wo_sign + + self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( + final_graph_wo_sign, cardinal.graph_hundred_component_at_least_one_none_zero_digit + ) + final_graph |= optional_graph_negative + get_quantity( + final_graph_wo_sign, cardinal.graph_hundred_component_at_least_one_none_zero_digit + ) + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/electronic.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/electronic.py new file mode 100644 index 0000000000..018efd6cf7 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/electronic.py @@ -0,0 +1,66 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import NEMO_ALPHA, GraphFst, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class ElectronicFst(GraphFst): + """ + Finite state transducer for classifying electronic: as URLs, email addresses, etc. + e.g. c d f one at a b c dot e d u -> tokens { electronic { username: "cdf1" domain: "abc.edu" } } + """ + + def __init__(self): + super().__init__(name="electronic", kind="classify") + + delete_extra_space = pynutil.delete(" ") + alpha_num = ( + NEMO_ALPHA + | pynini.string_file(get_abs_path("data/numbers/digit.tsv")) + | pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + ) + username = ( + pynutil.insert("username: \"") + + pynini.closure( + alpha_num + delete_extra_space + pynini.closure(pynini.cross("dot", '.') + delete_extra_space, 0, 1) + ) + + alpha_num + + pynutil.insert("\"") + ) + single_alphanum = pynini.closure(alpha_num + delete_extra_space) + alpha_num + server = single_alphanum | pynini.string_file(get_abs_path("data/electronic/server_name.tsv")).invert() + domain = single_alphanum | pynini.string_file(get_abs_path("data/electronic/domain.tsv")).invert() + domain_graph = ( + pynutil.insert("domain: \"") + + server + + delete_extra_space + + pynini.cross("dot", ".") + + delete_extra_space + + domain + + pynutil.insert("\"") + ) + graph = username + delete_extra_space + pynutil.delete("at") + insert_space + delete_extra_space + domain_graph + + final_graph = self.add_tokens(graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/fraction.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/fraction.py new file mode 100644 index 0000000000..6b2cd74948 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/fraction.py @@ -0,0 +1,26 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst + + +class FractionFst(GraphFst): + """ + Finite state transducer for classifying fraction + """ + + def __init__(self): + super().__init__(name="fraction", kind="classify") + # integer_part # numerator # denominator diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/measure.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/measure.py new file mode 100644 index 0000000000..9cb6755b05 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/measure.py @@ -0,0 +1,103 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_SIGMA, + GraphFst, + convert_space, + delete_extra_space, + delete_space, + get_singulars, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class MeasureFst(GraphFst): + """ + Finite state transducer for classifying measure + e.g. minus twelve kilograms -> measure { negative: "true" cardinal { integer: "12" } units: "kg" } + + Args: + cardinal: CardinalFst + decimal: DecimalFst + """ + + def __init__(self, cardinal: GraphFst, decimal: GraphFst): + super().__init__(name="measure", kind="classify") + + cardinal_graph = cardinal.graph_no_exception + + graph_unit = pynini.string_file(get_abs_path("data/measurements.tsv")) + graph_unit_singular = pynini.invert(graph_unit) # singular -> abbr + graph_unit_plural = get_singulars(graph_unit_singular) # plural -> abbr + + optional_graph_negative = pynini.closure( + pynutil.insert("negative: ") + pynini.cross("minus", "\"true\"") + delete_extra_space, 0, 1 + ) + + unit_singular = convert_space(graph_unit_singular) + unit_plural = convert_space(graph_unit_plural) + unit_misc = pynutil.insert("/") + pynutil.delete("per") + delete_space + convert_space(graph_unit_singular) + + unit_singular = ( + pynutil.insert("units: \"") + + (unit_singular | unit_misc | pynutil.add_weight(unit_singular + delete_space + unit_misc, 0.01)) + + pynutil.insert("\"") + ) + unit_plural = ( + pynutil.insert("units: \"") + + (unit_plural | unit_misc | pynutil.add_weight(unit_plural + delete_space + unit_misc, 0.01)) + + pynutil.insert("\"") + ) + + subgraph_decimal = ( + pynutil.insert("decimal { ") + + optional_graph_negative + + decimal.final_graph_wo_negative + + pynutil.insert(" }") + + delete_extra_space + + unit_plural + ) + subgraph_cardinal = ( + pynutil.insert("cardinal { ") + + optional_graph_negative + + pynutil.insert("integer: \"") + + ((NEMO_SIGMA - "one") @ cardinal_graph) + + pynutil.insert("\"") + + pynutil.insert(" }") + + delete_extra_space + + unit_plural + ) + subgraph_cardinal |= ( + pynutil.insert("cardinal { ") + + optional_graph_negative + + pynutil.insert("integer: \"") + + pynini.cross("one", "1") + + pynutil.insert("\"") + + pynutil.insert(" }") + + delete_extra_space + + unit_singular + ) + final_graph = subgraph_decimal | subgraph_cardinal + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/money.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/money.py new file mode 100644 index 0000000000..8787419fdb --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/money.py @@ -0,0 +1,112 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_DIGIT, + NEMO_SIGMA, + GraphFst, + convert_space, + delete_extra_space, + delete_space, + get_singulars, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class MoneyFst(GraphFst): + """ + Finite state transducer for classifying money + e.g. twelve dollars and five cents -> money { integer_part: "12" fractional_part: 05 currency: "$" } + + Args: + cardinal: CardinalFst + decimal: DecimalFst + """ + + def __init__(self, cardinal: GraphFst, decimal: GraphFst): + super().__init__(name="money", kind="classify") + # quantity, integer_part, fractional_part, currency + + cardinal_graph = cardinal.graph_no_exception + graph_decimal_final = decimal.final_graph_wo_negative + + unit = pynini.string_file(get_abs_path("data/currency.tsv")) + unit_singular = pynini.invert(unit) + unit_plural = get_singulars(unit_singular) + + graph_unit_singular = pynutil.insert("currency: \"") + convert_space(unit_singular) + pynutil.insert("\"") + graph_unit_plural = pynutil.insert("currency: \"") + convert_space(unit_plural) + pynutil.insert("\"") + + add_leading_zero_to_double_digit = (NEMO_DIGIT + NEMO_DIGIT) | (pynutil.insert("0") + NEMO_DIGIT) + # twelve dollars (and) fifty cents, zero cents + cents_standalone = ( + pynutil.insert("fractional_part: \"") + + pynini.union( + pynutil.add_weight(((NEMO_SIGMA - "one") @ cardinal_graph), -0.7) @ add_leading_zero_to_double_digit + + delete_space + + pynutil.delete("cents"), + pynini.cross("one", "01") + delete_space + pynutil.delete("cent"), + ) + + pynutil.insert("\"") + ) + + optional_cents_standalone = pynini.closure( + delete_space + + pynini.closure(pynutil.delete("and") + delete_space, 0, 1) + + insert_space + + cents_standalone, + 0, + 1, + ) + # twelve dollars fifty, only after integer + optional_cents_suffix = pynini.closure( + delete_extra_space + + pynutil.insert("fractional_part: \"") + + pynutil.add_weight(cardinal_graph @ add_leading_zero_to_double_digit, -0.7) + + pynutil.insert("\""), + 0, + 1, + ) + + graph_integer = ( + pynutil.insert("integer_part: \"") + + ((NEMO_SIGMA - "one") @ cardinal_graph) + + pynutil.insert("\"") + + delete_extra_space + + graph_unit_plural + + (optional_cents_standalone | optional_cents_suffix) + ) + graph_integer |= ( + pynutil.insert("integer_part: \"") + + pynini.cross("one", "1") + + pynutil.insert("\"") + + delete_extra_space + + graph_unit_singular + + (optional_cents_standalone | optional_cents_suffix) + ) + graph_decimal = graph_decimal_final + delete_extra_space + graph_unit_plural + graph_decimal |= pynutil.insert("currency: \"$\" integer_part: \"0\" ") + cents_standalone + final_graph = graph_integer | graph_decimal + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/ordinal.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/ordinal.py new file mode 100644 index 0000000000..54b48ae5dc --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/ordinal.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ImportError, ModuleNotFoundError): + PYNINI_AVAILABLE = False + + +class OrdinalFst(GraphFst): + """ + Finite state transducer for classifying ordinal + e.g. thirteenth -> ordinal { integer: "13" } + + Args: + cardinal: CardinalFst + """ + + def __init__(self, cardinal: GraphFst): + super().__init__(name="ordinal", kind="classify") + + cardinal_graph = cardinal.graph_no_exception + graph_digit = pynini.string_file(get_abs_path("data/ordinals/digit.tsv")) + graph_teens = pynini.string_file(get_abs_path("data/ordinals/teen.tsv")) + graph = pynini.closure(NEMO_CHAR) + pynini.union( + graph_digit, graph_teens, pynini.cross("tieth", "ty"), pynini.cross("th", "") + ) + + self.graph = graph @ cardinal_graph + final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/punctuation.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/punctuation.py new file mode 100644 index 0000000000..bd5a426b30 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/punctuation.py @@ -0,0 +1,41 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = False +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class PunctuationFst(GraphFst): + """ + Finite state transducer for classifying punctuation + e.g. a, -> tokens { name: "a" } tokens { name: "," } + """ + + def __init__(self): + super().__init__(name="punctuation", kind="classify") + + s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~" + punct = pynini.union(*s) + + graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/telephone.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/telephone.py new file mode 100644 index 0000000000..8705f0e752 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/telephone.py @@ -0,0 +1,67 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_space, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TelephoneFst(GraphFst): + """ + Finite state transducer for classifying telephone numbers, e.g. + one two three one two three five six seven eight -> { number_part: "123-123-5678" } + """ + + def __init__(self): + super().__init__(name="telephone", kind="classify") + delete_space = pynutil.delete(' ') + # country code, number_part, extension + add_separator = pynutil.insert(" ") # between components + digit = pynini.invert(pynini.string_file(get_abs_path("data/numbers/digit.tsv"))).optimize() | pynini.cross( + "0", pynini.union("o", "oh", "zero") + ) + + number_part = ( + ( + (pynini.closure(digit + insert_space, 2, 2) + digit + pynutil.delete("-")) + | ( + pynutil.delete("(") + + pynini.closure(digit + insert_space, 2, 2) + + digit + + pynutil.delete(")") + + pynini.closure(pynutil.delete("-"), 0, 1) + + delete_space + ) + ) + + add_separator + + pynini.closure(digit + insert_space, 2, 2) + + digit + + pynutil.delete("-") + + add_separator + + pynini.closure(digit + insert_space, 3, 3) + + digit + ) + number_part = pynutil.insert("number_part: \"") + pynini.invert(number_part) + pynutil.insert("\"") + + graph = number_part + final_graph = self.add_tokens(graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/time.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/time.py new file mode 100644 index 0000000000..1132ee9a13 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/time.py @@ -0,0 +1,124 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.inverse_text_normalization.taggers.cardinal import CardinalFst +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path, num_to_word +from nemo_text_processing.text_normalization.graph_utils import ( + GraphFst, + convert_space, + delete_extra_space, + delete_space, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TimeFst(GraphFst): + """ + Finite state transducer for classifying time + e.g. twelve thirty -> time { hours: "12" minutes: "30" } + e.g. twelve past one -> time { minutes: "12" hours: "1" } + e.g. two o clock a m -> time { hours: "2" suffix: "a.m." } + e.g. quarter to two -> time { hours: "1" minutes: "45" } + e.g. quarter past two -> time { hours: "2" minutes: "15" } + e.g. half past two -> time { hours: "2" minutes: "30" } + """ + + def __init__(self): + super().__init__(name="time", kind="classify") + # hours, minutes, seconds, suffix, zone, style, speak_period + + suffix_graph = pynini.string_file(get_abs_path("data/time/time_suffix.tsv")) + time_zone_graph = pynini.invert(pynini.string_file(get_abs_path("data/time/time_zone.tsv"))) + time_to_graph = pynini.string_file(get_abs_path("data/time/time_to.tsv")) + + # only used for < 1000 thousand -> 0 weight + cardinal = pynutil.add_weight(CardinalFst().graph_no_exception, weight=-0.7) + + labels_hour = [num_to_word(x) for x in range(0, 24)] + labels_minute_single = [num_to_word(x) for x in range(1, 10)] + labels_minute_double = [num_to_word(x) for x in range(10, 60)] + + graph_hour = pynini.union(*labels_hour) @ cardinal + + graph_minute_single = pynini.union(*labels_minute_single) @ cardinal + graph_minute_double = pynini.union(*labels_minute_double) @ cardinal + graph_minute_verbose = pynini.cross("half", "30") | pynini.cross("quarter", "15") + oclock = pynini.cross(pynini.union("o' clock", "o clock", "o'clock", "oclock"), "") + + final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_minute = ( + pynutil.insert("minutes: \"") + + ( + pynutil.insert("00") + | oclock + pynutil.insert("00") + | pynutil.delete("o") + delete_space + graph_minute_single + | graph_minute_double + ) + + pynutil.insert("\"") + ) + final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_suffix_optional = pynini.closure(delete_space + insert_space + final_suffix, 0, 1) + final_time_zone_optional = pynini.closure( + delete_space + + insert_space + + pynutil.insert("zone: \"") + + convert_space(time_zone_graph) + + pynutil.insert("\""), + 0, + 1, + ) + + # five o' clock + # two o eight, two thiry five (am/pm) + # two pm/am + graph_hm = final_graph_hour + delete_extra_space + final_graph_minute + # 10 past four, quarter past four, half past four + graph_mh = ( + pynutil.insert("minutes: \"") + + pynini.union(graph_minute_single, graph_minute_double, graph_minute_verbose) + + pynutil.insert("\"") + + delete_space + + pynutil.delete("past") + + delete_extra_space + + final_graph_hour + ) + + graph_quarter_time = ( + pynutil.insert("minutes: \"") + + pynini.cross("quarter", "45") + + pynutil.insert("\"") + + delete_space + + pynutil.delete(pynini.union("to", "till")) + + delete_extra_space + + pynutil.insert("hours: \"") + + time_to_graph + + pynutil.insert("\"") + ) + final_graph = ( + (graph_hm | graph_mh | graph_quarter_time) + final_suffix_optional + final_time_zone_optional + ).optimize() + + final_graph = self.add_tokens(final_graph) + + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/tokenize_and_classify.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/tokenize_and_classify.py new file mode 100644 index 0000000000..0c7718c499 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/tokenize_and_classify.py @@ -0,0 +1,91 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.taggers.cardinal import CardinalFst +from nemo_text_processing.inverse_text_normalization.taggers.date import DateFst +from nemo_text_processing.inverse_text_normalization.taggers.decimal import DecimalFst +from nemo_text_processing.inverse_text_normalization.taggers.electronic import ElectronicFst +from nemo_text_processing.inverse_text_normalization.taggers.measure import MeasureFst +from nemo_text_processing.inverse_text_normalization.taggers.money import MoneyFst +from nemo_text_processing.inverse_text_normalization.taggers.ordinal import OrdinalFst +from nemo_text_processing.inverse_text_normalization.taggers.punctuation import PunctuationFst +from nemo_text_processing.inverse_text_normalization.taggers.telephone import TelephoneFst +from nemo_text_processing.inverse_text_normalization.taggers.time import TimeFst +from nemo_text_processing.inverse_text_normalization.taggers.whitelist import WhiteListFst +from nemo_text_processing.inverse_text_normalization.taggers.word import WordFst +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_extra_space, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class ClassifyFst(GraphFst): + """ + Final class that composes all other classification grammars. This class can process an entire sentence, that is lower cased. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + More details to deployment at NeMo/tools/text_processing_deployment. + """ + + def __init__(self): + super().__init__(name="tokenize_and_classify", kind="classify") + + cardinal = CardinalFst() + cardinal_graph = cardinal.fst + + ordinal = OrdinalFst(cardinal) + ordinal_graph = ordinal.fst + + decimal = DecimalFst(cardinal) + decimal_graph = decimal.fst + + measure_graph = MeasureFst(cardinal=cardinal, decimal=decimal).fst + date_graph = DateFst(ordinal=ordinal).fst + word_graph = WordFst().fst + time_graph = TimeFst().fst + money_graph = MoneyFst(cardinal=cardinal, decimal=decimal).fst + whitelist_graph = WhiteListFst().fst + punct_graph = PunctuationFst().fst + electronic_graph = ElectronicFst().fst + telephone_graph = TelephoneFst().fst + + classify = ( + pynutil.add_weight(whitelist_graph, 1.01) + | pynutil.add_weight(time_graph, 1.1) + | pynutil.add_weight(date_graph, 1.09) + | pynutil.add_weight(decimal_graph, 1.1) + | pynutil.add_weight(measure_graph, 1.1) + | pynutil.add_weight(cardinal_graph, 1.1) + | pynutil.add_weight(ordinal_graph, 1.1) + | pynutil.add_weight(money_graph, 1.1) + | pynutil.add_weight(telephone_graph, 1.1) + | pynutil.add_weight(electronic_graph, 1.1) + | pynutil.add_weight(word_graph, 100) + ) + + punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") + token_plus_punct = ( + pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + ) + + graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) + graph = delete_space + graph + delete_space + + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/whitelist.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/whitelist.py new file mode 100644 index 0000000000..3ac3b9eb04 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/whitelist.py @@ -0,0 +1,40 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import GraphFst, convert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WhiteListFst(GraphFst): + """ + Finite state transducer for classifying whitelisted tokens + e.g. misses -> tokens { name: "mrs." } + This class has highest priority among all classifier grammars. Whitelisted tokens are defined and loaded from "data/whitelist.tsv". + """ + + def __init__(self): + super().__init__(name="whitelist", kind="classify") + + whitelist = pynini.string_file(get_abs_path("data/whitelist.tsv")).invert() + graph = pynutil.insert("name: \"") + convert_space(whitelist) + pynutil.insert("\"") + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/word.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/word.py new file mode 100644 index 0000000000..6be7360b78 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/taggers/word.py @@ -0,0 +1,36 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_SPACE, GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WordFst(GraphFst): + """ + Finite state transducer for classifying plain tokens, that do not belong to any special class. This can be considered as the default class. + e.g. sleep -> tokens { name: "sleep" } + """ + + def __init__(self): + super().__init__(name="word", kind="classify") + word = pynutil.insert("name: \"") + pynini.closure(NEMO_NOT_SPACE, 1) + pynutil.insert("\"") + self.fst = word.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/utils.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/utils.py new file mode 100644 index 0000000000..6a8bd48cda --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/utils.py @@ -0,0 +1,47 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os +from typing import Union + +import inflect + +_inflect = inflect.engine() + + +def num_to_word(x: Union[str, int]): + """ + converts integer to spoken representation + + Args + x: integer + + Returns: spoken representation + """ + if isinstance(x, int): + x = str(x) + x = _inflect.number_to_words(str(x)).replace("-", " ").replace(",", "") + return x + + +def get_abs_path(rel_path): + """ + Get absolute path + + Args: + rel_path: relative path to this file + + Returns absolute path + """ + return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/__init__.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/cardinal.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/cardinal.py new file mode 100644 index 0000000000..905da49895 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/cardinal.py @@ -0,0 +1,55 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class CardinalFst(GraphFst): + """ + Finite state transducer for verbalizing cardinal + e.g. cardinal { integer: "23" negative: "-" } -> -23 + """ + + def __init__(self): + super().__init__(name="cardinal", kind="verbalize") + optional_sign = pynini.closure( + pynutil.delete("negative:") + + delete_space + + pynutil.delete("\"") + + NEMO_NOT_QUOTE + + pynutil.delete("\"") + + delete_space, + 0, + 1, + ) + graph = ( + pynutil.delete("integer:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + self.numbers = graph + graph = optional_sign + graph + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/date.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/date.py new file mode 100644 index 0000000000..00ee9c5806 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/date.py @@ -0,0 +1,87 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_NOT_QUOTE, + GraphFst, + delete_extra_space, + delete_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class DateFst(GraphFst): + """ + Finite state transducer for verbalizing date, e.g. + date { month: "january" day: "5" year: "2012" preserve_order: true } -> february 5 2012 + date { day: "5" month: "january" year: "2012" preserve_order: true } -> 5 february 2012 + """ + + def __init__(self): + super().__init__(name="date", kind="verbalize") + month = ( + pynutil.delete("month:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + day = ( + pynutil.delete("day:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + year = ( + pynutil.delete("year:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + delete_space + + pynutil.delete("\"") + ) + + # month (day) year + graph_mdy = ( + month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + ) + + # (day) month year + graph_dmy = ( + pynini.closure(day + delete_extra_space, 0, 1) + month + pynini.closure(delete_extra_space + year, 0, 1) + ) + + optional_preserve_order = pynini.closure( + pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space + | pynutil.delete("field_order:") + + delete_space + + pynutil.delete("\"") + + NEMO_NOT_QUOTE + + pynutil.delete("\"") + + delete_space + ) + + final_graph = (graph_mdy | year | graph_dmy) + delete_space + optional_preserve_order + + delete_tokens = self.delete_tokens(final_graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/decimal.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/decimal.py new file mode 100644 index 0000000000..82645afc90 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/decimal.py @@ -0,0 +1,65 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class DecimalFst(GraphFst): + """ + Finite state transducer for verbalizing decimal, e.g. + decimal { negative: "true" integer_part: "12" fractional_part: "5006" quantity: "billion" } -> -12.5006 billion + """ + + def __init__(self): + super().__init__(name="decimal", kind="verbalize") + optionl_sign = pynini.closure(pynini.cross("negative: \"true\"", "-") + delete_space, 0, 1) + integer = ( + pynutil.delete("integer_part:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + optional_integer = pynini.closure(integer + delete_space, 0, 1) + fractional = ( + pynutil.insert(".") + + pynutil.delete("fractional_part:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + optional_fractional = pynini.closure(fractional + delete_space, 0, 1) + quantity = ( + pynutil.delete("quantity:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + optional_quantity = pynini.closure(pynutil.insert(" ") + quantity + delete_space, 0, 1) + graph = optional_integer + optional_fractional + optional_quantity + self.numbers = graph + graph = optionl_sign + graph + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/electronic.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/electronic.py new file mode 100644 index 0000000000..8a8cdf8e55 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/electronic.py @@ -0,0 +1,52 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class ElectronicFst(GraphFst): + """ + Finite state transducer for verbalizing electronic + e.g. tokens { electronic { username: "cdf1" domain: "abc.edu" } } -> cdf1@abc.edu + """ + + def __init__(self): + super().__init__(name="electronic", kind="verbalize") + user_name = ( + pynutil.delete("username:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + domain = ( + pynutil.delete("domain:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + + graph = user_name + delete_space + pynutil.insert("@") + domain + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/fraction.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/fraction.py new file mode 100644 index 0000000000..45a6bd03b5 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/fraction.py @@ -0,0 +1,25 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst + + +class FractionFst(GraphFst): + """ + Finite state transducer for verbalizing fraction, + """ + + def __init__(self): + super().__init__(name="fraction", kind="verbalize") diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/measure.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/measure.py new file mode 100644 index 0000000000..734beffa24 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/measure.py @@ -0,0 +1,69 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + + PYNINI_AVAILABLE = False + + +class MeasureFst(GraphFst): + """ + Finite state transducer for verbalizing measure, e.g. + measure { negative: "true" cardinal { integer: "12" } units: "kg" } -> -12 kg + + Args: + decimal: DecimalFst + cardinal: CardinalFst + """ + + def __init__(self, decimal: GraphFst, cardinal: GraphFst): + super().__init__(name="measure", kind="verbalize") + optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "-"), 0, 1) + unit = ( + pynutil.delete("units:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_CHAR - " ", 1) + + pynutil.delete("\"") + + delete_space + ) + graph_decimal = ( + pynutil.delete("decimal {") + + delete_space + + optional_sign + + delete_space + + decimal.numbers + + delete_space + + pynutil.delete("}") + ) + graph_cardinal = ( + pynutil.delete("cardinal {") + + delete_space + + optional_sign + + delete_space + + cardinal.numbers + + delete_space + + pynutil.delete("}") + ) + graph = (graph_cardinal | graph_decimal) + delete_space + pynutil.insert(" ") + unit + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/money.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/money.py new file mode 100644 index 0000000000..30dec90c76 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/money.py @@ -0,0 +1,48 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + + PYNINI_AVAILABLE = False + + +class MoneyFst(GraphFst): + """ + Finite state transducer for verbalizing money, e.g. + money { integer_part: "12" fractional_part: "05" currency: "$" } -> $12.05 + + Args: + decimal: DecimalFst + """ + + def __init__(self, decimal: GraphFst): + super().__init__(name="money", kind="verbalize") + unit = ( + pynutil.delete("currency:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_CHAR - " ", 1) + + pynutil.delete("\"") + ) + graph = unit + delete_space + decimal.numbers + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/ordinal.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/ordinal.py new file mode 100644 index 0000000000..8896067c1b --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/ordinal.py @@ -0,0 +1,64 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, NEMO_SIGMA, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class OrdinalFst(GraphFst): + """ + Finite state transducer for verbalizing ordinal, e.g. + ordinal { integer: "13" } -> 13th + """ + + def __init__(self): + super().__init__(name="ordinal", kind="verbalize") + graph = ( + pynutil.delete("integer:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + convert_eleven = pynini.cross("11", "11th") + convert_twelve = pynini.cross("12", "12th") + convert_thirteen = pynini.cross("13", "13th") + convert_one = pynini.cross("1", "1st") + convert_two = pynini.cross("2", "2nd") + convert_three = pynini.cross("3", "3rd") + convert_rest = pynutil.insert("th", weight=0.01) + + suffix = pynini.cdrewrite( + convert_eleven + | convert_twelve + | convert_thirteen + | convert_one + | convert_two + | convert_three + | convert_rest, + "", + "[EOS]", + NEMO_SIGMA, + ) + graph = graph @ suffix + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/telephone.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/telephone.py new file mode 100644 index 0000000000..5e6329e13c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/telephone.py @@ -0,0 +1,39 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TelephoneFst(GraphFst): + """ + Finite state transducer for verbalizing telephone, e.g. + telephone { number_part: "123-123-5678" } + -> 123-123-5678 + """ + + def __init__(self): + super().__init__(name="telephone", kind="verbalize") + + number_part = pynutil.delete("number_part: \"") + pynini.closure(NEMO_NOT_QUOTE, 1) + pynutil.delete("\"") + delete_tokens = self.delete_tokens(number_part) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/time.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/time.py new file mode 100644 index 0000000000..0319738f0e --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/time.py @@ -0,0 +1,87 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_CHAR, + NEMO_DIGIT, + GraphFst, + delete_space, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TimeFst(GraphFst): + """ + Finite state transducer for verbalizing time, e.g. + time { hours: "12" minutes: "30" } -> 12:30 + time { hours: "1" minutes: "12" } -> 01:12 + time { hours: "2" suffix: "a.m." } -> 02:00 a.m. + """ + + def __init__(self): + super().__init__(name="time", kind="verbalize") + add_leading_zero_to_double_digit = (NEMO_DIGIT + NEMO_DIGIT) | (pynutil.insert("0") + NEMO_DIGIT) + hour = ( + pynutil.delete("hours:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_DIGIT, 1) + + pynutil.delete("\"") + ) + minute = ( + pynutil.delete("minutes:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_DIGIT, 1) + + pynutil.delete("\"") + ) + suffix = ( + delete_space + + insert_space + + pynutil.delete("suffix:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_CHAR - " ", 1) + + pynutil.delete("\"") + ) + optional_suffix = pynini.closure(suffix, 0, 1) + zone = ( + delete_space + + insert_space + + pynutil.delete("zone:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_CHAR - " ", 1) + + pynutil.delete("\"") + ) + optional_zone = pynini.closure(zone, 0, 1) + graph = ( + hour @ add_leading_zero_to_double_digit + + delete_space + + pynutil.insert(":") + + (minute @ add_leading_zero_to_double_digit) + + optional_suffix + + optional_zone + ) + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/verbalize.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/verbalize.py new file mode 100644 index 0000000000..684a239d79 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/verbalize.py @@ -0,0 +1,62 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.verbalizers.cardinal import CardinalFst +from nemo_text_processing.inverse_text_normalization.verbalizers.date import DateFst +from nemo_text_processing.inverse_text_normalization.verbalizers.decimal import DecimalFst +from nemo_text_processing.inverse_text_normalization.verbalizers.electronic import ElectronicFst +from nemo_text_processing.inverse_text_normalization.verbalizers.measure import MeasureFst +from nemo_text_processing.inverse_text_normalization.verbalizers.money import MoneyFst +from nemo_text_processing.inverse_text_normalization.verbalizers.ordinal import OrdinalFst +from nemo_text_processing.inverse_text_normalization.verbalizers.telephone import TelephoneFst +from nemo_text_processing.inverse_text_normalization.verbalizers.time import TimeFst +from nemo_text_processing.inverse_text_normalization.verbalizers.whitelist import WhiteListFst +from nemo_text_processing.text_normalization.graph_utils import GraphFst + + +class VerbalizeFst(GraphFst): + """ + Composes other verbalizer grammars. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + More details to deployment at NeMo/tools/text_processing_deployment. + """ + + def __init__(self): + super().__init__(name="verbalize", kind="verbalize") + cardinal = CardinalFst() + cardinal_graph = cardinal.fst + ordinal_graph = OrdinalFst().fst + decimal = DecimalFst() + decimal_graph = decimal.fst + measure_graph = MeasureFst(decimal=decimal, cardinal=cardinal).fst + money_graph = MoneyFst(decimal=decimal).fst + time_graph = TimeFst().fst + date_graph = DateFst().fst + whitelist_graph = WhiteListFst().fst + telephone_graph = TelephoneFst().fst + electronic_graph = ElectronicFst().fst + graph = ( + time_graph + | date_graph + | money_graph + | measure_graph + | ordinal_graph + | decimal_graph + | cardinal_graph + | whitelist_graph + | telephone_graph + | electronic_graph + ) + self.fst = graph diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/verbalize_final.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/verbalize_final.py new file mode 100644 index 0000000000..e6f66103d7 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/verbalize_final.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.inverse_text_normalization.verbalizers.verbalize import VerbalizeFst +from nemo_text_processing.inverse_text_normalization.verbalizers.word import WordFst +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_extra_space, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class VerbalizeFinalFst(GraphFst): + """ + Finite state transducer that verbalizes an entire sentence, e.g. + tokens { name: "its" } tokens { time { hours: "12" minutes: "30" } } tokens { name: "now" } -> its 12:30 now + """ + + def __init__(self): + super().__init__(name="verbalize_final", kind="verbalize") + verbalize = VerbalizeFst().fst + word = WordFst().fst + types = verbalize | word + graph = ( + pynutil.delete("tokens") + + delete_space + + pynutil.delete("{") + + delete_space + + types + + delete_space + + pynutil.delete("}") + ) + graph = delete_space + pynini.closure(graph + delete_extra_space) + graph + delete_space + self.fst = graph diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/whitelist.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/whitelist.py new file mode 100644 index 0000000000..4bea2e50ab --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/whitelist.py @@ -0,0 +1,44 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, NEMO_SIGMA, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WhiteListFst(GraphFst): + """ + Finite state transducer for verbalizing whitelist + e.g. tokens { name: "mrs." } -> mrs. + """ + + def __init__(self): + super().__init__(name="whitelist", kind="verbalize") + graph = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_CHAR - " ", 1) + + pynutil.delete("\"") + ) + graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", NEMO_SIGMA) + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/word.py b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/word.py new file mode 100644 index 0000000000..24565800ee --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/inverse_text_normalization/verbalizers/word.py @@ -0,0 +1,39 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, NEMO_SIGMA, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WordFst(GraphFst): + """ + Finite state transducer for verbalizing plain tokens + e.g. tokens { name: "sleep" } -> sleep + """ + + def __init__(self): + super().__init__(name="word", kind="verbalize") + chars = pynini.closure(NEMO_CHAR - " ", 1) + char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") + graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", NEMO_SIGMA) + + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/setup.sh b/models/tts/NeMo/nemo_text_processing/setup.sh new file mode 100644 index 0000000000..e410b40fd7 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/setup.sh @@ -0,0 +1,2 @@ +#!/bin/bash +conda install -c conda-forge -y pynini=2.1.3 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/README.md b/models/tts/NeMo/nemo_text_processing/text_normalization/README.md new file mode 100644 index 0000000000..f4405dd6c1 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/README.md @@ -0,0 +1,11 @@ +Text Normalization system for english, e.g. `123 kg` -> `one hundred twenty three kilograms` +Offers prediction and evaluation on text normalization data, e.g. [Google text normalization dataset](https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish). + + +Install dependencies: +bash ../setup.sh + +Example prediction run: +python run_predict.py --input=`INPUT_FILE` --output=`OUTPUT_FILE` [--verbose] +Example evaluation run: +python run_evaluate.py --input=./en_with_types/output-00001-of-00100 [--cat CATEGORY] diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/__init__.py new file mode 100644 index 0000000000..97aabf71f3 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/__init__.py @@ -0,0 +1,33 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst +from nemo_text_processing.text_normalization.taggers.tokenize_and_classify import ClassifyFst +from nemo_text_processing.text_normalization.verbalizers.verbalize import VerbalizeFst +from nemo_text_processing.text_normalization.verbalizers.verbalize_final import VerbalizeFinalFst + +from nemo.utils import logging + +try: + import pynini + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + logging.warning( + "`pynini` is not installed ! \n" + "Please run the `nemo_text_processing/setup.sh` script" + "prior to usage of this toolkit." + ) + + PYNINI_AVAILABLE = False diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/clean_eval_data.py b/models/tts/NeMo/nemo_text_processing/text_normalization/clean_eval_data.py new file mode 100644 index 0000000000..8c33c4fa62 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/clean_eval_data.py @@ -0,0 +1,342 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser +from typing import List + +import regex as re +from nemo_text_processing.text_normalization.data_loader_utils import ( + EOS_TYPE, + Instance, + load_files, + training_data_to_sentences, +) + + +""" +This file is for evaluation purposes. +filter_loaded_data() cleans data (list of instances) for text normalization. Filters and cleaners can be specified for each semiotic class individually. +For example, normalized text should only include characters and whitespace characters but no punctuation. + Cardinal unnormalized instances should contain at least one integer and all other characters are removed. +""" + + +class Filter: + """ + Filter class + + Args: + class_type: semiotic class used in dataset + process_func: function to transform text + filter_func: function to filter text + + """ + + def __init__(self, class_type: str, process_func: object, filter_func: object): + self.class_type = class_type + self.process_func = process_func + self.filter_func = filter_func + + def filter(self, instance: Instance) -> bool: + """ + filter function + + Args: + filters given instance with filter function + + Returns: True if given instance fulfills criteria or does not belong to class type + """ + if instance.token_type != self.class_type: + return True + return self.filter_func(instance) + + def process(self, instance: Instance) -> Instance: + """ + process function + + Args: + processes given instance with process function + + Returns: processed instance if instance belongs to expected class type or original instance + """ + if instance.token_type != self.class_type: + return instance + return self.process_func(instance) + + +def filter_cardinal_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_cardinal_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r"[^0-9]", "", un_normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_ordinal_1(instance: Instance) -> bool: + ok = re.search(r"(st|nd|rd|th)\s*$", instance.un_normalized) + return ok + + +def process_ordinal_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r"[,\s]", "", un_normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_decimal_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_decimal_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + un_normalized = re.sub(r",", "", un_normalized) + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_measure_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_measure_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r",", "", un_normalized) + un_normalized = re.sub(r"m2", "m²", un_normalized) + un_normalized = re.sub(r"(\d)([^\d.\s])", r"\1 \2", un_normalized) + normalized = re.sub(r"[^a-z\s]", "", normalized) + normalized = re.sub(r"per ([a-z\s]*)s$", r"per \1", normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_money_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_money_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + un_normalized = re.sub(r",", "", un_normalized) + un_normalized = re.sub(r"a\$", r"$", un_normalized) + un_normalized = re.sub(r"us\$", r"$", un_normalized) + un_normalized = re.sub(r"(\d)m\s*$", r"\1 million", un_normalized) + un_normalized = re.sub(r"(\d)bn?\s*$", r"\1 billion", un_normalized) + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_time_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_time_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + un_normalized = re.sub(r": ", ":", un_normalized) + un_normalized = re.sub(r"(\d)\s?a\s?m\s?", r"\1 a.m.", un_normalized) + un_normalized = re.sub(r"(\d)\s?p\s?m\s?", r"\1 p.m.", un_normalized) + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_plain_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_plain_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_punct_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_punct_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_date_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_date_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + un_normalized = re.sub(r",", "", un_normalized) + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_letters_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_letters_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_verbatim_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_verbatim_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_digit_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_digit_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_telephone_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_telephone_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_electronic_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_electronic_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_fraction_1(instance: Instance) -> bool: + ok = re.search(r"[0-9]", instance.un_normalized) + return ok + + +def process_fraction_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +def filter_address_1(instance: Instance) -> bool: + ok = True + return ok + + +def process_address_1(instance: Instance) -> Instance: + un_normalized = instance.un_normalized + normalized = instance.normalized + normalized = re.sub(r"[^a-z ]", "", normalized) + return Instance(token_type=instance.token_type, un_normalized=un_normalized, normalized=normalized) + + +filters = [] +filters.append(Filter(class_type="CARDINAL", process_func=process_cardinal_1, filter_func=filter_cardinal_1)) +filters.append(Filter(class_type="ORDINAL", process_func=process_ordinal_1, filter_func=filter_ordinal_1)) +filters.append(Filter(class_type="DECIMAL", process_func=process_decimal_1, filter_func=filter_decimal_1)) +filters.append(Filter(class_type="MEASURE", process_func=process_measure_1, filter_func=filter_measure_1)) +filters.append(Filter(class_type="MONEY", process_func=process_money_1, filter_func=filter_money_1)) +filters.append(Filter(class_type="TIME", process_func=process_time_1, filter_func=filter_time_1)) + +filters.append(Filter(class_type="DATE", process_func=process_date_1, filter_func=filter_date_1)) +filters.append(Filter(class_type="PLAIN", process_func=process_plain_1, filter_func=filter_plain_1)) +filters.append(Filter(class_type="PUNCT", process_func=process_punct_1, filter_func=filter_punct_1)) +filters.append(Filter(class_type="LETTERS", process_func=process_letters_1, filter_func=filter_letters_1)) +filters.append(Filter(class_type="VERBATIM", process_func=process_verbatim_1, filter_func=filter_verbatim_1)) +filters.append(Filter(class_type="DIGIT", process_func=process_digit_1, filter_func=filter_digit_1)) +filters.append(Filter(class_type="TELEPHONE", process_func=process_telephone_1, filter_func=filter_telephone_1)) +filters.append(Filter(class_type="ELECTRONIC", process_func=process_electronic_1, filter_func=filter_electronic_1)) +filters.append(Filter(class_type="FRACTION", process_func=process_fraction_1, filter_func=filter_fraction_1)) +filters.append(Filter(class_type="ADDRESS", process_func=process_address_1, filter_func=filter_address_1)) +filters.append(Filter(class_type=EOS_TYPE, process_func=lambda x: x, filter_func=lambda x: True)) + + +def filter_loaded_data(data: List[Instance], verbose: bool = False) -> List[Instance]: + """ + Filters list of instances + + Args: + data: list of instances + + Returns: filtered and transformed list of instances + """ + updates_instances = [] + for instance in data: + updated_instance = False + for fil in filters: + if fil.class_type == instance.token_type and fil.filter(instance): + instance = fil.process(instance) + updated_instance = True + if updated_instance: + if verbose: + print(instance) + updates_instances.append(instance) + return updates_instances + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--input", help="input file path", type=str, default='./en_with_types/output-00001-of-00100') + parser.add_argument("--verbose", help="print filtered instances", action='store_true') + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + file_path = args.input + + print("Loading training data: " + file_path) + instance_list = load_files([file_path]) # List of instances + filtered_instance_list = filter_loaded_data(instance_list, args.verbose) + training_data_to_sentences(filtered_instance_list) diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency.tsv new file mode 100644 index 0000000000..1c300bf774 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency.tsv @@ -0,0 +1,35 @@ +$ dollar +$ us dollar +$ united states dollar +£ pound +€ euro +₩ won +nzd new zealand dollar +rs rupee +chf swiss franc +dkk danish kroner +fim finnish markka +aed arab emirates dirham +¥ yen +czk czech koruna +mro mauritanian ouguiya +pkr pakistani rupee +crc costa rican colon +hk$ hong kong dollar +npr nepalese rupee +awg aruban florin +nok norwegian kroner +tzs tanzanian shilling +sek swedish kronor +cyp cypriot pound +r real +sar saudi riyal +cve cape verde escudo +rsd serbian dinar +dm german mark +shp saint helena pounds +php philippine peso +cad canadian dollar +ssp south sudanese pound +scr seychelles rupee +mvr maldivian rufiyaa \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency_minor.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency_minor.tsv new file mode 100644 index 0000000000..61597b9203 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency_minor.tsv @@ -0,0 +1,2 @@ +cents +pence \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency_minor_one.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency_minor_one.tsv new file mode 100644 index 0000000000..d6ab85fdcf --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/currency/currency_minor_one.tsv @@ -0,0 +1,2 @@ +cent +penny \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/domain.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/domain.tsv new file mode 100644 index 0000000000..cb4bc3f8f4 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/domain.tsv @@ -0,0 +1 @@ +com com \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/server_name.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/server_name.tsv new file mode 100644 index 0000000000..c447cf4682 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/server_name.tsv @@ -0,0 +1,2 @@ +gmail g mail +nvidia n vidia \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/symbols.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/symbols.tsv new file mode 100644 index 0000000000..c14f465c39 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/electronic/symbols.tsv @@ -0,0 +1,20 @@ +. dot +- dash +_ underscore +! exclamation mark +# number sign +$ dollar sign +% percent sign +& ampersand +' quote +* asterisk ++ plus +/ slash += equal sign +? question mark +^ circumflex +` right single quote +{ left brace +| vertical bar +} right brace +~ tilde \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/letter_pronunciation.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/letter_pronunciation.tsv new file mode 100644 index 0000000000..3ec49acb9f --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/letter_pronunciation.tsv @@ -0,0 +1,2 @@ +k kay +j jay \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/magnitudes.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/magnitudes.tsv new file mode 100644 index 0000000000..3bfa4849e8 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/magnitudes.tsv @@ -0,0 +1,4 @@ +k thousand +m million +b billion +t trillion \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/measurements.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/measurements.tsv new file mode 100644 index 0000000000..b86459a1f6 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/measurements.tsv @@ -0,0 +1,115 @@ +f fahrenheit +c celsius +km kilometer +m meter +cm centimeter +mm millimeter +ha hectare +mi mile +m² square meter +m2 square meter +km² square kilometer +km2 square kilometer +ft foot +% percent +hz hertz +kw kilowatt +hp horsepower +mg milligram +kg kilogram +ghz gigahertz +khz kilohertz +mhz megahertz +lb pound +lbs pound +lbs. pound +v volt +h hour +mc mega coulomb +s second +nm nanometer +rpm revolution per minute +min minute +mA milli ampere +kwh kilo watt hour +m³ cubic meter +m3 cubic meter +mph mile per hour +tw tera watt +mv milli volt +mw megawatt +μm micrometer +" inch +tb terabyte +cc c c +g gram +da dalton +atm atmosphere +ω ohm +db decibel +ps peta second +oz ounce +hl hecto liter +μg microgram +pg petagram +gb gigabyte +kb kilobit +ev electron volt +mb megabyte +kb kilobyte +kbps kilobit per second +mbps megabit per second +kl kilo liter +tj tera joule +kv kilo volt +mv mega volt +kn kilonewton +mm megameter +au astronomical unit +yd yard +rad radian +lm lumen +hs hecto second +mol mole +gpa giga pascal +ml milliliter +gw gigawatt +ma mega ampere +kt knot +kgf kilogram force +ng nano gram +ns nanosecond +ms mega siemens +bar bar +gl giga liter +μs microsecond +da deci ampere +pa pascal +ds deci second +ms milli second +dm deci meter +dm³ cubic deci meter +dm3 cubic deci meter +amu atomic mass unit +mb megabit +mf mega farad +bq becquerel +pb petabit +mm² square millimeter +mm2 square millimeter +cm² square centimeter +cm2 square centimeter +sq mi square mile +sq ft square foot +kpa kilopascal +cd candela +tl tera liter +ms mega second +mpa megapascal +pb peta byte +gwh giga watt hour +kcal kilo calory +gy gray +sv sievert +cwt hundredweight +cc c c diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/abbr.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/abbr.tsv new file mode 100644 index 0000000000..5609e211d6 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/abbr.tsv @@ -0,0 +1,12 @@ +jan january +feb february +mar march +apr april +jun june +jul july +aug august +sep september +sept september +oct october +nov november +dec december \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/names.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/names.tsv new file mode 100644 index 0000000000..1c05ac02a3 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/names.tsv @@ -0,0 +1,12 @@ +january +february +march +april +may +june +july +august +september +october +november +december diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/numbers.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/numbers.tsv new file mode 100644 index 0000000000..f0cacd5d1e --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/months/numbers.tsv @@ -0,0 +1,24 @@ +1 january +2 february +3 march +4 april +5 may +6 june +7 july +8 august +9 september +10 october +11 november +12 december +01 january +02 february +03 march +04 april +05 may +06 june +07 july +08 august +09 september +10 october +11 november +12 december \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/cardinal_number_name.far b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/cardinal_number_name.far new file mode 100755 index 0000000000..2ec682541c Binary files /dev/null and b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/cardinal_number_name.far differ diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/digit.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/digit.tsv new file mode 100644 index 0000000000..fa329799db --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/digit.tsv @@ -0,0 +1,9 @@ +one 1 +two 2 +three 3 +four 4 +five 5 +six 6 +seven 7 +eight 8 +nine 9 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/hundred.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/hundred.tsv new file mode 100644 index 0000000000..10e2b9645a --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/hundred.tsv @@ -0,0 +1 @@ +hundred \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/teen.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/teen.tsv new file mode 100644 index 0000000000..8e60fa1d3c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/teen.tsv @@ -0,0 +1,10 @@ +ten 10 +eleven 11 +twelve 12 +thirteen 13 +fourteen 14 +fifteen 15 +sixteen 16 +seventeen 17 +eighteen 18 +nineteen 19 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/thousands.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/thousands.tsv new file mode 100644 index 0000000000..bf30542bb6 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/thousands.tsv @@ -0,0 +1,22 @@ +thousand +million +billion +trillion +quadrillion +quintillion +sextillion +septillion +octillion +nonillion +decillion +undecillion +duodecillion +tredecillion +quattuordecillion +quindecillion +sexdecillion +septendecillion +octodecillion +novemdecillion +vigintillion +centillion \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/ties.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/ties.tsv new file mode 100644 index 0000000000..65f9839fda --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/ties.tsv @@ -0,0 +1,8 @@ +twenty 2 +thirty 3 +forty 4 +fifty 5 +sixty 6 +seventy 7 +eighty 8 +ninety 9 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/zero.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/zero.tsv new file mode 100644 index 0000000000..c479272d40 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/numbers/zero.tsv @@ -0,0 +1 @@ +zero 0 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/digit.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/digit.tsv new file mode 100644 index 0000000000..4b7fc24f16 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/digit.tsv @@ -0,0 +1,9 @@ +first one +second two +third three +fourth four +fifth five +sixth sixth +seventh seven +eighth eight +ninth nine \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/teen.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/teen.tsv new file mode 100644 index 0000000000..496fefc086 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/ordinals/teen.tsv @@ -0,0 +1 @@ +twelfth twelve \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/digit_teen.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/digit_teen.tsv new file mode 100644 index 0000000000..cd0991331f --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/digit_teen.tsv @@ -0,0 +1,49 @@ +i 1 +ii 2 +iii 3 +iv 4 +v 5 +vi 6 +vii 7 +viii 8 +ix 9 +x 10 +xi 11 +xii 12 +xiii 13 +xiv 14 +xv 15 +xvi 16 +xvii 17 +xviii 18 +xix 19 +xx 20 +xxi 21 +xxii 22 +xxiii 23 +xxiv 24 +xxv 25 +xxvi 26 +xxvii 27 +xxviii 28 +xxix 29 +xxx 30 +xxxi 31 +xxxii 32 +xxxiii 33 +xxxiv 34 +xxxv 35 +xxxvi 36 +xxxvii 37 +xxxviii 38 +xxxix 39 +xl 40 +xli 41 +xlii 42 +xliii 43 +xliv 44 +xlv 45 +xlvi 46 +xlvii 47 +xlviii 48 +xlix 49 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/hundreds.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/hundreds.tsv new file mode 100644 index 0000000000..0aafad3049 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/hundreds.tsv @@ -0,0 +1,9 @@ +c 100 +cc 200 +ccc 300 +cd 400 +d 500 +dc 600 +dcc 700 +dccc 800 +cm 900 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/ties.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/ties.tsv new file mode 100644 index 0000000000..5516676f4b --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/roman/ties.tsv @@ -0,0 +1,5 @@ +l 50 +lx 60 +lxx 70 +lxxx 80 +xc 90 \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/suppletive.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/suppletive.tsv new file mode 100644 index 0000000000..dea620f37a --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/suppletive.tsv @@ -0,0 +1,37 @@ +deer +fish +sheep +foot feet +goose geese +man men +mouse mice +tooth teeth +woman women +won +child children +ox oxen +wife wives +wolf wolves +analysis analyses +criterion criteria +focus foci +percent +hertz +kroner krone +inch inches +calory calories +yen +megahertz +gigahertz +kilohertz +hertz +c c +horsepower +hundredweight +kilogram force kilograms force +mega siemens +revolution per minute revolutions per minute +mile per hour miles per hour +megabit per second megabits per second +square foot square feet +kilobit per second kilobits per second \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/time_suffix.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/time_suffix.tsv new file mode 100644 index 0000000000..98664c5a32 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/time_suffix.tsv @@ -0,0 +1,6 @@ +p.m. p m +p.m. p m +pm p m +a.m. a m +a.m a m +am a m \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/time_zone.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/time_zone.tsv new file mode 100644 index 0000000000..1d39472c84 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/time_zone.tsv @@ -0,0 +1,14 @@ +cst c s t +c.s.t c s t +cet c e t +c.e.t c e t +pst p s t +p.s.t p s t +est e s t +e.s.t e s t +pt p t +p.t p t +et e t +e.t e t +gmt g m t +g.m.t g m t diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/whitelist.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/whitelist.tsv new file mode 100644 index 0000000000..8539931904 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/whitelist.tsv @@ -0,0 +1,15924 @@ +Ph.D. p h d +Hon. honorable +& and +Mt. Mount +Maj. Major +Rev. Reverend +# hash +Gov. governor +7-eleven seven eleven +A. A. a a +A.A. a a +A.A.A. a a a +A&A a and a +AAAI a a a i +AAAM a a a m +AAAs a a a's +AAAS a a a s +AAAW a a a w +AABA a a b a +AaB a a b +A. A. B. a a b +AAB- a a b +AAB a a b +AABC a a b c +Aabo a a b o +AABP a a b p +AABW a a b w +aac a a c +AAC a a c +AACAR a a c a r +AACC a a c c +AACCUP a a c c u p +AACMI a a c m i +AACNo a a c n o +AACR a a c r +AACS a a c s +AACSB a a c s b +AACTA a a c t a +AACUPR a a c u p r +A.A.D. a a d +AADT a a d t +AADTs a a d t's +Aadu a a d u +AAEA a a e a +AAE a a e +AAERT a a e r t +AAF a a f +AAFCA a a f c a +AAFC a a f c +AAFC's a a f c's +AAFld a a f l d +AAFPOA a a f p o a +AAGPBL a a g p b l +AAHHE a a h h e +AAI a a i +AAIB a a i b +AAIP a a i p +AAK a a k +Aap a a p +AAP a a p +AAPB a a p b +AAPC a a p c +AAPEP a a p e p +AAPG a a p g +AAPM a a p m +AAPT a a p t +A.A.R.M. a a r m +A.B.A. a b a +A. B. a b +A.B. a b +A.B.A. J. a b a j +A.B.C. a b c +A. B. G. a b g +ABG a b g +ABH a b h +ABHD a b h d +Abhi a b h i +ABK a b k +abl a b l +Abl a b l +ABL a b l +ABLV a b l v +ABM a b m +ABMC a b m c +ABN a b n +ABN's a b n's +Abp a b p +ABP- a b p +ABP a b p +ABPD a b p d +ABPI a b p i +ABPN a b p n +ABPP a b p p +ABPRS a b p r s +ABPW a b p w +Abr a b r +ABR a b r +abv a b v +ABV- a b v +ABV a b v +ABW a b w +ABX a b x +Abz a b z +A.C.A. a c a +ACA a c a +A&C a and c +A. C. a c +A.C. a c +acac a c a c +Acad a c a d +ACAD a c a d +A.C.A.P. a c a p +ACAP a c a p +ACAS a c a s +ACAZ a c a z +ACBA a c b a +ACB a c b +ACBL a c b l +ACBS a c b s +ACBSP a c b s p +ACCA a c c a +ACC- a c c +ACC a c c +ACCC a c c c +ACCJ a c c j +ACC&S a c c and s +ACC's a c c's +ACCS a c c s +ACD a c d +ACDC a c d c +A.C. D.F.C. A.F.C. a c d f c a f c +A.C.E. a c e +ACF a c f +ACFE a c f e +A.C.G. a c g +ACG a c g +ACGIH's a c g i h's +A.C.H. a c h +ACH a c h +ACHP a c h p +ACHR a c h r +ACHRE a c h r e +Achs a c h's +ACHS a c h s +ACIA a c i a +ACI a c i +ACICS a c i c s +ACIGA a c i g a +ACIP a c i p +Acis a c i's +ACIS a c i s +ACLA a c l a +A.C.L. a c l +ACL a c l +A.C.L.N. a c l n +acls a c l s +ACL's a c l's +ACLS a c l s +ACLU a c l u +ACM- a c m +ACM a c m +ACMI a c m i +ACMs a c m's +ACN a c n +ACO's a c o's +AcpA a c p a +ACP- a c p +ACP a c p +AcpB a c p b +ACPI a c p i +ACPO a c p o +ACPP a c p p +ACPSEM a c p s e m +acq a c q +A&CR a and c r +ACR a c r +ACRPS a c r p s +ACSA a c s a +ACSAC a c s a c +ACSBR a c s b r +ACSCN a c s c n +ACSEL a c s e l +ACSIA a c s i a +ACSI a c s i +AcSOC a c s o c +ACSR a c s r +ACSS a c s s +ACSS's a c s s's +A.C.T. a c t +A.C.T.A.F.L. a c t a f l +acu a c u +ACU a c u +ACU's a c u's +ACUS a c u s +ACUV a c u v +ACV a c v +ACWA a c w a +AC&W a c and w +ACW a c w +ACWM a c w m +ACWS a c w s +Acy a c y +ADAA a d a a +A&D a and d +Adab a d a b +ADAC- a d a c +ADAC a d a c +A. D. a d +A.D. a d +A.D.A.M. a d a m +AdaSL a d a s l +ADATA a d a t a +ADBAC a d b a c +ADB- a d b +ADB a d b +ADBGR a d b g r +ADBICA a d b i c a +ADCA a d c a +AdC a d c +A.D.C. a d c +ADC a d c +ADCAF a d c a f +ADCC a d c c +ADC's a d c's +ADCs a d c's +ADCY a d c y +ADDA a d d a +AD&D a d and d +ADF a d f +ADF's a d f's +ADFS a d f s +ADG a d g +ADGB a d g b +ADH a d h +ADHD a d h d +ADHM a d h m +ADHS a d h s +ADIAC a d i a c +AdK a d k +ADK a d k +Adl a d l +ADL a d l +ADLL a d l l +ADLs a d l's +A.D.M.A.C. a d m a c +adm a d m +Adm a d m +ADM- a d m +ADM a d m +ADMK a d m k +admn a d m n +ADMN a d m n +ADNAC a d n a c +ADNEC a d n e c +ADNs a d n's +AdP a d p +A.D.P. a d p +ADP- a d p +ADP a d p +ADPCM a d p c m +ADPF a d p f +ADQ a d q +ADQ's a d q's +ADR a d r +ADRC a d r c +ADRP a d r p +ADRs a d r's +ADSB a d s b +ADSL a d s l +ADSP- a d s p +ADSR a d s r +ADSRs a d s r's +ADSs a d s's +ADT a d t +ADTG a d t g +Adwa a d w a +ADWC a d w c +ADX a d x +AEA a e a +A&E a and e +A. E. a e +A.E. a e +AEBN a e b n +A.E.C. a e c +AEC a e c +AECL a e c l +AECR a e c r +AEC's a e c's +Aed a e d +AED a e d +AEE a e e +AEF a e f +AEG a e g +A.E.G.I.S. a e g i s +aegte a e g t e +AEHL a e h l +AEIA a e i a +AEI a e i +AEIO- a e i o +AEIOU a e i o u +Aeka a e k a +AEK a e k +ael a e l +AEL a e l +Aema a e m a +AEM a e m +aen a e n +AEO a e o +AEPA a e p a +AEP a e p +aere a e r e +AEre a e r e +AErn a e r n +AERN a e r n +Aert a e r t +AESA a e s a +Aes a e's +AES- a e s +AES a e s +aet a e t +AET a e t +AETC a e t c +AEU a e u +AEV a e v +AEW a e w +AEX a e x +Afa a f a +AFA a f a +af a f +Af a f +A. F. a f +A.F. a f +AF- a f +AF a f +Afaf a f a f +AFAP a f a p +AFB a f b +AFCA- a f c a +AFCA a f c a +afc a f c +A.F.C. a f c +AFC a f c +AFCEC a f c e c +AFCO a f c o +AFC's a f c's +AFCS a f c s +AFCSThe a f c s t h e +Afd a f d +AfD a f d +AfDB a f d b +AFDD a f d d +Afe a f e +AFE a f e +Aff a f f +AFF a f f +AFFDL- a f f d l +AFG a f g +AFGM a f g m +AFH a f h +AFIA a f i a +Afi a f i +AFI a f i +AFIP a f i p +Afiq a f i q +AFI's a f i's +AFJ a f j +A.F.L. a f l +AFL- a f l +AFL a f l +AFLEG a f l e g +AFLPA a f l p a +AFLP a f l p +AFL's a f l's +AFLUA a f l u a +A. F.M.A. a f m a +AFM- a f m +AFM a f m +AFMC a f m c +AFMs a f m's +AFN a f n +AFNET a f n e t +AFNOR a f n o r +AfNS a f n s +Afo a f o +AFO a f o +AFOP a f o p +AFP a f p +AFRA a f r a +Afr a f r +AFR a f r +AFRC a f r c +AFRL a f r l +AFROTC a f r o t c +AFRTS a f r t s +AFS a f s +AFSC a f s c +AFSPC a f s p c +A. F. T. a f t +AFTRA a f t r a +AFTRA's a f t r a's +AFTRS a f t r s +AFV a f v +AFW a f w +AFWS a f w s +AFX a f x +A. G. a g +A.G. a g +AGB a g b +AGC a g c +Agco a g c o +AGCO a g c o +Agda a g d a +AGD a g d +A.G.E. a g e +AGF- a g f +AGF a g f +AG&G a g and g +agg a g g +Agi a g i +AGI a g i +AGID a g i d +AGIR a g i r +Agis a g i's +A.G.J. a g j +A. G. K. a g k +AGL a g l +agli a g l i +AGM- a g m +AGM a g m +AGMR- a g m r +Agn a g n +AgN a g n +AGN a g n +AGNs a g n's +A.G.P. a g p +AGP a g p +AGPL a g p l +AGPS a g p s +AGPW a g p w +AGRs a g r's +AGS a g s +AGSM a g s m +AGSS- a g s s +Agsu a g s u +agt a g t +AGT a g t +Agte a g t e +Agu a g u +AGU a g u +AGV a g v +A.H.A. a h a +A. H. a h +A.H. a h +A. H. C. a h c +AHC a h c +AHD a h d +Ahdhra a h d h r a +AHDR a h d r +ahe a h e +Ahe a h e +AHEC a h e c +AHF a h f +Ahirs a h i r's +AHJD- a h j d +Ahk a h k +AHL's a h l's +AHP a h p +AHRC a h r c +AHRS a h r s +AHSAA a h s a a +Ahsa'i a h s a i +Ahta a h t a +aht a h t +AHT a h t +AHTD a h t d +Ahu a h u +Ahva a h v a +AHV a h v +AIAA a i a a +A.I.A. a i a +A&I a and i +A. I. a i +A.I. a i +AI- a i +AI a i +AIAP a i a p +Aias a i a's +AIAS a i a s +AIATSIS a i a t s i s +AIAW a i a w +AIAWU a i a w u +A.I.B.A. a i b a +AIB a i b +AIBO's a i b o's +AIBS a i b s +AIC a i c +AICAR a i c a r +AICC a i c c +AICCCR a i c c c r +aici a i c i +AICL- a i c l +Aicme a i c m e +AICN a i c n +AICPA a i c p a +AICP a i c p +AICPS a i c p s +AICs a i c's +AICTE a i c t e +AICUF a i c u f +AIDA a i d a +AIDN a i d n +AIEC a i e c +AIEE- a i e e +AIEE a i e e +Aiel a i e l +AIEP a i e p +AIFA a i f a +AIF- a i f +AIF a i f +AIFB a i f b +Aife a i f e +AIFF a i f f +AIFM a i f m +AIFMD a i f m d +AiG a i g +AIG a i g +Aigis a i g i's +AIHA a i h a +AIH a i h +AIHL a i h l +AIIMS a i i m s +Aija a i j a +Aik a i k +AIK- a i k +AIK a i k +AIKR a i k r +A.I.M. a i m +AIP a i p +A.I.R. a i r +AIs a i's +AISC a i s c +AISD a i s d +AISRI a i s r i +Aist a i s t +AIST a i s t +AIT's a i t's +Aitu a i t u +Aiud a i u d +aius a i u s +AIVC a i v c +AIW a i w +AIX a i x +A. J. a j +A.J. a j +Ajba a j b a +A.J.B. a j b +AJC a j c +A.J.E. a j e +A.J.G.C. a j g c +A. J. J. a j j +A.J.K. a j k +AJK a j k +AJKF a j k f +AJL a j l +A. J. M. a j m +Ajmi a j m i +AJN a j n +AJO a j o +Ajok a j o k +A. J. P. a j p +A.J.P. a j p +AJPW a j p w +A.J.R. a j r +AJR- a j r +AJR a j r +Ajsa a j s a +A.J.S. a j s +AJS- a j s +AJS a j s +A. J. T. a j t +A.J.T. a j t +AJT a j t +Aju a j u +AJUFE a j u f e +AJUSCO a j u s c o +AJV a j v +AJW a j w +AJWRC a j w r c +Akaa a k a a +a.k.a. a k a +aka. a k a +aka a k a +Aka. a k a +Aka a k a +AKA- a k a +AKA a k a +A&K a and k +ak- a k +A. K. a k +A.K. a k +AK- a k +AK a k +AKAP a k a p +AKAPs a k a p's +akas a k a's +AKAs a k a's +AKB a k b +AKC a k c +AKCR a k c r +AKD a k d +AKG a k g +AKG's a k g's +akh a k h +AKH a k h +AKHS a k h s +AKK- a k k +AKM a k m +AKN a k n +Akpa a k p a +AKP a k p +AKPD a k p d +AKQA a k q a +AKR- a k r +AKR a k r +A.K.S. a k s +AKS a k s +Aku a k u +AKU a k u +AKW a k w +A&L a and l +A. L. a l +A.L. a l +ALDF a l d f +ALDH a l d h +ALDS a l d s +ALDT a l d t +ALFASID a l f a s i d +ALFIDI a l f i d i +ALFTP a l f t p +Alh a l h +ALH a l h +A. L. M. a l m +A.L.P. a l p +ALP- a l p +ALPG a l p g +ALQ- a l q +ALR a l r +ALSC a l s c +ALSF a l s f +ALSP a l s p +ALUs a l u's +A.L.V. a l v +ALVF a l v f +ALW a l w +ALWS a l w s +AMAA a m a a +A'ma a m a +A&M a and m +A. M. a m +A.M. a m +AMAP a m a p +AMARC a m a r c +AMAs a m a's +amb a m b +Amb a m b +AMB a m b +AmBX a m b x +AMc- a m c +A.M.C. a m c +AMC- a m c +AMC a m c +AMC's a m c's +AMCs a m c's +Amda a m d a +Amd a m d +AMD a m d +AMDH a m d h +AMDISA a m d i s a +AMD's a m d's +A.M.E. a m e +AMF a m f +AM&FM a m and f m +A. M. G. a m g +AMG- a m g +AMG a m g +AMGTV a m g t v +Amha a m h a +AMH a m h +A.M.H.S. a m h s +A. M. J. a m j +AMJ a m j +A. M. K. a m k +A.M.K. a m k +AMK a m k +aml a m l +AML- a m l +AML a m l +Amli a m l i +AMLS a m l s +Amlwch a m l w c h +Amm a m m +A. M. M. a m m +AMM a m m +AMMAYI a m m a y i +AMMB a m m b +AMMK a m m k +AMN a m n +AMNH a m n h +AMNRL a m n r l +AM&O a m and o +AMPK- a m p k +AMPK a m p k +AMPL a m p l +AMP's a m p's +Ampt a m p t +AMPTP a m p t p +amr a m r +A. M. R. a m r +AMREF a m r e f +AMSA a m s a +A&M's a and m's +Ams' a m's +AMS- a m s +amsl a m s l +AMSO's a m s o's +AMSRS a m s r s +AMS's a m s's +AMSS a m s s +AMSU- a m s u +AMSU a m s u +Amta a m t a +amt a m t +Amt a m t +A.M.T. a m t +AMT- a m t +AMT a m t +AMTB a m t b +AMTPAT a m t p a t +Amts a m t's +AMTV a m t v +Amu a m u +AMU a m u +AMU's a m u's +A.M.V.M. a m v m +AMW a m w +AMX- a m x +AMX a m x +AMYF a m y f +A. N. a n +A.N. a n +ANB a n b +ANBO a n b o +ANC a n c +ANCC a n c c +ANC's a n c's +andFHL a n d f h l +ANDOC a n d o c +ANDP a n d p +ANDPOP a n d p o p +andPPP a n d p p p +Anej a n e j +ANELFA a n e l f a +ANF a n f +ANFIS a n f i s +ANFP's a n f p's +ANGB a n g b +ANGPTL a n g p t l +ANGRAU a n g r a u +Angti a n g t i +Angu a n g u +Angul- a n g u l +Anhe a n h e +ANI a n i +ANL a n l +ANM a n m +ANP a n p +ANPP a n p p +ANPR a n p r +Anra a n r a +Anr a n r +ANR a n r +ANRC a n r c +ANREF's a n r e f's +Anrep a n r e p +ANRS a n r s +ANRW a n r w +Ans a n's +A.N.S. a n s +ANS a n s +ANSF a n s f +ANTB a n t b +ANTM a n t m +ANTV a n t v +ANUGA a n u g a +ANWB a n w b +ANX a n x +ANZ a n z +A&O a and o +AOA a o a +AOAC a o a c +A. O. a o +A.O. a o +AOB a o b +AOC a o c +AOCB a o c b +AOCCs a o c c's +AOCE a o c e +AOD a o d +Aodh a o d h +Aodla a o d l a +AODV a o d v +AOE a o e +AOGCC a o g c c +Aogo a o g o +AOH a o h +Aoke a o k e +Aoko a o k o +AOKP a o k p +Aola a o l a +Aol a o l +AOL- a o l +AOL a o l +AOL's a o l's +AOLTV a o l t v +AOM a o m +AOMC a o m c +AONB a o n b +Aone a o n e +Aoni a o n i +Aoos a o o's +AOPA a o p a +AOP a o p +aor a o r +AOR a o r +AORs a o r's +AORS a o r s +aos a o s +AoS a o s +AOS a o s +Aotus a o t u's +AOTW a o t w +aov a o v +aovf a o v f +AOWC a o w c +A&P a and p +A. P. a p +A.P. a p +APBA a p b a +APB a p b +APBL a p b l +APCA a p c a +Apc a p c +APC- a p c +APC a p c +APCCR a p c c r +APCh a p c h +APCR a p c r +APCRDA a p c r d a +APC's a p c's +APCs a p c's +APDA a p d a +APD- a p d +APD a p d +APDM a p d m +APEGBC a p e g b c +APFA a p f a +APF a p f +APFCT a p f c t +APFOL a p f o l +APFSDS a p f s d s +APFUTU a p f u t u +apg a p g +APG a p g +APGAW a p g a w +Aph a p h +A.P.H. a p h +APH a p h +APHEDA a p h e d a +APHL a p h l +Apiao a p i a o +Api a p i +APi a p i +API- a p i +API a p i +Apic a p i c +APICv a p i c v +API's a p i's +A.P.J. a p j +APJ a p j +Apl a p l +APL a p l +APLP a p l p +APLS a p l s +Aplu a p l u +APM a p m +APML a p m l +APMR a p m r +APMSO a p m s o +APN- a p n +APN a p n +A.P.N.C. a p n c +Apphttp a p p h t t p +Appl a p p l +APPO's a p p o's +appr a p p r +Appts a p p t's +appu a p p u +Appu a p p u +APRA's a p r a's +APRC a p r c +APR's a p r's +APRS a p r s +APRST a p r s t +A&P's a and p's +Aps a p's +APs a p's +APS- a p s +APS a p s +APSF a p s f +APSL a p s l +APSU- a p s u +APTA- a p t a +APTA a p t a +AP&T a p and t +APW a p w +APX a p x +APXS a p x s +Aqa a q a +AQA a q a +AQC a q c +AQI a q i +Aql a q l +AQP a q p +Aqr a q r +AQS- a q s +A&R a and r +ARA a r a +A. R. a r +A.R. a r +A.R.B. a r b +ARD a r d +ARDF a r d f +A.R.E. a r e +A.R.F. a r f +ARF a r f +A.R.G. a r g +ArgR a r g r +ARGs a r g's +ARGT a r g t +ArH a r h +ARJ a r j +ARL a r l +ARLFC a r l f c +A. R. M. a r m +ARMv a r m v +Arnd a r n d +ARNG a r n g +ARNT a r n t +A.R.P. a r p +ARPP- a r p p +ARQ- a r q +ARQ a r q +ARTL a r t l +A.R.U. a r u +arv a r v +Arv a r v +ARV a r v +ARVD a r v d +ARVN a r v n +ARW a r w +Arwi a r w i +ARWU a r w u +A.S.A. a s a +ASA a s a +As'ad a s a d +Asai a s a i +Asao a s a o +A.s a's +A. S. a s +A.S. a s +ASAS a s a s +Asasp a s a s p +ASAU a s a u +A.S.B. a s b +ASB a s b +ASBDA a s b d a +asbl a s b l +ASBM a s b m +ASBMH a s b m h +ASBO a s b o +A.S.C. a s c +ASC a s c +ASCB a s c b +ASCE a s c e +asci a s c i +Asci a s c i +ASCP a s c p +ASCW a s c w +A.S.D. a s d +ASD a s d +ASDIC a s d i c +Asdis a s d i's +ASE a s e +ASEC a s e c +ASEE a s e e +ASEF a s e f +ASFA a s f a +ASF a s f +Asfi a s f i +ASGA a s g a +ASG a s g +ASGC a s g c +ASGE a s g e +ASGS a s g s +A. S. H. a s h +A.S.I. a s i +ASL a s l +A.S.M. a s m +ASM a s m +ASMD a s m d +ASME a s m e +ASMIK a s m i k +ASML a s m l +ASMPH a s m p h +ASMSU a s m s u +Asn a s n +ASN- a s n +ASN a s n +Aso a s o +ASO a s o +ASP a s p +ASPCA a s p c a +Aspe a s p e +ASPTS a s p t s +ASQ- a s q +Asr a s r +ASR a s r +Assn a s s n +assoc a s s o c +Assoc a s s o c +ASSPs a s s p's +ASSR a s s r +Asst a s s t +ASSU a s s u +A. S. T. a s t +ASTAT a s t a t +ASTCL a s t c l +Aste a s t e +ASTE a s t e +ASTM a s t m +ASTR a s t r +Astt a s t t +A.S.U. a s u +ASU- a s u +ASU a s u +A.S.V. a s v +ASV a s v +Aswa a s w a +ASW- a s w +ASW a s w +ASX a s x +A&T a and t +atac a t a c +Atac a t a c +ATAC- a t a c +Atad a t a d +Ata's a t a's +Atas a t a's +ATA's a t a's +A. T. a t +A.T. a t +ATBF a t b f +ATBs a t b's +atc a t c +A.T.C. a t c +ATC a t c +atcc a t c c +ATCC- a t c c +ATCC a t c c +ATCDE a t c d e +ATCL a t c l +ATDC a t d c +ATF- a t f +ATF a t f +Atg a t g +ATG a t g +ATGM a t g m +ath a t h +Ath a t h +ATH a t h +ATHN a t h n +ATIA a t i a +Atia's a t i a's +Ati a t i +ATi a t i +ATI a t i +Atid a t i d +Atiiq a t i i q +Atil a t i l +ATINC a t i n c +ATIP a t i p +Atiq a t i q +atis a t i s +Ativ a t i v +ATJ a t j +Atka a t k a +ATK a t k +ATKN a t k n +ATK's a t k's +Atl a t l +ATL a t l +Atli a t l i +ATLY a t l y +A.T.M.A. a t m a +atm a t m +Atm a t m +ATM a t m +ATMs a t m's +ATN a t n +ATO- a t o +ATO a t o +Atos a t o's +ATO's a t o's +ATOs a t o's +ATOS a t o s +ATP- a t p +ATP a t p +atpB a t p b +A.T.Q. a t q +ATR- a t r +ATR a t r +atri a t r i +Atri a t r i +ATRP a t r p +A&T's a and t's +A.T.s a t's +A.T.S. a t s +ATS a t s +ATSC a t s c +ATSDR a t s d r +AT&SF a t and s f +ATSF a t s f +ATSIC a t s i c +AT&T- a t and t +AT&T a t and t +attd a t t d +ATTESA a t t e s a +ATTF a t t f +AT&T's a t and t's +ATTS a t t s +ATTWI a t t w i +ATTWX's a t t w x's +ATU's a t u's +atv a t v +ATV- a t v +ATV a t v +ATV's a t v's +ATVs a t v's +ATW a t w +AtxA a t x a +ATX a t x +ATXN a t x n +A. U. a u +A.U. a u +Aub a u b +Aubl a u b l +AUC a u c +AUVs a u v's +Auw a u w +avab a v a b +AVAC a v a c +A. V. a v +A.V. a v +AVCA a v c a +AVC- a v c +AVC a v c +AVCHD a v c h d +AVCs a v c's +AVCS a v c s +AVD a v d +avg a v g +AVG a v g +AVK a v k +AVL a v l +A. V. M. a v m +AVM a v m +Avn a v n +AVN a v n +Avo a v o +AvP a v p +AVP- a v p +AVP a v p +AVR a v r +AVRs a v r's +AVSA a v s a +Avs' a v's +Avs a v's +AVS a v s +AVSM a v s m +AVU a v u +avvo a v v o +AVX a v x +Awa a w a +A. W. A. a w a +AWA a w a +A. W. A. M. a w a m +AWAs a w a's +A. W. a w +A.W. a w +A. W. B. a w b +A.W.B. a w b +AWB a w b +AWB's a w b's +A.W.C. a w c +AWC a w c +AWD a w d +A. W. F. a w f +AWG a w g +AWGIE a w g i e +AWGN a w g n +A. W. H. a w h +AWHL a w h l +AWI a w i +AWM a w m +Awo a w o +AWP a w p +AWP's a w p's +AWR a w r +AWSA a w s a +AWU a w u +Awwa a w w a +AXAF a x a f +A. X. a x +AXS a x s +AXV a x v +A. Y. a y +A.Y. a y +AYF a y f +AYK a y k +AYP a y p +A. Z. a z +A.Z. a z +AZD- a z d +AZE a z e +Azg a z g +Azi a z i +AZI a z i +AZL a z l +azm a z m +AZN a z n +AZS a z s +AZSTA's a z s t a's +AZT a z t +AzTV a z t v +Baad b a a d +BAAG b a a g +Ba'al b a a l +Baal b a a l +BAAL b a a l +Baam b a a m +Baap b a a p +ba'as b a a's +Baat b a a t +Ba'ath b a a t h +B. A. b a +B.A. b a +B&A b and a +Bac b a c +B.A.C. b a c +BAC b a c +bae b a e +Bae b a e +BAe b a e +BAE b a e +BAFA b a f a +BAF b a f +BAMF b a m f +BAPED- b a p e d +BAPH- b a p h +B.A.R. b a r +BARV b a r v +BASCA b a s c a +BASEC b a s e c +BASF b a s f +B.A.S.P. b a s p +BATVG b a t v g +Bauw b a u w +B'Av b a v +BAV b a v +BBA b b a +BBAG b b a g +B&B b and b +bb b b +B. B. b b +B.B. b b +BB- b b +BB b b +BBB b b b +bbc b b c +B.B.C. b b c +BBC- b b c +BBC b b c +BBCBBC b b c b b c +BBCH- b b c h +BB&CI b b and c i +BBC&PJR b b c and p j r +BBC's b b c's +B.B.D. b b d +BBDO b b d o +BBE b b e +BBF b b f +BBFC b b f c +BBFF b b f f +BBG b b g +BBH b b h +BBI b b i +BBK b b k +BBKL b b k l +BBLB b b l b +BBL- b b l +BBL b b l +BBM b b m +BBMP b b m p +BBNG b b n g +BBN's b b n's +BBO b b o +BBP b b p +BBQ'er b b q e r +B&Bs b and b's +BB&S b b and s +BBS b b s +BBSes b b s e's +BBS's b b s's +BB&T b b and t +BBT b b t +BBTV b b t v +BBU b b u +BBVA b b v a +BbvCI b b v c i +BBV's b b v's +BBWAA b b w a a +BBWA b b w a +BBWR's b b w r's +BBYA b b y a +BBYO b b y o +BCA- b c a +BCA b c a +BCAD b c a d +BCAM b c a m +BCATP b c a t p +BCBA b c b a +BCB b c b +bc b c +B. C. b c +B.C. b c +BC- b c +BC, b c +BC b c +BCCA b c c a +BCC b c c +BCCCA b c c c a +Bcci b c c i +BCCI b c c i +BCCP b c c p +BCC's b c c's +BCDA b c d a +BCD b c d +BCEA b c e a +BCE- b c e +BCE, b c e +BCE b c e +BCE's b c e's +BCF b c f +BCG b c g +BCG's b c g's +BCHL b c h l +BCHR's b c h r's +BCHS b c h s +BCI b c i +B. C. J. b c j +bcl- b c l +Bcl- b c l +BCL b c l +BCLR b c l r +BCMA b c m a +BCM b c m +BCMG b c m g +BCMHS b c m h s +BCMS b c m s +BCN b c n +BCN's b c n's +BCP b c p +BCPM b c p m +BCR- b c r +BCR b c r +BCRF's b c r f's +B.C.'s b c's +BC's b c's +BCs b c's +BCS- b c s +BCS b c s +BCSC b c s c +BCSic b c s i c +BCSN b c s n +BCT b c t +BCTC- b c t c +BCTC b c t c +BCU b c u +BCYP b c y p +BdA b d a +BDA b d a +bd b d +B. D. b d +B.D. b d +BD- b d +BD b d +BDBL b d b l +BDC b d c +BDC's b d c's +BDD b d d +Bde b d e +B. D. E. b d e +BDE b d e +BDF b d f +BDI b d i +BDJ b d j +BDK b d k +BDK's b d k's +BDM- b d m +BDMI b d m i +BDN b d n +BDNF b d n f +BDO b d o +BDO's b d o's +BDOS b d o s +BDP b d p +BDR b d r +bds b d s +BDS b d s +BDSM b d s m +BdU b d u +BDU b d u +BDV b d v +B. D. W. b d w +B.E.A. b e a +B. E. b e +B.E. b e +Bedw b e d w +B.E.E. b e e +bef b e f +B.E.F. b e f +BEF b e f +BEF's b e f's +bei b e i +Bei b e i +Beih b e i h +BEIR b e i r +Beis b e i's +bej b e j +B.E.M. b e m +BEMs b e m's +BEP- b e p +BEP b e p +Ber b e r +BER b e r +Bes b e's +B.E.S. b e s +BES b e s +B.F.A. b f a +BFA b f a +bf b f +Bf- b f +B. F. b f +B.F. b f +BF b f +BFCA b f c a +BFC b f c +BFC's b f c's +BFDG b f d g +BFES b f e s +BFF b f f +BFFs b f f's +B. F. G. b f g +BFG b f g +BFG's b f g's +BFI b f i +BFI's b f i's +BFJA b f j a +BFKL b f k l +BFL b f l +BFM b f m +BFO b f o +BFRA b f r a +BFR b f r +BFSA b f s a +BFS b f s +BFT b f t +BFU b f u +bfy b f y +BGAB b g a b +BGA b g a +BGB b g b +B. G. b g +B.G. b g +BGC b g c +BGCI b g c i +BGD b g d +BGEA b g e a +B. G. E. b g e +BGH b g h +B. G. J. S. b g j s +BGL b g l +BglII b g l i i +BGM- b g m +BGM b g m +BGP- b g p +BGP b g p +BGRA b g r a +BGR b g r +BGRC b g r c +BGRSO b g r s o +bgs b g s +B.G.S. b g s +BGS b g s +BGSU b g s u +BGT b g t +Bgy b g y +Bha b h a +B&H b and h +B. H. b h +B.H. b h +BHCC b h c c +BHDP b h d p +BH&E b h and e +Bhe b h e +BHL b h l +BHMA b h m a +BHMO b h m o +BHMs b h m's +BHMT b h m t +BHO b h o +B. H. P. b h p +BHP b h p +BHRT b h r t +BHS b h s +BHSN b h s n +BHSU b h s u +bhttp b h t t p +BHU b h u +BHUSD b h u s d +BHVS b h v s +B. I. b i +B.I. b i +Bie b i e +BIE b i e +biedt b i e d t +Biem b i e m +B.I.G. b i g +B.I.G.'s b i g's +BiH b i h +BIH b i h +B.I.O.L.A. b i o l a +B. J. b j +B.J. b j +BJC- b j c +BJCC b j c c +BJCP b j c p +BJD b j d +B.J.F. b j f +B. J. I. b j i +BJJ b j j +BJP b j p +BJP's b j p's +BJPs b j p's +BJPS b j p s +B.J.'s b j's +B.J.T. b j t +BJT b j t +BJU b j u +BJY b j y +B.K.A. b k a +B&K b and k +BKB- b k b +B. K. b k +B.K. b k +BKCa b k c a +BKC b k c +BKF b k f +BKI b k i +BKN's b k n's +BKO b k o +B.K.R. b k r +BKR b k r +BK's b k's +BKS b k s +BKT b k t +BKTV b k t v +BKV b k v +B. L. b l +B.L. b l +BL- b l +BL b l +BLC b l c +BLCC b l c c +BLCN b l c n +ble b l e +bleg b l e g +blev b l e v +BLG- b l g +BLG b l g +BLH- b l h +bli b l i +BLI b l i +Blla b l l a +BLL- b l l +BLL b l l +BLM- b l m +BLM b l m +BLMC b l m c +BLOU's b l o u's +BLPP b l p p +BLR b l r +BLRC b l r c +BLR&D b l r and d +B.L.S. b l s +BLS b l s +BLT b l t +BLTF b l t f +BLTs b l t's +BLTTA- b l t t a +BLV b l v +BMAA b m a a +B.M.A. b m a +BMA b m a +BMAL b m a l +BMARC b m a r c +BMA's b m a's +B&M b and m +BMB b m b +B. M. b m +B.M. b m +BMCA b m c a +B.M.C. b m c +BMC b m c +bmd b m d +BMD b m d +BME b m e +BMF b m f +BMG b m g +BMHS b m h s +bmi b m i +BMI b m i +BMIC b m i c +BMIR b m i r +BMI's b m i's +BMIT b m i t +BMJ b m j +BMK b m k +BMKG b m k g +BMML- b m m l +BMMO b m m o +BMNA b m n a +BMNH b m n h +BMNP b m n p +BMO b m o +BMP- b m p +BMP b m p +BMPs b m p's +BMR- b m r +BMR b m r +BMRs b m r's +BMSA b m s a +B.M.S. b m s +BMS- b m s +BMS b m s +bmt b m t +B. M. T. b m t +B.M.T. b m t +BMT b m t +BMT's b m t's +BMU b m u +B.M.V. b m v +BMV b m v +B. M. W. b m w +BMW b m w +BMW M b m w +BMW's b m w's +BMX- b m x +BMX b m x +BNA b n a +B&N b and n +bnb b n b +BNB b n b +B. N. b n +B.N. b n +BNC b n c +BNCI b n c i +BNCM b n c m +BNCT b n c t +BND b n d +BNDY b n d y +BNET b n e t +BNF b n f +BNFL b n f l +BNI b n i +BNIC b n i c +BNL b n l +BNO b n o +BNP b n p +BNPH- b n p h +BNR b n r +BNRC b n r c +BNS b n s +BNSF b n s f +BNST b n s t +BNT b n t +BNU b n u +BNY b n y +BNZ b n z +Boac b o a c +BOAC- b o a c +BOAC b o a c +B&O b and o +B. O. b o +B.O. b o +Boc b o c +BoC b o c +BOC- b o c +BOC b o c +BOCs b o c's +BOCS b o c s +boj b o j +Boj b o j +BOJ b o j +B.O.M.B. b o m b +B&O's b and o's +B.O.S. b o s +BOYZZ b o y z z +BPA b p a +BPB b p b +bp b p +B. P. b p +B.P. b p +BP b p +BPC b p c +BPD b p d +BPER b p e r +B.P.H. b p h +BPH- b p h +BPI b p i +BPJ b p j +BPK b p k +bpl b p l +BPL b p l +B. P. M. b p m +BPM b p m +BPMN b p m n +B.P.N. b p n +BPN b p n +B.P.O. b p o +BPO b p o +BPPA b p p a +BPP b p p +BPR b p r +B.P.R.D. b p r d +BPRD b p r d +BPSA b p s a +BP's b p's +BPS b p s +BPSK b p s k +BPUP b p u p +BPV b p v +BPY b p y +B&Q b and q +BQB b q b +B. Q. b q +BQM- b q m +BRAC's b r a c's +BRAM b r a m +BRBDP b r b d p +Brbic b r b i c +B. R. b r +B.R. b r +BRBR b r b r +BRBs b r b's +BRCA b r c a +BRC b r c +BRCC b r c c +BRD b r d +BRF b r f +BRGF b r g f +BRGM b r g m +BRHS b r h s +B.R.I.C.K. b r i c k +BRK b r k +BRL b r l +BRM b r m +BRMS b r m s +Brn- b r n +BRN b r n +BRNC b r n c +Brne b r n e +Brno b r n o +BRNV b r n v +BRP b r p +BRPS b r p s +BRRA b r r a +BRS b r s +BRSCC's b r s c c's +BRT b r t +BRTs b r t's +BRTS b r t s +BRVM b r v m +BRW b r w +BRWK b r w k +Brza b r z a +BRZ b r z +brzu b r z u +brzy b r z y +Brzyk b r z y k +BSAA b s a a +B.S.A. b s a +BSA- b s a +BSA b s a +BSAC b s a c +BsaL b s a l +BSBA b s b a +BSB b s b +BSBI b s b i +B. s b's +B.'s b's +Bs b's +B. S. b s +B.S. b s +BS- b s +BS b s +B.S.C. b s c +BSC b s c +BSCL b s c l +bsd b s d +BSD- b s d +BSD b s d +BSDE b s d e +BSE b s e +BSET b s e t +BSFA b s f a +BSF b s f +BSG b s g +BSh- b s h +BSH b s h +BSIB b s i b +BSI b s i +BSK b s k +BSL- b s l +BSL b s l +BSM b s m +BSME b s m e +BSN b s n +BSNL b s n l +BSOG b s o g +BSPA b s p a +BSP b s p +BSPP b s p p +BSR b s r +BSRN b s r n +BSS b s s +BST- b s t +BSU b s u +BSV b s v +BSX b s x +BTA b t a +BTAF b t a f +BTB- b t b +BTB b t b +BTBD b t b d +bt b t +B. T. b t +B.T. b t +BT- b t +BT b t +BTCA b t c a +BTC- b t c +BTC b t c +BTCC b t c c +BTD b t d +BTF b t f +BTG b t g +BTH b t h +BTK b t k +BTL b t l +BTM b t m +BTN b t n +BTNK b t n k +BTOB b t o b +BTO b t o +BTP b t p +BTR- b t r +BTRC b t r c +BTRDA b t r d a +Btry b t r y +btsan b t s a n +BT's b t's +BTS b t s +BTT b t t +BTU b t u +btus b t u s +BTUs b t u's +BTV b t v +B.U. b u +B.U.M. b u m +BUV b u v +BVA b v a +BVB b v b +BVB's b v b's +B. V. b v +B.V. b v +BVC b v c +BVFB b v f b +BVI b v i +BVM b v m +BVO b v o +BVR b v r +BVRC b v r c +BVT b v t +BVVL b v v l +Bwa b w a +BWAF b w a f +B&W b and w +BWB b w b +BWBR b w b r +B. W. b w +B.W. b w +B.W.C. b w c +BWC b w c +BWF b w f +BWFC b w f c +BWHBC b w h b c +B. W. I. b w i +BWI b w i +BWIR b w i r +BWO b w o +BWR b w r +BWS b w s +BWT b w t +BWTs b w t's +BWV b w v +BWW b w w +bwwtv b w w t v +Bxa b x a +BX b x +Bxe b x e +BXML b x m l +BxPC- b x p c +BYA's b y a's +BYB b y b +B. Y. b y +byc b y c +BYC b y c +BYD b y d +BYFC b y f c +BYG b y g +bygd b y g d +Byk b y k +byn b y n +BYOB b y o b +Byo b y o +BYO b y o +Byou b y o u +BYR b y r +Byrl b y r l +bySLC b y s l c +BYST b y s t +byt b y t +B.Y.U. b y u +BYU b y u +BYU's b y u's +BYX b y x +BYZ b y z +B.Z. b z +BZ b z +BZD b z d +bzhed b z h e d +Bzik b z i k +BZK b z k +BZP b z p +Bzyb b z y b +CAA c a a +CAAC c a a c +CAAC's c a a c's +CAAHEP c a a h e p +CAAM c a a m +CAAS c a a s +CAASE c a a s e +Caat c a a t +CAAT c a a t +C. A. c a +C.A. c a +CA- c a +CA c a +C&A c and a +CAC c a c +CA&CC c a and c c +CACC c a c c +CACM c a c m +Cadw c a d w +CADW c a d w +caeca c a e c a +Cae c a e +CAE c a e +C.A.F.B. c a f b +CAFTT c a f t t +C. A. G. c a g +CAG c a g +CAGM c a g m +C. A. I. c a i +C. A. J. c a j +CALUX c a l u x +CAMC c a m c +CAMLG c a m l g +CAMPCO c a m p c o +C.A.M.'s c a m's +C. A. N. c a n +Capt. captain +C. A. R. c a r +cas c a s +Cas c a's +C.A.S. c a s +CAS c a s +Casc c a s c +CASC c a s c +CASD c a s d +CASF c a s f +CASG c a s g +CASQ c a s q +C. A. T. c a t +C.A.T. c a t +CATV c a t v +C.A.W. c a w +CBA c b a +CBBB c b b b +CBB c b b +CBBC c b b c +CBBS c b b s +CBCA c b c a +cb c b +C. B. c b +C.B. c b +CB- c b +CB c b +cbc c b c +CBC- c b c +CBC c b c +CB&CNS c b and c n s +CBC's c b c's +CBCS c b c s +CBCV- c b c v +CBDA c b d a +CBD c b d +CBDs c b d's +C.B.E. c b e +CBE c b e +CBEF c b e f +CBERS- c b e r s +CBE's c b e's +CBF- c b f +CBF c b f +CBFT- c b f t +CBFX- c b f x +CBGB c b g b +CBGB's c b g b's +CBGBs c b g b's +CBG c b g +CBH c b h +CBHG c b h g +CBI c b i +CBI's c b i's +CbiXS c b i x s +CBKB c b k b +CBK c b k +CBKMT c b k m t +CBKRT c b k r t +CBL c b l +CBL's c b l's +CBM- c b m +CBM c b m +CBME- c b m e +CBMS c b m s +CBMT- c b m t +CBN c b n +CBN's c b n's +CBNT c b n t +CBO- c b o +CBO c b o +CBON- c b o n +CBOT c b o t +CBP- c b p +CBP c b p +CBPE c b p e +CB&Q c b and q +CBR c b r +CBRE c b r e +CBRN c b r n +CBRNE c b r n e +CBSA- c b s a +CBSA c b s a +CBs c b's +CBS- c b s +CBS c b s +C.B.S.E. c b s e +CBSE c b s e +CBS's c b s's +CBT c b t +CBTC c b t c +CBTU c b t u +CBU c b u +CBUT c b u t +CBV- c b v +CBV c b v +CBWT c b w t +CBX c b x +CBY- c b y +CBYT- c b y t +CBYT c b y t +CBZH c b z h +'c c +CcaA c c a a +CCAA c c a a +CCAAT c c a a t +CCAB- c c a b +cca c c a +Cca c c a +CCA c c a +CCAF c c a f +CCAP c c a p +CCAR c c a r +CCAS c c a s +CCB c c b +CCBCC c c b c c +CCBE c c b e +CCCA c c c a +C&C c and c +C. C. c c +C.C. c c +CCCE c c c e +CCCF c c c f +CCCP- c c c p +CCCP c c c p +CCC's c c c's +CCDB c c d b +CCDC c c d c +ccd c c d +CCD c c d +CCDev c c d e v +CCDI c c d i +CCDI's c c d i's +CCDM c c d m +CCDR c c d r +CCD's c c d's +CCDs c c d's +CCEd c c e d +C.C.E.D. c c e d +CCF c c f +CCFO c c f o +CCG c c g +C.C.G.S. c c g s +CCGS c c g s +CCHA c c h a +CCHC c c h c +C.C.H. c c h +CCH c c h +CCHD c c h d +CCHS c c h s +C.C.I.A.A. c c i a a +CCIAA c c i a a +CCi c c i +CCI c c i +CCID c c i d +CCIE c c i e +CCIH c c i h +CCITT c c i t t +CCJ c c j +CCK c c k +CCLC c c l c +CCL c c l +CCMA c c m a +ccm c c m +CCM c c m +CCMM c c m m +CCMP c c m p +CCMS c c m s +CCN c c n +CCNY c c n y +CCO c c o +CCOKC c c o k c +CCOO c c o o +CCOP c c o p +C.C.O.W.E. E.T.F. c c o w e e t f +CCP c c p +CCPD c c p d +CCPL c c p l +CCPN c c p n +CCP's c c p's +CCQ c c q +CCRC c c r c +CCR c c r +CCRCs c c r c's +CCRH c c r h +CCRS c c r s +CCSA c c s a +CCSC c c s c +CC's c c's +CCS c c s +CCSD c c s d +C.C.S.D.N.Y. c c s d n y +C.C.S.M. c c s m +CCSN c c s n +CCSR c c s r +CCS's c c s's +cct c c t +CCT c c t +CCTF c c t f +CCTRN c c t r n +CCTT c c t t +CCTV- c c t v +CCTV c c t v +CCUA c c u a +CCU c c u +ccus c c u s +CCV c c v +CCVG c c v g +C. C. W. c c w +CCWD c c w d +CDAC- c d a c +CDA c d a +CDAT c d a t +CDB c d b +CDBS c d b s +CDCA c d c a +C&D c and d +cdc c d c +CDC- c d c +CDC c d c +cd c d +C. D. c d +C.D. c d +CD- c d +CD c d +CDC's c d c's +CDDB c d d b +CD&DR c d and d r +CDE c d e +CDF c d f +CDFI c d f i +CDFW c d f w +CDG- c d g +CDH c d h +CDi c d i +CDI c d i +CDISC c d i s c +CDI's c d i's +CDK c d k +CDKN c d k n +CDL c d l +CDLI c d l i +CDLS c d l s +CDMA- c d m a +CDMA c d m a +CDM- c d m +CDM c d m +cDNA c d n a +cdnas c d n a s +cDNAs c d n a's +CDN c d n +CDN's c d n's +CDO c d o +CDOs c d o's +CDPC c d p c +CDP c d p +CDP's c d p's +CDPs c d p's +CD&R c d and r +CDR c d r +cds c d s +CD's c d's +CDs, c d's +CDs c d's +CDS- c d s +CDS c d s +CDSG c d s g +CDSP c d s p +CDSPCo c d s p c o +CdtA c d t a +CDT c d t +CdTe c d t e +CDTi c d t i +CDTI c d t i +CDT's c d t's +CDTV c d t v +CDTV's c d t v's +CDU- c d u +CDU c d u +CDU's c d u's +CD&V c d and v +CDV c d v +CDW c d w +CDWP c d w p +Cec c e c +CEC c e c +c'e c e +ce c e +C'e c e +Ce c e +C. E. c e +C.E. c e +CE, c e +CE c e +CECP c e c p +CEDA c e d a +Ced- c e d +CED- c e d +CEDO c e d o +CEDR c e d r +CEDS c e d s +CEEBA c e e b a +CEEB c e e b +CEEH c e e h +Ceel c e e l +CEEOL c e e o l +CEESA c e e s a +CEFC c e f c +CEF c e f +Cefn c e f n +CEFP c e f p +CEFR c e f r +CEGB c e g b +C.E.G. c e g +C. E. H. c e h +CEI c e i +C&EI's c and e i's +C. E. M. c e m +C&EN c and e n +cen c e n +Cen c e n +CEN c e n +ceo c e o +C.E.O. c e o +CEO c e o +Ceol c e o l +CEOP c e o p +Ceorl c e o r l +Ceos c e o's +CEO's c e o's +CEOs' c e o's +CEOs c e o's +C.E.P. c e p +cer c e r +Cer c e r +CER c e r +CERP c e r p +CERPER c e r p e r +CERS c e r s +ces c e s +Ces c e's +C.E.S. c e s +CES c e s +CESL c e s l +CESO c e s o +CETB c e t b +C. E. T. c e t +CEV c e v +CEVCP c e v c p +CEVG c e v g +CEW c e w +Cex c e x +CeX c e x +CEX c e x +CEZ c e z +CEZMS c e z m s +CFAB c f a b +Cfa c f a +CFA c f a +CFAS c f a s +C. F. B. c f b +CFB c f b +CFBG c f b g +CFBISD c f b i s d +CFB's c f b's +C&F c and f +CFC- c f c +CFC c f c +cf c f +C. F. c f +C.F. c f +CF- c f +CF c f +CFCF- c f c f +CFC's c f c's +CFDA c f d a +CFD c f d +C.F.E. c f e +CFE c f e +CFFB- c f f b +CFFC c f f c +C. F. F. c f f +CFF- c f f +CFF c f f +C. F. H. c f h +CFH c f h +CFHS c f h s +CFI c f i +CFIT c f i t +CFJC- c f j c +CFJL c f j l +CFJP- c f j p +CFJR c f j r +CFL c f l +CFL's c f l's +CFM c f m +CFMEU c f m e u +CFMI c f m i +CFMM- c f m m +CFND c f n d +CFNY- c f n y +CFNY c f n y +C.F.O.A. c f o a +CFOA- c f o a +CFO c f o +CFOP c f o p +CFOs c f o's +CFPI c f p i +CFPL c f p l +CFQ c f q +CFRB c f r b +C.F.R. c f r +CFR- c f r +CFR c f r +CFRP c f r p +CFS c f s +CFSCI c f s c i +CFSP c f s p +CFTA c f t a +CFT c f t +CFTK- c f t k +CFTO- c f t o +CFTO c f t o +CFTR c f t r +CFTs c f t's +CFTU- c f t u +CFTXOP c f t x o p +CFU c f u +CFX- c f x +CFYN c f y n +CFZ c f z +CGA c g a +CGAP c g a p +CGB c g b +C&G c and g +CGCFAD c g c f a d +C. G. c g +C.G. c g +CGCRI c g c r i +CGCS c g c s +CGD c g d +CGDK c g d k +CGE c g e +CGF c g f +CGFP c g f p +CGG c g g +CGI- c g i +CGI c g i +CGIL c g i l +CGKD c g k d +cgl c g l +CGL c g l +CGP c g p +CGPM c g p m +CGPME c g p m e +CGPW c g p w +CGR c g r +CGRP c g r p +CGS c g s +CGSS c g s s +CGT c g t +CGTG c g t g +CGTIC c g t i c +CGU c g u +CGW c g w +CGW's c g w's +C.H.B. c h b +CHBM c h b m +CHBX c h b x +C&H c and h +C.H.C. c h c +C. H. c h +C.H. c h +CHCH- c h c h +CHCO c h c o +CHD c h d +CHDK c h d k +CHDS c h d s +CHF c h f +CHFI c h f i +C.H.G. c h g +CHGO c h g o +CHH c h h +CHHOTO c h h o t o +CHHS c h h s +Chhu c h h u +Chirs c h i r's +CHISMS c h i s m s +CHISZ c h i s z +CHJX- c h j x +CHK c h k +CHL c h l +CHLR c h l r +CHL's c h l's +C. H. M. c h m +CHM c h m +CHN c h n +CHP- c h p +CHP c h p +CHP'den c h p d e n +CHPs c h p's +CHPS c h p s +CHR c h r +CHRDI c h r d i +CHRGD c h r g d +C.H.S. c h s +CHS c h s +CHSE c h s e +CHT c h t +CHV c h v +CHX c h x +Chy c h y +CHYR c h y r +CIAC c i a c +Cia c i a +C.I.A. c i a +CIA- c i a +CIA, c i a +CIA c i a +CIC- c i c +CIC c i c +ci c i +Ci c i +C. I. c i +C.I. c i +CI c i +Cicic c i c i c +CICL c i c l +C.I.C.M. c i c m +CICO c i c o +CIC's c i c's +CICT- c i c t +C.I.D. c i d +CIFF c i f f +CIFL c i f l +CIHE c i h e +CIHR c i h r +cii c i i +CII c i i +CIID c i i d +CIJ c i j +CIK- c i k +CIKM c i k m +C.I.L. c i l +CIL c i l +cim c i m +Cim c i m +CIM c i m +Cio c i o +C.I.O. c i o +CIO c i o +CIOs c i o's +CIOS- c i o s +CIOT- c i o t +CIOT c i o t +CIPCA c i p c a +Cip c i p +C.I.P. c i p +CIP c i p +CIPD c i p d +CIPFA c i p f a +CIPM c i p m +CISA c i s a +CISC c i s c +cis- c i s +cis c i s +CIs c i's +CIS- c i s +CIS c i s +CITB c i t b +cit c i t +CIT- c i t +CIT c i t +CiTD c i t d +CITM c i t m +CITN c i t n +CITP c i t p +citS c i t s +CITV c i t v +CITWF c i t w f +CitX c i t x +Ciuc c i u c +Ciu c i u +CiU c i u +CIU c i u +CIUT c i u t +CIVD c i v d +CIW c i w +CIWS c i w s +CIX c i x +Cixi c i x i +CIY c i y +CIZN c i z n +CJA c j a +CJAD- c j a d +CJBHL c j b h l +CJCA c j c a +CJCB c j c b +CJC c j c +C. J. c j +C.J. c j +CJD c j d +C. J. E. c j e +CJFL c j f l +CJFT c j f t +CJFX c j f x +CJGC c j g c +CJGO c j g o +CJHL c j h l +CJIC- c j i c +CJK c j k +CJL c j l +CJLS c j l s +CJM c j m +CJNT c j n t +CJOC c j o c +CJOT- c j o t +CJP c j p +C. J. R. c j r +C.J.R. c j r +CJR c j r +CJSC c j s c +C. J. S. c j s +CJS c j s +CJT c j t +CJTF- c j t f +CJUF c j u f +CJWC c j w c +CKAC c k a c +CKC c k c +C. K. c k +C.K. c k +CK c k +CKD c k d +CKE c k e +CKLM c k l m +CKLN- c k l n +CKLQ c k l q +CKLW c k l w +CKND- c k n d +CKRY- c k r y +CKS c k s +CKSO- c k s o +CKSO c k s o +CKUE- c k u e +CKVL c k v l +C. K. W. c k w +CKWK c k w k +ckx c k x +CKXX c k x x +CKY c k y +CLAC c l a c +CLA c l a +CLAS c l a s +CLB- c l b +CLBs c l b's +clc c l c +CLC- c l c +CLC c l c +cl c l +C. L. c l +C.L. c l +CL- c l +CL c l +CLCN c l c n +CLD c l d +Cle c l e +CLE c l e +CLF c l f +C.L.G. c l g +CLG c l g +CLHIA c l h i a +CLHS c l h s +CLI c l i +clk- c l k +CLK c l k +cllr c l l r +CLM c l m +CLMD c l m d +CLMP c l m p +CL&N c l and n +Clo c l o +CL&P c l and p +CLP c l p +CLRC c l r c +C. L. R. c l r +C.L.R. c l r +CLR c l r +CLRK c l r k +CLSA c l s a +CLSC c l s c +CLs c l's +CLS- c l s +CLS c l s +CLT c l t +CLTD c l t d +CLTPA c l t p a +CLTs c l t's +Clu c l u +CLU c l u +CLUK c l u k +CLV c l v +Clwyd c l w y d +CLX c l x +C&MA c and m a +CMAC c m a c +C.M.A. c m a +CMA c m a +CMAL c m a l +CMAP's c m a p's +CMAS- c m a s +CMAS c m a s +CMAT c m a t +C.M.B. c m b +CMB c m b +CMCC c m c c +CMCCDI c m c c d i +C.M.C. c m c +CMC- c m c +CMC c m c +cm- c m +cm c m +C. M. c m +C.M. c m +CM- c m +CM c m +CMCs c m c's +CMCS c m c s +CMD c m d +Cmde c m d e +CMDISE c m d i s e +cmdr c m d r +CMEC c m e c +CME c m e +CMF c m f +CMFS c m f s +C. M. G. c m g +C.M.G. c m g +CMG c m g +CMHC c m h c +CMH c m h +CMHR c m h r +CMI c m i +CMJ c m j +CMLA c m l a +cml c m l +C.M.L. c m l +CML c m l +CMLDP c m l d p +CMLL c m l l +CMLL's c m l l's +CMM c m m +CMMI c m m i +CMMs c m m's +cmn c m n +CMO c m o +CMO's c m o's +CMPC c m p c +CMP c m p +CMPS c m p s +CMPSO c m p s o +CMPV c m p v +CMQ- c m q +CMQ c m q +CMQT c m q t +C. M. R. c m r +CMR c m r +CMSAF c m s a f +C.M.S. c m s +CMS c m s +C.M.S.L. c m s l +CMT c m t +Cmte c m t e +CMT's c m t's +CMTU c m t u +CMUCL c m u c l +CMU c m u +CMU's c m u's +CMV c m v +cmvIL- c m v i l +CMVM c m v m +CMVSS c m v s s +CMWF c m w f +CMXXII c m x x i i +CMYK c m y k +CNA c n a +CNBC- c n b c +CNBC c n b c +CNB c n b +CNC c n c +cn c n +C. N. c n +C.N. c n +CN c n +CNC's c n c's +CND c n d +CNDL c n d l +CNEB c n e b +CNE c n e +CNF c n f +CNFE c n f e +CNFK c n f k +CNG- c n g +CNG c n g +CNI c n i +CNJ c n j +CNJF c n j f +CNK c n k +C. N. L. c n l +C.N.L. c n l +CNL c n l +CNMs c n m's +CNN- c n n +CNN c n n +CNO c n o +CNPC c n p c +CNP c n p +cnr c n r +CNR c n r +CNRMA c n r m a +CNRS c n r s +CNSA c n s a +CNs c n's +CNS c n s +CNSS c n s s +CNSW c n s w +CNTA c n t a +CNT c n t +CNTE c n t e +CNTI c n t i +CNTK c n t k +CNTN c n t n +CNTs c n t's +CNTS c n t s +CNTV c n t v +CNU c n u +Cnut's c n u t's +CNV c n v +CNVs c n v's +CNW c n w +CNWS c n w s +CNZ c n z +COA c o a +C&O c and o +C. O. c o +C.O. c o +CO- c o +Co. company +C.O.D. c o d +Col. colonel +C.O.M.L. c o m l +Cov c o v +COX c o x +CPAC c p a c +CPA c p a +CPAP- c p a p +CPA's c p a's +CPB c p b +CPBL c p b l +C&P c and p +CPC- c p c +CPC c p c +cp c p +C. P. c p +C.P. c p +CP- c p +CP c p +CPC's c p c's +CPCs c p c's +CPCS c p c s +CPD c p d +CPDM c p d m +CPD's c p d's +CPEB c p e b +CPEC c p e c +C.P.E. c p e +CPE c p e +CPEO c p e o +CPF c p f +CPFO c p f o +CPG c p g +CPHC c p h c +CPH c p h +CPI- c p i +CPI c p i +CPIFL c p i f l +CPIM c p i m +CPI's c p i's +CPJ c p j +CPL c p l +CPLP c p l p +CPM c p m +CPN c p n +CPNI c p n i +cpo c p o +C.P.O. c p o +CPO c p o +CPOE c p o e +CPPA c p p a +CPPCC c p p c c +CPP c p p +CPPE c p p e +CPPIB c p p i b +CPPM c p p m +CPQ c p q +CPR c p r +CPRS c p r s +CPSA c p s a +CPSC c p s c +CPS- c p s +CPS c p s +CPSF c p s f +CPSK c p s k +CPSL c p s l +CPSU c p s u +CPTA c p t a +CPT c p t +CPTM c p t m +CPUC c p u c +CPU- c p u +CPU c p u +CPUID c p u i d +CPUSA c p u s a +CPUs c p u's +CPUT c p u t +CPVA c p v a +CPV c p v +CPW c p w +CQAC c q a c +CQBR c q b r +CQC c q c +cq c q +CQ c q +CQD c q d +CQDs c q d's +CQM c q m +CQO c q o +CQU c q u +CQW- c q w +CRABP- c r a b p +C. R. A. c r a +CRA c r a +C.R.A.Z.Y. c r a z y +CRBC c r b c +CRC c r c +cr- c r +cr c r +Cr- c r +C. R. c r +C.R. c r +CR- c r +CR c r +Cre c r e +CRE c r e +CRF- c r f +CRF c r f +CRG c r g +CRHK c r h k +CRH's c r h's +CRHS c r h s +CRJ c r j +crk c r k +Crkva c r k v a +crkve c r k v e +CRL- c r l +CRL c r l +CRLs c r l's +crm c r m +CRM c r m +CRNA c r n a +CRN- c r n +CRN c r n +Crne c r n e +Crni c r n i +Crno c r n o +CRP c r p +CRPF c r p f +CRRF c r r f +CRs c r's +CRS- c r s +CRS c r s +CRSP c r s p +CRST c r s t +CRTC c r t c +CRT c r t +CRTC's c r t c's +CRUA c r u a +CRWRC c r w r c +CRW's c r w's +CRX c r x +CSAC c s a c +csa c s a +Csa c s a +C.S.A. c s a +CSA c s a +CSAH c s a h +Csak c s a k +Csaky c s a k y +CSAR c s a r +CSAT c s a t +CSAV c s a v +CSB c s b +C&S c and s +C.S.C. c s c +CSC c s c +CSCD c s c d +CSCE c s c e +CSCR c s c r +cs c s +C.s c's +Cs c's +C. S. c s +C.S. c s +CS- c s +CS c s +CSC's c s c's +CSDA c s d a +CSDB c s d b +CSD c s d +CSDL c s d l +CSDN c s d n +CSDP c s d p +CSD's c s d's +CSDS c s d s +CSE c s e +C&SF c and s f +CSF c s f +CSFL c s f l +CSFSO c s f s o +CSGA c s g a +csg c s g +CSG c s g +CSG's c s g's +CSH c s h +CSHL c s h l +CSIC c s i c +CsI c s i +CSI c s i +CSIDC c s i d c +CSILUS c s i l u s +CSIO c s i o +CSIR c s i r +CSIRO c s i r o +CSis c s i's +CSIs c s i's +CSIS- c s i s +CSIS c s i s +CSIU c s i u +CSKA- c s k a +CSKA c s k a +CSKA's c s k a's +CSK c s k +CSK's c s k's +CSL- c s l +CSL c s l +CSLI c s l i +CSM c s m +CSNB c s n b +CSNBX c s n b x +CSN c s n +CSNJ- c s n j +CSN's c s n's +CSOB c s o b +CSO c s o +CSO's c s o's +CSOs c s o's +cspA c s p a +csp c s p +CSP c s p +CSPH c s p h +CSPI c s p i +CSPRA c s p r a +CSPs c s p's +CSRA c s r a +CSR c s r +CSRS c s r s +CSRT c s r t +CSRTT c s r t t +CSSA c s s a +CSS c s s +CSSD c s s d +CSSE c s s e +CSSL c s s l +CSSNCs c s s n c's +CSSR c s s r +CSSS c s s s +CSST c s s t +CSTB c s t b +CST c s t +CSTF c s t f +CSTI c s t i +CSTS c s t s +CSUB c s u b +CSU c s u +CSU's c s u's +Csuz c s u z +CSV c s v +CSW c s w +CSWF c s w f +CSWIP c s w i p +CSWS c s w s +CSX c s x +CSX's c s x's +CSXT c s x t +C.T.A. c t a +CTA c t a +CTAF c t a f +CTARL c t a r l +CTA's c t a's +CTAs c t a's +CTBT c t b t +CTBTO c t b t o +C&T c and t +CTC c t c +ct c t +C. T. c t +C.T. c t +CT- c t +CT c t +CTD c t d +cte c t e +CTE c t e +CTEQ c t e q +CTF c t f +CTGF c t g f +CTI c t i +CTLA- c t l a +CTLA c t l a +CTL c t l +CTLL- c t l l +CTLLS c t l l s +CTLs c t l's +CTM c t m +CTMP- c t m p +CTN c t n +CTO c t o +CTP c t p +ctr c t r +CTRG c t r g +Ctrip c t r i p +CTRL c t r l +CTRP c t r p +CTR's c t r's +CTS- c t s +CTS c t s +CTSD c t s d +CTSF c t s f +CTSG c t s g +CTUC c t u c +CTU c t u +CTUG c t u g +CTU's c t u's +CTVA- c t v a +CTVA c t v a +ctv c t v +CTV c t v +CTV's c t v's +CTX c t x +CTZ c t z +cu c u +C. U. c u +CU- c u +CU c u +C.U.L. c u l +C. U. S. c u s +C. U. T. c u t +CUW c u w +Cuyp c u y p +CVA- c v a +CVA c v a +CVAK c v a k +CVARG c v a r g +CVB c v b +C&V c and v +CVC c v c +cvcp c v c p +'cv c v +cv- c v +cv c v +C. V. c v +C.V. c v +CV- c v +CV c v +CVD c v d +CVEs c v e's +CVG c v g +CVI c v i +CVJM c v j m +CVM c v m +C.V.N. c v n +CVN- c v n +CVN c v n +C.V.O. c v o +CVO c v o +CVODE c v o d e +CVOs c v o's +CVP c v p +CVR c v r +CVs c v's +CVS c v s +CVSNT c v s n t +CVT c v t +CVTs c v t's +CVV c v v +CVVHDF c v v h d f +CVW- c v w +CVW c v w +CVZ c v z +CWA c w a +CWBI c w b i +CWBL c w b l +CWC c w c +C. W. c w +C.W. c w +CWD c w d +CWDS c w d s +C. W. E. c w e +CWE c w e +Cwej c w e j +Cwele c w e l e +CWF c w f +CWGC c w g c +CWG c w g +CWI c w i +C. W. K. c w k +cwm c w m +CWM c w m +CWO c w o +CWPA c w p a +CWR c w r +cwrt c w r t +CWRU c w r u +CWSAC c w s a c +CW's c w's +CWS c w s +CWTS c w t s +CWUR c w u r +C. W. Y. c w y +CWYFL c w y f l +CXAM- c x a m +CXC- c x c +CXC c x c +CXCL c x c l +CXCR c x c r +C.X. c x +C. Y. c y +C.Y. c y +cyl c y l +Cyn c y n +CYOG c y o g +CyP c y p +CYP c y p +Cyrl c y r l +Cys- c y s +Cys c y's +CYS c y s +CYSK c y s k +CYSY c y s y +czci c z c i +cz c z +C.Z. c z +CZ- c z +CZ c z +CZE c z e +CZM c z m +CZW c z w +czy c z y +DAAD d a a d +Daai d a a i +D. A. d a +D.A. d a +D&AD d and a d +D. A. H. d a h +D.A.M. d a m +D.A.N.C.E. d a n c e +D. A. T. d a t +Dav d a v +D.A.V. d a v +DAV d a v +DAV's d a v's +dBa d b a +DBA d b a +DBC d b c +DBCS d b c s +db d b +D. B. d b +D.B. d b +DB- d b +DB, d b +DB d b +DBE d b e +DBI d b i +DBL d b l +DBLE d b l e +DBLP d b l p +DBM d b m +DBMS d b m s +DBR d b r +DB's d b's +D. B. S. d b s +DBS d b s +DBU d b u +DBYC d b y c +DBZ d b z +DCA- d c a +DCA d c a +DCAF d c a f +DCAM d c a m +DCB d c b +DCCC d c c c +DCCCD d c c c d +DCC d c c +DCCG d c c g +DCCs d c c's +D&C d and c +dc d c +D. C. d c +D.C. d c +DC- d c +DC d c +DCE d c e +DCFC d c f c +DCF d c f +DCG- d c g +DCG d c g +DCH d c h +DCI d c i +DCK d c k +D.C.L. d c l +DCL d c l +DCMA d c m a +DCM d c m +DCMF d c m f +DCMI d c m i +DCMS d c m s +DCO- d c o +DCO d c o +DCom d c o m +DCOM d c o m +DCP d c p +DCR- d c r +DCR d c r +D.C.'s d c's +DC's d c's +DCs d c's +DCS d c s +DCSD d c s d +DCSO d c s o +DCSPER d c s p e r +DCSS d c s s +DCTA d c t a +DCT d c t +DCU d c u +DCUO d c u o +DCUs d c u's +DDA d d a +DDB d d b +D.D.C. d d c +DDC- d d c +DDC d d c +D&D d and d +D. D. d d +D.D. d d +DDF d d f +DDG- d d g +DDG d d g +DDHH d d h h +DDi d d i +DDI d d i +DDK d d k +DDLJ d d l j +DDM d d m +DDMS d d m s +DDN d d n +DDNOS d d n o s +Ddoc d d o c +DDO d d o +DDP d d p +DDR- d d r +DDR d d r +DDRMAX d d r m a x +DDSB d d s b +D.D.S. d d s +DDS d d s +DDT d d t +DDTV d d t v +Ddu d d u +DDU- d d u +DDWFTTW d d w f t t w +DDX d d x +DEAE- d e a e +D.E.A.F. d e a f +Deah d e a h +DEA's d e a's +DECC d e c c +DEC d e c +DECS d e c s +DEDD d e d d +D. E. d e +D.E. d e +deg d e g +D. E. G. d e g +DEG d e g +DEH d e h +DEHP d e h p +dei d e i +Dei d e i +DeI d e i +DEI d e i +Deijm d e i j m +Deip d e i p +D. E. J. d e j +D.E.R. d e r +dez d e z +Dez d e z +Dfa d f a +DFA d f a +DFB- d f b +DFB d f b +D.F.C. d f c +DFC d f c +DFCs d f c's +D. F. d f +D.F. d f +DFD's d f d's +DfE d f e +D.F.E. d f e +DFE d f e +DfES d f e s +DFF d f f +DFG d f g +DFJ d f j +DFL d f l +DFM d f m +DFOA d f o a +DFO d f o +DFP d f p +DFPs d f p's +DFRC d f r c +DFR d f r +DFs d f's +DFS d f s +DFT d f t +DFW d f w +dfx d f x +DFX d f x +DGAC d g a c +DGA d g a +DGALS d g a l s +DGAP d g a p +DGB d g b +DGCA d g c a +DGC d g c +DGCR d g c r +D&G d and g +D. G. d g +D.G. d g +D. G. E. d g e +DGE d g e +DGERT d g e r t +DGF d g f +DGFI d g f i +DGH d g h +DGK d g k +D. G. L. R. d g l r +DGM d g m +DGMO d g m o +DGP d g p +DGPS d g p s +DGR d g r +Dha d h a +DHA d h a +DHB d h b +DHBs d h b's +DHCP d h c p +D&H d and h +dh d h +D. H. d h +D.H. d h +DH- d h +DH d h +DHEA d h e a +d'HEC d h e c +dhe d h e +DHE d h e +Dheku d h e k u +DHFS d h f s +DHHC d h h c +DHH d h h +DHL d h l +DHMO d h m o +DHNC d h n c +D.H.N. d h n +DHNS d h n s +Dhod d h o d +dhol d h o l +Dhol d h o l +Dhoo d h o o +Dhor d h o r +Dhou d h o u +DHP d h p +DHPH d h p h +DHQ d h q +dhr d h r +D.H.R. d h r +DHRS d h r s +DHS d h s +DHSs d h s's +DHSS d h s s +DHT d h t +DHV d h v +DHW d h w +DHX d h x +D.I. d i +DIFC d i f c +Dif d i f +DIF d i f +Dir d i r +DIR d i r +Dirs d i r's +DISD d i s d +div d i v +Div d i v +DIV d i v +DIW d i w +DIX d i x +diy d i y +DiY d i y +D.I.Y. d i y +DIY d i y +Diyn d i y n +DIYs d i y's +D. J. A. d j a +Djam d j a m +dj d j +D. J. d j +D.J. d j +DJ- d j +DJ d j +DJ'ed d j e d +Djem d j e m +D.J.H. d j h +Dji d j i +DJI d j i +D.J.N. d j n +Djo d j o +D'Jok d j o k +djr d j r +DJR d j r +D.J.s d j's +DJ's d j's +DJs d j's +D. J. S. d j s +D.J.T. d j t +DjVu d j v u +DJVU d j v u +DjVus d j v u's +D. J. Y. d j y +DKC d k c +D. K. d k +D.K. d k +DKI d k i +DKIM d k i m +DKK d k k +DKNY d k n y +DKPP d k p p +D. K. R. d k r +DKW d k w +dla d l a +DLA d l a +DLC d l c +DLD d l d +dl d l +D. L. d l +D.L. d l +DL- d l +DL d l +DLE d l e +DLEG- d l e g +D.L.F. d l f +DLF- d l f +DLF d l f +DLFN d l f n +DLG d l g +D. L. H. d l h +DLHS- d l h s +dlia d l i a +DLI d l i +DLJ d l j +DLJ's d l j's +DLK d l k +DLL d l l +DL&LR d l and l r +DLL's d l l's +DLLs d l l's +DLM d l m +DLMPST d l m p s t +DLNA- d l n a +dlo d l o +DLP- d l p +DLP d l p +DLPFC d l p f c +DLP's d l p's +DLR d l r +DLS d l s +DLSS d l s s +DLSU d l s u +DLSU's d l s u's +DLT- d l t +DLT d l t +DL&W d l and w +DLZ d l z +DMAA d m a a +DMA d m a +DMBA- d m b a +DMB d m b +DMBT d m b t +DMCA d m c a +DMCC d m c c +DMC- d m c +DMC d m c +DMCM d m c m +D.M.C.'s d m c's +D&M d and m +DMD d m d +DMDK d m d k +D. M. d m +D.M. d m +DM- d m +DM d m +DMDS d m d s +DMed d m e d +DM&E d m and e +DME d m e +DMF d m f +DMG d m g +DMGT d m g t +DMHS d m h s +DMI d m i +DMK d m k +DML d m l +DMN d m n +DMO d m o +DMOFD d m o f d +DMOZ d m o z +DMP d m p +DMRC d m r c +DMR d m r +DMRG d m r g +DMs d m's +DMS- d m s +DMS d m s +DMSII d m s i i +DMT- d m t +DMT d m t +DMTF d m t f +DMU d m u +DMUs d m u's +DMV d m v +DMX d m x +DMZ d m z +Dna d n a +DnA d n a +DNA- d n a +DNA d n a +DNADTC d n a d t c +DNAITV d n a i t v +DNAJA d n a j a +DNAJC d n a j c +DnaJ d n a j +DNAN d n a n +DNB d n b +DNBY d n b y +DNC d n c +DNCG d n c g +DND d n d +D. N. d n +D.N. d n +Dnepr d n e p r +dnes d n e s +Dnes d n e's +DNES d n e s +DNF d n f +DNF's d n f's +DNFs d n f's +DNG d n g +DNJO d n j o +DNKA d n k a +DNK d n k +DNM d n m +DNMT d n m t +DNP d n p +DNQ d n q +DNR d n r +DNSBLs d n s b l's +dns d n s +D.N.S. d n s +DNS d n s +DNSS d n s s +DNT d n t +DNU d n u +DNV- d n v +DNVP d n v p +DNX- d n x +D.O.A. d o a +D.O.A.'s d o a's +DOAS d o a s +D.O.B. d o b +D.O.C. d o c +DOCG d o c g +dod d o d +Dod d o d +DoD d o d +DOD d o d +D. O. d o +D.O. d o +DOD's d o d's +D.O.E. d o e +DOGNY d o g n y +DOHC d o h c +D'oh d o h +DOJ d o j +DOJ's d o j's +dok d o k +Dok- d o k +Dok d o k +domt d o m t +Domt d o m t +Dop d o p +DOP d o p +DOPR d o p r +DOSAAF d o s a a f +DOTC d o t c +dotCMS d o t c m s +DOTD d o t d +DOTP d o t p +DotSVN d o t s v n +Douw d o u w +DOW d o w +DOXO d o x o +dozd d o z d +d'OZ d o z +Dozhd d o z h d +DPAA's d p a a's +DPA d p a +DPB d p b +DPD d p d +dp- d p +dp d p +D. P. d p +D.P. d p +DP- d p +DP d p +DPF- d p f +DPG d p g +DPJ- d p j +DPJ d p j +DPMA d p m a +DPM d p m +DPN d p n +DPNM d p n m +DPO d p o +dpon d p o n +DPP d p p +DPP's d p p's +DPPX d p p x +DPR d p r +DPRK- d p r k +DPRK d p r k +DPRK's d p r k's +DProf d p r o f +DPRP d p r p +DPRT d p r t +DPs d p's +DPS d p s +DPSS d p s s +DPSSL d p s s l +DPT d p t +DPW d p w +DPW's d p w's +DPX d p x +DPYS d p y s +D.Q. d q +DQOY d q o y +DQV d q v +DRB d r b +DRBG d r b g +DRB's d r b's +DRCA d r c a +drc d r c +DRC- d r c +DRC d r c +DRDA d r d a +DRDC d r d c +DRD d r d +Dr. doctor +D. R. d r +D.R. d r +DRG d r g +dri d r i +Dri d r i +DRI d r i +DRIs d r i's +DRITTES d r i t t e s +Driu d r i u +Driv d r i v +DRJTBC d r j t b c +DRK d r k +DRL d r l +DRLR d r l r +DRMC d r m c +D.R.M. d r m +DRM- d r m +DRM d r m +Drnis d r n i's +Dro d r o +DRO d r o +DRP d r p +Drs. doctors +DR's d r's +DRS d r s +Drska d r s k a +DRT d r t +druj d r u j +DSA d s a +DsbA d s b a +DS&BB d s and b b +DSB- d s b +DSB d s b +DSBL d s b l +D.S.C. d s c +DSC- d s c +DSC d s c +D.S.C.H. d s c h +DSCM d s c m +DSCR d s c r +DSCS d s c s +DSDB d s d b +DSD d s d +DSDM d s d m +dsDNA d s d n a +DSDP d s d p +ds d s +Ds d's +D. S. d s +D.S. d s +DS- d s +DS d s +DSDS- d s d s +DSDS d s d s +DSG d s g +DSGi d s g i +DSIC d s i c +dsi d s i +DSi d s i +DSI d s i +DSIF d s i f +DSIR d s i r +DSJ d s j +dsl d s l +D. S. L. d s l +DSL- d s l +DSL d s l +DSLR d s l r +DSMB d s m b +D.S.M. d s m +DSM- d s m +DSM d s m +DSMs d s m's +dsn d s n +D&SNGRR d and s n g r r +DSNP d s n p +D.S.O. d s o +DSO d s o +DSPA d s p a +DSP d s p +DSR- d s r +DSR d s r +DS's d s's +DSS d s s +DST d s t +dsu' d s u +DSU d s u +DSV d s v +DSW d s w +DSX d s x +DSZG d s z g +D.T.A. d t a +DTA- d t a +DTA d t a +DTC d t c +DTC's d t c's +DTCs d t c's +DTDC d t d c +DTD d t d +DTDs d t d's +D. T. d t +D.T. d t +DT&E d t and e +DTE d t e +DTEK d t e k +DTES d t e s +DTG d t g +DTH d t h +DTIC d t i c +DT&I d t and i +DTI d t i +DTLA d t l a +dtl d t l +DTL d t l +DTLLS d t l l s +DTLS d t l s +DTM d t m +DTO d t o +DTPA d t p a +DTPC d t p c +DTRA d t r a +DTR d t r +DTRS d t r s +DTs d t's +DTS- d t s +DTS d t s +DTT d t t +D.T.U. d t u +DTU d t u +DTV d t v +dty d t y +D.U. d u +dva d v a +Dva d v a +DVA d v a +DVB- d v b +DVB d v b +DVBE d v b e +DVBIC d v b i c +DVC d v c +DVCs d v c's +DVCS d v c s +dvd d v d +DVD- d v d +DVD, d v d +DVD d v d +DVD's d v d's +DVDs d v d's +dv d v +D. V. d v +D.V. d v +DV- d v +DV d v +dve d v e +dvfb d v f b +DVFS d v f s +DVGS d v g s +DVHS d v h s +dvi d v i +DVI d v i +dvije d v i j e +Dvin d v i n +Dvir d v i r +DVLA d v l a +DVLA's d v l a's +DVL- d v l +DVL d v l +DVM d v m +DVP d v p +DVR d v r +DVRJRBC d v r j r b c +DVR's d v r's +DVSC d v s c +DVS d v s +DVT d v t +DVV d v v +DVVT d v v t +DVWK d v w k +Dwa d w a +DWA d w a +D&W d and w +DWDM- d w d m +D. W. d w +D.W. d w +Dwedw d w e d w +DWET d w e t +DWF d w f +DWGD- d w g d +DWG d w g +D.W.H. d w h +Dwi d w i +DWI d w i +DWM d w m +DWN d w n +DWNU d w n u +DWP d w p +dwr d w r +DWR d w r +DWRFC d w r f c +DWS d w s +DWT d w t +DWTS d w t s +dwur d w u r +DWWX- d w w x +dxa d x a +DXA d x a +dx d x +D.X. d x +DX- d x +DX d x +DXED- d x e d +DXED d x e d +DXers d x e r's +DXII d x i i +DXK d x k +DXP- d x p +DXR d x r +DXZ d x z +Dydd d y d d +dy d y +Dy d y +D. Y. d y +DY d y +DYFI d y f i +DYIS d y i s +DYNC d y n c +dypl d y p l +dyr d y r +DYSP d y s p +DZBB d z b b +DZB d z b +dz d z +DZEC d z e c +DZEM d z e m +dzis d z i s +DZNE d z n e +Dzog d z o g +Dzor d z o r +DZRH d z r h +DZRH's d z r h's +DZS d z s +Dzus d z u's +DZXL d z x l +DZZ d z z +EAAB e a a b +EAA e a a +EAAP e a a p +E. A. C. e a c +EAC e a c +EACEF e a c e f +EACS e a c s +EAD e a d +E. A. e a +E.A. e a +EAF e a f +EAGE e a g e +E.A.H. e a h +EAH e a h +EAI e a i +EAIE e a i e +E. A. J. e a j +EAMs e a m's +EAP- e a p +EAP e a p +E.A.R.L. e a r l +E.A.R.T.H. e a r t h +EASA e a s a +EASB e a s b +eas e a s +Eas e a's +EA's e a's +EAS e a s +Eav e a v +EAX e a x +E. B. e b +E.B. e b +EBL e b l +Ebn e b n +EBNF e b n f +EBOV e b o v +E&BR e and b r +EBR- e b r +EBR e b r +Ebru e b r u +EBS e b s +EBSL e b s l +EBT e b t +EBU e b u +EBV e b v +E.B.W. e b w +EBWs e b w's +EBZR e b z r +ECAC- e c a c +ECAC e c a c +ECA e c a +ECAR e c a r +ECB, e c b +ECB e c b +ECBJ e c b j +ECBL e c b l +ECCA e c c a +ECCC e c c c +ecc e c c +ECC e c c +ECCI e c c i +Eccl e c c l +ECCRFA e c c r f a +ECCW e c c w +ECD e c d +ec e c +Ec e c +E. C. e c +E.C. e c +EC- e c +EC e c +ece e c e +Ece e c e +ECE e c e +Ecem e c e m +ECFA e c f a +ECF e c f +ECFR e c f r +ECFU e c f u +ecg e c g +ECG e c g +ECH e c h +ECHL e c h l +ECHR e c h r +ECI e c i +ECJ e c j +ECLAC e c l a c +ECLA e c l a +ECLC e c l c +ECL e c l +ECLI e c l i +ECLSS e c l s s +ecma e c m a +ECMA e c m a +ECMC e c m c +ECMDB e c m d b +E.C.M. e c m +ECM e c m +ECMI e c m i +ECML e c m l +ECMLRUS e c m l r u s +ECMWF e c m w f +ECN e c n +ECNP e c n p +ECNs e c n's +ECNS e c n s +ECNZ e c n z +ECP&DA's e c p and d a's +ECP e c p +ECPR e c p r +ECQ e c q +ECR e c r +eCRM e c r m +ECRR e c r r +ECR's e c r's +ECSA e c s a +EC's e c's +E. C. S. e c s +E.C.S. e c s +ECS e c s +ect e c t +E. C. T. e c t +ECT e c t +ECTLO e c t l o +ECTN e c t n +Ecto- e c t o +ECTS e c t s +ECV e c v +ECVs e c v's +ECWA e c w a +ECW e c w +edb e d b +EDB e d b +EDC e d c +E. D. e d +E.D. e d +EDF e d f +Edh e d h +EDH e d h +edhe e d h e +EDHF e d h f +EDIB e d i b +EDI e d i +EDK e d k +EDL e d l +EDMA e d m a +EDM- e d m +EDM e d m +EDMs e d m's +EDMW e d m w +Edmx e d m x +edn e d n +Edn e d n +EDN e d n +EDP e d p +EDQH e d q h +E.D.R. e d r +EDRP e d r p +EDSA e d s a +EDSC e d s c +eds e d s +Eds e d's +E. D. S. e d s +EDS- e d s +EDS e d s +EDSP e d s p +EDTA e d t a +EDT e d t +EDVAC e d v a c +EDWDO e d w d o +Edw e d w +EDXRD e d x r d +EEA e e a +EEAF e e a f +EEAS e e a s +EEBC e e b c +Eeb e e b +EEC e e c +EECS e e c s +E. E. e e +E.E. e e +Eef e e f +EEF e e f +Eega e e g a +EEI e e i +EEI's e e i's +EEJF e e j f +EELC e e l c +eene e e n e +EEOC e e o c +EEPCO e e p c o +eep e e p +EEP e e p +EERC e e r c +EERI e e r i +Ees e e's +EES e e s +EEST e e s t +EEUU e e u u +eeuw e e u w +Eeuw e e u w +EEZs e e z's +EFA e f a +EFAS e f a s +EFCC e f c c +efc e f c +EFC e f c +EFCJ e f c j +EFD e f d +ef e f +Ef e f +E. F. e f +E.F. e f +EF- e f +EF e f +E. F. J. e f j +E. F. K. e f k +EFL e f l +EFM e f m +EFNB e f n b +EFPA e f p a +EFP e f p +EFQUEL e f q u e l +EFR e f r +EFRU e f r u +EFSA e f s a +eFS e f s +Efs e f's +EFSET e f s e t +EFTPOS e f t p o s +EFUA e f u a +Efu e f u +EFV e f v +EFVR e f v r +EFVS e f v s +E. F. W. e f w +EFX e f x +EFY e f y +E. G. A. e g a +EGA e g a +Egba e g b a +EGBA e g b a +EGBT e g b t +EGC e g c +EGD e g d +EGF- e g f +EGF e g f +egfl e g f l +EGFL e g f l +e.g. for example +e. g. for example +EGFR e g f r +EG&G e g and g +E. G. H. e g h +Egi e g i +EGL e g l +EGM e g m +E&GR e and g r +EGR e g r +EGREM e g r e m +EGSC e g s c +EGX e g x +egy e g y +EGY e g y +Egyl e g y l +EHAAT e h a a t +EHAs e h a's +EHC e h c +EHCI e h c i +EHD e h d +E. H. e h +E.H. e h +ehf e h f +EHF e h f +E. H. M. e h m +EHP e h p +EHRC e h r c +EHR e h r +EHS e h s +EHSH e h s h +Ehttp e h t t p +EHV- e h v +EIA- e i a +EIA e i a +EIC e i c +EiCs e i c's +EIC's e i c's +E. I. e i +E.I. e i +EIES e i e s +eIF e i f +EIF e i f +EIFF e i f f +EIF's e i f's +EIT e i t +EITI e i t i +Eitr e i t r +EITs e i t's +EIU e i u +E. J. A. e j a +EJB e j b +EJBs e j b's +EJC e j c +ej e j +E. J. e j +E.J. e j +EJ e j +E. J. H. e j h +E.J.H. e j h +EJH e j h +EJHL e j h l +EJMA e j m a +E. J. M. e j m +EJML e j m l +E. J. R. e j r +E.J.S. e j s +EJSM e j s m +EJV e j v +EKB e k b +EKBO e k b o +EKCO's e k c o's +EKD e k d +E. K. e k +E.K. e k +EKF e k f +EKG e k g +eki e k i +EKIZ e k i z +EKL e k l +EKLR e k l r +Ekma e k m a +EKNU e k n u +Ekow e k o w +eks e k s +Eks e k's +EKS- e k s +EKS e k s +E. K. U. e k u +EKU e k u +E. L. C. e l c +E. L. e l +E.L. e l +ELH e l h +ELTs e l t's +ELUs e l u's +ELV e l v +EMA e m a +EMAS e m a s +Embd e m b d +EMB- e m b +EMB e m b +eMC e m c +EMC e m c +EMD e m d +EMD's e m d's +EMDs e m d's +EMDS e m d s +Emea e m e a +EMEA e m e a +E&M e and m +E. M. e m +E.M. e m +EMF e m f +Emge e m g e +EMG- e m g +EMG e m g +EML- e m l +EMLL e m l l +E.M.M. e m m +EMNLP e m n l p +EMP- e m p +EMP e m p +emra e m r a +EMR e m r +EMSC e m s c +E.M.S. e m s +EMTE e m t e +EMT e m t +EMV e m v +E&NA e and n a +Enas e n a's +ENAS e n a s +EnBW e n b w +ENC e n c +EncFS e n c f s +Encyc e n c y c +E.N.D. e n d +ENEC e n e c +E. N. e n +E.N. e n +Enes e n e's +engl e n g l +Engl e n g l +ENGM e n g m +Engr e n g r +Eni e n i +ENI e n i +enn e n n +Enn e n n +ENP e n p +ENR e n r +ens e n s +Ens e n's +ENS e n s +Enso e n s o +ENSO e n s o +ENTJs e n t j's +ENTP e n t p +ENTV e n t v +Env e n v +ENV e n v +EOC e o c +EOC's e o c's +EoD e o d +EOD e o d +E. O. e o +E.O. e o +EOFFTV e o f f t v +EOG e o g +Eois e o i's +EOKA e o k a +eok e o k +Eole e o l e +Eol e o l +EOL e o l +Eolss e o l s's +Eom e o m +EOOW e o o w +eop e o p +eoptta e o p t t a +Eora e o r a +EOR e o r +Eorl e o r l +Eors e o r's +EOS- e o s +EOSHD e o s h d +EOTC e o t c +EOTDC e o t d c +EOT- e o t +EOU e o u +E.O.W. e o w +E&P e and p +Epe e p e +E. P. E. e p e +EPEG e p e g +ep e p +Ep e p +E. P. e p +E.P. e p +EP- e p +EP, e p +EP e p +EPF e p f +EPFL e p f l +EPG e p g +EPHAR e p h a r +EPHB e p h b +Eph e p h +EPH e p h +Ephs e p h's +EPIA e p i a +epi e p i +Epi e p i +EPI e p i +E.P.J. e p j +EPL e p l +EPLF e p l f +EPM- e p m +EPM e p m +EPOC e p o c +ePO e p o +EPO e p o +epos' e p o's +epos e p o s +Epos e p o's +EPRDF e p r d f +EPR e p r +EPRI e p r i +EPRO e p r o +EPSC e p s c +E.P.s e p's +EPT e p t +EPV e p v +EPZA e p z a +EPZ e p z +EQAO e q a o +eq e q +Eq e q +EQ- e q +EQ e q +EQG e q g +EQN e q n +EQP e q p +EQS e q s +EQT e q t +E.R.A. e r a +ERAF e r a f +ERAU e r a u +ERCB e r c b +ERC e r c +ERCES e r c e s +ERCIM e r c i m +ERCS e r c s +E. R. D. e r d +E. R. e r +E.R. e r +ERF e r f +ERH e r h +ERJs e r j's +ERK e r k +ERP e r p +ERPM e r p m +ERPMRP e r p m r p +ERRC e r r c +E. R. R. e r r +E.R.R. e r r +ERSAR e r s a r +ERT e r t +Eru e r u +ERU e r u +ERUUF e r u u f +Eruv e r u v +Erv e r v +ERW e r w +ES3 e s three +Esa e s a +E.S.A. e s a +ESA- e s a +ESA e s a +Esam e s a m +Esa's e s a's +ESA's e s a's +Esat e s a t +ESB e s b +ESC e s c +ESCHL e s c h l +ESCs e s c's +Esd e s d +ESD e s d +ESDI e s d i +E.S.E. e s e +E. S. e s +E.S. e s +ESF e s f +ESFG e s f g +E.S.G. e s g +ESG e s g +ESGR e s g r +esi e s i +ESI e s i +ESL e s l +ESM e s m +ESMLC e s m l c +ESMT e s m t +ESNAGI e s n a g i +ESN e s n +ESPA e s p a +ESPC e s p c +ESPD e s p d +ESP e s p +espn e s p n +ESPN, e s p n +ESPN e s p n +ESPNhttp e s p n h t t p +ESPN's e s p n's +ESPNU e s p n u +ESPNW's e s p n w's +esq e s q +Esq e s q +Esq. esquire +Esraa e s r a a +Esra e s r a +ESRB e s r b +ESRC e s r c +Esref e s r e f +ESR- e s r +ESR e s r +ESRF e s r f +Esri e s r i +ESRI- e s r i +ESRI e s r i +ESSR e s s r +estd e s t d +Estd e s t d +EST e s t +ESTP e s t p +esu e s u +ESU e s u +ESU's e s u's +Esva e s v a +ESV e s v +ESWAT e s w a t +ESW e s w +Etad e t a d +ETA e t a +ETAF e t a f +ETA's e t a's +ETASU e t a s u +ETB e t b +ETBs e t b's +E.T.C. e t c +ETCO's e t c o's +ETCS e t c s +ETCSL e t c s l +ETD e t d +E. T. e t +E.T. e t +ETF e t f +ETFs e t f's +ETFS e t f s +Et'hem e t h e m +ETH e t h +Eti e t i +ETI e t i +ETIM e t i m +ETIM's e t i m's +etj e t j +ETJ e t j +ETL e t l +ETNZ e t n z +Eto e t o +ETO e t o +ETRM e t r m +ETRS e t r s +ETSEIB e t s e i b +ets e t s +Ets e t's +ET's e t's +ETs e t's +ETS- e t s +ETS e t s +Etsi e t s i +ETSI e t s i +ett e t t +Ett e t t +ETT e t t +Ettre e t t r e +Etts e t t's +ETU e t u +ETV e t v +Etz e t z +EUA e u a +EUB e u b +EUBs e u b's +EUBS e u b s +EUCAP e u c a p +EUCE e u c e +EUC- e u c +EUC e u c +EUCOM's e u c o m's +EUDEC e u d e c +EUD e u d +E. U. e u +E.U. e u +EU- e u +EU e u +eup e u p +EUP e u p +Eurwg e u r w g +EUSA e u s a +Euse e u s e +EU's e u's +EUS e u s +EUSEW e u s e w +Eusi e u s i +EUV e u v +EUX e u x +EVAAP e v a a p +EvaGT e v a g t +EVDC e v d c +EVD e v d +EVDO e v d o +E. V. e v +E.V. e v +EVEX e v e x +EVIT e v i t +evl e v l +EVL e v l +EVM e v m +Evna e v n a +EVN e v n +EVNS e v n s +evnt e v n t +E.V.O. e v o +E.V.O.L. e v o l +EVP e v p +EVR e v r +EVs e v's +EVS- e v s +EVS e v s +evv e v v +EVV e v v +Ewa's e w a's +EWD e w d +EWEB e w e b +E. W. e w +E.W. e w +EWF- e w f +EWF e w f +EWF's e w f's +Ewha e w h a +EWH e w h +E.W.K. e w k +EWLA e w l a +EWL e w l +EWO e w o +EWOs e w o's +EWP e w p +EW's e w's +EWS- e w s +EWS e w s +EWTN e w t n +Ewu e w u +EWU e w u +excl e x c l +exd e x d +E.X. e x +exhb e x h b +exh e x h +EXI e x i +EXIF e x i f +EXIN e x i n +EXL e x l +EXLM e x l m +E. X. M. e x m +Exptl e x p t l +EXR e x r +exsul e x s u l +EXT e x t +EYA e y a +EYCN e y c n +EYD e y d +E&Y e and y +E. Y. e y +EY e y +Eyk e y k +eyu e y u +Ezaa e z a a +ez e z +Ez e z +EZ- e z +EZ e z +Ezh e z h +EZH e z h +EZO e z o +EZR e z r +EZTV e z t v +Ezu e z u +EZW e z w +Ezy e z y +Faaa f a a a +faa f a a +Faa f a a +FAA- f a a +FAA f a a +faama f a a m a +FAANP f a a n p +FabH f a b h +FABM f a b m +FabR f a b r +F.A.B.'s f a b's +F.A.C.E. f a c e +FACR f a c r +F.A.C.S. f a c s +FADD f a d d +Fadl f a d l +FADS f a d s +FAEM f a e m +fa f a +Fa f a +F. A. f a +F.A. f a +FA- f a +FA f a +Faf f a f +FAFG f a f g +Fafhrd f a f h r d +FAFL f a f l +F. A. G. f a g +F.A.H. f a h +fai f a i +Fai f a i +FAI- f a i +FAI f a i +FAIM f a i m +FAIP f a i p +FAIPs f a i p's +Faiq f a i q +fajn f a j n +FAK f a k +Fal f a l +FAL f a l +F.A.M.E. f a m e +F.A.M. f a m +FAQ f a q +FARC's f a r c's +FARDC f a r d c +F.A.R. f a r +farw f a r w +FASA f a s a +FASB f a s b +fasc f a s c +Fasc f a s c +fas f a s +Fas f a's +FAs f a's +F.A.S. f a s +FAS f a s +FASG f a s g +FasL f a s l +FASRE f a s r e +FATF f a t f +FAUP f a u p +FAVC f a v c +FBA f b a +FBD f b d +FBFA f b f a +F. B. f b +F.B. f b +FB- f b +FB f b +F.B.G.s f b g's +F.B.I. f b i +FBI- f b i +FBI f b i +FBI's f b i's +FBK f b k +FBLA- f b l a +FBL f b l +FBMW f b m w +FBN f b n +FBO f b o +FBOs f b o's +FBPA f b p a +FBR f b r +FBS f b s +FBSN f b s n +FBSU f b s u +FBT f b t +FCA f c a +fcb f c b +FCB f c b +FCBF f c b f +F.C.B.I. f c b i +FCBL f c b l +FCBs f c b's +F.C.C. f c c +FCC f c c +FCD f c d +F&C f and c +fc f c +F. C. f c +F.C. f c +FCF f c f +FCGR f c g r +F. C. H. f c h +F.C.H.L. f c h l +FCIAC f c i a c +FCIC f c i c +FCI f c i +FCIP f c i p +FCK f c k +FCL f c l +FCN- f c n +FCNM f c n m +FCNZ f c n z +FCO f c o +FCP f c p +F.C.P.R.E.M.I.X. f c p r e m i x +FCPS f c p s +FCRLA f c r l a +F.C.'s f c's +FC's f c's +FCS f c s +FCT f c t +FCU f c u +FCUL f c u l +FCV- f c v +FCV f c v +F. C. W. f c w +FCW f c w +FCW's f c w's +FDA f d a +FDCC f d c c +FD&C f d and c +FDCPA f d c p a +FDCPAs f d c p a's +FDD f d d +FDDI f d d i +FDE f d e +F. D. f d +F.D. f d +FDFL f d f l +FDGB- f d g b +FDGB f d g b +FDG f d g +FDICs f d i c's +FdI f d i +FDI f d i +FDIO f d i o +FDJ f d j +FDL f d l +FDLT f d l t +FDM f d m +FDMNES f d m n e s +FDNY f d n y +FDP f d p +FDP's f d p's +FDQ, f d q +FDR f d r +FDU f d u +F. E. f e +F.E. f e +FEFK f e f k +FEFSI f e f s i +FEFY f e f y +FEG f e g +FEGs f e g's +FEP f e p +F.E.R.S. f e r s +Fes f e's +FES f e s +FFA f f a +FFAK f f a k +ffc f f c +FFC f f c +FFC's f f c's +FFDO f f d o +ff f f +F. F. f f +F.F. f f +FF- f f +FF f f +FFF f f f +FFG f f g +FFH f f h +FFI f f i +FFK- f f k +FFK f f k +FFL f f l +FFL's f f l's +FFME f f m e +ffm f f m +FFOM f f o m +FFP f f p +FFPIR f f p i r +FFR f f r +FFSA f f s a +FF's f f's +FFs f f's +FFS f f s +FFU f f u +FFV f f v +FFWD f f w d +FGCU f g c u +FGFA f g f a +FGF f g f +F. G. f g +F.G. f g +FGFR f g f r +FGID f g i d +fgk f g k +FGL f g l +FGM f g m +FGR f g r +FGs f g's +FGS f g s +FGU f g u +FGV f g v +F. H. A. f h a +FHA- f h a +FHA f h a +FHAR f h a r +F. H. C. f h c +FHFA f h f a +F. H. f h +F.H. f h +FHI's f h i's +FHL f h l +fhm f h m +FHM f h m +FHM's f h m's +FHQ f h q +FHSAA f h s a a +FHS f h s +FHWA f h w a +F. H. W. f h w +FIAC f i a c +fiadh f i a d h +fia f i a +FIA f i a +FIAPO f i a p o +FIA's f i a's +FIAS f i a s +FIBCO f i b c o +F. I. B. f i b +FiBL f i b l +FIBT f i b t +FIC f i c +Fi'd f i d +FIDH f i d h +F.I f +F&I f and i +F. I. f i +F.I. f i +FIFTA f i f t a +FIGC f i g c +Figl f i g l +FIH f i h +FIH's f i h's +fija f i j a +FIJL f i j l +FIK f i k +F.I.M. f i m +FIM f i m +FIMI- f i m i +FIMI f i m i +Fio f i o +FiO f i o +FIO f i o +F.I.R. f i r +F.I.T.T. f i t t +FIU f i u +FIVB f i v b +F. J. A. f j a +FJE f j e +F. J. f j +F.J. f j +F.J.G. f j g +FJK f j k +F. J. M. f j m +FKA f k a +FKBP f k b p +F. K. f k +F.K. f k +FKK f k k +FKL's f k l's +FKS f k s +FKY f k y +Fla f l a +FLB f l b +FLCL f l c l +FLCS f l c s +FLD f l d +FLDS f l d s +FLEC f l e c +FLE f l e +Fles f l e's +FLEST f l e s t +FLETC f l e t c +fl f l +F. L. f l +F.L. f l +FL- f l +FL f l +FLH f l h +fli f l i +FLI f l i +FLI's f l i's +FLL f l l +F.L.M. f l m +FLM f l m +FLNB f l n b +FLN f l n +FLP f l p +FLPP f l p p +FLQ f l q +FLRA f l r a +FLR f l r +FLRL f l r l +FLS f l s +FLSW f l s w +FLT f l t +FLTK f l t k +FLTR f l t r +FLTS f l t s +FLV f l v +FLW f l w +flyr f l y r +FMA f m a +FMARS f m a r s +FMC f m c +FMCG f m c g +FMCSA f m c s a +FMCS f m c s +FMF f m f +fm f m +F. M. f m +F.M. f m +FM- f m +FM f m +FMH f m h +FMHS f m h s +FMI f m i +FMJD f m j d +FMLN f m l n +FMM f m m +FMN f m n +FMNJ f m n j +FMNP f m n p +FMPA's f m p a's +FMP f m p +FMQB f m q b +FMR f m r +fMRI f m r i +FMRP f m r p +FMRU f m r u +FM's f m's +FMs f m's +FMS- f m s +FMS f m s +FMSF's f m s f's +FMSO f m s o +FMV f m v +FMW f m w +FMX f m x +FMY f m y +FNAC f n a c +F. N. A. f n a +FNA f n a +FNB f n b +FNC f n c +F. N. D. f n d +FNE f n e +F. N. f n +F.N. f n +FNIC f n i c +FNI f n i +FNLA f n l a +FNLC f n l c +FNLC's f n l c's +FNL f n l +FNO f n o +FNP f n p +FNPI f n p i +FNSEA f n s e a +FNU f n u d s a +FOF f o f +F. O. f o +FPA- f p a +FPA f p a +FPB f p b +FPC f p c +FPCP f p c p +F.P.D. f p d +FPD- f p d +FPDs f p d's +fp f p +F. P. f p +F.P. f p +FP- f p +FP f p +FPGA f p g a +FPGAs f p g a's +FPI f p i +FPJ f p j +FPK f p k +F. P. L. f p l +FPL f p l +FPMR f p m r +FPM&SA f p m and s a +FPO f p o +FPP f p p +FPPS f p p s +F.P.R. f p r +FPR f p r +FPSC f p s c +fps f p s +FPS- f p s +FPS f p s +FPSO f p s o +FPSs f p s's +FPTP- f p t p +FPTP f p t p +FPU f p u +FQAD f q a d +F.Q. f q +FRA f r a +FRA's f r a's +FRAXA f r a x a +FRB f r b +FRBNF f r b n f +FRCC f r c c +F. R. C. f r c +FRC f r c +FRCNA f r c n a +FRCP f r c p +FRCS f r c s +FRDF f r d f +F.R.D. f r d +fr f r +F. R. f r +F.R. f r +FR- f r +FR f r +FRFS f r f s +FRFT f r f t +FRG f r g +F.R.G.S. f r g s +F.R.H.S. f r h s +F.R.I.B.A. f r i b a +FRMN f r m n +FRMR f r m r +FRMS f r m s +F. R. P. f r p +FRP f r p +FRR f r r +F.R.S.A. f r s a +FRSA f r s a +FRSC f r s c +FRSE f r s e +frs f r s +F. R. S. f r s +F.R.S. f r s +FRS f r s +FRSL f r s l +Fru f r u +FRU f r u +FRV f r v +FRWT f r w t +Frwydr f r w y d r +FSA f s a +FSB f s b +FSC- f s c +FSC f s c +FSCS f s c s +FSE f s e +fsf f s f +FSF f s f +FSFH f s f h +FSFLA f s f l a +fs f s +F. s f's +Fs f's +F. S. f s +F.S. f s +FS f s +FSFS f s f s +FSG f s g +FSGS f s g s +FSH f s h +FSIC f s i c +FSI f s i +FSIL f s i l +FSK f s k +FSKN f s k n +FSLF f s l f +FSL f s l +FSLN f s l n +FSM f s m +FSM's f s m's +fsn f s n +FSN f s n +FSO f s o +FSP f s p +FSPS f s p s +FSR f s r +FSRU f s r u +FSSB f s s b +FSS f s s +FSSM f s s m +FSSRU f s s r u +FSSs f s s's +FST f s t +FSU f s u +FSU's f s u's +FSV f s v +FTA f t a +FTAs f t a's +F.T.B. f t b +FTC f t c +FTD f t d +FTDI f t d i +FTE f t e +FTF f t f +Ft. Fort +F. T. f t +F.T. f t +FTI f t i +FTII f t i i +FTJ f t j +FTK f t k +FTL f t l +FTM f t m +FTN f t n +FTO f t o +ftp f t p +FTP f t p +FTPI f t p i +FTPL f t p l +FTRA f t r a +FTR f t r +FtsA f t s a +FT's f t's +FTS f t s +FTT f t t +FTTH f t t h +FTUC f t u c +Ftuh f t u h +FTU's f t u's +FTV f t v +FVA- f v a +F. V. B. f v b +F. V. f v +F.V. f v +FVGM f v g m +FVHS f v h s +FVJC f v j c +FVL f v l +FVP f v p +FVW f v w +FVWM f v w m +FWAA f w a a +FWA f w a +FWCD f w c d +FWD f w d +FWE f w e +F&W f and w +F. W. f w +F.W. f w +FWH f w h +FWI f w i +FWMW f w m w +FWO f w o +fwr f w r +F.W.R. f w r +F. W. S. f w s +F.W.S. f w s +FWS f w s +FWX f w x +FXCM's f x c m's +FXE f x e +FXFL f x f l +fx f x +FX f x +FXI f x i +F.X.R. f x r +FX's f x's +FXS f x s +FXX f x x +FXX's f x x's +FXXT f x x t +Fyb f y b +FYCO f y c o +F.Y. f y +F.Y.I. f y i +FYI f y i +Fyn f y n +FYN f y n +FYP f y p +fyr f y r +F.Y.R. f y r +FYR f y r +FYT f y t +FYU f y u +Fyw f y w +FZCI f z c i +FZD f z d +F.Z. f z +FZG f z g +FZJD f z j d +FZKA f z k a +Gaac g a a c +GAAF g a a f +G.A.A. g a a +G. A. g a +G.A. g a +G.A.R. g a r +GASL g a s l +GATA g a t a +GATV g a t v +GATX g a t x +G. A. W. g a w +GBA g b a +GBAG g b a g +GBASE g b a s e +GBA's g b a's +GBB g b b +GBBI g b b i +GBCA g b c a +GBC g b c +GBCM g b c m +GBD g b d +Gbe g b e +GbE g b e +GBE g b e +G. B. g b +G.B. g b +GBGC g b g c +GBGM g b g m +GBH g b h +GBIB g b i b +GBIF g b i f +GBI g b i +GBK g b k +GBLA- g b l a +GBL g b l +GBMA g b m a +G. B. M. g b m +GB&NDR g b and n d r +GBP g b p +GBR g b r +GB's g b's +GBS g b s +GBV- g b v +GBV g b v +GBX g b x +GBYLA g b y l a +GCAC g c a c +GCAF g c a f +GCA g c a +G. C. B. g c b +GCB g c b +GCC g c c +GCCN g c c n +GCC's g c c's +GCCS g c c s +GCD g c d +G.C.E. g c e +GCEP g c e p +GCF g c f +G&C g and c +G. C. g c +G.C. g c +GCHQ g c h q +GCI g c i +GCIP g c i p +GCIT's g c i t's +GCIV g c i v +GCJ g c j +GCL g c l +GCMC g c m c +G.C.M. g c m +GCM g c m +GCMG g c m g +GCN g c n +GCNT g c n t +GCOEC g c o e c +GCOS g c o s +GCPD g c p d +GCPH g c p h +GCR- g c r +GCR g c r +GCRIO g c r i o +GCSE g c s e +GCSEs g c s e's +GC&SF g c and s f +GCS g c s +GCSI g c s i +GCT- g c t +GCU g c u +GCV g c v +GCVO g c v o +GCVS g c v s +GDA g d a +Gdal g d a l +GDAL g d a l +GDAP g d a p +GDB g d b +GDC g d c +GDD g d d +GDDM g d d m +GDDR g d d r +Gde g d e +G. D. F. g d f +GDF g d f +gd g d +G. D. g d +G.D. g d +GD g d +GDI g d i +G.D.L. g d l +GDNFR- g d n f r +GDN g d n +G.D.O.D. g d o d +Gdow g d o w +G.D.P. g d p +GDP- g d p +GDP g d p +GDPs g d p's +gdr g d r +GDR g d r +GDR's g d r's +GDS g d s +GDST g d s t +GDT g d t +GDV g d v +Gebr g e b r +Gebt g e b t +GEC g e c +Ged g e d +GED g e d +GEF g e f +GEGB g e g b +ge g e +G. E. g e +G.E. g e +GE- g e +GE g e +GEHA g e h a +GEHX g e h x +Gek g e k +G.E.M. g e m +Gen. general +G.E.O. g e o +G. E. P. g e p +geq g e q +Geu g e u +GEU g e u +Gev g e v +GeV g e v +GEV g e v +Gfa g f a +GFA g f a +GF&A g f and a +GFC g f c +GFD g f d +GFDL g f d l +GFE g f e +GFF g f f +G. F. g f +G.F. g f +GFM g f m +GFP- g f p +GFP g f p +G. F. R. g f r +GFR g f r +GFS g f s +GFSN g f s n +GFSR g f s r +GFT g f t +GFWC g f w c +GFW g f w +GFY g f y +G. G. g g +G.G. g g +G. G. J. g g j +GGL g g l +GGM g g m +GGMRC g g m r c +GGS g g s +G.H.A. g h a +GHB g h b +GHC g h c +GHD g h d +Ghe g h e +G. H. g h +G.H. g h +GHGs g h g's +Ghir g h i r +GHK g h k +GHMC g h m c +G.H.N. g h n +G.H.Q. g h q +GHQ g h q +ghra g h r a +GHR g h r +GH&SA g h and s a +GHSA g h s a +GHSFHA g h s f h a +GHS g h s +GHSR g h s r +GHSU g h s u +GHU g h u +G.H.W. g h w +GICL g i c l +Giei g i e i +GIFV g i f v +G. I. g i +G.I. g i +G.J.B. g j b +G. J. g j +G.J. g j +G. K. g k +G.K. g k +Gla g l a +GLA g l a +GLB g l b +GLBT g l b t +GLBTQ g l b t q +GLC- g l c +GLC g l c +G&L g and l +G. L. g l +G.L. g l +gli g l i +Gli g l i +GLI g l i +Glis g l i's +GLK g l k +GLL g l l +GLLI g l l i +GLM g l m +glnA g l n a +G.L.O.R.Y. g l o r y +Glos g l o's +GLP- g l p +GLRS g l r s +GLS g l s +GLT g l t +GLVC g l v c +GLX- g l x +GLX g l x +GLYX- g l y x +GMAC g m a c +gma g m a +GMA- g m a +GMA g m a +GMAM g m a m +GMA's g m a's +GMAT g m a t +GMB g m b +GmbH g m b h +GMBH g m b h +GMB's g m b's +GMCC g m c c +GMC g m c +GMCL g m c l +GMC's g m c's +GMDSS g m d s s +GME g m e +Gmel g m e l +GMES g m e s +GMFF g m f f +GMF g m f +GMG g m g +gm g m +G. M. g m +G.M. g m +GM- g m +GM g m +GMG's g m g's +gmin g m i n +GMM g m m +GMO g m o +GMOs g m o's +GMPC g m p c +GMPE g m p e +GMP g m p +GMPTE g m p t e +GMR g m r +GMRL g m r l +GM's g m's +GMs g m's +GMS g m s +gmt g m t +GMT g m t +GMTV g m t v +Gmul g m u l +GNAA g n a a +GNA g n a +gnb g n b +GNB g n b +GNCC g n c c +G. N. C. g n c +GNC g n c +G. N. g n +G.N. g n +GNL g n l +GNP&BR g n p and b r +GNPDA g n p d a +gnp g n p +GNP g n p +GNPs g n p's +GNRD g n r d +GNR g n r +GNR's g n r's +GNS g n s +GNSS g n s s +GNT g n t +G.O. g o +G.O.O.D. g o o d +GOPB g o p b +G.O.P. g o p +GOP- g o p +GOP g o p +GOPIO g o p i o +GOP's g o p's +GOQ g o q +GPA- g p a +GPA g p a +GPB g p b +G.P.C. g p c +GPC g p c +GPCR g p c r +GPCRs g p c r's +GPD g p d +GPdI g p d i +GPDSC g p d s c +GPF g p f +GPFS g p f s +G&P g and p +GPG g p g +gp g p +G. P. g p +G.P. g p +GP- g p +GP g p +GPGPU g p g p u +GPI g p i +GpIIb g p i i b +GPIO g p i o +GPL g p l +GPMGs g p m g's +G.P.O. g p o +GPO g p o +GPOPS- g p o p s +GPP g p p +GPRA g p r a +GPRF g p r f +GPR g p r +GPRs g p r's +GPRS g p r s +GPSC g p s c +GPS- g p s +GPS g p s +GPT g p t +GPU g p u +GPUs g p u's +GPX g p x +G. Q. g q +GQ g q +GQNYGK g q n y g k +GQ's g q's +GQT's g q t's +G.R.A.B. g r a b +grac g r a c +gra g r a +Gra g r a +GRA g r a +Grbac g r b a c +GRB g r b +GRBs g r b's +grc g r c +GRC g r c +GRCs g r c's +GRDA g r d a +GRDC g r d c +GRDF- g r d f +G.R.D. g r d +GRD g r d +gre g r e +Gre g r e +GRE g r e +GRF g r f +GRG g r g +gr g r +G. R. g r +G.R. g r +GR- g r +GR g r +GRK g r k +G.R.L. g r l +GRL g r l +GRMG g r m g +GRM- g r m +GRM g r m +GRO g r o +GRP g r p +grrl g r r l +G. R. S. g r s +GRS g r s +GRT g r t +GRU g r u +GRX g r x +Gryf g r y f +gry g r y +Gry g r y +GRY g r y +grz g r z +GSA g s a +GSAPP g s a p p +GSAS g s a s +GSB g s b +GSC g s c +GSDF g s d f +GSD g s d +G.S.E. g s e +GSE g s e +GSFC g s f c +GSF g s f +G&S g and s +GSG- g s g +GSG g s g +Gs g's +G. S. g s +G.S. g s +GS- g s +GS g s +GSHD g s h d +GSh- g s h +GSH g s h +GSI g s i +GSIs g s i's +GSIS g s i s +GSK- g s k +GSK g s k +GSL g s l +GSLV g s l v +GSMA g s m a +GSM- g s m +GSM g s m +GSNAS g s n a s +G. S. N. g s n +GSN g s n +GSNM g s n m +GSN's g s n's +GSOC- g s o c +GSO g s o +GSOM's g s o m's +GSOp g s o p +GSPC g s p c +G. S. P. g s p +GSP g s p +GSR g s r +GSSE g s s e +GSS- g s s +GSS g s s +GSSP g s s p +GSTA g s t a +G.S.T. g s t +GST- g s t +GST g s t +GSU g s u +GSUSA g s u s a +GSVD g s v d +GSVS g s v s +GSX g s x +GTAA g t a a +GTA g t a +GTAs g t a's +GTB g t b +GTBK g t b k +GTBP g t b p +GTC g t c +GTCR g t c r +GTD g t d +GTEC g t e c +GTE g t e +GTF g t f +GTFT g t f t +GTFU g t f u +gt g t +G. T. g t +G.T. g t +GT- g t +GT g t +G. T. H. g t h +GTi g t i +GTI g t i +GTK g t k +GTL g t l +GTMO g t m o +G. T. O. g t o +GTO g t o +GTPBP g t p b p +GTP- g t p +GTP g t p +GTPR g t p r +GTRA g t r a +GTR g t r +GTRI g t r i +GTs g t's +GTS g t s +GTSV g t s v +GTT g t t +GTTP g t t p +GTU g t u +GTV- g t v +GTV g t v +GTW- g t w +GTW g t w +GTW's g t w's +GTX g t x +G. U. g u +G.U. g u +GUI's g u i's +GUIs g u i's +G.U.Y. g u y +GVA g v a +GVAV g v a v +GVBA g v b a +GVC g v c +G. V. g v +G.V. g v +GVHSS g v h s s +GVJZ g v j z +GVK g v k +GvpA g v p a +GVP- g v p +GVP g v p +GVRD g v r d +GVSU g v s u +GVT g v t +GVU g v u +GVVV g v v v +GVW g v w +GVWR g v w r +GWA g w a +GWB g w b +G. W. C. g w c +GWe g w e +Gwet g w e t +G.W.F. g w f +GWF g w f +gw g w +G. W. g w +G.W. g w +GW g w +G. W. H. g w h +GWHS g w h s +Gwi g w i +GWK g w k +G.W.L. g w l +GWO g w o +GWOT g w o t +gwr g w r +GWR g w r +GWR's g w r's +GW's g w's +G. W. S. g w s +GWS g w s +GWU g w u +GXG g x g +GX- g x +GX g x +GXP g x p +GYAN's g y a n's +GYC g y c +GYF g y f +Gy g y +G. Y. g y +GY- g y +GY g y +Gyn g y n +GYN g y n +Gyps g y p's +gyu g y u +Gyu g y u +GZA g z a +GZ g z +GZR g z r +GZSZ- g z s z +HAFV h a f v +H. A. h a +H.A. h a +H&A h and a +HBA h b a +HBA's h b a's +HBBO h b b o +HBC h b c +HBCo h b c o +HBC's h b c's +HBCU h b c u +HBES h b e s +HBF h b f +HBGAs h b g a's +H. B. h b +H.B. h b +HB h b +HBI h b i +HBK- h b k +HBK h b k +HBLX h b l x +HBMP h b m p +HBO h b o +HBO's h b o's +HBPU h b p u +H&BR h and b r +HBR h b r +HBRN h b r n +HBRY h b r y +HB's h b's +HBS h b s +HBSS h b s s +HBU's h b u's +HBV h b v +HBW h b w +HCA h c a +HCAW h c a w +HCBE h c b e +HCB h c b +HCBS h c b s +HCCA h c c a +HCCB h c c b +HCCC h c c c +HCC h c c +HCDCH h c d c h +HCD- h c d +HCE h c e +HCEs h c e's +HCF h c f +H. C. G. h c g +HCGVN h c g v n +H. C. h c +H.C. h c +HC h c +HCH h c h +HCHS h c h s +HCI h c i +HCJB h c j b +hcl h c l +HCL h c l +HClO h c l o +HCMC's h c m c's +HCME h c m e +HCM h c m +HCN h c n +HCO h c o +HCoV h c o v +HCPA h c p a +H. C. P. h c p +HCP h c p +HCPs h c p's +HCR h c r +HCS h c s +HCT h c t +HCTV's h c t v's +HCUP h c u p +HCV h c v +HCVHDB- h c v h d b +Hcy h c y +HDAC h d a c +Hdad h d a d +hDAF h d a f +HDA h d a +HDB h d b +HDBK- h d b k +HDCAM h d c a m +HDC h d c +HDCP h d c p +HDD h d d +HDDs h d d's +HDFC h d f c +H.D.F. h d f +HDFM h d f m +H.D.G. h d g +HDG h d g +hd h d +H. D. h d +H.D. h d +HD- h d +HD h d +hdh h d h +HDi h d i +HDI h d i +HDLC h d l c +HDL- h d l +HDL h d l +HDLS h d l s +HDM h d m +HDMI h d m i +HDPE h d p e +HDP h d p +HDR h d r +HDRO h d r o +HDRR h d r r +H. E. h e +H.E. h e +hezb h e z b +Hezb h e z b +HFA h f a +HFB h f b +HFC h f c +HFCS h f c s +HFDF h f d f +HFF h f f +HFF's h f f's +H&F h and f +H. F. h f +H.F. h f +HFI h f i +H. F. J. h f j +HFNB h f n b +HFPA h f p a +HFP h f p +HFPPV h f p p v +HFRS h f r s +HFS h f s +HFTP h f t p +H.G.B. h g b +hgcA h g c a +HGC h g c +HGF h g f +hg h g +H. G. h g +H.G. h g +HG- h g +HG h g +HGH h g h +HGNC h g n c +HgO h g o +H. G. O. h g o +HGPIN h g p i n +HGPs h g p's +HGTV h g t v +HGTV's h g t v's +HgU h g u +HGVs h g v's +HGY h g y +HHA h h a +H. H. B. h h b +HHC h h c +HHD h h d +HHDP h h d p +HHG h h g +H&H h and h +HHHDB h h h d b +H. H. h h +H.H. h h +HHPC h h p c +HHSAA h h s a a +HHS h h s +H.H.T. h h t +H. I. h i +H.I. h i +H.I.V.E. h i v e +HIV- h i v +HIV h i v +HIY h i y +Hizb h i z b +hja h j a +HJC h j c +H. J. h j +H.J. h j +HJ- h j +HJ h j +HJK h j k +H.J.M. h j m +HJN h j n +H. J. R. h j r +H.J.R. h j r +HJT- h j t +H. J. W. h j w +HJW h j w +HKAA h k a a +Hkam h k a m +HKB h k b +HKCEC h k c e c +HKCEE h k c e e +HKDF h k d f +HKD h k d +HKDSE h k d s e +HKFC h k f c +H&K h and k +H. K. h k +H.K. h k +HK- h k +HK h k +H. K. J. h k j +HKK h k k +HKMA h k m a +HKN h k n +H.K.P. h k p +HKP h k p +HKSAR h k s a r +HKSCS- h k s c s +HK's h k's +HKTB- h k t b +HKT h k t +HKTWU h k t w u +HKU h k u +Hla h l a +HLA- h l a +HLA h l a +HLBN h l b n +HLD- h l d +HLE h l e +HLH h l h +H. L. h l +H.L. h l +HLL h l l +HLLS h l l s +HLN h l n +HLN's h l n's +HLS h l s +HLT h l t +HLTV h l t v +Hluk h l u k +hlutr h l u t r +HLV h l v +Hly h l y +HMAC h m a c +HMA- h m a +HMA h m a +Hman h m a n +H.M.A.S. h m a s +HMAS h m a s +HMB h m b +HMC h m c +HMCS h m c s +HMDA h m d a +HMD h m d +HMDI h m d i +HMDP h m d p +HME h m e +HMEL h m e l +HMF h m f +HMGA h m g a +HMG- h m g +HMG h m g +HMGIC h m g i c +HMGN h m g n +H&M h and m +HMH- h m h +HMH h m h +hm h m +H. M. h m +H.M. h m +HM- h m +HM h m +HMI h m i +HMIS h m i s +HMMC h m m c +HMM- h m m +HMM h m m +HMMWV h m m w v +HMN h m n +HMNZS h m n z s +HMNZT h m n z t +HMO h m o +HMOs h m o's +HMPC h m p c +HMP h m p +HMRC h m r c +hMRE h m r e +H&M's h and m's +H. M. S. h m s +H.M.S. h m s +HMS h m s +HMSN h m s n +H.M.S.O. h m s o +HMSO, h m s o +HMSO h m s o +HNF- h n f +H. N. h n +H.N. h n +HNI h n i +HNK h n k +HNL- h n l +HNL h n l +HNLMS h n l m s +HNMT h n m t +HNN h n n +HNO h n o +HNS h n s +H. O. h o +H.O. h o +HPA- h p a +HPA h p a +H.P.B. h p b +hpc h p c +HPC h p c +HPCT h p c t +HPDD- h p d d +H. P. D. h p d +HPD h p d +HPE h p e +HPH h p h +hp h p +H. P. h p +H.P. h p +HP- h p +HP h p +HPI h p i +HPLC h p l c +HPL h p l +HPP h p p +HPRA h p r a +HPR h p r +HPRT h p r t +HP's h p's +HPS h p s +H. P. T. h p t +HPV h p v +HQAA h q a a +HQDA h q d a +H.Q. h q +HQ- h q +HQ h q +HRG h r g +H&R h and r +H.R.H. h r h +HRH h r h +H. R. h r +H.R. h r +HR- h r +HR h r +HRK h r k +HRL h r l +HRM h r m +hroa h r o a +HRo h r o +HRO h r o +HRP h r p +HRSA h r s a +HRSA's h r s a's +HRSG h r s g +hrs h r s +HRs h r's +HRST h r s t +HRT h r t +HRV h r v +HRW h r w +HRX h r x +HSA h s a +HSAN h s a n +HSBC h s b c +HSB h s b +HSCA h s c a +HSC h s c +hscy h s c y +HSD h s d +HSDPA h s d p a +Hsee h s e e +HSE h s e +HSE's h s e's +HSF h s f +HSG h s g +HSHB h s h b +H.S.H. h s h +hs h s +Hs h's +H. S. h s +H.S. h s +HS- h s +HS h s +HSHS h s h s +hsi h s i +HSJ h s j +HSK h s k +HSL- h s l +HSL h s l +hSlo h s l o +H.S.M. h s m +HSM h s m +HSMR h s m r +HSN h s n +HSNY's h s n y's +HSPA h s p a +HSPA's h s p a's +HSPG h s p g +HSPH h s p h +HSP h s p +HSP's h s p's +HSPs h s p's +HSRC h s r c +HSR h s r +HSSCCT h s s c c t +HSSF h s s f +HSS h s s +HST h s t +HST's h s t's +HSTs h s t's +HSTS h s t s +hsv h s v +HSV- h s v +HSV h s v +HSX h s x +HSY h s y +HSYK's h s y k's +HSZ h s z +HTA- h t a +HTA h t a +hTAS h t a s +HTAs h t a's +htc h t c +HTC- h t c +HTC h t c +HTC's h t c's +HTGL h t g l +HTHC h t h c +ht h t +H. T. h t +H.T. h t +HT- h t +HT h t +HTIB h t i b +hti h t i +Htin h t i n +HTK h t k +HTLA h t l a +HTLV- h t l v +HTLV h t l v +HTML h t m l +HTM's h t m's +Htoo h t o o +HTPB h t p b +HTPC h t p c +HTP h t p +HtrA h t r a +HTR h t r +HTRs h t r's +HTS h t s +http h t t p +HTTP- h t t p +HTTP h t t p +https h t t p s +HTTPS h t t p s +HTV h t v +HTV's h t v's +Htwa h t w a +HTW h t w +HTX h t x +H. U. h u +HVC h v c +HVCs h v c's +hvcv h v c v +HVDC- h v d c +HVDC h v d c +hvem h v e m +Hvem h v e m +HVF h v f +HVHS h v h s +H. V. h v +H.V. h v +H.W.A. h w a +H. W. F. h w f +H. W. h w +H.W. h w +H.W.L. h w l +HWL h w l +HWM h w m +HWRT h w r t +HWT h w t +HWV h w v +hwy h w y +Hwy h w y +hxt h x t +H. Y. h y +H.Y. h y +H. Z. h z +HZ&PC h z and p c +iaaf i a a f +IAAF i a a f +IAAF's i a a f's +IAA i a a +IAAPA i a a p a +IAAP i a a p +IaaS i a a s +IAAS i a a s +IABSE i a b s e +Iaca i a c a +IACC i a c c +IACD i a c d +IAC i a c +Iacob i a c o b +IACR i a c r +IAC's i a c's +IACTs i a c t's +Iacub i a c u b +IACUC i a c u c +IACW i a c w +IAD i a d +IAD's i a d's +IAEA i a e a +IAEG i a e g +I.A.E. i a e +IAE i a e +IAF i a f +IAF's i a f's +IAG i a g +IAG's i a g's +iagt i a g t +IAHP i a h p +IAHRC i a h r c +I. A. i a +I.A. i a +I.A.M. i a m +IAO i a o +Iapa i a p a +IAPHC i a p h c +IAP- i a p +IAP i a p +iar i a r +Iar i a r +IAR- i a r +IAR i a r +IAST i a s t +IATA i a t a +I.A.T. i a t +IAT i a t +IATR i a t r +IAUC i a u c +IAUCs i a u c's +IAU i a u +IAV i a v +IAWTV i a w t v +Iax i a x +IAX i a x +Ibac i b a c +Ibaes i b a e's +IBA's i b a's +IBB i b b +IBC i b c +IBCs i b c's +IBCT i b c t +IBDA'A i b d a a +Ibda i b d a +ibdal i b d a l +IBDB i b d b +IBDE i b d e +IBD i b d +IBE i b e +IBF i b f +IBGE i b g e +IBH i b h +I. B. i b +I.B. i b +IBJJF i b j j f +IBL i b l +IBLT i b l t +IBMA's i b m a's +IBMC i b m c +Ibm i b m +I.B.M. i b m +IBM- i b m +IBM i b m +IBM's i b m's +IBMXF i b m x f +Ibne i b n e +ibn i b n +IBO i b o +IBPA i b p a +IBP i b p +Ibra i b r a +IBRA i b r a +IBRC i b r c +IBR i b r +IBRO i b r o +Ibsa i b s a +IBSA i b s a +IBSF i b s f +Ibs i b's +IBS i b s +IBSL i b s l +IBTC i b t c +iBT i b t +IBTS i b t s +Ibu i b u +IBU i b u +IBV i b v +IBWFF i b w f f +iby i b y +ICBA i c b a +ICBC i c b c +ICB i c b +ICBL i c b l +ICBM i c b m +ICBMs i c b m's +ICBP i c b p +ICBS i c b s +ICBTT i c b t t +ICCC i c c c +ICCC's i c c c's +ICCEC i c c e c +ICCF i c c f +Icche i c c h e +I.C.C. i c c +ICC- i c c +ICC i c c +ICCID i c c i d +ICCI i c c i +ICCR i c c r +ICC's i c c's +ICCS i c c s +ICCW i c c w +ICDC i c d c +ICD- i c d +ICD i c d +ICDS- i c d s +ICDS i c d s +I.C.E. i c e +ICFF i c f f +icf i c f +ICF- i c f +ICF i c f +ICFR i c f r +ICG i c g +ICHC i c h c +I&C i and c +ic i c +Ic i c +I. C. i c +I.C. i c +IC i c +Ici i c i +ICI i c i +ICJ i c j +ICL i c l +ICL's i c l's +ICMA i c m a +ICME i c m e +I.C.M. i c m +ICM i c m +iCN i c n +ICN i c n +Ico i c o +ICO i c o +ICOP i c o p +ICPC i c p c +ICPD i c p d +ICP i c p +ICP's i c p's +ICQ i c q +ICRC i c r c +ICRF i c r f +ICR i c r +IcRn i c r n +ICRP i c r p +ICRS i c r s +ICRW i c r w +ICSA i c s a +ICSB i c s b +ICSC i c s c +ICSD- i c s d +ICSE i c s e +ICSG i c s g +ics i c s +ICs i c's +ICS i c s +ICSS i c s s +ICSSR i c s s r +ICTA i c t a +ICT i c t +ICTM i c t m +ICTP i c t p +ICTSF i c t s f +ICTs i c t's +ICTS i c t s +ICTU i c t u +ICTUR i c t u r +ICTVdB i c t v d b +ICTV's i c t v's +ICTY i c t y +ICU i c u +ICUN i c u n +ICU's i c u's +ICv i c v +ICV i c v +ICVS i c v s +ICWA i c w a +ICWC i c w c +ICW i c w +ICW's i c w's +iCyt i c y t +ICZ i c z +ICZN i c z n +IDAF i d a f +IDBF i d b f +IDB i d b +IDBI i d b i +IDC i d c +Iddaa i d d a a +Idd i d d +ID'd i d d +IDEF i d e f +IDEG- i d e g +IDE i d e +IDFA i d f a +IDFG i d f g +IDF i d f +IDFL i d f l +IDF's i d f's +IDG i d g +I. D. H. i d h +IDH- i d h +IDH i d h +id i d +I. D. i d +I.D. i d +ID- i d +ID i d +Idi i d i +Idir i d i r +IDIS i d i s +IDLHs i d l h's +IDL i d l +Idm i d m +IDM i d m +IDMR i d m r +IDMS i d m s +IDNA i d n a +IDN i d n +IDNR i d n r +IDNR's i d n r's +IDPs i d p's +IDRC i d r c +IDRO i d r o +IDSA i d s a +IDSA's i d s a's +IDSF i d s f +ids i d s +ID's i d's +IDs i d's +IDV i d v +IDW i d w +IDWO i d w o +IDW's i d w's +IDx i d x +IDX i d x +IEA- i e a +IEA i e a +I.E.M. i e m +IEM i e m +I.E.S. i e s +IES i e s +i.e. that is +Ifa i f a +IFA- i f a +IFA i f a +IFBB i f b b +IFBC i f b c +I.F.B. i f b +IFCC- i f c c +IFCC i f c c +ifc i f c +IFC i f c +IFC's i f c's +IfF i f f +I. F. F. i f f +IFF i f f +IFFI i f f i +IFFK i f f k +Ifft i f f t +IFHA i f h a +IFHP i f h p +ifi i f i +Ifi i f i +IFI i f i +IFIP i f i p +IFIs i f i's +IFK i f k +IFLA i f l a +IFL i f l +IfM i f m +IFM i f m +IFN- i f n +IFN i f n +Ifni i f n i +IFNs i f n's +IFO i f o +Ifop i f o p +IFP i f p +IFPI i f p i +IFPRI i f p r i +IFRC i f r c +IFR i f r +IFRS i f r s +IFRSs i f r s's +IFSB i f s b +IFSC i f s c +IFSTA i f s t a +IFTA i f t a +IFTF i f t f +IFT i f t +Iga i g a +IgA i g a +IGA i g a +IGA's i g a's +IGBT i g b t +IGC i g c +IGCP i g c p +IGCSE i g c s e +IGD i g d +IGF- i g f +IGF i g f +IGFI i g f i +IGFs i g f's +IGIB i g i b +'ig i g +Ig i g +I.G. i g +IG i g +igi i g i +IGI i g i +IgM i g m +IGM i g m +IGMP i g m p +ign i g n +Ign i g n +IGN i g n +IHF i h f +IHGS i h g s +IHH i h h +Ih i h +I. H. i h +IH i h +ihi i h i +IHI i h i +IHL i h l +IHL's i h l's +IHMC i h m c +IHME i h m e +ihm i h m +ihn i h n +Ihn i h n +IHOM i h o m +Ihor i h o r +IHPK i h p k +IHRC i h r c +IHSAA- i h s a a +IHSAA i h s a a +IHSA i h s a +IHS i h s +Iht i h t +IIA i i a +IIb i i b +IIB- i i b +IIB i i b +IIc i i c +IIDC i i d c +IID i i d +IIDI i i d i +IIDS i i d s +IIDX i i d x +IIED i i e d +IIFA i i f a +IIF i i f +IIfx i i f x +Iiga i i g a +IIga i i g a +IIG i i g +IIgs i i g's +IIGs i i g's +IIGS i i g s +IIGS's i i g s's +IIHF i i h f +IIHR i i h r +IIHS i i h s +IIIA i i i a +IIIb i i i b +IIIB i i i b +IIIBy i i i b y +IIIc i i i c +IIId i i i d +IIIF i i i f +I. I. i i +I.I. i i +ija i j a +Ija i j a +IJA i j a +IJF i j f +IJGL i j g l +IJHCD i j h c d +IJHL i j h l +ij i j +I. J. i j +IJ- i j +IJ i j +IJN i j n +Ijok i j o k +IJSB i j s b +IJSE i j s e +IJSEM i j s e m +IJS i j s +Ijui i j u i +IJVM i j v m +IKBN i k b n +IKCO i k c o +IKF i k f +Ikh i k h +I.K. i k +IKK- i k k +IKK i k k +IKL i k l +IKr i k r +Iksa i k s a +IKTA i k t a +Iku i k u +Ilbe i l b e +ILB i l b +ILC i l c +ILCS i l c s +ILD i l d +ILDIS i l d i s +ILFC i l f c +ILF i l f +IL&FS i l and f s +ili i l i +Ili i l i +ILI i l i +I. L. i l +I.L. i l +Ilm i l m +ILM i l m +ILM's i l m's +ILN i l n +ILP i l p +ILR i l r +ILRS i l r s +ilu i l u +Ilu i l u +ILU i l u +IMbd i m b d +IMBD i m b d +IMB i m b +IMCB i m c b +IMCI i m c i +imc i m c +IMC i m c +imdb i m d b +IMDb. i m d b +IMDb i m d b +IMDB- i m d b +IMDB. i m d b +IMDB i m d b +IMD i m d +IMD's i m d's +ime i m e +IME i m e +IMF, i m f +IMF i m f +IMG i m g +I&M i and m +IMI i m i +I. M. i m +I.M. i m +IMIS i m i s +IMK i m k +Iml i m l +IMMEX i m m e x +Imm i m m +Imms' i m m's +IMMU- i m m u +imoa i m o a +IMO's i m o's +IMPA i m p a +IMPDH i m p d h +impr i m p r +Impr i m p r +IMPSA i m p s a +IMR i m r +IMSLP i m s l p +iMSNs i m s n's +Imst i m s t +IMTA i m t a +IMT- i m t +IMT i m t +IMTV i m t v +IMUK i m u k +IMU's i m u's +IMWF i m w f +INAC i n a c +Inba i n b a +INBA i n b a +INBF i n b f +inb i n b +incl i n c l +INCT- i n c t +INDC i n d c +ind i n d +Ind i n d +IND i n d +INFJ i n f j +I.N.F.O. i n f o +Infs i n f's +ING i n g +INGR i n g r +inHg i n h g +INH i n h +INI i n i +I. N. i n +I.N. i n +InlB i n l b +INL i n l +INM i n m +INPD i n p d +INPE i n p e +InP i n p +INP i n p +INPP i n p p +INRA- i n r a +Inre i n r e +I.N.R.I. i n r i +INR i n r +INSB i n s b +INSDC i n s d c +INSD i n s d +Inspx i n s p x +INSS i n s s +INTF i n t f +Ints i n t's +INTs i n t's +Intu i n t u +INVA i n v a +inv i n v +INXS i n x s +INXS's i n x s's +Ioba i o b a +IOCG i o c g +IOC- i o c +IOC i o c +IOCP i o c p +IODP i o d p +IODs i o d's +IOE i o e +IOFB i o f b +IOF i o f +io. i o +io i o +.I.o. i o +Io. i o +Io i o +.IO i o +I/O i o +IO- i o +IO i o +Ioka i o k a +Iok i o k +IOK- i o k +IOL i o l +IoM i o m +IOM i o m +IOPI i o p i +IOP i o p +IOPs i o p's +IORG i o r g +IOR- i o r +IOR i o r +IORM i o r m +IOSA i o s a +IOSB i o s b +iOS- i o s +iOS i o s +Io's i o's +I/Os i o's +I.O.'s i o's +IOS i o s +IOSR i o s r +Iosu i o s u +IoT i o t +IOT i o t +I.O.U. i o u +IOU i o u +IoW i o w +IOW i o w +IPA i p a +IPA's i p a's +IPB i p b +IPCC i p c c +IPCC's i p c c's +IPC i p c +IPCRI i p c r i +IPCs i p c's +IPDA i p d a +IPDI i p d i +IPD i p d +IPE i p e +IPFA i p f a +IPFI i p f i +IPF i p f +IPFP i p f p +IPFW i p f w +IPG i p g +IPH i p h +ip i p +Ip i p +I.P. i p +IP- i p +IP i p +IPL i p l +IPMG i p m g +IPM i p m +IPNI i p n i +IPN i p n +Ipo i p o +I.P.O. i p o +IPO i p o +IPP i p p +IPPL i p p l +IPPs i p p's +IPPT i p p t +Ippu i p p u +iPPV i p p v +IPR i p r +iPSC i p s c +IPSC i p s c +iPSCs i p s c's +ipse i p s e +IPSF i p s f +iPS i p s +IPs i p's +I.P.S. i p s +IPS i p s +IPTA i p t a +IPTC i p t c +IPTG i p t g +IPT i p t +IPTL i p t l +IPTV i p t v +IPU i p u +IPUM i p u m +IPv i p v +IPWA i p w a +IPW i p w +IPX i p x +iQ i q +I.Q. i q +IQ i q +IRAC i r a c +I.R.A. i r a +IRB i r b +IRBM i r b m +IRCA i r c a +IRC i r c +IRCM i r c m +IRCTC i r c t c +IRCT i r c t +IRD i r d +IRDS i r d s +IREC i r e c +I.R.E. i r e +IRE i r e +IRFCA i r f c a +IRFC i r f c +IRF i r f +IRFU i r f u +IRGC i r g c +Irgm i r g m +IRHG i r h g +IRHT i r h t +IRIB i r i b +Irig i r i g +IRI i r i +Iril i r i l +IRIN i r i n +i'r i r +ir i r +Ir i r +I. R. i r +I.R. i r +IR- i r +IR i r +IRIX i r i x +Irla i r l a +IRL i r l +IRM i r m +Irmis i r m i's +IRM's i r m's +IRNE i r n e +IRN- i r n +IRN i r n +iro i r o +Iro i r o +IRO i r o +IRP i r p +IRP's i r p's +IRPS i r p s +IRQ i r q +Irra i r r a +IRRI i r r i +IRR i r r +IRs i r's +I.R.S. i r s +IRS- i r s +IRS i r s +IRSN i r s n +I.R.T. i r t +IRT i r t +iru i r u +Iru i r u +ISA i s a +ISAv i s a v +Isba i s b a +ISBA i s b a +ISBER i s b e r +ISBHF i s b h f +ISB i s b +isbn i s b n +ISBN i s b n +ISBNs i s b n's +ISBO i s b o +ISCB i s c b +ISCC- i s c c +ISCE i s c e +isCf i s c f +ISCI i s c i +iSC i s c +Isc i s c +ISC- i s c +ISC i s c +ISCM i s c m +ISCM's i s c m's +ISCP i s c p +ISDA i s d a +ISDB- i s d b +Isd i s d +ISD i s d +isdn i s d n +ISDN i s d n +ISDS i s d s +ISDT i s d t +ISFI i s f i +ISF i s f +Isgec i s g e c +ISG i s g +ISHC i s h c +ishq i s h q +Ishq i s h q +ISHS i s h s +I. S. i s +I.S. i s +ISKB i s k b +ISK i s k +Isl i s l +ISL i s l +ISMF i s m f +ISMIR i s m i r +ISMN- i s m n +ISMNs i s m n's +isn i s n +ISN i s n +ISNTUC i s n t u c +ISOGG i s o g g +ISO- i s o +ISO i s o +ISO's i s o's +I.S.P.C.K. i s p c k +ISPF i s p f +ISPI i s p i +ISP i s p +ISPL i s p l +ISPR i s p r +ISPs i s p's +Isra'il i s r a i l +Isra i s r a +ISRA i s r a +ISR i s r +ISRN i s r n +ISRO i s r o +ISRP i s r p +ISSCC i s s c c +ISSF i s s f +iss i s s +Iss i s's +ISS i s s +ISSN- i s s n +ISSN i s s n +ISSNs i s s n's +ISSSSP i s s s s p +ISSUU i s s u u +ISTAF i s t a f +ISTD i s t d +Isu i s u +ISU i s u +ISU's i s u's +ISWC i s w c +ISWI i s w i +ISW i s w +I.T.A. i t a +ITA i t a +ITB i t b +ITC i t c +Ite i t e +ITE i t e +ITF i t f +ITG i t g +ITHF i t h f +I.T.H.M. i t h m +ITHs i t h's +iti i t i +Iti i t i +ITI i t i +I.T. i t +itk i t k +Itk i t k +ITK i t k +ITM i t m +ITMO i t m o +ITN i t n +ITNNs i t n n's +ITN's i t n's +Itoi i t o i +ITPF i t p f +iTP i t p +ITP i t p +ITRD i t r d +Itse i t s e +ITSF i t s f +ITTF i t t f +itt i t t +ITT i t t +ITUC i t u c +Itu i t u +ITU- i t u +ITU i t u +ITV i t v +ITV's i t v's +ITWF i t w f +ITW i t w +Itz i t z +IUB i u b +I.U.C.N.B. i u c n b +Iucn i u c n +IUCN- i u c n +IUCN i u c n +IUCN's i u c n's +IUD i u d +IUDs i u d's +IU i u +IUT i u t +IUU i u u +Iuz i u z +IVAA i v a a +IVAW i v a w +IVB i v b +I. V. i v +I.V. i v +IVL i v l +IVLP i v l p +IVM's i v m's +IVP i v p +IVPP i v p p +IVPTC i v p t c +IVRCL i v r c l +IVs i v's +Ivu i v u +IVWA i v w a +IWA i w a +Iwas i w a's +Iwaz i w a z +IWBs i w b's +IWC i w c +IWD i w d +IWF i w f +IWFL i w f l +IWF's i w f's +IWFS i w f s +IWG i w g +IWGP i w g p +iwi i w i +iw i w +I. W. i w +IW- i w +IW i w +IWL i w l +IWMF i w m f +IWM i w m +IWPR i w p r +iwrg i w r g +IWRG i w r g +IWR i w r +IWRS i w r s +IWSA i w s a +iwspy i w s p y +Iwuh i w u h +IWU's i w u's +I.W.W. i w w +IWW i w w +IXA i x a +IXb i x b +IXBs i x b's +IXC i x c +IXCs i x c's +IXe i x e +IX i x +iXL i x l +IXO i x o +ixtle i x t l e +IXV i x v +Iya i y a +IYA i y a +Iyar i y a r +IYC i y c +iyem i y e m +IYHF i y h f +Iyi i y i +iy i y +I. Y. i y +I.Y. i y +IY i y +J.A.A. j a a +J. A. B. j a b +Jaf j a f +J.A.G. j a g +JAG j a g +JAIR j a i r +J. A. j a +J.A.K.E. j a k e +J. A. N. j a n +J.A.P. j a p +J.A.S. j a s +JBA j b a +JBBA j b b a +JBB j b b +JBC j b c +JBG j b g +J&B j and b +jb j b +J. B. j b +J.B. j b +JB j b +J.B.L.D. j b l d +JBL j b l +J.B.M. j b m +JBM j b m +JB's j b's +J. B. S. j b s +J.B.S. j b s +JBT j b t +JBZ j b z +JCA j c a +JCB j c b +JCC j c c +JCF j c f +JCFS j c f s +J.C.G. j c g +JCG j c g +JCGM j c g m +JCIB j c i b +JCI j c i +jcis j c i s +jc j c +J. C. j c +J.C. j c +JC j c +JCJC j c j c +J. C. L. j c l +JCPA j c p a +J. C. P. j c p +JCP j c p +JCPV j c p v +JCRB j c r b +JCR j c r +JCRs j c r's +J.C.S. j c s +JCS j c s +JCSR j c s r +JCVI j c v i +J. C. W. j c w +JCW j c w +J. D. A. j d a +JDA j d a +JDBC j d b c +JDC j d c +JdeBP j d e b p +J.D.E. j d e +JDE j d e +JDENET j d e n e t +J.D.F. j d f +JDF j d f +jdi j d i +JDI j d i +jd j d +J. D. j d +J.D. j d +JD- j d +JD j d +jdk j d k +JDK j d k +J.D.M. j d m +JDM j d m +JDO j d o +JDRF j d r f +J&D's j and d's +J.D.s j d's +JDT j d t +JDU j d u +JDV j d v +J. E. A. j e a +J.E.A.L. j e a l +J.E.B. j e b +J. E. C. j e c +J. E. E. j e e +J. E. j e +J.E. j e +J.F.G. j f g +JFH j f h +jf j f +J. F. j f +J.F. j f +JF- j f +JF j f +JFK j f k +JFN j f n +JFQA j f q a +JFS j f s +J. G. A. j g a +J.G.B. j g b +JGI j g i +J. G. j g +J.G. j g +J. G. L. j g l +JGLS j g l s +J. G. M. j g m +J.G.M. j g m +JGO j g o +JGP j g p +JGR j g r +JGR's j g r's +JGSDF j g s d f +JGTC j g t c +JGU j g u +J.H.B. j h b +JHB j h b +J. H. C. j h c +JHC j h c +JHED j h e d +J.H.F. j h f +J.H.H. j h h +jh j h +J. H. j h +J.H. j h +JH- j h +JH j h +JHL j h l +J. H. M. j h m +J.H.R. j h r +JHs j h's +JHS j h s +JHSVs j h s v's +JHud j h u d +JHU- j h u +JHU j h u +J. I. j i +J.I. j i +J.J.A. j j a +JJB j j b +J. J. C. j j c +J.J.G. j j g +J&J j and j +J. J. j j +J.J. j j +JJ j j +JJK j j k +J. J. L. j j l +J.J.M. j j m +J.J.N. j j n +JJN j j n +Jka j k a +JKA j k a +JKC j k c +jkd j k d +JKH j k h +J&K j and k +J.K. J. j k j +J. K. j k +J.K. j k +JK j k +JKNPP j k n p p +J. K. S. j k s +JKT j k t +jkx j k x +JKX j k x +JLA j l a +J.L.B. j l b +JLB j l b +JLC j l c +J. L. E. j l e +J.L.E. j l e +JLG j l g +JLI j l i +J. L. j l +J.L. j l +J.L.L. j l l +JLM j l m +JLO's j l o's +JLPGA j l p g a +JLP j l p +JLS j l s +JLT j l t +JMA j m a +JMBG j m b g +J.M.C. j m c +JMC j m c +J. M. D. j m d +JMD j m d +J. M. E. j m e +JMI j m i +J. M. J. j m j +J.M.J. j m j +J. M. j m +J.M. j m +JMKR j m k r +J. M. L. j m l +JML j m l +JMM j m m +JMP j m p +JMS j m s +JMU j m u +JMV j m v +J. M. W. j m w +J.M.W. j m w +JMX j m x +JNA j n a +jnb j n b +JNF j n f +JNG- j n g +JNG j n g +JNIH- j n i h +JNI j n i +J. N. j n +J.N. j n +JN j n +JNK j n k +J. N. L. j n l +Jno j n o +JNPT j n p t +jnr j n r +JNR j n r +JNS's j n s's +JOJ j o j +J. O. j o +J.O. j o +J.P.D.B. j p d b +J. P. F. j p f +JPF j p f +jpg j p g +JPG j p g +JPH j p h +J&P j and p +jp j p +J.P. j p +JP- j p +JP j p +J.P.R. j p r +J.Q. j q +JQ j q +J.R.A. j r a +JRA j r a +J.R.C. j r c +JRC j r c +J.R.E. j r e +JRG j r g +JRHU j r h u +J. R. I. j r i +J&R j and r +J. R. j r +J.R. j r +JR j r +JRJ's j r j's +Jr. junior +JRPG j r p g +J. R. P. j r p +J. R. R. j r r +J.R.R. j r r +J.R.'s j r's +JRSP j r s p +J. R. T. j r t +JRU j r u +JRX j r x +JSA j s a +JSAP j s a p +JSA's j s a's +jsb j s b +JSCA j s c a +JSC j s c +J.S.D. j s d +JSD j s d +JSE j s e +JSF j s f +JSH j s h +Js j's +J. S. j s +J.S. j s +JS j s +J.S.K. j s k +J.S.L. j s l +JSL j s l +J. S. M. j s m +JSM j s m +JSNN j s n n +JSOC j s o c +JSO j s o +JSP j s p +JSPS j s p s +JSR j s r +JSS j s s +JSW j s w +JSW's j s w's +JSX j s x +JTA j t a +J.T.B. j t b +JTB j t b +JTC j t c +JTF- j t f +JTF j t f +JTG j t g +JTI j t i +J&T j and t +J. T. j t +J.T. j t +JT j t +JT&KW j t and k w +JTL j t l +JTM- j t m +JTO j t o +JTR j t r +JTRO j t r o +JTS j t s +JTTF j t t f +JTT j t t +JTV j t v +J. U. j u +J.U.L.I.A. j u l i a +J.V.B. j v b +JVB j v b +JVC j v c +J.V.E. j v e +JVG j v g +JVH j v h +J&V j and v +J. V. j v +J.V. j v +JV- j v +JV j v +JVM j v m +JVNW j v n w +JVP j v p +JVRA j v r a +JVS j v s +Jwa j w a +J. W. A. j w a +J. W. B. j w b +JWB j w b +J.W.C. U.S. j w c u s +JWHA j w h a +JWH- j w h +J.W.J. j w j +J. W. j w +J.W. j w +JWM j w m +J. W. P. j w p +JWP j w p +JWRC j w r c +J. W. S. j w s +J.W.S. j w s +JWS j w s +JWST j w s t +JWT j w t +jx j x +JX- j x +JX j x +J.X.W.P. j x w p +JYJ j y j +JYJ's j y j's +J. Y. j y +J.Y. j y +JY j y +Jym j y m +Jymn j y m n +JYP j y p +J. Y. S. j y s +J. Z. j z +J.Z. j z +JZ- j z +JZ j z +K.A.A. k a a +K. A. k a +K.A. k a +K.A.N. k a n +K. A. R. k a r +K.A.S. k a s +KBBI k b b i +KBB k b b +KBCI k b c i +KBC k b c +KBCO k b c o +KBD k b d +KBEH- k b e h +KBE k b e +KBEV- k b e v +KBFC k b f c +KBFL k b f l +KBFR k b f r +KBGD k b g d +KBG k b g +KBIG k b i g +kbi k b i +K&B k and b +K. B. k b +K.B. k b +KBKR k b k r +KBKS k b k s +KBKW k b k w +KBL k b l +KBMT- k b m t +KBMT k b m t +KBND k b n d +KBNZ- k b n z +KBNZ k b n z +KBO k b o +KBRC k b r c +KBR k b r +KBSC- k b s c +KBS k b s +KBSN k b s n +KCB k c b +KCBS- k c b s +KCBS k c b s +KCCC k c c c +KCCI k c c i +KCC k c c +KCCO k c c o +KCDC k c d c +K. C. D. k c d +KCD k c d +KCED k c e d +KCET k c e t +KCG k c g +KCGW- k c g w +KCHK- k c h k +KCHS k c h s +K.C.I.E. k c i e +K. C. k c +K.C. k c +KCK k c k +KCKM k c k m +KCLU k c l u +K.C.M.G. k c m g +KCMG k c m g +KCMO- k c m o +KCMU- k c m u +KCNA k c n a +KCNH k c n h +KCNJ k c n j +KCN k c n +KCNN k c n n +KCNQ k c n q +KCNV k c n v +KCPD k c p d +KCPI k c p i +KCPM k c p m +KCPQ k c p q +KCPQ's k c p q's +KCQL k c q l +KCRA k c r a +KCRC k c r c +KCRH k c r h +KCR k c r +KCRW k c r w +KCSD- k c s d +KCSD k c s d +KCSG k c s g +K.C.S.I. k c s i +KCSI k c s i +K.C.'s k c's +KCS k c s +KCSP k c s p +KCTD k c t d +KCT k c t +KCTL k c t l +KCTV k c t v +KCUB k c u b +KCU k c u +KCVO k c v o +KCYX k c y x +kDa k d a +KDA k d a +kdal k d a l +KDAL k d a l +KDB k d b +KDCE k d c e +KDCO k d c o +KDDG k d d g +KDDI k d d i +KDD k d d +KDE- k d e +KDE k d e +KDEX- k d e x +KDFW k d f w +KDFX k d f x +K. D. G. k d g +KDGS k d g s +KDHS k d h s +KDIS- k d i s +KDKA- k d k a +KDKA k d k a +K&D k and d +kd k d +K. D. k d +K.D. k d +KD- k d +KD k d +KDP k d p +KDPM k d p m +KDR k d r +K. D.'s k d's +K.D.'s k d's +KDS k d s +ke k e +Ke- k e +Ke k e +K. E. k e +K.E. k e +KE k e +KFAC k f a c +KFAED k f a e d +KFAI k f a i +KFA k f a +KFBK k f b k +KFBT k f b t +KFCC k f c c +KFCD k f c d +kfc k f c +K.F.C. k f c +KFC k f c +KFC's k f c's +KFDI- k f d i +KFD k f d +KFDM- k f d m +KFDM k f d m +KFEM k f e m +KFFL k f f l +KFH k f h +KFI k f i +KFJB k f j b +K.F. k f +KFLX k f l x +KFMB- k f m b +KFMB's k f m b's +KFMG- k f m g +KFMJ k f m j +K. F. M. k f m +K.F.M. k f m +KFOR k f o r +KFQC k f q c +KFQ k f q +KFRC k f r c +KFRC's k f r c's +KFRD k f r d +KFRG k f r g +KFRH k f r h +KFRO k f r o +KFSD k f s d +KFSN- k f s n +KFSN k f s n +KFST k f s t +KFTA- k f t a +KFT k f t +KFTL- k f t l +KFTL k f t l +KFUK- k f u k +KFUM- k f u m +KFUM's k f u m's +KFVD k f v d +KFV k f v +KFWB k f w b +KFXO k f x o +KFYO- k f y o +KFYO's k f y o's +KFYR- k f y r +KGA k g a +K.G.B. k g b +KGB- k g b +KGB k g b +KGBX- k g b x +KGC k g c +KGPH k g p h +KGP k g p +KGRI k g r i +kgr k g r +KGRK k g r k +KGS- k g s +KGS k g s +kgt k g t +KGT k g t +KGTV k g t v +KgU k g u +KGU k g u +KGUN k g u n +kgv k g v +K.G.V. k g v +KGWB k g w b +KGW k g w +KHHZ k h h z +KHI k h i +Khizr k h i z r +KHJ- k h j +KHJ k h j +Kh- k h +K. H. k h +K.H. k h +KH- k h +KH k h +KHK k h k +KHKV k h k v +Khlav k h l a v +KHL k h l +KHL's k h l's +Khmu k h m u +KHMY k h m y +KHOJ k h o j +KHOW k h o w +KHQA k h q a +KHQ k h q +Khri k h r i +Khru k h r u +KHRW k h r w +KHSAA k h s a a +KHSC- k h s c +KHS k h s +KHTL k h t l +KHTZ k h t z +KHUI k h u i +KHUT k h u t +KHWI k h w i +K.I.D. k i d +K. I. k i +K.I. k i +K.I.M. k i m +KIP k i p +K.I.S.S. k i s s +KJCE k j c e +KJCT k j c t +KJEF k j e f +Kjer k j e r +KJHK k j h k +KJI k j i +KJIN k j i n +kj k j +K. J. k j +K.J. k j +KJ k j +KJKP k j k p +KJMM k j m m +KJNE k j n e +KJNW k j n w +Kjop k j o p +KJo's k j o's +KJQY k j q y +KJV k j v +KJZI k j z i +KKBL k k b l +KKBQ k k b q +KKCR k k c r +KKDA k k d a +KKDS- k k d s +K.K.E. k k e +KKE k k e +K. K. k k +K.K. k k +KKMC k k m c +KKMI k k m i +KKM k k m +KKMX k k m x +KKRG k k r g +KKR k k r +KKSF k k s f +KK's k k's +KKSY k k s y +KKT k k t +KKWB- k k w b +KLA k l a +KLBK k l b k +KLBM k l b m +KLCI k l c i +KLC k l c +KLDS k l d s +Kle k l e +KLE k l e +KLF k l f +KLGH k l g h +KLH k l h +KLK k l k +kl k l +K. L. k l +K.L. k l +KL k l +KLKS k l k s +KLLP k l l p +KLLV k l l v +KLM k l m +KLN's k l n's +KLQQ k l q q +KLQT k l q t +KLRY k l r y +KLSC k l s c +KLSD k l s d +KLS- k l s +KLS k l s +KLSR- k l s r +KLSX k l s x +KLT k l t +KLX k l x +KMB k m b +KMCI- k m c i +KMC k m c +KMCX- k m c x +KMD k m d +KMGO k m g o +KMGT k m g t +KMIH k m i h +KMIR k m i r +KMITL k m i t l +KMJC- k m j c +KMJ k m j +K&M k and m +KMK k m k +km k m +K. M. k m +K.M. k m +KM- k m +KM k m +KMLA k m l a +KML k m l +KMME k m m e +KMML k m m l +KMRL k m r l +KMS k m s +KMSP k m s p +KMT k m t +KMTP k m t p +KMTV k m t v +KMVA k m v a +KMVQ- k m v q +KMXW k m x w +KMZG- k m z g +KMZ k m z +KNFM k n f m +KNHP k n h p +K&N k and n +kn k n +K. N. k n +K.N. k n +KN k n +KNKT k n k t +KNLA k n l a +KNMI k n m i +KNNV k n n v +KNPN k n p n +KNPR k n p r +KNPU k n p u +KNRE k n r e +KNRS k n r s +KNS k n s +KNSM k n s m +KNST- k n s t +KNST k n s t +KNTH k n t h +K.N.T. k n t +KNTN k n t n +KNTU k n t u +KNVA's k n v a's +KNVB k n v b +KNVN k n v n +KNWA- k n w a +KNX- k n x +KNX k n x +KNXT k n x t +KNXV- k n x v +K.O.D. k o d +k'o k o +K. O. k o +K.O. k o +KOTC k o t c +KOTD k o t d +K.O.T. k o t +KOWL k o w l +K. P. A. C. k p a c +KPAC- k p a c +KPAC k p a c +KPA k p a +KPAN k p a n +KPCC k p c c +KPC k p c +KPD k p d +KPDR- k p d r +KPE k p e +KPFT- k p f t +KPIF k p i f +KPI- k p i +KPI k p i +KPIs k p i's +K&P k and p +K. P. k p +K.P. k p +KP k p +KPLC k p l c +KPL k p l +KPLZ k p l z +KPMG k p m g +KPM k p m +KPNA k p n a +KPNB k p n b +KPO k p o +KPPC- k p p c +Kppen k p p e n +KPPV k p p v +KPRC- k p r c +KPS k p s +KPSL- k p s l +kptm k p t m +KPTV- k p t v +KPTV k p t v +KPU k p u +KPVD k p v d +KPVX k p v x +KPWR k p w r +KPXE k p x e +KPXO- k p x o +KPZK- k p z k +KQA k q a +KQBU- k q b u +KQCD- k q c d +KQED k q e d +KQEG- k q e g +KQHN k q h n +KQKK k q k k +KQKS k q k s +KQLT k q l t +KQLZ k q l z +KQMO k q m o +KQRS- k q r s +KQV k q v +KQXR k q x r +KRBM k r b m +krc k r c +KRCR k r c r +KRCs k r c's +KRF k r f +KRH k r h +K.R.I.T. k r i t +K.R.I.T.'s k r i t's +Krka k r k a +KRKC k r k c +KRK k r k +kr k r +K. R. k r +K.R. k r +KR k r +KRMT k r m t +Krne k r n e +KRNS k r n s +KRNV- k r n v +KRNV k r n v +KROQ- k r o q +KROQ k r o q +kroz k r o z +Kroz k r o z +KRPS k r p s +KRRT k r r t +K. R. S. k r s +KRS- k r s +KRS k r s +KRT k r t +Kru k r u +KRU k r u +KRWC k r w c +KRX k r x +K.R.Y. k r y +KRZA k r z a +ksa k s a +KSA k s a +K. S. B. k s b +KSBW- k s b w +KSBW k s b w +KSBY k s b y +KSCB k s c b +KSC k s c +KSCO k s c o +KSC's k s c's +KSCS k s c s +KSDB k s d b +KSD- k s d +KSDO k s d o +KSEB k s e b +KSEG k s e g +KSE k s e +KSEQ k s e q +kset k s e t +KSET k s e t +Kseur k s e u r +Ksevt k s e v t +KSFO k s f o +KSG k s g +KSHB k s h b +KSI's k s i's +KSJR k s j r +KSKJ k s k j +KSKK k s k k +K.S.K. k s k +KSK k s k +ks k s +K. S. k s +K.S. k s +KS k s +KSLA- k s l a +KSLC k s l c +KSL k s l +KSLM k s l m +KSLU k s l u +KSMB k s m b +KSNB- k s n b +KSNR k s n r +KSNV k s n v +KSP k s p +KSPS k s p s +KSTE k s t e +KSTH k s t h +KST k s t +KSTP- k s t p +KSTP k s t p +KSTV- k s t v +KSU- k s u +KSU k s u +K.S.V. k s v +KSWD k s w d +KSWT k s w t +KSWW- k s w w +KSWW k s w w +KSYL k s y l +KSZE- k s z e +KSZR k s z r +KTAC- k t a c +Ktav k t a v +KTAV k t a v +KTBS k t b s +KTCC k t c c +KTC k t c +KTCZ k t c z +KTDA k t d a +KTDD k t d d +KTE k t e +KTFA- k t f a +KTH k t h +KTHV- k t h v +KT&K k t and k +kt k t +K. T. k t +K.T. k t +KT- k t +KT k t +KTKZ k t k z +KTLA k t l a +K. T. M. k t m +KTM k t m +KTMS- k t m s +KTNE- k t n e +KTN k t n +KTOK k t o k +KTO k t o +KTRE k t r e +KTRH k t r h +KTRK k t r k +KTR k t r +KT's k t's +KTSM k t s m +KTT k t t +KTTN k t t n +KTTV k t t v +KTU k t u +KTUL- k t u l +KTUM k t u m +KTUU- k t u u +KTVA k t v a +KTVK k t v k +KTVS k t v s +KTVT k t v t +KTVU k t v u +KTVZ k t v z +KTWD's k t w d's +KTWN- k t w n +KTWO k t w o +KTXC- k t x c +KTX k t x +KTXL k t x l +KTXS- k t x s +KTZR- k t z r +Kuaa k u a a +Kud k u d +KUFM k u f m +KUGB k u g b +kuih k u i h +K. U. k u +K.U. k u +Kutb k u t b +KUTV k u t v +KUUU k u u u +KUWL k u w l +KUYI k u y i +KVAB k v a b +KVBC k v b c +KVB k v b +K.V.G.K. k v g k +KVHV k v h v +Kvik k v i k +KVI k v i +Kvit k v i t +KVIT k v i t +KVK k v k +K. V. k v +K.V. k v +KVLO k v l o +KVLY- k v l y +KVLY k v l y +KVM k v m +KVMRT k v m r t +KVMX k v m x +KVNG k v n g +KVN k v n +KVP k v p +KVTV k v t v +KVVS's k v v s's +KVVV k v v v +Kvyat k v y a t +KVZ k v z +kwa k w a +KWBM k w b m +KWC k w c +KWD k w d +KWEI k w e i +kwe k w e +KWG k w g +KWJJ k w j j +KWKH k w k h +KWKW k w k w +KWMR k w m r +KWMT k w m t +KWNA k w n a +KWNK k w n k +KWOA k w o a +KWP k w p +KWPT k w p t +KWQC k w q c +KWRE k w r e +KWRU k w r u +KWSB k w s b +KWSX k w s x +KWU k w u +KWVE- k w v e +KWWL's k w w l's +KXAL- k x a l +KXAN k x a n +KXAS k x a s +KXHD- k x h d +KXI k x i +KXJB- k x j b +KXJB k x j b +KXKS k x k s +KXKW- k x k w +KX k x +KXLX k x l x +KXLY k x l y +KXMC k x m c +KXNA k x n a +KXOL's k x o l's +KXTA k x t a +KXTN k x t n +KXTX k x t x +KYAY k y a y +KYBE k y b e +KYEZ k y e z +KYLX- k y l x +KYNM k y n m +Kyse k y s e +KYSL k y s l +KYTC k y t c +kyt k y t +Kyt k y t +KYW- k y w +KYW k y w +kz k z +K. Z. k z +KZ k z +KZLZ k z l z +KZM- k z m +KZMP k z m p +KZMU k z m u +KZQX k z q x +KZZQ k z z q +L. A. G. l a g +L. A. l a +L.A.'s l a's +Lay's l a y's +L.B.A. l b a +LBA l b a +LBCC l b c c +LBC l b c +LBC's l b c's +LBE- l b e +LBE l b e +LBF l b f +LBi l b i +LBJ l b j +LBJ's l b j's +L. B. l b +L.B. l b +LB l b +LBL l b l +LBM l b m +LBN l b n +LBNL l b n l +LBP l b p +LBPs l b p's +LBR l b r +LBS l b s +LBV l b v +LBW l b w +LCAC l c a c +LCA l c a +LCAs l c a's +LCAS l c a s +L.C.B. l c b +LCBL l c b l +LCBM l c b m +lcc l c c +LCC l c c +LCCN l c c n +LCCs l c c's +LCCS l c c s +LCDB l c d b +LCD- l c d +LCD l c d +LCDP l c d p +LC&DR l c and d r +LCDR l c d r +LCDs l c d's +L.C.E. l c e +LCFC l c f c +LCF l c f +LCH l c h +LCI l c i +LCIs l c i's +L. C. l c +L.C. l c +LC- l c +LC l c +LCL l c l +LCMC l c m c +LCM l c m +LCMR l c m r +LCMV l c m v +lcn l c n +LCO l c o +LCP l c p +LCQ l c q +LCR l c r +LCSB l c s b +LCSC l c s c +LC's l c's +LCS- l c s +LCS l c s +LCSs l c s's +LCST l c s t +LCT- l c t +LCT l c t +LCTs l c t's +LCU l c u +LCV l c v +LCVP l c v p +LCVPs l c v p's +LCZ l c z +Lda l d a +LDA l d a +LDAP l d a p +LDB l d b +LDBV l d b v +L.D.C. l d c +LDC l d c +LDCM l d c m +LDCs l d c's +LDD l d d +LDDP l d d p +LDF l d f +LDH l d h +L.D.I. l d i +LDK l d k +L. D. l d +L.D. l d +LDL- l d l +LDL l d l +LDLR l d l r +LDM l d m +LDMR l d m r +LDN l d n +LDP l d p +LDS l d s +LDU l d u +LDV l d v +LEB l e b +L. E. l e +L.E. l e +LFA- l f a +LFA l f a +L.F.C. l f c +LFC l f c +L.F.D. l f d +LFE l f e +LFFCs l f f c's +LFF l f f +LFG l f g +LFHC l f h c +LFH l f h +L. F. l f +L.F. l f +LFL l f l +LFM l f m +LFOC l f o c +LFO l f o +LFOM l f o m +LFPB l f p b +LFP l f p +LFQD l f q d +LFRJ l f r j +LFR l f r +LFs l f's +LFS l f s +LFSRs l f s r's +LFTR l f t r +LFXA l f x a +LGA l g a +LGB l g b +LGBTI l g b t i +LGBT- l g b t +LGBT l g b t +LGBTQ l g b t q +LGC l g c +LGD- l g d +LGD l g d +LGI l g i +L. G. l g +L.G. l g +LG- l g +LG l g +LGMB l g m b +LGM- l g m +LGMs l g m's +LGN l g n +LGPL l g p l +LGPZ l g p z +LG&RDD l g and r d d +LGR l g r +LGs l g's +LGS l g s +LGTB l g t b +LGT's l g t's +LGU l g u +LGUs l g u's +LGV l g v +LGVs l g v's +Lha l h a +LHB l h b +L. H. C. l h c +LHC l h c +LHD l h d +LHFP l h f p +LHICE l h i c e +LHI l h i +L. H. l h +L.H. l h +LHMC l h m c +L.H.M. l h m +LHO l h o +L. H. P. l h p +LHP l h p +LHS l h s +LHW l h w +LHX l h x +LHY l h y +Li- l i +L. I. l i +L.I. l i +LJBL l j b l +Lje l j e +LJJ l j j +L. J. K. l j k +L. J. l j +L.J. l j +L.J.V. l j v +LKB l k b +LKG l k g +LKK l k k +L. K. l k +LK l k +LKL l k l +LKML l k m l +LKPR l k p r +LKRN l k r n +lks l k s +LKS l k s +LKTI l k t i +Llapi l l a p i +LLAW l l a w +L.L.B. l l b +LLB l l b +LLCC l l c c +L.L.C. l l c +LLC- l l c +LLC l l c +LLD l l d +lle l l e +LLE l l e +LLHS l l h s +lli l l i +LLI l l i +L&L l and l +LLLE l l l e +L. L. l l +L.L. l l +LL l l +LLL- l l l +LLL l l l +LLLP l l l p +LLM l l m +LLMNR l l m n r +LLPFX l l p f x +llp l l p +LLP l l p +LLPX l l p x +LLRW l l r w +LLT l l t +LLTV l l t v +llu l l u +LLVM l l v m +LLWS l l w s +LMA l m a +LMC l m c +LMCT l m c t +LMDB l m d b +LMDC l m d c +LMD l m d +LME l m e +LMFAO's l m f a o's +LMFF l m f f +LMG l m g +LMGs l m g's +LMGTE l m g t e +LMH l m h +LMHS l m h s +lm l m +L. M. l m +L.M. l m +LM- l m +LM l m +LMM l m m +LMP l m p +LMPs l m p's +LMQ l m q +LMQs l m q's +L.M.S. l m s +LMS l m s +LMT l m t +LMTP l m t p +LMU l m u +LMWH l m w h +LMX l m x +LnAIB l n a i b +LNA l n a +LNAV l n a v +LNBF l n b f +LNB l n b +L.N.C. l n c +LNCS l n c s +L.N.E.R. l n e r +LNER l n e r +LNFS l n f s +LNG l n g +LNH l n h +L&N l and n +L. N. l n +L.N. l n +L.O.C.'s l o c's +Lokk l o k k +l'OL l o l +L.O.L. l o l +L. O. l o +L.O. l o +L.O.V.E. l o v e +L.O.V. l o v +LOXL l o x l +LPA l p a +L.P.A.M. l p a m +LPAM l p a m +LPARs l p a r's +LPAVS l p a v s +LPBG l p b g +L.P.B. l p b +LPB l p b +LPC l p c +LPCM l p c m +LPDA l p d a +LPDR l p d r +LPE l p e +LPF l p f +LPFM l p f m +LPGA l p g a +LPG l p g +L.P.H. l p h +LPH- l p h +LPH l p h +LPI l p i +LPK- l p k +LPLA l p l a +L&P l and p +LPL- l p l +LPL l p l +lp l p +L. P. l p +L.P. l p +LP- l p +LP, l p +LP l p +L&PM l and p m +LPMN l p m n +LPMud l p m u d +LPMUD l p m u d +LPN l p n +LPO l p o +LPRP l p r p +LPSC l p s c +LP's l p's +LPs l p's +LPS l p s +LPSN l p s n +LPTB l p t b +LPThe l p t h e +LPTV l p t v +LQ l q +LRAD l r a d +LRA l r a +LRCD l r c d +LRC l r c +L.R.C.P. l r c p +LRCP l r c p +LRDG l r d g +LRG- l r g +LRG l r g +lr l r +L. R. l r +L.R. l r +LR l r +LRO l r o +L.R.P.C. l r p c +LRP l r p +LRPPRC l r p p r c +LRRI l r r i +LRRP l r r p +LRS l r s +LRTA l r t a +LRT- l r t +LRT l r t +LRTR l r t r +LRTs l r t's +LRVs l r v's +LRY l r y +LSAC l s a c +LSA l s a +LSBC l s b c +LSB l s b +LSC- l s c +LSC l s c +LSCS l s c s +LSDHH l s d h h +L.S.D. l s d +LSD- l s d +LSD l s d +LSDP l s d p +LSDs l s d's +LSE l s e +LSF l s f +LSG l s g +LSH l s h +LSi l s i +LSI- l s i +LSI l s i +LSJ l s j +LSK l s k +L.S.L. l s l +ls l s +Ls l's +L. S. l s +L.S. l s +LSM l s m +LSO l s o +LSP l s p +LSPN l s p n +LSQC l s q c +LSRI l s r i +L. S. R. l s r +LSS l s s +LSSP l s s p +LSSR l s s r +LST- l s t +LST l s t +LSTM l s t m +LSTs l s t's +LSUA l s u a +LSU l s u +LSU's l s u's +LSVCCs l s v c c's +LSV l s v +LSWR l s w r +LTAF l t a f +L. T. B. l t b +LTB l t b +LTBP- l t b p +LTCF l t c f +LTCI l t c i +L. T. C. l t c +L.T.C. l t c +LTC l t c +Ltda l t d a +LTDA l t d a +Ltd. limited +LTE- l t e +LTE l t e +LTFA l t f a +LTHS l t h s +LTI l t i +LTK l t k +L&T l and t +Lt. lieutenant +LTL l t l +lt l t +L. T. l t +L.T. l t +L.T.M. l t m +LTM l t m +LTMPS l t m p s +LTN l t n +ltoh l t o h +LTP l t p +LTRPC l t r p c +LTTE l t t e +LTTR l t t r +LTU l t u +LTV l t v +L. U. l u +lv l v +L. V. l v +L.V. l v +LV- l v +LV l v +LVMH l v m h +LWDB l w d b +LWD l w d +LWDS l w d s +LWE l w e +LWF l w f +LWH l w h +L. W. l w +L.W. l w +L.Y. l y +L&YR l and y r +M.A.D.E. m a d e +M.A.D. m a d +M.A.J. m a j +M.A.K. m a k +M. A. m a +M.A. m a +M&A m and a +M. A. O. m a o +M.A.O. m a o +M.A.R. m a r +M.A.S.K. m a s k +M.A.S. m a s +M.A.S.S. m a s s +M.B.A. m b a +M.B.B.S. m b b s +MBBS m b b s +MBC- m b c +MBC m b c +MBC's m b c's +MBDA m b d a +MBD m b d +M.B. D.P.M. m b d p m +M.B.E. m b e +MBE m b e +Mbewu m b e w u +MBF m b f +MBFW m b f w +MBGN m b g n +MBH- m b h +MBH m b h +mbi m b i +MBI m b i +M. B. J. m b j +MBJ m b j +MBK m b k +MBL m b l +M. B. m b +M.B. m b +MBM m b m +MBNA m b n a +mBo m b o +MBO m b o +MBPJ m b p j +MBP m b p +Mbre m b r e +mbr m b r +MBR m b r +MBSE m b s e +MBSI m b s i +MBS- m b s +MBS m b s +MBTI m b t i +MBT- m b t +MBT m b t +MBU's m b u's +MBX m b x +MCAF m c a f +MCAL m c a l +M.C.A. m c a +MCA- m c a +MCA m c a +MCA's m c a's +MCAS m c a s +MCB- m c b +MCB m c b +M. C. C. m c c +M.C.C. m c c +MCC m c c +MCC's m c c's +MCCs m c c's +MCCU m c c u +MCDA m c d a +MCDC m c d c +MCD m c d +MCDM m c d m +MCE m c e +MCFA m c f a +M. C. F. m c f +MCF m c f +MCFM m c f m +MCGJCW m c g j c w +MCG m c g +MCI m c i +MCJ m c j +MCLA m c l a +Mcl- m c l +MCL m c l +M. C. m c +M.C. m c +MCMC m c m c +M. C. M. m c m +MCM m c m +MCN m c n +MCO m c o +MCOT's m c o t's +MCPC m c p c +MCPI m c p i +MCP- m c p +MCP m c p +MCPON m c p o n +MCPP m c p p +MCQ m c q +MCRD m c r d +MCR m c r +MCRP m c r p +MCSBA m c s b a +MCs m c's +M.C. S. m c s +MCS- m c s +MCS m c s +MCST m c s t +MCTFS m c t f s +MCT m c t +MCTU m c t u +MCTV m c t v +MCU m c u +MCV m c v +MCVTS m c v t s +MCW m c w +mcyG m c y g +MCYO m c y o +MCZ m c z +MDAA m d a a +MDAH m d a h +mda- m d a +MDA m d a +mdb m d b +MDB m d b +MDC- m d c +MDC m d c +MDC's m d c's +MDCs m d c's +MDCT m d c t +mdDA m d d a +MD&DI m d and d i +Mde m d e +MDE m d e +MDF m d f +MDG m d g +MDGs m d g's +MDH m d h +MDHUs m d h u's +mDia m d i a +MDIB m d i b +MDic m d i c +MDI m d i +MDJT's m d j t's +MDK m d k +mdla m d l a +MDL m d l +MDMA- m d m a +MDMA m d m a +md m d +M. D. m d +M.D. m d +MD- m d +MD m d +MDM m d m +MDNA m d n a +M.D.N. m d n +MDPD's m d p d's +MDP m d p +MDPPP m d p p p +MDPS m d p s +MDPV m d p v +MDQ m d q +MDRC m d r c +mdr m d r +MDR- m d r +MDR m d r +MDT m d t +MDV m d v +MDX m d x +M. E. m e +M.E. m e +M.E.N. m e n +M. E. P. m e p +M.F.A. m f a +MFA m f a +MFB m f b +MFCC m f c c +MFC m f c +MFCS m f c s +MFD m f d +MFDs m f d's +MFe m f e +MFF m f f +MFG m f g +MFI m f i +MFJ m f j +MFK m f k +M. F. m f +M.F. m f +MFMF m f m f +MFM m f m +MFN m f n +MFNW m f n w +MFP m f p +MFR m f r +MFS m f s +MFS's m f s's +MFTBC m f t b c +MFT m f t +MFTs m f t's +Mgadla m g a d l a +Mgal m g a l +mga m g a +Mga m g a +MGA m g a +MGB m g b +Mgbo m g b o +MGCCC m g c c c +MGC m g c +MGE m g e +MGen m g e n +MGG m g g +MGH m g h +MGIMO m g i m o +MGIT m g i t +MGJH m g j h +MGK m g k +M.G.L. m g l +M&G m and g +M. G. m g +M.G. m g +M. G. M. m g m +MGM- m g m +MGM m g m +MGM's m g m's +MGMT m g m t +M&GN m and g n +MGN m g n +MGO m g o +MGP m g p +M&GR's m and g r's +Mha m h a +M.H.A. m h a +MHA m h a +MHB's m h b's +MHC- m h c +MHC m h c +MHCs m h c's +MHD m h d +MHI m h i +MHK m h k +MH&L m h and l +MHL m h l +M&H m and h +M. H. m h +M.H. m h +Mhor m h o r +Mhow m h o w +MHPL m h p l +M.H.P. m h p +MHP m h p +MHP's m h p's +MHRA m h r a +M. H. R. m h r +MHSA m h s a +MHSC m h s c +MHS m h s +MHT m h t +M.I.A. m i a +M.I.A.'s m i a's +M.I.C.M. m i c m +M.I.H. m i h +M&I m and i +M. I. m i +M.I. m i +M.I.N. m i n +M.I.R.V. m i r v +M.I.S. m i s +M.I.T. m i t +MJAHL's m j a h l's +M.J.A. m j a +MJA m j a +MJBHA m j b h a +MJB m j b +MJC m j c +M.J.F. m j f +MJF m j f +MJG m j g +MJHL m j h l +MJHL's m j h l's +M&J m and j +M. J. m j +M.J. m j +M. J. Y. m j y +MKBHD m k b h d +MKB m k b +MKDE m k d e +MKE m k e +MKF- m k f +MKFM m k f m +MKG m k g +MKK m k k +MKMF m k m f +M. K. m k +M.K. m k +MKNG m k n g +MKP m k p +MKRN m k r n +MKs m k's +MKS- m k s +MKS m k s +MKT m k t +MKTO m k t o +MKTV m k t v +MKZ's m k z's +M.L.A. m l a +MLA m l a +MLANA m l a n a +MLAs m l a's +MLB- m l b +MLB m l b +MLBPA m l b p a +MLCA m l c a +mlc m l c +M.L.C. m l c +MLC m l c +MLCs m l c's +MLD m l d +MLE m l e +MLF m l f +mlg m l g +MLG m l g +MLH m l h +MLIA m l i a +mli m l i +MLI m l i +MLIM m l i m +MLK m l k +MLK's m l k's +Mlle m l l e +MLL m l l +MLLT m l l t +M&L m and l +M. L. m l +M.L. m l +MLM m l m +MLND m l n d +M.L.N. m l n +MLN m l n +MLPH m l p h +MLP m l p +MLPs m l p's +M. L. R. m l r +MLR m l r +MLSE m l s e +M. L. S. m l s +M.L.S. m l s +MLS m l s +MLT m l t +MLW- m l w +MLW m l w +MLWS m l w s +MM&A m m and a +M. M. J. m m j +MMK m m k +MML m m l +M&M m and m +M. M. m m +M.M. m m +MMSA m m s a +M&M's m and m's +M. M. S. m m s +MMS- m m s +MMS m m s +MMST m m s t +MMTB's m m t b's +MMTS m m t s +MMTV m m t v +MMWR m m w r +MMX m m x +MNA m n a +MNBA m n b a +MNC m n c +MNCPPC m n c p p c +MNCs m n c's +MNDM m n d m +MND- m n d +MNDNR m n d n r +MNDO m n d o +MNDR m n d r +M&NF m and n f +MNF- m n f +MNF m n f +MNG m n g +MNI- m n i +MNI m n i +MNLA m n l a +MNLA's m n l a's +MNLF m n l f +MNM m n m +M. N. m n +M.N. m n +MNNA m n n a +MNN m n n +MNP m n p +MNPP m n p p +MNRG m n r g +MNSD m n s d +MNS m n s +MNTC m n t c +mnt m n t +MNT m n t +MNZM m n z m +MNZ m n z +MOBKL m o b k l +MOBK m o b k +M.O.B. m o b +M.O.D. m o d +M.O.D.O.K. m o d o k +M.O.G.U.E.R.A.'s m o g u e r a's +M. O. H. m o h +M. O. m o +M.O. m o +M.O.P. m o p +M.O.R. m o r +M.O.T. m o t +MPAA m p a a +MpA m p a +MPA m p a +MPAs m p a's +MPB m p b +MPBN m p b n +MPBu m p b u +MPCA m p c a +MPCI m p c i +MPC- m p c +MPC m p c +MPD m p d +MPE m p e +MPF m p f +MPG m p g +M.P.H. m p h +MPH m p h +MP&I m p and i +MPi m p i +MPI m p i +MPIO m p i o +MPKAB m p k a b +MPLAD m p l a d +MPLA m p l a +MPL m p l +MPLMs m p l m's +MPLR m p l r +MPLS m p l s +M&P m and p +MPM m p m +M. P. m p +M.P. m p +MPO m p o +MPPJ m p p j +M.P.P. m p p +MPP m p p +MPPSC m p p s c +MPPs m p p's +MPQC m p q c +MPR m p r +MPRO m p r o +MPRP m p r p +MPRS m p r s +MPSE m p s e +MPSF m p s f +M.P.'s m p's +MP's m p's +MPs m p's +MPS- m p s +MPS m p s +MPThe m p t h e +MPT m p t +MPTP- m p t p +Mpu m p u +MPU m p u +MPV m p v +MPW m p w +MPZ m p z +MQM m q m +M. Q. m q +Mra m r a +MRA m r a +M. R. B. m r b +MRB m r b +MRBs m r b's +MRCB m r c b +MRCK m r c k +MRCM m r c m +MRC- m r c +MRC m r c +M.R.C.S. m r c s +MRCS m r c s +MRDC m r d c +MRD m r d +mre m r e +MRF m r f +MR&LE m r and l e +MRL m r l +Mr. mister +MRM m r m +M. R. m r +M.R. m r +mRNA m r n a +mRNAs m r n a's +MRPGM m r p g m +MRP m r p +MRPS m r p s +MRR m r r +MRSA m r s a +M.R.S.C. m r s c +Mrs. misses +MRSM m r s m +MRT m r t +MRTs m r t's +MRTS m r t s +MRTT m r t t +MRU- m r u +MRU m r u +MSAA m s a a +MSAC m s a c +MSAD m s a d +MSA m s a +M.S.A.S. m s a s +MSAW m s a w +MSB m s b +MSBO m s b o +MSBSD m s b s d +MSBS- m s b s +MSCC m s c c +M.S.C. m s c +MSC- m s c +MSC m s c +msd m s d +MSD m s d +MSDN m s d n +MSDS m s d s +Mse m s e +MSFC m s f c +M.S.F. m s f +MSF m s f +MSFT m s f t +MSG m s g +MSI m s i +MSJ m s j +MSK m s k +MSL m s l +M&S m and s +Ms. miss +MSML m s m l +M.S.M. m s m +MSM m s m +M. S. m s +M.S. m s +MSNBC m s n b c +MSNBC's m s n b c's +MSN m s n +MSOB m s o b +MSOM m s o m +M.S.P. m s p +MSP m s p +MSPs m s p's +MSR m s r +MSRP m s r p +MSRTC m s r t c +MSSH m s s h +MSSK m s s k +MSSMLP m s s m l p +MSSM m s s m +MSS m s s +MSTA m s t a +mst m s t +MST m s t +MSTS m s t s +MSU's m s u's +MSVCC m s v c c +MSV m s v +MSVU m s v u +MSX m s x +MSY m s y +MSZ- m s z +MSZP m s z p +MTAC m t a c +Mta m t a +M.T.A. m t a +MTA- m t a +MTA m t a +MTA's m t a's +MTAs m t a's +MTBE m t b e +MTBI m t b i +MTB- m t b +MTB m t b +MTC m t c +MTDB m t d b +M.T.D. m t d +MTD m t d +MTD's m t d's +MTG m t g +MTHFR m t h f r +MTHL m t h l +MTIC m t i c +MTI m t i +MTIP- m t i p +MTJ m t j +MTK m t k +MTKO m t k o +M&T m and t +MTM m t m +MTMR m t m r +M. T. m t +M.T. m t +Mtor m t o r +MTOSI m t o s i +M.T.O.W. m t o w +MTPA m t p a +MTPC m t p c +MTP m t p +MTQ m t q +MTRCB m t r c b +MTRC m t r c +MTRJ m t r j +MTR- m t r +MTR m t r +M. T. S. m t s +MTS- m t s +MTS m t s +MTSU m t s u +MTTM m t t m +MTT m t t +MTU m t u +MTV- m t v +MTV m t v +MTVR m t v r +MTV's m t v's +mtvU m t v u +MTVu m t v u +MTVU m t v u +MTY m t y +MTZ- m t z +MTZ m t z +MUKW m u k w +M. U. m u +MUSL m u s l +M.U.s m u's +M. U.S. m u s +MVA m v a +M. V. C. m v c +M.V.C. m v c +MVC m v c +M. V. D. m v d +MVD m v d +MVH m v h +MVK m v k +MVL- m v l +M.V.M. m v m +MVM m v m +MVMs m v m's +M. V. m v +M.V. m v +M.V.O. m v o +MVO m v o +MVP- m v p +MVP m v p +MVPs m v p's +MVPS m v p s +M.V. P.T. m v p t +MVs m v's +MVS m v s +MVSN m v s n +MVSR m v s r +MVT m v t +MVV m v v +M.W.A.M. m w a m +MWA m w a +MWC m w c +MWe m w e +M. W. E. m w e +MWF m w f +MWHL m w h l +MWIFF m w i f f +MWI m w i +MWJCHL m w j c h l +MWJHL m w j h l +M. W. J. m w j +M. W. m w +M.W. m w +M. X. m x +myb m y b +myc m y c +Myc m y c +MYC m y c +MYCN- m y c n +MYCN m y c n +MyDD m y d d +M. Y. m y +M.Y. m y +MySQL m y s q l +MYSQL m y s q l +MZC m z c +MZH m z h +MZM m z m +M. Z. m z +MZP m z p +MZT m z t +NAACCR n a a c c r +NAACL n a a c l +NAAC n a a c +N.A.A.C.P. n a a c p +NAACP n a a c p +NAACP's n a a c p's +N.A.C.L. n a c l +NAFBL n a f b l +NAFC n a f c +NAFC's n a f c's +NAFH n a f h +NAFI n a f i +N. A. n a +N.A. n a +N.A.S.A. n a s a +NBADL n b a d l +N.B.A. n b a +NBA- n b a +NBA n b a +N.B.A.'s n b a's +NBA's n b a's +NBAs n b a's +NBCC n b c c +NBC- n b c +NBC n b c +NBC's n b c's +NBCSN n b c s n +NBDL n b d l +NBD n b d +NBFA's n b f a's +NBF n b f +NBG n b g +NBI n b i +NBK n b k +NBL n b l +N. B. n b +N.B. n b +N.B.N. n b n +NBN n b n +NBP n b p +NBQ n b q +NBR n b r +NBR's n b r's +NBSK n b s k +nbs n b s +NBS n b s +NBSP n b s p +NBTE n b t e +NBT n b t +NBTwo n b t w o +NBTY n b t y +NBW n b w +NCAAs n c a a's +NCAM n c a m +NCA- n c a +NCA n c a +NCAP n c a p +NCBC n c b c +NCBI- n c b i +NCBI n c b i +NCCC n c c c +NCCF n c c f +NCCMH n c c m h +NCC n c c +NCCP n c c p +NCCs n c c's +NCCS n c c s +NCCU n c c u +NCDC n c d c +NCDD's n c d d's +NCD n c d +NCES n c e s +NCFA n c f a +N.C.F. n c f +NCF n c f +NCG n c g +NCGS n c g s +NCHC n c h c +NCH n c h +NCHU n c h u +NCID n c i d +NCI n c i +NCIS n c i s +NCIS's n c i s's +NCKU n c k u +NCLB n c l b +NCLC n c l c +NCL n c l +NCLR n c l r +NCLT n c l t +NCMA n c m a +NCMC n c m c +NCMM n c m m +NCM n c m +NCMP n c m p +NCNB n c n b +N. C. n c +N.C. n c +N.C.O. n c o +NCO n c o +NCO's n c o's +NCOs n c o's +NCPA n c p a +NCPC n c p c +NCP n c p +NCPO n c p o +NCRC n c r c +NCR n c r +NCSU n c s u +NCSY n c s y +NCTA n c t a +NCTC n c t c +NCTE n c t e +NCT n c t +NCUA n c u a +NCVO n c v o +NCVS n c v s +NCVT n c v t +NCWM n c w m +NCW n c w +NDA n d a +NDE n d e +NDEP n d e p +NDFB n d f b +NDFT n d f t +N.D.G. n d g +N. D. n d +N.D. n d +N.D.N.Y. n d n y +NDPH n d p h +NDPK n d p k +N.D.P. n d p +NDP- n d p +NDP n d p +NDP's n d p's +NDRC n d r c +NDRE n d r e +NDRF n d r f +N.E.A.R. n e a r +NECW n e c w +N. E. n e +N.E. n e +N.E.R.D. n e r d +NERFU n e r f u +ner n e r +N.E.R. n e r +NER n e r +N.E.W.S.T. n e w s t +NFA n f a +NFATc n f a t c +NFAT n f a t +NFB n f b +NFB's n f b's +NFC n f c +NFC's n f c's +NFD n f d +NFEA n f e a +nfed n f e d +NFFC n f f c +NFF n f f +NFH n f h +NFHS n f h s +NFIB n f i b +NFI n f i +NFISD n f i s d +N.F.L. n f l +NFL n f l +NFLPA n f l p a +NFL's n f l's +N. F. n f +N.F. n f +NFO n f o +NFPA n f p a +NFPF n f p f +NFP n f p +NFPW n f p w +NFRS n f r s +NFS n f s +NGC n g c +Nge n g e +NGF- n g f +NGF n g f +N.G.L. n g l +N. G. n g +N.G. n g +NGO n g o +NGO's n g o's +NGOs n g o's +NGRC n g r c +N.G.R. n g r +NGR n g r +NGRR n g r r +NGS- n g s +NGS n g s +NGSS n g s s +NGST n g s t +NGTC n g t c +Nha n h a +NHA n h a +NHCEs n h c e's +NHC n h c +NHCP n h c p +NHCs n h c's +NHD n h d +NHHD n h h d +NHHI n h h i +nhi n h i +NHI n h i +NHK n h k +NHK's n h k's +N.H.L. n h l +NHM&W n h m and w +N. H. n h +N.H. n h +NHP n h p +NHPs n h p's +NHPS n h p s +NHRA n h r a +NHRC n h r c +NHSCT n h s c t +NHS- n h s +NHS n h s +NHST n h s t +NHW n h w +N.I.B.B.L.E. n i b b l e +NIBC n i b c +NIBR n i b r +N. I. n i +N.I. n i +N'I n i +N.J.A.C. n j a c +N. J. A. n j a +NJA n j a +N. J. C. n j c +NJC n j c +NJIT n j i t +N. J. M. n j m +NJM n j m +N. J. n j +N.J. n j +NJN n j n +Njoo n j o o +NJPS n j p s +NJPW n j p w +N.J.S.A. n j s a +N.J.'s n j's +NKL n k l +N. K. n k +N.K. n k +N. K. N. n k n +NKP n k p +NKP's n k p's +NKR n k r +NKT n k t +NKU n k u +nkvd n k v d +NKVD n k v d +NKVM n k v m +nkv n k v +NKX n k x +NLAES n l a e s +NLAI n l a i +NLA n l a +NLB n l b +NLCF n l c f +NLC n l c +NLCS n l c s +NLDC n l d c +NLD n l d +NLD's n l d's +NLDS n l d s +N. L. n l +N.L. n l +NLP n l p +NLRB n l r b +NLRP n l r p +NLRs n l r's +NLW n l w +NLX n l x +NMAA n m a a +NMA n m a +NMBS n m b s +NMBU n m b u +NMCA n m c a +NMCB n m c b +NMC n m c +N. M. n m +N.M. n m +NMP n m p +NMR n m r +NMRW n m r w +NMT n m t +NMU n m u +NMV n m v +NMW n m w +NNCL n n c l +NNFL n n f l +NNF n n f +N. N. n n +N.N. n n +NNSA n n s a +NNSS n n s s +NNSU n n s u +NNSW n n s w +NNW- n n w +No. number +N.O.R. n o r +NPA- n p a +NPA n p a +NPBD n p b d +NPBL n p b l +NPB n p b +NPC n p c +NPCs n p c's +NPDC n p d c +NPD n p d +NPF- n p f +NPF n p f +NPF's n p f's +NPGL n p g l +NPG n p g +NPGS n p g s +NPH n p h +NPHS n p h s +NPIM n p i m +NPI n p i +NPJ n p j +NPK n p k +NPL n p l +NPMA n p m a +NPM n p m +NPMs n p m's +N. P. n p +N.P. n p +NP- n p +NP n p +NPO n p o +NPP n p p +NPP's n p p's +NPRL n p r l +NPRN n p r n +NPR- n p r +NPR n p r +NPRR n p r r +NPR's n p r's +NPRSO n p r s o +NPSC n p s c +NPSG n p s g +NPSL n p s l +NP's n p's +NPs n p's +NPS n p s +NPTI n p t i +NPT n p t +NPV n p v +NPWE n p w e +NPWS n p w s +NPY n p y +NQAI n q a i +NQEA n q e a +NQHS n q h s +N. Q. n q +nri n r i +Nri n r i +NRI n r i +NRIs n r i's +NRJ n r j +NRJs n r j's +NRK n r k +NRLA n r l a +NRL n r l +NRL's n r l's +NRMA n r m a +NRM n r m +NRN n r n +nr n r +N. R. n r +N.R. n r +NSA- n s a +NSA n s a +NSAP n s a p +NSA's n s a's +NSB- n s b +NSB n s b +NSC- n s c +NSC n s c +NSD n s d +NSE n s e +NSFE n s f e +NSF n s f +NSFW n s f w +NSG n s g +NSI n s i +NSL n s l +NSM n s m +NSN n s n +N. S. n s +N.S. n s +NSO n s o +NSOs n s o's +NSPCC n s p c c +NSPC n s p c +NSP n s p +NSRI n s r i +NSRL n s r l +NSR n s r +NSSCD n s s c d +NSSDC n s s d c +NSS- n s s +NSS n s s +NSTA n s t a +NST n s t +NSU n s u +NSV n s v +NSVT n s v t +NSWBLF n s w b l f +NSWC- n s w c +NSWC n s w c +NSWCV n s w c v +NSWEC n s w e c +NSWGR n s w g r +N.S.W. n s w +NSW- n s w +NSW n s w +Nta n t a +NTA n t a +NTA's n t a's +ntb n t b +NTB n t b +NTC n t c +NTD n t d +NTFL n t f l +NTF n t f +NTI n t i +NTKF n t k f +NTL n t l +N. T. n t +N.T. n t +NTRK- n t r k +NTR n t r +NTR's n t r's +NTSB n t s b +NTSB's n t s b's +NTSC n t s c +NTS n t s +NTTF n t t f +NTT n t t +NTTR n t t r +NTUC n t u c +NTU n t u +NTVB n t v b +NTV n t v +ntw n t w +N.U. n u +NVA n v a +NVAO n v a o +NVC n v c +NVDA n v d a +NvDA's n v d a's +NVFC n v f c +NVGOP n v g o p +NVHJ's n v h j's +NVH n v h +NVIDIA's n v i d i a's +NVIDIA n v i d i a +NVI n v i +NVL n v l +N. V. n v +N.V. n v +NWAFU n w a f u +nwa n w a +N.W.A. n w a +NWA- n w a +NWA n w a +NWA's n w a's +N.W.F.P. n w f p +NWFP n w f p +NWMP n w m p +N&W n and w +NWN n w n +N. W. n w +N.W. n w +NWPD n w p d +N.W.P. n w p +NWP n w p +NWS n w s +N.W.T. n w t +NWT n w t +nyc n y c +N.Y.C. n y c +NYC- n y c +NYC n y c +NYC's n y c's +N.Y. G.O.P. n y g o p +N.Y.L. n y l +N. Y. n y +N.Y. n y +N.Y.P.D. n y p d +NYPD n y p d +NYPD's n y p d's +NYPL n y p l +N.Y.P. n y p +NY's n y's +N.Y.S. n y s +N.Y.S.V. n y s v +N.Y.U. n y u +NYU- n y u +NYU n y u +NYU's n y u's +NYW&B- n y w and b +NYYC n y y c +NZAID n z a i d +NZAOD n z a o d +NZCA n z c a +NZDT n z d t +NZEF n z e f +NZETC n z e t c +NZFC n z f c +NZF n z f +N'Zif n z i f +NZiK n z i k +N'Zi n z i +NZIV n z i v +NZL n z l +NZLP n z l p +N.Z. n z +NZPA n z p a +NZR- n z r +NZSAS n z s a s +NZ's n z's +NZS n z s +NZTA n z t a +N.Z.W.P.W. n z w p w +O.A.C. o a c +O. A. o a +O.A. o a +O.A.R. o a r +O.A.S. o a s +O.B.E. o b e +obl o b l +Oblt o b l t +O.B. o b +O. C. o c +O.C. o c +Octl o c t l +O.C.T. o c t +O. D. o d +O.E. o e +O.F.M. o f m +O.F. o f +ofr o f r +O.F.R. o f r +OFR o f r +O.F.T.B. o f t b +OGL o g l +O&G o and g +O. G. o g +O.G. o g +O.G.S. o g s +OGS o g s +O.H.A. o h a +O. H. o h +O.H. o h +OHV o h v +O. I. o i +O.I. o i +O. J. o j +O.J. o j +OJ o j +OJSC o j s c +OKBM o k b m +OKB- o k b +OKB o k b +OKC o k c +OK'd o k d +OKd o k d +OKD o k d +oke o k e +Oke o k e +OKH o k h +O.K.I. o k i +OKK- o k k +OKK o k k +OKM o k m +O&K o and k +O. K. o k +O.K. o k +O.L.F.A.L. o l f a l +O. L. K. o l k +O. L. o l +O.L. o l +Olo o l o +Olov o l o v +OLPC o l p c +OLPH o l p h +ols o l s +OLs o l's +OLS o l s +OLSR o l s r +OLTL o l t l +OLVT o l v t +olvwm o l v w m +olwm o l w m +OMB o m b +OMC o m c +OMCS o m c s +OMD o m d +O&MFL o and m f l +OMF o m f +omg o m g +OMG o m g +OMGs o m g's +O. M. o m +O.M. o m +OmOm o m o m +Om's o m's +OMTP o m t p +OMVG o m v g +OMW o m w +OMX o m x +ONCHR o n c h r +ONC o n c +oncu o n c u +ONDH o n d h +ond o n d +onf o n f +ONF o n f +ONGC o n g c +ONIR o n i r +ONJSC o n j s c +ONM o n m +ONMR o n m r +O. N. o n +O.N. o n +O&O o and o +OOOA o o o a +O. O. o o +O.O. o o +OO o o +OOO's o o o's +Oop o o p +OOP o o p +OOPSLA o o p s l a +oor o o r +oose o o s e +Oo's o o's +OOTP o o t p +Ootw o o t w +OOUR o o u r +OPAC o p a c +OPAG o p a g +O.P.A. o p a +OPBF o p b f +OPB o p b +OPC o p c +OPCS o p c s +OPCW o p c w +OPD- o p d +OPD o p d +O.P.I. o p i +OPIRG o p i r g +OPJHL o p j h l +OPLC o p l c +OPL o p l +OPMB o p m b +OPM o p m +OPMs o p m's +OPN o p n +O. P. o p +O.P. o p +opr o p r +Opr o p r +OPR o p r +O.R.C. o r c +ORECA o r e c a +ORFs o r f's +ORMO o r m o +ORNL o r n l +OR&N o r and n +O. R. o r +osaa o s a a +OSAA o s a a +osa o s a +Osa o s a +O.S.A. o s a +OSA o s a +O.S.B. o s b +OSB o s b +OSBs o s b's +O.S.N. o s n +OSN o s n +O. S. o s +O.S. o s +OSP o s p +OSR o s r +OSTM o s t m +O.S.T. o s t +OSTP o s t p +OSTs o s t's +osv o s v +OSV- o s v +OSV o s v +OTB o t b +OTC o t c +OTJ o t j +OTL o t l +OTMH o t m h +OTO o t o +O.T. o t +OTP o t p +OTR o t r +OTs o t's +OTS- o t s +OTS o t s +Otu o t u +OTU o t u +otv o t v +O.U. o u +OU- o u +OU o u +OUP o u p +OUSA o u s a +OVA's o v a's +ovca o v c a +Ovca o v c a +OVC o v c +OVC's o v c's +Ovda o v d a +OVF o v f +OVM o v m +Ovo o v o +OVO o v o +ov o v +Ov o v +O. V. o v +O.V. o v +OV- o v +OV o v +OVP o v p +OVS o v s +OVT o v t +OVW o v w +O.W.A. o w a +OWBT o w b t +OWCs o w c's +OWGR o w g r +OWHA's o w h a's +OWIU o w i u +OWM o w m +O. W. o w +O.W. o w +P. A. C. p a c +P. A. J. p a j +P.A.J. p a j +P. A. M. p a m +P.A.M. p a m +pa p a +P. A. p a +P.A. p a +PA- p a +PA p a +PAQ p a q +P.A.R. p a r +P. A. S. p a s +PBAA p b a a +PBA- p b a +PBA p b a +PBA's p b a's +PBB p b b +PBCC p b c c +pBCE p b c e +pbc p b c +P.B.C. p b c +PBF p b f +PBG p b g +PBIL p b i l +PBI p b i +PBK p b k +PBL p b l +PBM p b m +PBMR p b m r +PBO p b o +pb p b +pB- p b +P. B. p b +P.B. p b +PB- p b +PB p b +PBP p b p +PBR p b r +PBS p b s +PBT p b t +PBX p b x +PBY- p b y +PBY p b y +P. C. A. p c a +PCA p c a +PCB p c b +PCB's p c b's +PCBs p c b's +PCBS p c b s +PCC- p c c +PCC p c c +PCCS p c c s +PCCW p c c w +PCCW's p c c w's +PCDDs p c d d's +PCDHB p c d h b +PCDH p c d h +PCD p c d +PCE p c e +PCeU p c e u +PCFCL p c f c l +PCFL p c f l +PCF p c f +PCFS p c f s +PCG p c g +PCGS p c g s +PCHA p c h a +PCHA's p c h a's +P.C.H. p c h +PCH p c h +PCHR p c h r +PCIe p c i e +pci p c i +PCI- p c i +PCI p c i +PCJHL p c j h l +P. C. J. p c j +PCJ p c j +PCK p c k +PCLM p c l m +pcl p c l +PCL p c l +PCMCIA p c m c i a +PCMNO p c m n o +PCM p c m +PCMS p c m s +PCNA- p c n a +PCNA p c n a +PCN p c n +PcoA p c o a +PCOE p c o e +PCO p c o +pc p c +P. C. p c +P.C. p c +PC- p c +PC p c +PCPFL p c p f l +PCP p c p +PCPV p c p v +PCRC p c r c +PCRE p c r e +PCRev p c r e v +PCRM p c r m +PCR- p c r +PCR p c r +PCSK p c s k +PCSOM p c s o m +pcs p c s +P.C.s p c's +PC's p c's +PCTL p c t l +PCT p c t +PCTV p c t v +PCU p c u +PCU's p c u's +PCW p c w +PCX p c x +pDAB p d a b +PDAB p d a b +PDA p d a +PDAs p d a's +PDBML p d b m l +PDB p d b +pdbp p d b p +PDBsum p d b s u m +PDCI p d c i +PDC- p d c +PDC p d c +PDCPD p d c p d +PDC's p d c's +PDCs p d c's +PD&D p d and d +PDD- p d d +PDD p d d +PDE p d e +PDEs p d e's +PDES p d e s +pdf p d f +PDF p d f +PDFs p d f's +PDFT p d f t +PDGF- p d g f +PDGF p d g f +PDG p d g +PDHJ p d h j +PDH p d h +PDI- p d i +PDI p d i +PDJ p d j +PDK p d k +PDL p d l +PDN p d n +PDO p d o +P&D p and d +P. D. p d +P.D. p d +PD- p d +PD p d +PDP- p d p +PDP p d p +PDP's p d p's +PDPs p d p's +P. D. Q. p d q +PDR p d r +PDSA p d s a +PDSI p d s i +PD's p d's +PDs p d's +P.D.S. p d s +PDS p d s +PDSP p d s p +PDT p d t +P.E.I. p e i +P.E.I.'s p e i's +P.E.N. p e n +P. E. p e +P.E. p e +P.F.L. p f l +PF's p f's +PFs p f's +PFS p f s +PFU p f u +PFV p f v +PFW&C p f w and c +P.G.A. p g a +PGA p g a +PGM p g m +PGMs p g m's +PG&N p g and n +P&G p and g +P. G. p g +P.G. p g +P.G.T. p g t +PHD p h d +P. H. G. p h g +PHH p h h +PHHS p h h s +PHLF p h l f +Phlo p h l o +PHL p h l +PHN p h n +P&H p and h +P. H. p h +P.H. p h +php p h p +PHP- p h p +PHP p h p +PHPs p h p's +P.I.D.E. p i d e +P&I p and i +P. I. p i +P.I. p i +P. I. W. p i w +Piz p i z +PJB p j b +PJC p j c +PJD p j d +P. J. F. p j f +P.J.K. p j k +pj p j +P. J. p j +P.J. p j +PJ- p j +PJ p j +PJ's p j's +PJs p j's +P.J.T. p j t +P. K. p k +P.K. p k +PK- p k +PK p k +PLCC p l c c +plc p l c +PLC- p l c +PLC p l c +PLC's p l c's +PLCs p l c's +P.L.O. p l o +P. L. p l +P.L. p l +PLP p l p +PLX- p l x +PLX p l x +P. M. A. p m a +PMC- p m c +PMC p m c +PMG p m g +PML- p m l +PML p m l +PMLP p m l p +PMMA p m m a +PMMoV p m m o v +PMM p m m +PMMS p m m s +PMMT p m m t +PMOI p m o i +PMO p m o +PMPC p m p c +PMPF p m p f +pm p m +P. M. p m +P.M. p m +PM- p m +PM p m +PMP p m p +PMP's p m p's +PMRN p m r n +PMRO p m r o +PMR p m r +PMRR p m r r +PMSE p m s e +PMSL p m s l +PMSM p m s m +PM's p m's +PMs p m's +P. M. S. p m s +PMS p m s +PMSSY p m s s y +PMT p m t +PMVY p m v y +PMWA p m w a +PNaCl p n a c l +pna p n a +PNA p n a +Pnau p n a u +pnb p n b +PNB p n b +PNC p n c +PNE p n e +PNE's p n e's +PNEs p n e's +PNETs p n e t's +pneus p n e u s +PNFA p n f a +PNF p n f +PNGIA p n g i a +png p n g +PNG p n g +PNH p n h +PNI p n i +PNL p n l +PNL's p n l's +PNMT p n m t +PNNL p n n l +PNNs p n n's +PNoy p n o y +pn p n +P. N. p n +P.N. p n +PN- p n +PN p n +PNP p n p +P.N.R.A. p n r a +P.O.D. p o d +P.O.D.'s p o d's +P&O p and o +P. O. p o +P.O. p o +P.O.S. p o s +P.O.V. p o v +P.O.W. p o w +P.O.W.'s p o w's +PPACA p p a c a +PPA p p a +PPBS p p b s +PPC p p c +PPD- p p d +PPD p p d +PPE p p e +PPF p p f +PPG p p g +P&PH p and p h +PPH p p h +PPi p p i +P.P.I. p p i +PPI p p i +PPJ p p j +PPKM p p k m +PPK p p k +PPL p p l +PPM p p m +PPNB p p n b +PPN p p n +PPO p p o +P. P. p p +P.P. p p +PPRP p p r p +PPR p p r +PP's p p's +PPS p p s +PPT p p t +PPTV p p t v +PPV p p v +PPVs p p v's +PPy p p y +PQBP- p q b p +PQDT p q d t +PQM- p q m +PQ- p q +PQ p q +PQQ p q q +PQ's p q's +PQS p q s +PRB- p r b +PRB p r b +PRCA p r c a +PRC- p r c +PRC p r c +PRCS p r c s +PRDM p r d m +PRD p r d +PRD's p r d's +P.R.I.M.E. p r i m e +pr p r +P. R. p r +P.R. p r +PR- p r +PR p r +P.S.A. p s a +P.S.C. p s c +P.S.K. p s k +PSP- p s p +PSP p s p +P. s p's +P.'s p's +Ps- p s +Ps p's +P. S. p s +P.S. p s +PSQL p s q l +PSR p s r +PSRU p s r u +PSSA- p s s a +PSSA p s s a +PSSA's p s s a's +PSSAs p s s a's +PSSI p s s i +PSS p s s +PSS's p s s's +PST p s t +PSU p s u +PSV p s v +PSW p s w +pTA p t a +PTA p t a +PTA's p t a's +PTAs p t a's +PTB p t b +PTBT p t b t +PTCH p t c h +PTC p t c +Pte p t e +PTEs p t e's +PTFE p t f e +ptf p t f +PTF p t f +PTH p t h +Ptie p t i e +PTI- p t i +PTI p t i +PTI's p t i's +PTK p t k +PTLLS p t l l s +PTL p t l +PTLs p t l's +PTNNT p t n n t +Ptol p t o l +PTO p t o +PTPN p t p n +PTP p t p +pt p t +P. T. p t +P.T. p t +PT- p t +PT p t +PTSD p t s d +PTSE p t s e +PTS- p t s +PTS p t s +PTTGC p t t g c +PTTG p t t g +PTTOW p t t o w +PTT p t t +PTTs p t t's +PTUN p t u n +PTV p t v +PTV's p t v's +P'Twa p t w a +PTY p t y +P.U.F. p u f +Puiu p u i u +PUKKE p u k k e +PUK p u k +Pul p u l +P. U. M. p u m +PVA p v a +PVAs p v a's +PVB p v b +P.V.C. p v c +PVC p v c +PvdA p v d a +PVEM p v e m +pve p v e +PVH p v h +P.V.L. p v l +PVL p v l +PVN p v n +PVO p v o +PVP p v p +Pv- p v +P. V. p v +P.V. p v +PV- p v +PV p v +pvr p v r +PVR p v r +PVRS p v r s +PVs p v's +PVS p v s +PVU p v u +PWA p w a +PWB p w b +PWC p w c +PWD p w d +PWF p w f +P. W. G. p w g +P.W.G. p w g +PWG p w g +PWI p w i +pwll p w l l +PWM p w m +P.W.O. p w o +P&W p and w +P. W. p w +P.W. p w +PW- p w +PW p w +Pyi p y i +PYI p y i +Pyk p y k +PYK p y k +Pyl p y l +PyL p y l +PYP p y p +P. Y. p y +Pyu p y u +Pyw p y w +Pyx p y x +PZB p z b +PZL- p z l +PZL p z l +PZP p z p +P. Z. p z +Q. A. P. q a p +QAP q a p +Q. A. q a +QA q a +Q&A q and a +Q&A's q and a's +Q&As q and a's +QbA q b a +QBE q b e +QBH q b h +QBL q b l +Q.B. q b +QB q b +Q.C.B. q c b +QCD q c d +Q.C. q c +QD q d +QDS q d s +QEA q e a +QEC q e c +Q.E.D. q e d +QED q e d +Q.E.H. q e h +Q. E. q e +Q.H.C. q h c +QHP q h p +Q.H. q h +Q.I. q i +Q. J. q j +QJ q j +Q. J. R. q j r +Q. N. q n +QPF q p f +QPFS q p f s +QPM q p m +QPO q p o +QPOs q p o's +QPP q p p +QP q p +Q.P.R. q p r +QPR q p r +QQ q q +QRL q r l +QRL's q r l's +QRNA q r n a +QROPS q r o p s +QRP q r p +QR q r +QRS q r s +QRT q r t +QRV q r v +QSAR q s a r +QSES q s e s +QSI q s i +QSK- q s k +QSL q s l +QSM q s m +QSO q s o +QST q s t +QSY q s y +QSZ- q s z +QTE- q t e +QVC q v c +QVD q v d +QVM q v m +Q.V. q v +R.A.B. r a b +R. A. E. r a e +R.A.F. r a f +R.A.J. r a j +R. A. r a +R.A. r a +R&A r and a +RAV r a v +R&AW r and a w +RAZR r a z r +RBAC r b a c +RbAg r b a g +RBA r b a +RBBP r b b p +RBC r b c +RBCs r b c's +RBD r b d +RBE r b e +RBGE r b g e +RBG r b g +RBI r b i +RBIs r b i's +R. B. J. r b j +R.B.J. r b j +RBK- r b k +RBK r b k +R.B.L. r b l +RBL r b l +RBMG r b m g +RBMG's r b m g's +RBM r b m +RBMY r b m y +RBO r b o +RBP r b p +R&B- r and b +R&B r and b +R. B. r b +R.B. r b +RBR r b r +RBS r b s +RBTCO's r b t c o's +RBTH r b t h +RBWH r b w h +RBW r b w +RBX r b x +RBZ r b z +RCAC r c a c +RCAF r c a f +rca r c a +R.C.A. r c a +RCA- r c a +RCA r c a +RCA's r c a's +RCBC r c b c +RCB r c b +RCCA r c c a +RCCD r c c d +RCCL r c c l +RCCM r c c m +RCC r c c +RCI r c i +R. C. K. r c k +RCL r c l +R.C.M. r c m +RCM r c m +RCN r c n +RCPO r c p o +RCP r c p +RCPT r c p t +R&C r and c +rc r c +R. C. r c +R.C. r c +RCRD r c r d +RCR- r c r +RCR r c r +RCSB r c s b +RCS r c s +RCT r c t +RCTs r c t's +RCTS r c t s +RCTV r c t v +RCV r c v +RCW r c w +RDB r d b +RDC- r d c +RDC r d c +RDD r d d +RdE r d e +RDFC r d f c +RDF r d f +R.D.G. r d g +R&D- r and d +R&D r and d +R. D. r d +R.D. r d +R.E.A. r e a +R. E. B. r e b +R.E.D. r e d +R.E.L. r e l +R.E.M. r e m +R.E.M.'s r e m's +R.E.O. r e o +REPL r e p l +REPLs r e p l's +R.E.P. r e p +R. E. r e +R.E. r e +R.E.R. r e r +R.F.A. r f a +R.F.C. r f c +R. F. r f +R.F. r f +R. F.S. r f s +R.F.W. r f w +RFW r f w +RFX r f x +RGB r g b +RGD r g d +RGF r g f +RGG r g g +RGI r g i +RGK r g k +RGMA r g m a +RGNL r g n l +RGPH r g p h +R&G r and g +R. G. r g +R.G. r g +RGR r g r +RGS- r g s +RGS r g s +RGS's r g s's +RGTP r g t p +RGU r g u +RGV r g v +RGX r g x +R. H. C. r h c +RHC r h c +RHD r h d +RHHF r h h f +RHH r h h +RHHS r h h s +RHIBs r h i b's +RHIC r h i c +RHI r h i +Rhiw r h i w +RHK r h k +rhl r h l +RHM r h m +RhoG r h o g +Rho's r h o's +RHP r h p +R. H. r h +R.H. r h +R.H.S.J. r h s j +R. I. C. r i c +R.I.C. r i c +R. I. P. r i p +R.I.P. r i p +R&I r and i +R. I. r i +R.I. r i +riu r i u +Riu r i u +rivs r i v s +Rivu r i v u +RIXS r i x s +Rixt r i x t +RJB r j b +RJD- r j d +RJD r j d +RJE r j e +rjf r j f +RJHS r j h s +R. J. J. r j j +RJJ r j j +RJL r j l +RJN r j n +R. J. r j +R.J. r j +RJ r j +RJR r j r +RJ's r j's +RJTD r j t d +R. K. B. r k b +RKD r k d +RKG- r k g +RKI r k i +RKKA r k k a +RKL r k l +R.K.O. r k o +RKO- r k o +RKO r k o +RKO's r k o's +R. K. r k +R.K. r k +RLM r l m +RLM's r l m's +R. L. r l +R.L. r l +RL r l +RLS- r l s +RLV r l v +rly r l y +RMAF r m a f +RMAG r m a g +RMA r m a +RMB r m b +RMCH r m c h +RMCL r m c l +R.M.C. r m c +RMC r m c +RMDs r m d's +R.M.E.S. r m e s +RMFL r m f l +R.M.F. r m f +RMF r m f +RMG r m g +R. M. H. r m h +RMI r m i +RMIT r m i t +RMIT's r m i t's +RMJM r m j m +R.M.K. r m k +RMK r m k +R. M. L. r m l +RML r m l +RMLs r m l's +RMM r m m +RMP r m p +rm r m +R. M. r m +R.M. r m +RM r m +RMR r m r +RMRS r m r s +RMSD r m s d +RMSDs r m s d's +RMSE r m s e +RMs r m's +RMS r m s +R. M. W. r m w +R.M.W. r m w +RMW r m w +RMX r m x +RNAO r n a o +rna r n a +RNA- r n a +RNA r n a +RNA's r n a's +RNAs r n a's +RNAS r n a s +RNC r n c +RND r n d +RNE r n e +Rnet r n e t +RNF r n f +RNG r n g +Rnic r n i c +RNJD r n j d +RNK r n k +RNLAF r n l a f +RNLI r n l i +RNLI's r n l i's +RNNs r n n's +RNOH r n o h +RNP r n p +RNPs r n p's +RNPS r n p s +R. N. r n +R.N. r n +RN- r n +RN r n +R.N.R. r n r +RNZ r n z +R.O.C.K. r o c k +R.O.C. r o c +ROKMC r o k m c +ROKN r o k n +rOmpB r o m p b +Rooi r o o i +ROP r o p +R. O. r o +R.O. r o +ROVs r o v's +Roxb r o x b +Roxx r o x x +RPA r p a +RPAYC r p a y c +RPB r p b +RP&C r p and c +rpc r p c +RPC r p c +Rpe r p e +RPE r p e +RPF r p f +RPGA r p g a +rpg r p g +RPG- r p g +RPG r p g +RPG's r p g's +RPGs r p g's +RPI r p i +RPI's r p i's +RPK r p k +RPL r p l +rpm r p m +RPM r p m +RPMS r p m s +RPO r p o +RPP r p p +R. P. r p +R.P. r p +RP- r p +RP r p +RPs r p's +R.P.S. r p s +RPS r p s +RPT r p t +RPVE r p v e +RPV r p v +R. Q. r q +R.Q. r q +RQ- r q +RQW r q w +rra r r a +RRA r r a +RRC r r c +RRDE r r d e +rrd r r d +Rreli r r e l i +rre r r e +RRE r r e +RRG r r g +R. R. H. r r h +RRH r r h +RRKM r r k m +R.R.K. r r k +rRNA r r n a +rRNAs r r n a's +RRN r r n +RRP r r p +R&R r and r +rr r r +R. R. r r +R.R. r r +RR- r r +RR r r +RRR r r r +RRSR r r s r +RRS r r s +RSAF r s a f +R. S. A. r s a +RSA r s a +RSCG r s c g +RSCJ r s c j +R.S.C. r s c +RSC r s c +RSCS r s c s +RSD r s d +RSE r s e +RSF r s f +RSFSR r s f s r +RSG- r s g +RSG r s g +RSHA r s h a +rsh r s h +RSHS r s h s +RSICC r s i c c +RSID r s i d +RSI r s i +Rsis r s i's +RSI's r s i's +RSIS r s i s +RSK r s k +RSL r s l +rsly r s l y +RsmA r s m a +RSMC r s m c +RSMI r s m i +RSML r s m l +RSM r s m +RSM's r s m's +RSNO r s n o +RSN r s n +RSNZ r s n z +R. S. O. r s o +RSO r s o +RSPB r s p b +RSPCA r s p c a +RSpec r s p e c +RSPK r s p k +RSPO r s p o +R. S. P. r s p +RSP r s p +RSR r s r +rs r s +Rs r's +R. S. r s +R.S. r s +RS- r s +RS r s +RSSI r s s i +R. S. S. r s s +RSS- r s s +RSS r s s +RSSSF r s s s f +RSTC r s t c +RST r s t +RSu r s u +R.S.U. r s u +RSU r s u +RSU's r s u's +RSVP r s v p +RSV r s v +RSX r s x +RTAFB r t a f b +RTA r t a +RTBF r t b f +rtb r t b +RTB r t b +RTCG's r t c g's +RTC r t c +RTCs r t c's +RTD r t d +RTEC r t e c +RTEjr r t e j r +RTeOR r t e o r +RTE r t e +RTE's r t e's +RTFB r t f b +RTFM r t f m +RTF r t f +RTF's r t f's +RTFS r t f s +RTHK r t h k +RTHK's r t h k's +RTHL r t h l +RTIP r t i p +RTI r t i +RTKL r t k l +RTK r t k +RTLM r t l m +RTL r t l +RTL's r t l's +RTML r t m l +RTMP r t m p +RTM r t m +RTNDA- r t n d a +RTN r t n +RTO r t o +RTOs r t o's +RTOS r t o s +rtPA r t p a +RTP r t p +RTR r t r +rt r t +R. T. r t +R.T. r t +RT- r t +RT r t +rts r t s +RTS- r t s +RTS r t s +RTS's r t s's +RTTOV r t t o v +RTT r t t +RTTY r t t y +RTUK r t u k +RTU r t u +RTVC r t v c +RTVE r t v e +RTVFBiH r t v f b i h +RTV r t v +RTVV r t v v +RTW r t w +RTXC r t x c +RTX r t x +RTZ r t z +Ruao r u a o +RUC r u c +RUC's r u c's +RUFC r u f c +Rukn r u k n +Ruk r u k +RUMC r u m c +RUNX r u n x +Rupf r u p f +Rupr r u p r +rup r u p +Rup r u p +RUP r u p +R.U.R. r u r +RUR r u r +R. U. r u +R.U. r u +RU r u +RUSD r u s d +RutB r u t b +Ruu r u u +RUV r u v +Ruwa r u w a +Ruy r u y +Ruyt r u y t +RVAH- r v a h +RVAR r v a r +Rvat r v a t +RVC r v c +RVCT r v c t +RVD r v d +R. V. E. r v e +RVE r v e +RVGK r v g k +R.V.G. r v g +R. V. J. r v j +R.V.J. r v j +RVM r v m +RVNG r v n g +Rvo r v o +RVO r v o +RVR r v r +R. V. r v +R.V. r v +RV r v +rvs r v s +RVTD's r v t d's +RVU r v u +RVV- r v v +rwa r w a +Rwa r w a +RWA r w a +rwb r w b +RWB r w b +RWC r w c +RWD r w d +RWDSU r w d s u +RWE r w e +RWEs r w e's +RWFC r w f c +RWIS r w i s +rwjf r w j f +R.W.R.J. r w r j +R. W. r w +R.W. r w +RW r w +RWSL r w s l +R. W. S. r w s +RWS r w s +Rxa r x a +RXL r x l +RX- r x +RX r x +RZR r z r +rz r z +R. Z. r z +RZ r z +RZS r z s +S.A.B. s a b +S.A.D. s a d +SAIC- s a i c +SAIC s a i c +Sa'id s a i d +S.A.I. s a i +SAIT's s a i t's +SAKEC s a k e c +S. A. L. s a l +S.A.P.A. s a p a +S.A.P.I. s a p i +S. A. R. s a r +S.A.R. s a r +S. A. s a +S.A. s a +SA s a +S.A.S. s a s +SAS s a s +S.A.V.A. s a v a +Saxl s a x l +Sa'yo s a y o +SAZU s a z u +SBAC's s b a c's +Sbai s b a i +S.B.A. s b a +SBA s b a +SBBK s b b k +SBB- s b b +SBB s b b +SBCM s b c m +SBCMT s b c m t +sbc s b c +SBC s b c +SBCs s b c's +SBDE s b d e +SBD s b d +SBE s b e +SBF s b f +SBI s b i +SBKP s b k p +sbk s b k +SBK s b k +SBL s b l +SBML s b m l +SBMNH s b m n h +SBM s b m +SBN s b n +SBOE s b o e +SBOL s b o l +SBOP s b o p +Sborz s b o r z +SBP- s b p +SBP s b p +SBRJ s b r j +sb s b +S. B. s b +S.B. s b +SB- s b +SB s b +SBSE s b s e +SBS s b s +SBS's s b s's +SBTDC s b t d c +SBT s b t +SBTU s b t u +Sbu s b u +SBU s b u +SBY s b y +SBZ- s b z +SBZ s b z +S.C.A. s c a +SCA s c a +SCBA s c b a +SCB s c b +SCBWI s c b w i +SCCA s c c a +SCCA's s c c a's +SCCC s c c c +SCCI s c c i +SCCP s c c p +S.C.C. s c c +SCC s c c +scr s c r +SCR s c r +SCRs s c r's +SCRS s c r s +SCRTC s c r t c +scry s c r y +sc s c +S. C. s c +S.C. s c +SC- s c +SC s c +S.C.S.C. s c s c +SC's s c's +SCS s c s +SCTC s c t c +SCTE s c t e +SCTO s c t o +SCTP s c t p +SCT s c t +SCTS s c t s +SCTV s c t v +SCUAA s c u a a +Scuf s c u f +SCU's s c u's +SCW s c w +Scymn s c y m n +SD&AE s d and a e +SDAP s d a p +SDAPS s d a p s +SDA s d a +SDAS s d a s +SDASS s d a s s +SDAT s d a t +SDAX s d a x +SDB s d b +SDCC s d c c +SDCG s d c g +SDC s d c +SDDI s d d i +SDDOT s d d o t +SDD s d d +SDDS s d d s +Sdei s d e i +sde s d e +Sde s d e +SDF s d f +SDG&E s d g and e +SDHA- s d h a +SDHA s d h a +SDHC s d h c +SDHD s d h d +SDHHD s d h h d +SDH s d h +SDI, s d i +SDI s d i +SDJ s d j +sdk s d k +SDK s d k +SDK's s d k's +SDLC s d l c +SDLP s d l p +SDLP's s d l p's +S. D. L. s d l +S.D.N.Y. s d n y +S. D. s d +S.D. s d +SDSM&T's s d s m and t's +S. D. S. s d s +S.E.C. s e c +Sejms s e j m's +sejr s e j r +Sekl s e k l +Sek s e k +SEK s e k +S. E. s e +S.E. s e +SESL s e s l +ses s e s +Ses s e's +SEs s e's +S.E.S. s e s +SES- s e s +SES s e s +S.E.S.'s s e s's +S.E.X. s e x +Seyh s e y h +SFACS- s f a c s +sfadb s f a d b +SFA s f a +SFA's s f a's +Sfax s f a x +S.F.B.J. s f b j +SFB s f b +SFCC s f c c +SFCH s f c h +SFC s f c +SFDR s f d r +SFD s f d +SFE s f e +SFES s f e s +SFFAS s f f a s +SFFCo s f f c o +SFFH s f f h +SFF s f f +SFG s f g +SFI s f i +S. F. L. s f l +SFL s f l +sfn s f n +SFN s f n +SFOD- s f o d +SFOR s f o r +SFO s f o +SFPA s f p a +SFPD s f p d +SFP s f p +SFRA s f r a +SFRJ s f r j +sFRP- s f r p +SFR s f r +'sf s f +sf s f +S. F. s f +S.F. s f +SF- s f +SF s f +SFSG s f s g +SFSR s f s r +sfs s f s +S.F.'s s f's +SFS- s f s +SFS s f s +SFSS s f s s +SFSU s f s u +SFU s f u +SFWA- s f w a +SFWA s f w a +SFX s f x +SFX's s f x's +SGAE s g a e +SGAP s g a p +SGA s g a +SGA's s g a's +sgb s g b +SGB s g b +S.G.C. s g c +SGC s g c +SGD s g d +SGE s g e +S. G. F. s g f +SGF s g f +SGH- s g h +SGH s g h +Sgip s g i p +SGI s g i +SGL s g l +sgml s g m l +SGML- s g m l +SGML s g m l +SGP s g p +SGRAM s g r a m +sgra s g r a +SGR s g r +S. G. s g +S.G. s g +SGSN s g s n +SGSNs s g s n's +SGS s g s +SGSY s g s y +Sgt. sergeant +sgt s g t +SGT s g t +SGU s g u +SGV s g v +SGX s g x +SHBG s h b g +SHB s h b +SHBT s h b t +SHCA s h c a +Shch- s h c h +SHC s h c +SHG s h g +SHH s h h +S.H.I.E.L.D.'s s h i e l d's +S.H.I.E.L. s h i e l +SHL- s h l +SHL s h l +Shma s h m a +SHMD s h m d +SHM s h m +SHN s h n +S. H. s h +S.H. s h +S.I.D. s i d +SIF s i f +S.I.R. s i r +S. I. s i +S.I. s i +S.I.T. s i t +SJAM s j a m +SJCH s j c h +SJC s j c +SJDA s j d a +S.J.D. s j d +sjef s j e f +SJEM s j e m +S. J. J. F. s j j f +SJK s j k +S.J.L. s j l +SJL s j l +SJM s j m +Sjon s j o n +SJPCD s j p c d +SJP s j p +SJPT s j p t +SJR s j r +SJR's s j r's +S. J. s j +S.J. s j +SJ s j +SKB- s k b +SKB s k b +SKC s k c +SKD s k d +S. K. F. s k f +SKF- s k f +SKF s k f +SKG s k g +SKH s k h +SKJ s k j +SKPC s k p c +SKP s k p +SKR s k r +SKSD s k s d +S. K. s k +S.K. s k +SK- s k +SK s k +S. K. S. s k s +S.L.A.A.'s s l a a's +SLAF s l a f +sla s l a +SLA s l a +SLAs s l a's +SLCO s l c o +SLC- s l c +SLC s l c +SLDL s l d l +SLD- s l d +SLD s l d +S.L.E. s l e +SLF s l f +SLG s l g +SLHS s l h s +S.L.I.F.E.R. s l i f e r +SLL s l l +SLMC s l m c +S. L. M. s l m +SLM s l m +SLN s l n +SLPIM s l p i m +SLP s l p +SLRC s l r c +SLR s l r +S&L s and l +SLS&E s l s and e +SLSF s l s f +SLSK s l s k +S. L. s l +S.L. s l +SL- s l +SL s l +SLS s l s +SLVR s l v r +SLV s l v +S.M.A.R.T. s m a r t +S.M.A.S.H. s m a s h +SMA s m a +SMB s m b +SMe s m e +SME s m e +SME's s m e's +SMEs s m e's +SMF s m f +SMG- s m g +SMG s m g +SMHI s m h i +SMH s m h +SmI s m i +SMI s m i +SMJR s m j r +SMK s m k +SMLS s m l s +SMMT s m m t +SMN s m n +SMNS s m n s +SMP s m p +SMP's s m p's +SMPSs s m p s's +SMPTE s m p t e +SMRJ s m r j +smr s m r +SMR- s m r +SMR s m r +S&M s and m +sm s m +S. M. s m +S.M. s m +SM s m +S.M.'s s m's +SMs s m s +S.M.S. s m s +SMS. s m s +SMS s m s +SMS's s m s's +SMSU s m s u +SMTP s m t p +S.M.T. s m t +SMT s m t +SMTV s m t v +SMU s m u +SMW s m w +SMX s m x +Smyl s m y l +SMYS s m y s +SNA s n a +SNBA s n b a +SNB s n b +SNCB s n c b +SNCC s n c c +SNCF s n c f +SNC- s n c +SNC s n c +SNDC s n d c +SND s n d +SNEP- s n e p +SNEP s n e p +SNESjr s n e s j r +SNES s n e s +SNET s n e t +SNF s n f +SNFU s n f u +sngle s n g l e +SNGM- s n g m +SNG s n g +SNJ- s n j +SNK s n k +SNLA s n l a +SNL s n l +SNLS s n l s +SNMCMG s n m c m g +SNMC s n m c +SNMMA s n m m a +SNMP s n m p +SNNPR s n n p r +SNNR s n n r +SNPJ s n p j +SNP s n p +SNPs s n p's +SNRI s n r i +SNR s n r +SNSD's s n s d's +SnSe s n s e +sn- s n +S. N. s n +S.N. s n +SN- s n +SN s n +SNS s n s +S.O.E. s o e +S. O. s o +S.O. s o +Sos s o's +S.O.S. s o s +SOS s o s +Sovn s o v n +Sov s o v +SOX s o x +S. O. Y. s o y +S&P 500 s and p five hundred +S.P.A.L. s p a l +SPB s p b +SPCAs s p c a's +SPCG s p c g +SPCK s p c k +S.P.C. s p c +SPC s p c +SPDI s p d i +SPD- s p d +SPD s p d +Spe s p e +SPE s p e +SPES s p e s +SPFH s p f h +SPFL s p f l +SPF s p f +SPG s p g +SPHL s p h l +SPH- s p h +SPH s p h +SPINE's s p i n e's +SPIN's s p i n's +SPIR s p i r +SPI s p i +S.P.I.T. s p i t +SPJA s p j a +SPK s p k +SPLA- s p l a +SPLA s p l a +SPLC s p l c +SPLM- s p l m +SPL s p l +SPME s p m e +SPML s p m l +SPMRL s p m r l +SPM s p m +SPNJ s p n j +SPNM s p n m +SPNN s p n n +SPN s p n +Spoa s p o a +SPOC s p o c +SPOF s p o f +SPO s p o +SPOU s p o u +SpPIn s p p i n +SPP s p p +SPP's s p p's +SPQA s p q a +SPRL s p r l +SPRM s p r m +SPR s p r +SPRU s p r u +S&P s and p +SPSA s p s a +SPSL s p s l +S. P. s p +S.P. s p +SP- s p +SP s p +SP's s p's +SPs s p's +SPS- s p s +SPS s p s +SPSS s p s s +SPTA s p t a +SPT s p t +SPUC s p u c +SPUP s p u p +SPU s p u +SPUs s p u's +SPV s p v +SPX- s p x +SQA s q a +SQBB s q b b +SQI- s q i +SQI s q i +sql s q l +SQL- s q l +SQL s q l +SQM s q m +SQQ- s q q +sq s q +S. Q. s q +SQ s q +Sra s r a +SRA s r a +SRAs s r a's +Srba s r b a +SRBIJA s r b i j a +Srbi s r b i +SRBOC s r b o c +SRBP s r b p +SRB s r b +SRBs s r b's +SRBY s r b y +SRCC s r c c +SRCL s r c l +SRC s r c +SRC's s r c's +SRCS s r c s +S.R.E. s r e +SRE s r e +SRFC s r f c +S.R.F. s r f +SRF s r f +S.R.G. s r g +SRG s r g +SRGs s r g's +Srhir s r h i r +SRH s r h +S. R. J. s r j +SRK- s r k +SRK s r k +SRK's s r k's +SRLGs s r l g's +S.R.L. s r l +SRL s r l +SRMC s r m c +S.R.M. s r m +SRM s r m +SRMs s r m's +Srni s r n i +SRN s r n +srp s r p +SRPT s r p t +S. R. R. s r r +SRR s r r +sr s r +S. R. s r +S.R. s r +SR- s r +SR s r +SRS s r s +SRT- s r t +SRT s r t +SRT's s r t's +Srul s r u l +SRU s r u +SRU's s r u's +SRV s r v +srx s r x +SRY s r y +SRZ- s r z +SSAA s s a a +SSAB s s a b +SSAC s s a c +SSAR s s a r +SSA s s a +SSAs s s a's +SSBSE s s b s e +S.S.B. s s b +SSB s s b +SSCCC s s c c c +S.S.C. s s c +SSC- s s c +SSC s s c +SSCV s s c v +SSDB s s d b +SSDL s s d l +SSDP s s d p +S.S.D. s s d +SSD s s d +SSDs s s d's +SSE s s e +SSFL s s f l +SSF s s f +SSGRC s s g r c +SSG- s s g +SSG s s g +SSHSA's s s h s a's +SSH s s h +SSI s s i +SSIS s s i s +SSJA s s j a +SSKI's s s k i's +SSK s s k +SSLAM s s l a m +ssl s s l +SSL s s l +SSLT s s l t +SSME s s m e +SSM- s s m +SSM s s m +SSN- s s n +SSN s s n +SSOD s s o d +SSOF s s o f +SSO s s o +SSoSV s s o s v +SSPA s s p a +SSPH s s p h +SSPR s s p r +SSP s s p +SSPX s s p x +SSQ s s q +SSRAA s s r a a +SSRF s s r f +SSRI s s r i +SSRIs s s r i's +SSRN s s r n +SSRP s s r p +SSR s s r +S&S s and s +SSSC s s s c +SSSI s s s i +SSSR s s s r +ss s s +S. s s's +Ss s's +S. S. s s +S.S. s s +SS- s s +SSS s s s +SSTC- s s t c +SSTH s s t h +SSTL- s s t l +SSTL s s t l +SSTR s s t r +SST s s t +Ssu s s u +SSU s s u +SSVC s s v c +SSV s s v +SSWAHS s s w a h s +SSWC s s w c +ssw s s w +SSW s s w +SSX s s x +Ssy s s y +ST&AJ s t and a j +S.T.A.L.K.E.R. s t a l k e r +S.T.A.T.U.S. s t a t u s +STB s t b +STCC- s t c c +STCC s t c c +S.T.C. s t c +STC s t c +STCW s t c w +STDP s t d p +STDs s t d's +STD- s t d +STD s t d +STF s t f +stfv s t f v +STGs s t g's +STG s t g +STIs s t i's +STIS s t i s +STi s t i +STI- s t i +STI s t i +STK s t k +stl s t l +S.T.L. s t l +STL s t l +STMB s t m b +STMIK s t m i k +STM's s t m's +STM s t m +STNS s t n s +STN- s t n +STN s t n +STPI s t p i +STPNS s t p n s +STPRI s t p r i +STP s t p +S.T.R.I.D.E. s t r i d e +STRV s t r v +Stryj s t r y j +STRZ s t r z +St. saint +S&T s and t +STSAT- s t s a t +STScI s t s c i +ST's s t's +STS- s t s +STS s t s +S. T. s t +S.T. s t +ST- s t +STT s t t +STVL s t v l +STVS s t v s +STV s t v +STXBP s t x b p +STX s t x +S.U.C. s u c +SUDV s u d v +suo s u o +Suo s u o +S. U. s u +SUSV s u s v +Susz s u s z +SUTs s u t's +SUT s u t +Suu s u u +SUVF s u v f +SUVs s u v's +SUV- s u v +SUV s u v +SUWN s u w n +SUW s u w +Suy s u y +Svac s v a c +svar s v a r +SVA s v a +SVB s v b +SVC s v c +SVD s v d +sve s v e +Sve s v e +SVGA s v g a +SVG s v g +SVGT s v g t +SVIA s v i a +SVM's s v m's +SVMs s v m's +SVM s v m +svn s v n +svom s v o m +SVOPC s v o p c +SVO s v o +SVP s v p +SVPW s v p w +SVSCEP s v s c e p +SVSC s v s c +SVS s v s +sv s v +S. V. s v +S.V. s v +SV- s v +SV s v +SVTs s v t's +SVT s v t +S.V.U. s v u +SVU s v u +SVV s v v +SVW s v w +SVYASA s v y a s a +SVZ s v z +SWABC s w a b c +SWAC s w a c +S.W.A. s w a +SWC s w c +SWD- s w d +SWD s w d +S.W.E.A.T. s w e a t +Swe s w e +SWE s w e +SWF s w f +SWG s w g +SWHL s w h l +SWHS s w h s +SWH s w h +SWJN s w j n +S. W. K. s w k +SWMRS s w m r s +SWM s w m +SWNH s w n h +SWOC s w o c +S.W.O.R.D. s w o r d +SWPA s w p a +SWPL s w p l +SWP s w p +SwRI s w r i +SWR- s w r +SWR s w r +S&W s and w +SWS s w s +sw s w +S. W. s w +S.W. s w +SW- s w +SW s w +SWTPC s w t p c +SWT s w t +SWWTP s w w t p +SXCT s x c t +SXN s x n +SXSW s x s w +sx s x +S.X. s x +SX- s x +SX s x +SXY s x y +SXZ s x z +SysML s y s m l +SysRq s y s r q +Sys s y's +SYS s y s +SysV s y s v +S. Y. s y +S.Y. s y +SyT s y t +syv s y v +Syxx s y x x +SZA s z a +Szasz s z a s z +SZD- s z d +SZDSZ s z d s z +Szer s z e r +szkic s z k i c +Szklo s z k l o +Szlak s z l a k +SzMME s z m m e +Szpir s z p i r +SZSE s z s e +sz s z +S. Z. s z +SZ s z +Szu s z u +Szyk s z y k +SZZ s z z +taf t a f +Taf t a f +TAF t a f +T.A.P. t a p +T.A.S. t a s +T. A. t a +T.A. t a +T&A t and a +TAZ t a z +TBA t b a +TBCA t b c a +TBCs t b c's +TBCS t b c s +TBD t b d +TBE t b e +TBF t b f +TBHS t b h s +TBI t b i +TBK- t b k +TBK t b k +TBL t b l +TBMA t b m a +TBMM t b m m +TBMs t b m's +T.B.M. t b m +TBM t b m +TBNJ t b n j +TBN t b n +TBP- t b p +TBP t b p +TBSC t b s c +TBS t b s +T. B. t b +T.B. t b +TB- t b +TB t b +TBTG t b t g +TBT t b t +TBV t b v +TBWA t b w a +TBX t b x +TCAR t c a r +TCA t c a +TCC t c c +TCDB t c d b +TCDC t c d c +TCDD t c d d +TCDOG t c d o g +TCDSB t c d s b +TCEB t c e b +T.C.E. t c e +TCE t c e +TCFAQ t c f a q +TCF t c f +TCGA t c g a +TCG- t c g +TCG t c g +TCHC t c h c +TCHRD t c h r d +TCH t c h +TCiAP t c i a p +TCI t c i +TCKs t c k's +TCLP t c l p +TCL t c l +TCM t c m +TCO t c o +TCPL t c p l +TCPMP t c p m p +TCP's t c p's +TCP- t c p +TCP t c p +TCR t c r +T. C. S. t c s +TCS t c s +tc t c +Tc- t c +T. C. t c +T.C. t c +TC- t c +TC t c +TCT t c t +TCU's t c u's +TCUs t c u's +tcu t c u +TCU t c u +TCV t c v +TCWC t c w c +TCWP t c w p +TCW t c w +TDA t d a +tdb t d b +TDCC t d c c +TDCi t d c i +TDCJ t d c j +TDC's t d c's +TDC t d c +TDD t d d +TDE t d e +TDF t d f +TDH t d h +TdIF t d i f +TDI t d i +TDK t d k +TDLAS t d l a s +TDMA- t d m a +TDMA t d m a +TDME t d m e +TDM- t d m +TDM t d m +TDO t d o +TDP t d p +TDRS t d r s +TDR t d r +TDSB's t d s b's +TD's t d's +TDs t d's +T. D. t d +T.D. t d +TDT t d t +Teatr t e a t r +TEBD t e b d +T. E. t e +T.E. t e +TEVT t e v t +TFAP t f a p +TFAS t f a s +TFA t f a +TFC t f c +TFCU t f c u +TFF t f f +TFG t f g +T.F.H. t f h +TFH t f h +TFIH t f i h +TFIID t f i i d +TFK t f k +TFL t f l +TFM t f m +TFN t f n +TFO t f o +TFP t f p +TFR t f r +TFSA t f s a +TFSI t f s i +TFSP t f s p +TFS t f s +T&F t and f +TFTC t f t c +tf t f +T. F. t f +T.F. t f +TF- t f +TF t f +TFTR t f t r +TFT t f t +TFVC t f v c +TFWA t f w a +TFW t f w +TGA t g a +T. G. C. t g c +TGC t g c +tge t g e +T.G.E. t g e +TGE t g e +tgf t g f +TGF- t g f +TGF t g f +TGG t g g +TGIF t g i f +TGL t g l +TGM's t g m's +TGM t g m +TGP t g p +TGS t g s +tg t g +T. G. t g +T.G. t g +TG t g +TGVs t g v's +tgv t g v +TGV- t g v +TGV t g v +THB t h b +THCD- t h c d +THCF t h c f +THC- t h c +THC t h c +T.H.E.M. t h e m +T.H.F. t h f +THF t h f +THG t h g +THHK t h h k +T.H.I.N.K. B.I.G. t h i n k b i g +THJ- t h j +THLL t h l l +THL t h l +THMP t h m p +THM t h m +THNN t h n n +Thok t h o k +T.H.P. t h p +THQ's t h q's +THQ t h q +THR's t h r's +THR t h r +ThSe t h s e +THSR t h s r +T. H. S. t h s +T. H. t h +T.H. t h +T.H.U.N.D.E.R. t h u n d e r +THW t h w +THX t h x +TIAA- t i a a +TIAA t i a a +T. I. A. t i a +TIFR t i f r +TIF t i f +T.I.'s t i's +T.I. t i +Tiu t i u +tiv t i v +Tiv t i v +TIYM t i y m +Tiy's t i y's +TJD- t j d +TJFR t j f r +T.J.F. t j f +T. J. J. t j j +TJJ t j j +TJP t j p +TJRC t j r c +T. J. S. t j s +T. J. t j +T.J. t j +TJ t j +t'ju t j u +TJX t j x +tjz t j z +TKAG t k a g +TKA's t k a's +TKA t k a +T. K. B. t k b +TKB- t k b +TKCA- t k c a +TKD t k d +TKE t k e +TKI t k i +tko t k o +Tko t k o +TKO- t k o +TKO t k o +TKP t k p +TK's t k's +TKS t k s +tk t k +T. K. t k +T.K. t k +TK- t k +TK t k +TKT t k t +Tlas t l a's +TLA t l a +TLC's t l c's +TLC t l c +TlCu t l c u +TLDs t l d's +TLDS t l d s +TLD t l d +TLE t l e +TLF t l f +TLG t l g +TLH t l h +tli t l i +TLI t l i +TLM t l m +TLN t l n +tlp t l p +TLP t l p +TLRs t l r's +TLR t l r +TLs t l's +TLS t l s +TLTB t l t b +tl t l +T. L. t l +T.L. t l +TL- t l +TL t l +Tluk t l u k +TLU's t l u's +Tluszcz t l u s z c z +TLV- t l v +TLV t l v +TMA- t m a +TMA t m a +TMBG t m b g +TMB t m b +tmc t m c +TMC t m c +TMDL t m d l +TMD t m d +TMEM t m e m +TMF t m f +TMGs t m g's +TMG t m g +TMI t m i +TMJ t m j +T. M. M. t m m +TMMTX t m m t x +TMNT t m n t +TMOK t m o k +tmos t m o s +TMO's t m o's +TMOS t m o s +TMO t m o +TMPRSS t m p r s s +TMP t m p +tmRNA t m r n a +TMRS t m r s +TMR t m r +TMSI t m s i +TMSs t m s's +T.M.s t m's +TM's t m's +TMS t m s +tm t m +T. M. t m +T.M. t m +TM- t m +TM t m +T.M.T. t m t +TMT t m t +TMU t m u +TMX t m x +TMZ t m z +TNA's t n a's +TNAS t n a s +tna t n a +TNA t n a +tnbc t n b c +TNCA t n c a +TNCC t n c c +TNCDSB t n c d s b +TNC- t n c +TNC t n c +TNFA t n f a +TNFRSF t n f r s f +TNFSF t n f s f +TNF- t n f +TNF t n f +TNGHT t n g h t +TNG t n g +TNIP t n i p +TNI t n i +TNK- t n k +TNL t n l +TNMM t n m m +TNMS t n m s +TNM t n m +TNNC t n n c +TNN t n n +TNO t n o +TNPL t n p l +TNP t n p +TNQ t n q +TNRIS t n r i s +TNRP t n r p +TNR t n r +TNSDL t n s d l +TNSTC- t n s t c +TNSTC t n s t c +TNS- t n s +TNS t n s +TNTAs t n t a's +tn t n +T. N. t n +T.N. t n +TN- t n +TN t n +TNT's t n t's +TNT t n t +TNX- t n x +T.O.'s t o's +T. O. t o +T.O. t o +TOTP t o t p +tou t o u +Tou t o u +ToU t o u +TOU t o u +TPAC t p a c +TPAO t p a o +tPA t p a +TPA- t p a +TPA t p a +TPB t p b +TPC- t p c +TPC t p c +TPEC t p e c +TPF t p f +TPG- t p g +TPG t p g +TPH t p h +TPI t p i +TPLF t p l f +TPMS t p m s +TPM t p m +TPO t p o +TPP t p p +tpr t p r +TPR t p r +TPS- t p s +TPS t p s +tp t p +T. P. t p +T.P. t p +TP- t p +TP t p +TPU t p u +TPVs t p v's +TPWS t p w s +TPW t p w +TQR t q r +TQS t q s +T.Q. t q +TQ t q +tra t r a +Tra t r a +TRA t r a +TRAXX t r a x x +Trbic t r b i c +T. R. B. t r b +TRB t r b +Trcek t r c e k +TRC t r c +TRDB- t r d b +TRD t r d +TRFA t r f a +TRF t r f +TRG t r g +TRHS t r h s +TRH t r h +T.R.I.B.E. t r i b e +TRIZ t r i z +TrkA t r k a +TRL- t r l +TRL t r l +Trmcic t r m c i c +TRMM t r m m +T.R.M.P.A. t r m p a +T. R. M. t r m +TRM- t r m +tRNA t r n a +TRNC t r n c +TRPC t r p c +TRPM t r p m +TRPP t r p p +TRP's t r p's +TRPs t r p's +TRP t r p +TRPV t r p v +TRPY t r p y +TRSC t r s c +TRSI t r s i +trs t r s +T.R.S. t r s +TRS t r s +TRSV t r s v +tr t r +T. R. t r +T.R. t r +TR- t r +TR t r +TRT t r t +T.R.U. t r u +truTV t r u t v +TruTV t r u t v +Trve t r v e +TRVL t r v l +TRV t r v +TRW t r w +TRX t r x +TRZ t r z +Tsa t s a +TSA t s a +TSB t s b +TSCC t s c c +TSC- t s c +TSC t s c +TSF t s f +TSG- t s g +TSG t s g +TSH t s h +TSIO- t s i o +TSI t s i +tsit t s i t +Tsiv t s i v +TSJDF t s j d f +TSLA t s l a +TSLC t s l c +TSL t s l +TSN t s n +T.S.O.L. t s o l +tso t s o +Tso t s o +T.S.O. t s o +TSO t s o +TSP t s p +TSQ- t s q +TSRA t s r a +TSRTC t s r t c +tsr t s r +TSR t s r +TSSA t s s a +TSS t s s +TSTA t s t a +TSTC t s t c +Ts t's +T. S. t s +T.S. t s +TS- t s +TS t s +TST t s t +tsus t s u s +Tsvi t s v i +TSV t s v +TSW t s w +TSX t s x +TTAB t t a b +Tta t t a +TTA t t a +TTBB t t b b +TTCP t t c p +TTC's t t c's +TTC t t c +TTDC t t d c +TTD t t d +tteok t t e o k +tte t t e +TTFF t t f f +TTHS t t h s +TTH t t h +TTI t t i +TTKG t t k g +TTKST t t k s t +TTL- t t l +TTL t t l +T.T.N. t t n +ttp t t p +TTP t t p +TTR t t r +TTS t t s +T&T t and t +tt- t t +T. T. t t +T.T. t t +TT- t t +TT t t +TTU t t u +TTV t t v +TTW t t w +TUAW t u a w +TUBB t u b b +TUDCA t u d c a +TUHUS t u h u s +TUKP's t u k p's +T'uqu t u q u +TUSCC t u s c c +TUSC t u s c +T.U. t u +TU- t u +TU t u +Tuzk t u z k +tvaan t v a a n +Tvam t v a m +TVARK t v a r k +tva t v a +TVA t v a +TVB's t v b's +TVB t v b +TVCM t v c m +TVC t v c +TVD t v d +TVETs t v e t's +TVE t v e +TVF t v f +TVGN t v g n +TVGOS t v g o s +TVii t v i i +TVIn t v i n +TVi t v i +TVI t v i +TVL t v l +TVMK t v m k +TVM t v m +TVN t v n +TVNZ's t v n z's +TVNZ t v n z +tvo t v o +Tvo t v o +TVO t v o +TVP t v p +TVRi t v r i +tvr t v r +TVR t v r +TVSpy t v s p y +tvs t v s +TV's t v's +TVs t v's +TVS t v s +tVTA t v t a +TVT- t v t +TVT t v t +tv t v +T. V. t v +T.V. t v +TV- t v +TV t v +TVXQ's t v x q's +TVXQ t v x q +TVX t v x +TWAP t w a p +TWA's t w a's +Twa t w a +TWA t w a +T. W. B. t w b +TWCS t w c s +TWC t w c +TWF t w f +Twi t w i +TWI t w i +TWK t w k +T.W.P. t w p +TWP t w p +TWR- t w r +TWR t w r +TWTs t w t's +TWT t w t +T. W. t w +T.W. t w +TW t w +T.W.U. t w u +TWU t w u +TXC t x c +TxDOT t x d o t +TXII t x i i +TXJS t x j s +TXK t x k +TXMM t x m m +TXNDC t x n d c +TXP t x p +TXTC t x t c +TXT t x t +T.X. t x +TX- t x +TX t x +TXU t x u +tya t y a +TYC t y c +TYIB t y i b +TYIN t y i n +Tyk t y k +Tza'ar t z a a r +Tze t z e +TZL t z l +Tzrif t z r i f +Tzuh t z u h +TZUM t z u m +Tzvi t z v i +UAAP's u a a p's +UAAP u a a p +UAB's u a b's +UAB u a b +UAC- u a c +UAC u a c +UADY u a d y +UAE's u a e's +U.A.E. u a e +UAE- u a e +UAE u a e +UAF u a f +UAHC u a h c +UAIC u a i c +UAI u a i +UALR u a l r +UAL u a l +UALVP u a l v p +UAM u a m +UANL u a n l +UAN u a n +UAOC u a o c +UAP u a p +UAR's u a r's +UARS u a r s +UAR u a r +UASLP u a s l p +UASL u a s l +Ua's u a's +UAs u a's +UAS u a s +uat u a t +U.A. u a +UA u a +UAVs u a v's +UAV u a v +U.A.W. u a w +UAW- u a w +UAW u a w +Uba u b a +UBA u b a +UBBC u b b c +UBB u b b +UbcM u b c m +UBC's u b c's +UBC u b c +UBD u b d +Ube u b e +UBE u b e +ubi u b i +Ubi u b i +U.B.I. u b i +UBI u b i +UBM's u b m's +UBM u b m +UBOS u b o s +UBO u b o +UBPR u b p r +UBP u b p +UBRE u b r e +UBRS u b r s +UB's u b's +UBS u b s +UB- u b +UB u b +UBU u b u +UBV u b v +UBX u b x +UCAC u c a c +Ucar u c a r +UCAR u c a r +Uca u c a +UCA u c a +UCAV u c a v +UCBH u c b h +UCBTLA u c b t l a +UCB u c b +UCCE u c c e +UCCJ u c c j +UCCS u c c s +UCC- u c c +UCC u c c +UCDA u c d a +UCDP u c d p +UCD u c d +UCF's u c f's +UCF u c f +UCG u c g +UCIEP u c i e p +UCISA- u c i s a +UCI's u c i's +UCI u c i +UCLA's u c l a's +UCLA- u c l a +UCLA u c l a +UCLH u c l h +UCL u c l +Ucmak u c m a k +UCMJ u c m j +UCMMA u c m m a +UCMP u c m p +UCM u c m +Ucn u c n +UCOS u c o s +UCO u c o +UCP u c p +UCREF u c r e f +UCR u c r +UCSB u c s b +UCSC u c s c +UCSD's u c s d's +UCSD u c s d +UCSF u c s f +UCSN u c s n +uCs u c's +UCS u c s +UCTI u c t i +UCT u c t +Uc u c +U.C. u c +UC- u c +UC u c +Ucuncu u c u n c u +UCU u c u +UCVTS u c v t s +UCW u c w +Uczta u c z t a +UDAR u d a r +UDA's u d a's +Uda u d a +UDA u d a +UDC u d c +UDF u d f +UDG u d g +UDHR u d h r +UDICHI u d i c h i +Udit u d i t +UDK u d k +UDN u d n +Udny u d n y +UDP- u d p +UDP u d p +UDRP u d r p +UDR u d r +UDS u d s +UDTs u d t's +UDT u d t +ud u d +Ud u d +U. D. u d +UD u d +udu u d u +Udu u d u +UDYCO u d y c o +UEA u e a +UEBT u e b t +UEC u e c +Ueda's u e d a's +Uéda u e acute d a +UEFA's u e f a's +UEFA u e f a +UEFI u e f i +UEFS u e f s +UEF u e f +UEGCL u e g c l +UEI u e i +UET u e t +U. E. u e +UE u e +UFABC u f a b c +UFA's u f a's +UFCO u f c o +UFC's u f c's +UFC u f c +UFCW u f c w +UFD u f d +UFE u f e +UFF u f f +UFIA u f i a +UFIP u f i p +UFI u f i +UFJ u f j +UFL's u f l's +UFL u f l +UFM u f m +UFOFU u f o f u +UFORE u f o r e +UFO's u f o's +UFOs u f o's +Ufot u f o t +ufo u f o +Ufo u f o +UFO- u f o +UFO u f o +UFPE u f p e +UFPLS u f p l s +UFPR u f p r +UFRGS u f r g s +UFRJ u f r j +UFSB u f s b +UFSM u f s m +UF's u f's +UFS u f s +UFTP u f t p +UFT u f t +U. F. u f +UF- u f +UF u f +UFU u f u +UFW u f w +Uga u g a +UGA u g a +UGCCWA u g c c w a +UGC- u g c +UGC u g c +UGEED- u g e e d +Uge u g e +UGF- u g f +UGK u g k +UGME u g m e +Ugni u g n i +UGP u g p +Ugra u g r a +Ugrszke u g r s z k e +UGS u g s +UGT u g t +Ug u g +UG u g +UHD u h d +UHF u h f +Uhha u h h a +UHMW u h m w +uhn u h n +UHOA u h o a +UHPT u h p t +UHP- u h p +UHP u h p +UHRA u h r a +UHRF u h r f +UHSAA u h s a a +UHTCs u h t c's +UHTI u h t i +Uhud u h u d +U.H. u h +uhur u h u r +UIAA u i a a +Uiam u i a m +UIA u i a +UIAW u i a w +Uibh u i b h +Uible u i b l e +UIC u i c +UID u i d +UIFL u i f l +Uig u i g +UIL u i l +UIMC u i m c +UIM u i m +UIN u i n +UIP u i p +uisae u i s a e +UISD u i s d +UiS u i s +UI's u i's +UIS u i s +UiTM u i t m +uit u i t +Uit u i t +UIT u i t +UI- u i +UI u i +UIV u i v +Uiy u i y +UJA u j a +UJC u j c +Ujed u j e d +UJEP u j e p +Ujsag u j s a g +UJS u j s +uj u j +U.J. u j +UJ u j +Uka u k a +UKA u k a +UKCC u k c c +UKC u k c +Ukhra u k h r a +Ukic u k i c +UKIP's u k i p's +Uki u k i +UKM u k m +UKNC u k n c +Ukoh u k o h +Uko u k o +UKPC u k p c +UKRDA's u k r d a's +UKRDA u k r d a +ukr u k r +Ukr u k r +UKR u k r +UK's u k's +UKTV's u k t v's +UKTV u k t v +Ukui u k u i +u'k u k +uk u k +Uk u k +U.K. u k +UK- u k +UK u k +UKUP u k u p +UKVRN u k v r n +UKW u k w +Ukwu u k w u +ULC- u l c +ULC u l c +ULDB u l d b +ULEB u l e b +ULEV u l e v +ULFA u l f a +Ulic u l i c +Ull u l l +Ulms u l m's +UL's u l's +ULS u l s +UlSU u l s u +Uluj u l u j +U.L. u l +UL u l +ulus u l u s +Ulus u l u's +ULV u l v +Ulwa u l w a +Ulwe u l w e +UMBC u m b c +Umbr u m b r +UMCC u m c c +UMC u m c +UMD u m d +umelcu u m e l c u +UMe u m e +UMF u m f +UMG u m g +UMHB u m h b +UMH u m h +UMKC's u m k c's +UMKC u m k c +UMK u m k +Umla u m l a +UML u m l +Umme u m m e +UMMSM u m m s m +UMNO's u m n o's +Umno u m n o +UMNO u m n o +UMN u m n +UMO u m o +UMPP u m p p +UMP's u m p's +Umptee- u m p t e e +Umri u m r i +UM's u m's +UMTS u m t s +UMTV u m t v +UMUC u m u c +U. M. u m +U.M. u m +UMWA's u m w a's +UMWA u m w a +UMW's u m w's +UNAB u n a b +UNAFF u n a f f +UNAH u n a h +UNAMIR u n a m i r +UNAPACK's u n a p a c k's +UNB u n b +UNCAC u n c a c +UNCAF u n c a f +UNCG's u n c g's +UNCG u n c g +UNCHE u n c h e +UNCHS u n c h s +UNCMAC u n c m a c +UNCRO u n c r o +UNC's u n c's +UNCTAD u n c t a d +UNC- u n c +UNC u n c +UNCW u n c w +UNFCCC u n f c c c +UNFICYP u n f i c y p +UNFPA u n f p a +UNFP u n f p +UNF- u n f +UNF u n f +UNFWP u n f w p +UNGC u n g c +UNHCR u n h c r +UNHRC's u n h r c's +UNH u n h +UNJLC u n j l c +Unli u n l i +UNL u n l +Unlu u n l u +UNLV u n l v +UNMCK u n m c k +UNMC- u n m c +UNMC u n m c +UNMEE's u n m e e's +UNMEE u n m e e +UNMIH u n m i h +UNMIK u n m i k +UNMIL u n m i l +UNMOVIC u n m o v i c +unm u n m +UNM u n m +UNNC u n n c +unnd u n n d +UNODC u n o d c +UNOMIG u n o m i g +UNPA u n p a +UNPD u n p d +UNPO u n p o +UNP u n p +UNRA u n r a +UNRSGN u n r s g n +UNRWA u n r w a +UNSCOM u n s c o m +UNSCR u n s c r +UNSC's u n s c's +UNSC u n s c +UNSh u n s h +Unst u n s t +Uns u n's +U.N.'s u n's +UNSW u n s w +UNTSO u n t s o +UNTS u n t s +Unt u n t +UNT u n t +U.N. u n +UNU u n u +UNV u n v +UNWE u n w e +UNWTO u n w t o +Unz u n z +UOCl u o c l +UODA u o d a +UofM u o f m +UOHI u o h i +Uoho u o h o +UOJCA u o j c a +UOL u o l +UOMINI u o m i n i +UoMs u o m's +UOP u o p +uORF u o r f +UOR u o r +UO's u o's +UO u o +UoW u o w +UPABA u p a b a +UPBs u p b's +UPCI u p c i +UPC u p c +UPDRS u p d r s +UPDS u p d s +Upd u p d +UPD u p d +UPEC u p e c +UPFA- u p f a +UPFA u p f a +UPF's u p f's +Upf u p f +UPF u p f +UPGMA u p g m a +UPHSD u p h s d +UPIN u p i n +Upir u p i r +UPI's u p i's +UPI u p i +UPLA u p l a +UPLB u p l b +UPL u p l +UPMC u p m c +UPMSS u p m s s +UPND u p n d +UPNE u p n e +UPnP u p n p +UPN u p n +UPOV u p o v +UPPP u p p p +Uppu u p p u +UPRR u p r r +UPR u p r +UPSC u p s c +U.P.S.E.B. u p s e b +UPSIDC u p s i d c +UPSID u p s i d +UPSI u p s i +UPSMF u p s m f +UPTI u p t i +UPTU u p t u +U.P. u p +UPW u p w +UPX u p x +UPyD u p y d +UQAIB u q a i b +UQAM's u q a m's +UQC u q c +U.Q.P. u q p +UQP u q p +uq u q +UQ u q +URAA u r a a +URAC u r a c +URAP- u r a p +URAP u r a p +URBACT u r b a c t +URBED u r b e d +Urbz u r b z +URCL u r c l +URCSA u r c s a +URCs u r c's +URC u r c +Urdd u r d d +Urei u r e i +urf u r f +Urla u r l a +URLhttp u r l h t t p +URLs u r l's +URL u r l +URNU u r n u +UROC u r o c +URRACA u r r a c a +URSB u r s b +URSEC u r s e c +URSS u r s s +URW's u r w's +URW u r w +URZ u r z +USAACE u s a a c e +USAAC u s a a c +U.S.A.A.F. u s a a f +USAAF u s a a f +USAB u s a b +USACE u s a c e +USAC u s a c +USAFE u s a f e +USAF u s a f +USAHA u s a h a +USAID's u s a i d's +USAID- u s a i d +USAID u s a i d +Usal u s a l +usan u s a n +Usan u s a n +USAN u s a n +usao u s a o +USAO u s a o +USAPA u s a p a +USARPS u s a r p s +USASA u s a s a +USA's u s a's +usata u s a t a +USATF u s a t f +USAT u s a t +usa u s a +.USA u s a +U.S.A.. u s a +U.S.A. u s a +U.S.A u s a +USA- u s a +USA. u s a +USA u s a +USAya u s a y a +USBA u s b a +USBC's u s b c's +USBC u s b c +USBs u s b's +USB- u s b +USB u s b +USBWA u s b w a +USCCB u s c c b +USCF u s c f +USCGAUX u s c g a u x +USCGC u s c g c +USCG u s c g +USCHS u s c h s +Uscie u s c i e +USCIRF u s c i r f +USCIS u s c i s +USCITA u s c i t a +USCOB u s c o b +USCRP u s c r p +U.S.C.'s u s c's +USC's u s c's +USCT u s c t +U.S.C. u s c +USC- u s c +USC u s c +USDA's u s d a's +U.S.D.A. u s d a +USDA- u s d a +USDA u s d a +USDoE u s d o e +USDOJ u s d o j +USDP u s d p +USDTV u s d t v +usd u s d +USD u s d +USEER u s e e r +USENIX u s e n i x +USFbA u s f b a +USFE u s f e +USFF u s f f +USFLPA u s f l p a +USFL u s f l +USFS u s f s +USF u s f +USF&WS u s f and w s +USFWS u s f w s +USFW u s f w +USGA u s g a +U.S.G.S. u s g s +USGS- u s g s +USGS u s g s +U.S. H.R. u s h r +UShs u s h's +Usia u s i a +USIA u s i a +usih u s i h +USIP's u s i p's +USIP u s i p +USISL u s i s l +USIS u s i s +USIU u s i u +USJA u s j a +USJ u s j +USLES u s l e s +USLHE u s l h e +USLMRA u s l m r a +USLS- u s l s +USLTA u s l t a +USL- u s l +USL u s l +Uslu u s l u +Usmar u s m a r +USMA u s m a +USMC M u s m c +USMCR u s m c r +USMC's u s m c's +USMC u s m c +USMLE u s m l e +USML u s m l +USMMA u s m m a +USM u s m +Usna u s n a +USNA u s n a +USNG u s n g +Usnic u s n i c +USNM u s n m +USNO u s n o +USNR u s n r +USNSCS u s n s c s +USNS u s n s +USNTDP u s n t d p +U.S.N. u s n +USN u s n +USOC's u s o c's +USOC u s o c +Usoi u s o i +Usos u s o's +Uso u s o +USO u s o +USPHL u s p h l +USPO u s p o +USPPD u s p p d +U.S.P.Q. u s p q +USPSA u s p s a +USPs u s p's +USPS- u s p s +USPS u s p s +USPTA u s p t a +USPTO u s p t o +U.S.P. u s p +USP u s p +USRA u s r a +USRC u s r c +USSA u s s a +USSBA u s s b a +USSB u s s b +USSD u s s d +USSF's u s s f's +USSF u s s f +USSR's u s s r's +USSR u s s r +U.S.'s u s's +US's u s's +U.S.S. u s s +USS u s s +Usti u s t i +Ustka u s t k a +USTR's u s t r's +USTR u s t r +UST's u s t's +Usui u s u i +usum u s u m +U. S. u s +U.S. u s +Usut u s u t +USVC u s v c +Usvit u s v i t +USVI u s v i +U.S.V. u s v +USWA u s w a +USWNT's u s w n t's +USWNT u s w n t +U.S.W. u s w +USW u s w +Uta's u t a's +uta u t a +Uta u t a +UTA u t a +UTB u t b +Utca u t c a +UTC- u t c +UTC u t c +Utd's u t d's +Utd u t d +utea u t e a +UTEP u t e p +Utes u t e's +UTET u t e t +UTF- u t f +UTF u t f +Uth u t h +uti u t i +Uti u t i +UTJ u t j +UTMB u t m b +UTMF u t m f +UTM u t m +Utne u t n e +Utnur u t n u r +Uto u t o +UTO u t o +UTPB u t p b +UTP u t p +utrci u t r c i +UTRGV u t r g v +UTR u t r +utsav u t s a v +Utsav u t s a v +UTSC u t s c +UTSI u t s i +UT's u t's +UTS u t s +Utsu u t s u +UTTR u t t r +U.T. u t +Utu u t u +Utva u t v a +UTVA u t v a +UTV u t v +UUA u u a +UUCP u u c p +Uuh u u h +UUKKY u u k k y +Uul u u l +Uulu u u l u +UUP u u p +Uusi u u s i +UUs u u's +Uuto u u t o +UUT u u t +Uvac u v a c +UVB u v b +UVC u v c +Uvea u v e a +UVF u v f +UVIMB u v i m b +UVI u v i +uvnitr u v n i t r +Uvo u v o +UVT u v t +UVU u v u +U.V. u v +UV u v +UWAP u w a p +Uwasa u w a s a +UWASA u w a s a +U.W.A. u w a +UWA u w a +UWB u w b +UWC u w c +UWFi u w f i +UWF u w f +UWG u w g +UWH u w h +UWI u w i +UWM's u w m's +UWP u w p +UWRF u w r f +UWSP's u w s p's +UWSU u w s u +UW's u w's +UWS u w s +UW- u w +UW u w +UWW u w w +VASP v a s p +VA's v a's +V. A. v a +V.A. v a +VA- v a +VA v a +V&A v and a +Vav v a v +VAV v a v +VAW- v a w +Vay v a y +VBA v b a +VBBS v b b s +VBE v b e +VBF- v b f +VBI v b i +VBM v b m +VBR v b r +vb v b +V. B. v b +V.B. v b +VB- v b +VB v b +VCA's v c a's +VCA v c a +VCCS v c c s +VCDs v c d's +VCD v c d +VCE v c e +VCH v c h +VCIT v c i t +VCI v c i +VCJ v c j +V. C. M. v c m +VCO v c o +VCP v c p +VCQ v c q +VCRs v c r's +VCR v c r +VCSELs v c s e l's +VCSEL v c s e l +VCs v c's +VCS v c s +VCTs v c t's +VCT v c t +VCU v c u +vc v c +V. C. v c +V.C. v c +VC- v c +VC v c +VCXIV v c x i v +VCX- v c x +VCY v c y +VDAC v d a c +vda v d a +Vda v d a +VDCs v d c's +VDC v d c +vdiq v d i q +VDJ v d j +VDL's v d l's +VDL- v d l +VDL v d l +vdm v d m +VDM v d m +VDP v d p +VDR v d r +VDSL v d s l +VDUP v d u p +VDU v d u +V. D. v d +V.D. v d +VD- v d +VD v d +VDV v d v +V. D. W. v d w +ve v e +Ve v e +V. E. v e +VE v e +vez v e z +Vez v e z +VEZ v e z +V.F.A. v f a +VFA- v f a +VFA v f a +VFC v f c +VFDS v f d s +V.F.D. v f d +VFD v f d +VFIL v f i l +VFLI v f l i +VFL's v f l's +VFL v f l +VFMC v f m c +VFM v f m +vfp v f p +VFP- v f p +vfr v f r +VFR v f r +VF's v f's +VFS v f s +VFTS v f t s +VFTX v f t x +VFU v f u +V. F. v f +V.F. v f +VF- v f +VF v f +VFW- v f w +VFW v f w +VFX v f x +VGA v g a +VGC- v g c +V'Ger v g e r +VGIK v g i k +VGKO v g k o +VGMA v g m a +vgmdb v g m d b +VGMDB v g m d b +VGM v g m +VGo v g o +VGP v g p +Vgr- v g r +VGSCs v g s c's +VGSoM v g s o m +VGTU v g t u +VGU v g u +V. G. v g +V.G. v g +VG- v g +VG v g +VGV v g v +VGX v g x +VHB v h b +VHC v h c +VHD v h d +VHF v h f +VHL v h l +VHP v h p +VHR v h r +VHSL v h s l +VHS v h s +V. H. v h +VH- v h +VH v h +V.I.C.'s v i c's +Vict v i c t +VICT v i c t +viita v i i t a +vijf v i j f +vij v i j +Vij v i j +V.I.P.s v i p's +VIPs v i p's +V.I.P. v i p +VIP- v i p +VIP v i p +V. I. v i +V.I. v i +V. J. J. v j j +VJL v j l +VJs v j's +VJTF v j t f +V. J. v j +V.J. v j +VJ- v j +VJ v j +VKG v k g +VKh- v k h +VKIBC's v k i b c's +V. K. N. v k n +VKOS v k o s +VKS v k s +V. K. v k +V.K. v k +VK v k +Vlah v l a h +VLA v l a +VLBA v l b a +VLBI v l b i +VLB v l b +VLCCs v l c c's +VLCS v l c s +VLC v l c +VLDB v l d b +VLDL v l d l +vlei v l e i +VLE v l e +VLEX v l e x +VLF- v l f +VLF v l f +Vlijt v l i j t +VLN v l n +VLQ v l q +VLR v l r +VLSI v l s i +VLS v l s +VLT v l t +V. L. v l +V.L. v l +VL v l +VMAs v m a's +VMAT- v m a t +VMAT v m a t +VMA- v m a +VMA v m a +VMBO- v m b o +VMCAS v m c a s +VMCM v m c m +VME v m e +VMFM v m f m +VMF- v m f +VMF v m f +VMG v m g +VMI's v m i's +VMI v m i +VMK v m k +VMM- v m m +VMM v m m +VMO- v m o +VMPS v m p s +VMS v m s +VMT v m t +vm v m +V. M. v m +V.M. v m +VM- v m +VM v m +VMX- v m x +VNAF v n a f +VNAV v n a v +VNCH v n c h +VNC v n c +vner v n e r +V.N.E. v n e +VNI v n i +VNO- v n o +VNO v n o +VNPT v n p t +VNQDD v n q d d +VNR v n r +VNSA v n s a +VNS v n s +VNTH v n t h +VNTR v n t r +VNU v n u +V. N. v n +V.N. v n +VN- v n +VN v n +V.O. v o +VO v o +Vov v o v +VOX's v o x's +Voz v o z +VPAL v p a l +VPB- v p b +VPD v p d +VPH v p h +vpis v p i s +V. P. I. v p i +VPI v p i +VPLS v p l s +VPL v p l +VPM v p m +VPNs v p n's +VPN v p n +VPO's v p o's +VPP v p p +VPRO's v p r o's +VPRO v p r o +VPS v p s +vpu v p u +V. P. v p +V.P. v p +VP- v p +VP v p +VQT v q t +vq v q +VQ v q +vrak v r a k +VRA v r a +Vrba v r b a +Vrbuv v r b u v +VRB- v r b +VRD v r d +Vrej v r e j +vrem v r e m +VRE v r e +VRF v r f +Vrin v r i n +VRI v r i +vrj v r j +VRLA v r l a +vrn v r n +vroee v r o e e +vrou v r o u +vrouw v r o u w +Vrouw v r o u w +VRO v r o +VRR v r r +Vrsac v r s a c +VRSA v r s a +VR's v r's +VRS v r s +Vrtis v r t i's +VRT's v r t's +VRT v r t +V&R v and r +vr v r +V. R. v r +V.R. v r +VR- v r +VR v r +VRV v r v +VSANs v s a n's +VSAN v s a n +VSAP's v s a p's +VSATs v s a t's +VSA- v s a +VSB's v s b's +VSB v s b +VSCP v s c p +VSC v s c +VSD v s d +Vsekh v s e k h +vse v s e +VSE v s e +VSFV v s f v +VSG v s g +VSI v s i +VSM v s m +VSNL v s n l +vso v s o +VSO v s o +VSPs v s p's +VSP v s p +VSS v s s +VSTM v s t m +VST v s t +VSU v s u +vs. versus +_vs._ versus +VSV v s v +VSW v s w +VSX v s x +vsyo v s y o +VSZ v s z +VTA v t a +VTB v t b +VTCO v t c o +VTD v t d +VTEC v t e c +VTE v t eL +Vtic v t i c +VTi v t i +VTL v t l +VTM- v t m +VTM v t m +VTR v t r +VTSK v t s k +V&T's v and t's +VTS v t s +VTU v t u +V&T v and t +VTVL v t v l +V.T. v t +VT- v t +VT v t +VTV v t v +Vuur v u u r +VUU v u u +V.V.B. v v b +VVCS v v c s +VVD v v d +VVIPs v v i p's +VVI v v i +VVMC v v m c +VVN- v v n +VVO v v o +VVS v v s +VVT v v t +V. V. v v +V.V. v v +VV v v +VWI- v w i +VWM v w m +VWR v w r +VW's v w's +VWs v w's +VWS v w s +VW- v w +VW v w +VXL v x l +VXR v x r +VX- v x +VX v x +Vyg v y g +vyr v y r +vy v y +VY v y +vz v z +V. Z. v z +WAAAAAGH w a a a a a g h +WAAAF w a a a f +WAAC w a a c +WAAF w a a f +WAAI w a a i +WAAV w a a v +Waay w a a y +WABB w a b b +WABC- w a b c +WABC w a b c +WABD w a b d +WABG's w a b g's +WABG w a b g +W. A. C. w a c +W.A.C. w a c +Wa'il w a i l +WAIO w a i o +W.A.J. w a j +Wakf w a k f +W.A.K.O. w a k o +WAKR- w a k r +wa'l w a l +WANFL w a n f l +WAOK w a o k +waqf w a q f +waqt w a q t +Waqt w a q t +WAQZ w a q z +WARH w a r h +W.A.R. w a r +WASC w a s c +Wasl w a s l +WASL w a s l +W. A. S. w a s +W.A.S. w a s +Watfa w a t f a +WATL w a t l +WATW w a t w +wau w a u +Wau w a u +WAU w a u +WAVF w a v f +WAVM w a v m +W. A. w a +W.A. w a +WA w a +WAWLI w a w l i +waw w a w +WAWZ w a w z +WAYH w a y h +waza w a z a +Waza w a z a +WBAK- w b a k +WBAL's w b a l's +WBAL- w b a l +WBAM w b a m +WBA w b a +WBBL w b b l +WBBM- w b b m +WBBM w b b m +WBBP w b b p +WBCA w b c a +WBCN- w b c n +WBCN w b c n +W.B.C. w b c +WBC w b c +WBCY w b c y +WBEV w b e v +WBEZ w b e z +WBFF- w b f f +WBFJ w b f j +WBFO w b f o +WBF w b f +WBGU- w b g u +WBG w b g +WBHY w b h y +WBIC w b i c +WBIL w b i l +WBIR- w b i r +WBJEE w b j e e +WBKI- w b k i +WBKO- w b k o +WBKO w b k o +WBK w b k +WBLQ w b l q +WBME w b m e +WBNI w b n i +WBNX w b n x +WBON w b o n +WBO w b o +WBQC w b q c +WBRC w b r c +WBRE- w b r e +WBRE w b r e +WBRO w b r o +WBRZ w b r z +WBs w b's +WBS w b s +WBTI w b t i +WBT's w b t's +WBT w b t +WBTW w b t w +WBTZ w b t z +WBUR- w b u r +WBUR w b u r +WBVE w b v e +WBV w b v +W. B. w b +W.B. w b +WB- w b +WB w b +WBWT- w b w t +WBXH w b x h +WBZB w b z b +WBZT- w b z t +WBZ- w b z +WBZ w b z +WCAA w c a a +WCAC- w c a c +WCAG w c a g +WCAU- w c a u +WCAU w c a u +WCA w c a +WCBA w c b a +WCBI- w c b i +WCBN w c b n +W&CBR w and c b r +WCBR w c b r +WCBS- w c b s +WCBS w c b s +WCB w c b +WCCES w c c e s +WCCG w c c g +WCCO- w c c o +WCCO w c c o +WCCP- w c c p +WCC's w c c's +WCC w c c +WCDA w c d a +WCEO w c e o +WCFL w c f l +WCF w c f +WCG w c g +WCHA w c h a +WCHL's w c h l's +WCHL w c h l +WCHS w c h s +WCHV w c h v +WCIU w c i u +WCKD w c k d +WCKY w c k y +WCLN w c l n +WCLU w c l u +WCLV w c l v +WCL w c l +WCMA w c m a +WCMC w c m c +WCMFA w c m f a +WCMH w c m h +WCMR w c m r +WCM w c m +WCNB w c n b +WCNG- w c n g +WCOZ w c o z +WCPO w c p o +WCPW w c p w +WCRC w c r c +WCRE w c r e +WCRL w c r l +WCSA w c s a +WCSM- w c s m +WCSM w c s m +WCS w c s +WCSY- w c s y +WCTD w c t d +WCTG w c t g +WCTI- w c t i +W.C.T.U. w c t u +WCTU w c t u +WCTV w c t v +WCT w c t +WCTX w c t x +WCVB's w c v b's +WCVB w c v b +WCVE w c v e +WCVP- w c v p +WCWA w c w a +wc w c +W. C. w c +W.C. w c +WC- w c +WC w c +WCWC w c w c +WCW w c w +WCWW- w c w w +WCXR w c x r +WCYB w c y b +WCZE w c z e +WCZU- w c z u +WDAF- w d a f +WDAF w d a f +WDA w d a +WDCN- w d c n +WDC w d c +WDDM w d d m +WDFN w d f n +WDF w d f +WDGY w d g y +WDH w d h +WDIA w d i a +WDIF w d i f +WDIG w d i g +WDIS w d i s +WDIV w d i v +WDJM- w d j m +WDJT w d j t +WDLA w d l a +WDL w d l +WDM- w d m +WDM w d m +WDNN w d n n +WDNS w d n s +WDNT w d n t +WDOD- w d o d +WDOK- w d o k +WDOK w d o k +WDO w d o +WDRE w d r e +WDR w d r +WDSE w d s e +WDSN w d s n +WDSO w d s o +WDSU w d s u +WDS- w d s +WDS w d s +WDTM w d t m +WDUQ w d u q +WDVE w d v e +WDVR w d v r +W. D. w d +W.D. w d +WD- w d +WDWS w d w s +WDW w d w +WDXZ w d x z +WDYZ w d y z +W. E. B. w e b +W.E.B. w e b +W. E. H. w e h +W.E.H. w e h +WEVD- w e v d +WEVD w e v d +WEVV w e v v +W. E. w e +W.E. w e +WEZQ w e z q +WFAA- w f a a +WFAA w f a a +WFAL w f a l +WFAN w f a n +WFBC- w f b c +WFB w f b +WFCA w f c a +WFCB- w f c b +WFC w f c +WFDC w f d c +WFDD w f d d +WFDF w f d f +WFDU w f d u +WFD w f d +WFFF- w f f f +WFGF w f g f +WFHE w f h e +W.F.H. w f h +WFH w f h +WFIKKN w f i k k n +WFIL w f i l +WFIRST w f i r s t +WFLA- w f l a +WFLN- w f l n +WFL's w f l's +WFL w f l +WFLX w f l x +WFLZ- w f l z +WFMT w f m t +WFMU w f m u +W. F. M. w f m +WFM w f m +WFMZ- w f m z +WF&NW w f and n w +WFNZ w f n z +WFOY w f o y +WFPC w f p c +WFPF w f p f +WFPMA w f p m a +W.F.P. w f p +WFP w f p +WFRX- w f r x +WFSBP w f s b p +WFS w f s +WFTDA w f t d a +WFTM w f t m +WFTV w f t v +WFT w f t +WFVA w f v a +W&F w and f +WFWA w f w a +W. F. w f +W.F. w f +WF w f +WFWM w f w m +WFXR w f x r +WFXW- w f x w +wga w g a +WGA w g a +WGBH- w g b h +WGBH w g b h +WGBI w g b i +W. G. B. w g b +WGB w g b +WGCB w g c b +W.G.C. w g c +WGC- w g c +WGC w g c +WGEM- w g e m +WGEM w g e m +WGFT w g f t +WGGH w g g h +WGHP w g h p +WGH w g h +WGHW w g h w +WGI w g i +WGN's w g n's +WGNT w g n t +wgn w g n +WGN- w g n +WGN w g n +WGOK w g o k +WGO w g o +WGP w g p +WGPX w g p x +WGRM- w g r m +WGRP w g r p +WGRT w g r t +WGRV w g r v +WGRZ- w g r z +WGSA w g s a +WGSS w g s s +WGS- w g s +WGS w g s +WGTN- w g t n +WGUN w g u n +WGVU- w g v u +wg w g +W. G. w g +W.G. w g +WG w g +WGWG- w g w g +WGWW w g w w +WGZR w g z r +W.H.B. w h b +WHCA w h c a +WHCB w h c b +WHCC w h c c +W. H. C. w h c +W.H.C. w h c +WHC w h c +WHDG w h d g +WHDH- w h d h +WHDI w h d i +W.H.D. w h d +WHEC- w h e c +WHEC w h e c +WHFA w h f a +WHFM w h f m +WHFS- w h f s +WHFS w h f s +WHHS w h h s +W.H.H. w h h +WHJA w h j a +W.H.J. w h j +WHJY w h j y +WHKT w h k t +WHL's w h l's +WHL w h l +WHNE- w h n e +WHNO w h n o +WHP- w h p +WHP w h p +WHRC w h r c +WHSG- w h s g +WHSN w h s n +WHSS w h s s +W.H.S. w h s +WHS w h s +W. H. T. w h t +WHTZ w h t z +W&H w and h +W. H. w h +W.H. w h +WH w h +Whyld w h y l d +WHYN- w h y n +WHYP- w h y p +WHYT w h y t +W.I.B.C. w i b c +Wica w i c a +WICA- w i c a +wici w i c i +WIC's w i c's +WICU w i c u +Wif w i f +WIF w i f +WIHB- w i h b +WIJD w i j d +Wijk w i j k +WIKQ- w i k q +WIP w i p +W.I.S.E. w i s e +W.I.T.C.H. w i t c h +Wiwa w i w a +W. I. w i +WIXY's w i x y's +WJAC- w j a c +WJAL w j a l +WJAR w j a r +wjaz w j a z +WJBB- w j b b +WJBC w j b c +WJBF w j b f +WJBK- w j b k +WJBK w j b k +WJBR w j b r +WJCI w j c i +WJC w j c +WJEC w j e c +WJEF w j e f +WJER w j e r +WJET w j e t +W. J. J. O. w j j o +W.J.J.O. w j j o +WJJZ w j j z +WJKA's w j k a's +WJKA w j k a +WJKS w j k s +WJLA- w j l a +WJLA w j l a +WJMK w j m k +WJMN- w j m n +WJM w j m +WJOI w j o i +WJRD w j r d +WJR w j r +WJSC w j s c +W&J's w and j's +WJTD w j t d +WJUX w j u x +W&J w and j +W. J. w j +W.J. w j +WJ w j +WJXX w j x x +WJYL w j y l +WJYM w j y m +WJYS- w j y s +WJZA w j z a +WJZJ w j z j +WJZ w j z +WKAG w k a g +WKAQ- w k a q +WKAQ w k a q +WKAR- w k a r +WKAR w k a r +WKA w k a +WKBD w k b d +WKBF's w k b f's +WKB w k b +WKBW w k b w +W. K. C. w k c +WKDM w k d m +WKDN w k d n +WKDZ w k d z +WKFS w k f s +WKHG- w k h g +W.K.H. w k h +WKH w k h +WKIK w k i k +WKIP w k i p +WKKB w k k b +WKKD w k k d +WKKZ w k k z +WKLH w k l h +WKLK- w k l k +WKNA- w k n a +WKNDS w k n d s +WKNR- w k n r +WKN w k n +WKQI w k q i +WKQX- w k q x +WKRC- w k r c +WKRK- w k r k +WKRR w k r r +WKSE w k s e +WKSM w k s m +WKSU w k s u +WKTU's w k t u's +WKTU w k t u +W.K.V. w k v +W. K. w k +W.K. w k +WKXJ w k x j +WKX w k x +WKYS's w k y s's +WKYT w k y t +WKY w k y +W.L.A.K. w l a k +WLB w l b +WLCH w l c h +WLDM- w l d m +W&LE w and l e +Wley w l e y +WLFM- w l f m +WLGH w l g h +WLH w l h +WLHZ- w l h z +WLIE- w l i e +WLII w l i i +WLIR w l i r +WLIX- w l i x +WLJY w l j y +WLKT w l k t +WLKY w l k y +WLLM w l l m +WLNE- w l n e +WLNP w l n p +WLNQ w l n q +WLNS w l n s +WLPA w l p a +WLPR w l p r +WLPWR w l p w r +wlrs w l r s +WLSG w l s g +WLS- w l s +WLS w l s +WLTM w l t m +WLTQ- w l t q +WLUK's w l u k's +WLUK w l u k +WLUP- w l u p +WLWC w l w c +W. L. w l +W.L. w l +WL- w l +WL w l +WLW- w l w +WLW w l w +WLYH's w l y h's +WLYN w l y n +WLY w l y +WLZK w l z k +WMAC w m a c +WMAG w m a g +WMAN- w m a n +WMAP w m a p +WMAQ- w m a q +WMAQ w m a q +WMAR- w m a r +WMATA's w m a t a's +WMATA w m a t a +W.M.A. w m a +WMA w m a +WMAZ w m a z +WMBD- w m b d +WMBG w m b g +WMCAT w m c a t +WMCA w m c a +WMCT- w m c t +WMDB w m d b +WMDI w m d i +WMD w m d +WMEC- w m e c +WMEE w m e e +WMEN w m e n +WMET w m e t +WME w m e +WMEX w m e x +WMFC w m f c +WMFM w m f m +WMFP w m f p +W.M.F. w m f +WMF w m f +WMGF w m g f +WMG w m g +WMGY w m g y +WMHG w m h g +WMHT w m h t +WMIX w m i x +WMJR w m j r +WMJZ w m j z +WMKG- w m k g +WMKQ w m k q +WMKW w m k w +wml w m l +WMLW w m l w +WMLY w m l y +WMLZ- w m l z +WMMS w m m s +W. M. M. w m m +WMNF w m n f +WMNI- w m n i +WMN w m n +WMOB w m o b +WMO w m o +WMRA w m r a +WMRO w m r o +WMSG w m s g +W.M.S. w m s +WMS w m s +WMTL- w m t l +WMTS w m t s +WMT w m t +WMTW w m t w +WMUB w m u b +WMU w m u +WMVM w m v m +WMVP- w m v p +WMV w m v +W&M w and m +W. M. w m +W.M. w m +WM- w m +WM w m +WMWM w m w m +WMXB w m x b +WMX w m x +WMYE w m y e +WMYL w m y l +WNAK w n a k +W.N.B.A. w n b a +WNBA w n b a +WNBC w n b c +WNBL w n b l +WNBR w n b r +WNB w n b +WNCN w n c n +WNCT- w n c t +WNC w n c +WNCX w n c x +WND w n d +WNED w n e d +WNEG- w n e g +WNEL w n e l +WNEM w n e m +WNEP w n e p +WNET w n e t +WNEW's w n e w's +WNEW w n e w +WNFM w n f m +WNGL w n g l +WNHC- w n h c +WNIC- w n i c +WNIT w n i t +WNJN w n j n +WNJO w n j o +WNJY w n j y +WNK w n k +WNLO w n l o +WNLU w n l u +WNL w n l +WNMX w n m x +WNNS- w n n s +WNO w n o +WNPT w n p t +WNRN w n r n +WNR's w n r's +WNRV w n r v +WNSH w n s h +WNSR w n s r +WNTM w n t m +Wnt- w n t +WNT w n t +WNUA w n u a +WNUF w n u f +WNUSP w n u s p +WNWBL w n w b l +W. N. w n +W.N. w n +WN w n +WNWO w n w o +WNWS w n w s +WNW w n w +WNYA w n y a +WNYC's w n y c's +WNYC w n y c +WNYM w n y m +WNYO- w n y o +WNY w n y +WNYW- w n y w +WNYZ- w n y z +WNYZ w n y z +wnzaa w n z a a +WOAA w o a a +WOAK w o a k +WOA w o a +WOBC- w o b c +WOC w o c +W&OD w and o d +W.O.G. w o g +WOGX's w o g x's +WOJG w o j g +WOKZ- w o k z +WOLH w o l h +WOOM w o o m +WOPC w o p c +Worh w o r h +WOSM w o s m +WOTM- w o t m +W. O. w o +W.O. w o +WOWOW w o w o w +WOXY- w o x y +WOXY w o x y +WOZN w o z n +WPAC w p a c +WPA- w p a +WPA w p a +WPBC w p b c +WPBF- w p b f +WPBSA w p b s a +WPBS w p b s +WPBT w p b t +WPCH- w p c h +WPCQ- w p c q +W. P. C. w p c +WPC w p c +WPDH w p d h +WPD w p d +WPFF w p f f +WPFL w p f l +WPF w p f +WPGC w p g c +WPGG- w p g g +WPHG- w p h g +WPHL w p h l +WPH w p h +WPIAL w p i a l +WPI w p i +WPIX- w p i x +WPIX w p i x +WPLA w p l a +WPLB- w p l b +WPLG w p l g +WPLS w p l s +WPMW w p m w +WPO w p o +WPPT w p p t +WPP- w p p +WPP w p p +WPRI w p r i +WPRO w p r o +WPRP w p r p +WPSC w p s c +WPSL w p s l +WPS w p s +WPTB w p t b +WPTD's w p t d's +wpt w p t +WPT w p t +WPTY w p t y +W. P. U. w p u +wp w p +W. P. w p +W.P. w p +WP- w p +WP w p +WPXN w p x n +WPYR w p y r +WQDE w q d e +WQDR- w q d r +WQHT w q h t +WQHU- w q h u +WQIV w q i v +WQLR w q l r +WQLT- w q l t +WQLZ w q l z +WQMS w q m s +WQOK w q o k +WQSJ- w q s j +WQSO w q s o +WQSX w q s x +WQVC- w q v c +WQ w q +WQXI w q x i +WQXR w q x r +W. R. A. S. w r a s +WRAT w r a t +WRBU w r b u +WRCB w r c b +WRCN w r c n +WRC- w r c +WRC w r c +WRDL w r d l +WRDW- w r d w +Wrec w r e c +WRFD w r f d +WRFK w r f k +WRFM w r f m +WRFX w r f x +WRGO w r g o +WRGP w r g p +WRGT- w r g t +W. R. G. w r g +WRHA w r h a +WRHC w r h c +WRHI w r h i +W.R.I. w r i +WRJK- w r j k +WRJN- w r j n +WRKA w r k a +WRKO w r k o +WRKT w r k t +W.R.K. w r k +WRKY w r k y +WRNS w r n s +WRN w r n +WRNY w r n y +WRNZ w r n z +WROJ w r o j +WROK w r o k +WROO w r o o +WRPT w r p t +WRQQ w r q q +WRRK- w r r k +WRSW- w r s w +WRTA w r t a +WRTC w r t c +WRUV w r u v +WRU w r u +WRVL w r v l +WRVS w r v s +WRWB w r w b +W. R. w r +W.R. w r +WR- w r +WR w r +W.R.W. w r w +WRXT w r x t +WSAR w s a r +WSAV w s a v +W.S.A. w s a +WSA w s a +WSAZ w s a z +WSBG w s b g +WSBM w s b m +WSBT w s b t +WSB- w s b +WSB w s b +WSCA's w s c a's +WSCR w s c r +W. S. C. w s c +WSC w s c +WSDD w s d d +WSDG w s d g +WSDOT w s d o t +WSDP w s d p +WSDR w s d r +WSEE- w s e e +WSEG w s e g +WSEK w s e k +WSER- w s e r +WSE w s e +WSFA w s f a +WSFL w s f l +WSFM w s f m +WSFS w s f s +WSFX w s f x +WSGE w s g e +WSGW w s g w +WSHL w s h l +WSHU w s h u +WSH w s h +WSIPC w s i p c +WSIS w s i s +WSJM w s j m +WSJO w s j o +WSJ w s j +WSKG- w s k g +WSK- w s k +WSLU w s l u +WSL w s l +WSMG- w s m g +WSMK w s m k +W.S.M. w s m +WSM- w s m +WSM w s m +WSNJ w s n j +WSNS w s n s +WSNT w s n t +WSN w s n +WSNX- w s n x +WSOF w s o f +WSOPE w s o p e +WSOP's w s o p's +WSOP, w s o p +WSOP w s o p +WSORR w s o r r +WSO w s o +WSOY w s o y +WSPA- w s p a +WSPA w s p a +WSPC w s p c +WSPD w s p d +WSPN w s p n +WSPS w s p s +WSPU w s p u +WSP w s p +WSRA w s r a +WSRE's w s r e's +WSRF w s r f +WSSA w s s a +WSSCC w s s c c +WSSSA- w s s s a +WSSSA w s s s a +WSS w s s +WSTE w s t e +WSTR- w s t r +WSTR w s t r +WSUA w s u a +WSUD w s u d +WSUE w s u e +WSUS w s u s +WSU w s u +WSVF- w s v f +WSV w s v +WSVZ w s v z +WSWI w s w i +ws w s +W. S. w s +W.S. w s +WS- w s +WS w s +WSW w s w +WSYN w s y n +WSYX w s y x +WSYY- w s y y +WSZE- w s z e +WTAA w t a a +WTAB- w t a b +WTAE- w t a e +WTAE w t a e +WTAF w t a f +WTAG w t a g +WTAM's w t a m's +W. T. A. w t a +WTA w t a +WTAW w t a w +WTCC w t c c +W. T. C. w t c +W.T.C. w t c +WTC w t c +WTDY w t d y +WTEL w t e l +WTEV's w t e v's +WTEV w t e v +WTF w t f +WTGB- w t g b +WTGB w t g b +WTHR's w t h r's +WTHS w t h s +WTIU w t i u +WTI w t i +WTIX w t i x +WTKG w t k g +WTKO w t k o +WTKR w t k r +WTKS- w t k s +WTKS w t k s +WTLS w t l s +WTMA w t m a +WTNH w t n h +WTNO- w t n o +WTN w t n +WTOK- w t o k +WTOP- w t o p +WTOP w t o p +WTO's w t o's +WTOS w t o s +WTO w t o +WTPA w t p a +WTPL- w t p l +WTPM- w t p m +WTPQ w t p q +WTRF- w t r f +WTR w t r +WTSP- w t s p +WTSQ- w t s q +WTs w t's +WTS w t s +WTTF w t t f +WTTM w t t m +WTTs w t t's +WTT w t t +WTTW w t t w +WTTZ w t t z +WTVA- w t v a +WTVA w t v a +WTVC w t v c +WTVD w t v d +WTVF- w t v f +WTVI w t v i +WTVM w t v m +WTVU w t v u +WTVW w t v w +WTVX w t v x +WTWF w t w f +W. T. w t +W.T. w t +WT- w t +WT w t +WTX w t x +WUAB w u a b +WUAP w u a p +WUAs w u a's +WUC w u c +Wudl w u d l +Wuhr w u h r +Wuhu w u h u +WUKB w u k b +WUPG w u p g +WUPN w u p n +WUPZ w u p z +Wu's w u's +WUSW w u s w +WUTK w u t k +WUVC w u v c +wuv w u v +W. U. w u +Wuz w u z +WVAN- w v a n +WVa w v a +WVBX w v b x +WVCA w v c a +WVFJ- w v f j +WVIR- w v i r +WVIT w v i t +WVIZ- w v i z +WVJC w v j c +WVJP w v j p +WVLZ w v l z +WVNI w v n i +WVNS- w v n s +WVNY w v n y +W. V. O. w v o +WVO w v o +WVPA w v p a +WVPO w v p o +WVPX- w v p x +WVSSAC w v s s a c +WVST- w v s t +WVSU w v s u +WVTF- w v t f +WVTU's w v t u's +WVTU w v t u +WVUE w v u e +WVU w v u +WVUW w v u w +W. V. w v +W.V. w v +WV- w v +WV w v +W- w +WWAA w w a a +WWAC- w w a c +WWAMI w w a m i +WWASP w w a s p +WWA w w a +WWAX w w a x +WWBZ w w b z +WWCD w w c d +WWC w w c +WWDB's w w d b's +WWDC- w w d c +WWDC w w d c +WWDR w w d r +WWD w w d +WWEN- w w e n +WWE's w w e's +Wwe w w e +WWE w w e +WWF's w w f's +WWF- w w f +WWF w w f +WWGC w w g c +WWHT w w h t +WWHW w w h w +WWIL- w w i l +WWIS w w i s +WWJP w w j p +WWJ- w w j +WWKB w w k b +WWKL w w k l +WWML w w m l +WWMM w w m m +WWMX w w m x +WWMY w w m y +WWN w w n +WWNW w w n w +WWNY w w n y +WWOX w w o x +WWPG w w p g +WWP w w p +WWPW w w p w +WWRFC w w r f c +WWRV w w r v +WWSS w w s s +WWS w w s +WWTI's w w t i's +WWTI w w t i +WWTN w w t n +WWU w w u +WWVA w w v a +WWVR w w v r +WWV w w v +WWWF w w w f +WWWQ's w w w q's +W. W. w w +W.W. w w +WW- w w +WW w w +WWW w w w +WWWW w w w w +WWYD w w y d +WXB w x b +WXCN w x c n +WXCY- w x c y +WXDU w x d u +WXEL- w x e l +WXEL w x e l +WXII- w x i i +WXIL w x i l +WXIN's w x i n's +WXIN w x i n +WXJM's w x j m's +WXJ w x j +WXL- w x l +WXN w x n +WXQW w x q w +WXRK- w x r k +WXSR- w x s r +WXTG w x t g +WXTH- w x t h +WXW w x w +wx w x +WX w x +W.Y. w y +WY w y +wyzc w y z c +WZBC- w z b c +WZBC w z b c +WZFG w z f g +WZFT w z f t +WZGC w z g c +WZJQ w z j q +WZKC w z k c +WZLR w z l r +WZLX- w z l x +WZLX w z l x +WZRB w z r b +WZRD w z r d +WZST w z s t +WZT- w z t +W. Z. w z +WZ w z +WZX w z x +WZZM- w z z m +WZZM w z z m +WZZX w z z x +Xav x a v +XAV x a v +xa x a +Xa x a +XA- x a +XA x a +xbg x b g +XBMC x b m c +XBM x b m +Xbra x b r a +XBRL x b r l +XBV x b v +XB- x b +XB x b +XCB x c b +XCE x c e +XCMS x c m s +XCOM x c o m +XCOR x c o r +xCo x c o +XCR x c r +XCT x c t +XCVI x c v i +XCV x c v +X. C. x c +XC- x c +XC x c +XCX x c x +xda x d a +XDDM x d d m +XDE- x d e +XDH x d h +XDK x d k +XDRs x d r's +XDR- x d r +XDR x d r +XD's x d's +XDS x d s +X. D. x d +XD x d +Xfce x f c e +XFC x f c +XFDL x f d l +XFD x f d +XFL x f l +XFM x f m +XFR- x f r +XFS x f s +XFU x f u +xf x f +X.F. x f +XF- x f +XF x f +XHVFC x h v f c +xh x h +X. H. x h +XH x h +xix x i x +XIX- x i x +XJD x j d +XJ- x j +XJ x j +XKE x k e +XKL x k l +xk x k +XK x k +XL x l +XMB x m b +xmc x m c +xml x m l +XML- x m l +XML x m l +XMMS x m m s +XMM- x m m +XMPP x m p p +XMP x m p +XMRV x m r v +XMSN x m s n +XM's x m's +XMU x m u +X.M. x m +XM x m +XN x n +xO x o +X.O. x o +XO- x o +XO x o +XPA x p a +XPD x p d +XPe x p e +XPL x p l +XPN x p n +XP's x p's +XPS x p s +XPT x p t +XPW x p w +XP- x p +XP x p +X.Q. x q +XQ x q +XRCC x r c c +XRCO x r c o +XRDS x r d s +XRD x r d +XRE x r e +XR x r +XSD x s d +xSi- x s i +XSi x s i +XSI x s i +XSLT x s l t +XSL- x s l +xsr x s r +XS x s +XTB x t b +XTC x t c +XTG x t g +XTwas x t w a's +xt x t +XT x t +XVAS x v a +XV- x v +XV x v +XWA x w a +XWB x w b +XWF x w f +XWIS x w i s +X. W. x w +XYNU x y n u +XYP x y p +xy x y +XY x y +xyz x y z +xyZ x y z +XYZ x y z +XZ x z +Y. A. y a +Y.A. y a +YA- y a +YBA y b a +YBL y b l +YBNL y b n l +YBN y b n +YBR y b r +YBT- y b t +YBX y b x +Y. B. y b +Y.B. y b +YB- y b +YB y b +YBYS y b y s +YCC y c c +YCD y c d +ycia y c i a +ycie y c i e +YCL y c l +YCO y c o +YCU y c u +YCV y c v +yc y c +Y. C. y c +Y.C. y c +YC- y c +YC y c +YCYW y c y w +YDA y d a +Ydby y d b y +Yde y d e +YDG y d g +YdiB y d i b +YDK y d k +YDM- y d m +YDNC y d n c +yd y d +Y. D. y d +Y.D. y d +YD y d +Y.E. y e +YFCMD y f c m d +YFCs y f c's +YFCU y f c u +YFC- y f c +YFC y f c +YFD- y f d +yfle y f l e +YF&R y f and r +YF- y f +YF y f +YGB y g b +YGEX y g e x +YGLA y g l a +YGL y g l +YG's y g's +yg y g +Y. G. y g +Y.G. y g +YG- y g +YG y g +YHA y h a +YHWH y h w h +Y. H. y h +YH- y h +YH y h +YIF y i f +YIG y i g +YITP y i t p +Y. I. y i +Y.I. y i +YJO y j o +yj y j +Y. J. y j +YJ- y j +YJ y j +Ykt y k t +Y. K. y k +Y.K. y k +YK y k +YKY y k y +Yle y l e +YLE y l e +YLF y l f +yl- y l +yl y l +Y. L. y l +YL- y l +YL y l +YMCA's y m c a's +Y.M.C.A. y m c a +YMCA. y m c a +YMCA y m c a +YMF y m f +Ymke y m k e +Y.M.M. y m m +YMM y m m +YMO y m o +YMS- y m s +ym y m +Y. M. y m +YM y m +yndi y n d i +Yndi y n d i +yne y n e +Yngve y n g v e +Ynis y n i's +YNK y n k +YNU y n u +yn y n +Yn y n +Y. N. y n +Y.N. y n +YN- y n +YN y n +Y.O. y o +YO y o +ypa y p a +YPC y p c +YPF y p f +YPG y p g +YPM y p m +YPO y p o +Y.P.P. y p p +YPP y p p +Y.P.S. y p s +YPT- y p t +YPT y p t +YPU y p u +Y.P. y p +YP- y p +YP y p +Y. Q. y q +YQ y q +yra y r a +Yra y r a +YRC y r c +YRF y r f +YRK y r k +YRM y r m +YRT y r t +Y&R y and r +YSD- y s d +YSK'da y s k d a +YSK'ya y s k y a +YSK y s k +YSL y s l +YSRCP y s r c p +YSR y s r +YSU's y s u's +YSU y s u +Ys y's +Y.S. y s +YS- y s +YS y s +YSY y s y +yta y t a +YTB y t b +Ytre y t r e +YTS y t s +YTV's y t v's +YTV y t v +Y&T y and t +Y. T. y t +Y.T. y t +YT y t +Y. V. y v +YWCA- y w c a +YWCA. y w c a +YWCA y w c a +YWC y w c +Y. W. y w +Y.Y. y y +Z. A. B. z a b +Z. A. z a +Z.A. z a +ZBC z b c +ZBS z b s +ZB z b +ZCE z c e +ZCGWM z c g w m +ZCMI z c m i +ZCTA z c t a +ZCY z c y +ZC z c +Zdar z d a r +zda z d a +ZDF z d f +ZDR z d r +Z.E.R.O. z e r o +ZFC z f c +ZFK z f k +ZFM z f m +ZFP z f p +ZFS z f s +ZFX z f x +ZF z f +ZGE z g e +ZGI z g i +zg z g +Z.G. z g +ZG z g +zh z h +Z. H. z h +Z.H. z h +ZH z h +zijn z i j n +zij z i j +Zij z i j +Z. I. z i +Z. J. z j +ZKK z k k +ZKM z k m +Z.K. z k +Z.L. z l +Z.M. z m +ZNA z n a +ZNF z n f +ZNK z n k +ZnO z n o +ZNW z n w +ZPAP z p a p +Zpav z p a v +ZPAV z p a v +ZPA z p a +ZPG- z p g +ZPK z p k +ZPU- z p u +Z.P. z p +ZP- z p +ZP z p +ZQ z q +ZRC z r c +ZrI z r i +ZRK z r k +ZSC z s c +ZSE z s e +ZSF z s f +ZSG's z s g's +ZSIS z s i s +ZSL z s l +ZSO z s o +ZSSS z s s s +Zsuzsa z s u z s a +ZSU- z s u +ZTE z t e +ZWYCI z w y c i +Z.W. z w +Z.X.V. z x v +Z.X. z x +ZYF z y f +Z. Y. z y +Z.Y. z y +Z. Z. z z +Z.Z. z z diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data/whitelist_alternatives.tsv b/models/tts/NeMo/nemo_text_processing/text_normalization/data/whitelist_alternatives.tsv new file mode 100644 index 0000000000..8cea774c58 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data/whitelist_alternatives.tsv @@ -0,0 +1,17 @@ +Hon. Honorable +Hon. Honourable +St. Street +St Street +St. Saint +St Saint +Dr. Drive +Dr. Doctor +Mr. Mister +Mrs. Misses +Ms. Miss +Mr Mister +Mrs Misses +Ms Miss +&Co. and Co. +§ section += equals \ No newline at end of file diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/data_loader_utils.py b/models/tts/NeMo/nemo_text_processing/text_normalization/data_loader_utils.py new file mode 100644 index 0000000000..1aaac7a762 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/data_loader_utils.py @@ -0,0 +1,269 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 csv +import json +import os +import re +from collections import defaultdict, namedtuple +from typing import Dict, List, Optional, Set, Tuple + +EOS_TYPE = "EOS" +PUNCT_TYPE = "PUNCT" +PLAIN_TYPE = "PLAIN" +Instance = namedtuple('Instance', 'token_type un_normalized normalized') +known_types = [ + "PLAIN", + "DATE", + "CARDINAL", + "LETTERS", + "VERBATIM", + "MEASURE", + "DECIMAL", + "ORDINAL", + "DIGIT", + "MONEY", + "TELEPHONE", + "ELECTRONIC", + "FRACTION", + "TIME", + "ADDRESS", +] + + +def load_kaggle_text_norm_file(file_path: str) -> List[Instance]: + """ + https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish + Loads text file in the Kaggle Google text normalization file format: \t\t<`self` if trivial class or normalized text> + E.g. + PLAIN Brillantaisia + PLAIN is + PLAIN a + PLAIN genus + PLAIN of + PLAIN plant + PLAIN in + PLAIN family + PLAIN Acanthaceae + PUNCT . sil + + + Args: + file_path: file path to text file + + Returns: flat list of instances + """ + res = [] + with open(file_path, 'r') as fp: + for line in fp: + parts = line.strip().split("\t") + if parts[0] == "": + res.append(Instance(token_type=EOS_TYPE, un_normalized="", normalized="")) + else: + l_type, l_token, l_normalized = parts + l_token = l_token.lower() + l_normalized = l_normalized.lower() + + if l_type == PLAIN_TYPE: + res.append(Instance(token_type=l_type, un_normalized=l_token, normalized=l_token)) + elif l_type != PUNCT_TYPE: + res.append(Instance(token_type=l_type, un_normalized=l_token, normalized=l_normalized)) + return res + + +def load_files(file_paths: List[str], load_func=load_kaggle_text_norm_file) -> List[Instance]: + """ + Load given list of text files using the `load_func` function. + + Args: + file_paths: list of file paths + load_func: loading function + + Returns: flat list of instances + """ + res = [] + for file_path in file_paths: + res.extend(load_func(file_path=file_path)) + return res + + +def clean_generic(text: str) -> str: + """ + Cleans text without affecting semiotic classes. + + Args: + text: string + + Returns: cleaned string + """ + text = text.strip() + text = text.lower() + return text + + +def evaluate(preds: List[str], labels: List[str], input: Optional[List[str]] = None, verbose: bool = True) -> float: + """ + Evaluates accuracy given predictions and labels. + + Args: + preds: predictions + labels: labels + input: optional, only needed for verbosity + verbose: if true prints [input], golden labels and predictions + + Returns accuracy + """ + acc = 0 + nums = len(preds) + for i in range(nums): + pred_norm = clean_generic(preds[i]) + label_norm = clean_generic(labels[i]) + if pred_norm == label_norm: + acc = acc + 1 + else: + if input: + print(f"inpu: {json.dumps(input[i])}") + print(f"gold: {json.dumps(label_norm)}") + print(f"pred: {json.dumps(pred_norm)}") + return acc / nums + + +def training_data_to_tokens( + data: List[Instance], category: Optional[str] = None +) -> Dict[str, Tuple[List[str], List[str]]]: + """ + Filters the instance list by category if provided and converts it into a map from token type to list of un_normalized and normalized strings + + Args: + data: list of instances + category: optional semiotic class category name + + Returns Dict: token type -> (list of un_normalized strings, list of normalized strings) + """ + result = defaultdict(lambda: ([], [])) + for instance in data: + if instance.token_type != EOS_TYPE: + if category is None or instance.token_type == category: + result[instance.token_type][0].append(instance.un_normalized) + result[instance.token_type][1].append(instance.normalized) + return result + + +def training_data_to_sentences(data: List[Instance]) -> Tuple[List[str], List[str], List[Set[str]]]: + """ + Takes instance list, creates list of sentences split by EOS_Token + Args: + data: list of instances + Returns (list of unnormalized sentences, list of normalized sentences, list of sets of categories in a sentence) + """ + # split data at EOS boundaries + sentences = [] + sentence = [] + categories = [] + sentence_categories = set() + + for instance in data: + if instance.token_type == EOS_TYPE: + sentences.append(sentence) + sentence = [] + categories.append(sentence_categories) + sentence_categories = set() + else: + sentence.append(instance) + sentence_categories.update([instance.token_type]) + un_normalized = [" ".join([instance.un_normalized for instance in sentence]) for sentence in sentences] + normalized = [" ".join([instance.normalized for instance in sentence]) for sentence in sentences] + return un_normalized, normalized, categories + + +def load_labels(abs_path): + """ + loads relative path file as dictionary + + Args: + abs_path: absolute path + + Returns dictionary of mappings + """ + label_tsv = open(abs_path) + labels = list(csv.reader(label_tsv, delimiter="\t")) + return labels + + +def get_abs_path(rel_path): + """ + Get absolute path + + Args: + rel_path: relative path to this file + + Returns absolute path + """ + return os.path.dirname(os.path.abspath(__file__)) + '/' + rel_path + + +def post_process_punctuation(text: str) -> str: + """ + Normalized quotes and spaces + + Args: + text: text + + Returns: text with normalized spaces and quotes + """ + text = ( + text.replace('( ', '(') + .replace(' )', ')') + .replace('{ ', '{') + .replace(' }', '}') + .replace('[ ', '[') + .replace(' ]', ']') + .replace(' ', ' ') + .replace('”', '"') + .replace("’", "'") + .replace("»", '"') + .replace("«", '"') + .replace("\\", "") + .replace("„", '"') + .replace("´", "'") + .replace("’", "'") + .replace('“', '"') + .replace("‘", "'") + .replace('`', "'") + .replace('- -', "--") + ) + + for punct in "!,.:;?": + text = text.replace(f' {punct}', punct) + return text.strip() + + +def pre_process(text: str) -> str: + """ + Adds space around punctuation marks + + Args: + text: string that may include semiotic classes + + Returns: text with spaces around punctuation marks + """ + space_both = '*<=>^[]{}' + for punct in space_both: + text = text.replace(punct, ' ' + punct + ' ') + + text = text.replace('--', ' ' + '--' + ' ') + # remove extra space + text = re.sub(r' +', ' ', text) + return text diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/graph_utils.py b/models/tts/NeMo/nemo_text_processing/text_normalization/graph_utils.py new file mode 100644 index 0000000000..71cdc62917 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/graph_utils.py @@ -0,0 +1,216 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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 os +import string +from pathlib import Path + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path + +try: + import pynini + from pynini import Far + from pynini.examples import plurals + from pynini.lib import byte, pynutil, utf8 + + NEMO_CHAR = utf8.VALID_UTF8_CHAR + + NEMO_DIGIT = byte.DIGIT + NEMO_LOWER = pynini.union(*string.ascii_lowercase).optimize() + NEMO_UPPER = pynini.union(*string.ascii_uppercase).optimize() + NEMO_ALPHA = pynini.union(NEMO_LOWER, NEMO_UPPER).optimize() + NEMO_ALNUM = pynini.union(NEMO_DIGIT, NEMO_ALPHA).optimize() + NEMO_HEX = pynini.union(*string.hexdigits).optimize() + NEMO_NON_BREAKING_SPACE = u"\u00A0" + NEMO_SPACE = " " + NEMO_WHITE_SPACE = pynini.union(" ", "\t", "\n", "\r", u"\u00A0").optimize() + NEMO_NOT_SPACE = pynini.difference(NEMO_CHAR, NEMO_WHITE_SPACE).optimize() + NEMO_NOT_QUOTE = pynini.difference(NEMO_CHAR, r'"').optimize() + + NEMO_PUNCT = pynini.union(*map(pynini.escape, string.punctuation)).optimize() + NEMO_GRAPH = pynini.union(NEMO_ALNUM, NEMO_PUNCT).optimize() + + NEMO_SIGMA = pynini.closure(NEMO_CHAR) + + delete_space = pynutil.delete(pynini.closure(NEMO_WHITE_SPACE)) + insert_space = pynutil.insert(" ") + delete_extra_space = pynini.cross(pynini.closure(NEMO_WHITE_SPACE, 1), " ") + + suppletive = pynini.string_file(get_abs_path("data/suppletive.tsv")) + # _v = pynini.union("a", "e", "i", "o", "u") + _c = pynini.union( + "b", "c", "d", "f", "g", "h", "j", "k", "l", "m", "n", "p", "q", "r", "s", "t", "v", "w", "x", "y", "z" + ) + _ies = NEMO_SIGMA + _c + pynini.cross("y", "ies") + _es = NEMO_SIGMA + pynini.union("s", "sh", "ch", "x", "z") + pynutil.insert("es") + _s = NEMO_SIGMA + pynutil.insert("s") + + graph_plural = plurals._priority_union( + suppletive, plurals._priority_union(_ies, plurals._priority_union(_es, _s, NEMO_SIGMA), NEMO_SIGMA), NEMO_SIGMA + ).optimize() + + SINGULAR_TO_PLURAL = graph_plural + PLURAL_TO_SINGULAR = pynini.invert(graph_plural) + TO_LOWER = pynini.union(*[pynini.cross(x, y) for x, y in zip(string.ascii_uppercase, string.ascii_lowercase)]) + TO_UPPER = pynini.invert(TO_LOWER) + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + # Create placeholders + NEMO_CHAR = None + + NEMO_DIGIT = None + NEMO_LOWER = None + NEMO_UPPER = None + NEMO_ALPHA = None + NEMO_ALNUM = None + NEMO_HEX = None + NEMO_NON_BREAKING_SPACE = u"\u00A0" + NEMO_SPACE = " " + NEMO_WHITE_SPACE = None + NEMO_NOT_SPACE = None + NEMO_NOT_QUOTE = None + + NEMO_PUNCT = None + NEMO_GRAPH = None + + NEMO_SIGMA = None + + delete_space = None + insert_space = None + delete_extra_space = None + + suppletive = None + # _v = pynini.union("a", "e", "i", "o", "u") + _c = None + _ies = None + _es = None + _s = None + + graph_plural = None + + SINGULAR_TO_PLURAL = None + PLURAL_TO_SINGULAR = None + TO_LOWER = None + TO_UPPER = None + + PYNINI_AVAILABLE = False + + +def get_plurals(fst): + """ + Given singular returns plurals + + Args: + fst: Fst + + Returns plurals to given singular forms + """ + return SINGULAR_TO_PLURAL @ fst + + +def get_singulars(fst): + """ + Given plural returns singulars + + Args: + fst: Fst + + Returns singulars to given plural forms + """ + return PLURAL_TO_SINGULAR @ fst + + +def convert_space(fst) -> 'pynini.FstLike': + """ + Converts space to nonbreaking space. + Used only in tagger grammars for transducing token values within quotes, e.g. name: "hello kitty" + This is making transducer significantly slower, so only use when there could be potential spaces within quotes, otherwise leave it. + + Args: + fst: input fst + + Returns output fst where breaking spaces are converted to non breaking spaces + """ + return fst @ pynini.cdrewrite(pynini.cross(NEMO_SPACE, NEMO_NON_BREAKING_SPACE), "", "", NEMO_SIGMA) + + +class GraphFst: + """ + Base class for all grammar fsts. + + Args: + name: name of grammar class + kind: either 'classify' or 'verbalize' + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, name: str, kind: str, deterministic: bool = True): + self.name = name + self.kind = str + self._fst = None + self.deterministic = deterministic + + self.far_path = Path(os.path.dirname(__file__) + '/grammars/' + kind + '/' + name + '.far') + if self.far_exist(): + self._fst = Far(self.far_path, mode="r", arc_type="standard", far_type="default").get_fst() + + def far_exist(self) -> bool: + """ + Returns true if FAR can be loaded + """ + return self.far_path.exists() + + @property + def fst(self) -> 'pynini.FstLike': + return self._fst + + @fst.setter + def fst(self, fst): + self._fst = fst + + def add_tokens(self, fst) -> 'pynini.FstLike': + """ + Wraps class name around to given fst + + Args: + fst: input fst + + Returns: + Fst: fst + """ + return pynutil.insert(f"{self.name} {{ ") + fst + pynutil.insert(" }") + + def delete_tokens(self, fst) -> 'pynini.FstLike': + """ + Deletes class name wrap around output of given fst + + Args: + fst: input fst + + Returns: + Fst: fst + """ + res = ( + pynutil.delete(f"{self.name}") + + delete_space + + pynutil.delete("{") + + delete_space + + fst + + delete_space + + pynutil.delete("}") + ) + return res @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", NEMO_SIGMA) diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/normalize.py b/models/tts/NeMo/nemo_text_processing/text_normalization/normalize.py new file mode 100644 index 0000000000..353025a3a9 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/normalize.py @@ -0,0 +1,246 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 itertools +from argparse import ArgumentParser +from collections import OrderedDict +from typing import List + +from nemo_text_processing.text_normalization.data_loader_utils import post_process_punctuation, pre_process +from nemo_text_processing.text_normalization.taggers.tokenize_and_classify import ClassifyFst +from nemo_text_processing.text_normalization.token_parser import PRESERVE_ORDER_KEY, TokenParser +from nemo_text_processing.text_normalization.verbalizers.verbalize_final import VerbalizeFinalFst +from tqdm import tqdm + +try: + import pynini + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class Normalizer: + """ + Normalizer class that converts text from written to spoken form. + Useful for TTS preprocessing. + + Args: + input_case: expected input capitalization + """ + + def __init__(self, input_case: str): + assert input_case in ["lower_cased", "cased"] + + self.tagger = ClassifyFst(input_case=input_case, deterministic=True) + self.verbalizer = VerbalizeFinalFst(deterministic=True) + self.parser = TokenParser() + + def normalize_list(self, texts: List[str], verbose=False) -> List[str]: + """ + NeMo text normalizer + + Args: + texts: list of input strings + verbose: whether to print intermediate meta information + + Returns converted list input strings + """ + res = [] + for input in tqdm(texts): + try: + text = self.normalize(input, verbose=verbose) + except: + print(input) + raise Exception + res.append(text) + return res + + def normalize( + self, text: str, verbose: bool, punct_pre_process: bool = False, punct_post_process: bool = False + ) -> str: + """ + Main function. Normalizes tokens from written to spoken form + e.g. 12 kg -> twelve kilograms + + Args: + text: string that may include semiotic classes + verbose: whether to print intermediate meta information + punct_pre_process: whether to perform punctuation pre-processing, for example, [25] -> [ 25 ] + punct_post_process: whether to normalize punctuation + + Returns: spoken form + """ + if punct_pre_process: + text = pre_process(text) + text = text.strip() + if not text: + if verbose: + print(text) + return text + text = pynini.escape(text) + tagged_lattice = self.find_tags(text) + tagged_text = self.select_tag(tagged_lattice) + if verbose: + print(tagged_text) + self.parser(tagged_text) + tokens = self.parser.parse() + tags_reordered = self.generate_permutations(tokens) + for tagged_text in tags_reordered: + tagged_text = pynini.escape(tagged_text) + verbalizer_lattice = self.find_verbalizer(tagged_text) + if verbalizer_lattice.num_states() == 0: + continue + output = self.select_verbalizer(verbalizer_lattice) + if punct_post_process: + output = post_process_punctuation(output) + return output + raise ValueError() + + def _permute(self, d: OrderedDict) -> List[str]: + """ + Creates reorderings of dictionary elements and serializes as strings + + Args: + d: (nested) dictionary of key value pairs + + Return permutations of different string serializations of key value pairs + """ + l = [] + if PRESERVE_ORDER_KEY in d.keys(): + d_permutations = [d.items()] + else: + d_permutations = itertools.permutations(d.items()) + for perm in d_permutations: + subl = [""] + for k, v in perm: + if isinstance(v, str): + subl = ["".join(x) for x in itertools.product(subl, [f"{k}: \"{v}\" "])] + elif isinstance(v, OrderedDict): + rec = self._permute(v) + subl = ["".join(x) for x in itertools.product(subl, [f" {k} {{ "], rec, [f" }} "])] + elif isinstance(v, bool): + subl = ["".join(x) for x in itertools.product(subl, [f"{k}: true "])] + else: + raise ValueError() + l.extend(subl) + return l + + def generate_permutations(self, tokens: List[dict]): + """ + Generates permutations of string serializations of list of dictionaries + + Args: + tokens: list of dictionaries + + Returns string serialization of list of dictionaries + """ + + def _helper(prefix: str, tokens: List[dict], idx: int): + """ + Generates permutations of string serializations of given dictionary + + Args: + tokens: list of dictionaries + prefix: prefix string + idx: index of next dictionary + + Returns string serialization of dictionary + """ + if idx == len(tokens): + yield prefix + return + token_options = self._permute(tokens[idx]) + for token_option in token_options: + yield from _helper(prefix + token_option, tokens, idx + 1) + + return _helper("", tokens, 0) + + def find_tags(self, text: str) -> 'pynini.FstLike': + """ + Given text use tagger Fst to tag text + + Args: + text: sentence + + Returns: tagged lattice + """ + lattice = text @ self.tagger.fst + return lattice + + def select_tag(self, lattice: 'pynini.FstLike') -> str: + """ + Given tagged lattice return shortest path + + Args: + tagged_text: tagged text + + Returns: shortest path + """ + tagged_text = pynini.shortestpath(lattice, nshortest=1, unique=True).string() + return tagged_text + + def find_verbalizer(self, tagged_text: str) -> 'pynini.FstLike': + """ + Given tagged text creates verbalization lattice + This is context-independent. + + Args: + tagged_text: input text + + Returns: verbalized lattice + """ + lattice = tagged_text @ self.verbalizer.fst + return lattice + + def select_verbalizer(self, lattice: 'pynini.FstLike') -> str: + """ + Given verbalized lattice return shortest path + + Args: + lattice: verbalization lattice + + Returns: shortest path + """ + output = pynini.shortestpath(lattice, nshortest=1, unique=True).string() + return output + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("input_string", help="input string", type=str) + parser.add_argument( + "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + ) + parser.add_argument("--verbose", help="print info for debugging", action='store_true') + parser.add_argument( + "--punct_post_process", help="set to True to enable punctuation post processing", action="store_true" + ) + parser.add_argument( + "--punct_pre_process", help="set to True to enable punctuation pre processing", action="store_true" + ) + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + normalizer = Normalizer(input_case=args.input_case) + print( + normalizer.normalize( + args.input_string, + verbose=args.verbose, + punct_pre_process=args.punct_pre_process, + punct_post_process=args.punct_post_process, + ) + ) diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/normalize_with_audio.py b/models/tts/NeMo/nemo_text_processing/text_normalization/normalize_with_audio.py new file mode 100644 index 0000000000..4df248f2d4 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/normalize_with_audio.py @@ -0,0 +1,327 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 json +import os +import time +from argparse import ArgumentParser +from typing import List, Tuple + +from joblib import Parallel, delayed +from nemo_text_processing.text_normalization.data_loader_utils import post_process_punctuation, pre_process +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.taggers.tokenize_and_classify import ClassifyFst +from nemo_text_processing.text_normalization.verbalizers.verbalize_final import VerbalizeFinalFst + +from nemo.collections.asr.metrics.wer import word_error_rate +from nemo.collections.asr.models import ASRModel + +try: + import pynini + from pynini.lib import rewrite + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + +""" +The script provides multiple normalization options and chooses the best one that minimizes CER of the ASR output +(most of the semiotic classes use deterministic=False flag). + +To run this script with a .json manifest file: + python normalize_with_audio.py \ + --audio_data PATH/TO/MANIFEST.JSON \ + --model QuartzNet15x5Base-En \ + --verbose + + The manifest file should contain the following fields: + "audio_filepath" - path to the audio file + "text" - raw text + "transcript" - ASR model prediction (optional) + + +To run with a single audio file, specify path to audio and text with: + python normalize_with_audio.py \ + --audio_data PATH/TO/AUDIO.WAV \ + --text raw text OR PATH/TO/.TXT/FILE + --model QuartzNet15x5Base-En \ + --verbose + +To see possible normalization options for a text input without an audio file (could be used for debugging), run: + python python normalize_with_audio.py --text "RAW TEXT" +""" + + +class NormalizerWithAudio(Normalizer): + """ + Normalizer class that converts text from written to spoken form. + Useful for TTS preprocessing. + + Args: + input_case: expected input capitalization + """ + + def __init__(self, input_case: str): + super().__init__(input_case) + + self.tagger = ClassifyFst(input_case=input_case, deterministic=False) + self.verbalizer = VerbalizeFinalFst(deterministic=False) + + def normalize( + self, + text: str, + n_tagged: int, + punct_pre_process: bool = True, + punct_post_process: bool = True, + verbose: bool = False, + ) -> str: + """ + Main function. Normalizes tokens from written to spoken form + e.g. 12 kg -> twelve kilograms + + Args: + text: string that may include semiotic classes + n_tagged: number of tagged options to consider, -1 - to get all possible tagged options + punct_pre_process: whether to perform punctuation pre-processing, for example, [25] -> [ 25 ] + punct_post_process: whether to normalize punctuation + verbose: whether to print intermediate meta information + + Returns: + normalized text options (usually there are multiple ways of normalizing a given semiotic class) + """ + if punct_pre_process: + text = pre_process(text) + text = text.strip() + if not text: + if verbose: + print(text) + return text + + text = pynini.escape(text) + if n_tagged == -1: + tagged_texts = rewrite.rewrites(text, self.tagger.fst) + else: + tagged_texts = rewrite.top_rewrites(text, self.tagger.fst, nshortest=n_tagged) + + normalized_texts = [] + for tagged_text in tagged_texts: + self._verbalize(tagged_text, normalized_texts) + if len(normalized_texts) == 0: + raise ValueError() + if punct_post_process: + normalized_texts = [post_process_punctuation(t) for t in normalized_texts] + normalized_texts = set(normalized_texts) + return normalized_texts + + def _verbalize(self, tagged_text: str, normalized_texts: List[str]): + """ + Verbalizes tagged text + + Args: + tagged_text: text with tags + normalized_texts: list of possible normalization options + """ + + def get_verbalized_text(tagged_text): + tagged_text = pynini.escape(tagged_text) + return rewrite.rewrites(tagged_text, self.verbalizer.fst) + + try: + normalized_texts.extend(get_verbalized_text(tagged_text)) + except pynini.lib.rewrite.Error: + self.parser(tagged_text) + tokens = self.parser.parse() + tags_reordered = self.generate_permutations(tokens) + for tagged_text_reordered in tags_reordered: + try: + normalized_texts.extend(get_verbalized_text(tagged_text_reordered)) + except pynini.lib.rewrite.Error: + continue + + def select_best_match( + self, normalized_texts: List[str], transcript: str, verbose: bool = False, remove_punct: bool = False + ): + """ + Selects the best normalization option based on the lowest CER + + Args: + normalized_texts: normalized text options + transcript: ASR model transcript of the audio file corresponding to the normalized text + verbose: whether to print intermediate meta information + remove_punct: whether to remove punctuation before calculating CER + + Returns: + normalized text with the lowest CER and CER value + """ + normalized_texts = calculate_cer(normalized_texts, transcript, remove_punct) + normalized_texts = sorted(normalized_texts, key=lambda x: x[1]) + normalized_text, cer = normalized_texts[0] + + if verbose: + print('-' * 30) + for option in normalized_texts: + print(option) + print('-' * 30) + return normalized_text, cer + + +def calculate_cer(normalized_texts: List[str], transcript: str, remove_punct=False) -> List[Tuple[str, float]]: + """ + Calculates character error rate (CER) + + Args: + normalized_texts: normalized text options + transcript: ASR model output + + Returns: normalized options with corresponding CER + """ + normalized_options = [] + for text in normalized_texts: + text_clean = text.replace('-', ' ').lower() + if remove_punct: + for punct in "!?:;,.-()*+-/<=>@^_": + text_clean = text_clean.replace(punct, "") + cer = round(word_error_rate([transcript], [text_clean], use_cer=True) * 100, 2) + normalized_options.append((text, cer)) + return normalized_options + + +def get_asr_model(asr_model: ASRModel): + """ + Returns ASR Model + + Args: + asr_model: NeMo ASR model + """ + if os.path.exists(args.model): + asr_model = ASRModel.restore_from(asr_model) + elif args.model in ASRModel.get_available_model_names(): + asr_model = ASRModel.from_pretrained(asr_model) + else: + raise ValueError( + f'Provide path to the pretrained checkpoint or choose from {ASRModel.get_available_model_names()}' + ) + return asr_model + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--text", help="input string or path to a .txt file", default=None, type=str) + parser.add_argument( + "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + ) + parser.add_argument("--audio_data", help="path to an audio file or .json manifest") + parser.add_argument( + '--model', type=str, default='QuartzNet15x5Base-En', help='Pre-trained model name or path to model checkpoint' + ) + parser.add_argument( + "--n_tagged", + type=int, + default=1000, + help="number of tagged options to consider, -1 - return all possible tagged options", + ) + parser.add_argument("--verbose", help="print info for debugging", action="store_true") + parser.add_argument("--remove_punct", help="remove punctuation before calculating cer", action="store_true") + parser.add_argument( + "--no_punct_pre_process", help="set to True to disable punctuation pre processing", action="store_true" + ) + parser.add_argument( + "--no_punct_post_process", help="set to True to disable punctuation post processing", action="store_true" + ) + return parser.parse_args() + + +def _normalize_line(normalizer: NormalizerWithAudio, line: str, asr_model: ASRModel = None): + line = json.loads(line) + audio = line['audio_filepath'] + if 'transcript' in line: + transcript = line['transcript'] + else: + transcript = asr_model.transcribe([audio])[0] + + normalized_texts = normalizer.normalize( + text=line['text'], + verbose=args.verbose, + n_tagged=args.n_tagged, + punct_pre_process=not args.no_punct_pre_process, + punct_post_process=not args.no_punct_post_process, + ) + normalized_text, cer = normalizer.select_best_match(normalized_texts, transcript, args.verbose, args.remove_punct) + line['nemo_normalized'] = normalized_text + line['CER_nemo_normalized'] = cer + return line + + +def normalize_manifest(args): + """ + Args: + args.audio_data: path to .json manifest file. + """ + normalizer = NormalizerWithAudio(input_case=args.input_case) + manifest_out = args.audio_data.replace('.json', '_normalized.json') + asr_model = None + with open(args.audio_data, 'r') as f: + with open(manifest_out, 'w') as f_out: + lines = f.readlines() + first_line = json.loads(lines[0]) + if 'transcript' not in first_line: + asr_model = get_asr_model(args.model) + normalized_lines = Parallel(n_jobs=-1)( + delayed(_normalize_line)(normalizer, line, asr_model) for line in lines + ) + + for line in normalized_lines: + f_out.write(json.dumps(line, ensure_ascii=False) + '\n') + print(f'Normalized version saved at {manifest_out}') + + +if __name__ == "__main__": + args = parse_args() + + start = time.time() + if args.text: + normalizer = NormalizerWithAudio(input_case=args.input_case) + if os.path.exists(args.text): + with open(args.text, 'r') as f: + args.text = f.read().strip() + normalized_texts = normalizer.normalize( + text=args.text, + verbose=args.verbose, + n_tagged=args.n_tagged, + punct_pre_process=not args.no_punct_pre_process, + punct_post_process=not args.no_punct_post_process, + ) + if args.audio_data: + asr_model = get_asr_model(args.model) + transcript = asr_model.transcribe([args.audio_data])[0] + normalized_text, cer = normalizer.select_best_match( + normalized_texts, transcript, args.verbose, args.remove_punct + ) + print(f'Transcript: {transcript}') + print(f'Normalized: {normalized_text}') + else: + print('Normalization options:') + for norm_text in normalized_texts: + print(norm_text) + elif not os.path.exists(args.audio_data): + raise ValueError(f'{args.audio_data} not found.') + elif args.audio_data.endswith('.json'): + normalize_manifest(args) + else: + raise ValueError( + "Provide either path to .json manifest in '--audio_data' OR " + + "'--audio_data' path to audio file and '--text' path to a text file OR" + "'--text' string text (for debugging without audio)" + ) + print(f'Execution time: {round((time.time() - start)/60, 2)} min.') diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/run_evaluate.py b/models/tts/NeMo/nemo_text_processing/text_normalization/run_evaluate.py new file mode 100644 index 0000000000..3065a09005 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/run_evaluate.py @@ -0,0 +1,115 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser + +from nemo_text_processing.text_normalization.clean_eval_data import filter_loaded_data +from nemo_text_processing.text_normalization.data_loader_utils import ( + evaluate, + known_types, + load_files, + training_data_to_sentences, + training_data_to_tokens, +) +from nemo_text_processing.text_normalization.normalize import Normalizer + + +''' +Runs Evaluation on data in the format of : \t\t<`self` if trivial class or normalized text> +like the Google text normalization data https://www.kaggle.com/richardwilliamsproat/text-normalization-for-english-russian-and-polish +''' + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--input", help="input file path", type=str) + parser.add_argument( + "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + ) + parser.add_argument( + "--cat", + dest="category", + help="focus on class only (" + ", ".join(known_types) + ")", + type=str, + default=None, + choices=known_types, + ) + parser.add_argument("--filter", action='store_true', help="clean data for normalization purposes") + return parser.parse_args() + + +if __name__ == "__main__": + # Example usage: + # python run_evaluate.py --input= --cat= --filter + args = parse_args() + file_path = args.input + normalizer = Normalizer(input_case=args.input_case) + + print("Loading training data: " + file_path) + training_data = load_files([file_path]) + + if args.filter: + training_data = filter_loaded_data(training_data) + + if args.category is None: + print("Sentence level evaluation...") + sentences_un_normalized, sentences_normalized, _ = training_data_to_sentences(training_data) + print("- Data: " + str(len(sentences_normalized)) + " sentences") + sentences_prediction = normalizer.normalize_list(sentences_un_normalized) + print("- Normalized. Evaluating...") + sentences_accuracy = evaluate( + preds=sentences_prediction, labels=sentences_normalized, input=sentences_un_normalized + ) + print("- Accuracy: " + str(sentences_accuracy)) + + print("Token level evaluation...") + tokens_per_type = training_data_to_tokens(training_data, category=args.category) + token_accuracy = {} + for token_type in tokens_per_type: + print("- Token type: " + token_type) + tokens_un_normalized, tokens_normalized = tokens_per_type[token_type] + print(" - Data: " + str(len(tokens_normalized)) + " tokens") + tokens_prediction = normalizer.normalize_list(tokens_un_normalized) + print(" - Denormalized. Evaluating...") + token_accuracy[token_type] = evaluate( + preds=tokens_prediction, labels=tokens_normalized, input=tokens_un_normalized + ) + print(" - Accuracy: " + str(token_accuracy[token_type])) + token_count_per_type = {token_type: len(tokens_per_type[token_type][0]) for token_type in tokens_per_type} + token_weighted_accuracy = [ + token_count_per_type[token_type] * accuracy for token_type, accuracy in token_accuracy.items() + ] + print("- Accuracy: " + str(sum(token_weighted_accuracy) / sum(token_count_per_type.values()))) + print(" - Total: " + str(sum(token_count_per_type.values())), '\n') + + print(" - Total: " + str(sum(token_count_per_type.values())), '\n') + + for token_type in token_accuracy: + if token_type not in known_types: + raise ValueError("Unexpected token type: " + token_type) + + if args.category is None: + c1 = ['Class', 'sent level'] + known_types + c2 = ['Num Tokens', len(sentences_normalized)] + [ + token_count_per_type[known_type] if known_type in tokens_per_type else '0' for known_type in known_types + ] + c3 = ['Normalization', sentences_accuracy] + [ + token_accuracy[known_type] if known_type in token_accuracy else '0' for known_type in known_types + ] + + for i in range(len(c1)): + print(f'{str(c1[i]):10s} | {str(c2[i]):10s} | {str(c3[i]):5s}') + else: + print(f'numbers\t{token_count_per_type[args.category]}') + print(f'Normalization\t{token_accuracy[args.category]}') diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/run_predict.py b/models/tts/NeMo/nemo_text_processing/text_normalization/run_predict.py new file mode 100644 index 0000000000..3059dfe07f --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/run_predict.py @@ -0,0 +1,78 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from argparse import ArgumentParser +from typing import List + +from nemo_text_processing.text_normalization.normalize import Normalizer + + +''' +Runs normalization prediction on text data +''' + + +def load_file(file_path: str) -> List[str]: + """ + Load given text file into list of string. + + Args: + file_path: file path + + Returns: flat list of string + """ + res = [] + with open(file_path, 'r') as fp: + for line in fp: + res.append(line) + return res + + +def write_file(file_path: str, data: List[str]): + """ + Writes out list of string to file. + + Args: + file_path: file path + data: list of string + + """ + with open(file_path, 'w') as fp: + for line in data: + fp.write(line + '\n') + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--input", help="input file path", required=True, type=str) + parser.add_argument("--output", help="output file path", required=True, type=str) + parser.add_argument( + "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + ) + parser.add_argument("--verbose", help="print meta info for debugging", action='store_true') + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + file_path = args.input + normalizer = Normalizer(input_case=args.input_case) + + print("Loading data: " + file_path) + data = load_file(file_path) + + print("- Data: " + str(len(data)) + " sentences") + normalizer_prediction = normalizer.normalize_list(data, verbose=args.verbose) + write_file(args.output, normalizer_prediction) + print(f"- Normalized. Writing out to {args.output}") diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/cardinal.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/cardinal.py new file mode 100644 index 0000000000..2d01a49357 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/cardinal.py @@ -0,0 +1,115 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path, load_labels +from nemo_text_processing.text_normalization.graph_utils import NEMO_ALPHA, NEMO_DIGIT, GraphFst, insert_space +from nemo_text_processing.text_normalization.taggers.date import get_hundreds_graph + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class CardinalFst(GraphFst): + """ + Finite state transducer for classifying cardinals, e.g. + -23 -> cardinal { negative: "true" integer: "twenty three" } } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="cardinal", kind="classify", deterministic=deterministic) + + graph = pynini.Far(get_abs_path("data/numbers/cardinal_number_name.far")).get_fst() + self.graph_hundred_component_at_least_one_none_zero_digit = ( + pynini.closure(NEMO_DIGIT, 2, 3) | pynini.difference(NEMO_DIGIT, pynini.accep("0")) + ) @ graph + self.graph = ( + pynini.closure(NEMO_DIGIT, 1, 3) + + pynini.closure(pynini.closure(pynutil.delete(","), 0, 1) + NEMO_DIGIT + NEMO_DIGIT + NEMO_DIGIT) + ) @ graph + + graph_digit = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) + graph_zero = pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + single_digits_graph = pynutil.add_weight(pynini.invert(graph_digit | graph_zero), 1.2) | pynutil.add_weight( + pynini.cross("0", "oh"), 1.1 + ) + self.single_digits_graph = single_digits_graph + pynini.closure(pynutil.insert(" ") + single_digits_graph) + + if not deterministic: + single_digits_graph_with_commas = pynini.closure( + self.single_digits_graph + pynutil.insert(" "), 1, 3 + ) + pynini.closure( + pynutil.delete(",") + + single_digits_graph + + pynutil.insert(" ") + + single_digits_graph + + pynutil.insert(" ") + + single_digits_graph, + 1, + ) + self.graph |= self.single_digits_graph | get_hundreds_graph() | single_digits_graph_with_commas + self.range_graph = ( + pynini.closure(pynutil.insert("from "), 0, 1) + + self.graph + + (pynini.cross("-", " to ") | pynini.cross("-", " ")) + + self.graph + ) + + self.range_graph |= self.graph + (pynini.cross("x", " by ") | pynini.cross(" x ", " by ")) + self.graph + self.range_graph = self.range_graph.optimize() + + optional_minus_graph = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + final_graph = self.graph | pynutil.add_weight(self.get_serial_graph(), 1.2) + + if not deterministic: + final_graph |= self.range_graph + + final_graph = optional_minus_graph + pynutil.insert("integer: \"") + final_graph + pynutil.insert("\"") + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() + + def get_serial_graph(self): + """ + Finite state transducer for classifying serial. + The serial is a combination of digits, letters and dashes, e.g.: + c325-b -> tokens { serial { value: "c three two five b" } } + """ + alpha = NEMO_ALPHA + + if self.deterministic: + num_graph = self.single_digits_graph + else: + num_graph = self.graph + letter_pronunciation = pynini.string_map(load_labels(get_abs_path("data/letter_pronunciation.tsv"))) + alpha |= letter_pronunciation + + delimiter = insert_space | pynini.cross("-", " ") | pynini.cross("/", " ") + letter_num = pynini.closure(alpha + delimiter, 1) + num_graph + num_letter = pynini.closure(num_graph + delimiter, 1) + alpha + next_alpha_or_num = pynini.closure(delimiter + (alpha | num_graph)) + serial_graph = (letter_num | num_letter) + next_alpha_or_num + + if not self.deterministic: + serial_graph += pynini.closure(pynini.accep("s") | pynini.cross("s", "es"), 0, 1) + return serial_graph diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/date.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/date.py new file mode 100644 index 0000000000..8c7c6a713a --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/date.py @@ -0,0 +1,201 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path, load_labels +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_CHAR, + NEMO_DIGIT, + NEMO_SIGMA, + TO_LOWER, + GraphFst, + delete_extra_space, + delete_space, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + graph_teen = pynini.invert(pynini.string_file(get_abs_path("data/numbers/teen.tsv"))).optimize() + graph_digit = pynini.invert(pynini.string_file(get_abs_path("data/numbers/digit.tsv"))).optimize() + ties_graph = pynini.invert(pynini.string_file(get_abs_path("data/numbers/ties.tsv"))).optimize() + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + # Add placeholders for global variables + graph_teen = None + graph_digit = None + ties_graph = None + + PYNINI_AVAILABLE = True + + +def get_ties_graph(deterministic: bool = True): + """ + Returns two digit transducer, e.g. + 03 -> o three + 12 -> thirteen + 20 -> twenty + """ + graph = graph_teen | ties_graph + pynutil.delete("0") | ties_graph + insert_space + graph_digit + + if deterministic: + graph = graph | pynini.cross("0", "o") + insert_space + graph_digit + else: + graph = ( + graph + | (pynini.cross("0", "oh") | pynini.cross("0", "o") | pynini.cross("0", "zero")) + + insert_space + + graph_digit + ) + + return graph.optimize() + + +def get_hundreds_graph(deterministic: bool = True): + """ + Returns a four digit transducer which is combination of ties/teen or digits + (using hundred instead of thousand format), e.g. + 1219 -> twelve nineteen + 3900 -> thirty nine hundred + """ + graph_ties = get_ties_graph(deterministic) + graph = ( + graph_ties + insert_space + graph_ties + | graph_teen + insert_space + pynini.cross("00", "hundred") + | (graph_teen + insert_space + (ties_graph | pynini.cross("1", "ten")) + pynutil.delete("0s")) + @ pynini.cdrewrite(pynini.cross("y", "ies") | pynutil.insert("s"), "", "[EOS]", NEMO_SIGMA) + | pynutil.add_weight( + graph_digit + + insert_space + + pynini.cross("00", "thousand") + + (pynutil.delete("0") | insert_space + graph_digit), + weight=-0.001, + ) + ) + return graph + + +def _get_year_graph(deterministic: bool = True): + """ + Transducer for year, only from 1000 - 2999 e.g. + 1290 -> twelve nineteen + 2000 - 2009 will be verbalized as two thousand. + """ + graph = get_hundreds_graph(deterministic) + graph = ( + pynini.union("1", "2") + + NEMO_DIGIT + + NEMO_DIGIT + + NEMO_DIGIT + + pynini.closure(pynini.cross(" s", "s") | "s", 0, 1) + ) @ graph + return graph + + +class DateFst(GraphFst): + """ + Finite state transducer for classifying date, e.g. + jan. 5, 2012 -> date { month: "january" day: "five" year: "twenty twelve" preserve_order: true } + jan. 5 -> date { month: "january" day: "five" preserve_order: true } + 5 january 2012 -> date { day: "five" month: "january" year: "twenty twelve" preserve_order: true } + 2012-01-05 -> date { year: "twenty twelve" month: "january" day: "five" } + 2012.01.05 -> date { year: "twenty twelve" month: "january" day: "five" } + 2012/01/05 -> date { year: "twenty twelve" month: "january" day: "five" } + 2012 -> date { year: "twenty twelve" } + + Args: + ordinal: OrdinalFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal: GraphFst, deterministic: bool): + super().__init__(name="date", kind="classify", deterministic=deterministic) + + month_graph = pynini.string_file(get_abs_path("data/months/names.tsv")).optimize() + month_graph |= (TO_LOWER + pynini.closure(NEMO_CHAR)) @ month_graph + month_abbr_graph = pynini.string_file(get_abs_path("data/months/abbr.tsv")).optimize() + month_abbr_graph = ( + month_abbr_graph | (TO_LOWER + pynini.closure(NEMO_CHAR)) @ month_abbr_graph + ) + pynini.closure(pynutil.delete("."), 0, 1) + month_graph |= month_abbr_graph + + # to support all caps names + names_all_caps = [[x[0].upper()] for x in load_labels(get_abs_path("data/months/names.tsv"))] + abbr_all_caps = [(x.upper(), y) for x, y in load_labels(get_abs_path("data/months/abbr.tsv"))] + month_graph |= pynini.string_map(names_all_caps) | ( + pynini.string_map(abbr_all_caps) + pynini.closure(pynutil.delete("."), 0, 1) + ) + + month_numbers_graph = pynini.string_file(get_abs_path("data/months/numbers.tsv")).optimize() + cardinal_graph = cardinal.graph_hundred_component_at_least_one_none_zero_digit + + year_graph = _get_year_graph(deterministic) + + YEAR_WEIGHT = 0.001 + year_graph_standalone = ( + pynutil.insert("year: \"") + pynutil.add_weight(year_graph, YEAR_WEIGHT) + pynutil.insert("\"") + ) + + month_graph = pynutil.insert("month: \"") + month_graph + pynutil.insert("\"") + month_numbers_graph = pynutil.insert("month: \"") + month_numbers_graph + pynutil.insert("\"") + + day_graph = ( + pynutil.insert("day: \"") + + ((pynini.union("1", "2", "3") + NEMO_DIGIT) | NEMO_DIGIT) @ cardinal_graph + + pynutil.insert("\"") + ) + optional_day_graph = pynini.closure(delete_extra_space + day_graph, 0, 1) + + year_graph = pynutil.insert("year: \"") + year_graph + pynutil.insert("\"") + optional_graph_year = pynini.closure(delete_extra_space + year_graph, 0, 1,) + graph_mdy = ( + month_graph + + optional_day_graph + + delete_space + + pynini.closure(pynutil.delete(","), 0, 1) + + optional_graph_year + ) + + delete_sep = pynutil.delete(pynini.union("-", "/", ".")) + graph_mdy |= ( + month_numbers_graph + + delete_sep + + insert_space + + pynini.closure(pynutil.delete("0"), 0, 1) + + day_graph + + delete_sep + + insert_space + + year_graph + ) + + graph_dmy = day_graph + delete_extra_space + month_graph + optional_graph_year + graph_ymd = ( + year_graph + + delete_sep + + insert_space + + month_numbers_graph + + delete_sep + + insert_space + + pynini.closure(pynutil.delete("0"), 0, 1) + + day_graph + ) + + final_graph = (graph_mdy | graph_dmy) + pynutil.insert(" preserve_order: true") + final_graph |= graph_ymd | year_graph_standalone + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/decimal.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/decimal.py new file mode 100644 index 0000000000..c4eba3fe65 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/decimal.py @@ -0,0 +1,101 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_extra_space, delete_space + +try: + import pynini + from pynini.lib import pynutil + + delete_space = pynutil.delete(" ") + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +def get_quantity(decimal: 'pynini.FstLike', cardinal_up_to_hundred: 'pynini.FstLike') -> 'pynini.FstLike': + """ + Returns FST that transforms either a cardinal or decimal followed by a quantity into a numeral, + e.g. 1 million -> integer_part: "one" quantity: "million" + e.g. 1.5 million -> integer_part: "one" fractional_part: "five" quantity: "million" + + Args: + decimal: decimal FST + cardinal_up_to_hundred: cardinal FST + """ + numbers = cardinal_up_to_hundred + suffix = pynini.union("million", "billion", "trillion", "quadrillion", "quintillion", "sextillion") + res = ( + pynutil.insert("integer_part: \"") + + numbers + + pynutil.insert("\"") + + delete_extra_space + + pynutil.insert("quantity: \"") + + suffix + + pynutil.insert("\"") + ) + res |= decimal + delete_extra_space + pynutil.insert("quantity: \"") + (suffix | "thousand") + pynutil.insert("\"") + return res + + +class DecimalFst(GraphFst): + """ + Finite state transducer for classifying decimal, e.g. + -12.5006 billion -> decimal { negative: "true" integer_part: "12" fractional_part: "five o o six" quantity: "billion" } + 1 billion -> decimal { integer_part: "one" quantity: "billion" } + + cardinal: CardinalFst + """ + + def __init__(self, cardinal: GraphFst, deterministic: bool): + super().__init__(name="decimal", kind="classify", deterministic=deterministic) + + cardinal_graph = cardinal.graph + cardinal_graph_hundred_component_at_least_one_none_zero_digit = ( + cardinal.graph_hundred_component_at_least_one_none_zero_digit + ) + + graph_decimal = pynini.string_file(get_abs_path("data/numbers/digit.tsv")) + graph_decimal |= pynini.string_file(get_abs_path("data/numbers/zero.tsv")) + + graph_decimal = ( + pynini.cross("zero", "0") + | graph_decimal + | (graph_decimal | pynini.cross("o", "0")) + + pynini.closure(delete_space + (graph_decimal | pynini.cross("o", "0")), 1) + ) + self.graph = pynini.invert(graph_decimal).optimize() + if not deterministic: + self.graph = self.graph | cardinal_graph + + point = pynutil.delete(".") + optional_graph_negative = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + + graph_fractional = pynutil.insert("fractional_part: \"") + self.graph + pynutil.insert("\"") + graph_integer = pynutil.insert("integer_part: \"") + cardinal_graph + pynutil.insert("\"") + final_graph_wo_sign = ( + pynini.closure(graph_integer + pynutil.insert(" "), 0, 1) + point + pynutil.insert(" ") + graph_fractional + ) + + self.final_graph_wo_negative = final_graph_wo_sign | get_quantity( + final_graph_wo_sign, cardinal_graph_hundred_component_at_least_one_none_zero_digit + ) + + final_graph = optional_graph_negative + self.final_graph_wo_negative + + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/electronic.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/electronic.py new file mode 100644 index 0000000000..097a8cd918 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/electronic.py @@ -0,0 +1,63 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.graph_utils import NEMO_ALPHA, NEMO_DIGIT, GraphFst, get_abs_path + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class ElectronicFst(GraphFst): + """ + Finite state transducer for classifying electronic: as URLs, email addresses, etc. + e.g. cdf1@abc.edu -> tokens { electronic { username: "cdf1" domain: "abc.edu" } } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="electronic", kind="classify", deterministic=deterministic) + + accepted_symbols = [] + with open(get_abs_path("data/electronic/symbols.tsv"), 'r') as f: + for line in f: + symbol, _ = line.split('\t') + accepted_symbols.append(pynini.accep(symbol)) + + username = ( + pynutil.insert("username: \"") + + NEMO_ALPHA + + pynini.closure(NEMO_ALPHA | NEMO_DIGIT | pynini.union(*accepted_symbols)) + + pynutil.insert("\"") + + pynini.cross('@', ' ') + ) + domain_graph = ( + NEMO_ALPHA + + (pynini.closure(NEMO_ALPHA | NEMO_DIGIT | pynini.accep('-') | pynini.accep('.'))) + + (NEMO_ALPHA | NEMO_DIGIT) + ) + domain_graph = pynutil.insert("domain: \"") + domain_graph + pynutil.insert("\"") + graph = username + domain_graph + + final_graph = self.add_tokens(graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/fraction.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/fraction.py new file mode 100644 index 0000000000..0feb5ce1dc --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/fraction.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class FractionFst(GraphFst): + """ + Finite state transducer for classifying fraction + "23 4/5" -> + tokens { fraction { integer: "twenty three" numerator: "four" denominator: "five" } } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal, deterministic: bool = True): + super().__init__(name="fraction", kind="classify", deterministic=deterministic) + cardinal_graph = cardinal.graph + + integer = pynutil.insert("integer: \"") + cardinal_graph + pynutil.insert("\"") + pynini.accep(" ") + numerator = ( + pynutil.insert("numerator: \"") + cardinal_graph + (pynini.cross("/", "\" ") | pynini.cross(" / ", "\" ")) + ) + denominator = pynutil.insert("denominator: \"") + cardinal_graph + pynutil.insert("\"") + + self.graph = pynini.closure(integer, 0, 1) + numerator + denominator + final_graph = self.add_tokens(self.graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/measure.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/measure.py new file mode 100644 index 0000000000..a0475c5b10 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/measure.py @@ -0,0 +1,162 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_ALPHA, + NEMO_NON_BREAKING_SPACE, + NEMO_SIGMA, + SINGULAR_TO_PLURAL, + GraphFst, + convert_space, + delete_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class MeasureFst(GraphFst): + """ + Finite state transducer for classifying measure, suppletive aware, e.g. + -12kg -> measure { negative: "true" cardinal { integer: "twelve" } units: "kilograms" } + 1kg -> measure { cardinal { integer: "one" } units: "kilogram" } + .5kg -> measure { decimal { fractional_part: "five" } units: "kilograms" } + + Args: + cardinal: CardinalFst + decimal: DecimalFst + fraction: FractionFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal: GraphFst, decimal: GraphFst, fraction: GraphFst, deterministic: bool = True): + super().__init__(name="measure", kind="classify", deterministic=deterministic) + cardinal_graph = cardinal.graph + + if not deterministic: + cardinal_graph |= cardinal.range_graph + + graph_unit = pynini.string_file(get_abs_path("data/measurements.tsv")) + graph_unit_plural = convert_space(graph_unit @ SINGULAR_TO_PLURAL) + graph_unit = convert_space(graph_unit) + optional_graph_negative = pynini.closure(pynutil.insert("negative: ") + pynini.cross("-", "\"true\" "), 0, 1) + + graph_unit2 = pynini.cross("/", "per") + delete_space + pynutil.insert(NEMO_NON_BREAKING_SPACE) + graph_unit + + optional_graph_unit2 = pynini.closure( + delete_space + pynutil.insert(NEMO_NON_BREAKING_SPACE) + graph_unit2, 0, 1, + ) + + unit_plural = ( + pynutil.insert("units: \"") + + (graph_unit_plural + optional_graph_unit2 | graph_unit2) + + pynutil.insert("\"") + ) + + unit_singular = ( + pynutil.insert("units: \"") + (graph_unit + optional_graph_unit2 | graph_unit2) + pynutil.insert("\"") + ) + + subgraph_decimal = ( + pynutil.insert("decimal { ") + + optional_graph_negative + + decimal.final_graph_wo_negative + + delete_space + + pynutil.insert(" } ") + + unit_plural + ) + + subgraph_cardinal = ( + pynutil.insert("cardinal { ") + + optional_graph_negative + + pynutil.insert("integer: \"") + + ((NEMO_SIGMA - "1") @ cardinal_graph) + + delete_space + + pynutil.insert("\"") + + pynutil.insert(" } ") + + unit_plural + ) + + subgraph_cardinal |= ( + pynutil.insert("cardinal { ") + + optional_graph_negative + + pynutil.insert("integer: \"") + + pynini.cross("1", "one") + + delete_space + + pynutil.insert("\"") + + pynutil.insert(" } ") + + unit_singular + ) + + cardinal_dash_alpha = ( + pynutil.insert("cardinal { integer: \"") + + cardinal_graph + + pynini.cross('-', '') + + pynutil.insert("\" } units: \"") + + pynini.closure(NEMO_ALPHA, 1) + + pynutil.insert("\"") + ) + + alpha_dash_cardinal = ( + pynutil.insert("units: \"") + + pynini.closure(NEMO_ALPHA, 1) + + pynini.cross('-', '') + + pynutil.insert("\"") + + pynutil.insert(" cardinal { integer: \"") + + cardinal_graph + + pynutil.insert("\" } preserve_order: true") + ) + + decimal_dash_alpha = ( + pynutil.insert("decimal { ") + + decimal.final_graph_wo_negative + + pynini.cross('-', '') + + pynutil.insert(" } units: \"") + + pynini.closure(NEMO_ALPHA, 1) + + pynutil.insert("\"") + ) + + alpha_dash_decimal = ( + pynutil.insert("units: \"") + + pynini.closure(NEMO_ALPHA, 1) + + pynini.cross('-', '') + + pynutil.insert("\"") + + pynutil.insert(" decimal { ") + + decimal.final_graph_wo_negative + + pynutil.insert(" } preserve_order: true") + ) + + subgraph_fraction = ( + pynutil.insert("fraction { ") + fraction.graph + delete_space + pynutil.insert(" } ") + unit_plural + ) + + final_graph = ( + subgraph_decimal + | subgraph_cardinal + | cardinal_dash_alpha + | alpha_dash_cardinal + | decimal_dash_alpha + | alpha_dash_decimal + | subgraph_fraction + ) + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/money.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/money.py new file mode 100644 index 0000000000..5965934e3b --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/money.py @@ -0,0 +1,86 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_SIGMA, + SINGULAR_TO_PLURAL, + GraphFst, + convert_space, + insert_space, +) +from nemo_text_processing.text_normalization.taggers.date import get_hundreds_graph + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class MoneyFst(GraphFst): + """ + Finite state transducer for classifying money, suppletive aware, e.g. + $12.05 -> money { currency: "dollars" integer_part: "twelve" fractional_part: "o five" } + $1 -> money { currency: "dollar" integer_part: "one" } + + Args: + cardinal: CardinalFst + decimal: DecimalFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal: GraphFst, decimal: GraphFst, deterministic: bool = True): + super().__init__(name="money", kind="classify", deterministic=deterministic) + cardinal_graph = cardinal.graph + graph_decimal_final = decimal.final_graph_wo_negative + + unit_singular = pynini.string_file(get_abs_path("data/currency/currency.tsv")) + unit_plural = convert_space(unit_singular @ SINGULAR_TO_PLURAL) + unit_singular = convert_space(unit_singular) + + graph_unit_singular = pynutil.insert("currency: \"") + unit_singular + pynutil.insert("\"") + graph_unit_plural = pynutil.insert("currency: \"") + unit_plural + pynutil.insert("\"") + + singular_graph = ( + graph_unit_singular + pynutil.insert(" integer_part: \"") + pynini.cross("1", "one") + pynutil.insert("\"") + ) + + graph_decimal = graph_unit_plural + insert_space + graph_decimal_final + + if deterministic: + graph_integer = ( + graph_unit_plural + + pynutil.insert(" integer_part: \"") + + ((NEMO_SIGMA - "1") @ cardinal_graph) + + pynutil.insert("\"") + ) + else: + graph_integer = ( + graph_unit_plural + + pynutil.insert(" integer_part: \"") + + ((NEMO_SIGMA - "1") @ (get_hundreds_graph(deterministic) | cardinal_graph)) + + pynutil.insert("\"") + ) + graph_decimal |= singular_graph + insert_space + graph_decimal_final + + graph_integer |= singular_graph + + final_graph = graph_integer | graph_decimal + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/ordinal.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/ordinal.py new file mode 100644 index 0000000000..49e452729d --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/ordinal.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_DIGIT, GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ImportError, ModuleNotFoundError): + PYNINI_AVAILABLE = False + + +class OrdinalFst(GraphFst): + """ + Finite state transducer for classifying ordinal, e.g. + 13th -> ordinal { integer: "thirteen" } + + Args: + cardinal: CardinalFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal: GraphFst, deterministic: bool = True): + super().__init__(name="ordinal", kind="classify", deterministic=deterministic) + + cardinal_graph = cardinal.graph + endings = ["rd", "th", "st", "nd"] + endings += [x.upper() for x in endings] + self.graph = ( + (pynini.closure(NEMO_DIGIT | pynini.accep(",")) + pynutil.delete(pynini.union(*endings))) @ cardinal_graph + ).optimize() + final_graph = pynutil.insert("integer: \"") + self.graph + pynutil.insert("\"") + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/punctuation.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/punctuation.py new file mode 100644 index 0000000000..6a42e2d7f4 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/punctuation.py @@ -0,0 +1,46 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = False +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class PunctuationFst(GraphFst): + """ + Finite state transducer for classifying punctuation + e.g. a, -> tokens { name: "a" } tokens { name: "," } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="punctuation", kind="classify", deterministic=deterministic) + + s = "!#$%&\'()*+,-./:;<=>?@^_`{|}~\"" + punct = pynini.union(*s) + + graph = pynutil.insert("name: \"") + punct + pynutil.insert("\"") + + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/roman.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/roman.py new file mode 100644 index 0000000000..136d0fb520 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/roman.py @@ -0,0 +1,61 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.data_loader_utils import load_labels +from nemo_text_processing.text_normalization.graph_utils import GraphFst, get_abs_path, insert_space +from nemo_text_processing.text_normalization.taggers.cardinal import CardinalFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class RomanFst(GraphFst): + """ + Finite state transducer for classifying electronic: as URLs, email addresses, etc. + e.g. cdf1@abc.edu -> tokens { electronic { username: "cdf1" domain: "abc.edu" } } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="roman", kind="classify", deterministic=deterministic) + + def _load_roman(file: str): + roman = load_labels(get_abs_path(file)) + roman_numerals = [(x, y) for x, y in roman] + [(x.upper(), y) for x, y in roman] + return pynini.string_map(roman_numerals) + + cardinal_graph = CardinalFst(deterministic=True).graph + digit_teen = _load_roman("data/roman/digit_teen.tsv") @ cardinal_graph + ties = _load_roman("data/roman/ties.tsv") @ cardinal_graph + hundreds = _load_roman("data/roman/hundreds.tsv") @ cardinal_graph + + graph = ( + (ties | digit_teen | hundreds) + | (ties + insert_space + digit_teen) + | (hundreds + pynini.closure(insert_space + ties, 0, 1) + pynini.closure(insert_space + digit_teen, 0, 1)) + ).optimize() + + graph = pynutil.insert("integer: \"") + graph + pynutil.insert("\"") + graph = self.add_tokens(graph) + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/telephone.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/telephone.py new file mode 100644 index 0000000000..fa2ec7848b --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/telephone.py @@ -0,0 +1,91 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_space, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TelephoneFst(GraphFst): + """ + Finite state transducer for classifying telephone, which includes country code, number part and extension + country code optional: +*** + number part: ***-***-****, or (***) ***-**** + extension optional: 1-9999 + E.g + +1 123-123-5678-1 -> telephone { country_code: "one" number_part: "one two three, one two three, five six seven eight" extension: "one" } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="telephone", kind="classify", deterministic=deterministic) + + add_separator = pynutil.insert(", ") # between components + digit = pynini.invert(pynini.string_file(get_abs_path("data/numbers/digit.tsv"))).optimize() | pynini.cross( + "0", "o" + ) + + country_code = ( + pynutil.insert("country_code: \"") + + pynutil.delete("+") + + pynini.closure(digit + insert_space, 0, 2) + + digit + + pynutil.insert("\"") + ) + optional_country_code = pynini.closure( + country_code + pynini.closure(pynutil.delete("-"), 0, 1) + delete_space + insert_space, 0, 1 + ) + number_part = ( + ( + (pynini.closure(digit + insert_space, 2, 2) + digit + pynutil.delete("-")) + | ( + pynutil.delete("(") + + pynini.closure(digit + insert_space, 2, 2) + + digit + + pynutil.delete(")") + + pynini.closure(pynutil.delete("-"), 0, 1) + + delete_space + ) + ) + + add_separator + + pynini.closure(digit + insert_space, 2, 2) + + digit + + pynutil.delete("-") + + add_separator + + pynini.closure(digit + insert_space, 3, 3) + + digit + ) + number_part = pynutil.insert("number_part: \"") + number_part + pynutil.insert("\"") + extension = ( + pynutil.insert("extension : \"") + + pynini.closure(digit + insert_space, 0, 3) + + digit + + pynutil.insert("\"") + ) + optional_extension = pynini.closure(insert_space + pynutil.delete("-") + extension, 0, 1) + + graph = optional_country_code + number_part + optional_extension + final_graph = self.add_tokens(graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/time.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/time.py new file mode 100644 index 0000000000..7da3f56763 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/time.py @@ -0,0 +1,115 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_DIGIT, + GraphFst, + convert_space, + delete_space, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TimeFst(GraphFst): + """ + Finite state transducer for classifying time, e.g. + 12:30 a.m. est -> time { hours: "twelve" minutes: "thirty" suffix: "a m" zone: "e s t" } + 2.30 a.m. -> time { hours: "two" minutes: "thirty" suffix: "a m" } + 02.30 a.m. -> time { hours: "two" minutes: "thirty" suffix: "a m" } + 2.00 a.m. -> time { hours: "two" suffix: "a m" } + 2 a.m. -> time { hours: "two" suffix: "a m" } + 02:00 -> time { hours: "two" } + 2:00 -> time { hours: "two" } + + Args: + cardinal: CardinalFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal: GraphFst, deterministic: bool = True): + super().__init__(name="time", kind="classify", deterministic=deterministic) + suffix_graph = pynini.string_file(get_abs_path("data/time_suffix.tsv")) + time_zone_graph = pynini.string_file(get_abs_path("data/time_zone.tsv")) + + # only used for < 1000 thousand -> 0 weight + cardinal = cardinal.graph + + labels_hour = [str(x) for x in range(0, 24)] + labels_minute_single = [str(x) for x in range(1, 10)] + labels_minute_double = [str(x) for x in range(10, 60)] + + delete_leading_zero_to_double_digit = (NEMO_DIGIT + NEMO_DIGIT) | ( + pynini.closure(pynutil.delete("0"), 0, 1) + NEMO_DIGIT + ) + + graph_hour = delete_leading_zero_to_double_digit @ pynini.union(*labels_hour) @ cardinal + + graph_minute_single = pynini.union(*labels_minute_single) @ cardinal + graph_minute_double = pynini.union(*labels_minute_double) @ cardinal + + final_graph_hour = pynutil.insert("hours: \"") + graph_hour + pynutil.insert("\"") + final_graph_minute = ( + pynutil.insert("minutes: \"") + + (pynini.cross("0", "o") + insert_space + graph_minute_single | graph_minute_double) + + pynutil.insert("\"") + ) + final_suffix = pynutil.insert("suffix: \"") + convert_space(suffix_graph) + pynutil.insert("\"") + final_suffix_optional = pynini.closure(delete_space + insert_space + final_suffix, 0, 1) + final_time_zone_optional = pynini.closure( + delete_space + + insert_space + + pynutil.insert("zone: \"") + + convert_space(time_zone_graph) + + pynutil.insert("\""), + 0, + 1, + ) + + # 2:30 pm, 02:30, 2:00 + graph_hm = ( + final_graph_hour + + pynutil.delete(":") + + (pynutil.delete("00") | insert_space + final_graph_minute) + + final_suffix_optional + + final_time_zone_optional + ) + + # 2.xx pm/am + graph_hm2 = ( + final_graph_hour + + pynutil.delete(".") + + (pynutil.delete("00") | insert_space + final_graph_minute) + + delete_space + + insert_space + + final_suffix + + final_time_zone_optional + ) + # 2 pm est + graph_h = final_graph_hour + delete_space + insert_space + final_suffix + final_time_zone_optional + final_graph = (graph_hm | graph_h | graph_hm2).optimize() + + final_graph = self.add_tokens(final_graph) + self.fst = final_graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/tokenize_and_classify.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/tokenize_and_classify.py new file mode 100644 index 0000000000..1bef666a11 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/tokenize_and_classify.py @@ -0,0 +1,107 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_extra_space, delete_space +from nemo_text_processing.text_normalization.taggers.cardinal import CardinalFst +from nemo_text_processing.text_normalization.taggers.date import DateFst +from nemo_text_processing.text_normalization.taggers.decimal import DecimalFst +from nemo_text_processing.text_normalization.taggers.electronic import ElectronicFst +from nemo_text_processing.text_normalization.taggers.fraction import FractionFst +from nemo_text_processing.text_normalization.taggers.measure import MeasureFst +from nemo_text_processing.text_normalization.taggers.money import MoneyFst +from nemo_text_processing.text_normalization.taggers.ordinal import OrdinalFst +from nemo_text_processing.text_normalization.taggers.punctuation import PunctuationFst +from nemo_text_processing.text_normalization.taggers.roman import RomanFst +from nemo_text_processing.text_normalization.taggers.telephone import TelephoneFst +from nemo_text_processing.text_normalization.taggers.time import TimeFst +from nemo_text_processing.text_normalization.taggers.whitelist import WhiteListFst +from nemo_text_processing.text_normalization.taggers.word import WordFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class ClassifyFst(GraphFst): + """ + Final class that composes all other classification grammars. This class can process an entire sentence including punctuation. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + More details to deployment at NeMo/tools/text_processing_deployment. + + Args: + input_case: accepting either "lower_cased" or "cased" input. + deterministic: if True will provide a single transduction option, + for False multiple options (used for audio-based normalization) + """ + + def __init__(self, input_case: str, deterministic: bool = True): + super().__init__(name="tokenize_and_classify", kind="classify", deterministic=deterministic) + + cardinal = CardinalFst(deterministic=deterministic) + cardinal_graph = cardinal.fst + + ordinal = OrdinalFst(cardinal=cardinal, deterministic=deterministic) + ordinal_graph = ordinal.fst + + decimal = DecimalFst(cardinal=cardinal, deterministic=deterministic) + decimal_graph = decimal.fst + fraction = FractionFst(deterministic=deterministic, cardinal=cardinal) + fraction_graph = fraction.fst + + measure = MeasureFst(cardinal=cardinal, decimal=decimal, fraction=fraction, deterministic=deterministic) + measure_graph = measure.fst + date_graph = DateFst(cardinal=cardinal, deterministic=deterministic).fst + word_graph = WordFst(deterministic=deterministic).fst + time_graph = TimeFst(cardinal=cardinal, deterministic=deterministic).fst + telephone_graph = TelephoneFst(deterministic=deterministic).fst + electonic_graph = ElectronicFst(deterministic=deterministic).fst + money_graph = MoneyFst(cardinal=cardinal, decimal=decimal, deterministic=deterministic).fst + whitelist_graph = WhiteListFst(input_case=input_case, deterministic=deterministic).fst + punct_graph = PunctuationFst(deterministic=deterministic).fst + + classify = ( + pynutil.add_weight(whitelist_graph, 1.01) + | pynutil.add_weight(time_graph, 1.1) + | pynutil.add_weight(date_graph, 1.09) + | pynutil.add_weight(decimal_graph, 1.1) + | pynutil.add_weight(measure_graph, 1.1) + | pynutil.add_weight(cardinal_graph, 1.1) + | pynutil.add_weight(ordinal_graph, 1.1) + | pynutil.add_weight(money_graph, 1.1) + | pynutil.add_weight(telephone_graph, 1.1) + | pynutil.add_weight(electonic_graph, 1.1) + | pynutil.add_weight(fraction_graph, 1.1) + | pynutil.add_weight(word_graph, 100) + ) + + if not deterministic: + roman_graph = RomanFst(deterministic=deterministic).fst + # the weight matches the word_graph weight for "I" cases in long sentences with multiple semiotic tokens + classify |= pynutil.add_weight(roman_graph, 100) + + punct = pynutil.insert("tokens { ") + pynutil.add_weight(punct_graph, weight=1.1) + pynutil.insert(" }") + token = pynutil.insert("tokens { ") + classify + pynutil.insert(" }") + token_plus_punct = ( + pynini.closure(punct + pynutil.insert(" ")) + token + pynini.closure(pynutil.insert(" ") + punct) + ) + + graph = token_plus_punct + pynini.closure(delete_extra_space + token_plus_punct) + graph = delete_space + graph + delete_space + + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/whitelist.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/whitelist.py new file mode 100644 index 0000000000..50c883f803 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/whitelist.py @@ -0,0 +1,68 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path, load_labels +from nemo_text_processing.text_normalization.graph_utils import GraphFst, convert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WhiteListFst(GraphFst): + """ + Finite state transducer for classifying whitelist, e.g. + misses -> tokens { name: "mrs" } + for non-deterministic case: "Dr. Abc" -> + tokens { name: "drive" } tokens { name: "Abc" } + tokens { name: "doctor" } tokens { name: "Abc" } + tokens { name: "Dr." } tokens { name: "Abc" } + This class has highest priority among all classifier grammars. Whitelisted tokens are defined and loaded from "data/whitelist.tsv". + + Args: + input_case: accepting either "lower_cased" or "cased" input. + deterministic: if True will provide a single transduction option, + for False multiple options (used for audio-based normalization) + """ + + def __init__(self, input_case: str, deterministic: bool = True): + super().__init__(name="whitelist", kind="classify") + + def _get_whitelist_graph(input_case, file="data/whitelist.tsv"): + whitelist = load_labels(get_abs_path(file)) + if input_case == "lower_cased": + whitelist = [(x.lower(), y) for x, y in whitelist] + else: + whitelist = [(x, y) for x, y in whitelist] + graph = pynini.string_map(whitelist) + return graph + + def _get_whitelist_non_deterministic_graph(file="data/whitelist_alternatives.tsv"): + whitelist = load_labels(get_abs_path(file)) + whitelist_lower = [(x.lower(), y.lower()) for x, y in whitelist] + whitelist_cased = [(x, y) for x, y in whitelist] + graph = pynini.string_map(whitelist_lower + whitelist_cased) + return graph + + graph = _get_whitelist_graph(input_case) + if not deterministic: + graph |= _get_whitelist_graph("lower_cased") | _get_whitelist_non_deterministic_graph() + + graph = pynutil.insert("name: \"") + convert_space(graph) + pynutil.insert("\"") + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/word.py b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/word.py new file mode 100644 index 0000000000..6ff35ee1f7 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/taggers/word.py @@ -0,0 +1,41 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_SPACE, GraphFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WordFst(GraphFst): + """ + Finite state transducer for classifying word. Considers sentence boundary exceptions. + e.g. sleep -> tokens { name: "sleep" } + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="word", kind="classify", deterministic=deterministic) + + word = pynutil.insert("name: \"") + pynini.closure(NEMO_NOT_SPACE, 1) + pynutil.insert("\"") + self.fst = word.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/token_parser.py b/models/tts/NeMo/nemo_text_processing/text_normalization/token_parser.py new file mode 100644 index 0000000000..d3f7fd95ad --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/token_parser.py @@ -0,0 +1,192 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 string +from collections import OrderedDict +from typing import Dict, List, Union + +PRESERVE_ORDER_KEY = "preserve_order" +EOS = "" + + +class TokenParser: + """ + Parses tokenized/classified text, e.g. 'tokens { money { integer: "20" currency: "$" } } tokens { name: "left"}' + + Args + text: tokenized text + """ + + def __call__(self, text): + """ + Setup function + + Args: + text: text to be parsed + + """ + self.text = text + self.len_text = len(text) + self.char = text[0] # cannot handle empty string + self.index = 0 + + def parse(self) -> List[dict]: + """ + Main function. Implements grammar: + A -> space F space F space F ... space + + Returns list of dictionaries + """ + l = list() + while self.parse_ws(): + token = self.parse_token() + if not token: + break + l.append(token) + return l + + def parse_token(self) -> Dict[str, Union[str, dict]]: + """ + Implements grammar: + F-> no_space KG no_space + + Returns: K, G as dictionary values + """ + d = OrderedDict() + key = self.parse_string_key() + if key is None: + return None + self.parse_ws() + if key == PRESERVE_ORDER_KEY: + self.parse_char(":") + self.parse_ws() + value = self.parse_chars("true") + else: + value = self.parse_token_value() + + d[key] = value + return d + + def parse_token_value(self) -> Union[str, dict]: + """ + Implements grammar: + G-> no_space :"VALUE" no_space | no_space {A} no_space + + Returns: string or dictionary + """ + if self.char == ":": + self.parse_char(":") + self.parse_ws() + self.parse_char("\"") + value_string = self.parse_string_value() + self.parse_char("\"") + return value_string + elif self.char == "{": + d = OrderedDict() + self.parse_char("{") + list_token_dicts = self.parse() + # flatten tokens + for tok_dict in list_token_dicts: + for k, v in tok_dict.items(): + d[k] = v + self.parse_char("}") + return d + else: + raise ValueError() + + def parse_char(self, exp) -> bool: + """ + Parses character + + Args: + exp: character to read in + + Returns true if successful + """ + assert self.char == exp + self.read() + return True + + def parse_chars(self, exp) -> bool: + """ + Parses characters + + Args: + exp: characters to read in + + Returns true if successful + """ + ok = False + for x in exp: + ok |= self.parse_char(x) + return ok + + def parse_string_key(self) -> str: + """ + Parses string key, can only contain ascii and '_' characters + + Returns parsed string key + """ + assert self.char not in string.whitespace and self.char != EOS + incl_criterium = string.ascii_letters + "_" + l = [] + while self.char in incl_criterium: + l.append(self.char) + if not self.read(): + raise ValueError() + + if not l: + return None + return "".join(l) + + def parse_string_value(self) -> str: + """ + Parses string value, ends with quote followed by space + + Returns parsed string value + """ + assert self.char not in string.whitespace and self.char != EOS + l = [] + while self.char != "\"" or self.text[self.index + 1] != " ": + l.append(self.char) + if not self.read(): + raise ValueError() + + if not l: + return None + return "".join(l) + + def parse_ws(self): + """ + Deletes whitespaces. + + Returns true if not EOS after parsing + """ + not_eos = self.char != EOS + while not_eos and self.char == " ": + not_eos = self.read() + return not_eos + + def read(self): + """ + Reads in next char. + + Returns true if not EOS + """ + if self.index < self.len_text - 1: # should be unique + self.index += 1 + self.char = self.text[self.index] + return True + self.char = EOS + return False diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/__init__.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/__init__.py new file mode 100644 index 0000000000..bc443be41c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/cardinal.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/cardinal.py new file mode 100644 index 0000000000..bdf7a2ab38 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/cardinal.py @@ -0,0 +1,56 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class CardinalFst(GraphFst): + """ + Finite state transducer for verbalizing cardinal, e.g. + cardinal { negative: "true" integer: "23" } -> minus twenty three + + Args: + deterministic: if True will provide a single transduction option, + for False multiple options (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="cardinal", kind="verbalize", deterministic=deterministic) + + self.optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "minus ") + delete_space, 0, 1) + + if deterministic: + integer = pynini.closure(NEMO_NOT_QUOTE, 1) + else: + integer = ( + pynini.closure(NEMO_NOT_QUOTE) + + pynini.closure(pynini.cross("hundred ", "") | pynini.cross("hundred ", "hundred and "), 0, 1) + + pynini.closure(NEMO_NOT_QUOTE) + ) + + self.integer = delete_space + pynutil.delete("\"") + integer + pynutil.delete("\"") + integer = pynutil.delete("integer:") + self.integer + + self.numbers = self.optional_sign + integer + delete_tokens = self.delete_tokens(self.numbers) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/date.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/date.py new file mode 100644 index 0000000000..08195927df --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/date.py @@ -0,0 +1,100 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_NOT_QUOTE, + GraphFst, + delete_extra_space, + delete_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class DateFst(GraphFst): + """ + Finite state transducer for verbalizing date, e.g. + date { month: "february" day: "five" year: "twenty twelve" preserve_order: true } -> february fifth twenty twelve + date { day: "five" month: "february" year: "twenty twelve" preserve_order: true } -> the fifth of february twenty twelve + + Args: + ordinal: OrdinalFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, ordinal: GraphFst, deterministic: bool = True): + super().__init__(name="date", kind="verbalize", deterministic=deterministic) + + month = pynini.closure(NEMO_NOT_QUOTE, 1) + day_cardinal = ( + pynutil.delete("day:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + day = day_cardinal @ ordinal.suffix + + if not deterministic: + month |= pynutil.insert(" of ") + month + day |= day_cardinal + + month = pynutil.delete("month:") + delete_space + pynutil.delete("\"") + month + pynutil.delete("\"") + + year = ( + pynutil.delete("year:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + delete_space + + pynutil.delete("\"") + ) + + # month (day) year + graph_mdy = ( + month + pynini.closure(delete_extra_space + day, 0, 1) + pynini.closure(delete_extra_space + year, 0, 1) + ) + + # day month year + graph_dmy = ( + pynutil.insert("the ") + + day + + delete_extra_space + + pynutil.insert("of ") + + month + + pynini.closure(delete_extra_space + year, 0, 1) + ) + + optional_preserve_order = pynini.closure( + pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space + | pynutil.delete("field_order:") + + delete_space + + pynutil.delete("\"") + + NEMO_NOT_QUOTE + + pynutil.delete("\"") + + delete_space + ) + + final_graph = (graph_mdy | year | graph_dmy) + delete_space + optional_preserve_order + + delete_tokens = self.delete_tokens(final_graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/decimal.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/decimal.py new file mode 100644 index 0000000000..9ebefebce9 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/decimal.py @@ -0,0 +1,74 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class DecimalFst(GraphFst): + """ + Finite state transducer for verbalizing decimal, e.g. + decimal { negative: "true" integer_part: "twelve" fractional_part: "five o o six" quantity: "billion" } -> minus twelve point five o o six billion + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, cardinal, deterministic: bool = True): + super().__init__(name="decimal", kind="verbalize", deterministic=deterministic) + + self.optional_sign = pynini.closure(pynini.cross("negative: \"true\"", "minus ") + delete_space, 0, 1) + self.integer = pynutil.delete("integer_part:") + cardinal.integer + self.optional_integer = pynini.closure(self.integer + delete_space + insert_space, 0, 1) + self.fractional_default = ( + pynutil.delete("fractional_part:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + + if deterministic: + self.fractional = pynutil.insert("point ") + self.fractional_default + else: + self.fractional = pynini.closure(pynutil.insert("point "), 0, 1) + self.fractional_default + + self.quantity = ( + delete_space + + insert_space + + pynutil.delete("quantity:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + self.optional_quantity = pynini.closure(self.quantity, 0, 1) + + graph = self.optional_sign + ( + self.integer + | (self.integer + self.quantity) + | (self.optional_integer + self.fractional + self.optional_quantity) + ) + self.numbers = graph + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/electronic.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/electronic.py new file mode 100644 index 0000000000..c51f84a2ac --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/electronic.py @@ -0,0 +1,91 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_ALPHA, + NEMO_NOT_QUOTE, + GraphFst, + delete_space, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class ElectronicFst(GraphFst): + """ + Finite state transducer for verbalizing electronic + e.g. tokens { electronic { username: "cdf1" domain: "abc.edu" } } -> c d f one at a b c dot e d u + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="electronic", kind="verbalize", deterministic=deterministic) + graph_digit = pynini.invert(pynini.string_file(get_abs_path("data/numbers/digit.tsv"))).optimize() + graph_symbols = pynini.string_file(get_abs_path("data/electronic/symbols.tsv")).optimize() + user_name = ( + pynutil.delete("username:") + + delete_space + + pynutil.delete("\"") + + ( + pynini.closure( + pynutil.add_weight(graph_digit + insert_space, 1.09) + | pynutil.add_weight(pynini.closure(graph_symbols + pynutil.insert(" ")), 1.09) + | pynutil.add_weight(NEMO_NOT_QUOTE + insert_space, 1.1) + ) + ) + + pynutil.delete("\"") + ) + + domain_default = ( + pynini.closure(NEMO_NOT_QUOTE + insert_space) + + pynini.cross(".", "dot ") + + NEMO_NOT_QUOTE + + pynini.closure(insert_space + NEMO_NOT_QUOTE) + ) + + server_default = ( + pynini.closure((graph_digit | NEMO_ALPHA) + insert_space, 1) + + pynini.closure(graph_symbols + insert_space) + + pynini.closure((graph_digit | NEMO_ALPHA) + insert_space, 1) + ) + server_common = pynini.string_file(get_abs_path("data/electronic/server_name.tsv")) + insert_space + + domain_common = pynini.cross(".", "dot ") + pynini.string_file(get_abs_path("data/electronic/domain.tsv")) + + domain = ( + pynutil.delete("domain:") + + delete_space + + pynutil.delete("\"") + + (pynutil.add_weight(server_common, 1.09) | pynutil.add_weight(server_default, 1.1)) + + (pynutil.add_weight(domain_common, 1.09) | pynutil.add_weight(domain_default, 1.1)) + + delete_space + + pynutil.delete("\"") + ) + + graph = user_name + delete_space + pynutil.insert("at ") + delete_space + domain + delete_space + + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/fraction.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/fraction.py new file mode 100644 index 0000000000..d49bd4bf5f --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/fraction.py @@ -0,0 +1,72 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, insert_space +from nemo_text_processing.text_normalization.verbalizers.ordinal import OrdinalFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class FractionFst(GraphFst): + """ + Finite state transducer for verbalizing fraction + e.g. tokens { fraction { integer: "twenty three" numerator: "four" denominator: "five" } } -> + twenty three four fifth + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="fraction", kind="verbalize", deterministic=deterministic) + suffix = OrdinalFst().suffix + + integer = pynutil.delete("integer: \"") + pynini.closure(NEMO_NOT_QUOTE) + pynutil.delete("\" ") + numerator = pynutil.delete("numerator: \"") + pynini.closure(NEMO_NOT_QUOTE) + pynutil.delete("\" ") + numerator_one = pynutil.delete("numerator: \"") + pynini.accep("one") + pynutil.delete("\" ") + denominator = pynutil.delete("denominator: \"") + ( + pynini.closure(NEMO_NOT_QUOTE) @ suffix | pynini.cross('four', 'quarter') + ) + conjunction = pynutil.insert("and ") + if not deterministic: + conjunction = pynini.closure(conjunction, 0, 1) + + integer = pynini.closure(integer + insert_space + conjunction, 0, 1) + + denominator_half = pynini.cross("numerator: \"one\" denominator: \"two\"", "a half") + denominator_one_two = pynini.cross("denominator: \"one\"", "over one") | pynini.cross( + "denominator: \"two\"", "halves" + ) + fraction_default = pynutil.add_weight( + numerator + insert_space + denominator + pynutil.insert("s") + pynutil.delete("\""), 0.001 + ) + fraction_with_one = pynutil.add_weight( + numerator_one + insert_space + denominator + pynutil.delete("\""), 0.0001 + ) + + graph = integer + denominator_half | (fraction_with_one | fraction_default) + graph |= pynini.cross("numerator: \"one\" denominator: \"two\"", "one half") + graph |= (numerator | numerator_one) + insert_space + denominator_one_two + + self.graph = graph + delete_tokens = self.delete_tokens(self.graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/measure.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/measure.py new file mode 100644 index 0000000000..3937efe26c --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/measure.py @@ -0,0 +1,76 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, GraphFst, delete_space, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class MeasureFst(GraphFst): + """ + Finite state transducer for verbalizing measure, e.g. + measure { negative: "true" cardinal { integer: "twelve" } units: "kilograms" } -> minus twelve kilograms + measure { decimal { integer_part: "twelve" fractional_part: "five" } units: "kilograms" } -> twelve point five kilograms + tokens { measure { units: "covid" decimal { integer_part: "nineteen" fractional_part: "five" } } } -> covid nineteen point five + + Args: + decimal: DecimalFst + cardinal: CardinalFst + fraction: FractionFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, decimal: GraphFst, cardinal: GraphFst, fraction: GraphFst, deterministic: bool = True): + super().__init__(name="measure", kind="verbalize", deterministic=deterministic) + optional_sign = cardinal.optional_sign + unit = pynutil.delete("units: \"") + pynini.closure(NEMO_CHAR - " ", 1) + pynutil.delete("\"") + delete_space + + graph_decimal = ( + pynutil.delete("decimal {") + + delete_space + + optional_sign + + delete_space + + decimal.numbers + + delete_space + + pynutil.delete("}") + ) + graph_cardinal = ( + pynutil.delete("cardinal {") + + delete_space + + optional_sign + + delete_space + + cardinal.numbers + + delete_space + + pynutil.delete("}") + ) + + graph_fraction = ( + pynutil.delete("fraction {") + delete_space + fraction.graph + delete_space + pynutil.delete("}") + ) + + graph = (graph_cardinal | graph_decimal | graph_fraction) + delete_space + insert_space + unit + + # SH adds "preserve_order: true" by default + preserve_order = pynutil.delete("preserve_order:") + delete_space + pynutil.delete("true") + delete_space + graph |= unit + insert_space + (graph_cardinal | graph_decimal) + delete_space + pynini.closure(preserve_order) + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/money.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/money.py new file mode 100644 index 0000000000..97b95080bf --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/money.py @@ -0,0 +1,117 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import ( + NEMO_NOT_QUOTE, + GraphFst, + delete_space, + get_abs_path, + insert_space, +) + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + + PYNINI_AVAILABLE = False + + +class MoneyFst(GraphFst): + """ + Finite state transducer for verbalizing money, e.g. + money { integer_part: "twelve" fractional_part: "o five" currency: "dollars" } -> twelve o five dollars + + Args: + decimal: DecimalFst + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, decimal: GraphFst, deterministic: bool = True): + super().__init__(name="money", kind="verbalize", deterministic=deterministic) + + def _get_minor_currencies(file): + minor_currencies = [] + with open(get_abs_path(file), 'r') as f: + for line in f: + min_cur = line.strip() + minor_currencies.append(pynutil.insert(min_cur)) + return minor_currencies + + unit = ( + pynutil.delete("currency:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + graph = decimal.numbers + delete_space + pynutil.insert(" ") + unit + + if not deterministic: + minor_currencies_singular = _get_minor_currencies("data/currency/currency_minor_one.tsv") + minor_currencies_singular = pynini.union(*minor_currencies_singular) + minor_currencies_singular = ( + pynini.closure(NEMO_NOT_QUOTE) + + ( + pynini.accep("one") + | pynini.cross("zero one", "one") + | pynini.cross("oh one", "one") + | pynini.cross(" o one", " one") + ) + + insert_space + + minor_currencies_singular + ) + + minor_currencies_plural = _get_minor_currencies("data/currency/currency_minor.tsv") + minor_currencies_plural = insert_space + pynini.union(*minor_currencies_plural) + + fractional_default = ( + pynutil.delete("fractional_part:") + + delete_space + + pynutil.delete("\"") + + ((pynini.closure(NEMO_NOT_QUOTE, 1) + minor_currencies_plural) | minor_currencies_singular) + + pynutil.delete("\"") + ) + + # $2.00 {two zero zero dollars} -> two dollars + fractional_with_zeros = ( + pynutil.delete("fractional_part:") + + delete_space + + pynutil.delete("\"") + + pynini.cross('zero', '') + + pynini.closure(pynini.cross(' zero', '')) + + delete_space + + pynutil.delete("\"") + + delete_space + ) + + fractional = fractional_with_zeros | fractional_default + + graph |= ( + decimal.integer + + delete_space + + insert_space + + unit + + delete_space + + insert_space + + pynini.closure(pynutil.insert("and "), 0, 1) + + fractional + ) + + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/ordinal.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/ordinal.py new file mode 100644 index 0000000000..20ec399a29 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/ordinal.py @@ -0,0 +1,63 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.data_loader_utils import get_abs_path +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, NEMO_SIGMA, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class OrdinalFst(GraphFst): + """ + Finite state transducer for verbalizing ordinal, e.g. + ordinal { integer: "thirteen" } } -> thirteenth + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="ordinal", kind="verbalize", deterministic=deterministic) + + graph_digit = pynini.string_file(get_abs_path("data/ordinals/digit.tsv")).invert() + graph_teens = pynini.string_file(get_abs_path("data/ordinals/teen.tsv")).invert() + + graph = ( + pynutil.delete("integer:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + convert_rest = pynutil.insert("th", weight=0.01) + + suffix = pynini.cdrewrite( + graph_digit | graph_teens | pynutil.add_weight(pynini.cross("ty", "tieth"), weight=0.001) | convert_rest, + "", + "[EOS]", + NEMO_SIGMA, + ).optimize() + graph = graph @ suffix + self.suffix = suffix + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/roman.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/roman.py new file mode 100644 index 0000000000..bb42f3c522 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/roman.py @@ -0,0 +1,46 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst +from nemo_text_processing.text_normalization.verbalizers.ordinal import OrdinalFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class RomanFst(GraphFst): + """ + Finite state transducer for verbalizing roman numerals + e.g. tokens { roman { integer: "one" } } -> one + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="roman", kind="verbalize", deterministic=deterministic) + suffix = OrdinalFst().suffix + + integer = pynini.closure(NEMO_NOT_QUOTE) + integer |= pynini.closure(pynutil.insert("the "), 0, 1) + integer @ suffix + graph = pynutil.delete("integer: \"") + integer + pynutil.delete("\"") + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/telephone.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/telephone.py new file mode 100644 index 0000000000..3920819adb --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/telephone.py @@ -0,0 +1,69 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TelephoneFst(GraphFst): + """ + Finite state transducer for verbalizing telephone numbers, e.g. + telephone { country_code: "one" number_part: "one two three, one two three, five six seven eight" extension: "one" } + -> one, one two three, one two three, five six seven eight, one + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="telephone", kind="verbalize", deterministic=deterministic) + + add_separator = pynutil.insert(",") # between components + + optional_country_code = pynini.closure( + pynutil.delete("country_code: \"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + + delete_space + + add_separator + + insert_space, + 0, + 1, + ) + + number_part = pynutil.delete("number_part: \"") + pynini.closure(NEMO_NOT_QUOTE, 1) + pynutil.delete("\"") + + optional_extension = pynini.closure( + delete_space + + add_separator + + insert_space + + pynutil.delete("extension: \"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\""), + 0, + 1, + ) + + graph = optional_country_code + number_part + optional_extension + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/time.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/time.py new file mode 100644 index 0000000000..c5881c3ee7 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/time.py @@ -0,0 +1,74 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_NOT_QUOTE, GraphFst, delete_space, insert_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class TimeFst(GraphFst): + """ + Finite state transducer for verbalizing time, e.g. + time { hours: "twelve" minutes: "thirty" suffix: "a m" zone: "e s t" } -> twelve thirty a m e s t + time { hours: "twelve" } -> twelve o'clock + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="time", kind="verbalize", deterministic=deterministic) + hour = ( + pynutil.delete("hours:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + minute = ( + pynutil.delete("minutes:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + suffix = ( + pynutil.delete("suffix:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + optional_suffix = pynini.closure(delete_space + insert_space + suffix, 0, 1) + zone = ( + pynutil.delete("zone:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_NOT_QUOTE, 1) + + pynutil.delete("\"") + ) + optional_zone = pynini.closure(delete_space + insert_space + zone, 0, 1) + graph = hour + delete_space + insert_space + minute + optional_suffix + optional_zone + graph |= hour + insert_space + pynutil.insert("o'clock") + optional_zone + graph |= hour + delete_space + insert_space + suffix + optional_zone + delete_tokens = self.delete_tokens(graph) + self.fst = delete_tokens.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/verbalize.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/verbalize.py new file mode 100644 index 0000000000..b14468a8f0 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/verbalize.py @@ -0,0 +1,79 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst +from nemo_text_processing.text_normalization.verbalizers.cardinal import CardinalFst +from nemo_text_processing.text_normalization.verbalizers.date import DateFst +from nemo_text_processing.text_normalization.verbalizers.decimal import DecimalFst +from nemo_text_processing.text_normalization.verbalizers.electronic import ElectronicFst +from nemo_text_processing.text_normalization.verbalizers.fraction import FractionFst +from nemo_text_processing.text_normalization.verbalizers.measure import MeasureFst +from nemo_text_processing.text_normalization.verbalizers.money import MoneyFst +from nemo_text_processing.text_normalization.verbalizers.ordinal import OrdinalFst +from nemo_text_processing.text_normalization.verbalizers.roman import RomanFst +from nemo_text_processing.text_normalization.verbalizers.telephone import TelephoneFst +from nemo_text_processing.text_normalization.verbalizers.time import TimeFst +from nemo_text_processing.text_normalization.verbalizers.whitelist import WhiteListFst + + +class VerbalizeFst(GraphFst): + """ + Composes other verbalizer grammars. + For deployment, this grammar will be compiled and exported to OpenFst Finate State Archiv (FAR) File. + More details to deployment at NeMo/tools/text_processing_deployment. + + Args: + deterministic: if True will provide a single transduction option, + for False multiple options (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="verbalize", kind="verbalize", deterministic=deterministic) + cardinal = CardinalFst(deterministic=deterministic) + cardinal_graph = cardinal.fst + decimal = DecimalFst(cardinal=cardinal, deterministic=deterministic) + decimal_graph = decimal.fst + ordinal = OrdinalFst(deterministic=deterministic) + ordinal_graph = ordinal.fst + fraction = FractionFst(deterministic=deterministic) + fraction_graph = fraction.fst + telephone_graph = TelephoneFst(deterministic=deterministic).fst + electronic_graph = ElectronicFst(deterministic=deterministic).fst + measure = MeasureFst(decimal=decimal, cardinal=cardinal, fraction=fraction, deterministic=deterministic) + measure_graph = measure.fst + time_graph = TimeFst(deterministic=deterministic).fst + date_graph = DateFst(ordinal=ordinal, deterministic=deterministic).fst + money_graph = MoneyFst(decimal=decimal, deterministic=deterministic).fst + whitelist_graph = WhiteListFst(deterministic=deterministic).fst + + graph = ( + time_graph + | date_graph + | money_graph + | measure_graph + | ordinal_graph + | decimal_graph + | cardinal_graph + | telephone_graph + | electronic_graph + | fraction_graph + | whitelist_graph + ) + + if not deterministic: + roman_graph = RomanFst(deterministic=deterministic).fst + graph |= roman_graph + + self.fst = graph diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/verbalize_final.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/verbalize_final.py new file mode 100644 index 0000000000..c2b2ac8647 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/verbalize_final.py @@ -0,0 +1,54 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import GraphFst, delete_extra_space, delete_space +from nemo_text_processing.text_normalization.verbalizers.verbalize import VerbalizeFst +from nemo_text_processing.text_normalization.verbalizers.word import WordFst + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class VerbalizeFinalFst(GraphFst): + """ + Finite state transducer that verbalizes an entire sentence, e.g. + tokens { name: "its" } tokens { time { hours: "twelve" minutes: "thirty" } } tokens { name: "now" } tokens { name: "." } -> its twelve thirty now . + + Args: + deterministic: if True will provide a single transduction option, + for False multiple options (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="verbalize_final", kind="verbalize", deterministic=deterministic) + verbalize = VerbalizeFst(deterministic=deterministic).fst + word = WordFst(deterministic=deterministic).fst + types = verbalize | word + graph = ( + pynutil.delete("tokens") + + delete_space + + pynutil.delete("{") + + delete_space + + types + + delete_space + + pynutil.delete("}") + ) + graph = delete_space + pynini.closure(graph + delete_extra_space) + graph + delete_space + self.fst = graph diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/whitelist.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/whitelist.py new file mode 100644 index 0000000000..5d1889d9b1 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/whitelist.py @@ -0,0 +1,48 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, NEMO_SIGMA, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WhiteListFst(GraphFst): + """ + Finite state transducer for verbalizing whitelist + e.g. tokens { name: "misses" } } -> misses + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="whitelist", kind="verbalize", deterministic=deterministic) + graph = ( + pynutil.delete("name:") + + delete_space + + pynutil.delete("\"") + + pynini.closure(NEMO_CHAR - " ", 1) + + pynutil.delete("\"") + ) + graph = graph @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", NEMO_SIGMA) + self.fst = graph.optimize() diff --git a/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/word.py b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/word.py new file mode 100644 index 0000000000..8b1fcac449 --- /dev/null +++ b/models/tts/NeMo/nemo_text_processing/text_normalization/verbalizers/word.py @@ -0,0 +1,43 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +from nemo_text_processing.text_normalization.graph_utils import NEMO_CHAR, NEMO_SIGMA, GraphFst, delete_space + +try: + import pynini + from pynini.lib import pynutil + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + PYNINI_AVAILABLE = False + + +class WordFst(GraphFst): + """ + Finite state transducer for verbalizing word + e.g. tokens { name: "sleep" } -> sleep + + Args: + deterministic: if True will provide a single transduction option, + for False multiple transduction are generated (used for audio-based normalization) + """ + + def __init__(self, deterministic: bool = True): + super().__init__(name="word", kind="verbalize", deterministic=deterministic) + chars = pynini.closure(NEMO_CHAR - " ", 1) + char = pynutil.delete("name:") + delete_space + pynutil.delete("\"") + chars + pynutil.delete("\"") + graph = char @ pynini.cdrewrite(pynini.cross(u"\u00A0", " "), "", "", NEMO_SIGMA) + + self.fst = graph.optimize() diff --git a/models/tts/NeMo/reinstall.sh b/models/tts/NeMo/reinstall.sh new file mode 100755 index 0000000000..8e2eb10792 --- /dev/null +++ b/models/tts/NeMo/reinstall.sh @@ -0,0 +1,40 @@ +#!/usr/bin/env bash +set -e + +INSTALL_OPTION=${1:-"dev"} + + +PIP=pip + +echo 'Uninstalling stuff' +${PIP} uninstall -y nemo_toolkit +${PIP} uninstall -y sacrebleu + +# Kept for legacy purposes +${PIP} uninstall -y nemo_asr +${PIP} uninstall -y nemo_nlp +${PIP} uninstall -y nemo_tts +${PIP} uninstall -y nemo_simple_gan +${PIP} uninstall -y nemo_cv + +${PIP} install -U setuptools + +echo 'Installing nemo and nemo_text_processing' +if [[ "$INSTALL_OPTION" == "dev" ]]; then + ${PIP} install --editable ".[all]" +else + rm -rf dist/ + python setup.py bdist_wheel + DIST_FILE=$(find ./dist -name "*.whl" | head -n 1) + ${PIP} install "${DIST_FILE}[all]" +fi + +echo 'Installing additional nemo_text_processing conda dependency' +bash nemo_text_processing/setup.sh > /dev/null 2>&1 && echo "nemo_text_processing installed!" || echo "nemo_text_processing could not be installed!" + +if [ -x "$(command -v conda)" ]; then + echo 'Attempting update to numba installation via conda' + conda update -c numba numba -y > /dev/null 2>&1 && echo "Numba updated!" || echo "Numba could not be updated!" +fi + +echo 'All done!' diff --git a/models/tts/NeMo/requirements/requirements.txt b/models/tts/NeMo/requirements/requirements.txt new file mode 100644 index 0000000000..ac7d1a3dd4 --- /dev/null +++ b/models/tts/NeMo/requirements/requirements.txt @@ -0,0 +1,19 @@ +numpy>=1.18.2 +onnx>=1.7.0 +pytorch-lightning>=1.3.0 +python-dateutil +torch +wget +wrapt +ruamel.yaml +scikit-learn +omegaconf>=2.1.0 +hydra-core>=1.1.0 +transformers>=4.0.1 +sentencepiece<1.0.0 +webdataset>=0.1.48,<=0.1.62 +tqdm>=4.41.0 +opencc +pangu +jieba +numba diff --git a/models/tts/NeMo/requirements/requirements_asr.txt b/models/tts/NeMo/requirements/requirements_asr.txt new file mode 100644 index 0000000000..1f8b8419db --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_asr.txt @@ -0,0 +1,24 @@ +braceexpand +editdistance +frozendict +inflect +kaldi-io +librosa +marshmallow +packaging +ruamel.yaml +soundfile +sox +unidecode +kaldi-python-io +kaldiio +scipy +pandas +g2p_en +pydub +pyannote.core +pyannote.metrics +pyamg +torch-stft +ipywidgets +matplotlib diff --git a/models/tts/NeMo/requirements/requirements_cv.txt b/models/tts/NeMo/requirements/requirements_cv.txt new file mode 100644 index 0000000000..f0eabe59f4 --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_cv.txt @@ -0,0 +1,2 @@ +pillow +torchvision diff --git a/models/tts/NeMo/requirements/requirements_docs.txt b/models/tts/NeMo/requirements/requirements_docs.txt new file mode 100644 index 0000000000..8a380b7020 --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_docs.txt @@ -0,0 +1,6 @@ +latexcodec +sphinx>=3.0 +sphinx_rtd_theme +sphinxcontrib-bibtex +wrapt +numpy diff --git a/models/tts/NeMo/requirements/requirements_nlp.txt b/models/tts/NeMo/requirements/requirements_nlp.txt new file mode 100644 index 0000000000..b2a0d4b4f7 --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_nlp.txt @@ -0,0 +1,13 @@ +boto3 +h5py +matplotlib>=3.3.2 +sentencepiece +unidecode +youtokentome>=1.0.5 +numpy +rapidfuzz +gdown +megatron-lm==2.2.0 +inflect +sacrebleu[ja] +sacremoses>=0.0.43 diff --git a/models/tts/NeMo/requirements/requirements_test.txt b/models/tts/NeMo/requirements/requirements_test.txt new file mode 100644 index 0000000000..7c629176cb --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_test.txt @@ -0,0 +1,11 @@ +black==19.10b0 +isort[requirements] < 5 +parameterized +pytest +pytest-runner +ruamel.yaml +sphinx +sphinxcontrib-bibtex +wrapt +wget +wandb diff --git a/models/tts/NeMo/requirements/requirements_text_processing.txt b/models/tts/NeMo/requirements/requirements_text_processing.txt new file mode 100644 index 0000000000..9b7085f3ed --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_text_processing.txt @@ -0,0 +1,2 @@ +inflect +regex diff --git a/models/tts/NeMo/requirements/requirements_tts.txt b/models/tts/NeMo/requirements/requirements_tts.txt new file mode 100644 index 0000000000..f8721050e4 --- /dev/null +++ b/models/tts/NeMo/requirements/requirements_tts.txt @@ -0,0 +1,7 @@ +matplotlib +pypinyin +attrdict +pystoi +pesq +g2p_en +nltk diff --git a/models/tts/NeMo/scripts/asr_language_modeling/neural_rescorer/create_tarred_transformer_lm_dataset.py b/models/tts/NeMo/scripts/asr_language_modeling/neural_rescorer/create_tarred_transformer_lm_dataset.py new file mode 100644 index 0000000000..6c9a836914 --- /dev/null +++ b/models/tts/NeMo/scripts/asr_language_modeling/neural_rescorer/create_tarred_transformer_lm_dataset.py @@ -0,0 +1,318 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Execution : + +python create_tarred_tokenized_text_lm_dataset.py \ + --text_path= \ + --data_root= \ + --tokenizer_name="bert-base-cased" \ + --tokenizer_vocab_file= \ + --num_shards=64 \ + --chunk_size=8192 \ + --chunk_write_buffer=512 \ + --lower_case \ + --log +""" + + +import argparse +import glob +import json +import logging +import os +import tarfile + +import joblib +import numpy as np +from tqdm import tqdm + +from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer + +parser = argparse.ArgumentParser(description='Tarred Tokenized dataset for text language modelling') + +# Data path arguments +parser.add_argument('--text_path', required=True, default=None, type=str, help='Text paths, seperated by commas') +parser.add_argument('--data_root', required=True, default=None, type=str, help='Output directory') + +# General arguments +parser.add_argument( + '--chunk_write_buffer', + default=128, + type=int, + help='Number of chunks of `chunk_size` to buffer for parallel tokenization and serial write to disk', +) +parser.add_argument('--lower_case', action='store_true', help='Whether to lower case the corpus') +parser.add_argument('--log', action='store_true', help='Whether to print logs to terminal') + +# Tokenizer arguments +parser.add_argument('--tokenizer_name', required=False, default=None, type=str, help='Tokenizer name for resolution') +parser.add_argument( + '--tokenizer_model', required=False, default=None, type=str, help='Path to tokenizer model for sentencepiece' +) +parser.add_argument('--tokenizer_vocab_file', required=False, type=str, default=None, help='Path to a vocab file') +parser.add_argument( + '--tokenizer_special_tokens', default=None, type=str, nargs='+', help='List of special tokens for the tokenizer' +) + +# Tarred dataset arguments +parser.add_argument('--num_shards', default=1, type=int, help='Number of shards for the tarfile') +parser.add_argument('--chunk_size', default=8192, type=int, help='Number of rows of data concatenated into a vector') + +parser.set_defaults(log=False, lower_case=False) +args = parser.parse_args() + + +def __build_dataset_from_text(texts: str, lower_case: bool, chunk_size: int): + if ',' in texts: + texts = texts.split(',') + else: + texts = [texts] + + num_lines = 0 + text_dataset = [] + + for text in texts: + with open(text, 'r') as in_reader: + reader = tqdm(iter(lambda: in_reader.readline(), ''), desc="Read 0 lines", unit=' lines') + + for i, line in enumerate(reader): + # Clean text line + line = line.replace("\n", "").strip() + + if lower_case: + line = line.lower() + + if line: + text_dataset.append(line) + + num_lines += 1 + if num_lines % 100000 == 0: + reader.set_description(f"Read {num_lines} lines") + + if num_lines % chunk_size == 0: + yield text_dataset, num_lines + + # Empty cache + text_dataset = [] + + logging.info(f"Finished extracting manifest : {text}") + + logging.info("Finished extracting all manifests ! Number of sentences : {}".format(num_lines)) + + if len(text_dataset) != 0: + yield text_dataset, num_lines + + +def __tokenize_str(texts, tokenizer): + tokenized_text = [] + for text in texts: + tok_text = tokenizer.text_to_ids(text) + tokenized_text.extend(tok_text) + return tokenized_text + + +def __tokenize_text( + text_paths, tokenizer, tokenized_cachedir, lower_case: bool = False, chunk_size=8192, write_buffer: int = -1 +): + if write_buffer < 1: + write_buffer = max(os.cpu_count() - write_buffer, 1) + + logging.info(f"Using write chunk buffer of size {write_buffer}") + + if not os.path.exists(tokenized_cachedir): + os.makedirs(tokenized_cachedir) + + # global parameters + global_chunk_idx = 0 + chunk_paths = [] + chunk_lens = [] + + # buffer parameters + data_cache = [] + chunk_idx = 0 + + text_generator = iter(__build_dataset_from_text(text_paths, lower_case=lower_case, chunk_size=chunk_size)) + global_num_lines = 0 + last_batch = False + + with joblib.Parallel(n_jobs=-2, verbose=10) as parallel: + + while True: + try: + data, num_lines = next(text_generator) + data_cache.append(data) + + global_num_lines = num_lines + except StopIteration: + last_batch = True + + # Update counters + chunk_idx += 1 + + if (chunk_idx == write_buffer) or last_batch: + # write the chunks into disk after parallel tokenization + tokenized_data_list = parallel( + joblib.delayed(__tokenize_str)(chunk, tokenizer) for chunk in data_cache + ) + + # Sequential write cache + for chunk in tokenized_data_list: + fp = os.path.join(tokenized_cachedir, f"chunk_{global_chunk_idx}.npy") + chunk = np.asarray(chunk, dtype=np.int64) + np.save(fp, chunk, allow_pickle=False) + + chunk_paths.append(fp) + chunk_lens.append(len(chunk)) + + global_chunk_idx += 1 + + logging.info(f"Wrote a total of {global_chunk_idx} chunks to file...") + + # reset buffers + data_cache.clear() + del data_cache + + data_cache = [] + chunk_idx = 0 + + if last_batch: + logging.info("Finished tokenizing last chunk") + break + + logging.info( + f"Chunking {global_num_lines} rows into {global_num_lines // chunk_size} tasks (each chunk contains {chunk_size} elements)" + ) + return chunk_paths, chunk_lens + + +def __create_chunk(data_root, chunk_path, shard_id, compute_metrics=False): + """Creates a tarball containing the tokenized text chunks. + """ + tar = tarfile.open(os.path.join(data_root, f'text_{shard_id}.tar'), mode='a') + + # We squash the filename since we do not preserve directory structure of tokenized text in the tarball. + base, ext = os.path.splitext(chunk_path) + base = base.replace(os.pathsep, '_') + # Need the following replacement as long as WebDataset splits on first period + base = base.replace('.', '_') + squashed_filename = f'{base}{ext}' + tar.add(chunk_path, arcname=squashed_filename) + + tar.close() + + if compute_metrics: + data = np.load(chunk_path, allow_pickle=False) + chunk_len = len(data) + return (chunk_len,) + else: + return None + + +def __write_tarred_tokenized_text_dataset(data_root, num_shards, chunk_paths, chunk_lens): + num_chunks = len(chunk_paths) + + if chunk_lens is not None: + num_text = sum(chunk_lens) + shard_counts = {chunk_id: chunk_len for chunk_id, chunk_len in enumerate(chunk_lens)} + compute_metrics = False + + else: + num_text = 0 + shard_counts = {} + compute_metrics = True + + for chunk_id, chunk_path in enumerate(tqdm(chunk_paths, desc='Writing chunk ', unit=' chunks')): + shard_id = chunk_id % num_shards + metrics = __create_chunk(data_root, chunk_path, shard_id, compute_metrics=compute_metrics) + + if metrics is not None: + num_text += metrics[0] + shard_counts[chunk_id] = metrics[0] + + # write metadata + metadata_path = os.path.join(data_root, 'metadata.json') + with open(metadata_path, 'w') as f: + metadata = {'num_chunks': num_chunks, 'num_text': num_text, 'shard_count': shard_counts} + json.dump(metadata, f, indent=4) + logging.info("Metadata writen..") + + +def main(): + text_path = args.text_path + data_root = args.data_root + + if args.log: + logging.basicConfig(level=logging.INFO) + + tokenized_cachedir = os.path.join(data_root, '_tokenized_dataset_cachedir') + + if os.path.exists(tokenized_cachedir): + logging.warning( + f'Tokenized cache directory {tokenized_cachedir} already potentially contains files.' + f'In such a case, please be aware that the tarfiles will be **appended** instead of overridden!' + ) + + if not os.path.exists(data_root): + os.makedirs(data_root) + + chunk_paths = None + chunk_lens = None + + if os.path.exists(tokenized_cachedir): + paths = glob.glob(os.path.join(tokenized_cachedir, "*.npy")) + if len(paths) > 0: + logging.info("Cached tokenized numpy files found, skipping re-tokenization of dataset") + + chunk_paths = paths + chunk_lens = None + + if chunk_paths is None: + if args.tokenizer_name is None: + raise ValueError("`tokenizer_name` name is required when tokenizing the dataset for the first time.") + + if args.tokenizer_vocab_file is None: + raise ValueError("`tokenizer_vocab_file` is required when constructing the tokenized dataset") + + tokenizer = get_tokenizer( + tokenizer_name=args.tokenizer_name, + tokenizer_model=args.tokenizer_model, + vocab_file=args.tokenizer_vocab_file, + special_tokens=args.tokenizer_special_tokens, + ) + + logging.info("Built tokenizer") + + # tokenize text data into sub-words + chunk_paths, chunk_lens = __tokenize_text( + text_paths=text_path, + tokenizer=tokenizer, + tokenized_cachedir=tokenized_cachedir, + lower_case=args.lower_case, + chunk_size=args.chunk_size, + write_buffer=args.chunk_write_buffer, + ) + logging.info(f"Tokenized dataset into sub-words and serialized cache at {tokenized_cachedir}") + + # Write tarred dataset + __write_tarred_tokenized_text_dataset( + data_root, num_shards=args.num_shards, chunk_paths=chunk_paths, chunk_lens=chunk_lens + ) + + logging.info('Done preparing tokenized dataset!') + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py b/models/tts/NeMo/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py new file mode 100644 index 0000000000..f713af6049 --- /dev/null +++ b/models/tts/NeMo/scripts/asr_language_modeling/neural_rescorer/eval_neural_rescorer.py @@ -0,0 +1,288 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +""" +This script would evaluate a neural language model (Transformer) trained with +`examples/nlp/language_modeling/transformer_lm.py' as a rescorer for ASR systems. +Given a trained TransformerLMModel `.nemo` file, this script can be used to re-score the beams obtained from a beam +search decoder of an ASR model. + +USAGE: +1. Obtain `.tsv` file with beams and their corresponding scores. Scores can be from a regular beam search decoder or + in fusion with an N-gram LM scores. For a given beam size `beam_size` and a number of examples + for evaluation `num_eval_examples`, it should contain (`beam_size` x `num_eval_examples`) lines of + form `beam_candidate_text \t score`. This file can be generated by `scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py`. + +2. Rescore the candidates: + python eval_neural_rescorer.py + --lm_model=[path to .nemo file of the LM] + --beams_file=[path to beams .tsv file] + --beam_size=[size of the beams] + --eval_manifest=[path to eval manifest .json file] + --batch_size=[batch size used for inference on the LM model] + --alpha=[the value for the parameter rescorer_alpha] + --beta=[the value for the parameter rescorer_beta] + +You may find more info on how to use this script at: +https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html + +""" + +import contextlib +import json +from argparse import ArgumentParser + +import editdistance +import matplotlib.pyplot as plt +import numpy as np +import pandas as pd +import torch +import tqdm + +from nemo.collections.nlp.models.language_modeling import TransformerLMModel +from nemo.utils import logging + + +class BeamScoresDataset(torch.utils.data.Dataset): + """ + Dataset to read the score file containing the beams and their score + + Args: + data_path: path to the beams file + tokenizer: tokenizer of the LM model + manifest_path: manifest `.json` file which contains the ground truths transcripts + beam_size: the number of beams per sample + max_seq_length: the maximum length of sequences + """ + + def __init__(self, data_path, tokenizer, manifest_path, beam_size=128, max_seq_length=256): + self.data = pd.read_csv(data_path, delimiter="\t", header=None) + self.tokenizer = tokenizer + self.ground_truths = [] + with open(manifest_path, 'r') as f_orig: + for line in f_orig: + item = json.loads(line) + self.ground_truths.append(item['text']) + self.beam_size = beam_size + self.max_seq_length = max_seq_length + + def __len__(self): + return len(self.data) + + def __getitem__(self, idx): + text = str(self.data[0][idx]) + tokens = [self.tokenizer.bos_id] + self.tokenizer.text_to_ids(text) + [self.tokenizer.eos_id] + input_ids = [self.tokenizer.pad_id] * self.max_seq_length + input_ids[: len(tokens)] = tokens + input_ids = np.array(input_ids) + input_mask = (input_ids != self.tokenizer.pad_id).astype(np.float32) + acoustic_score = self.data[1][idx] + dist = editdistance.eval(text.split(), self.ground_truths[idx // self.beam_size].split()) + ref_len = len(self.ground_truths[idx // self.beam_size].split()) + len_in_chars = len(str(self.data[0][idx])) + return input_ids, input_mask, acoustic_score, dist, ref_len, len_in_chars, idx + + +def linear_search_wer( + dists, scores1, scores2, total_len, coef_range=[0, 10], coef_steps=10000, param_name='parameter' +): + """ + performs linear search to find the best coefficient when two set of scores are getting linearly fused. + + Args: + dists: Tesnor of the distances between the ground truth and the candidates with shape of [number of samples, beam size] + scores1: Tensor of the first set of scores with shape of [number of samples, beam size] + scores2: Tensor of the second set of scores with shape of [number of samples, beam size] + total_len: The total length of all samples + coef_range: the search range for the coefficient + coef_steps: the number of steps that the search range would get divided into + param_name: the name of the parameter to be used in the figure + + Output: + (best coefficient found, best WER achived) + """ + scale = scores1.mean().abs().item() / scores2.mean().abs().item() + left = coef_range[0] * scale + right = coef_range[1] * scale + coefs = np.linspace(left, right, coef_steps) + + best_wer = 10000 + best_coef = left + wers = [] + for coef in coefs: + scores = scores1 + coef * scores2 + wer = compute_wer(dists, scores, total_len) + wers.append(wer) + if wer < best_wer: + best_wer = wer + best_coef = coef + + plt.plot(coefs, wers) + plt.title(f'WER% after rescoring with different values of {param_name}') + plt.ylabel('WER%') + plt.xlabel(param_name) + plt.show() + return best_coef, best_wer + + +def compute_wer(dists, scores, total_len): + """ + Sorts the candidates based on the scores and calculates the WER with the new top candidates. + + Args: + dists: Tensor of the distances between the ground truth and the candidates with shape of [number of samples, beam size] + scores: Tensor of the scores for candidates with shape of [number of samples, beam size] + total_len: The total length of all samples + + Output: + WER with the new scores + """ + indices = scores.max(dim=1, keepdim=True)[1] + wer = dists.gather(dim=1, index=indices).sum() / total_len + wer = wer.item() + return wer + + +def main(): + parser = ArgumentParser() + parser.add_argument("--lm_model_file", type=str, required=True, help="path to LM model .nemo file") + parser.add_argument("--beams_file", type=str, required=True, help="path to beams .tsv file") + parser.add_argument( + "--eval_manifest", type=str, required=True, help="path to the evaluation `.json` manifest file" + ) + parser.add_argument("--beam_size", type=int, required=True, help="number of beams per candidate") + parser.add_argument("--batch_size", type=int, default=256, help="inference batch size") + parser.add_argument("--alpha", type=float, default=None, help="parameter alpha of the fusion") + parser.add_argument("--beta", type=float, default=None, help="parameter beta of the fusion") + parser.add_argument( + "--scores_output_file", default=None, type=str, help="The optional path to store the rescored beams" + ) + parser.add_argument( + "--device", default="cuda", type=str, help="The device to load the model onto to calculate the scores" + ) + parser.add_argument( + "--use_amp", action="store_true", help="Whether to use AMP if available to calculate the scores" + ) + args = parser.parse_args() + + device = args.device + if device.startswith("cuda") and not torch.cuda.is_available(): + logging.info(f"cuda is not available! switched to cpu.") + device = "cpu" + + if args.lm_model_file.endswith(".nemo"): + logging.info("Attempting to initialize from .nemo file") + model = TransformerLMModel.restore_from( + restore_path=args.lm_model_file, map_location=torch.device(device) + ).eval() + else: + raise NotImplementedError(f"Only supports .nemo files, but got: {args.model}") + + max_seq_length = model.encoder._embedding.position_embedding.pos_enc.shape[0] + dataset = BeamScoresDataset(args.beams_file, model.tokenizer, args.eval_manifest, args.beam_size, max_seq_length) + data_loader = torch.utils.data.DataLoader(dataset=dataset, batch_size=args.batch_size) + + if args.use_amp: + if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'): + logging.info("AMP is enabled!\n") + autocast = torch.cuda.amp.autocast + else: + + @contextlib.contextmanager + def autocast(): + yield + + logging.info(f"Rescoring with beam_size: {args.beam_size}") + logging.info("Calculating the scores...") + with autocast(): + with torch.no_grad(): + am_scores, lm_scores, dists, ref_lens, lens_in_chars = [], [], [], [], [] + for batch in tqdm.tqdm(data_loader): + input_ids, input_mask, acoustic_score, dist, ref_len, len_in_chars, idx = batch + + max_len_in_batch = input_mask.sum(dim=0).argmin().item() + input_ids, input_mask = input_ids[:, :max_len_in_batch], input_mask[:, :max_len_in_batch] + if torch.cuda.is_available(): + input_ids, input_mask = input_ids.to(device), input_mask.to(device) + dist, acoustic_score, len_in_chars = ( + dist.to(device), + acoustic_score.to(device), + len_in_chars.to(device), + ) + + log_probs = model.forward(input_ids[:, :-1], input_mask[:, :-1]) + target_log_probs = log_probs.gather(2, input_ids[:, 1:].unsqueeze(2)).squeeze(2) + neural_lm_score = torch.sum(target_log_probs * input_mask[:, 1:], dim=-1) + + am_scores.append(acoustic_score) + lm_scores.append(neural_lm_score) + dists.append(dist) + ref_lens.append(ref_len) + lens_in_chars.append(len_in_chars) + + am_scores = torch.cat(am_scores).view(-1, args.beam_size) + lm_scores = torch.cat(lm_scores).view(-1, args.beam_size) + dists = torch.cat(dists).view(-1, args.beam_size) + ref_lens = torch.cat(ref_lens).view(-1, args.beam_size) + lens_in_chars = torch.cat(lens_in_chars).view(-1, args.beam_size).to(am_scores.dtype) + + total_len = ref_lens[:, 0].sum() + model_wer = dists[:, 0].sum() / total_len + ideal_wer = dists.min(dim=1)[0].sum() / total_len + + if args.alpha is None: + logging.info("Linear search for alpha...") + coef1, _ = linear_search_wer( + dists=dists, scores1=am_scores, scores2=lm_scores, total_len=total_len, param_name='alpha' + ) + coef1 = np.round(coef1, 3) + logging.info(f"alpha={coef1} achieved the best WER.") + logging.info(f"------------------------------------------------") + else: + coef1 = args.alpha + + scores = am_scores + coef1 * lm_scores + + if args.beta is None: + logging.info("Linear search for beta...") + coef2, _ = linear_search_wer( + dists=dists, scores1=scores, scores2=lens_in_chars, total_len=total_len, param_name='beta' + ) + coef2 = np.round(coef2, 3) + logging.info(f"beta={coef2} achieved the best WER.") + logging.info(f"------------------------------------------------") + else: + coef2 = args.beta + + new_scores = am_scores + coef1 * lm_scores + coef2 * lens_in_chars + rescored_wer = compute_wer(dists, new_scores, total_len) + + logging.info(f"Input beams WER: {np.round(model_wer.item() * 100, 2)}%") + logging.info(f"------------------------------------------------") + logging.info(f" +LM rescoring WER: {np.round(rescored_wer * 100, 2)}%") + logging.info(f" with alpha={coef1}, beta={coef2}") + logging.info(f"------------------------------------------------") + logging.info(f"Best possible WER: {np.round(ideal_wer.item() * 100, 2)}%") + logging.info(f"------------------------------------------------") + + new_scores_flatten = new_scores.flatten() + if args.scores_output_file is not None: + logging.info(f'Saving the candidates with their new scores at `{args.scores_output_file}`...') + with open(args.scores_output_file, "w") as fout: + for sample_id in range(len(dataset)): + fout.write(f"{dataset.data[0][sample_id]}\t{new_scores_flatten[sample_id]}\n") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py new file mode 100644 index 0000000000..6f5b224bef --- /dev/null +++ b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/eval_beamsearch_ngram.py @@ -0,0 +1,345 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +# + +# This script would evaluate an N-gram language model trained with KenLM library (https://github.com/kpu/kenlm) in +# fusion with beam search decoders on top of a trained ASR model. NeMo's beam search decoders are capable of using the +# KenLM's N-gram models to find the best candidates. This script supports both character level and BPE level +# encodings and models which is detected automatically from the type of the model. +# You may train the LM model with 'scripts/ngram_lm/train_kenlm.py'. +# +# USAGE: python eval_beamsearch_ngram.py --nemo_model_file \ +# --input_manifest \ +# --beam_width \ +# --beam_alpha \ +# --beam_width \ +# --preds_output_folder \ +# --decoding_mode beam_search_ngram +# ... +# +# You may find more info on how to use this script at: +# https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html + + +# Please check train_kenlm.py to find out why we need TOKEN_OFFSET for BPE-based models +TOKEN_OFFSET = 100 + +import argparse +import contextlib +import json +import os +import pickle + +import editdistance +import kenlm_utils +import numpy as np +import torch +from sklearn.model_selection import ParameterGrid +from tqdm.auto import tqdm + +import nemo +import nemo.collections.asr as nemo_asr +from nemo.utils import logging + + +def beam_search_eval( + all_probs, + target_transcripts, + vocab, + ids_to_text_func=None, + preds_output_file=None, + lm_path=None, + beam_alpha=1.0, + beam_beta=0.0, + beam_width=128, + beam_batch_size=128, + progress_bar=True, +): + # creating the beam search decoder + beam_search_lm = nemo_asr.modules.BeamSearchDecoderWithLM( + vocab=vocab, + beam_width=beam_width, + alpha=beam_alpha, + beta=beam_beta, + lm_path=lm_path, + num_cpus=max(os.cpu_count(), 1), + input_tensor=False, + ) + + wer_dist_first = cer_dist_first = 0 + wer_dist_best = cer_dist_best = 0 + words_count = 0 + chars_count = 0 + sample_idx = 0 + if preds_output_file: + out_file = open(preds_output_file, 'w') + + if progress_bar: + it = tqdm( + range(int(np.ceil(len(all_probs) / beam_batch_size))), + desc=f"Beam search decoding with width={beam_width}, alpha={beam_alpha}, beta={beam_beta}", + ncols=120, + ) + else: + it = range(int(np.ceil(len(all_probs) / beam_batch_size))) + for batch_idx in it: + # disabling type checking + with nemo.core.typecheck.disable_checks(): + probs_batch = all_probs[batch_idx * beam_batch_size : (batch_idx + 1) * beam_batch_size] + beams_batch = beam_search_lm.forward(log_probs=probs_batch, log_probs_length=None,) + + for beams_idx, beams in enumerate(beams_batch): + target = target_transcripts[sample_idx + beams_idx] + target_split_w = target.split() + target_split_c = list(target) + words_count += len(target_split_w) + chars_count += len(target_split_c) + wer_dist_min = cer_dist_min = 10000 + for candidate_idx, candidate in enumerate(beams): + if ids_to_text_func is not None: + # For BPE encodings, need to shift by TOKEN_OFFSET to retrieve the original sub-word ids + pred_text = ids_to_text_func([ord(c) - TOKEN_OFFSET for c in candidate[1]]) + else: + pred_text = candidate[1] + pred_split_w = pred_text.split() + wer_dist = editdistance.eval(target_split_w, pred_split_w) + pred_split_c = list(pred_text) + cer_dist = editdistance.eval(target_split_c, pred_split_c) + + wer_dist_min = min(wer_dist_min, wer_dist) + cer_dist_min = min(cer_dist_min, cer_dist) + + if candidate_idx == 0: + # first candidate + wer_dist_first += wer_dist + cer_dist_first += cer_dist + + score = candidate[0] + if preds_output_file: + out_file.write('{}\t{}\n'.format(pred_text, score)) + wer_dist_best += wer_dist_min + cer_dist_best += cer_dist_min + sample_idx += len(probs_batch) + + if preds_output_file: + out_file.close() + logging.info(f"Stored the predictions of beam search decoding at '{preds_output_file}'.") + + if lm_path: + logging.info( + 'WER/CER with beam search decoding and N-gram model = {:.2%}/{:.2%}'.format( + wer_dist_first / words_count, cer_dist_first / chars_count + ) + ) + else: + logging.info( + 'WER/CER with beam search decoding = {:.2%}/{:.2%}'.format( + wer_dist_first / words_count, cer_dist_first / chars_count + ) + ) + logging.info( + 'Best WER/CER in candidates = {:.2%}/{:.2%}'.format(wer_dist_best / words_count, cer_dist_best / chars_count) + ) + logging.info(f"=================================================================================") + + +def main(): + parser = argparse.ArgumentParser( + description='Evaluate an ASR model with beam search decoding and n-gram KenLM language model.' + ) + parser.add_argument( + "--nemo_model_file", required=True, type=str, help="The path of the '.nemo' file of the ASR model" + ) + parser.add_argument( + "--kenlm_model_file", required=False, default=None, type=str, help="The path of the KenLM binary model file" + ) + parser.add_argument("--input_manifest", required=True, type=str, help="The manifest file of the evaluation set") + parser.add_argument( + "--preds_output_folder", default=None, type=str, help="The optional folder where the predictions are stored" + ) + parser.add_argument( + "--probs_cache_file", default=None, type=str, help="The cache file for storing the outputs of the model" + ) + parser.add_argument( + "--acoustic_batch_size", default=16, type=int, help="The batch size to calculate log probabilities" + ) + parser.add_argument( + "--device", default="cuda", type=str, help="The device to load the model onto to calculate log probabilities" + ) + parser.add_argument( + "--use_amp", action="store_true", help="Whether to use AMP if available to calculate log probabilities" + ) + parser.add_argument( + "--decoding_mode", + choices=["greedy", "beamsearch", "beamsearch_ngram"], + default="beamsearch_ngram", + type=str, + help="The decoding scheme to be used for evaluation.", + ) + parser.add_argument( + "--beam_width", + required=True, + type=int, + nargs="+", + help="The width or list of the widths for the beam search decoding", + ) + parser.add_argument( + "--beam_alpha", + required=True, + type=float, + nargs="+", + help="The alpha parameter or list of the alphas for the beam search decoding", + ) + parser.add_argument( + "--beam_beta", + required=True, + type=float, + nargs="+", + help="The beta parameter or list of the betas for the beam search decoding", + ) + parser.add_argument( + "--beam_batch_size", default=128, type=int, help="The batch size to be used for beam search decoding" + ) + args = parser.parse_args() + + asr_model = nemo_asr.models.EncDecCTCModelBPE.restore_from( + args.nemo_model_file, map_location=torch.device(args.device) + ) + + target_transcripts = [] + with open(args.input_manifest, 'r') as manifest_file: + audio_file_paths = [] + for line in tqdm(manifest_file, desc=f"Reading Manifest {args.input_manifest} ...", ncols=120): + data = json.loads(line) + target_transcripts.append(data['text']) + audio_file_paths.append(data['audio_filepath']) + + if args.probs_cache_file and os.path.exists(args.probs_cache_file): + logging.info(f"Found a pickle file of probabilities at '{args.probs_cache_file}'.") + logging.info(f"Loading the cached pickle file of probabilities from '{args.probs_cache_file}' ...") + with open(args.probs_cache_file, 'rb') as probs_file: + all_probs = pickle.load(probs_file) + + if len(all_probs) != len(audio_file_paths): + raise ValueError( + f"The number of samples in the probabilities file '{args.probs_cache_file}' does not " + f"match the manifest file. You may need to delete the probabilities cached file." + ) + else: + if args.use_amp: + if torch.cuda.is_available() and hasattr(torch.cuda, 'amp') and hasattr(torch.cuda.amp, 'autocast'): + logging.info("AMP is enabled!\n") + autocast = torch.cuda.amp.autocast + else: + + @contextlib.contextmanager + def autocast(): + yield + + with autocast(): + with torch.no_grad(): + all_logits = asr_model.transcribe(audio_file_paths, batch_size=args.acoustic_batch_size, logprobs=True) + all_probs = [kenlm_utils.softmax(logits) for logits in all_logits] + if args.probs_cache_file: + logging.info(f"Writing pickle files of probabilities at '{args.probs_cache_file}'...") + with open(args.probs_cache_file, 'wb') as f_dump: + pickle.dump(all_probs, f_dump) + + wer_dist_greedy = 0 + cer_dist_greedy = 0 + words_count = 0 + chars_count = 0 + for batch_idx, probs in enumerate(all_probs): + preds = np.argmax(probs, axis=1) + preds_tensor = torch.tensor(preds, device='cpu').unsqueeze(0) + pred_text = asr_model._wer.ctc_decoder_predictions_tensor(preds_tensor)[0] + + pred_split_w = pred_text.split() + target_split_w = target_transcripts[batch_idx].split() + pred_split_c = list(pred_text) + target_split_c = list(target_transcripts[batch_idx]) + + wer_dist = editdistance.eval(target_split_w, pred_split_w) + cer_dist = editdistance.eval(target_split_c, pred_split_c) + + wer_dist_greedy += wer_dist + cer_dist_greedy += cer_dist + words_count += len(target_split_w) + chars_count += len(target_split_c) + + logging.info('Greedy WER/CER = {:.2%}/{:.2%}'.format(wer_dist_greedy / words_count, cer_dist_greedy / chars_count)) + + encoding_level = kenlm_utils.SUPPORTED_MODELS.get(type(asr_model).__name__, None) + if not encoding_level: + logging.warning( + f"Model type '{type(asr_model).__name__}' may not be supported. Would try to train a char-level LM." + ) + encoding_level = 'char' + + vocab = asr_model.decoder.vocabulary + ids_to_text_func = None + if encoding_level == "subword": + vocab = [chr(idx + TOKEN_OFFSET) for idx in range(len(vocab))] + ids_to_text_func = asr_model.tokenizer.ids_to_text + # delete the model to free the memory + del asr_model + + if args.decoding_mode == "beamsearch_ngram": + if not os.path.exists(args.kenlm_model_file): + raise FileNotFoundError(f"Could not find the KenLM model file '{args.kenlm_model_file}'.") + lm_path = args.kenlm_model_file + else: + lm_path = None + + # 'greedy' decoding_mode would skip the beam search decoding + if args.decoding_mode in ["beamsearch_ngram", "beamsearch"]: + + params = {'beam_width': args.beam_width, 'beam_alpha': args.beam_alpha, 'beam_beta': args.beam_beta} + hp_grid = ParameterGrid(params) + hp_grid = list(hp_grid) + + logging.info(f"==============================Starting the beam search decoding===============================") + logging.info(f"Grid search size: {len(hp_grid)}") + logging.info(f"It may take some time...") + logging.info(f"==============================================================================================") + + if args.preds_output_folder and not os.path.exists(args.preds_output_folder): + os.mkdir(args.preds_output_folder) + for hp in hp_grid: + if args.preds_output_folder: + preds_output_file = os.path.join( + args.preds_output_folder, + f"preds_out_width{hp['beam_width']}_alpha{hp['beam_alpha']}_beta{hp['beam_beta']}.tsv", + ) + else: + preds_output_file = None + + beam_search_eval( + all_probs=all_probs, + target_transcripts=target_transcripts, + vocab=vocab, + ids_to_text_func=ids_to_text_func, + preds_output_file=preds_output_file, + lm_path=lm_path, + beam_width=hp["beam_width"], + beam_alpha=hp["beam_alpha"], + beam_beta=hp["beam_beta"], + beam_batch_size=args.beam_batch_size, + progress_bar=True, + ) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh new file mode 100644 index 0000000000..44f5fa82fa --- /dev/null +++ b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/install_beamsearch_decoders.sh @@ -0,0 +1,6 @@ +git clone https://github.com/NVIDIA/OpenSeq2Seq -b ctc-decoders +mv OpenSeq2Seq/decoders . +rm -rf OpenSeq2Seq +cd decoders +./setup.sh +cd .. diff --git a/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/kenlm_utils.py b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/kenlm_utils.py new file mode 100644 index 0000000000..ed967ced33 --- /dev/null +++ b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/kenlm_utils.py @@ -0,0 +1,112 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +# + +""" +Utility methods to be used for training N-gram LM with KenLM in train_kenlm.py +""" + +import json +import os + +import numpy as np +from joblib import Parallel, delayed +from tqdm.auto import tqdm + +# List of the supported models to be used with N-gram LM and beam search decoding +SUPPORTED_MODELS = { + 'EncDecCTCModelBPE': 'subword', + 'EncDecCTCModel': 'char', +} + + +def softmax(x): + e = np.exp(x - np.max(x)) + return e / e.sum(axis=-1).reshape([x.shape[0], 1]) + + +def read_train_file(path, lowercase: bool = False): + lines_read = 0 + text_dataset = [] + + with open(path, 'r') as f: + reader = tqdm(iter(lambda: f.readline(), ''), desc="Read 0 lines", unit=' lines') + for i, line in enumerate(reader): + if path.endswith('.json'): + line = json.loads(line)['text'] + + line = line.replace("\n", "").strip() + if lowercase: + line = line.lower() + + if line: + text_dataset.append(line) + + lines_read += 1 + if lines_read % 100000 == 0: + reader.set_description(f"Read {lines_read} lines") + + return text_dataset + + +def tokenize_str(texts, tokenizer, offset): + tokenized_text = [] + for text in texts: + tok_text = tokenizer.text_to_ids(text) + tok_text = [chr(token + offset) for token in tok_text] + tokenized_text.append(tok_text) + return tokenized_text + + +def tokenize_text(data, tokenizer, path, chunk_size=8192, buffer_size=32, token_offset=100): + dataset_len = len(data) + print( + f"Chunking {dataset_len} rows into {dataset_len / float(chunk_size):0.4f} tasks (each chunk contains {chunk_size} elements)" + ) + + current_step = 0 + if os.path.exists(path): + print(f"Deleting previous file : {path}") + os.remove(path) + + with Parallel(n_jobs=-2, verbose=10) as parallel: + while True: + start = current_step * chunk_size + end = min((current_step + buffer_size) * chunk_size, dataset_len) + + tokenized_data = parallel( + delayed(tokenize_str)(data[start : start + chunk_size], tokenizer, token_offset) + for start in range(start, end, chunk_size) + ) + + # Write dataset + write_dataset(tokenized_data, path) + current_step += len(tokenized_data) + print(f"Finished writing {len(tokenized_data)} chunks to {path}. Current chunk index = {current_step}") + del tokenized_data + if end >= dataset_len: + break + + +def write_dataset(chunks, path): + basedir = os.path.dirname(path) + + if not os.path.exists(basedir): + os.makedirs(basedir, exist_ok=True) + + with open(path, 'a+', encoding='utf-8') as f: + for chunk_idx in tqdm(range(len(chunks)), desc='Chunk ', total=len(chunks), unit=' chunks'): + for text in chunks[chunk_idx]: + line = ' '.join(text) + f.write(f"{line}\n") diff --git a/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/train_kenlm.py b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/train_kenlm.py new file mode 100644 index 0000000000..a28ecfe4b2 --- /dev/null +++ b/models/tts/NeMo/scripts/asr_language_modeling/ngram_lm/train_kenlm.py @@ -0,0 +1,154 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +# + + +# This script would train an N-gram language model with KenLM library (https://github.com/kpu/kenlm) which can be used +# with the beam search decoders on top of the ASR models. This script supports both character level and BPE level +# encodings and models which is detected automatically from the type of the model. +# After the N-gram model is trained, and stored in the binary format, you may use +# 'scripts/ngram_lm/eval_beamsearch_ngram.py' to evaluate it on an ASR model. +# +# You need to install the KenLM library and also the beam search decoders to use this feature. Please refer +# to 'scripts/ngram_lm/install_beamsearch_decoders.sh' on how to install them. +# +# USAGE: python train_kenlm.py --nemo_model_file \ +# --train_file \ +# --kenlm_model_file \ +# --ngram_length +# +# After training is done, the binary LM model is stored at the path specified by '--kenlm_model_file'. +# You may find more info on how to use this script at: +# https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/asr/asr_language_modeling.html + +import argparse +import logging +import os +import subprocess +import sys + +import kenlm_utils +import torch + +import nemo.collections.asr as nemo_asr +from nemo.utils import logging + +""" +NeMo's beam search decoders only support char-level encodings. In order to make it work with BPE-level encodings, we +use a trick to encode the sub-word tokens of the training data as unicode characters and train a char-level KenLM. +TOKEN_OFFSET is the offset in the unicode table to be used to encode the BPE sub-words. This encoding scheme reduces +the required memory significantly, and the LM and its binary blob format require less storage space. +""" +TOKEN_OFFSET = 100 + +CHUNK_SIZE = 8192 +CHUNK_BUFFER_SIZE = 512 + + +def main(): + parser = argparse.ArgumentParser( + description='Train an N-gram language model with KenLM to be used with beam search decoder of ASR models.' + ) + parser.add_argument( + "--train_file", + required=True, + type=str, + help="Path to the training file, it can be a text file or JSON manifest", + ) + parser.add_argument( + "--nemo_model_file", required=True, type=str, help="The path of the '.nemo' file of the ASR model" + ) + parser.add_argument( + "--kenlm_model_file", required=True, type=str, help="The path to store the KenLM binary model file" + ) + parser.add_argument("--ngram_length", required=True, type=int, help="The order of N-gram LM") + parser.add_argument("--kenlm_bin_path", required=True, type=str, help="The path to the bin folder of KenLM") + parser.add_argument( + "--do_lowercase", action='store_true', help="Whether to apply lower case conversion on the training text" + ) + args = parser.parse_args() + + """ TOKENIZER SETUP """ + logging.info(f"Loading nemo model '{args.nemo_model_file}' ...") + model = nemo_asr.models.ASRModel.restore_from(args.nemo_model_file, map_location=torch.device('cpu')) + + encoding_level = kenlm_utils.SUPPORTED_MODELS.get(type(model).__name__, None) + if not encoding_level: + logging.warning( + f"Model type '{type(model).__name__}' may not be supported. Would try to train a char-level LM." + ) + encoding_level = 'char' + + """ DATASET SETUP """ + logging.info(f"Encoding the train file '{args.train_file}' ...") + dataset = kenlm_utils.read_train_file(args.train_file, lowercase=args.do_lowercase) + encoded_train_file = f"{args.kenlm_model_file}.tmp.txt" + if encoding_level == "subword": + kenlm_utils.tokenize_text( + dataset, + model.tokenizer, + path=encoded_train_file, + chunk_size=CHUNK_SIZE, + buffer_size=CHUNK_BUFFER_SIZE, + token_offset=TOKEN_OFFSET, + ) + # --discount_fallback is needed for training KenLM for BPE-based models + discount_arg = "--discount_fallback" + else: + with open(encoded_train_file, 'w', encoding='utf-8') as f: + for line in dataset: + f.write(f"{line}\n") + + discount_arg = "" + + del model + + arpa_file = f"{args.kenlm_model_file}.tmp.arpa" + """ LMPLZ ARGUMENT SETUP """ + kenlm_args = [ + os.path.join(args.kenlm_bin_path, 'lmplz'), + "-o", + f"{args.ngram_length}", + "--text", + encoded_train_file, + "--arpa", + arpa_file, + discount_arg, + ] + + ret = subprocess.run(kenlm_args, capture_output=False, text=True, stdout=sys.stdout, stderr=sys.stderr) + if ret.returncode != 0: + raise RuntimeError("Training KenLM was not successful!") + """ BINARY BUILD """ + logging.info(f"Running binary_build command \n\n{' '.join(kenlm_args)}\n\n") + kenlm_args = [ + os.path.join(args.kenlm_bin_path, "build_binary"), + "trie", + arpa_file, + args.kenlm_model_file, + ] + ret = subprocess.run(kenlm_args, capture_output=False, text=True, stdout=sys.stdout, stderr=sys.stderr) + + if ret.returncode != 0: + raise RuntimeError("Training KenLM was not successful!") + + os.remove(encoded_train_file) + logging.info(f"Deleted the temporary encoded training file '{encoded_train_file}'.") + os.remove(arpa_file) + logging.info(f"Deleted the arpa file '{arpa_file}'.") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/average_model_checkpoints.py b/models/tts/NeMo/scripts/average_model_checkpoints.py new file mode 100644 index 0000000000..06c522f1e1 --- /dev/null +++ b/models/tts/NeMo/scripts/average_model_checkpoints.py @@ -0,0 +1,171 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# Copyright 2017 Johns Hopkins University (Shinji Watanabe) +# +# 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. + +""" +# Changes to script +Change the script to import the NeMo model class you would like to load a checkpoint for, +then update the model constructor to use this model class. This can be found by the line: + +<<< Change model class here ! >>> + +By default, this script imports and creates the `EncDecCTCModelBPE` class but it can be +changed to any NeMo Model. + +# Run the script + +## Saving a .nemo model file (loaded with ModelPT.restore_from(...)) +HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \ + --config-path="" \ + --config-name="" \ + name= \ + +checkpoint_dir= \ + +checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\" + + +## Saving an averaged pytorch checkpoint (loaded with torch.load(...)) +HYDRA_FULL_ERROR=1 python average_model_checkpoints.py \ + --config-path="" \ + --config-name="" \ + name= \ + +checkpoint_dir= \ + +checkpoint_paths=\"[/path/to/ptl_1.ckpt,/path/to/ptl_2.ckpt,/path/to/ptl_3.ckpt,...]\" \ + +save_ckpt_only=true + +""" + +import os + +import pytorch_lightning as pl +import torch +from omegaconf import OmegaConf, open_dict + +# Change this import to the model you would like to average +from nemo.collections.asr.models import EncDecCTCModelBPE +from nemo.core.config import hydra_runner +from nemo.utils import logging + + +def process_config(cfg: OmegaConf): + if 'name' not in cfg or cfg.name is None: + raise ValueError("`cfg.name` must be provided to save a model checkpoint") + + if 'checkpoint_paths' not in cfg or cfg.checkpoint_paths is None: + raise ValueError( + "`cfg.checkpoint_paths` must be provided as a list of one or more str paths to " + "pytorch lightning checkpoints" + ) + + save_ckpt_only = False + + with open_dict(cfg): + name_prefix = cfg.name + checkpoint_paths = cfg.pop('checkpoint_paths') + + if 'checkpoint_dir' in cfg: + checkpoint_dir = cfg.pop('checkpoint_dir') + else: + checkpoint_dir = None + + if 'save_ckpt_only' in cfg: + save_ckpt_only = cfg.pop('save_ckpt_only') + + if type(checkpoint_paths) not in (list, tuple): + checkpoint_paths = str(checkpoint_paths).replace("[", "").replace("]", "") + checkpoint_paths = checkpoint_paths.split(",") + checkpoint_paths = [ckpt_path.strip() for ckpt_path in checkpoint_paths] + + if checkpoint_dir is not None: + checkpoint_paths = [os.path.join(checkpoint_dir, path) for path in checkpoint_paths] + + return name_prefix, checkpoint_paths, save_ckpt_only + + +@hydra_runner(config_path=None, config_name=None) +def main(cfg): + name_prefix, checkpoint_paths, save_ckpt_only = process_config(cfg) + + if not save_ckpt_only: + trainer = pl.Trainer(**cfg.trainer) + + # <<< Change model class here ! >>> + # Model architecture which will contain the averaged checkpoints + # Change the model constructor to the one you would like (if needed) + model = EncDecCTCModelBPE(cfg=cfg.model, trainer=trainer) + + """ < Checkpoint Averaging Logic > """ + # load state dicts + n = len(checkpoint_paths) + avg_state = None + + logging.info(f"Averaging {n} checkpoints ...") + + for ix, path in enumerate(checkpoint_paths): + checkpoint = torch.load(path, map_location='cpu') + + if 'state_dict' in checkpoint: + checkpoint = checkpoint['state_dict'] + + if ix == 0: + # Initial state + avg_state = checkpoint + + logging.info(f"Initialized average state dict with checkpoint : {path}") + else: + # Accumulated state + for k in avg_state: + avg_state[k] = avg_state[k] + checkpoint[k] + + logging.info(f"Updated average state dict with state from checkpoint : {path}") + + for k in avg_state: + if str(avg_state[k].dtype).startswith("torch.int"): + # For int type, not averaged, but only accumulated. + # e.g. BatchNorm.num_batches_tracked + pass + else: + avg_state[k] = avg_state[k] / n + + # Save model + if save_ckpt_only: + ckpt_name = name_prefix + '-averaged.ckpt' + torch.save(avg_state, ckpt_name) + + logging.info(f"Averaged pytorch checkpoint saved as : {ckpt_name}") + else: + # Set model state + logging.info("Loading averaged state dict in provided model") + model.load_state_dict(avg_state, strict=True) + + ckpt_name = name_prefix + '-averaged.nemo' + model.save_to(ckpt_name) + + logging.info(f"Averaged model saved as : {ckpt_name}") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/get_commonvoice_data.py b/models/tts/NeMo/scripts/dataset_processing/get_commonvoice_data.py new file mode 100644 index 0000000000..b53415603b --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/get_commonvoice_data.py @@ -0,0 +1,178 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +# + +# Copyright (c) 2020, SeanNaren. 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. +# + +# To convert mp3 files to wav using sox, you must have installed sox with mp3 support +# For example sudo apt-get install libsox-fmt-mp3 +import argparse +import csv +import json +import logging +import multiprocessing +import os +import subprocess +import sys +import tarfile +from multiprocessing.pool import ThreadPool +from pathlib import Path +from typing import List + +import sox +from sox import Transformer +from tqdm import tqdm + +parser = argparse.ArgumentParser(description='Downloads and processes Mozilla Common Voice dataset.') +parser.add_argument("--data_root", default='CommonVoice_dataset/', type=str, help="Directory to store the dataset.") +parser.add_argument('--manifest_dir', default='./', type=str, help='Output directory for manifests') +parser.add_argument("--num_workers", default=multiprocessing.cpu_count(), type=int, help="Workers to process dataset.") +parser.add_argument('--sample_rate', default=16000, type=int, help='Sample rate') +parser.add_argument('--n_channels', default=1, type=int, help='Number of channels for output wav files') +parser.add_argument( + '--files_to_process', + nargs='+', + default=['test.tsv', 'dev.tsv', 'train.tsv'], + type=str, + help='list of *.csv file names to process', +) +parser.add_argument( + '--version', + default='cv-corpus-5.1-2020-06-22', + type=str, + help='Version of the dataset (obtainable via https://commonvoice.mozilla.org/en/datasets', +) +parser.add_argument( + '--language', + default='en', + type=str, + help='Which language to download.(default english,' + 'check https://commonvoice.mozilla.org/en/datasets for more language codes', +) +args = parser.parse_args() +COMMON_VOICE_URL = ( + f"https://voice-prod-bundler-ee1969a6ce8178826482b88e843c335139bd3fb4.s3.amazonaws.com/" + "{}/{}.tar.gz".format(args.version, args.language) +) + + +def create_manifest(data: List[tuple], output_name: str, manifest_path: str): + output_file = Path(manifest_path) / output_name + output_file.parent.mkdir(exist_ok=True, parents=True) + + with output_file.open(mode='w') as f: + for wav_path, duration, text in tqdm(data, total=len(data)): + f.write( + json.dumps({'audio_filepath': os.path.abspath(wav_path), "duration": duration, 'text': text}) + '\n' + ) + + +def process_files(csv_file, data_root, num_workers): + """ Read *.csv file description, convert mp3 to wav, process text. + Save results to data_root. + + Args: + csv_file: str, path to *.csv file with data description, usually start from 'cv-' + data_root: str, path to dir to save results; wav/ dir will be created + """ + wav_dir = os.path.join(data_root, 'wav/') + os.makedirs(wav_dir, exist_ok=True) + audio_clips_path = os.path.dirname(csv_file) + '/clips/' + + def process(x): + file_path, text = x + file_name = os.path.splitext(os.path.basename(file_path))[0] + text = text.lower().strip() + audio_path = os.path.join(audio_clips_path, file_path) + output_wav_path = os.path.join(wav_dir, file_name + '.wav') + + tfm = Transformer() + tfm.rate(samplerate=args.sample_rate) + tfm.channels(n_channels=args.n_channels) + tfm.build(input_filepath=audio_path, output_filepath=output_wav_path) + duration = sox.file_info.duration(output_wav_path) + return output_wav_path, duration, text + + logging.info('Converting mp3 to wav for {}.'.format(csv_file)) + with open(csv_file) as csvfile: + reader = csv.DictReader(csvfile, delimiter='\t') + next(reader, None) # skip the headers + data = [(row['path'], row['sentence']) for row in reader] + with ThreadPool(num_workers) as pool: + data = list(tqdm(pool.imap(process, data), total=len(data))) + return data + + +def main(): + data_root = args.data_root + os.makedirs(data_root, exist_ok=True) + + target_unpacked_dir = os.path.join(data_root, "CV_unpacked") + + if os.path.exists(target_unpacked_dir): + logging.info('Find existing folder {}'.format(target_unpacked_dir)) + else: + logging.info("Could not find Common Voice, Downloading corpus...") + + commands = [ + 'wget', + '--user-agent', + '"Mozilla/5.0 (Windows NT 10.0; WOW64) ' + 'AppleWebKit/537.36 (KHTML, like Gecko) Chrome/51.0.2704.103 Safari/537.36"', + '-P', + data_root, + f'{COMMON_VOICE_URL}', + ] + commands = " ".join(commands) + subprocess.run(commands, shell=True, stderr=sys.stderr, stdout=sys.stdout, capture_output=False) + filename = f"{args.language}.tar.gz" + target_file = os.path.join(data_root, os.path.basename(filename)) + + os.makedirs(target_unpacked_dir, exist_ok=True) + logging.info("Unpacking corpus to {} ...".format(target_unpacked_dir)) + tar = tarfile.open(target_file) + tar.extractall(target_unpacked_dir) + tar.close() + + folder_path = os.path.join(target_unpacked_dir, args.version + f'/{args.language}/') + + for csv_file in args.files_to_process: + data = process_files( + csv_file=os.path.join(folder_path, csv_file), + data_root=os.path.join(data_root, os.path.splitext(csv_file)[0]), + num_workers=args.num_workers, + ) + logging.info('Creating manifests...') + create_manifest( + data=data, + output_name=f'commonvoice_{os.path.splitext(csv_file)[0]}_manifest.json', + manifest_path=args.manifest_dir, + ) + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/get_hi-mia_data.py b/models/tts/NeMo/scripts/dataset_processing/get_hi-mia_data.py new file mode 100644 index 0000000000..59dd629c04 --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/get_hi-mia_data.py @@ -0,0 +1,185 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# USAGE: python get_aishell_data.py --data_root= + +import argparse +import json +import logging as _logging +import os +import tarfile +import urllib.request +from glob import glob + +import librosa as l +from sklearn.model_selection import StratifiedShuffleSplit +from tqdm import tqdm + +parser = argparse.ArgumentParser(description='HI-MIA Data download') +parser.add_argument("--data_root", required=True, default=None, type=str) +parser.add_argument("--log_level", default=20, type=int) +args = parser.parse_args() +logging = _logging.getLogger(__name__) +logging.addHandler(_logging.StreamHandler()) +logging.setLevel(args.log_level) + +URL = { + 'dev': "http://www.openslr.org/resources/85/dev.tar.gz", + 'test': "http://www.openslr.org/resources/85/test.tar.gz", + 'train': "http://www.openslr.org/resources/85/train.tar.gz", +} + + +def __maybe_download_file(destination: str, source: str): + """ + Downloads source to destination if it doesn't exist. + If exists, skips download + Args: + destination: local filepath + source: url of resource + + Returns: + + """ + source = URL[source] + if not os.path.exists(destination) and not os.path.exists(os.path.splitext(destination)[0]): + logging.info("{0} does not exist. Downloading ...".format(destination)) + urllib.request.urlretrieve(source, filename=destination + '.tmp') + os.rename(destination + '.tmp', destination) + logging.info("Downloaded {0}.".format(destination)) + elif os.path.exists(destination): + logging.info("Destination {0} exists. Skipping.".format(destination)) + elif os.path.exists(os.path.splitext(destination)[0]): + logging.warning( + "Assuming extracted folder %s contains the extracted files from %s. Will not download.", + os.path.basename(destination), + destination, + ) + return destination + + +def __extract_all_files(filepath: str, data_root: str, data_dir: str): + if not os.path.exists(data_dir): + extract_file(filepath, data_root) + audio_dir = os.path.join(data_dir, 'wav') + for subfolder, _, filelist in os.walk(audio_dir): + for ftar in filelist: + extract_file(os.path.join(subfolder, ftar), subfolder) + else: + logging.info('Skipping extracting. Data already there %s' % data_dir) + + +def extract_file(filepath: str, data_dir: str): + try: + tar = tarfile.open(filepath) + tar.extractall(data_dir) + tar.close() + except Exception: + logging.info('Not extracting. Maybe already there?') + + +def __remove_tarred_files(filepath: str, data_dir: str): + if os.path.exists(data_dir) and os.path.isfile(filepath): + logging.info('Deleting %s' % filepath) + os.remove(filepath) + + +def write_file(name, lines, idx): + with open(name, 'w') as fout: + for i in idx: + dic = lines[i] + json.dump(dic, fout) + fout.write('\n') + logging.info("wrote %s", name) + + +def __process_data(data_folder: str, data_set: str): + """ + To generate manifest + Args: + data_folder: source with wav files + Returns: + + """ + fullpath = os.path.abspath(data_folder) + scp = glob(fullpath + '/**/*.wav', recursive=True) + out = os.path.join(fullpath, data_set + '_all.json') + utt2spk = os.path.join(fullpath, 'utt2spk') + utt2spk_file = open(utt2spk, 'w') + id = -2 # speaker id + + if os.path.exists(out): + logging.warning( + "%s already exists and is assumed to be processed. If not, please delete %s and rerun this script", + out, + out, + ) + return + + speakers = [] + lines = [] + with open(out, 'w') as outfile: + for line in tqdm(scp): + line = line.strip() + y, sr = l.load(line, sr=None) + if sr != 16000: + y, sr = l.load(line, sr=16000) + l.output.write_wav(line, y, sr) + dur = l.get_duration(y=y, sr=sr) + if data_set == 'test': + speaker = line.split('/')[-1].split('.')[0].split('_')[0] + else: + speaker = line.split('/')[id] + speaker = list(speaker) + speaker = ''.join(speaker) + speakers.append(speaker) + meta = {"audio_filepath": line, "duration": float(dur), "label": speaker} + lines.append(meta) + json.dump(meta, outfile) + outfile.write("\n") + utt2spk_file.write(line.split('/')[-1] + "\t" + speaker + "\n") + + utt2spk_file.close() + + if data_set != 'test': + sss = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42) + for train_idx, test_idx in sss.split(speakers, speakers): + print(len(train_idx)) + + out = os.path.join(fullpath, 'train.json') + write_file(out, lines, train_idx) + out = os.path.join(fullpath, 'dev.json') + write_file(out, lines, test_idx) + + +def main(): + data_root = args.data_root + for data_set in URL.keys(): + + # data_set = 'data_aishell' + logging.info("\n\nWorking on: {0}".format(data_set)) + file_path = os.path.join(data_root, data_set + ".tgz") + logging.info("Getting {0}".format(data_set)) + __maybe_download_file(file_path, data_set) + logging.info("Extracting {0}".format(data_set)) + data_folder = os.path.join(data_root, data_set) + __extract_all_files(file_path, data_root, data_folder) + __remove_tarred_files(file_path, data_folder) + logging.info("Processing {0}".format(data_set)) + __process_data(data_folder, data_set) + logging.info('Done!') + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/get_librispeech_data.py b/models/tts/NeMo/scripts/dataset_processing/get_librispeech_data.py new file mode 100644 index 0000000000..e009a2322f --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/get_librispeech_data.py @@ -0,0 +1,171 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +# +# USAGE: python get_librispeech_data.py --data_root= +# --data_set= --num_workers= +# where can be: dev_clean, dev_other, test_clean, +# test_other, train_clean_100, train_clean_360, train_other_500 or ALL +# You can also put more than one data_set comma-separated: +# --data_set=dev_clean,train_clean_100 +import argparse +import fnmatch +import functools +import json +import logging +import multiprocessing +import os +import subprocess +import tarfile +import urllib.request + +from sox import Transformer +from tqdm import tqdm + +parser = argparse.ArgumentParser(description='LibriSpeech Data download') +parser.add_argument("--data_root", required=True, default=None, type=str) +parser.add_argument("--data_sets", default="dev_clean", type=str) +parser.add_argument("--num_workers", default=4, type=int) +args = parser.parse_args() + +URLS = { + 'TRAIN_CLEAN_100': ("http://www.openslr.org/resources/12/train-clean-100.tar.gz"), + 'TRAIN_CLEAN_360': ("http://www.openslr.org/resources/12/train-clean-360.tar.gz"), + 'TRAIN_OTHER_500': ("http://www.openslr.org/resources/12/train-other-500.tar.gz"), + 'DEV_CLEAN': "http://www.openslr.org/resources/12/dev-clean.tar.gz", + 'DEV_OTHER': "http://www.openslr.org/resources/12/dev-other.tar.gz", + 'TEST_CLEAN': "http://www.openslr.org/resources/12/test-clean.tar.gz", + 'TEST_OTHER': "http://www.openslr.org/resources/12/test-other.tar.gz", +} + + +def __maybe_download_file(destination: str, source: str): + """ + Downloads source to destination if it doesn't exist. + If exists, skips download + Args: + destination: local filepath + source: url of resource + Returns: + """ + source = URLS[source] + if not os.path.exists(destination): + logging.info("{0} does not exist. Downloading ...".format(destination)) + urllib.request.urlretrieve(source, filename=destination + '.tmp') + os.rename(destination + '.tmp', destination) + logging.info("Downloaded {0}.".format(destination)) + else: + logging.info("Destination {0} exists. Skipping.".format(destination)) + return destination + + +def __extract_file(filepath: str, data_dir: str): + try: + tar = tarfile.open(filepath) + tar.extractall(data_dir) + tar.close() + except Exception: + logging.info('Not extracting. Maybe already there?') + + +def __process_transcript(file_path: str, dst_folder: str): + """ + Converts flac files to wav from a given transcript, capturing the metadata. + Args: + file_path: path to a source transcript with flac sources + dst_folder: path where wav files will be stored + Returns: + a list of metadata entries for processed files. + """ + entries = [] + root = os.path.dirname(file_path) + with open(file_path, encoding="utf-8") as fin: + for line in fin: + id, text = line[: line.index(" ")], line[line.index(" ") + 1 :] + transcript_text = text.lower().strip() + + # Convert FLAC file to WAV + flac_file = os.path.join(root, id + ".flac") + wav_file = os.path.join(dst_folder, id + ".wav") + if not os.path.exists(wav_file): + Transformer().build(flac_file, wav_file) + # check duration + duration = subprocess.check_output("soxi -D {0}".format(wav_file), shell=True) + + entry = {} + entry['audio_filepath'] = os.path.abspath(wav_file) + entry['duration'] = float(duration) + entry['text'] = transcript_text + entries.append(entry) + return entries + + +def __process_data(data_folder: str, dst_folder: str, manifest_file: str, num_workers: int): + """ + Converts flac to wav and build manifests's json + Args: + data_folder: source with flac files + dst_folder: where wav files will be stored + manifest_file: where to store manifest + num_workers: number of parallel workers processing files + Returns: + """ + + if not os.path.exists(dst_folder): + os.makedirs(dst_folder) + + files = [] + entries = [] + + for root, dirnames, filenames in os.walk(data_folder): + for filename in fnmatch.filter(filenames, '*.trans.txt'): + files.append(os.path.join(root, filename)) + + with multiprocessing.Pool(num_workers) as p: + processing_func = functools.partial(__process_transcript, dst_folder=dst_folder) + results = p.imap(processing_func, files) + for result in tqdm(results, total=len(files)): + entries.extend(result) + + with open(manifest_file, 'w') as fout: + for m in entries: + fout.write(json.dumps(m) + '\n') + + +def main(): + data_root = args.data_root + data_sets = args.data_sets + num_workers = args.num_workers + + if data_sets == "ALL": + data_sets = "dev_clean,dev_other,train_clean_100,train_clean_360,train_other_500,test_clean,test_other" + + for data_set in data_sets.split(','): + logging.info("\n\nWorking on: {0}".format(data_set)) + filepath = os.path.join(data_root, data_set + ".tar.gz") + logging.info("Getting {0}".format(data_set)) + __maybe_download_file(filepath, data_set.upper()) + logging.info("Extracting {0}".format(data_set)) + __extract_file(filepath, data_root) + logging.info("Processing {0}".format(data_set)) + __process_data( + os.path.join(os.path.join(data_root, "LibriSpeech"), data_set.replace("_", "-"),), + os.path.join(os.path.join(data_root, "LibriSpeech"), data_set.replace("_", "-"),) + "-processed", + os.path.join(data_root, data_set + ".json"), + num_workers=num_workers, + ) + logging.info('Done!') + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/get_openslr_rir_data.py b/models/tts/NeMo/scripts/dataset_processing/get_openslr_rir_data.py new file mode 100644 index 0000000000..f1663061ba --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/get_openslr_rir_data.py @@ -0,0 +1,140 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# USAGE: python get_openslr_rir_data.py --data_root= +# Data is downloaded from OpenSLR's "Room Impulse Response and Noise Database" +# RIRs in multichannel files are separated into single channel files and +# a json file that can be used as in input to NeMo is created + +import argparse +import glob +import json +import logging +import os +import subprocess +import urllib.request +from shutil import copy, move +from zipfile import ZipFile + +parser = argparse.ArgumentParser(description='OpenSLR RIR Data download and process') +parser.add_argument("--data_root", required=True, default=None, type=str) +args = parser.parse_args() + +URLS = { + 'SLR28': ("http://www.openslr.org/resources/28/rirs_noises.zip"), +} + + +def __maybe_download_file(destination: str, source: str): + """ + Downloads source to destination if it doesn't exist. + If exists, skips download + Args: + destination: local filepath + source: url of resource + Returns: + """ + source = URLS[source] + if not os.path.exists(destination): + logging.info("{0} does not exist. Downloading ...".format(destination)) + urllib.request.urlretrieve(source, filename=destination + '.tmp') + os.rename(destination + '.tmp', destination) + logging.info("Downloaded {0}.".format(destination)) + else: + logging.info("Destination {0} exists. Skipping.".format(destination)) + return destination + + +def __extract_file(filepath: str, data_dir: str): + try: + with ZipFile(filepath, 'r') as zipObj: + zipObj.extractall(data_dir) + except Exception: + logging.info('Not extracting. Maybe already there?') + + +def __process_data(data_folder: str, dst_folder: str, manifest_file: str): + """ + Converts flac to wav and build manifests's json + Args: + data_folder: source with flac files + dst_folder: where wav files will be stored + manifest_file: where to store manifest + Returns: + """ + if not os.path.exists(dst_folder): + os.makedirs(dst_folder) + + real_rir_list = os.path.join(data_folder, "RIRS_NOISES", "real_rirs_isotropic_noises", "rir_list") + rirfiles = [] + with open(real_rir_list, "r") as rir_f: + for line in rir_f: + rirfiles.append(os.path.join(data_folder, line.rstrip().split(" ")[4])) + + real_rir_folder = os.path.join(dst_folder, "real_rirs") + if not os.path.exists(real_rir_folder): + os.makedirs(real_rir_folder) + # split multi-channel rir files to single channel + for rir_f in rirfiles: + n_chans = int(subprocess.check_output("soxi -c {0}".format(rir_f), shell=True)) + if n_chans == 1: + copy(rir_f, real_rir_folder) + else: + for chan in range(1, n_chans + 1): + chan_file_name = os.path.join( + real_rir_folder, os.path.splitext(os.path.basename(rir_f))[0] + "-" + str(chan) + ".wav" + ) + _ = subprocess.check_output(f"sox {rir_f} {chan_file_name} remix {chan}", shell=True) + + # move simulated rirs to processed + if not os.path.exists(os.path.join(dst_folder, "simulated_rirs")): + move(os.path.join(data_folder, "RIRS_NOISES", "simulated_rirs"), dst_folder) + + os.chdir(dst_folder) + all_rirs = glob.glob("**/*.wav", recursive=True) + with open(manifest_file, 'w') as man_f: + entry = {} + for rir in all_rirs: + rir_file = os.path.join(dst_folder, rir) + duration = subprocess.check_output("soxi -D {0}".format(rir_file), shell=True) + entry['audio_filepath'] = rir_file + entry['duration'] = float(duration) + entry['offset'] = 0 + entry['text'] = '_' + man_f.write(json.dumps(entry) + "\n") + + print('Done!') + + +def main(): + data_root = os.path.abspath(args.data_root) + data_set = "slr28" + logging.getLogger().setLevel(logging.INFO) + logging.info("\n\nWorking on: {0}".format(data_set)) + filepath = os.path.join(data_root, data_set + ".zip") + logging.info("Getting {0}".format(data_set)) + __maybe_download_file(filepath, data_set.upper()) + logging.info("Extracting {0}".format(data_set)) + __extract_file(filepath, data_root) + logging.info("Processing {0}".format(data_set)) + __process_data( + data_root, + os.path.join(os.path.join(data_root, "processed")), + os.path.join(os.path.join(data_root, "processed", "rir.json")), + ) + logging.info('Done!') + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/ljspeech/README b/models/tts/NeMo/scripts/dataset_processing/ljspeech/README new file mode 100644 index 0000000000..d16fdac55a --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/ljspeech/README @@ -0,0 +1,14 @@ +These scripts are used for extracting features from LJSpeech. + +You will need to run the following scripts (note that these will call the other +scripts as necessary): + +- `create_manifests_and_textfiles.py`: Creates the manifest and text files that + are used by the Montreal Forced Aligner (MFA) library. +- `extract_ljspeech_phonemes_and_durs.sh`: Extracts phonemes and alignments via + MFA and calculates phoneme durations. + Needs to be run with Conda active since it creates a Conda environment. + Also note that you may need to install OpenBlas manually (e.g. sudo apt-get + install libopenblas-dev), sometimes MFA can't find it if installed via Conda. + (This script calls `create_token2idx_dict.py` and `calculate_durs.py`.) +- `extract_ljspeech_energy_pitch.py` diff --git a/models/tts/NeMo/scripts/dataset_processing/ljspeech/calculate_durs.py b/models/tts/NeMo/scripts/dataset_processing/ljspeech/calculate_durs.py new file mode 100644 index 0000000000..f7a37c2a2e --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/ljspeech/calculate_durs.py @@ -0,0 +1,126 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. +""" +Calculates durations for LJSpeech based on MFA TextGrid alignments. +""" +import argparse +import glob +import json +import os +import pickle +from math import ceil + +import numpy as np +import tgt +import torch +from tqdm import tqdm + +parser = argparse.ArgumentParser(description="Calculates phoneme durations for LJSpeech from TextGrids.") +parser.add_argument('--ljspeech_dir', required=True, default=None, type=str) +parser.add_argument( + '--mappings', + required=False, + default=None, + type=str, + help='JSON file of mappings created with create_token2idx_dict.py', +) +parser.add_argument('--sr', required=False, default=22050, type=int) +parser.add_argument('--hop_length', required=False, default=256, type=int) +args = parser.parse_args() + + +def calculate_durations(textgrid, phone2idx): + tokens = [] + durs = [] + + frames_per_second = args.sr / args.hop_length + tg = tgt.read_textgrid(textgrid, include_empty_intervals=True) + data_tier = tg.get_tier_by_name("phones") + + # Get total frames + total_frames = ceil((data_tier.end_time - data_tier.start_time) * frames_per_second) + + # Find start and end frames of each token + se_in_frames = np.array([(frames_per_second * d.start_time, frames_per_second * d.end_time) for d in data_tier]) + se_in_frames = np.round(se_in_frames) + durs = (se_in_frames[:, 1] - se_in_frames[:, 0]).astype(int) + blank_set = ('sil', 'sp', 'spn', '', '') + blank_token = " " + + # merge repeated blank tokens + tokens, durations = [], [] + for i in range(len(data_tier)): + x = data_tier[i].text + if x == 'spn': + return None, None, None + x = blank_token if x in blank_set else x + + if len(tokens) and tokens[-1] == blank_token and x == blank_token: + durations[-1] += durs[i] + else: + tokens.append(x) + durations.append(durs[i]) + + tokens_enc = [phone2idx[token] for token in tokens] + tokens_enc, durations = torch.LongTensor(tokens_enc), torch.LongTensor(durations) + + # Add rounding error to final token + durations[-1] += total_frames - durations.sum() + + return tokens, tokens_enc, durations + + +def main(): + textgrid_list = glob.glob(os.path.join(args.ljspeech_dir, 'alignments/wavs/*.TextGrid')) + + # Create target_dir if necessary + target_dir = os.path.join(args.ljspeech_dir, 'phoneme_durations/') + print(f"Calculating phoneme durations, files will be in: {target_dir}") + + if not os.path.exists(target_dir): + print(f"Creating {target_dir}") + os.mkdir(target_dir) + + # Read phoneme to idx mappings + phone2idx = None + if args.mappings: + with open(args.mappings, 'r') as f: + mappings = json.load(f) + phone2idx = mappings['phone2idx'] + + oov_samples = [] + + # Iterate through all TextGrid files + for textgrid in tqdm(textgrid_list): + basename = os.path.splitext(os.path.basename(textgrid))[0][5:] # Chop off 'wavs_' prefix + + phones_mfa, tokens_mfa, durs = calculate_durations(textgrid, phone2idx) + + if phones_mfa is None: + oov_samples.append(basename) + continue + + # Save to file + target_path = os.path.join(target_dir, f'{basename}.pt') + torch.save({'text_encoded': tokens_mfa, 'token_duration': durs}, target_path) + + print(f"Getting rid of {len(oov_samples)} samples with OOV words.") + oov_target = os.path.join(args.ljspeech_dir, 'wavs_to_ignore.pkl') + with open(oov_target, 'wb') as f: + pickle.dump(oov_samples, f) + print(f"List of OOV samples written to: {oov_target}") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/ljspeech/create_manifests_and_textfiles.py b/models/tts/NeMo/scripts/dataset_processing/ljspeech/create_manifests_and_textfiles.py new file mode 100644 index 0000000000..aede9f828d --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/ljspeech/create_manifests_and_textfiles.py @@ -0,0 +1,71 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. +""" +Creates the manifest and .txt transcript files for an LJSpeech split. +The manifest will be used for training, and the .txt files are for the MFA library to find. +""" + +import argparse +import json +import os + +import sox +import wget + +from nemo.collections.common.parts.preprocessing import parsers + +parser = argparse.ArgumentParser() +parser.add_argument('--ljspeech_base', required=True, default=None, type=str) +args = parser.parse_args() + + +def main(): + filelist_base = 'https://raw.githubusercontent.com/NVIDIA/tacotron2/master/filelists/' + filelists = ['train', 'val', 'test'] + + # NeMo parser for text normalization + text_parser = parsers.make_parser(name='en') + + for split in filelists: + # Download file list if necessary + filelist_path = os.path.join(args.ljspeech_base, f"ljs_audio_text_{split}_filelist.txt") + if not os.path.exists(filelist_path): + wget.download(f"{filelist_base}/ljs_audio_text_{split}_filelist.txt", out=args.ljspeech_base) + + manifest_target = os.path.join(args.ljspeech_base, f"ljspeech_{split}.json") + with open(manifest_target, 'w') as f_out: + with open(filelist_path, 'r') as filelist: + print(f"\nCreating {manifest_target}...") + for line in filelist: + basename = line[6:16] + text = text_parser._normalize(line[21:].strip()) + + # Make sure corresponding wavfile exists and write .txt transcript + wav_path = os.path.join(args.ljspeech_base, 'wavs/', basename + '.wav') + assert os.path.exists(wav_path) + txt_path = os.path.join(args.ljspeech_base, 'wavs/', basename + '.txt') + with open(txt_path, 'w') as f_txt: + f_txt.write(text) + + # Write manifest entry + entry = { + 'audio_filepath': wav_path, + 'duration': sox.file_info.duration(wav_path), + 'text': text, + } + f_out.write(json.dumps(entry) + '\n') + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/ljspeech/create_token2idx_dict.py b/models/tts/NeMo/scripts/dataset_processing/ljspeech/create_token2idx_dict.py new file mode 100644 index 0000000000..2d3e73a38a --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/ljspeech/create_token2idx_dict.py @@ -0,0 +1,68 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. +""" +Creates a dictionary from token to index based on dictionary .txt given. +""" +import argparse +import json +import os + +parser = argparse.ArgumentParser() +parser.add_argument('--dictionary', required=True, default=None, type=str) +parser.add_argument('--dict_out', required=True, default=None, type=str) +args = parser.parse_args() + + +def main(): + if not os.path.exists(args.dictionary): + raise FileNotFoundError(f"Could not find dictionary file {args.dictionary}") + + phonemes = set() + word2phones = {} + with open(args.dictionary, 'r') as f: + for line in f: + line = line.split() + word = line[0] + tokens = line[1:] + + word2phones[word] = tokens + phonemes.update(tokens) + + # Small list of additional punctuation + word2phones[','] = [' '] + word2phones[';'] = [' '] + word2phones['.'] = [' '] + word2phones['!'] = [' '] + word2phones['?'] = [' '] + word2phones['"'] = [' '] + word2phones['-'] = [' '] + + phone2idx = {k: i for i, k in enumerate(phonemes)} + phone2idx[' '] = len(phone2idx) + phone2idx['sil'] = phone2idx[' '] # Silence + phone2idx['sp'] = phone2idx[' '] # Space + phone2idx['spn'] = phone2idx[' '] # OOV/unk + + dicts = { + 'phone2idx': phone2idx, + 'word2phones': word2phones, + } + with open(args.dict_out, 'w') as f: + json.dump(dicts, f) + + print(f"Total number of phone indices: {len(phone2idx)}") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_energy_pitch.py b/models/tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_energy_pitch.py new file mode 100644 index 0000000000..a76ce6dd9f --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_energy_pitch.py @@ -0,0 +1,88 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. +""" +Extracts energy and pitch data from LJSpeech wav files and writes them to a directory as LJxxx-xxxx.npy files. +Assuming that the wavs are located at `/wavs`, this script will write the +energy files to `/energies/` and the pitch files to `/pitches`, +creating the relevant directories if necessary. + +USAGE: python extract_ljspeech_energy_pitch.py --ljspeech_dir= +""" +import argparse +from pathlib import Path + +import librosa +import numpy as np + +tqdm = None + +try: + from tqdm import tqdm +except ModuleNotFoundError: + pass + +parser = argparse.ArgumentParser( + description="Extracts energies (L2-norm of STFT frame amplitudes) and pitches (F0) from LJSpeech data." +) +parser.add_argument("--ljspeech_dir", required=True, default=None, type=str) +args = parser.parse_args() + + +def main(): + wavfile_dir = Path(args.ljspeech_dir) / "wavs" + wavfile_list = list(wavfile_dir.glob('*.wav')) + + target_dir = Path(args.ljspeech_dir) + # Create target dir /energies and /pitches if necessary + if not Path(target_dir / "energies").exists(): + print(f"Creating target directory: {target_dir/'energies'}") + Path(target_dir / "energies").mkdir() + if not Path(target_dir / "pitches").exists(): + print(f"Creating target directory: {target_dir/'pitches'}") + Path(target_dir / "pitches").mkdir() + + if tqdm is not None: + wavfile_list = tqdm(wavfile_list) + for count, file_ in enumerate(wavfile_list): + basename = Path(file_).stem + pitch_path = target_dir / "pitches" / f"{basename}.npy" + energy_path = target_dir / "energies" / f"{basename}.npy" + if pitch_path.exists() and energy_path.exists(): + continue + audio, sr = librosa.load(file_, sr=22050) + + # Calculate f0 + # Please note that fmin and fmax are good approximates for the speaker in LJSpeech and may not generalize to + # other speakers + f0, _, _ = librosa.pyin(audio, fmin=80, fmax=800, frame_length=1024, sr=sr, fill_na=0.0) + + # Save to new file + np.save(pitch_path, f0) + + # Calculate energy + stft_amplitude = np.abs(librosa.stft(audio, n_fft=1024, hop_length=256, win_length=1024)) + energy = np.linalg.norm(stft_amplitude, axis=0) # axis=0 since librosa.stft -> (freq bins, frames) + + # Save to new file + np.save(energy_path, energy) + + assert energy.shape == f0.shape + if tqdm is None and count % 1000 == 0: + print(f"Finished processing {count} wav files...") + + print(f"Finished energy extraction for a total of {len(wavfile_list)} wav files.") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_phonemes_and_durs.sh b/models/tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_phonemes_and_durs.sh new file mode 100755 index 0000000000..d5ee2778a5 --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/ljspeech/extract_ljspeech_phonemes_and_durs.sh @@ -0,0 +1,150 @@ +#!/bin/bash + +# Extracts the phonemes and alignments for the LJSpeech dataset via the +# Montreal Forced Aligner (MFA) library, and computes durations per phoneme. +# Assumes you have downloaded and expanded the LJSpeech dataset, and that they +# are located at the directory specified. +# +# This script will create: +# - /mappings.json: Contains word->phone and phone->idx mappings +# - /wavs_to_ignore.pkl: A pickled list of wavs to ignore +# because of OOV words in their transcripts +# - /phoneme_durations/LJ*.npz: Numpy files for each wavfile +# with fields 'tokens' (token indices) and 'durs' (durations per token) +# +# It will also create intermediate files: +# - /alignments/wavs/LJ*.TextGrid: MFA output +# +# Note: This script will create a conda environment to set up MFA, so please +# make sure that conda is active before running this script. +# +# Note: You may need to install OpenBlas yourself if the MFA script can't find +# it. (e.g. sudo apt-get install libopenblas-dev) +# +# Example Usage: +# ./extract_ljspeech_phonemes_and_durs.sh /data/LJSpeech-1.1 + +ENV_NAME='aligner' + +SAMPLE_RATE=22050 +WINDOW_STRIDE=256 + +CMUDICT_URL='https://raw.githubusercontent.com/cmusphinx/cmudict/master/cmudict.dict' +SPLITS_BASE_URL='https://raw.githubusercontent.com/NVIDIA/tacotron2/master/filelists/' + +# Usage info +show_help() { +cat << EOF +Usage: $(basename "$0") [-h] \ + [--skip_env_setup] \ + [--g2p_dict=] \ + +Extracts phonemes and their respective durations for the LJSpeech dataset using +the Montreal Forced Aligner (MFA). + + -h Help message + --skip_env_setup (Optional) Skips setting up the MFA conda environment + "aligner". Use only if you already have this set up. + --g2p_dict (Optional) Path to the grapheme to phoneme dictionary + text file, if already generated. If set, skips the G2P + step. +EOF +} + +SKIP_ENV_SETUP=false +G2P_DICT='' + +while :; do + case $1 in + -h|-\?|--help) + show_help + exit + ;; + --g2p_dict=?*) + G2P_DICT=${1#*=} + ;; + --skip_env_setup) + SKIP_ENV_SETUP=true + ;; + *) + break + esac + shift +done + +if [[ -z $1 ]]; then + echo "Must specify a LJSpeech base directory." + exit 1 +fi +LJSPEECH_BASE=$1 + +# Check for conda +read -r -d '' CONDA_MESSAGE << EOM +This script requires either Anaconda or Miniconda to be active, as it installs the Montreal Forced Aligner via a conda environment. +See their documentation (https://montreal-forced-aligner.readthedocs.io/en/latest/installation.html) for more details. +EOM +if [ -z "$(which conda)" ]; then + echo $CONDA_MESSAGE + exit +fi + +CONDA_PREFIX=$(conda info --base) +source $CONDA_PREFIX/etc/profile.d/conda.sh + +# Set up env for MFA (env name "aligner") and install +if $SKIP_ENV_SETUP; then + echo "Skipping environment setup. Assuming env name "aligner" exists." +else + echo "Setting up conda environment for MFA (env name \"aligner\")..." + conda create -n $ENV_NAME -c conda-forge openblas python=3.8 openfst pynini ngram baumwelch + conda activate $ENV_NAME + pip install montreal-forced-aligner tgt torch + mfa thirdparty download + echo "Conda environment \"$ENV_NAME\" set up." +fi + +if ! conda activate $ENV_NAME; then + echo "Could not activate environment, see Conda output above." + exit +fi + +# Download CMU word-to-phoneme dict and clean out comments so they're not mistaken for tokens +if [ -z $G2P_DICT ]; then + if [ ! -f $LJSPEECH_BASE/cmudict.dict ]; then + echo "Downloading CMU dict." + wget -P $LJSPEECH_BASE $CMUDICT_URL + fi + if [ ! -f $LJSPEECH_BASE/uncommented_cmudict.dict ]; then + echo "Creating uncommented version of CMUdict." + sed 's/\ \#.*//' $LJSPEECH_BASE/cmudict.dict > $LJSPEECH_BASE/uncommented_cmudict.dict + fi + G2P_DICT=$LJSPEECH_BASE/uncommented_cmudict.dict +fi + +# Run alignment +echo "Starting MFA with dictionary at: $G2P_DICT" + +read -r -d '' MFA_ERROR_MSG << EOM +Could not run MFA. If it could not find OpenBlas, you may need to install it manually. +(e.g. sudo apt-get install libopenblas-dev) +EOM + +if ! mfa download acoustic english; then + echo $MFA_ERROR_MSG +fi +mfa align --clean $LJSPEECH_BASE $G2P_DICT english $LJSPEECH_BASE/alignments + +# Create JSON mappings from word to phonemes and phonemes to indices +echo "Creating word->phone and phone->idx mappings at $LJSPEECH_BASE/mappings.json..." +python create_token2idx_dict.py \ + --dictionary=$G2P_DICT \ + --dict_out=$LJSPEECH_BASE/mappings.json + +# Calculate phoneme durations +echo "Calculating phoneme durations..." +python calculate_durs.py \ + --ljspeech_dir=$LJSPEECH_BASE \ + --mappings=$LJSPEECH_BASE/mappings.json \ + --sr=$SAMPLE_RATE \ + --hop_length=$WINDOW_STRIDE +echo "Phoneme durations and tokens written." diff --git a/models/tts/NeMo/scripts/dataset_processing/process_an4_data.py b/models/tts/NeMo/scripts/dataset_processing/process_an4_data.py new file mode 100644 index 0000000000..5eaace2a7c --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/process_an4_data.py @@ -0,0 +1,88 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import glob +import json +import logging +import os +import subprocess +import tarfile + +import librosa +import wget + +parser = argparse.ArgumentParser(description="AN4 dataset download and processing") +parser.add_argument("--data_root", required=True, default=None, type=str) +args = parser.parse_args() + + +def build_manifest(data_root, transcripts_path, manifest_path, wav_path): + with open(transcripts_path, 'r') as fin: + with open(manifest_path, 'w') as fout: + for line in fin: + # Lines look like this: + # transcript (fileID) + transcript = line[: line.find('(') - 1].lower() + transcript = transcript.replace('', '').replace('', '') + transcript = transcript.strip() + + file_id = line[line.find('(') + 1 : -2] # e.g. "cen4-fash-b" + audio_path = os.path.join( + data_root, wav_path, file_id[file_id.find('-') + 1 : file_id.rfind('-')], file_id + '.wav', + ) + + duration = librosa.core.get_duration(filename=audio_path) + + # Write the metadata to the manifest + metadata = { + "audio_filepath": audio_path, + "duration": duration, + "text": transcript, + } + json.dump(metadata, fout) + fout.write('\n') + + +def main(): + data_root = os.path.abspath(args.data_root) + + # Convert from .sph to .wav + logging.info("Converting audio files to .wav...") + sph_list = glob.glob(os.path.join(data_root, 'an4/**/*.sph'), recursive=True) + for sph_path in sph_list: + wav_path = sph_path[:-4] + '.wav' + cmd = ['sox', sph_path, wav_path] + subprocess.run(cmd) + logging.info("Finished conversion.") + + # Build manifests + logging.info("Building training manifest...") + train_transcripts = os.path.join(data_root, 'an4/etc/an4_train.transcription') + train_manifest = os.path.join(data_root, 'an4/train_manifest.json') + train_wavs = os.path.join(data_root, 'an4/wav/an4_clstk') + build_manifest(data_root, train_transcripts, train_manifest, train_wavs) + logging.info("Training manifests created.") + + logging.info("Building test manifest...") + test_transcripts = os.path.join(data_root, 'an4/etc/an4_test.transcription') + test_manifest = os.path.join(data_root, 'an4/test_manifest.json') + test_wavs = os.path.join(data_root, 'an4/wav/an4test_clstk') + build_manifest(data_root, test_transcripts, test_manifest, test_wavs) + logging.info("Test manifest created.") + + logging.info("Done with AN4 processing!") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/process_degli_data.py b/models/tts/NeMo/scripts/dataset_processing/process_degli_data.py new file mode 100644 index 0000000000..ba1b26a3b4 --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/process_degli_data.py @@ -0,0 +1,70 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import sys + +from nemo.collections.tts.data.datalayers import preprocess_linear_specs_dataset + +sys.path.insert(0, '../') + + +def main(): + parser = argparse.ArgumentParser( + description='Create dataset fitted for training and validating deep griffin iteration from wavefiles' + ) + parser.add_argument( + "-v", + "--valid_filelist", + help="Filelist for validation set, with all validation audio files listed", + required=True, + default=None, + type=str, + ) + parser.add_argument( + "-t", + "--train_filelist", + help="Filelist for train set, with all train audio files listed", + required=True, + default=None, + type=str, + ) + parser.add_argument( + "-n", + "--n_fft", + help="Value for the n_fft parameter, and the filter length for the STFT", + default=512, + type=int, + ) + parser.add_argument("--hop_length", help="STFT parameter", default=256, type=int) + parser.add_argument( + "-d", "--destination", help="Destination to save the preprocessed data set to", default="/tmp", type=str + ) + parser.add_argument( + "-s", + "--num_snr", + help="Number of distinctive noisy samples to generate for each clear sample at the file list", + default=1, + type=int, + ) + + args = parser.parse_args() + + preprocess_linear_specs_dataset( + args.valid_filelist, args.train_filelist, args.n_fft, args.hop_length, args.num_snr, args.destination + ) + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/process_hub5_data.py b/models/tts/NeMo/scripts/dataset_processing/process_hub5_data.py new file mode 100644 index 0000000000..e67a0d04a0 --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/process_hub5_data.py @@ -0,0 +1,256 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +# This script is heavily derived from the Patter HUB5 processing script written +# by Ryan Leary + +import argparse +import glob +import json +import os +import re +import subprocess +import sys +from collections import namedtuple +from math import ceil, floor +from operator import attrgetter + +import numpy as np +import scipy.io.wavfile as wavfile +from tqdm import tqdm + +parser = argparse.ArgumentParser(description="Prepare HUB5 data for training/eval") +parser.add_argument( + "--data_root", default=None, type=str, required=True, help="The path to the root LDC HUB5 dataset directory.", +) +parser.add_argument( + "--dest_root", + default=None, + type=str, + required=True, + help="Path to the destination root directory for processed files.", +) + +# Optional arguments +parser.add_argument( + "--min_slice_duration", default=10.0, type=float, help="Minimum audio slice duration after processing.", +) + +args = parser.parse_args() + +StmUtterance = namedtuple( + 'StmUtterance', ['filename', 'channel', 'speaker_id', 'begin', 'end', 'label', 'transcript',], +) +STM_LINE_FMT = re.compile(r"^(\w+)\s+(\w+)\s+(\w+)\s+([0-9.]+)\s+([0-9.]+)\s+(<.*>)?\s+(.+)$") + +# Transcription errors and their fixes +TRANSCRIPT_BUGS = {"en_4622-B-12079-12187": "KIND OF WEIRD BUT"} + + +def get_utt_id(segment): + """ + Gives utterance IDs in a form like: en_4156-a-36558-37113 + """ + return "{}-{}-{}-{}".format(segment.filename, segment.channel, int(segment.begin * 100), int(segment.end * 100),) + + +def convert_utterances(sph_path, wav_path): + """ + Converts a sphere audio file to wav. + """ + cmd = ["sph2pipe", "-f", "wav", "-p", sph_path, wav_path] + subprocess.run(cmd) + + +def create_wavs(data_root, dest_root): + """ + Converts the English sph files to wav using sph2pipe. + """ + sph_root = os.path.join(data_root, "hub5e_00", "english") + sph_list = glob.glob(os.path.join(sph_root, "*.sph")) + + # Iterate over each sphere file and conver to wav + for sph_path in tqdm(sph_list, desc="Converting to wav", unit="file"): + sph_name, _ = os.path.splitext(os.path.basename(sph_path)) + wav_path = os.path.join(dest_root, 'full_audio_wav', sph_name + ".wav") + cmd = ["sph2pipe", "-f", "wav", "-p", sph_path, wav_path] + subprocess.run(cmd) + + +def process_transcripts(dataset_root): + """ + Reads in transcripts for each audio segment and processes them. + """ + stm_path = os.path.join(dataset_root, "2000_hub5_eng_eval_tr", "reference", "hub5e00.english.000405.stm",) + results = [] + chars = set() + + with open(stm_path, "r") as fh: + for line in fh: + # lines with ';;' are comments + if line.startswith(";;"): + continue + + if "IGNORE_TIME_SEGMENT_" in line: + continue + line = line.replace("", "").replace("", "") + line = line.replace("(%HESITATION)", "UH") + line = line.replace("-", "") + line = line.replace("(%UH)", "UH") + line = line.replace("(%AH)", "UH") + line = line.replace("(", "").replace(")", "") + + line = line.lower() + + m = STM_LINE_FMT.search(line.strip()) + utt = StmUtterance(*m.groups()) + + # Convert begin/end times to float + utt = utt._replace(begin=float(utt.begin)) + utt = utt._replace(end=float(utt.end)) + + # Check for utterance in dict of transcript mistakes + transcript_update = TRANSCRIPT_BUGS.get(get_utt_id(utt)) + if transcript_update is not None: + utt = utt._replace(transcript=transcript_update) + + results.append(utt) + chars.update(list(utt.transcript)) + return results, chars + + +def write_one_segment(dest_root, speaker_id, count, audio, sr, duration, transcript): + """ + Writes out one segment of audio, and writes its corresponding transcript + in the manifest. + + Args: + dest_root: the path to the output directory root + speaker_id: ID of the speaker, used in file naming + count: number of segments from this speaker so far + audio: the segment's audio data + sr: sample rate of the audio + duration: duration of the audio + transcript: the corresponding transcript + """ + audio_path = os.path.join(dest_root, "audio", f"{speaker_id}_{count:03}.wav") + + manifest_path = os.path.join(dest_root, "manifest_hub5.json") + + # Write audio + wavfile.write(audio_path, sr, audio) + + # Write transcript + transcript = { + "audio_filepath": audio_path, + "duration": duration, + "text": transcript, + } + with open(manifest_path, 'a') as f: + json.dump(transcript, f) + f.write('\n') + + +def segment_audio(info_list, dest_root, min_slice_duration): + """ + Combines audio into >= min_slice_duration segments of the same speaker, + and writes the combined transcripts into a manifest. + + Args: + info_list: list of StmUtterance objects with transcript information. + dest_root: path to output destination + min_slice_duration: min number of seconds per output audio slice + """ + info_list = sorted(info_list, key=attrgetter('speaker_id', 'begin')) + + prev_id = None # For checking audio concatenation + id_count = 0 + + sample_rate, audio_data = None, None + transcript_buffer = '' + audio_buffer = [] + buffer_duration = 0.0 + + # Iterate through utterances to build segments + for info in info_list: + if info.speaker_id != prev_id: + # Scrap the remainder in the buffers and start next segment + prev_id = info.speaker_id + id_count = 0 + + sample_rate, audio_data = wavfile.read(os.path.join(dest_root, 'full_audio_wav', info.filename + '.wav')) + transcript_buffer = '' + audio_buffer = [] + buffer_duration = 0.0 + + # Append utterance info to buffers + transcript_buffer += info.transcript + channel = 0 if info.channel.lower() == 'a' else 1 + audio_buffer.append( + audio_data[floor(info.begin * sample_rate) : ceil(info.end * sample_rate), channel,] + ) + buffer_duration += info.end - info.begin + + if buffer_duration < min_slice_duration: + transcript_buffer += ' ' + else: + # Write out segment and transcript + id_count += 1 + write_one_segment( + dest_root, + info.speaker_id, + id_count, + np.concatenate(audio_buffer, axis=0), + sample_rate, + buffer_duration, + transcript_buffer, + ) + + transcript_buffer = '' + audio_buffer = [] + buffer_duration = 0.0 + + +def main(): + data_root = args.data_root + dest_root = args.dest_root + + min_slice_duration = args.min_slice_duration + + if not os.path.exists(os.path.join(dest_root, 'full_audio_wav')): + os.makedirs(os.path.join(dest_root, 'full_audio_wav')) + if not os.path.exists(os.path.join(dest_root, 'audio')): + os.makedirs(os.path.join(dest_root, 'audio')) + + # Create/wipe manifest contents + open(os.path.join(dest_root, "manifest_hub5.json"), 'w').close() + + # Convert full audio files from .sph to .wav + create_wavs(data_root, dest_root) + + # Get each audio transcript from transcript file + info_list, chars = process_transcripts(data_root) + + print("Writing out vocab file", file=sys.stderr) + with open(os.path.join(dest_root, "vocab.txt"), 'w') as fh: + for x in sorted(list(chars)): + fh.write(x + "\n") + + # Segment the audio data + print("Segmenting audio and writing manifest") + segment_audio(info_list, dest_root, min_slice_duration) + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/process_speech_commands_data.py b/models/tts/NeMo/scripts/dataset_processing/process_speech_commands_data.py new file mode 100644 index 0000000000..9080b334fa --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/process_speech_commands_data.py @@ -0,0 +1,402 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +Usage: + +python process_speech_commands_data.py \ + --data_root= \ + --data_version= \ + --class_split= \ + --rebalance \ + --log +""" + +import argparse +import glob +import json +import logging +import os +import re +import tarfile +import urllib.request + +import librosa +import numpy as np + +URL_v1 = "http://download.tensorflow.org/data/speech_commands_v0.01.tar.gz" +URL_v2 = "http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz" + + +def __maybe_download_file(destination: str, source: str): + """ + Downloads source to destination if it doesn't exist. + If exists, skips download + Args: + destination: local filepath + source: url of resource + + Returns: + + """ + if not os.path.exists(destination): + logging.info(f"{destination} does not exist. Downloading ...") + urllib.request.urlretrieve(source, filename=destination + '.tmp') + os.rename(destination + '.tmp', destination) + logging.info(f"Downloaded {destination}.") + else: + logging.info(f"Destination {destination} exists. Skipping.") + return destination + + +def __extract_all_files(filepath: str, data_root: str, data_dir: str): + if not os.path.exists(data_dir): + extract_file(filepath, data_dir) + else: + logging.info(f'Skipping extracting. Data already there {data_dir}') + + +def extract_file(filepath: str, data_dir: str): + try: + tar = tarfile.open(filepath) + tar.extractall(data_dir) + tar.close() + except Exception: + logging.info('Not extracting. Maybe already there?') + + +def __process_data( + data_folder: str, dst_folder: str, rebalance: bool = False, class_split: str = "all", skip_duration: bool = False +): + """ + To generate manifest + + Args: + data_folder: source with wav files and validation / test lists + dst_folder: where manifest files will be stored + rebalance: rebalance the classes to have same number of samples. + class_split: whether to use all classes as distinct labels, or to use + 10 classes subset and rest of the classes as noise or background. + skip_duration: Bool whether to skip duration computation. Use this only for + colab notebooks where knowing duration is not necessary for demonstration. + + Returns: + + """ + + if not os.path.exists(dst_folder): + os.makedirs(dst_folder) + + # Used for 10 classes + silence + unknown class setup + class_subset = ["yes", "no", "up", "down", "left", "right", "on", "off", "stop", "go"] + + pattern = re.compile(r"(.+\/)?(\w+)\/([^_]+)_.+wav") + all_files = glob.glob(os.path.join(data_folder, '*/*wav')) + + with open(os.path.join(data_folder, 'validation_list.txt'), 'r') as fin: + validation_files = fin.readlines() + + valset = set() + for entry in validation_files: + r = re.match(pattern, entry) + if r: + valset.add(r.group(3)) + + with open(os.path.join(data_folder, 'testing_list.txt'), 'r') as fin: + testing_files = fin.readlines() + + testset = set() + for entry in testing_files: + r = re.match(pattern, entry) + if r: + testset.add(r.group(3)) + + logging.info("Validation and test set lists extracted") + + label_count = {} + label_filepaths = {} + unknown_val_filepaths = [] + unknown_test_filepaths = [] + + train, val, test = [], [], [] + for entry in all_files: + r = re.match(pattern, entry) + if r: + label, uid = r.group(2), r.group(3) + + if label == '_background_noise_' or label == 'silence': + continue + + if class_split == "sub": + if label not in class_subset: + label = "unknown" # replace label + + if uid in valset: + unknown_val_filepaths.append((label, entry)) + elif uid in testset: + unknown_test_filepaths.append((label, entry)) + + sample = (label, entry) + + if uid not in valset and uid not in testset: + if label in label_count: + label_count[label] += 1 + else: + label_count[label] = 1 + + if label in label_filepaths: + label_filepaths[label] += [sample] + else: + label_filepaths[label] = [sample] + + if label == 'unknown': + continue + + if uid in valset: + val.append(sample) + elif uid in testset: + test.append(sample) + else: + train.append(sample) + + logging.info("Prepared filepaths for dataset") + + # Add silence and unknown class label samples + if class_split == "sub": + logging.info("Perforiming 10+2 class subsplit") + + silence_path = os.path.join(data_folder, "silence") + if not os.path.exists(silence_path): + os.mkdir(silence_path) + + silence_stride = 1000 # 0.0625 second stride + sampling_rate = 16000 + folder = os.path.join(data_folder, "_background_noise_") + + silence_files = [] + rng = np.random.RandomState(0) + + for file in os.listdir(folder): + if ".wav" in file: + load_path = os.path.join(folder, file) + y, sr = librosa.load(load_path, sr=sampling_rate) + + for i in range(0, len(y) - sampling_rate, silence_stride): + file_path = "silence/{}_{}.wav".format(file[:-4], i) + y_slice = y[i : i + sampling_rate] + magnitude = rng.uniform(0.0, 1.0) + y_slice *= magnitude + out_file_path = os.path.join(data_folder, file_path) + librosa.output.write_wav(out_file_path, y_slice, sr) + + silence_files.append(('silence', out_file_path)) + + rng = np.random.RandomState(0) + rng.shuffle(silence_files) + logging.info(f"Constructed silence set of {len(silence_files)}") + + # Create the splits + rng = np.random.RandomState(0) + silence_split = 0.1 + unknown_split = 0.1 + + # train split + num_total_samples = sum([label_count[cls] for cls in class_subset]) + num_silence_samples = int(np.ceil(silence_split * num_total_samples)) + + # initialize sample + label_count['silence'] = 0 + + for silence_id in range(num_silence_samples): + label_count['silence'] += 1 + + if 'silence' in label_filepaths: + label_filepaths['silence'] += [silence_files[silence_id]] + else: + label_filepaths['silence'] = [silence_files[silence_id]] + + train.extend(label_filepaths['silence']) + + # Update train unknown set + unknown_train_samples = label_filepaths['unknown'] + + rng.shuffle(unknown_train_samples) + unknown_size = int(np.ceil(unknown_split * num_total_samples)) + + label_count['unknown'] = unknown_size + label_filepaths['unknown'] = unknown_train_samples[:unknown_size] + + train.extend(label_filepaths['unknown']) + + logging.info("Train set prepared") + + # val set silence + num_val_samples = len(val) + num_silence_samples = int(np.ceil(silence_split * num_val_samples)) + + val_idx = label_count['silence'] + 1 + for silence_id in range(num_silence_samples): + val.append(silence_files[val_idx + silence_id]) + + # Update val unknown set + rng.shuffle(unknown_val_filepaths) + unknown_size = int(np.ceil(unknown_split * num_val_samples)) + + val.extend(unknown_val_filepaths[:unknown_size]) + + logging.info("Validation set prepared") + + # test set silence + num_test_samples = len(test) + num_silence_samples = int(np.ceil(silence_split * num_test_samples)) + + test_idx = val_idx + num_silence_samples + 1 + for silence_id in range(num_silence_samples): + test.append(silence_files[test_idx + silence_id]) + + # Update test unknown set + rng.shuffle(unknown_test_filepaths) + unknown_size = int(np.ceil(unknown_split * num_test_samples)) + + test.extend(unknown_test_filepaths[:unknown_size]) + + logging.info("Test set prepared") + + max_command = None + max_count = -1 + for command, count in label_count.items(): + if command == 'unknown': + continue + + if count > max_count: + max_count = count + max_command = command + + if rebalance: + logging.info(f"Command with maximum number of samples = {max_command} with {max_count} samples") + logging.info(f"Rebalancing dataset by duplicating classes with less than {max_count} samples...") + + for command, samples in label_filepaths.items(): + filepaths = [sample[1] for sample in samples] + + rng = np.random.RandomState(0) + filepaths = np.asarray(filepaths) + num_samples = len(filepaths) + + if num_samples < max_count: + difference = max_count - num_samples + duplication_ids = rng.choice(num_samples, difference, replace=True) + + filepaths = np.append(filepaths, filepaths[duplication_ids], axis=0) + + logging.info(f"Extended class label {command} from {num_samples} samples to {len(filepaths)} samples") + + label_filepaths[command] = [(command, filepath) for filepath in filepaths] + + del train + train = [] + for label, samples in label_filepaths.items(): + train.extend(samples) + + manifests = [ + ('train_manifest.json', train), + ('validation_manifest.json', val), + ('test_manifest.json', test), + ] + + for manifest_filename, dataset in manifests: + with open(os.path.join(dst_folder, manifest_filename), 'w') as fout: + num_files = len(dataset) + pct_file = num_files // 100 + file_count = 0 + + logging.info(f"Preparing manifest : {manifest_filename} with #{num_files} files") + + for label, audio_path in dataset: + if not skip_duration: + duration = librosa.core.get_duration(filename=audio_path) + else: + duration = 0.0 + + # Write the metadata to the manifest + metadata = { + "audio_filepath": audio_path, + "duration": duration, + "command": label, + } + json.dump(metadata, fout) + fout.write('\n') + fout.flush() + + file_count += 1 + if file_count % pct_file == 0: + if not skip_duration: + logging.info(f"Finished serializing {file_count} / {num_files} into {manifest_filename}") + + logging.info(f"Finished construction of manifest : {manifest_filename}") + + if skip_duration: + logging.info( + f"\n<> Duration computation was skipped for demonstration purposes on Colaboratory.\n" + f"In order to replicate paper results and properly perform data augmentation, \n" + f"please recompute the manifest file without the `--skip_duration` flag !\n" + ) + + +def main(): + parser = argparse.ArgumentParser(description='Google Speech Command Data download') + parser.add_argument("--data_root", required=True, default=None, type=str) + parser.add_argument('--data_version', required=True, default=1, type=int, choices=[1, 2]) + parser.add_argument('--class_split', required=False, default='all', type=str, choices=['all', 'sub']) + parser.add_argument('--rebalance', required=False, action='store_true') + parser.add_argument('--skip_duration', required=False, action='store_true') + parser.add_argument('--log', required=False, action='store_true') + parser.set_defaults(log=False, rebalance=False, skip_duration=False) + args = parser.parse_args() + + if args.log: + logging.basicConfig(level=logging.DEBUG) + + data_root = args.data_root + data_set = "google_speech_recognition_v{0}".format(args.data_version) + data_folder = os.path.join(data_root, data_set) + + logging.info(f"Working on: {data_set}") + + if args.data_version == 1: + URL = URL_v1 + else: + URL = URL_v2 + + # Download and extract + if not os.path.exists(data_folder): + file_path = os.path.join(data_root, data_set + ".tar.bz2") + logging.info(f"Getting {data_set}") + __maybe_download_file(file_path, URL) + logging.info(f"Extracting {data_set}") + __extract_all_files(file_path, data_root, data_folder) + + logging.info(f"Processing {data_set}") + __process_data( + data_folder, + data_folder, + rebalance=args.rebalance, + class_split=args.class_split, + skip_duration=args.skip_duration, + ) + logging.info('Done!') + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/dataset_processing/process_vad_data.py b/models/tts/NeMo/scripts/dataset_processing/process_vad_data.py new file mode 100644 index 0000000000..ba98113d9b --- /dev/null +++ b/models/tts/NeMo/scripts/dataset_processing/process_vad_data.py @@ -0,0 +1,486 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +""" +Usage: + +python process_vad_data.py \ + --out_dir= \ + --speech_data_root= \ + --background_data_root= \ + --rebalance_method=<'under' or 'over' or 'fixed'> \ + --log + (Optional --demo (for demonstration in tutorial). If you want to use your own background noise data, make sure to delete --demo) +""" +import argparse +import glob +import json +import logging +import os +import tarfile +import urllib.request + +import librosa +import numpy as np +import soundfile as sf +from sklearn.model_selection import train_test_split + +sr = 16000 + +# google speech command v2 +URL = "http://download.tensorflow.org/data/speech_commands_v0.02.tar.gz" + + +def __maybe_download_file(destination: str, source: str): + """ + Downloads source to destination if it doesn't exist. + If exists, skips download + Args: + destination: local filepath + source: url of resource + + Returns: + + """ + if not os.path.exists(destination): + logging.info(f"{destination} does not exist. Downloading ...") + urllib.request.urlretrieve(source, filename=destination + '.tmp') + os.rename(destination + '.tmp', destination) + logging.info(f"Downloaded {destination}.") + else: + logging.info(f"Destination {destination} exists. Skipping.") + return destination + + +def extract_file(filepath: str, data_dir: str): + try: + tar = tarfile.open(filepath) + tar.extractall(data_dir) + tar.close() + except Exception: + logging.info('Not extracting. Maybe already there?') + + +def __extract_all_files(filepath: str, data_root: str, data_dir: str): + if not os.path.exists(data_dir): + extract_file(filepath, data_dir) + else: + logging.info(f'Skipping extracting. Data already there {data_dir}') + + +def split_train_val_test(data_dir, file_type, test_size=0.1, val_size=0.1, demo=False): + X = [] + if file_type == "speech": + for o in os.listdir(data_dir): + if os.path.isdir(os.path.join(data_dir, o)) and o.split("/")[-1] != "_background_noise_": + X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav')) + + if demo: + logging.info( + f"For Demonstration, we use {int(len(X)/100)}/{len(X)} speech data. Make sure to remove --demo flag when you actually train your model!" + ) + X = np.random.choice(X, int(len(X) / 100), replace=False) + + else: + for o in os.listdir(data_dir): + if os.path.isdir(os.path.join(data_dir, o)): + X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav')) + else: # for using "_background_noise_" from google speech commands as background data + if o.endswith(".wav"): + X.append(os.path.join(data_dir, o)) + + X_train, X_test = train_test_split(X, test_size=test_size, random_state=1) + val_size_tmp = val_size / (1 - test_size) + X_train, X_val = train_test_split(X_train, test_size=val_size_tmp, random_state=1) + + with open(os.path.join(data_dir, file_type + "_training_list.txt"), "w") as outfile: + outfile.write("\n".join(X_train)) + with open(os.path.join(data_dir, file_type + "_testing_list.txt"), "w") as outfile: + outfile.write("\n".join(X_test)) + with open(os.path.join(data_dir, file_type + "_validation_list.txt"), "w") as outfile: + outfile.write("\n".join(X_val)) + + logging.info(f'Overall: {len(X)}, Train: {len(X_train)}, Validatoin: {len(X_val)}, Test: {len(X_test)}') + logging.info(f"Finish spliting train, val and test for {file_type}. Write to files!") + + +def process_google_speech_train(data_dir): + X = [] + for o in os.listdir(data_dir): + if os.path.isdir(os.path.join(data_dir, o)) and o.split("/")[-1] != "_background_noise_": + X.extend(glob.glob(os.path.join(data_dir, o) + '/*.wav')) + + short_files = [i.split(data_dir)[1] for i in files] + + with open(os.path.join(data_dir, 'testing_list.txt'), 'r') as allfile: + testing_list = allfile.read().splitlines() + + with open(os.path.join(data_dir, 'validation_list.txt'), 'r') as allfile: + validation_list = allfile.read().splitlines() + + exist_set = set(testing_list).copy() + exist_set.update(set(validation_list)) + + training_list = [i for i in short_files if i not in exist_set] + + with open(os.path.join(data_dir, "training_list.txt"), "w") as outfile: + outfile.write("\n".join(training_list)) + + logging.info( + f'Overall: {len(files)}, Train: {len(training_list)}, Validatoin: {len(validation_list)}, Test: {len(testing_list)}' + ) + + +def write_manifest( + out_dir, + files, + prefix, + manifest_name, + start=0.0, + end=None, + duration_stride=1.0, + duration_max=None, + duration_limit=100.0, + filter_long=False, +): + """ + Given a list of files, segment each file and write them to manifest with restrictions. + Args: + out_dir: directory of generated manifest + files: list of files to be processed + prefix: label of samples + manifest_name: name of generated manifest + start: beginning of audio of generating segment + end: end of audio of generating segment + duration_stride: stride for segmenting audio samples + duration_max: duration for each segment + duration_limit: duration threshold for filtering out long audio samples + filter_long: boolean to determine whether to filter out long audio samples + Returns: + """ + seg_num = 0 + skip_num = 0 + if duration_max is None: + duration_max = 1e9 + + if not os.path.exists(out_dir): + logging.info(f'Outdir {out_dir} does not exist. Creat directory.') + os.mkdir(out_dir) + + output_path = os.path.join(out_dir, manifest_name + '.json') + with open(output_path, 'w') as fout: + for file in files: + label = prefix + + try: + x, _sr = librosa.load(file, sr=sr) + duration = librosa.get_duration(x, sr=sr) + + except Exception: + continue + + if filter_long and duration > duration_limit: + skip_num += 1 + continue + + offsets = [] + durations = [] + + if duration > duration_max: + current_offset = start + + while current_offset < duration: + if end is not None and current_offset > end: + break + + difference = duration - current_offset + + if difference < duration_max: + break + + offsets.append(current_offset) + durations.append(duration_max) + + current_offset += duration_stride + + else: + # Duration is not long enough! Skip + skip_num += 1 + + for duration, offset in zip(durations, offsets): + metadata = { + 'audio_filepath': file, + 'duration': duration, + 'label': label, + 'text': '_', # for compatibility with ASRAudioText + 'offset': offset, + } + json.dump(metadata, fout) + fout.write('\n') + fout.flush() + seg_num += 1 + return skip_num, seg_num, output_path + + +def load_list_write_manifest( + data_dir, + out_dir, + filename, + prefix, + start, + end, + duration_stride=1.0, + duration_max=1.0, + duration_limit=100.0, + filter_long=True, +): + + filename = prefix + '_' + filename + file_path = os.path.join(data_dir, filename) + + with open(file_path, 'r') as allfile: + files = allfile.read().splitlines() + + manifest_name = filename.split('_list.txt')[0] + '_manifest' + skip_num, seg_num, output_path = write_manifest( + out_dir, + files, + prefix, + manifest_name, + start, + end, + duration_stride, + duration_max, + duration_limit, + filter_long=True, + ) + return skip_num, seg_num, output_path + + +def rebalance_json(data_dir, data_json, num, prefix): + data = [] + seg = 0 + for line in open(data_json, 'r'): + data.append(json.loads(line)) + filename = data_json.split('/')[-1] + fout_path = os.path.join(data_dir, prefix + "_" + filename) + + if len(data) >= num: + selected_sample = np.random.choice(data, num, replace=False) + else: + selected_sample = np.random.choice(data, num, replace=True) + + with open(fout_path, 'a') as fout: + for i in selected_sample: + seg += 1 + json.dump(i, fout) + fout.write('\n') + fout.flush() + + logging.info(f'Get {seg}/{num} to {fout_path} from {data_json}') + return fout_path + + +def generate_variety_noise(data_dir, filename, prefix): + + curr_dir = data_dir.split("_background_noise_")[0] + silence_path = os.path.join(curr_dir, "_background_noise_more") + + if not os.path.exists(silence_path): + os.mkdir(silence_path) + + silence_stride = 1000 # stride = 1/16 seconds + sampling_rate = 16000 + + silence_files = [] + rng = np.random.RandomState(0) + + filename = prefix + '_' + filename + file_path = os.path.join(data_dir, filename) + + with open(file_path, 'r') as allfile: + files = allfile.read().splitlines() + + for file in files: + y, sr = librosa.load(file, sr=sampling_rate) + + for i in range( + 0, len(y) - sampling_rate, silence_stride * 100 + ): # stride * 100 to generate less samples for demo + file_name = "{}_{}.wav".format(file.split("/")[-1], i) + y_slice = y[i : i + sampling_rate] + magnitude = rng.uniform(0.0, 1.0) + y_slice *= magnitude + out_file_path = os.path.join(silence_path, file_name) + sf.write(out_file_path, y_slice, sr) + + silence_files.append(out_file_path) + + new_list_file = os.path.join(silence_path, filename) + with open(new_list_file, "w") as outfile: + outfile.write("\n".join(silence_files)) + + logging.info(f"Generate {len(out_file_path)} background files for {file_path}. => {new_list_file} !") + return len(silence_files) + + +def main(): + parser = argparse.ArgumentParser(description='Speech and backgound data download and preprocess') + parser.add_argument("--out_dir", required=False, default='./manifest/', type=str) + parser.add_argument("--speech_data_root", required=True, default=None, type=str) + parser.add_argument("--background_data_root", required=True, default=None, type=str) + parser.add_argument('--test_size', required=False, default=0.1, type=float) + parser.add_argument('--val_size', required=False, default=0.1, type=float) + parser.add_argument('--seg_len', required=False, default=0.63, type=float) + parser.add_argument('--log', required=False, action='store_true') + parser.add_argument('--rebalance_method', required=False, default=None, type=str) + parser.add_argument('--demo', required=False, action='store_true') + parser.set_defaults(log=False, generate=False) + args = parser.parse_args() + + if not args.rebalance_method: + rebalance = False + else: + if args.rebalance_method != 'over' and args.rebalance_method != 'under' and args.rebalance_method != 'fixed': + raise NameError("Please select a valid sampling method: over/under/fixed.") + else: + rebalance = True + + if args.log: + logging.basicConfig(level=logging.DEBUG) + + # Download speech data + speech_data_root = args.speech_data_root + data_set = "google_speech_recognition_v2" + speech_data_folder = os.path.join(speech_data_root, data_set) + + background_data_folder = args.background_data_root + logging.info(f"Working on: {data_set}") + + # Download and extract speech data + if not os.path.exists(speech_data_folder): + file_path = os.path.join(speech_data_root, data_set + ".tar.bz2") + logging.info(f"Getting {data_set}") + __maybe_download_file(file_path, URL) + logging.info(f"Extracting {data_set}") + __extract_all_files(file_path, speech_data_root, speech_data_folder) + + logging.info(f"Split speech data!") + # dataset provide testing.txt and validation.txt feel free to split data using that with process_google_speech_train + split_train_val_test(speech_data_folder, "speech", args.test_size, args.val_size, args.demo) + + logging.info(f"Split background data!") + split_train_val_test(background_data_folder, "background", args.test_size, args.val_size) + + out_dir = args.out_dir + + # Process Speech manifest + logging.info(f"=== Write speech data to manifest!") + skip_num_val, speech_seg_num_val, speech_val = load_list_write_manifest( + speech_data_folder, out_dir, 'validation_list.txt', 'speech', 0.2, 0.8, args.seg_len, args.seg_len + ) + skip_num_test, speech_seg_num_test, speech_test = load_list_write_manifest( + speech_data_folder, out_dir, 'testing_list.txt', 'speech', 0.2, 0.8, 0.01, args.seg_len + ) + skip_num_train, speech_seg_num_train, speech_train = load_list_write_manifest( + speech_data_folder, out_dir, 'training_list.txt', 'speech', 0.2, 0.8, args.seg_len, args.seg_len + ) + + logging.info(f'Val: Skip {skip_num_val} samples. Get {speech_seg_num_val} segments! => {speech_val} ') + logging.info(f'Test: Skip {skip_num_test} samples. Get {speech_seg_num_test} segments! => {speech_test}') + logging.info(f'Train: Skip {skip_num_train} samples. Get {speech_seg_num_train} segments!=> {speech_train}') + + # Process background manifest + # if we select to generate more background noise data + if args.demo: + logging.info("Start generating more background noise data") + generate_variety_noise(background_data_folder, 'validation_list.txt', 'background') + generate_variety_noise(background_data_folder, 'training_list.txt', 'background') + generate_variety_noise(background_data_folder, 'testing_list.txt', 'background') + background_data_folder = os.path.join( + background_data_folder.split("_background_noise_")[0], "_background_noise_more" + ) + + logging.info(f"=== Write background data to manifest!") + skip_num_val, background_seg_num_val, background_val = load_list_write_manifest( + background_data_folder, out_dir, 'validation_list.txt', 'background', 0, None, 0.15, args.seg_len + ) + skip_num_test, background_seg_num_test, background_test = load_list_write_manifest( + background_data_folder, out_dir, 'testing_list.txt', 'background', 0, None, 0.01, args.seg_len + ) + skip_num_train, background_seg_num_train, background_train = load_list_write_manifest( + background_data_folder, out_dir, 'training_list.txt', 'background', 0, None, 0.15, args.seg_len + ) + + logging.info(f'Val: Skip {skip_num_val} samples. Get {background_seg_num_val} segments! => {background_val}') + logging.info(f'Test: Skip {skip_num_test} samples. Get {background_seg_num_test} segments! => {background_test}') + logging.info( + f'Train: Skip {skip_num_train} samples. Get {background_seg_num_train} segments! => {background_train}' + ) + min_val, max_val = min(speech_seg_num_val, background_seg_num_val), max(speech_seg_num_val, background_seg_num_val) + min_test, max_test = ( + min(speech_seg_num_test, background_seg_num_test), + max(speech_seg_num_test, background_seg_num_test), + ) + min_train, max_train = ( + min(speech_seg_num_train, background_seg_num_train), + max(speech_seg_num_train, background_seg_num_train), + ) + + logging.info('Finish generating manifest!') + + if rebalance: + # Random Oversampling: Randomly duplicate examples in the minority class. + # Random Undersampling: Randomly delete examples in the majority class. + if args.rebalance_method == 'under': + logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.") + logging.info(f'Val: {min_val} Test: {min_test} Train: {min_train}!') + + rebalance_json(out_dir, background_val, min_val, 'balanced') + rebalance_json(out_dir, background_test, min_test, 'balanced') + rebalance_json(out_dir, background_train, min_train, 'balanced') + + rebalance_json(out_dir, speech_val, min_val, 'balanced') + rebalance_json(out_dir, speech_test, min_test, 'balanced') + rebalance_json(out_dir, speech_train, min_train, 'balanced') + + if args.rebalance_method == 'over': + logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.") + logging.info(f'Val: {max_val} Test: {max_test} Train: {max_train}!') + + rebalance_json(out_dir, background_val, max_val, 'balanced') + rebalance_json(out_dir, background_test, max_test, 'balanced') + rebalance_json(out_dir, background_train, max_train, 'balanced') + + rebalance_json(out_dir, speech_val, max_val, 'balanced') + rebalance_json(out_dir, speech_test, max_test, 'balanced') + rebalance_json(out_dir, speech_train, max_train, 'balanced') + + if args.rebalance_method == 'fixed': + fixed_test, fixed_val, fixed_train = 200, 100, 500 + logging.info(f"Rebalancing number of samples in classes using {args.rebalance_method} sampling.") + logging.info(f'Val: {fixed_val} Test: {fixed_test} Train: {fixed_train}!') + + rebalance_json(out_dir, background_val, fixed_val, 'balanced') + rebalance_json(out_dir, background_test, fixed_test, 'balanced') + rebalance_json(out_dir, background_train, fixed_train, 'balanced') + + rebalance_json(out_dir, speech_val, fixed_val, 'balanced') + rebalance_json(out_dir, speech_test, fixed_test, 'balanced') + rebalance_json(out_dir, speech_train, fixed_train, 'balanced') + else: + logging.info("Don't rebalance number of samples in classes.") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/scripts/export/convasr_to_enemo.py b/models/tts/NeMo/scripts/export/convasr_to_enemo.py new file mode 100644 index 0000000000..908236f418 --- /dev/null +++ b/models/tts/NeMo/scripts/export/convasr_to_enemo.py @@ -0,0 +1,76 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os +import tarfile +from shutil import copyfile + +from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel +from nemo.utils import logging + + +def get_parser(): + parser = argparse.ArgumentParser(description="Convert .nemo file to onnx file (encoder and decoder fused)") + parser.add_argument( + "--nemo_file", default=None, type=str, required=True, help="Path to .nemo file", + ) + parser.add_argument( + "--onnx_file", default=None, type=str, required=True, help="Path to the onnx output", + ) + parser.add_argument( + "--enemo_file", default=None, type=str, required=True, help="Path to the enemo output, consiting of onnx file", + ) + parser.add_argument( + "--model_type", + default="asr", + type=str, + choices=["asr", "speech_label", "speaker"], + help="Type of decoder used by the model.", + ) + return parser + + +def main( + nemo_file, enemo_file, onnx_file, model_type="asr", +): + if model_type == "asr": + logging.info("Preparing ASR model") + model = EncDecCTCModel.restore_from(nemo_file) + elif model_type == "speech_label": + logging.info("Preparing Speech Label Classification model") + model = EncDecClassificationModel.restore_from(nemo_file) + elif model_type == "speaker": + logging.info("Preparing Speaker Recognition model") + model = EncDecSpeakerLabelModel.restore_from(nemo_file) + else: + raise NameError("Available model names are asr, speech_label and speaker") + + logging.info("Writing onnx file") + model.export(onnx_file, onnx_opset_version=12) + logging.info("succesfully ported onnx file") + + with tarfile.open(nemo_file, "r") as archive: + archive.extract("./model_config.yaml") + with tarfile.open(enemo_file, "w") as enemo_archive: + enemo_archive.add("./model_config.yaml") + copyfile(onnx_file, "model_graph.onnx") + enemo_archive.add("model_graph.onnx") + os.remove("model_graph.onnx") # cleanup extra file + + +if __name__ == "__main__": + args = get_parser().parse_args() + main( + args.nemo_file, args.enemo_file, args.onnx_file, model_type=args.model_type, + ) diff --git a/models/tts/NeMo/scripts/export/convasr_to_onnx.py b/models/tts/NeMo/scripts/export/convasr_to_onnx.py new file mode 100644 index 0000000000..2b92e98835 --- /dev/null +++ b/models/tts/NeMo/scripts/export/convasr_to_onnx.py @@ -0,0 +1,66 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse + +from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel +from nemo.utils import logging + + +def get_parser(): + parser = argparse.ArgumentParser(description="Convert .nemo file to encoder decoder onnx files") + parser.add_argument( + "--nemo_file", default=None, type=str, required=True, help="Path to .nemo file", + ) + parser.add_argument( + "--onnx_encoder", default=None, type=str, required=True, help="Path to the onnx encoder output.", + ) + parser.add_argument( + "--onnx_decoder", default=None, type=str, required=True, help="Path to the onnx decoder output.", + ) + parser.add_argument( + "--model_type", + default='asr', + type=str, + choices=['asr', 'speech_label', 'speaker'], + help="Type of decoder used by the model.", + ) + return parser + + +def main( + nemo_file, onnx_encoder, onnx_decoder, model_type='asr', +): + if model_type == 'asr': + logging.info("Preparing encoder decoder for ASR model") + model = EncDecCTCModel.restore_from(nemo_file) + elif model_type == 'speech_label': + logging.info("Preparing encoder decoder for Speech Label Classification model") + model = EncDecClassificationModel.restore_from(nemo_file) + elif model_type == 'speaker': + logging.info("Preparing encoder decoder for Speaker Recognition model") + model = EncDecSpeakerLabelModel.restore_from(nemo_file) + else: + raise NameError("Available model names are asr, speech_label and speaker ") + + logging.info("Writing onnx encoder and decoder onnx files") + model.encoder.export(onnx_encoder) + model.decoder.export(onnx_decoder, onnx_opset_version=10) + logging.info("succesfully ported onnx files") + + +if __name__ == "__main__": + args = get_parser().parse_args() + main( + args.nemo_file, args.onnx_encoder, args.onnx_decoder, model_type=args.model_type, + ) diff --git a/models/tts/NeMo/scripts/export/convasr_to_single_onnx.py b/models/tts/NeMo/scripts/export/convasr_to_single_onnx.py new file mode 100644 index 0000000000..7c7321f723 --- /dev/null +++ b/models/tts/NeMo/scripts/export/convasr_to_single_onnx.py @@ -0,0 +1,62 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse + +from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel +from nemo.utils import logging + + +def get_parser(): + parser = argparse.ArgumentParser(description="Convert .nemo file to onnx file (encoder and decoder fused)") + parser.add_argument( + "--nemo_file", default=None, type=str, required=True, help="Path to .nemo file", + ) + parser.add_argument( + "--onnx_file", default=None, type=str, required=True, help="Path to the onnx output", + ) + parser.add_argument( + "--model_type", + default='asr', + type=str, + choices=['asr', 'speech_label', 'speaker'], + help="Type of decoder used by the model.", + ) + return parser + + +def main( + nemo_file, onnx_file, model_type='asr', +): + if model_type == 'asr': + logging.info("Preparing ASR model") + model = EncDecCTCModel.restore_from(nemo_file) + elif model_type == 'speech_label': + logging.info("Preparing Speech Label Classification model") + model = EncDecClassificationModel.restore_from(nemo_file) + elif model_type == 'speaker': + logging.info("Preparing Speaker Recognition model") + model = EncDecSpeakerLabelModel.restore_from(nemo_file) + else: + raise NameError("Available model names are asr, speech_label and speaker") + + logging.info("Writing onnx file") + model.export(onnx_file, onnx_opset_version=12) + logging.info("succesfully ported onnx file") + + +if __name__ == "__main__": + args = get_parser().parse_args() + main( + args.nemo_file, args.onnx_file, model_type=args.model_type, + ) diff --git a/models/tts/NeMo/scripts/freesound_download_resample/download_resample_freesound.sh b/models/tts/NeMo/scripts/freesound_download_resample/download_resample_freesound.sh new file mode 100644 index 0000000000..af0a07f44d --- /dev/null +++ b/models/tts/NeMo/scripts/freesound_download_resample/download_resample_freesound.sh @@ -0,0 +1,115 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +#!/bin/bash + +# This is bash script actually run the downloading and resampling script. +# See instructions in freesound_download.py + +# Change this arguments if you want +page_size=100 # Number of sounds per page +max_samples=200 # Maximum number of sound samples +min_filesize=0 # Minimum filesize allowed (in MB) +max_filesize=100 # Maximum filesize allowed (in MB) + +if [[ $# -ne 3 ]]; then + echo "Require number of all files | data directory | resample data directory as arguments to the script" + exit 2 +fi + + +NUM_ALL_FILES=$1 +DATADIR=$2 +RESAMPLE_DATADIR=$3 + + +if [ ! -d "$DATADIR" ]; then + echo "Creating dir $DATADIR" + mkdir -p "$DATADIR" +fi + +if [ ! -d "$RESAMPLE_DATADIR" ]; then + echo "Creating dir $RESAMPLE_DATADIR" + mkdir -p "$RESAMPLE_DATADIR" +fi + +# we just need background categories for constructing dataset, feel free to include other (speech) categories for testing and training your VAD model +# background +categories=( + "Air brake" + "Static" + "Acoustic environment" + "Distortion" + "Tape hiss" + "Hubbub" + "Vibration" + "Cacophony" + "Throbbing" + "Reverberation" + "Inside, public space" + "Inside, small room" + "Echo" + "Outside, rural" + "Outside, natural" + "Outside, urban" + "Outside, manmade" + "Car" + "Bus" + "Traffic noise" + "Roadway noise" + "Truck" + "Emergency vehicle" + "Motorcycle" + "Aircraft engine" + "Aircraft" + "Helicopter" + "Bicycle" + "Skateboard" + "Subway, metro, underground" + "Railroad car" + "Train wagon" + "Train" + "Sailboat" + "Rowboat" + "Ship" +) + + +WAV_FILECOUNT="$(find $DATADIR -name '*.wav' -type f | wc -l)" +FLAC_FILECOUNT="$(find $DATADIR -name '*.flac' -type f | wc -l)" +FILECOUNT="$((WAV_FILECOUNT + FLAC_FILECOUNT))" +echo "File count: " $FILECOUNT + + +while((FILECOUNT <= NUM_ALL_FILES)) +do + for category in "${categories[@]}" + do + python freesound_download.py --data_dir "${DATADIR}" --category "${category}" --page_size "${page_size}" --max_samples "${max_samples}" --min_filesize "${min_filesize}" --max_filesize "${max_filesize}" + ret=$? + if [ $ret -ne 0 ]; then + exit 1 + fi + done + + WAV_FILECOUNT="$(find $DATADIR -name '*.wav' -type f | wc -l)" + FLAC_FILECOUNT="$(find $DATADIR -name '*.flac' -type f | wc -l)" + FILECOUNT="$((WAV_FILECOUNT + FLAC_FILECOUNT))" + echo "Current file count is: " $FILECOUNT +done + +# RESAMPLE +echo "Got enough data. Start resample!" +python freesound_resample.py --data_dir="${DATADIR}" --resampled_dir="${RESAMPLE_DATADIR}" + +echo "Done resample data!" diff --git a/models/tts/NeMo/scripts/freesound_download_resample/freesound_download.py b/models/tts/NeMo/scripts/freesound_download_resample/freesound_download.py new file mode 100644 index 0000000000..076c69bfa6 --- /dev/null +++ b/models/tts/NeMo/scripts/freesound_download_resample/freesound_download.py @@ -0,0 +1,585 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import os +import pickle +import time + +try: + import librosa + import requests + import requests_oauthlib + from joblib import Parallel, delayed + from oauthlib.oauth2 import TokenExpiredError +except (ModuleNotFoundError, ImportError) as e: + raise e + +try: + import freesound +except ModuleNotFoundError as e: + raise ModuleNotFoundError( + "freesound is not installed. Execute `pip install --no-cache-dir git+https://github.com/MTG/freesound-python.git` in terminal" + ) + + +""" +Instructions +1. We will need some requirements including freesound, requests, requests_oauthlib, joblib, librosa and sox. If they are not installed, please run `pip install -r freesound_requirements.txt` +2. Create an API key for freesound.org at https://freesound.org/help/developers/ +3. Create a python file called `freesound_private_apikey.py` and add lined `api_key = ` and `client_id = ` +4. Authorize by run `python freesound_download.py --authorize` and visit website, and paste response code +5. Feel free to change any arguments in download_resample_freesound.sh such as max_samples and max_filesize +6. Run `bash download_resample_freesound.sh ` +""" + +# Import the API Key +try: + from freesound_private_apikey import api_key, client_id + + print("API Key found !") +except ImportError: + raise ImportError( + "Create a python file called `freesound_private_apikey.py` and add lined `api_key = ` and `client_id = `" + ) + +auth_url = 'https://freesound.org/apiv2/oauth2/authorize/' +redirect_url = 'https://freesound.org/home/app_permissions/permission_granted/' +token_url = 'https://freesound.org/apiv2/oauth2/access_token/' +scope = ["read", "write"] + +BACKGROUND_CLASSES = [ + "Air brake", + "Static", + "Acoustic environment", + "Distortion", + "Tape hiss", + "Hubbub", + "Vibration", + "Cacophony", + "Throbbing", + "Reverberation", + "Inside, public space", + "Inside, small room", + "Echo", + "Outside, rural", + "Outside, natural", + "Outside, urban", + "Outside, manmade", + "Car", + "Bus", + "Traffic noise", + "Roadway noise", + "Truck", + "Emergency vehicle", + "Motorcycle", + "Aircraft engine", + "Aircraft", + "Helicopter", + "Bicycle", + "Skateboard", + "Subway, metro, underground", + "Railroad car", + "Train wagon", + "Train", + "Sailboat", + "Rowboat", + "Ship", +] + +SPEECH_CLASSES = [ + "Male speech", + "Female speech", + "Speech synthesizer", + "Babbling", + "Conversation", + "Child speech", + "Narration", + "Laughter", + "Yawn", + "Whispering", + "Whimper", + "Baby cry", + "Sigh", + "Groan", + "Humming", + "Male singing", + "Female singing", + "Child singing", + "Children shouting", +] + + +def initialize_oauth(): + # If token already exists, then just load it + if os.path.exists('_token.pkl'): + token = unpickle_object('_token') + oauth = requests_oauthlib.OAuth2Session(client_id, redirect_uri=redirect_url, scope=scope, token=token) + + else: + # Construct a new token after OAuth2 flow + # Initialize a OAuth2 session + oauth = requests_oauthlib.OAuth2Session(client_id, redirect_uri=redirect_url, scope=scope) + + authorization_url, state = oauth.authorization_url(auth_url) + print(f"Visit below website and paste access token below : \n\n{authorization_url}\n") + + authorization_response = input("Paste authorization response code here :\n") + + token = oauth.fetch_token( + token_url, + authorization_response=authorization_response, + code=authorization_response, + client_secret=api_key, + ) + + # Save the token generated + pickle_object(token, '_token') + + return oauth, token + + +def instantiate_session(): + # Reconstruct session in process, and force singular execution thread to reduce session + # connections to server + token = unpickle_object('_token') + session = requests_oauthlib.OAuth2Session(client_id, redirect_uri=redirect_url, scope=scope, token=token) + adapter = requests.adapters.HTTPAdapter(pool_connections=1, pool_maxsize=1) + session.mount('http://', adapter) + return session + + +def refresh_token(session): + print("Refreshing tokens...") + # Token expired, perform token refresh + extras = {'client_id': client_id, 'client_secret': api_key} + token = session.refresh_token(token_url, **extras) + print("Token refresh performed...") + # Save the refreshed token + pickle_object(token, '_token') + return session + + +def pickle_object(token, name): + with open(name + '.pkl', 'wb') as f: + pickle.dump(token, f) + + +def unpickle_object(name): + fp = name + '.pkl' + if os.path.exists(fp): + with open(fp, 'rb') as f: + token = pickle.load(f) + + return token + else: + raise FileNotFoundError('Token not found!') + + +def is_resource_limited(e: freesound.FreesoundException): + """ + Test if the reason for a freesound exception was either rate limit + or daily limit. + + If it was for either reason, sleep for an appropriate delay and return + to try again. + + Args: + e: Freesound Exception object + + Returns: + A boolean which describes whether the error was due to some + api limit issue, or if it was some other reason. + + If false is returned, then the user should carefully check the cause + and log it. + """ + detail = e.detail['detail'] + + if '2000' in detail: + # This is the request limit, hold off for 1 hour and try again + print(f"Hit daily limit, sleeping for 20 minutes.") + time.sleep(60 * 20) + return True + + elif '60' in detail: + # This is the request limit per minute, hold off for 1 minute and try again + print(f"Hit rate limit, sleeping for 1 minute.") + time.sleep(60) + return True + + else: + return False + + +def prepare_client(client: freesound.FreesoundClient, token) -> freesound.FreesoundClient: + # Initialize the client with token auth + client.set_token(token['access_token'], auth_type='oauth') + print("Client ready !") + return client + + +def get_text_query_with_resource_limit_checks(client, query: str, filters: list, fields: str, page_size: int): + """ + Performs a text query, checks for rate / api limits, and retries. + + Args: + client: FreesoundAPI client + query: query string (either exact or inexact) + filters: list of string filters + fields: String of values to recover + page_size: samples per page returned + + Returns: + + """ + pages = None + attempts = 20 + + while pages is None: + try: + pages = client.text_search(query=query, filter=" ".join(filters), fields=fields, page_size=str(page_size),) + + except freesound.FreesoundException as e: + # Most probably a rate limit or a request limit + # Check if that was the case, and wait appropriate ammount of time + # for retry + was_resource_limited = is_resource_limited(e) + + # If result of test False, it means that failure was due to some other reason. + # Log it, then break loop + if not was_resource_limited: + print(e.with_traceback(None)) + break + + attempts -= 1 + + # Attempt to refresh tokens if it fails multiple times + if attempts % 5 == 0 and attempts > 0: + session = instantiate_session() + refresh_token(session) + session.close() + token = unpickle_object('_token') + client = prepare_client(client, token) + + if attempts <= 0: + print(f"Failed to query pages for '{query}' after 10 attempts, skipping query") + break + + if pages is None: + print(f"Query attempts remaining = {attempts}") + + return client, pages + + +def get_resource_with_auto_refresh(session, download_url): + """ + Attempts download of audio with a token refresh if necessary. + """ + try: + result = session.get(download_url) + + except TokenExpiredError as e: + session = refresh_token(session) + result = session.get(download_url) + + except Exception as e: + result = None + + print(f"Skipping file {download_url} due to exception below\n\n") + print(e) + + return result.content + + +def download_song(basepath, id, name, download_url): + # Cleanup name + name = name.replace(":", "-") + name = name.replace(",", "-") + name = name.replace("/", "-") + name = name.replace("\\", "-") + name = name.replace(".", "-") + name = name.replace(" ", "") + + # Correct last `.` for filetype + name = name[:-4] + '.wav' + + # Add file id to filename + name = f"id_{id} " + name + + fp = os.path.join(basepath, name) + + # Check if file, if exists already, can be loaded by librosa + # If it cannot be loaded, possibly corrupted file. + # Delete and then re-download + if os.path.exists(fp): + try: + _ = librosa.load(fp) + except Exception: + # File is currupted, delete and re-download. + os.remove(fp) + + print(f"Pre-existing file {fp} was corrupt and was deleted, will be re-downloaded.") + + if not os.path.exists(fp): + print("Downloading file :", name) + + session = instantiate_session() + + data = None + attempts = 10 + + try: + while data is None: + + try: + # Get the sound data + data = get_resource_with_auto_refresh(session, download_url) + + except freesound.FreesoundException as e: + # Most probably a rate limit or a request limit + # Check if that was the case, and wait appropriate amount of time + # for retry + was_resource_limited = is_resource_limited(e) + + # If result of test False, it means that failure was due to some other reason. + # Log it, then break loop + if not was_resource_limited: + print(e) + break + + attempts -= 1 + + if attempts <= 0: + print(f"Failed to download file {fp} after 10 attempts, skipping file") + break + + if data is None: + print(f"Download attempts remaining = {attempts}") + + finally: + session.close() + + # Write the data to file + if data is not None: + print("Downloaded file :", name) + + with open(fp, 'wb') as f: + f.write(data) + + # If file size is less than 89, then this probably is a text format and not an actual audio file. + if os.path.getsize(fp) > 89: + print(f"File written : {fp}") + + else: + os.remove(fp) + print(f"File corrupted and has been deleted: {fp}") + + else: + print(f"File [{fp}] corrupted or faced some issue when downloading, skipped.") + + # Sleep to avoid hitting rate limits + time.sleep(5) + + else: + print(f"File [{fp}] already exists in dataset, skipping re-download.") + + +def get_songs_by_category( + client: freesound.FreesoundClient, + category: str, + data_dir: str, + max_num_samples=100, + page_size=100, + min_filesize_in_mb=0, + max_filesize_in_mb=10, + n_jobs=None, +): + """ + Download songs of a category with restrictions + + Args: + client: FreesoundAPI client + category: category to be downloaded + data_dir: directory of downloaded songs + max_num_samples: maximum number of samples of this category + page_size: samples per page returned + min_filesize_in_mb: minimum filesize of the song in MB + max_filesize_in_mb: maximum filesize of the song in MB + n_jobs: number of jobs for parallel processing + + Returns: + + """ + # quote string to force exact match + query = f'"{category}"' + print(f"Query : {query}") + + page_size = min(page_size, 150) + max_filesize = int(max_filesize_in_mb * (2 ** 20)) + + if min_filesize_in_mb == 0: + min_filesize_in_mb = 1 + else: + min_filesize_in_mb = int(min_filesize_in_mb * (2 ** 20)) + + if max_num_samples < 0: + max_num_samples = int(1e6) + + filters = [ + 'type:(wav OR flac)', + 'license:("Attribution" OR "Creative Commons 0")', + f'filesize:[{min_filesize_in_mb} TO {max_filesize}]', + ] + + fields = "id,name,download,license" + + client, pages = get_text_query_with_resource_limit_checks( + client, query=query, filters=filters, fields=fields, page_size=page_size + ) + + if pages is None: + print(f"Number of attempts exceeded limit, skipping query {query}") + return + + num_pages = pages.count + + # Check if returned empty result; if so, fallback to inexact category search + if num_pages == 0: + print(f"Found 0 samples of results for query '{query}'") + print(f"Trying less restricted query : {category}") + + client, pages = get_text_query_with_resource_limit_checks( + client, query=category, filters=filters, fields=fields, page_size=page_size + ) + + if pages is None: + print(f"Number of attempts exceeded limit, skipping query {query}") + return + + num_pages = pages.count + + print(f"Found {num_pages} samples of results for query '{query}'") + + category = category.replace(' ', '_') + basepath = os.path.join(data_dir, category) + + if not os.path.exists(basepath): + os.makedirs(basepath) + + sounds = [] + sample_count = 0 + + # Retrieve sound license information + with open(os.path.join(basepath, 'licenses.txt'), 'w') as f: + f.write("ID,LICENSE\n") + f.flush() + + while True: + for sound in pages: + if sample_count >= max_num_samples: + print( + f"Collected {sample_count} samples, which is >= max number of samples requested " + f"{max_num_samples}. Stopping for this category : {category}" + ) + break + + sounds.append(sound) + sample_count += 1 + + f.write(f"{sound.id},{sound.license}\n") + f.flush() + + if sample_count >= max_num_samples: + break + + try: + pages = pages.next_page() + except ValueError: + break + + if n_jobs is None: + n_jobs = max(1, len(sounds)) + + # Parallel download all songs + with Parallel(n_jobs=n_jobs, verbose=10) as parallel: + _ = parallel(delayed(download_song)(basepath, sound.id, sound.name, sound.download) for sound in sounds) + + +if __name__ == '__main__': + + parser = argparse.ArgumentParser(description="Freesound download script") + + parser.add_argument( + '--authorize', action='store_true', dest='auth', help='Flag to only perform OAuth2 authorization step' + ) + + parser.add_argument('-c', '--category', default='', type=str, help='Category required to download') + + parser.add_argument('-d', '--data_dir', default='', type=str, help='Destination folder to store data') + + parser.add_argument('--page_size', default=100, type=int, help='Number of sounds per page') + + parser.add_argument('--max_samples', default=100, type=int, help='Maximum number of sound samples') + + parser.add_argument('--min_filesize', default=0, type=int, help='Maximum filesize allowed (in MB)') + + parser.add_argument('--max_filesize', default=20, type=int, help='Maximum filesize allowed (in MB)') + + parser.set_defaults(auth=False) + + args = parser.parse_args() + + if args.auth: + """ Initialize oauth token to be used by all """ + oauth, token = initialize_oauth() + oauth.close() + + print("Authentication suceeded ! Token stored in `_token.pkl`") + exit(0) + + if not os.path.exists('_token.pkl'): + raise FileNotFoundError( + "Please authorize the application first using " "`python freesound_download.py --authorize`" + ) + if args.data_dir == '': + raise ValueError("Data dir must be passed as an argument using `--data_dir`") + + data_dir = args.data_dir + + page_size = args.page_size + max_num_samples = args.max_samples + min_filesize_in_mb = args.min_filesize + max_filesize_in_mb = args.max_filesize + + # Initialize and authenticate client + token = unpickle_object('_token') + freesound_client = freesound.FreesoundClient() + client = prepare_client(freesound_client, token) + + category = args.category + + if category == '': + raise ValueError("Cannot pass empty string as it will select all of FreeSound data !") + + print(f"Downloading category : {category}") + get_songs_by_category( + client, + category, + data_dir=data_dir, + max_num_samples=max_num_samples, + page_size=page_size, + min_filesize_in_mb=min_filesize_in_mb, + max_filesize_in_mb=max_filesize_in_mb, + n_jobs=30, + ) diff --git a/models/tts/NeMo/scripts/freesound_download_resample/freesound_requirements.txt b/models/tts/NeMo/scripts/freesound_download_resample/freesound_requirements.txt new file mode 100644 index 0000000000..127e50992a --- /dev/null +++ b/models/tts/NeMo/scripts/freesound_download_resample/freesound_requirements.txt @@ -0,0 +1,6 @@ +git+git://github.com/MTG/freesound-python.git +requests +requests_oauthlib +joblib +librosa +sox \ No newline at end of file diff --git a/models/tts/NeMo/scripts/freesound_download_resample/freesound_resample.py b/models/tts/NeMo/scripts/freesound_download_resample/freesound_resample.py new file mode 100644 index 0000000000..79a01c0946 --- /dev/null +++ b/models/tts/NeMo/scripts/freesound_download_resample/freesound_resample.py @@ -0,0 +1,119 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import glob +import os +import time + +import librosa +import sox +from joblib import Parallel, delayed + + +def resample_file(resampled_dir, filepath, ext, sample_rate): + """ + Resample an audio file to 16kHZ and transform to monochannel + Remove incompatible files. + + Args: + resampled_dir: Directory of transformed files. + filepath: Filepath of Audio + ext: File type e.g. "wav", "flac" + + Returns: + + """ + head, filename = os.path.split(filepath) + _, clsname = os.path.split(head) + + filename, _ = os.path.splitext(filename) + + new_dir = os.path.join(resampled_dir, clsname) + if not os.path.exists(new_dir): + os.makedirs(new_dir) + + new_path = os.path.join(new_dir, filename + f'.{ext}') + + # check if the resampled data exists. + if os.path.exists(new_path): + print(f"Resampled file {filepath} exists. Skip it.") + return None + + transform = sox.Transformer() + transform.set_output_format(file_type='wav') + transform.convert(samplerate=sample_rate, n_channels=1) + + try: + transform.build(filepath, new_path) + print(f"Finished converting file {filepath}.") + + return None + + except sox.core.SoxError as e: + + try: + # Check if the file is readable + librosa.load(filepath) + + # if it is, force input format and try again + transform.set_input_format(file_type=ext) + transform.build(filepath, new_path) + return None + + except Exception: + return filepath + + +def main(): + start = time.time() + parser = argparse.ArgumentParser(description='Freesound data resample') + parser.add_argument("--data_dir", required=True, default=None, type=str) + parser.add_argument('--resampled_dir', required=True, default=None, type=str) + parser.add_argument('--sample_rate', default=16000, type=int) + args = parser.parse_args() + + data_dir = args.data_dir + resampled_dir = args.resampled_dir + sample_rate = args.sample_rate + + wav_files = sorted(glob.glob(os.path.join(data_dir, '*/*.wav'))) + flac_files = sorted(glob.glob(os.path.join(data_dir, '*/*.flac'))) + + with Parallel(n_jobs=-1, verbose=10) as parallel: + wav_files_failed = parallel( + delayed(resample_file)(resampled_dir, filepath, ext='wav', sample_rate=sample_rate) + for filepath in wav_files + ) + + flac_files_failed = parallel( + delayed(resample_file)(resampled_dir, filepath, ext='flac', sample_rate=sample_rate) + for filepath in flac_files + ) + + with open('dataset_conversion_logs.txt', 'w') as f: + for file in wav_files_failed: + if file is not None: + f.write(f"{file}\n") + + for file in flac_files_failed: + if file is not None: + f.write(f"{file}\n") + + end = time.time() + print(f'Resample data in {data_dir} and save to {resampled_dir} takes {end-start} seconds.') + + +if __name__ == '__main__': + + main() diff --git a/models/tts/NeMo/scripts/nemo_legacy_import/asr_checkpoint_port.py b/models/tts/NeMo/scripts/nemo_legacy_import/asr_checkpoint_port.py new file mode 100644 index 0000000000..f9b7387ffb --- /dev/null +++ b/models/tts/NeMo/scripts/nemo_legacy_import/asr_checkpoint_port.py @@ -0,0 +1,72 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +"""This script converts old Jasper/QuartzNet models from NeMo 0.11.* to NeMo v1.0.0* +""" + +import argparse + +import torch +from omegaconf import DictConfig +from ruamel.yaml import YAML + +import nemo.collections.asr as nemo_asr +from nemo.utils import logging + + +def get_parser(): + parser = argparse.ArgumentParser(description="Converts old Jasper/QuartzNet models to NeMo v1.0beta") + parser.add_argument("--config_path", default=None, required=True, help="Path to model config (NeMo v1.0beta)") + parser.add_argument("--encoder_ckpt", default=None, required=True, help="Encoder checkpoint path") + parser.add_argument("--decoder_ckpt", default=None, required=True, help="Decoder checkpoint path") + parser.add_argument("--output_path", default=None, required=True, help="Output checkpoint path (should be .nemo)") + parser.add_argument( + "--model_type", + default='asr', + type=str, + choices=['asr', 'speech_label', 'speaker'], + help="Type of decoder used by the model.", + ) + + return parser + + +def main(config_path, encoder_ckpt, decoder_ckpt, output_path, model_type): + + yaml = YAML(typ='safe') + with open(config_path) as f: + params = yaml.load(f) + + model = None + if model_type == 'asr': + logging.info("Creating ASR NeMo 1.0 model") + model = nemo_asr.models.EncDecCTCModel(cfg=DictConfig(params['model'])) + elif model_type == 'speech_label': + logging.info("Creating speech label NeMo 1.0 model") + model = nemo_asr.models.EncDecClassificationModel(cfg=DictConfig(params['model'])) + else: + logging.info("Creating Speaker Recognition NeMo 1.0 model") + model = nemo_asr.models.EncDecSpeakerLabelModel(cfg=DictConfig(params['model'])) + + model.encoder.load_state_dict(torch.load(encoder_ckpt)) + model.decoder.load_state_dict(torch.load(decoder_ckpt)) + logging.info("Succesfully ported old checkpoint") + + model.save_to(output_path) + logging.info("new model saved at {}".format(output_path)) + + +if __name__ == "__main__": + args = get_parser().parse_args() + main(args.config_path, args.encoder_ckpt, args.decoder_ckpt, args.output_path, args.model_type) diff --git a/models/tts/NeMo/scripts/neural_machine_translation/filter_langs_nmt.py b/models/tts/NeMo/scripts/neural_machine_translation/filter_langs_nmt.py new file mode 100644 index 0000000000..73634063d2 --- /dev/null +++ b/models/tts/NeMo/scripts/neural_machine_translation/filter_langs_nmt.py @@ -0,0 +1,403 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import logging +import multiprocessing as mp +import re +import shutil +import warnings +from pathlib import Path +from time import sleep + +import fasttext +from tqdm import tqdm + +""" +Usage: +python filter_by_language.py --input-src train.en \ + --input-tgt train.de \ + --output-src train_lang_filtered.en \ + --output-tgt train_lang_filtered.de \ + --source-lang en \ + --target-lang de \ + --removed-src train_garbage.en \ + --removed-tgt train_garbage.de \ + --fasttext-model lid.176.bin +""" + + +logging.basicConfig(level=logging.INFO) + + +def get_args(): + parser = argparse.ArgumentParser( + description="It is a script for verifying language in machine translation data sets. If the script is used on " + "a parallel corpus, it verifies both a source and a target language. If number of jobs `--num-jobs` is bigger " + "than 1 than lines in an input file (or files if parallel corpus is checked) split equally between workers. " + "If `num_jobs > 1` is used, the best performance is achieved if dataset is shuffled and lines with different " + "lengths are distributed evenly in the input file. Filtered data is stored into `output_src`[, `--output-tgt`]" + " and removed lines are put into `removed_src`[, `--removed-tgt`] files. If language cannot be detected " + "(e.g. date), the line is removed. Working time on en-de wikimatrix (6.23M pairs: 700 MB German and 625 MB " + "English) from wmt20 on machine with 20 CPU cores: less than 1 minute." + ) + parser.add_argument( + "--input-src", + "-s", + help="Path to the input file which has to contain text in language `source_lang`.", + required=True, + type=Path, + ) + parser.add_argument( + "--input-tgt", + "-t", + help="Path to the input file which has to contain text in language `target_lang`. If not provided, data is " + "processed as monolingual.", + type=Path, + ) + parser.add_argument( + "--output-src", + "-S", + help="Path to the file where filtered `input_src` will be saved.", + required=True, + type=Path, + ) + parser.add_argument( + "--output-tgt", "-T", help="Path to the output target file", type=Path, + ) + parser.add_argument( + "--source-lang", + "-l", + required=True, + help="Input language. For options see https://fasttext.cc/docs/en/language-identification.html.", + ) + parser.add_argument( + "--target-lang", + "-L", + help="Output language. For options see https://fasttext.cc/docs/en/language-identification.html.", + ) + parser.add_argument( + "--removed-src", "-r", required=True, help="Path to file where removed source lines will be saved", type=Path, + ) + parser.add_argument( + "--removed-tgt", "-R", help="Path to file where removed target lines will be saved", type=Path, + ) + parser.add_argument( + "--num-jobs", + "-j", + type=int, + help="Number of jobs. By default, the number of jobs is equal to the number of CPU cores.", + ) + parser.add_argument( + "--fasttext-model", + "-m", + help="Path to fasttext model. The description and download links are here " + "https://fasttext.cc/docs/en/language-identification.html", + type=Path, + ) + args = parser.parse_args() + if not ( + args.output_tgt is None + and args.input_tgt is None + and args.target_lang is None + and args.removed_tgt is None + or args.output_tgt is not None + and args.input_tgt is not None + and args.target_lang is not None + and args.removed_tgt is not None + ): + raise ValueError( + f"Arguments `input_tgt`, `output_tgt`, `target_lang`, `removed_tgt` have to be either `None` " + f"simultaneously or not `None` simultaneously. Given " + f"input_tgt={args.input_tgt}, output_tgt={args.output_tgt}, target_lang={args.target_lang}, " + f"removed_tgt={args.removed_tgt}" + ) + args.input_src = args.input_src.expanduser() + if args.input_tgt is not None: + args.input_tgt = args.input_tgt.expanduser() + args.output_src = args.output_src.expanduser() + if args.output_tgt is not None: + args.output_tgt = args.output_tgt.expanduser() + args.removed_src = args.removed_src.expanduser() + if args.removed_tgt is not None: + args.removed_tgt = args.removed_tgt.expanduser() + args.fasttext_model = args.fasttext_model.expanduser() + return args + + +def get_lang(line, fasttext_model): + labels, _ = fasttext_model.predict(line, k=1) + lang = labels[0].split('__')[-1] + return lang + + +def get_edges_in_1_file(fn, num_parts): + num_lines = 0 + edges = [0] + with open(fn) as f: + i = 0 + for l in f: + i += len(l.encode('utf-8')) + edges.append(i) + num_lines += 1 + return [edges[int(i * num_lines / num_parts)] for i in range(num_parts)] + [edges[-1]], num_lines + + +def get_edges_and_num_lines(src_fn, tgt_fn, num_parts): + src_edges, src_num_lines = get_edges_in_1_file(src_fn, num_parts) + assert num_parts + 1 == len(src_edges) + src_edges = [(src_edges[i], src_edges[i + 1]) for i in range(len(src_edges) - 1)] + if tgt_fn is not None: + tgt_edges, tgt_num_lines = get_edges_in_1_file(tgt_fn, num_parts) + tgt_edges = [(tgt_edges[i], tgt_edges[i + 1]) for i in range(len(tgt_edges) - 1)] + if tgt_num_lines != src_num_lines: + raise ValueError( + f"Source {repr(src_fn)} and target {repr(tgt_fn)} files have different number of lines " + f"{src_num_lines} and {tgt_num_lines} correspondingly." + ) + + else: + tgt_edges = [None] * num_parts + assert len(src_edges) == num_parts + return src_edges, tgt_edges, src_num_lines + + +def filter_pairs( + src_edges, + tgt_edges, + input_src, + input_tgt, + filtered_dir_src, + filtered_dir_tgt, + removed_dir_src, + removed_dir_tgt, + source_lang, + target_lang, + fasttext_model, + rank, +): + global counter + fasttext_model = fasttext.load_model(str(fasttext_model)) + output_src = filtered_dir_src / Path(f"rank{rank}") + output_src_removed = removed_dir_src / Path(f"rank{rank}") + output_tgt = filtered_dir_tgt / Path(f"rank{rank}") + output_tgt_removed = removed_dir_tgt / Path(f"rank{rank}") + with open(input_src) as in_src, open(input_tgt) as in_tgt, open(output_src, 'w') as out_src, open( + output_tgt, 'w' + ) as out_tgt, open(output_src_removed, 'w') as out_r_src, open(output_tgt_removed, 'w') as out_r_tgt: + in_src.seek(src_edges[0]) + in_tgt.seek(tgt_edges[0]) + src_l, tgt_l, i = in_src.readline(), in_tgt.readline(), 0 + if in_src.tell() > src_edges[1] or in_tgt.tell() > tgt_edges[1]: + return + while src_l and tgt_l: + with counter.get_lock(): + counter.value += 1 + src_l = src_l.strip() + tgt_l = tgt_l.strip() + src_lang = get_lang(src_l, fasttext_model) + if src_lang is not None: + tgt_lang = get_lang(tgt_l, fasttext_model) + if src_lang is None or tgt_lang is None or src_lang != source_lang or tgt_lang != target_lang: + out_r_src.write(src_l + '\n') + out_r_tgt.write(tgt_l + '\n') + else: + out_src.write(src_l + '\n') + out_tgt.write(tgt_l + '\n') + if in_src.tell() >= src_edges[1]: + if in_tgt.tell() < tgt_edges[1]: + raise ValueError( + f"Edges of target and source has to be reached simultaneously, whereas " + f"in_src.tell()={in_src.tell()}, in_tgt.tell()={in_tgt.tell()}, " + f"src_edges[1]={src_edges[1]}, tgt_edges[1]={tgt_edges[1]}." + ) + break + if in_tgt.tell() >= tgt_edges[1]: + raise ValueError( + f"Edges of target and source has to be reached simultaneously, whereas " + f"in_src.tell()={in_src.tell()}, in_tgt.tell()={in_tgt.tell()}, " + f"src_edges[1]={src_edges[1]}, tgt_edges[1]={tgt_edges[1]}." + ) + src_l, tgt_l, i = in_src.readline(), in_tgt.readline(), i + 1 + with counter.get_lock(): + counter.value += 1 + + +def filter_singles( + src_edges, input_src, filtered_dir_src, removed_dir_src, source_lang, fasttext_model, rank, +): + logging.debug("filter singles") + global counter + fasttext_model = fasttext.load_model(str(fasttext_model)) + output_src = filtered_dir_src / Path(f"rank{rank}") + output_src_removed = removed_dir_src / Path(f"rank{rank}") + with open(input_src) as in_f, open(output_src, 'w') as out_f, open(output_src_removed, 'w') as out_r_f: + in_f.seek(src_edges[0]) + i, line = 0, in_f.readline() + if in_f.tell() > src_edges[1]: + return + while line: + with counter.get_lock(): + counter.value += 1 + line = line.strip() + in_lang = get_lang(line, fasttext_model) + if in_lang is None or in_lang != source_lang: + out_r_f.write(line + '\n') + else: + out_f.write(line + '\n') + if in_f.tell() >= src_edges[1]: + break + i, line = i + 1, in_f.readline() + with counter.get_lock(): + counter.value += 1 + + +def filter_by_lang(args): + ( + src_edges, + tgt_edges, + input_src, + input_tgt, + filtered_dir_src, + filtered_dir_tgt, + removed_dir_src, + removed_dir_tgt, + source_lang, + target_lang, + fasttext_model, + rank, + ) = args + logging.debug(f"filter by lang input_tgt: {input_tgt}") + if input_tgt is None: + if tgt_edges is not None: + warnings.warn("If input target is not provided `tgt_edges` argument is expected to be `None`") + filter_singles( + src_edges, input_src, filtered_dir_src, removed_dir_src, source_lang, fasttext_model, rank, + ) + else: + filter_pairs( + src_edges, + tgt_edges, + input_src, + input_tgt, + filtered_dir_src, + filtered_dir_tgt, + removed_dir_src, + removed_dir_tgt, + source_lang, + target_lang, + fasttext_model, + rank, + ) + + +def _cat_results(out_file, tmp_dir): + file_name_pattern = re.compile(r"/rank([1-9][0-9]*)|0$") + with out_file.open('w') as out_f: + for f in sorted(tmp_dir.iterdir()): + if not f.is_file(): + warnings.warn(f"Unexpected not file {f}") + elif not file_name_pattern.search(str(f)): + warnings.warn(f"Unexpected file {f}") + else: + with f.open('r') as in_f: + for l in in_f: + out_f.write(l) + + +def cat_results(out_files, tmp_dirs): + for o_f, t_d in zip(out_files, tmp_dirs): + if o_f is None or t_d is None: + if o_f is not None or t_d is not None: + warnings.warn( + f"Output file and tmp directory are expected to be `None` simultaneously whereas tmp directory " + f"is {t_d} and output file is {o_f}." + ) + else: + _cat_results(o_f, t_d) + + +counter = None + + +def init(args): + global counter + counter = args + + +def main(): + args = get_args() + tmp_dir = Path("tmp") + i = 0 + while tmp_dir.exists(): + tmp_dir = Path("tmp" + str(i)) + i += 1 + tmp_filtered = tmp_dir / Path("filtered") + tmp_filtered_src = tmp_filtered / Path("src") + tmp_filtered_src.mkdir(parents=True, exist_ok=True) + if args.input_tgt is None: + tmp_filtered_tgt = None + else: + tmp_filtered_tgt = tmp_filtered / Path("tgt") + tmp_filtered_tgt.mkdir(parents=True, exist_ok=True) + tmp_removed = tmp_dir / Path("removed") + tmp_removed_src = tmp_removed / Path("src") + tmp_removed_src.mkdir(parents=True, exist_ok=True) + if args.input_tgt is None: + tmp_removed_tgt = None + else: + tmp_removed_tgt = tmp_removed / Path("tgt") + tmp_removed_tgt.mkdir(parents=True, exist_ok=True) + num_jobs = mp.cpu_count() if args.num_jobs is None else args.num_jobs + src_edges, tgt_edges, num_lines = get_edges_and_num_lines(args.input_src, args.input_tgt, num_jobs) + global counter + counter = mp.Value('i', 0) + t = tqdm(total=num_lines, desc="processed lines / total number of lines") + with mp.Pool(num_jobs, initializer=init, initargs=(counter,)) as pool: + async_result = pool.map_async( + filter_by_lang, + [ + ( + se, + te, + args.input_src, + args.input_tgt, + tmp_filtered_src, + tmp_filtered_tgt, + tmp_removed_src, + tmp_removed_tgt, + args.source_lang, + args.target_lang, + args.fasttext_model, + rank, + ) + for rank, (se, te) in enumerate(zip(src_edges, tgt_edges)) + ], + ) + while not async_result.ready(): + t.update(counter.value) + with counter.get_lock(): + counter.value = 0 + sleep(0.1) + t.update(counter.value) + + cat_results( + [args.output_src, args.output_tgt, args.removed_src, args.removed_tgt], + [tmp_filtered_src, tmp_filtered_tgt, tmp_removed_src, tmp_removed_tgt], + ) + shutil.rmtree(tmp_dir) + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/speaker_recognition/rttm_to_manifest.py b/models/tts/NeMo/scripts/speaker_recognition/rttm_to_manifest.py new file mode 100644 index 0000000000..3d7b772b05 --- /dev/null +++ b/models/tts/NeMo/scripts/speaker_recognition/rttm_to_manifest.py @@ -0,0 +1,40 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse + +from nemo.collections.asr.parts.utils.speaker_utils import write_rttm2manifest +from nemo.utils import logging + +""" +This file converts vad outputs to manifest file for speaker diarization purposes +present in vad output directory. +every vad line consists of start_time, end_time , speech/non-speech +""" + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--paths2rttm_files", help="path to vad output rttm-like files. Could be a list or a text file", required=True + ) + parser.add_argument( + "--paths2audio_files", + help="path to audio files that vad was computed. Could be a list or a text file", + required=True, + ) + parser.add_argument("--manifest_file", help="output manifest file name", type=str, required=True) + args = parser.parse_args() + + write_rttm2manifest(args.paths2audio_files, args.paths2rttm_files, args.manifest_file) + logging.info("wrote {} file from vad output files present in {}".format(args.manifest_file, args.paths2rttm_files)) diff --git a/models/tts/NeMo/scripts/speaker_recognition/scp_to_manifest.py b/models/tts/NeMo/scripts/speaker_recognition/scp_to_manifest.py new file mode 100644 index 0000000000..1261db9fd5 --- /dev/null +++ b/models/tts/NeMo/scripts/speaker_recognition/scp_to_manifest.py @@ -0,0 +1,96 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import json +import logging +import os + +import librosa as l +from sklearn.model_selection import StratifiedShuffleSplit +from tqdm import tqdm + + +""" +This scipt converts a scp file where each line contains + +to a manifest json file. +Args: +--scp: scp file name +--id: index of speaker label in filename present in scp file that is separated by '/' +--out: output manifest file name +--split: True / False if you would want to split the manifest file for training purposes + you may not need this for test set. output file names is _.json + Defaults to False +""" + + +def write_file(name, lines, idx): + with open(name, 'w') as fout: + for i in idx: + dic = lines[i] + json.dump(dic, fout) + fout.write('\n') + logging.info("wrote", name) + + +def main(scp, id, out, split=False): + if os.path.exists(out): + os.remove(out) + scp_file = open(scp, 'r').readlines() + + lines = [] + speakers = [] + with open(out, 'w') as outfile: + for line in tqdm(scp_file): + line = line.strip() + y, sr = l.load(line, sr=None) + dur = l.get_duration(y=y, sr=sr) + speaker = line.split('/')[id] + speaker = list(speaker) + speaker = ''.join(speaker) + speakers.append(speaker) + meta = {"audio_filepath": line, "offset": 0, "duration": float(dur), "label": speaker} + lines.append(meta) + json.dump(meta, outfile) + outfile.write("\n") + + path = os.path.dirname(out) + if split: + sss = StratifiedShuffleSplit(n_splits=1, test_size=0.1, random_state=42) + for train_idx, test_idx in sss.split(speakers, speakers): + logging.info(len(train_idx)) + + out = os.path.join(path, 'train.json') + write_file(out, lines, train_idx) + out = os.path.join(path, 'dev.json') + write_file(out, lines, test_idx) + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("--scp", help="scp file name", type=str, required=True) + parser.add_argument( + "--id", help="field num seperated by '/' to be considered as speaker label", type=int, required=True + ) + parser.add_argument("--out", help="manifest_file name", type=str, required=True) + parser.add_argument( + "--split", + help="bool if you would want to split the manifest file for training purposes", + required=False, + action='store_true', + ) + args = parser.parse_args() + + main(args.scp, args.id, args.out, args.split) diff --git a/models/tts/NeMo/scripts/speech_recognition/convert_to_tarred_audio_dataset.py b/models/tts/NeMo/scripts/speech_recognition/convert_to_tarred_audio_dataset.py new file mode 100644 index 0000000000..4e13c0536c --- /dev/null +++ b/models/tts/NeMo/scripts/speech_recognition/convert_to_tarred_audio_dataset.py @@ -0,0 +1,601 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. +""" +# This script converts an existing audio dataset with a manifest to +# a tarred and sharded audio dataset that can be read by the +# TarredAudioToTextDataLayer. + +# Please make sure your audio_filepath DOES NOT CONTAIN '-sub'! +# Because we will use it to handle files which have duplicate filenames but with different offsets +# (see function create_shard for details) + +# Usage: +1) Creating a new tarfile dataset + +python convert_to_tarred_audio_dataset.py \ + --manifest_path= \ + --target_dir= \ + --num_shards= + --max_duration= \ + --min_duration= \ + --shuffle --shuffle_seed=1 + + +2) Concatenating more tarfiles to a pre-existing tarred dataset + +python convert_to_tarred_audio_dataset.py \ + --manifest_path= \ + --metadata_path= \ + --target_dir= \ + --max_duration= \ + --min_duration= \ + --shuffle --shuffle_seed=1 \ + --concat_manifest_paths \ + + +3) Writing an empty metadata file + +python convert_to_tarred_audio_dataset.py \ + --target_dir= \ + # any other optional argument + --num_shards=8 \ + --max_duration=16.7 \ + --min_duration=0.01 \ + --shuffle \ + --shuffle_seed=1 \ + --write_metadata + +""" +import argparse +import copy +import json +import os +import random +import tarfile +from dataclasses import dataclass, field +from datetime import datetime +from typing import Any, List, Optional + +from joblib import Parallel, delayed +from omegaconf import DictConfig, OmegaConf, open_dict + +parser = argparse.ArgumentParser( + description="Convert an existing ASR dataset to tarballs compatible with TarredAudioToTextDataLayer." +) +parser.add_argument( + "--manifest_path", default=None, type=str, required=False, help="Path to the existing dataset's manifest." +) + +parser.add_argument( + '--concat_manifest_paths', + nargs='+', + default=None, + type=str, + required=False, + help="Path to the additional dataset's manifests that will be concatenated with base dataset.", +) + +# Optional arguments +parser.add_argument( + "--target_dir", + default='./tarred', + type=str, + help="Target directory for resulting tarballs and manifest. Defaults to `./tarred`. Creates the path if necessary.", +) + +parser.add_argument( + "--metadata_path", required=False, default=None, type=str, help="Path to metadata file for the dataset.", +) + +parser.add_argument( + "--num_shards", + default=-1, + type=int, + help="Number of shards (tarballs) to create. Used for partitioning data among workers.", +) +parser.add_argument( + '--max_duration', + default=None, + type=float, + help='Maximum duration of audio clip in the dataset. By default, it is None and will not filter files.', +) +parser.add_argument( + '--min_duration', + default=None, + type=float, + help='Minimum duration of audio clip in the dataset. By default, it is None and will not filter files.', +) +parser.add_argument( + "--shuffle", + action='store_true', + help="Whether or not to randomly shuffle the samples in the manifest before tarring/sharding.", +) +parser.add_argument("--shuffle_seed", type=int, help="Random seed for use if shuffling is enabled.") +parser.add_argument( + '--write_metadata', + action='store_true', + help=( + "Flag to write a blank metadata with the current call config. " + "Note that the metadata will not contain the number of shards, " + "and it must be filled out by the user." + ), +) +parser.add_argument('--workers', type=int, default=1, help='Number of worker processes') +args = parser.parse_args() + + +@dataclass +class ASRTarredDatasetConfig: + num_shards: int = -1 + shuffle: bool = False + max_duration: Optional[float] = None + min_duration: Optional[float] = None + shuffle_seed: Optional[int] = None + + +@dataclass +class ASRTarredDatasetMetadata: + created_datetime: Optional[str] = None + version: int = 0 + num_samples_per_shard: Optional[int] = None + is_concatenated_manifest: bool = False + + dataset_config: Optional[ASRTarredDatasetConfig] = ASRTarredDatasetConfig() + history: Optional[List[Any]] = field(default_factory=lambda: []) + + def __post_init__(self): + self.created_datetime = self.get_current_datetime() + + def get_current_datetime(self): + return datetime.now().strftime("%m-%d-%Y %H-%M-%S") + + @classmethod + def from_config(cls, config: DictConfig): + obj = cls() + obj.__dict__.update(**config) + return obj + + @classmethod + def from_file(cls, filepath: str): + config = OmegaConf.load(filepath) + return ASRTarredDatasetMetadata.from_config(config=config) + + +class ASRTarredDatasetBuilder: + """ + Helper class that constructs a tarred dataset from scratch, or concatenates tarred datasets + together and constructs manifests for them. + """ + + def __init__(self): + self.config = None + + def configure(self, config: ASRTarredDatasetConfig): + """ + Sets the config generated from command line overrides. + + Args: + config: ASRTarredDatasetConfig dataclass object. + """ + self.config = config # type: ASRTarredDatasetConfig + + if self.config.num_shards < 0: + raise ValueError("`num_shards` must be > 0. Please fill in the metadata information correctly.") + + def create_new_dataset(self, manifest_path: str, target_dir: str = "./tarred/", num_workers: int = 0): + """ + Creates a new tarred dataset from a given manifest file. + + Args: + manifest_path: Path to the original ASR manifest. + target_dir: Output directory. + num_workers: Integer denoting number of parallel worker processes which will write tarfiles. + Defaults to 1 - which denotes sequential worker process. + + Output: + Writes tarfiles, along with the tarred dataset compatible manifest file. + Also preserves a record of the metadata used to construct this tarred dataset. + """ + if self.config is None: + raise ValueError("Config has not been set. Please call `configure(config: ASRTarredDatasetConfig)`") + + if manifest_path is None: + raise FileNotFoundError("Manifest filepath cannot be None !") + + config = self.config # type: ASRTarredDatasetConfig + + if not os.path.exists(target_dir): + os.makedirs(target_dir) + + # Read the existing manifest + entries, filtered_entries, filtered_duration = self._read_manifest(manifest_path, config) + + if len(filtered_entries) > 0: + print(f"Filtered {len(filtered_entries)} files which amounts to {filtered_duration} seconds of audio.") + + if config.shuffle: + random.seed(config.shuffle_seed) + print("Shuffling...") + random.shuffle(entries) + + # Create shards and updated manifest entries + print(f"Number of samples added : {len(entries)}") + print(f"Remainder: {len(entries) % config.num_shards}") + + start_indices = [] + end_indices = [] + # Build indices + for i in range(config.num_shards): + start_idx = (len(entries) // config.num_shards) * i + end_idx = start_idx + (len(entries) // config.num_shards) + print(f"Shard {i} has entries {start_idx} ~ {end_idx}") + if i == config.num_shards - 1: + # We discard in order to have the same number of entries per shard. + print(f"Have {len(entries) - end_idx} entries left over that will be discarded.") + + start_indices.append(start_idx) + end_indices.append(end_idx) + + with Parallel(n_jobs=num_workers, verbose=config.num_shards) as parallel: + # Call parallel tarfile construction + new_entries_list = parallel( + delayed(self._create_shard)(entries[start_idx:end_idx], target_dir, i) + for i, (start_idx, end_idx) in enumerate(zip(start_indices, end_indices)) + ) + + # Flatten the list of list of entries to a list of entries + new_entries = [sample for manifest in new_entries_list for sample in manifest] + del new_entries_list + + print("Total number of files in manifest :", len(new_entries)) + + # Write manifest + new_manifest_path = os.path.join(target_dir, 'tarred_audio_manifest.json') + with open(new_manifest_path, 'w') as m2: + for entry in new_entries: + json.dump(entry, m2) + m2.write('\n') + + # Write metadata (default metadata for new datasets) + new_metadata_path = os.path.join(target_dir, 'metadata.yaml') + metadata = ASRTarredDatasetMetadata() + + # Update metadata + metadata.dataset_config = config + metadata.num_samples_per_shard = len(new_entries) // config.num_shards + + # Write metadata + metadata_yaml = OmegaConf.structured(metadata) + OmegaConf.save(metadata_yaml, new_metadata_path, resolve=True) + + def create_concatenated_dataset( + self, + base_manifest_path: str, + manifest_paths: List[str], + metadata: ASRTarredDatasetMetadata, + target_dir: str = "./tarred_concatenated/", + num_workers: int = 1, + ): + """ + Creates new tarfiles in order to create a concatenated dataset, whose manifest contains the data for + both the original dataset as well as the new data submitted in manifest paths. + + Args: + base_manifest_path: Path to the manifest file which contains the information for the original + tarred dataset (with flattened paths). + manifest_paths: List of one or more paths to manifest files that will be concatenated with above + base tarred dataset. + metadata: ASRTarredDatasetMetadata dataclass instance with overrides from command line. + target_dir: Output directory + + Output: + Writes tarfiles which with indices mapping to a "concatenated" tarred dataset, + along with the tarred dataset compatible manifest file which includes information + about all the datasets that comprise the concatenated dataset. + + Also preserves a record of the metadata used to construct this tarred dataset. + """ + if not os.path.exists(target_dir): + os.makedirs(target_dir) + + if base_manifest_path is None: + raise FileNotFoundError("Base manifest filepath cannot be None !") + + if manifest_paths is None or len(manifest_paths) == 0: + raise FileNotFoundError("List of additional manifest filepaths cannot be None !") + + config = ASRTarredDatasetConfig(**(metadata.dataset_config)) + + # Read the existing manifest (no filtering here) + base_entries, _, _ = self._read_manifest(base_manifest_path, config) + print(f"Read base manifest containing {len(base_entries)} samples.") + + # Precompute number of samples per shard + if metadata.num_samples_per_shard is None: + num_samples_per_shard = len(base_entries) // config.num_shards + else: + num_samples_per_shard = metadata.num_samples_per_shard + + print("Number of samples per shard :", num_samples_per_shard) + + # Compute min and max duration and update config (if no metadata passed) + print(f"Selected max duration : {config.max_duration}") + print(f"Selected min duration : {config.min_duration}") + + entries = [] + for new_manifest_idx in range(len(manifest_paths)): + new_entries, filtered_new_entries, filtered_duration = self._read_manifest( + manifest_paths[new_manifest_idx], config + ) + + if len(filtered_new_entries) > 0: + print( + f"Filtered {len(filtered_new_entries)} files which amounts to {filtered_duration:0.2f}" + f" seconds of audio from manifest {manifest_paths[new_manifest_idx]}." + ) + + entries.extend(new_entries) + + if config.shuffle: + random.seed(config.shuffle_seed) + print("Shuffling...") + random.shuffle(entries) + + # Drop last section of samples that cannot be added onto a chunk + drop_count = len(entries) % num_samples_per_shard + total_new_entries = len(entries) + entries = entries[:-drop_count] + + print( + f"Dropping {drop_count} samples from total new samples {total_new_entries} since they cannot " + f"be added into a uniformly sized chunk." + ) + + # Create shards and updated manifest entries + num_added_shards = len(entries) // num_samples_per_shard + + print(f"Number of samples in base dataset : {len(base_entries)}") + print(f"Number of samples in additional datasets : {len(entries)}") + print(f"Number of added shards : {num_added_shards}") + print(f"Remainder: {len(entries) % num_samples_per_shard}") + + start_indices = [] + end_indices = [] + shard_indices = [] + for i in range(num_added_shards): + start_idx = (len(entries) // num_added_shards) * i + end_idx = start_idx + (len(entries) // num_added_shards) + shard_idx = i + config.num_shards + print(f"Shard {shard_idx} has entries {start_idx + len(base_entries)} ~ {end_idx + len(base_entries)}") + + start_indices.append(start_idx) + end_indices.append(end_idx) + shard_indices.append(shard_idx) + + with Parallel(n_jobs=num_workers, verbose=num_added_shards) as parallel: + # Call parallel tarfile construction + new_entries_list = parallel( + delayed(self._create_shard)(entries[start_idx:end_idx], target_dir, shard_idx) + for i, (start_idx, end_idx, shard_idx) in enumerate(zip(start_indices, end_indices, shard_indices)) + ) + + # Flatten the list of list of entries to a list of entries + new_entries = [sample for manifest in new_entries_list for sample in manifest] + del new_entries_list + + # Write manifest + if metadata is None: + new_version = 1 # start with `1`, where `0` indicates the base manifest + dataset + else: + new_version = metadata.version + 1 + + print("Total number of files in manifest :", len(base_entries) + len(new_entries)) + + new_manifest_path = os.path.join(target_dir, f'tarred_audio_manifest_version_{new_version}.json') + with open(new_manifest_path, 'w') as m2: + # First write all the entries of base manifest + for entry in base_entries: + json.dump(entry, m2) + m2.write('\n') + + # Finally write the new entries + for entry in new_entries: + json.dump(entry, m2) + m2.write('\n') + + # Preserve historical metadata + base_metadata = metadata + + # Write metadata (updated metadata for concatenated datasets) + new_metadata_path = os.path.join(target_dir, f'metadata_version_{new_version}.yaml') + metadata = ASRTarredDatasetMetadata() + + # Update config + config.num_shards = config.num_shards + num_added_shards + + # Update metadata + metadata.version = new_version + metadata.dataset_config = config + metadata.num_samples_per_shard = num_samples_per_shard + metadata.is_concatenated_manifest = True + metadata.created_datetime = metadata.get_current_datetime() + + # Attach history + current_metadata = OmegaConf.structured(base_metadata.history) + metadata.history = current_metadata + + # Write metadata + metadata_yaml = OmegaConf.structured(metadata) + OmegaConf.save(metadata_yaml, new_metadata_path, resolve=True) + + def _read_manifest(self, manifest_path: str, config: ASRTarredDatasetConfig): + """Read and filters data from the manifest""" + # Read the existing manifest + entries = [] + filtered_entries = [] + filtered_duration = 0.0 + with open(manifest_path, 'r') as m: + for line in m: + entry = json.loads(line) + if (config.max_duration is None or entry['duration'] < config.max_duration) and ( + config.min_duration is None or entry['duration'] > config.min_duration + ): + entries.append(entry) + else: + filtered_entries.append(entry) + filtered_duration += entry['duration'] + + return entries, filtered_entries, filtered_duration + + def _create_shard(self, entries, target_dir, shard_id): + """Creates a tarball containing the audio files from `entries`. + """ + new_entries = [] + tar = tarfile.open(os.path.join(target_dir, f'audio_{shard_id}.tar'), mode='w', dereference=True) + + count = dict() + for entry in entries: + # We squash the filename since we do not preserve directory structure of audio files in the tarball. + base, ext = os.path.splitext(entry['audio_filepath']) + base = base.replace('/', '_') + # Need the following replacement as long as WebDataset splits on first period + base = base.replace('.', '_') + squashed_filename = f'{base}{ext}' + if squashed_filename not in count: + tar.add(entry['audio_filepath'], arcname=squashed_filename) + + if 'label' in entry: + base, ext = os.path.splitext(squashed_filename) + # no suffix if it's single sample or starting sub parts, -sub1 for the second subpart -sub2 -sub3 ,etc. + if squashed_filename not in count: + to_write = squashed_filename + count[squashed_filename] = 1 + else: + to_write = base + "-sub" + str(count[squashed_filename]) + ext + count[squashed_filename] += 1 + + new_entry = { + 'audio_filepath': to_write, + 'duration': entry['duration'], + 'text': entry['text'], + 'label': entry['label'], + 'offset': entry['offset'], + 'shard_id': shard_id, # Keep shard ID for recordkeeping + } + else: + count[squashed_filename] = 1 + new_entry = { + 'audio_filepath': squashed_filename, + 'duration': entry['duration'], + 'text': entry['text'], + 'shard_id': shard_id, # Keep shard ID for recordkeeping + } + + new_entries.append(new_entry) + + tar.close() + return new_entries + + @classmethod + def setup_history(cls, base_metadata: ASRTarredDatasetMetadata, history: List[Any]): + if 'history' in base_metadata.keys(): + for history_val in base_metadata.history: + cls.setup_history(history_val, history) + + if base_metadata is not None: + metadata_copy = copy.deepcopy(base_metadata) + with open_dict(metadata_copy): + metadata_copy.pop('history', None) + history.append(metadata_copy) + + +def main(): + manifest_path = args.manifest_path + concat_manifest_paths = args.concat_manifest_paths + target_dir = args.target_dir + metadata_path = args.metadata_path + num_shards = args.num_shards + max_duration = args.max_duration + min_duration = args.min_duration + shuffle = args.shuffle + seed = args.shuffle_seed if args.shuffle_seed else None + write_metadata = args.write_metadata + num_workers = args.workers + + builder = ASRTarredDatasetBuilder() + + if write_metadata: + metadata = ASRTarredDatasetMetadata() + dataset_cfg = ASRTarredDatasetConfig( + num_shards=num_shards, + shuffle=shuffle, + max_duration=max_duration, + min_duration=min_duration, + shuffle_seed=seed, + ) + metadata.dataset_config = dataset_cfg + + output_path = os.path.join(target_dir, 'default_metadata.yaml') + OmegaConf.save(metadata, output_path, resolve=True) + print(f"Default metadata written to {output_path}") + exit(0) + + if concat_manifest_paths is None or len(concat_manifest_paths) == 0: + print("Creating new tarred dataset ...") + + # Create a tarred dataset from scratch + config = ASRTarredDatasetConfig( + num_shards=num_shards, + shuffle=shuffle, + max_duration=max_duration, + min_duration=min_duration, + shuffle_seed=seed, + ) + builder.configure(config) + builder.create_new_dataset(manifest_path=manifest_path, target_dir=target_dir, num_workers=num_workers) + + else: + print("Concatenating multiple tarred datasets ...") + + # Implicitly update config from base details + if metadata_path is not None: + metadata = ASRTarredDatasetMetadata.from_file(metadata_path) + else: + raise ValueError("`metadata` yaml file path must be provided!") + + # Preserve history + history = [] + builder.setup_history(OmegaConf.structured(metadata), history) + metadata.history = history + + # Add command line overrides (everything other than num_shards) + metadata.dataset_config.max_duration = max_duration + metadata.dataset_config.min_duration = min_duration + metadata.dataset_config.shuffle = shuffle + metadata.dataset_config.shuffle_seed = seed + + builder.configure(metadata.dataset_config) + + # Concatenate a tarred dataset onto a previous one + builder.create_concatenated_dataset( + base_manifest_path=manifest_path, + manifest_paths=concat_manifest_paths, + metadata=metadata, + target_dir=target_dir, + num_workers=num_workers, + ) + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/tokenizers/process_asr_text_tokenizer.py b/models/tts/NeMo/scripts/tokenizers/process_asr_text_tokenizer.py new file mode 100644 index 0000000000..8c64507174 --- /dev/null +++ b/models/tts/NeMo/scripts/tokenizers/process_asr_text_tokenizer.py @@ -0,0 +1,282 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# USAGE: python process_asr_text_tokenizer.py --manifest= \ +# --data_root="" \ +# --vocab_size= \ +# --tokenizer=<"spe" or "wpe"> \ +# --log +# where can be: train_clean_100, train_clean_360, train_other_500 +# You can also put more than one data_set comma-separated: +# --manifest="train_clean_100,train_clean_360,train_other_500" +# or +# python process_asr_text_tokenizer.py --data_file= \ +# --data_root="" \ +# --vocab_size= \ +# --tokenizer=<"bpe" or "wpe"> \ +# --log +# where can be: train_clean_100, train_clean_360, train_other_500 +# You can also put more than one data_set comma-separated: +# --manifest="train_clean_100,train_clean_360,train_other_500" +# +# Args: +# --manifest or --data_file: If your text data lies inside of an ASR manifest file, +# then use the --manifest path. If instead the text data is inside a file with separate lines +# corresponding to different text lines, then use --data_file. +# In either case, you can add commas to concatenate different manifests or different data files. +# +# --data_root: The output directory (whose subdirectories will be created if not present) where +# the tokenizers will be placed. +# +# --vocab_size: The size of the tokenizer vocabulary. Larger vocabularies can accommodate almost entire, +# words but the decoder size of any model will grow proportionally. +# +# --tokenizer: Can be either spe or wpe . spe refers to the Google sentencepiece library tokenizer. +# wpe refers to the HuggingFace BERT Word Piece tokenizer. +# +# --no_lower_case: When this flag is passed, it will force the tokenizer to create seperate tokens for +# upper and lower case characters. By default, the script will turn all the text to lower case +# before tokenization (and if upper case characters are passed during training/inference, the +# tokenizer will emit a token equivalent to Out-Of-Vocabulary). Used primarily for the +# English language. +# +# --spe_type: The sentencepiece library has a few implementations of the tokenization technique, and +# spe_type refers to these implementations. Currently supported types are unigram, bpe, char, word. +# Defaults to bpe. +# +# --spe_character_coverage: The sentencepiece library considers how much of the original vocabulary it +# should cover in its "base set" of tokens (akin to the lower and upper case characters of the +# English language). For almost all languages with small base token sets (<1000 tokens), this +# should be kept at its default of 1.0. For languages with larger vocabularies (say Japanese, +# Mandarin, Korean etc), the suggested value is 0.9995. +# +# --spe_sample_size: If the dataset is too large, consider using a sampled dataset indicated by a +# positive integer. By default, any negative value (default = -1) will use the entire dataset. +# +# --spe_train_extremely_large_corpus: When training a sentencepiece tokenizer on very large amounts of text, +# sometimes the tokenizer will run out of memory or wont be able to process so much data on RAM. +# At some point you might receive the following error - "Input corpus too large, try with +# train_extremely_large_corpus=true". If your machine has large amounts of RAM, it might still be possible +# to build the tokenizer using the above flag. Will silently fail if it runs out of RAM. +# +# --spe_max_sentencepiece_length: Limits the maximum length that any any SentencePiece subword can be. +# Using this will change the subword tokens generated. +# +# --log: Whether the script should display log messages + + +import argparse +import json +import logging +import os + +import tokenizers + +from nemo.collections.common.tokenizers.sentencepiece_tokenizer import create_spt_model + +parser = argparse.ArgumentParser(description='Create tokenizer') +group = parser.add_mutually_exclusive_group(required=True) +group.add_argument("--manifest", default=None, type=str, help='Comma separated list of manifest files') +group.add_argument("--data_file", default=None, help='data file from which to create tokenizer model') +parser.add_argument("--data_root", required=True, default=None, type=str, help='Output directory') +parser.add_argument("--vocab_size", default=1024, type=int, help='Vocabulary size') +parser.add_argument("--tokenizer", default="wpe", choices=["spe", "wpe"], help='Type of tokenization to perform') +parser.add_argument( + "--spe_type", + default="bpe", + choices=['bpe', 'unigram', 'char', 'word'], + help='Type of the SentencePiece model. Can be `bpe`, `unigram`, `char` or `word`.' + 'Used only if --tokenizer == `spe`', +) +parser.add_argument( + '--spe_character_coverage', + type=float, + default=1.0, + help="Character coverage percentage for SentencePiece tokenization. For languages " + "with large vocabulary, should be close to 0.9995, otherwise kept as 1.0", +) +parser.add_argument( + '--spe_sample_size', + type=int, + default=-1, + help="Samples the dataset by `sample_size` if positive integer, otherwise uses whole dataset", +) +parser.add_argument('--spe_train_extremely_large_corpus', action='store_true', help='') +parser.add_argument( + '--spe_max_sentencepiece_length', + type=int, + default=-1, + help='Limit the maximum number of tokens in each SentencePiece subword. ' + 'Must be a positive integer > 0. By default places no limit on subword length.', +) +parser.add_argument('--no_lower_case', dest='lower_case', action='store_false') +parser.add_argument("--log", action='store_true') +parser.set_defaults(log=False, lower_case=True, spe_train_extremely_large_corpus=False) +args = parser.parse_args() + + +def __build_document_from_manifests( + data_root: str, manifests: str, +): + if ',' in manifests: + manifests = manifests.split(',') + else: + manifests = [manifests] + + document_dir = os.path.join(data_root, 'text_corpus') + if not os.path.exists(document_dir): + os.makedirs(document_dir) + + document_path = os.path.join(document_dir, 'document.txt') + + if os.path.exists(document_path): + logging.info('Corpus already exists at path : %s', document_path) + return document_path + + num_lines = 0 + with open(document_path, 'w') as out_writer: + for manifest in manifests: + with open(manifest, 'r') as in_reader: + for line in in_reader: + item = json.loads(line) + text = item['text'] + + out_writer.write(text + '\n') + out_writer.flush() + + num_lines += 1 + + logging.info(f"Finished extracting manifest : {manifest}") + + logging.info("Finished extracting all manifests ! Number of sentences : {}".format(num_lines)) + return document_path + + +def __process_data( + text_path: str, + dst_folder: str, + vocab_size: int, + tokenizer_type: str, + spe_type: str, + spe_character_coverage: float, + spe_train_extremely_large_corpus: bool, + spe_sample_size: int, + spe_max_sentencepiece_length: int, + lower_case: bool, +): + """ + Converts flac to wav and build manifests's json + Args: + text_path: source with text lines + dst_folder: where wav files will be stored + vocab_size: vocabular size used in encoding the text + tokenizer_type: type of tokenization to perform - wpe or spe + spe_type: type of tokenization model used for spe. + spe_character_coverage: float value between 0 and 1 (as a percentage). For languages with a vast charset, + can be < 1.0, but for all other languages, it should be set as 1.0 + spe_sample_size: int, default of -1. If positive integer is used, samples the dataset + by given sample size. + spe_train_extremely_large_corpus: bool. If dataset is too large, and user has sufficient RAM, + this flag can be set to try to trained the tokenizer. Will silently fail if it runs out of RAM. + spe_max_sentencepiece_length: Limits the maximum length of the SentencePiece subword that can be constructed. + By default, no limit is placed. + lower_case: whether to tokenize with lower case character set only (for english) + + Returns: + """ + if tokenizer_type == 'spe': + if spe_max_sentencepiece_length > 0: + tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_{}_v{}_max{}').format( + tokenizer_type, spe_type, vocab_size, spe_max_sentencepiece_length + ) + else: + tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_{}_v{}').format( + tokenizer_type, spe_type, vocab_size + ) + + if not os.path.exists(tokenizer_dir): + os.makedirs(tokenizer_dir) + + if os.path.exists(os.path.join(tokenizer_dir, 'tokenizer.model')): + logging.warning("Model file already exists, overriding old model file !") + os.remove(os.path.join(tokenizer_dir, 'tokenizer.model')) + + tokenizer_path, vocab_path = create_spt_model( + data_file=text_path, + vocab_size=vocab_size, + sample_size=spe_sample_size, + do_lower_case=lower_case, + output_dir=tokenizer_dir, + tokenizer_type=spe_type, + character_coverage=spe_character_coverage, + train_extremely_large_corpus=spe_train_extremely_large_corpus, + max_sentencepiece_length=spe_max_sentencepiece_length, + ) + + else: + tokenizer_dir = os.path.join(dst_folder, 'tokenizer_{}_v{}').format(tokenizer_type, vocab_size) + + if not os.path.exists(tokenizer_dir): + os.makedirs(tokenizer_dir) + + tokenizer = tokenizers.BertWordPieceTokenizer(lowercase=lower_case) + + tokenizer.train(text_path, vocab_size=vocab_size) + tokenizer.save_model(tokenizer_dir) + + return tokenizer_dir + + +def main(): + data_root = args.data_root + manifests = args.manifest + data_file = args.data_file + vocab_size = args.vocab_size + tokenizer = args.tokenizer + spe_type = args.spe_type + spe_character_coverage = args.spe_character_coverage + spe_sample_size = args.spe_sample_size + spe_train_extremely_large_corpus = args.spe_train_extremely_large_corpus + spe_max_sentencepiece_length = args.spe_max_sentencepiece_length + lower_case = args.lower_case + + if not os.path.exists(data_root): + os.makedirs(data_root) + + if args.log: + logging.basicConfig(level=logging.INFO) + + if manifests: + text_corpus_path = __build_document_from_manifests(data_root, manifests) + else: + text_corpus_path = data_file + tokenizer_path = __process_data( + text_corpus_path, + data_root, + vocab_size, + tokenizer, + spe_type, + lower_case=lower_case, + spe_character_coverage=spe_character_coverage, + spe_sample_size=spe_sample_size, + spe_train_extremely_large_corpus=spe_train_extremely_large_corpus, + spe_max_sentencepiece_length=spe_max_sentencepiece_length, + ) + + print("Serialized tokenizer at location :", tokenizer_path) + logging.info('Done!') + + +if __name__ == "__main__": + main() diff --git a/models/tts/NeMo/scripts/voice_activity_detection/vad_overlap_posterior.py b/models/tts/NeMo/scripts/voice_activity_detection/vad_overlap_posterior.py new file mode 100644 index 0000000000..0849a6515f --- /dev/null +++ b/models/tts/NeMo/scripts/voice_activity_detection/vad_overlap_posterior.py @@ -0,0 +1,259 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 glob +import os +import time +from argparse import ArgumentParser +from itertools import repeat +from multiprocessing import Pool + +import numpy as np +import pandas as pd + + +""" +This script serves two purposes: + + 1) gen_overlap_seq: + Generate predictions with overlapping input segments by using the frame level prediction from NeMo/examples/asr/vad_infer.py. + Then a smoothing filter is applied to decide the label for a frame spanned by multiple segments. + + 2）gen_seg_table: + Converting frame level prediction to speech/no-speech segment in start and end times format. + +Usage: + +python vad_overlap_posterior.py --gen_overlap_seq --gen_seg_table --frame_folder= --method='median' --overlap=0.875 --num_workers=20 --threshold=0.8 + +""" + + +def gen_overlap_seq(frame_filepath, per_args): + """ + Given a frame level prediction, generate predictions with overlapping input segments by using it + Args: + frame_filepath : frame prediction file to be processed. + per_args: + method : Median or mean smoothing filter. + overlap : Amounts of overlap. + seg_len : Length of window for generating the frame. + shift_len : Amount of shift of window for generating the frame. + out_dir : Output dir of generated prediction. + """ + + try: + method = per_args['method'] + overlap = per_args['overlap'] + seg_len = per_args['seg_len'] + shift_len = per_args['shift_len'] + out_dir = per_args['out_dir'] + + frame = np.loadtxt(frame_filepath) + name = os.path.basename(frame_filepath).split(".frame")[0] + "." + method + overlap_filepath = os.path.join(out_dir, name) + + shift = int(shift_len / 0.01) # number of units of shift + seg = int((seg_len / 0.01 + 1)) # number of units of each window/segment + + jump_on_target = int(seg * (1 - overlap)) # jump on target generated sequence + jump_on_frame = int(jump_on_target / shift) # jump on input frame sequence + + if jump_on_frame < 1: + raise ValueError( + f"Note we jump over frame sequence to generate overlapping input segments. \n \ + Your input makes jump_on_fram={jump_on_frame} < 1 which is invalid because it cannot jump and will stuck.\n \ + Please try different seg_len, shift_len and overlap choices. \n \ + jump_on_target = int(seg * (1 - overlap)) \n \ + jump_on_frame = int(jump_on_frame/shift) " + ) + + target_len = int(len(frame) * shift) + + if method == 'mean': + preds = np.zeros(target_len) + pred_count = np.zeros(target_len) + + for i, og_pred in enumerate(frame): + if i % jump_on_frame != 0: + continue + start = i * shift + end = start + seg + preds[start:end] = preds[start:end] + og_pred + pred_count[start:end] = pred_count[start:end] + 1 + + preds = preds / pred_count + last_non_zero_pred = preds[pred_count != 0][-1] + preds[pred_count == 0] = last_non_zero_pred + + elif method == 'median': + preds = [[] for _ in range(target_len)] + for i, og_pred in enumerate(frame): + if i % jump_on_frame != 0: + continue + + start = i * shift + end = start + seg + for j in range(start, end): + if j <= target_len - 1: + preds[j].append(og_pred) + + preds = np.array([np.median(l) for l in preds]) + nan_idx = np.isnan(preds) + last_non_nan_pred = preds[~nan_idx][-1] + preds[nan_idx] = last_non_nan_pred + + else: + raise ValueError("method should be either mean or median") + + round_final = np.round(preds, 4) + np.savetxt(overlap_filepath, round_final, delimiter='\n') + print(f"Finished! {overlap_filepath}!") + + except Exception as e: + raise (e) + + +def gen_seg_table(frame_filepath, per_args): + + """ + Convert frame level prediction to speech/no-speech segment in start and end times format. + And save to csv file in rttm-like format + 0, 10, speech + 10,12, no-speech + Args: + frame_filepath : frame prediction file to be processed. + per_args : + threshold : threshold for prediction score (from 0 to 1). + shift_len : Amount of shift of window for generating the frame. + out_dir : Output dir of generated table/csv file. + """ + threshold = per_args['threshold'] + shift_len = per_args['shift_len'] + out_dir = per_args['out_dir'] + + print(f"process {frame_filepath}") + name = frame_filepath.split("/")[-1].rsplit(".", 1)[0] + + sequence = np.loadtxt(frame_filepath) + start = 0 + end = 0 + start_list = [0] + end_list = [] + state_list = [] + + for i in range(len(sequence) - 1): + current_sate = "non-speech" if sequence[i] <= threshold else "speech" + next_state = "non-speech" if sequence[i + 1] <= threshold else "speech" + if next_state != current_sate: + end = i * shift_len + shift_len # shift_len for handling joint + state_list.append(current_sate) + end_list.append(end) + + start = (i + 1) * shift_len + start_list.append(start) + + end_list.append((i + 1) * shift_len + shift_len) + state_list.append(current_sate) + + seg_table = pd.DataFrame({'start': start_list, 'end': end_list, 'vad': state_list}) + + save_name = name + ".txt" + save_path = os.path.join(out_dir, save_name) + seg_table.to_csv(save_path, sep='\t', index=False, header=False) + + +if __name__ == '__main__': + start = time.time() + parser = ArgumentParser() + parser.add_argument("--gen_overlap_seq", default=False, action='store_true') + parser.add_argument("--gen_seg_table", default=False, action='store_true') + parser.add_argument("--frame_folder", type=str, required=True) + parser.add_argument( + "--method", + type=str, + required=True, + help="Use mean/median for overlapped prediction. Use frame for gen_seg_table of frame prediction", + ) + parser.add_argument("--overlap_out_dir", type=str) + parser.add_argument("--table_out_dir", type=str) + parser.add_argument("--overlap", type=float, default=0.875, help="Overlap percentatge. Default is 0.875") + parser.add_argument("--seg_len", type=float, default=0.63) + parser.add_argument("--shift_len", type=float, default=0.01) + parser.add_argument("--threshold", type=float, default=0.5) + parser.add_argument("--num_workers", type=int, default=4) + args = parser.parse_args() + + if args.gen_overlap_seq: + p = Pool(processes=args.num_workers) + print("Generating predictions with overlapping input segments") + frame_filepathlist = glob.glob(args.frame_folder + "/*.frame") + + if not args.overlap_out_dir: + overlap_out_dir = "./overlap_smoothing_output" + "_" + args.method + "_" + str(args.overlap) + else: + overlap_out_dir = args.overlap_out_dir + if not os.path.exists(overlap_out_dir): + print(f"Creating output dir {overlap_out_dir}") + os.mkdir(overlap_out_dir) + + per_args = { + "method": args.method, + "overlap": args.overlap, + "seg_len": args.seg_len, + "shift_len": args.shift_len, + "out_dir": overlap_out_dir, + } + + p.starmap(gen_overlap_seq, zip(frame_filepathlist, repeat(per_args))) + p.close() + p.join() + + end = time.time() + print(f"Generate overlapped prediction takes {end-start} seconds!\n Save to {overlap_out_dir}") + + if args.gen_seg_table: + p = Pool(processes=args.num_workers) + print("Converting frame level prediction to speech/no-speech segment in start and end times format.") + + if args.gen_overlap_seq: + print("Use overlap prediction. Change if you want to use basic frame level prediction") + frame_filepath = overlap_out_dir + shift_len = 0.01 + else: + print("Use basic frame level prediction") + frame_filepath = args.frame_folder + shift_len = args.shift_len + + frame_filepathlist = glob.glob(frame_filepath + "/*." + args.method) + + if not args.table_out_dir: + table_out_dir = "table_output_" + str(args.threshold) + else: + table_out_dir = args.table_out_dir + if not os.path.exists(table_out_dir): + print(f"Creating rttm-like table output dir {table_out_dir}") + os.mkdir(table_out_dir) + + per_args = { + "threshold": args.threshold, + "shift_len": shift_len, + "out_dir": table_out_dir, + } + + p.starmap(gen_seg_table, zip(frame_filepathlist, repeat(per_args))) + p.close() + p.join() + + end = time.time() + print(f"Generate rttm-like table for {frame_filepath} takes {end-start} seconds!\n Save to {table_out_dir}") diff --git a/models/tts/NeMo/scripts/voice_activity_detection/vad_tune_threshold.py b/models/tts/NeMo/scripts/voice_activity_detection/vad_tune_threshold.py new file mode 100644 index 0000000000..1486ae730c --- /dev/null +++ b/models/tts/NeMo/scripts/voice_activity_detection/vad_tune_threshold.py @@ -0,0 +1,58 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse + +import numpy as np + +from nemo.collections.asr.parts.utils.vad_utils import vad_tune_threshold_on_dev +from nemo.utils import logging + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--threshold_range", help="range of threshold in list 'START,END,STEP' to be tuned on", required=True + ) + parser.add_argument( + "--vad_pred", help="Directory of vad predictions or a file contains the paths of them.", required=True + ) + parser.add_argument( + "--groundtruth_RTTM", + help="Directory of groundtruch rttm files or a file contains the paths of them", + type=str, + required=True, + ) + parser.add_argument( + "--vad_pred_method", + help="suffix of prediction file. Should be either in 'frame', 'mean' or 'median'", + required=True, + ) + parser.add_argument( + "--focus_metric", + help="metrics we care most when tuning threshold. Should be either in 'DetER', 'FA', 'MISS' ", + type=str, + default='DetER', + ) + args = parser.parse_args() + + try: + start, stop, step = [float(i) for i in args.threshold_range.split(",")] + thresholds = np.arange(start, stop, step) + except: + raise ValueError("Theshold input is invalid! Please enter it as a 'START,STOP,STEP' ") + + best_threhsold = vad_tune_threshold_on_dev( + thresholds, args.vad_pred, args.groundtruth_RTTM, args.vad_pred_method, args.focus_metric + ) + logging.info(f"Best threshold selected from {thresholds} is {best_threhsold}!") diff --git a/models/tts/NeMo/scripts/voice_activity_detection/write_long_audio_manifest.py b/models/tts/NeMo/scripts/voice_activity_detection/write_long_audio_manifest.py new file mode 100644 index 0000000000..f4501224d4 --- /dev/null +++ b/models/tts/NeMo/scripts/voice_activity_detection/write_long_audio_manifest.py @@ -0,0 +1,173 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 json +import logging +import multiprocessing +import os +from argparse import ArgumentParser +from itertools import repeat + +import librosa +import numpy as np + + +""" +This script is designed for inference of frame level Voice Activity Detection (VAD) + +This script serves three goals: + (1) Write audio files to manifest + (2) Split audio file for avoiding CUDA memory issue + (3) Take care of joint of seperate json line for an audio file + +Usage: +python write_long_audio_manifest.py --inp_dir= --manifest_name= --split_duration=300 --time_length=0.63 --num_worker=10 + +""" + + +def write_manifest(file, args_func): + """ + Given a list of files, write them to manifest with restrictions. + Args: + files : file to be processed + label : label for audio snippet. + split_duration : Max duration of each audio clip (each line in json) + time_length : Used for taking care of joint. + Length of window for generating the frame. + Returns: + res : list of generated metadata line of json for file + """ + + label = args_func['label'] + split_duration = args_func['split_duration'] + time_length = args_func['time_length'] + + res = [] + try: + sr = 16000 + x, _sr = librosa.load(file, sr=sr) + duration = librosa.get_duration(x, sr=sr) + + left = duration + current_offset = 0 + status = 'single' + + while left > 0: + if left <= split_duration: + if status == 'single': + write_duration = left + current_offset = 0 + else: + status = 'end' + write_duration = left + time_length + current_offset -= time_length + offset_inc = left + left = 0 + else: + if status == 'start' or status == 'next': + status = 'next' + else: + status = 'start' + + if status == 'start': + write_duration = split_duration + offset_inc = split_duration + else: + write_duration = split_duration + time_length + current_offset -= time_length + offset_inc = split_duration + time_length + + left -= split_duration + + metadata = { + 'audio_filepath': file, + 'duration': write_duration, + 'label': label, + 'text': '_', + 'offset': current_offset, + } + res.append(metadata) + + current_offset += offset_inc + + except Exception as e: + err_file = "error.log" + with open(err_file, 'w') as fout: + fout.write(file + ":" + str(e)) + + return res + + +def main(): + parser = ArgumentParser() + parser.add_argument("--inp_dir", type=str, required=True, help="(full path) folder of files to be processed") + parser.add_argument( + "--inp_list", type=str, help="(full path) a file contains NAME of files inside inp_dir to be processed" + ) + parser.add_argument( + "--out_dir", type=str, default=".", help="[Optional](full path) location to store generated json file" + ) + parser.add_argument("--manifest_name", type=str, required=True, help="name of generated json file") + parser.add_argument("--split_duration", type=int, required=True, help="max duration of each audio clip/line") + parser.add_argument( + "--time_length", type=float, default=0.63, help="[Optional] time length of segment, default is 0.63s" + ) + parser.add_argument("--num_workers", type=int, default=4, help="[Optional] number of workers for multiprocessing") + args = parser.parse_args() + + if not args.inp_list: + input_audios = [] + for root, dirs, files in os.walk(args.inp_dir): + for basename in files: + if basename.endswith('.wav'): + filename = os.path.join(root, basename) + input_audios.append(filename) + else: + name_list = np.loadtxt(args.inp_list, dtype='str') + input_audios = [os.path.join(args.inp_dir, name + ".wav") for name in name_list] + print(f"Number of wav files to be processed: {len(input_audios)}") + + output_path = os.path.join(args.out_dir, args.manifest_name + '.json') + print(f"Save generate manifest to {output_path}!") + if not os.path.exists(args.out_dir): + logging.info(f'Outdir {args.out_dir} does not exist. Creat directory.') + os.mkdir(args.out_dir) + if os.path.exists(output_path): + print(f"Manifest {output_path} exists! Overwriting") + os.remove(output_path) + + p = multiprocessing.Pool(processes=args.num_workers) + + args_func = { + 'label': 'infer', + 'split_duration': args.split_duration, + 'time_length': args.time_length, + } + + results = p.starmap(write_manifest, zip(input_audios, repeat(args_func))) + + with open(output_path, 'a') as fout: + for res in results: + for r in res: + json.dump(r, fout) + fout.write('\n') + fout.flush() + + p.close() + + print(f"Done! Save to {output_path}") + + +if __name__ == '__main__': + main() diff --git a/models/tts/NeMo/setup.cfg b/models/tts/NeMo/setup.cfg new file mode 100644 index 0000000000..694e91590d --- /dev/null +++ b/models/tts/NeMo/setup.cfg @@ -0,0 +1,36 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +[aliases] +test=pytest + +# durations=0 will display all tests execution time, sorted in ascending order starting from from the slowest one. +# -vv will also display tests with durration = 0.00s +[tool:pytest] +addopts = --verbose --pyargs --durations=0 +markers = + unit: marks unit test, i.e. testing a single, well isolated functionality (deselect with '-m "not unit"') + integration: marks test checking the elements when integrated into subsystems (deselect with '-m "not integration"') + system: marks test working at the highest integration level (deselect with '-m "not system"') + acceptance: marks test checking whether the developed product/model passes the user defined acceptance criteria (deselect with '-m "not acceptance"') + docs: mark tests related to documentation (deselect with '-m "not docs"') + skipduringci: marks tests that are skipped ci as they are addressed by Jenkins jobs but should be run to test user setups + pleasefixme: marks tests that are broken and need fixing + +[isort] +known_localfolder = nemo,tests +sections = FUTURE,STDLIB,THIRDPARTY,LOCALFOLDER +default_section = THIRDPARTY +#TODO tests/unit/core/test_deploy_export.py gets screwed by isort +skip = setup.py, docs/source/conf.py, nemo/utils/__init__.py, tests/unit/core/test_deploy_export.py diff --git a/models/tts/NeMo/setup.py b/models/tts/NeMo/setup.py new file mode 100644 index 0000000000..77c8a56bd8 --- /dev/null +++ b/models/tts/NeMo/setup.py @@ -0,0 +1,238 @@ +# ! /usr/bin/python +# -*- coding: utf-8 -*- + +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +"""Setup for pip package.""" + +import codecs +import os +import subprocess +from distutils import cmd as distutils_cmd +from distutils import log as distutils_log +from itertools import chain +import imp + +import setuptools + + +package_info = imp.load_source('package_info', 'nemo/package_info.py') +from package_info import ( + __contact_emails__, + __contact_names__, + __description__, + __download_url__, + __homepage__, + __keywords__, + __license__, + __package_name__, + __repository_url__, + __version__, +) + +if os.path.exists('nemo/README.md'): + with open("nemo/README.md", "r") as fh: + long_description = fh.read() + long_description_content_type = "text/markdown" + +elif os.path.exists('README.rst'): + # codec is used for consistent encoding + long_description = codecs.open( + os.path.join(os.path.abspath(os.path.dirname(__file__)), 'README.rst'), 'r', 'utf-8', + ).read() + long_description_content_type = "text/x-rst" + +else: + long_description = 'See ' + __homepage__ + + +############################################################################### +# Dependency Loading # +# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # + + +def req_file(filename, folder="requirements"): + with open(os.path.join(folder, filename)) as f: + content = f.readlines() + # you may also want to remove whitespace characters + # Example: `\n` at the end of each line + return [x.strip() for x in content] + + +install_requires = req_file("requirements.txt") + +extras_require = { + # User packages + 'test': req_file("requirements_test.txt"), + # Collections Packages + 'asr': req_file("requirements_asr.txt"), + 'cv': req_file("requirements_cv.txt"), + 'nlp': req_file("requirements_nlp.txt"), + 'tts': req_file("requirements_tts.txt"), + 'text_processing': req_file("requirements_text_processing.txt"), +} + +extras_require['all'] = list(chain(extras_require.values())) + +# TTS depends on ASR +extras_require['tts'] = list(chain([extras_require['tts'], extras_require['asr']])) + +tests_requirements = extras_require["test"] + + +############################################################################### +# Code style checkers # +# %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%% # + + +class StyleCommand(distutils_cmd.Command): + __LINE_WIDTH = 119 + __ISORT_BASE = ( + 'isort ' + # These two lines makes isort compatible with black. + '--multi-line=3 --trailing-comma --force-grid-wrap=0 ' + f'--use-parentheses --line-width={__LINE_WIDTH} -rc -ws' + ) + __BLACK_BASE = f'black --skip-string-normalization --line-length={__LINE_WIDTH}' + description = 'Checks overall project code style.' + user_options = [ + ('scope=', None, 'Folder of file to operate within.'), + ('fix', None, 'True if tries to fix issues in-place.'), + ] + + def __call_checker(self, base_command, scope, check): + command = list(base_command) + + command.append(scope) + + if check: + command.extend(['--check', '--diff']) + + self.announce( + msg='Running command: %s' % str(' '.join(command)), level=distutils_log.INFO, + ) + + return_code = subprocess.call(command) + + return return_code + + def _isort(self, scope, check): + return self.__call_checker(base_command=self.__ISORT_BASE.split(), scope=scope, check=check,) + + def _black(self, scope, check): + return self.__call_checker(base_command=self.__BLACK_BASE.split(), scope=scope, check=check,) + + def _pass(self): + self.announce(msg='\033[32mPASS\x1b[0m', level=distutils_log.INFO) + + def _fail(self): + self.announce(msg='\033[31mFAIL\x1b[0m', level=distutils_log.INFO) + + # noinspection PyAttributeOutsideInit + def initialize_options(self): + self.scope = '.' + self.fix = '' + + def run(self): + scope, check = self.scope, not self.fix + isort_return = self._isort(scope=scope, check=check) + black_return = self._black(scope=scope, check=check) + + if isort_return == 0 and black_return == 0: + self._pass() + else: + self._fail() + exit(isort_return if isort_return != 0 else black_return) + + def finalize_options(self): + pass + + +############################################################################### + +setuptools.setup( + name=__package_name__, + # Versions should comply with PEP440. For a discussion on single-sourcing + # the version across setup.py and the project code, see + # https://packaging.python.org/en/latest/single_source_version.html + version=__version__, + description=__description__, + long_description=long_description, + long_description_content_type=long_description_content_type, + # The project's main homepage. + url=__repository_url__, + download_url=__download_url__, + # Author details + author=__contact_names__, + author_email=__contact_emails__, + # maintainer Details + maintainer=__contact_names__, + maintainer_email=__contact_emails__, + # The licence under which the project is released + license=__license__, + classifiers=[ + # How mature is this project? Common values are + # 1 - Planning + # 2 - Pre-Alpha + # 3 - Alpha + # 4 - Beta + # 5 - Production/Stable + # 6 - Mature + # 7 - Inactive + 'Development Status :: 4 - Beta', + # Indicate who your project is intended for + 'Intended Audience :: Developers', + 'Intended Audience :: Science/Research', + 'Intended Audience :: Information Technology', + # Indicate what your project relates to + 'Topic :: Scientific/Engineering', + 'Topic :: Scientific/Engineering :: Mathematics', + 'Topic :: Scientific/Engineering :: Image Recognition', + 'Topic :: Scientific/Engineering :: Artificial Intelligence', + 'Topic :: Software Development :: Libraries', + 'Topic :: Software Development :: Libraries :: Python Modules', + 'Topic :: Utilities', + # Pick your license as you wish (should match "license" above) + 'License :: OSI Approved :: Apache Software License', + # Supported python versions + 'Programming Language :: Python :: 3', + 'Programming Language :: Python :: 3.5', + 'Programming Language :: Python :: 3.6', + 'Programming Language :: Python :: 3.7', + 'Programming Language :: Python :: 3.8', + # Additional Setting + 'Environment :: Console', + 'Natural Language :: English', + 'Operating System :: OS Independent', + ], + packages=setuptools.find_packages(), + install_requires=install_requires, + setup_requires=['pytest-runner'], + tests_require=tests_requirements, + # List additional groups of dependencies here (e.g. development + # dependencies). You can install these using the following syntax, + # $ pip install -e ".[all]" + # $ pip install nemo_toolkit[all] + extras_require=extras_require, + # Add in any packaged data. + include_package_data=True, + exclude=['tools', 'tests'], + package_data={'': ['*.tsv', '*.txt', '*.far']}, + zip_safe=False, + # PyPI package information. + keywords=__keywords__, + # Custom commands. + cmdclass={'style': StyleCommand}, +) diff --git a/models/tts/NeMo/tests/__init__.py b/models/tts/NeMo/tests/__init__.py new file mode 100644 index 0000000000..9e32500719 --- /dev/null +++ b/models/tts/NeMo/tests/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py new file mode 100644 index 0000000000..c8f55021b3 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/test_rnnt_pytorch.py @@ -0,0 +1,256 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 numpy as np +import pytest +import torch + +from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_numpy import RNNTLoss as RNNTLoss_Numpy +from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import RNNTLossNumba +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ + +DEVICES = ['cpu'] + +if torch.cuda.is_available(): + DEVICES.append('cuda') + + +def wrap_and_call(fn, acts, labels, device): + acts = torch.FloatTensor(acts) + if 'cuda' in device: + acts = acts.cuda() + acts.requires_grad = True + + lengths = [acts.shape[1]] * acts.shape[0] + label_lengths = [len(l) for l in labels] + labels = torch.IntTensor(labels) + lengths = torch.IntTensor(lengths) + label_lengths = torch.IntTensor(label_lengths) + if 'cuda' in device: + labels = labels.cuda() + lengths = lengths.cuda() + label_lengths = label_lengths.cuda() + + costs = fn(acts, labels, lengths, label_lengths) + cost = torch.sum(costs) + cost.backward() + return costs.data.cpu().numpy(), acts.grad.data.cpu().numpy() + + +class TestRNNTLossPytorch: + @pytest.mark.unit + @pytest.mark.parametrize('device', DEVICES) + def test_case_small(self, device): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + acts = np.array( + [ + [ + [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.6, 0.1, 0.1], [0.1, 0.1, 0.2, 0.8, 0.1]], + [[0.1, 0.6, 0.1, 0.1, 0.1], [0.1, 0.1, 0.2, 0.1, 0.1], [0.7, 0.1, 0.2, 0.1, 0.1]], + ] + ] + ) + labels = [[1, 2]] + + fn_pt = RNNTLossNumba(blank=0, reduction='sum') + pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device) + + fn_np = RNNTLoss_Numpy() + np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device) + + expected_cost = 4.495666 + expected_grads = np.array( + [ + [ + [ + [-0.13116688, -0.3999269, 0.17703125, 0.17703125, 0.17703125], + [-0.18572757, 0.12247056, -0.18168412, 0.12247056, 0.12247056], + [-0.32091254, 0.06269141, 0.06928472, 0.12624499, 0.06269141], + ], + [ + [0.05456069, -0.21824276, 0.05456069, 0.05456069, 0.05456069], + [0.12073959, 0.12073959, -0.48295835, 0.12073959, 0.12073959], + [-0.6925882, 0.16871116, 0.18645467, 0.16871116, 0.16871116], + ], + ] + ] + ) + assert np.allclose(pt_cost, expected_cost, rtol=1e-6), "small_test costs mismatch." + assert np.allclose(pt_grads, expected_grads), "small_test gradient mismatch." + + assert np.allclose(pt_cost, np_cost, rtol=1e-6), "small_test costs mismatch." + assert np.allclose(pt_grads, np_grads), "small_test gradient mismatch." + + @pytest.mark.unit + @pytest.mark.parametrize('device', DEVICES) + def test_case_small_random(self, device): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + rng = np.random.RandomState(0) + acts = rng.randn(1, 4, 3, 3) + labels = [[1, 2]] + + fn_pt = RNNTLossNumba(blank=0, reduction='sum') + pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device) + + fn_np = RNNTLoss_Numpy() + np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device) + + assert np.allclose(pt_cost, np_cost, rtol=1e-6), "small_random_test costs mismatch." + assert np.allclose(pt_grads, np_grads), "small_random_test gradient mismatch." + + @pytest.mark.unit + @pytest.mark.parametrize('device', DEVICES) + def big_test(self, device): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # minibatch x T x U x alphabet_size + activations = [ + [ + [ + [0.06535690384862791, 0.7875301411923206, 0.08159176605666074], + [0.5297155426466327, 0.7506749639230854, 0.7541348379087998], + [0.6097641124736383, 0.8681404965673826, 0.6225318186056529], + ], + [ + [0.6685222872103057, 0.8580392805336061, 0.16453892311765583], + [0.989779515236694, 0.944298460961015, 0.6031678586829663], + [0.9467833543605416, 0.666202507295747, 0.28688179752461884], + ], + [ + [0.09418426230195986, 0.3666735970751962, 0.736168049462793], + [0.1666804425271342, 0.7141542198635192, 0.3993997272216727], + [0.5359823524146038, 0.29182076440286386, 0.6126422611507932], + ], + [ + [0.3242405528768486, 0.8007644367291621, 0.5241057606558068], + [0.779194617063042, 0.18331417220174862, 0.113745182072432], + [0.24022162381327106, 0.3394695622533106, 0.1341595066017014], + ], + ], + [ + [ + [0.5055615569388828, 0.051597282072282646, 0.6402903936686337], + [0.43073311517251, 0.8294731834714112, 0.1774668847323424], + [0.3207001991262245, 0.04288308912457006, 0.30280282975568984], + ], + [ + [0.6751777088333762, 0.569537369330242, 0.5584738347504452], + [0.08313242153985256, 0.06016544344162322, 0.10795752845152584], + [0.7486153608562472, 0.943918041459349, 0.4863558118797222], + ], + [ + [0.4181986264486809, 0.6524078485043804, 0.024242983423721887], + [0.13458171554507403, 0.3663418070512402, 0.2958297395361563], + [0.9236695822497084, 0.6899291482654177, 0.7418981733448822], + ], + [ + [0.25000547599982104, 0.6034295486281007, 0.9872887878887768], + [0.5926057265215715, 0.8846724004467684, 0.5434495396894328], + [0.6607698886038497, 0.3771277082495921, 0.3580209022231813], + ], + ], + ] + + expected_costs = [4.2806528590890736, 3.9384369822503591] + expected_grads = [ + [ + [ + [-1.86843902e-01, -6.25548810e-02, 2.49398798e-01], + [-2.03376666e-01, 2.02399328e-01, 9.77333169e-04], + [-1.41016081e-01, 7.91234672e-02, 6.18926100e-02], + ], + [ + [-1.15517676e-02, -8.12802389e-02, 9.28319991e-02], + [-1.54257029e-01, 2.29432687e-01, -7.51756504e-02], + [-2.46593088e-01, 1.46404594e-01, 1.00188486e-01], + ], + [ + [-1.29182907e-02, -6.15932420e-02, 7.45115355e-02], + [-5.59857301e-02, 2.19830811e-01, -1.63845062e-01], + [-4.97626871e-01, 2.09239945e-01, 2.88386941e-01], + ], + [ + [1.36048580e-02, -3.02196294e-02, 1.66147724e-02], + [1.13924511e-01, 6.27811998e-02, -1.76705718e-01], + [-6.67078257e-01, 3.67658824e-01, 2.99419403e-01], + ], + ], + [ + [ + [-3.56343776e-01, -5.53474613e-02, 4.11691219e-01], + [-9.69219357e-02, 2.94591039e-02, 6.74628317e-02], + [-6.35175705e-02, 2.76544970e-02, 3.58630717e-02], + ], + [ + [-1.54499024e-01, -7.39420280e-02, 2.28441030e-01], + [-1.66789949e-01, -8.78955179e-05, 1.66877866e-01], + [-1.72369644e-01, 1.05565332e-01, 6.68043196e-02], + ], + [ + [2.38748826e-02, -1.18255816e-01, 9.43809375e-02], + [-1.04707085e-01, -1.08934477e-01, 2.13641584e-01], + [-3.69844258e-01, 1.80118099e-01, 1.89726159e-01], + ], + [ + [2.57137045e-02, -7.94617534e-02, 5.37480488e-02], + [1.22328237e-01, -2.38788679e-01, 1.16460443e-01], + [-5.98686993e-01, 3.02203178e-01, 2.96483815e-01], + ], + ], + ] + + activations = np.array(activations) + labels = [[1, 2], [1, 1]] + + fn_pt = RNNTLossNumba(blank=0, reduction='sum') + pt_costs, pt_grads = wrap_and_call(fn_pt, activations, labels, device) + + fn_np = RNNTLoss_Numpy() + np_costs, np_grads = wrap_and_call(fn_np, activations, labels, device) + + assert np.allclose(pt_costs, sum(expected_costs)), "big_test average costs mismatch." + assert np.allclose(pt_grads, expected_grads, rtol=1e-3), "big_test grads for average cost mismatch." + + assert np.allclose(pt_costs, np_costs), "big_test average costs mismatch." + assert np.allclose(pt_grads, np_grads, rtol=1e-3), "big_test grads for average cost mismatch." + + @pytest.mark.unit + @pytest.mark.parametrize('device', DEVICES) + def test_case_large_random(self, device): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + rng = np.random.RandomState(0) + acts = rng.randn(4, 8, 11, 5) + labels = [ + [1, 2, 4, 3, 2, 2, 1, 1, 1, 1], + [3, 2, 2, 3, 4, 1, 1, 1, 1, 1], + [4, 4, 1, 2, 1, 3, 4, 3, 1, 2], + [1, 1, 2, 1, 2, 3, 3, 1, 1, 1], + ] + + fn_pt = RNNTLossNumba(blank=0, reduction='sum') + pt_cost, pt_grads = wrap_and_call(fn_pt, acts, labels, device) + + fn_np = RNNTLoss_Numpy() + np_cost, np_grads = wrap_and_call(fn_np, acts, labels, device) + + assert np.allclose(pt_cost, np_cost, atol=1e-5, rtol=1e-3), "large_random_test costs mismatch." + assert np.allclose(pt_grads, np_grads, atol=1e-5, rtol=1e-3), "large_random_test gradient mismatch." + + +if __name__ == "__main__": + pytest.main([__file__]) diff --git a/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_gpu_rnnt_kernel.py b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_gpu_rnnt_kernel.py new file mode 100644 index 0000000000..a0cb84cbb7 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_gpu_rnnt_kernel.py @@ -0,0 +1,169 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 numpy as np +import pytest +import torch +from numba import cuda + +from nemo.collections.asr.parts.numba.rnnt_loss import rnnt_numpy +from nemo.collections.asr.parts.numba.rnnt_loss.rnnt_pytorch import certify_inputs +from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import gpu_rnnt_kernel, reduce +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ + + +def log_softmax(x, axis=-1): + x = torch.from_numpy(x) # zero-copy + x = torch.log_softmax(x, axis) + x = x.numpy() + return x + + +class TestRNNTCUDAKernels: + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available") + @pytest.mark.unit + def test_compute_alphas_kernel(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + random = np.random.RandomState(0) + original_shape = [1, 5, 11, 3] + B, T, U, V = original_shape + + # Numpy kernel + x = random.randn(*original_shape) + labels = np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]]) # [1, 10] + label_len = len(labels[0]) + 1 + blank_idx = 0 + + x_np = log_softmax(x, axis=-1) + ground_alphas, ground_log_likelihood = rnnt_numpy.forward_pass( + x_np[0, :, :label_len, :], labels[0, : label_len - 1], blank_idx + ) + + # Pytorch kernel + device = torch.device('cuda') + if hasattr(cuda, 'external_stream'): + stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream) + else: + stream = cuda.default_stream() + + x_c = torch.tensor(x, device=device, dtype=torch.float32) + labels_c = torch.tensor(labels, device=device, dtype=torch.int32) + + # Allocate workspace memory + denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype) + alphas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype) + llForward = torch.zeros(B, device=device, dtype=x_c.dtype) + input_lengths = torch.tensor([T], dtype=torch.int32, device=device) + label_lengths = torch.tensor([len(labels[0])], dtype=torch.int32, device=device) + + # certify input data + certify_inputs(x_c, labels_c, input_lengths, label_lengths) + + # flatten activation tensor (for pointer based indexing) + x_c = x_c.view([-1]) + + # call kernel + # log softmax reduction + reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream) + reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream) + + # alpha kernel + gpu_rnnt_kernel.compute_alphas_kernel[B, U, stream, 0]( + x_c, denom, alphas, llForward, input_lengths, label_lengths, labels_c, B, T, U, V, blank_idx, + ) + + # sync kernel + stream.synchronize() + + # reshape alphas + alphas = alphas.view([B, T, U]) + diff = ground_alphas - alphas[0].cpu().numpy() + + assert np.abs(diff).mean() <= 1e-5 + assert np.square(diff).mean() <= 1e-10 + + ll_diff = ground_log_likelihood - llForward[0].cpu().numpy() + + assert np.abs(ll_diff).mean() <= 1e-5 + assert np.square(ll_diff).mean() <= 1e-10 + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available") + @pytest.mark.unit + def test_compute_betas_kernel(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + random = np.random.RandomState(0) + original_shape = [1, 5, 11, 3] + B, T, U, V = original_shape + + # Numpy kernel + x = random.randn(*original_shape) + labels = np.array([[1, 1, 1, 2, 2, 2, 1, 2, 2, 1]]) # [1, 10] + label_len = len(labels[0]) + 1 + blank_idx = 0 + + x_np = log_softmax(x, axis=-1) + ground_alphas, ground_log_likelihood = rnnt_numpy.backward_pass( + x_np[0, :, :label_len, :], labels[0, : label_len - 1], blank_idx + ) + + # Pytorch kernel + device = torch.device('cuda') + if hasattr(cuda, 'external_stream'): + stream = cuda.external_stream(torch.cuda.current_stream(device).cuda_stream) + else: + stream = cuda.default_stream() + + x_c = torch.tensor(x, device=device, dtype=torch.float32) + labels_c = torch.tensor(labels, device=device, dtype=torch.int32) + + # Allocate workspace memory + denom = torch.zeros(B * T * U, device=device, dtype=x_c.dtype) + betas = torch.zeros(B * T * U, device=device, dtype=x_c.dtype) + llBackward = torch.zeros(B, device=device, dtype=x_c.dtype) + input_lengths = torch.tensor([T], dtype=torch.int32, device=device) + label_lengths = torch.tensor([len(labels[0])], dtype=torch.int32, device=device) + + # certify input data + certify_inputs(x_c, labels_c, input_lengths, label_lengths) + + # flatten activation tensor (for pointer based indexing) + x_c = x_c.view([-1]) + + # call kernel + # log softmax reduction + reduce.reduce_max(x_c, denom, rows=V, cols=B * T * U, minus=False, stream=stream) + reduce.reduce_exp(x_c, denom, rows=V, cols=B * T * U, minus=True, stream=stream) + + # beta kernel + gpu_rnnt_kernel.compute_betas_kernel[B, U, stream, 0]( + x_c, denom, betas, llBackward, input_lengths, label_lengths, labels_c, B, T, U, V, blank_idx, + ) + + # sync kernel + stream.synchronize() + + # reshape alphas + betas = betas.view([B, T, U]) + diff = ground_alphas - betas[0].cpu().numpy() + + assert np.abs(diff).mean() <= 1e-5 + assert np.square(diff).mean() <= 1e-10 + + ll_diff = ground_log_likelihood - llBackward[0].cpu().numpy() + + assert np.abs(ll_diff).mean() <= 1e-5 + assert np.square(ll_diff).mean() <= 1e-10 diff --git a/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_reduce.py b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_reduce.py new file mode 100644 index 0000000000..7c2ba6a412 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_reduce.py @@ -0,0 +1,86 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 numpy as np +import pytest +from numba import cuda + +from nemo.collections.asr.parts.numba.rnnt_loss.utils.cuda_utils import reduce +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ + + +class TestRNNTCUDAReductions: + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available") + @pytest.mark.unit + def test_reduce_max(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + random = np.random.RandomState(0) + original_shape = [1, 5, 4, 3] + x = random.randn(*original_shape).reshape([-1]) + dx = random.randn(*x.shape) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + dx_c = cuda.to_device(dx, stream=stream) + + # call kernel + cols = np.prod(original_shape[:3]) + reduce.reduce_max(x_c, dx_c, rows=original_shape[-1], cols=cols, minus=False, stream=stream) + + # sync kernel + stream.synchronize() + + dx_result = dx_c.copy_to_host(stream=stream) + del x_c, dx_c + + # collect results in first [B * T * U] values; for all V + assert np.abs(dx_result[cols:] - dx[cols:]).sum() <= 1e-7 + # make sure dx_result updates the [B * T * U] values + assert np.abs(dx_result[:cols] - dx[:cols]).sum() > 0 + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Reductions can only be run when CUDA is available") + @pytest.mark.unit + def test_reduce_exp(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + random = np.random.RandomState(0) + original_shape = [1, 5, 4, 2] + x = random.randn(*original_shape).reshape([-1]) + dx = np.zeros_like(x) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + dx_c = cuda.to_device(dx, stream=stream) + + # call kernel + cols = np.prod(original_shape[:3]) + reduce.reduce_exp(x_c, dx_c, rows=original_shape[-1], cols=cols, minus=False, stream=stream) + + # sync kernel + stream.synchronize() + + dx_result = dx_c.copy_to_host(stream=stream) + del x_c, dx_c + + # collect results in first [B * T * U] values; for all V + assert (dx_result[cols:] - dx[cols:]).sum() <= 1e-7 + + # make sure dx_result updates the [B * T * U] values + assert np.abs(dx_result[:cols] - dx[:cols]).sum() > 0 + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_rnnt_helper.py b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_rnnt_helper.py new file mode 100644 index 0000000000..d054bae761 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/numba/rnnt_loss/utils/test_rnnt_helper.py @@ -0,0 +1,319 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 numpy as np +import pytest +from numba import cuda + +from nemo.collections.asr.parts.numba.rnnt_loss.utils import global_constants, rnnt_helper +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ + + +class TestRNNTHelper: + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_log_sum_exp(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x, y): + x_pos = cuda.grid(1) + if x_pos < x.shape[0] and x_pos < y.shape[0]: + x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos]) + + x = np.zeros([8]) # np.random.rand(8192) + y = np.ones([8]) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + y_c = cuda.to_device(y, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c, y_c + + assert (x_new.sum() - 10.506093500145782) <= 1e-5 + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_log_sum_exp_neg_inf(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x, y): + x_pos = cuda.grid(1) + if x_pos < x.shape[0] and x_pos < y.shape[0]: + x[x_pos] = rnnt_helper.log_sum_exp(x[x_pos], y[x_pos]) + + x = np.asarray([global_constants.FP32_NEG_INF] * 8) + y = np.ones([len(x)]) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + y_c = cuda.to_device(y, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c, y_c + + assert np.allclose(x_new, np.ones_like(x_new), atol=1e-5) + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_div_up(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x, y): + x_pos = cuda.grid(1) + if x_pos < x.shape[0] and x_pos < y.shape[0]: + x[x_pos] = rnnt_helper.div_up(x[x_pos], y[x_pos]) + + x = np.full([8], fill_value=10) # np.random.rand(8192) + y = np.full([8], fill_value=2) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + y_c = cuda.to_device(y, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c, y_c + + for i in range(len(x_new)): + assert x_new[i] == ((10 + 2 - 1) // 2) + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_add(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x, y): + x_pos = cuda.grid(1) + if x_pos < x.shape[0] and x_pos < y.shape[0]: + x[x_pos] = rnnt_helper.add(x[x_pos], y[x_pos]) + + x = np.full([8], fill_value=10) # np.random.rand(8192) + y = np.full([8], fill_value=2) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + y_c = cuda.to_device(y, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c, y_c + + for i in range(len(x_new)): + assert x_new[i] == 12 + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_maximum(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x, y): + x_pos = cuda.grid(1) + if x_pos < x.shape[0] and x_pos < y.shape[0]: + x[x_pos] = rnnt_helper.maximum(x[x_pos], y[x_pos]) + + x = np.full([8], fill_value=10) # np.random.rand(8192) + y = np.full([8], fill_value=2) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + y_c = cuda.to_device(y, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c, y_c + + for i in range(len(x_new)): + assert x_new[i] == 10 + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_identity(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x): + x_pos = cuda.grid(1) + if x_pos < x.shape[0]: + x[x_pos] = rnnt_helper.identity(x[x_pos]) + + x = np.full([8], fill_value=10) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c + + for i in range(len(x_new)): + assert x_new[i] == x[i] + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_negate(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x): + x_pos = cuda.grid(1) + if x_pos < x.shape[0]: + x[x_pos] = rnnt_helper.negate(x[x_pos]) + + x = np.full([8], fill_value=10) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c + + for i in range(len(x_new)): + assert x_new[i] == -x[i] + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_exponential(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x): + x_pos = cuda.grid(1) + if x_pos < x.shape[0]: + x[x_pos] = rnnt_helper.exponential(x[x_pos]) + + x = np.random.rand(8) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c + + y = np.exp(x) + for i in range(len(x_new)): + assert (x_new[i] - y[i]) < 1e-4 + + @pytest.mark.skipif(not cuda.is_available(), reason="CUDA Helpers can only be run when CUDA is available") + @pytest.mark.unit + def test_log_plus(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + # wrapper kernel for device function that is tested + @cuda.jit + def _kernel(x, y): + x_pos = cuda.grid(1) + if x_pos < x.shape[0] and x_pos < y.shape[0]: + x[x_pos] = rnnt_helper.log_plus(x[x_pos], y[x_pos]) + + x = np.full([8], fill_value=10.0) # np.random.rand(8192) + y = np.full([8], fill_value=2.0) # np.random.rand(8192) + + stream = cuda.stream() + x_c = cuda.to_device(x, stream=stream) + y_c = cuda.to_device(y, stream=stream) + + # call kernel + threads_per_block = global_constants.threads_per_block() + blocks_per_grid = (x.shape[0] + threads_per_block - 1) // threads_per_block + _kernel[blocks_per_grid, threads_per_block, stream](x_c, y_c) + + # sync kernel + stream.synchronize() + + x_new = x_c.copy_to_host(stream=stream) + del x_c, y_c + + z = np.log1p(np.exp(-np.fabs(x - y))) + np.maximum(x, y) + + for i in range(len(x_new)): + assert x_new[i] == z[i] + + +if __name__ == '__main__': + pytest.main([__file__]) diff --git a/models/tts/NeMo/tests/collections/asr/numba/spec_augment/test_spec_aug_numba.py b/models/tts/NeMo/tests/collections/asr/numba/spec_augment/test_spec_aug_numba.py new file mode 100644 index 0000000000..17b6d6c9cc --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/numba/spec_augment/test_spec_aug_numba.py @@ -0,0 +1,252 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +import torch +from omegaconf import OmegaConf + +from nemo.collections.asr.parts.numba.spec_augment import spec_aug_numba +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ + + +def get_cfg(seed=0, dtype='float32'): + # fmt: off + cfg = OmegaConf.create( + dict(b=2, f=80, t=20, device='cuda', + freq_masks=2, time_masks=2, freq_width=27, time_width=0.05, mask_value=0.0, + seed=seed, dtype=dtype) + ) + # fmt: on + return cfg + + +# fmt: off +def prepare_data(b, f, t, device='cuda', freq_masks=0, time_masks=0, freq_width=10, time_width=0.1, + seed=0, dtype='float32', + **kwargs): + torch.manual_seed(seed) + + if dtype == 'float16': + dtype = torch.float16 + else: + dtype = torch.float + + x = torch.randn([b, f, t], dtype=dtype, device=device) + x_len = torch.randint(t, size=[b], device=x.device) + + sh = x.shape + bs = sh[0] + + if isinstance(time_width, int): + adaptive_temporal_width = False + else: + if time_width > 1.0 or time_width < 0.0: + raise ValueError('If `time_width` is a float value, must be in range [0, 1]') + + adaptive_temporal_width = True + + if adaptive_temporal_width: + time_width = max(1, int(sh[2] * time_width)) + else: + time_width = time_width + + # Construct the freq and time masks as well as start positions + if freq_masks > 0: + freq_starts = torch.randint(0, sh[1] - freq_width + 1, size=[bs, freq_masks], device=x.device) + freq_lengths = torch.randint(0, freq_width + 1, size=[bs, freq_masks], device=x.device) + else: + freq_starts = torch.zeros([bs, 1], dtype=torch.int64, device=x.device) + freq_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=x.device) + + if time_masks > 0: + time_starts = torch.randint(0, sh[2] - time_width + 1, size=[bs, time_masks], device=x.device) + time_lengths = torch.randint(0, time_width + 1, size=[bs, time_masks], device=x.device) + else: + time_starts = torch.zeros([bs, 1], dtype=torch.int64, device=x.device) + time_lengths = torch.zeros([bs, 1], dtype=torch.int64, device=x.device) + + output = dict( + x=x, + x_len=x_len, + freq_starts=freq_starts, + freq_lengths=freq_lengths, + time_starts=time_starts, + time_lengths=time_lengths, + sh=sh, + ) + return output +# fmt: on + + +def launch_kernel(data, cfg): + # Launch CUDA kernel + # fmt: off + data['x'] = spec_aug_numba.launch_spec_augment_kernel( + x=data['x'], x_len=data['x_len'], + freq_starts=data['freq_starts'], freq_lengths=data['freq_lengths'], + time_starts=data['time_starts'], time_lengths=data['time_lengths'], + freq_masks=cfg.freq_masks, time_masks=cfg.time_masks, mask_value=cfg.mask_value + ) + # fmt: on + + +def freq_mask_check(x, x_len, f_start, f_len, mask_value, bidx): + check_result = True + for fidx in range(f_start, f_start + f_len): + if not (x[bidx, fidx, :] == mask_value).all(): + check_result = False + break + assert check_result + + +def time_mask_check(x, x_len, t_start, t_len, mask_value, bidx): + check_result = True + for tidx in range(t_start, t_start + t_len): + if tidx >= x_len[bidx]: + # this sample has smaller length than the time index of mask, ignore + continue + + if not (x[bidx, :, tidx] == mask_value).all(): + check_result = False + break + assert check_result + + +class TestSpecAugmentNumba: + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + @pytest.mark.parametrize('dtype', ['float16', 'float32']) + def test_spec_aug_kernel(self, dtype): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + cfg = get_cfg(seed=0, dtype=dtype) + cfg.freq_masks = 2 + cfg.time_masks = 10 + + data = prepare_data(**cfg) + + launch_kernel(data, cfg) + x, x_len, sh = data['x'], data['x_len'], data['sh'] + + # Assert freq masks are correct + for bidx in range(sh[0]): + for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]): + freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx) + + # Assert time masks are correct + for bidx in range(sh[0]): + for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]): + time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx) + + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + def test_spec_aug_kernel_mask_value(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + cfg = get_cfg(seed=0) + cfg.freq_masks = 2 + cfg.time_masks = 10 + cfg.mask_value = -1.0 + + data = prepare_data(**cfg) + + launch_kernel(data, cfg) + x, x_len, sh = data['x'], data['x_len'], data['sh'] + + # Assert freq masks are correct + for bidx in range(sh[0]): + for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]): + freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx) + + # Assert time masks are correct + for bidx in range(sh[0]): + for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]): + time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx) + + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + def test_spec_aug_kernel_grad(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + cfg = get_cfg(seed=0) + cfg.freq_masks = 2 + cfg.time_masks = 10 + + data = prepare_data(**cfg) + + launch_kernel(data, cfg) + + result = data['x'] # inplace modification via kernel + y = torch.ones_like(result, requires_grad=True) + z = y + result + z.mean().backward() + + assert y.grad is not None + + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + def test_spec_aug_kernel_no_freq_mask(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + cfg = get_cfg(seed=0) + cfg.freq_masks = 0 + cfg.time_masks = 10 + + data = prepare_data(**cfg) + + launch_kernel(data, cfg) + x, x_len, sh = data['x'], data['x_len'], data['sh'] + + # Assert time masks are correct + for bidx in range(sh[0]): + for t_start, t_len in zip(data['time_starts'][bidx], data['time_lengths'][bidx]): + time_mask_check(x, x_len, t_start, t_len, mask_value=cfg.mask_value, bidx=bidx) + + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + def test_spec_aug_kernel_no_time_mask(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + cfg = get_cfg(seed=0) + cfg.freq_masks = 2 + cfg.time_masks = 0 + + data = prepare_data(**cfg) + + launch_kernel(data, cfg) + x, x_len, sh = data['x'], data['x_len'], data['sh'] + + # Assert freq masks are correct + for bidx in range(sh[0]): + for f_start, f_len in zip(data['freq_starts'][bidx], data['freq_lengths'][bidx]): + freq_mask_check(x, x_len, f_start, f_len, mask_value=cfg.mask_value, bidx=bidx) + + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + def test_spec_aug_kernel_no_freq_time_mask(self): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + cfg = get_cfg(seed=0) + cfg.freq_masks = 0 + cfg.time_masks = 0 + + data = prepare_data(**cfg) + + x, x_len, sh = data['x'], data['x_len'], data['sh'] + x_copy = x.clone() + launch_kernel(data, cfg) + + # Assert no data edits occured + assert (data['x'] - x_copy).abs().mean() <= 1e-9 diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_classification_model.py b/models/tts/NeMo/tests/collections/asr/test_asr_classification_model.py new file mode 100644 index 0000000000..24c88d9beb --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_classification_model.py @@ -0,0 +1,208 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 copy +import os + +import pytest +import torch +from omegaconf import DictConfig, ListConfig + +from nemo.collections.asr.data import audio_to_label +from nemo.collections.asr.models import EncDecClassificationModel, configs +from nemo.utils.config_utils import assert_dataclass_signature_match + + +@pytest.fixture() +def speech_classification_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 32, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification', + 'params': {'feat_in': 32, 'num_classes': 30,}, + } + + modelConfig = DictConfig( + { + 'preprocessor': DictConfig(preprocessor), + 'encoder': DictConfig(encoder), + 'decoder': DictConfig(decoder), + 'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]), + } + ) + model = EncDecClassificationModel(cfg=modelConfig) + return model + + +class TestEncDecClassificationModel: + @pytest.mark.unit + def test_constructor(self, speech_classification_model): + asr_model = speech_classification_model.train() + + conv_cnt = (64 * 32 * 1 + 32) + (64 * 1 * 1 + 32) # separable kernel + bias + pointwise kernel + bias + bn_cnt = (4 * 32) * 2 # 2 * moving averages + dec_cnt = 32 * 30 + 30 # fc + bias + + param_count = conv_cnt + bn_cnt + dec_cnt + assert asr_model.num_weights == param_count + + # Check to/from config_dict: + confdict = asr_model.to_config_dict() + instance2 = EncDecClassificationModel.from_config_dict(confdict) + + assert isinstance(instance2, EncDecClassificationModel) + + @pytest.mark.unit + def test_forward(self, speech_classification_model): + asr_model = speech_classification_model.eval() + + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + + with torch.no_grad(): + # batch size 1 + logprobs_instance = [] + for i in range(input_signal.size(0)): + logprobs_ins = asr_model.forward( + input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1] + ) + logprobs_instance.append(logprobs_ins) + logprobs_instance = torch.cat(logprobs_instance, 0) + + # batch size 4 + logprobs_batch = asr_model.forward(input_signal=input_signal, input_signal_length=length) + + assert logprobs_instance.shape == logprobs_batch.shape + diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + diff = torch.max(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + + @pytest.mark.unit + def test_vocab_change(self, speech_classification_model): + asr_model = speech_classification_model.train() + + old_labels = copy.deepcopy(asr_model._cfg.labels) + nw1 = asr_model.num_weights + asr_model.change_labels(new_labels=old_labels) + # No change + assert nw1 == asr_model.num_weights + new_labels = copy.deepcopy(old_labels) + new_labels.append('dummy_cls_31') + new_labels.append('dummy_cls_32') + new_labels.append('dummy_cls_33') + asr_model.change_labels(new_labels=new_labels) + # fully connected + bias + assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1) + + @pytest.mark.unit + def test_transcription(self, speech_classification_model, test_data_dir): + # Ground truth labels = ["yes", "no"] + audio_filenames = ['an22-flrp-b.wav', 'an90-fbbh-b.wav'] + audio_paths = [os.path.join(test_data_dir, "asr", "train", "an4", "wav", fp) for fp in audio_filenames] + + model = speech_classification_model.eval() + + # Test Top 1 classification transcription + results = model.transcribe(audio_paths, batch_size=2) + assert len(results) == 2 + assert results[0].shape == torch.Size([1]) + + # Test Top 5 classification transcription + model._accuracy.top_k = [5] # set top k to 5 for accuracy calculation + results = model.transcribe(audio_paths, batch_size=2) + assert len(results) == 2 + assert results[0].shape == torch.Size([5]) + + # Test Top 1 and Top 5 classification transcription + model._accuracy.top_k = [1, 5] + results = model.transcribe(audio_paths, batch_size=2) + assert len(results) == 2 + assert results[0].shape == torch.Size([2, 1]) + assert results[1].shape == torch.Size([2, 5]) + assert model._accuracy.top_k == [1, 5] + + # Test log probs extraction + model._accuracy.top_k = [1] + results = model.transcribe(audio_paths, batch_size=2, logprobs=True) + assert len(results) == 2 + assert results[0].shape == torch.Size([len(model.cfg.labels)]) + + # Test log probs extraction remains same for any top_k + model._accuracy.top_k = [5] + results = model.transcribe(audio_paths, batch_size=2, logprobs=True) + assert len(results) == 2 + assert results[0].shape == torch.Size([len(model.cfg.labels)]) + + @pytest.mark.unit + def test_EncDecClassificationDatasetConfig_for_AudioToSpeechLabelDataset(self): + # ignore some additional arguments as dataclass is generic + IGNORE_ARGS = [ + 'is_tarred', + 'num_workers', + 'batch_size', + 'tarred_audio_filepaths', + 'shuffle', + 'pin_memory', + 'drop_last', + 'tarred_shard_strategy', + 'shuffle_n', + # `featurizer` is supplied at runtime + 'featurizer', + # additional ignored arguments + 'vad_stream', + 'int_values', + 'sample_rate', + 'normalize_audio', + 'augmentor', + ] + + REMAP_ARGS = {'trim_silence': 'trim'} + + result = assert_dataclass_signature_match( + audio_to_label.AudioToSpeechLabelDataset, + configs.EncDecClassificationDatasetConfig, + ignore_args=IGNORE_ARGS, + remap_args=REMAP_ARGS, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_ctc_encoder_model_bpe.py b/models/tts/NeMo/tests/collections/asr/test_asr_ctc_encoder_model_bpe.py new file mode 100644 index 0000000000..b42a06d9a0 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_ctc_encoder_model_bpe.py @@ -0,0 +1,313 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 copy +import os +import shutil +import tempfile + +import pytest +import torch +from omegaconf import DictConfig + +from nemo.collections.asr.data import audio_to_text +from nemo.collections.asr.models import configs +from nemo.collections.asr.models.ctc_bpe_models import EncDecCTCModelBPE +from nemo.collections.common import tokenizers +from nemo.utils.config_utils import assert_dataclass_signature_match + + +@pytest.fixture() +def asr_model(test_data_dir): + preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'} + encoder = { + '_target_': 'nemo.collections.asr.modules.ConvASREncoder', + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 1024, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + } + + decoder = { + '_target_': 'nemo.collections.asr.modules.ConvASRDecoder', + 'feat_in': 1024, + 'num_classes': -1, + 'vocabulary': None, + } + + tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'} + + modelConfig = DictConfig( + { + 'preprocessor': DictConfig(preprocessor), + 'encoder': DictConfig(encoder), + 'decoder': DictConfig(decoder), + 'tokenizer': DictConfig(tokenizer), + } + ) + + model_instance = EncDecCTCModelBPE(cfg=modelConfig) + return model_instance + + +class TestEncDecCTCModel: + @pytest.mark.unit + def test_constructor(self, asr_model): + asr_model.train() + # TODO: make proper config and assert correct number of weights + # Check to/from config_dict: + confdict = asr_model.to_config_dict() + instance2 = EncDecCTCModelBPE.from_config_dict(confdict) + assert isinstance(instance2, EncDecCTCModelBPE) + + @pytest.mark.unit + def test_forward(self, asr_model): + asr_model = asr_model.eval() + + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + + with torch.no_grad(): + # batch size 1 + logprobs_instance = [] + for i in range(input_signal.size(0)): + logprobs_ins, _, _ = asr_model.forward( + input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1] + ) + logprobs_instance.append(logprobs_ins) + print(len(logprobs_ins)) + logprobs_instance = torch.cat(logprobs_instance, 0) + + # batch size 4 + logprobs_batch, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length) + + assert logprobs_instance.shape == logprobs_batch.shape + diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + diff = torch.max(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + + @pytest.mark.unit + def test_save_restore_artifact(self, asr_model): + with tempfile.TemporaryDirectory() as tmpdir: + save_path = os.path.join(tmpdir, 'ctc_bpe.nemo') + asr_model.train() + asr_model.save_to(save_path) + + new_model = EncDecCTCModelBPE.restore_from(save_path) + assert isinstance(new_model, type(asr_model)) + assert new_model.vocab_path.endswith('_vocab.txt') + + assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128 + + @pytest.mark.unit + def test_save_restore_artifact_spe(self, asr_model, test_data_dir): + with tempfile.TemporaryDirectory() as tmpdir: + tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128") + asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe') + + save_path = os.path.join(tmpdir, 'ctc_bpe.nemo') + asr_model.train() + asr_model.save_to(save_path) + + new_model = EncDecCTCModelBPE.restore_from(save_path) + assert isinstance(new_model, type(asr_model)) + assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer) + assert new_model.model_path.endswith('_tokenizer.model') + assert new_model.vocab_path.endswith('_vocab.txt') + assert new_model.spe_vocab_path.endswith('_tokenizer.vocab') + + assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128 + + @pytest.mark.unit + def test_vocab_change(self, test_data_dir, asr_model): + old_vocab = copy.deepcopy(asr_model.decoder.vocabulary) + + with tempfile.TemporaryDirectory() as save_dir: + save_path = os.path.join(save_dir, 'temp.nemo') + + with tempfile.TemporaryDirectory() as tmpdir: + old_tmpdir_path = tmpdir + + old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt') + new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer') + + os.makedirs(new_tokenizer_dir, exist_ok=True) + shutil.copy2(old_tokenizer_dir, new_tokenizer_dir) + + nw1 = asr_model.num_weights + asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe') + # No change + assert nw1 == asr_model.num_weights + + with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f: + f.write("!\n") + f.write('$\n') + f.write('@\n') + + asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe') + # fully connected + bias + assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1) + + new_vocab = copy.deepcopy(asr_model.decoder.vocabulary) + assert len(old_vocab) != len(new_vocab) + + # save the model (after change of vocabulary) + asr_model.save_to(save_path) + assert os.path.exists(save_path) + # delete copied version of the vocabulary from nested tmpdir (by scope) + + # assert copied vocab no longer exists + assert not os.path.exists(os.path.join(old_tmpdir_path, 'tokenizer', 'vocab.txt')) + + # make a copy of the tokenizer before renaming + try: + os.rename(old_tokenizer_dir, old_tokenizer_dir + '.bkp') + assert not os.path.exists(old_tokenizer_dir) + + # restore model from .nemo + asr_model2 = EncDecCTCModelBPE.restore_from(save_path) + assert isinstance(asr_model2, EncDecCTCModelBPE) + + # Check if vocabulary size is same + assert asr_model.tokenizer.tokenizer.vocab_size == asr_model2.tokenizer.tokenizer.vocab_size + + # Make a copy of the tokenizer + new_tokenizer_dir = os.path.join(save_dir, 'tokenizer') + + os.makedirs(new_tokenizer_dir, exist_ok=True) + new_tokenizer_path = os.path.join(new_tokenizer_dir, 'vocab.txt') + with open(new_tokenizer_path, 'w') as f: + for v in asr_model2.tokenizer.tokenizer.get_vocab(): + f.write(f"{v}\n") + + # Add some new tokens too + f.write("^\n") + f.write("^^\n") + f.write("^^^\n") + + assert os.path.exists(new_tokenizer_path) + + # change vocabulary + asr_model2.change_vocabulary(new_tokenizer_dir, new_tokenizer_type='wpe') + assert asr_model.tokenizer.vocab_size != asr_model2.tokenizer.vocab_size + + new_save_path = os.path.join(save_dir, 'temp2.nemo') + asr_model2.save_to(new_save_path) + + asr_model3 = EncDecCTCModelBPE.restore_from(new_save_path) + assert isinstance(asr_model3, EncDecCTCModelBPE) + + # Check if vocabulary size is same + assert asr_model2.tokenizer.tokenizer.vocab_size == asr_model3.tokenizer.tokenizer.vocab_size + assert asr_model2.vocab_path != asr_model3.vocab_path + + # Model PT level checks + assert len(asr_model2.artifacts) == 1 + + finally: + os.rename(old_tokenizer_dir + '.bkp', old_tokenizer_dir) + + @pytest.mark.unit + def test_EncDecCTCDatasetConfig_for_AudioToBPEDataset(self): + # ignore some additional arguments as dataclass is generic + IGNORE_ARGS = [ + 'is_tarred', + 'num_workers', + 'batch_size', + 'tarred_audio_filepaths', + 'shuffle', + 'pin_memory', + 'drop_last', + 'tarred_shard_strategy', + 'shuffle_n', + 'parser', + 'normalize', + 'unk_index', + 'pad_id', + 'bos_id', + 'eos_id', + 'blank_index', + ] + + REMAP_ARGS = {'trim_silence': 'trim', 'labels': 'tokenizer'} + + result = assert_dataclass_signature_match( + audio_to_text.AudioToBPEDataset, + configs.EncDecCTCDatasetConfig, + ignore_args=IGNORE_ARGS, + remap_args=REMAP_ARGS, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_EncDecCTCDatasetConfig_for_TarredAudioToBPEDataset(self): + # ignore some additional arguments as dataclass is generic + IGNORE_ARGS = [ + 'is_tarred', + 'num_workers', + 'batch_size', + 'shuffle', + 'pin_memory', + 'drop_last', + 'parser', + 'normalize', + 'unk_index', + 'pad_id', + 'bos_id', + 'eos_id', + 'blank_index', + 'global_rank', + 'world_size', + ] + + REMAP_ARGS = { + 'trim_silence': 'trim', + 'tarred_audio_filepaths': 'audio_tar_filepaths', + 'tarred_shard_strategy': 'shard_strategy', + 'shuffle_n': 'shuffle', + 'labels': 'tokenizer', + } + + result = assert_dataclass_signature_match( + audio_to_text.TarredAudioToBPEDataset, + configs.EncDecCTCDatasetConfig, + ignore_args=IGNORE_ARGS, + remap_args=REMAP_ARGS, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_ctcencdec_model.py b/models/tts/NeMo/tests/collections/asr/test_asr_ctcencdec_model.py new file mode 100644 index 0000000000..be9c502f23 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_ctcencdec_model.py @@ -0,0 +1,305 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 copy + +import pytest +import torch +from omegaconf import DictConfig, OmegaConf, open_dict + +import nemo.collections.asr as nemo_asr +from nemo.collections.asr.data import audio_to_text +from nemo.collections.asr.models import EncDecCTCModel, configs +from nemo.utils.config_utils import assert_dataclass_signature_match, update_model_config + + +@pytest.fixture() +def asr_model(): + preprocessor = {'_target_': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor'} + encoder = { + '_target_': 'nemo.collections.asr.modules.ConvASREncoder', + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 1024, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + } + + decoder = { + '_target_': 'nemo.collections.asr.modules.ConvASRDecoder', + 'feat_in': 1024, + 'num_classes': 28, + 'vocabulary': [ + ' ', + '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', + "'", + ], + } + modelConfig = DictConfig( + {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)} + ) + + model_instance = EncDecCTCModel(cfg=modelConfig) + return model_instance + + +class TestEncDecCTCModel: + @pytest.mark.unit + def test_constructor(self, asr_model): + asr_model.train() + # TODO: make proper config and assert correct number of weights + # Check to/from config_dict: + confdict = asr_model.to_config_dict() + instance2 = EncDecCTCModel.from_config_dict(confdict) + assert isinstance(instance2, EncDecCTCModel) + + @pytest.mark.unit + def test_forward(self, asr_model): + asr_model = asr_model.eval() + + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + + with torch.no_grad(): + # batch size 1 + logprobs_instance = [] + for i in range(input_signal.size(0)): + logprobs_ins, _, _ = asr_model.forward( + input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1] + ) + logprobs_instance.append(logprobs_ins) + print(len(logprobs_ins)) + logprobs_instance = torch.cat(logprobs_instance, 0) + + # batch size 4 + logprobs_batch, _, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length) + + assert logprobs_instance.shape == logprobs_batch.shape + diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + diff = torch.max(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + + @pytest.mark.unit + def test_vocab_change(self, asr_model): + old_vocab = copy.deepcopy(asr_model.decoder.vocabulary) + nw1 = asr_model.num_weights + asr_model.change_vocabulary(new_vocabulary=old_vocab) + # No change + assert nw1 == asr_model.num_weights + new_vocab = copy.deepcopy(old_vocab) + new_vocab.append('!') + new_vocab.append('$') + new_vocab.append('@') + asr_model.change_vocabulary(new_vocabulary=new_vocab) + # fully connected + bias + assert asr_model.num_weights == nw1 + 3 * (asr_model.decoder._feat_in + 1) + + @pytest.mark.unit + def test_change_conv_asr_se_context_window(self, asr_model): + old_cfg = copy.deepcopy(asr_model.cfg) + asr_model.change_conv_asr_se_context_window(context_window=32) # 32 * 0.01s context + new_config = asr_model.cfg + + assert old_cfg.encoder.jasper[0].se_context_size == -1 + assert new_config.encoder.jasper[0].se_context_size == 32 + + for name, m in asr_model.encoder.named_modules(): + if type(m).__class__.__name__ == 'SqueezeExcite': + assert m.context_window == 32 + + @pytest.mark.unit + def test_change_conv_asr_se_context_window_no_config_update(self, asr_model): + old_cfg = copy.deepcopy(asr_model.cfg) + asr_model.change_conv_asr_se_context_window(context_window=32, update_config=False) # 32 * 0.01s context + new_config = asr_model.cfg + + assert old_cfg.encoder.jasper[0].se_context_size == -1 + assert new_config.encoder.jasper[0].se_context_size == -1 # no change + + for name, m in asr_model.encoder.named_modules(): + if type(m).__class__.__name__ == 'SqueezeExcite': + assert m.context_window == 32 + + @pytest.mark.unit + def test_dataclass_instantiation(self, asr_model): + model_cfg = configs.EncDecCTCModelConfig() + + # Update mandatory values + vocabulary = asr_model.decoder.vocabulary + model_cfg.model.labels = vocabulary + + # Update encoder + model_cfg.model.encoder.activation = 'relu' + model_cfg.model.encoder.feat_in = 64 + model_cfg.model.encoder.jasper = [ + nemo_asr.modules.conv_asr.JasperEncoderConfig( + filters=1024, + repeat=1, + kernel=[1], + stride=[1], + dilation=[1], + dropout=0.0, + residual=False, + se=True, + se_context_size=-1, + ) + ] + + # Update decoder + model_cfg.model.decoder.feat_in = 1024 + model_cfg.model.decoder.num_classes = 28 + model_cfg.model.decoder.vocabulary = vocabulary + + # Construct the model + asr_cfg = OmegaConf.create({'model': asr_model.cfg}) + model_cfg_v1 = update_model_config(model_cfg, asr_cfg) + new_model = EncDecCTCModel(cfg=model_cfg_v1.model) + + assert new_model.num_weights == asr_model.num_weights + # trainer and exp manager should be there + # assert 'trainer' in model_cfg_v1 + # assert 'exp_manager' in model_cfg_v1 + # datasets and optim/sched should not be there after ModelPT.update_model_dataclass() + assert 'train_ds' not in model_cfg_v1.model + assert 'validation_ds' not in model_cfg_v1.model + assert 'test_ds' not in model_cfg_v1.model + assert 'optim' not in model_cfg_v1.model + + # Construct the model, without dropping additional keys + asr_cfg = OmegaConf.create({'model': asr_model.cfg}) + model_cfg_v2 = update_model_config(model_cfg, asr_cfg, drop_missing_subconfigs=False) + + # Assert all components are in config + # assert 'trainer' in model_cfg_v2 + # assert 'exp_manager' in model_cfg_v2 + assert 'train_ds' in model_cfg_v2.model + assert 'validation_ds' in model_cfg_v2.model + assert 'test_ds' in model_cfg_v2.model + assert 'optim' in model_cfg_v2.model + + # Remove extra components (optim and sched can be kept without issue) + with open_dict(model_cfg_v2.model): + model_cfg_v2.model.pop('train_ds') + model_cfg_v2.model.pop('validation_ds') + model_cfg_v2.model.pop('test_ds') + + new_model = EncDecCTCModel(cfg=model_cfg_v2.model) + + assert new_model.num_weights == asr_model.num_weights + # trainer and exp manager should be there + + @pytest.mark.unit + def test_EncDecCTCDatasetConfig_for_AudioToCharDataset(self): + # ignore some additional arguments as dataclass is generic + IGNORE_ARGS = [ + 'is_tarred', + 'num_workers', + 'batch_size', + 'tarred_audio_filepaths', + 'shuffle', + 'pin_memory', + 'drop_last', + 'tarred_shard_strategy', + 'shuffle_n', + 'use_start_end_token', + 'use_start_end_token', + ] + + REMAP_ARGS = {'trim_silence': 'trim'} + + result = assert_dataclass_signature_match( + audio_to_text.AudioToCharDataset, + configs.EncDecCTCDatasetConfig, + ignore_args=IGNORE_ARGS, + remap_args=REMAP_ARGS, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_EncDecCTCDatasetConfig_for_TarredAudioToCharDataset(self): + # ignore some additional arguments as dataclass is generic + IGNORE_ARGS = [ + 'is_tarred', + 'num_workers', + 'batch_size', + 'shuffle', + 'pin_memory', + 'drop_last', + 'global_rank', + 'world_size', + 'use_start_end_token', + ] + + REMAP_ARGS = { + 'trim_silence': 'trim', + 'tarred_audio_filepaths': 'audio_tar_filepaths', + 'tarred_shard_strategy': 'shard_strategy', + 'shuffle_n': 'shuffle', + } + + result = assert_dataclass_signature_match( + audio_to_text.TarredAudioToCharDataset, + configs.EncDecCTCDatasetConfig, + ignore_args=IGNORE_ARGS, + remap_args=REMAP_ARGS, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_datasets.py b/models/tts/NeMo/tests/collections/asr/test_asr_datasets.py new file mode 100644 index 0000000000..b2a3ccfd09 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_datasets.py @@ -0,0 +1,127 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 copy +import os + +import pytest +from omegaconf import OmegaConf + +from nemo.collections.asr.data.audio_to_text import TarredAudioToBPEDataset, TarredAudioToCharDataset +from nemo.collections.asr.data.audio_to_text_dataset import inject_dataloader_value_from_model_config +from nemo.collections.common import tokenizers +from nemo.utils import logging + + +class TestASRDatasets: + labels = [ + " ", + "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", + "'", + ] + + @pytest.mark.unit + def test_tarred_dataset(self, test_data_dir): + manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json')) + + # Test braceexpand loading + tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar')) + ds_braceexpand = TarredAudioToCharDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000 + ) + assert len(ds_braceexpand) == 32 + count = 0 + for _ in ds_braceexpand: + count += 1 + assert count == 32 + + # Test loading via list + tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)] + ds_list_load = TarredAudioToCharDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, sample_rate=16000 + ) + count = 0 + for _ in ds_list_load: + count += 1 + assert count == 32 + + @pytest.mark.unit + def test_mismatch_in_model_dataloader_config(self, caplog): + logging._logger.propagate = True + caplog.set_level(logging.WARNING) + + model_cfg = OmegaConf.create(dict(labels=OmegaConf.create(["a", "b", "c"]))) + dataloader_cfg = OmegaConf.create(dict(labels=copy.deepcopy(self.labels))) + + inject_dataloader_value_from_model_config(model_cfg, dataloader_cfg, key='labels') + + assert ( + """`labels` is explicitly provided to the data loader, and is different from the `labels` provided at the model level config.""" + in caplog.text + ) + + logging._logger.propagate = False + + @pytest.mark.unit + def test_tarred_bpe_dataset(self, test_data_dir): + manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json')) + + tokenizer_path = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt') + tokenizer = tokenizers.AutoTokenizer(pretrained_model_name='bert-base-cased', vocab_file=tokenizer_path) + + # Test braceexpand loading + tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar')) + ds_braceexpand = TarredAudioToBPEDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, tokenizer=tokenizer, sample_rate=16000 + ) + assert len(ds_braceexpand) == 32 + count = 0 + for _ in ds_braceexpand: + count += 1 + assert count == 32 + + # Test loading via list + tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)] + ds_list_load = TarredAudioToBPEDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, tokenizer=tokenizer, sample_rate=16000 + ) + count = 0 + for _ in ds_list_load: + count += 1 + assert count == 32 diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_exportables.py b/models/tts/NeMo/tests/collections/asr/test_asr_exportables.py new file mode 100644 index 0000000000..0ebd23345a --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_exportables.py @@ -0,0 +1,301 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import tempfile + +import onnx +import pytest +from omegaconf import DictConfig, ListConfig + +from nemo.collections.asr.models import EncDecClassificationModel, EncDecCTCModel, EncDecSpeakerLabelModel +from nemo.collections.asr.modules import ConvASRDecoder, ConvASREncoder + + +class TestExportable: + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_EncDecCTCModel_export_to_onnx(self): + model_config = DictConfig( + { + 'preprocessor': DictConfig(self.preprocessor), + 'encoder': DictConfig(self.encoder_dict), + 'decoder': DictConfig(self.decoder_dict), + } + ) + model = EncDecCTCModel(cfg=model_config) + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'qn.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'audio_signal' + assert onnx_model.graph.output[0].name == 'logprobs' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_EncDecClassificationModel_export_to_onnx(self, speech_classification_model): + model = speech_classification_model.train() + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'edc.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'audio_signal' + assert onnx_model.graph.output[0].name == 'logits' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_EncDecSpeakerLabelModel_export_to_onnx(self, speaker_label_model): + model = speaker_label_model.train() + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'sl.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'audio_signal' + assert onnx_model.graph.output[0].name == 'logits' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_EncDecCitrinetModel_export_to_onnx(self, citrinet_model): + model = citrinet_model.train() + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join('.', 'citri.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'audio_signal' + assert onnx_model.graph.output[0].name == 'logprobs' + + def setup_method(self): + self.preprocessor = { + 'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', + 'params': dict({}), + } + + self.encoder_dict = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 1024, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + self.decoder_dict = { + 'cls': 'nemo.collections.asr.modules.ConvASRDecoder', + 'params': { + 'feat_in': 1024, + 'num_classes': 28, + 'vocabulary': [ + ' ', + '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', + "'", + ], + }, + } + + +@pytest.fixture() +def speech_classification_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 32, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.ConvASRDecoderClassification', + 'params': {'feat_in': 32, 'num_classes': 30,}, + } + + modelConfig = DictConfig( + { + 'preprocessor': DictConfig(preprocessor), + 'encoder': DictConfig(encoder), + 'decoder': DictConfig(decoder), + 'labels': ListConfig(["dummy_cls_{}".format(i + 1) for i in range(30)]), + } + ) + model = EncDecClassificationModel(cfg=modelConfig) + return model + + +@pytest.fixture() +def speaker_label_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 512, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': False, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.SpeakerDecoder', + 'params': {'feat_in': 512, 'num_classes': 2, 'pool_mode': 'xvector', 'emb_sizes': [1024]}, + } + + modelConfig = DictConfig( + {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)} + ) + speaker_model = EncDecSpeakerLabelModel(cfg=modelConfig) + return speaker_model + + +@pytest.fixture() +def citrinet_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 80, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 512, + 'repeat': 1, + 'kernel': [5], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + }, + { + 'filters': 512, + 'repeat': 5, + 'kernel': [11], + 'stride': [2], + 'dilation': [1], + 'dropout': 0.1, + 'residual': True, + 'separable': True, + 'se': True, + 'se_context_size': -1, + 'stride_last': True, + 'residual_mode': 'stride_add', + }, + { + 'filters': 512, + 'repeat': 5, + 'kernel': [13], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.1, + 'residual': True, + 'separable': True, + 'se': True, + 'se_context_size': -1, + }, + { + 'filters': 640, + 'repeat': 1, + 'kernel': [41], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': True, + 'separable': True, + 'se': True, + 'se_context_size': -1, + }, + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.ConvASRDecoder', + 'params': {'feat_in': 640, 'num_classes': 1024, 'vocabulary': list(chr(i % 28) for i in range(0, 1024))}, + } + + modelConfig = DictConfig( + {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)} + ) + citri_model = EncDecSpeakerLabelModel(cfg=modelConfig) + return citri_model diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_filterbankfeatures_seq_len.py b/models/tts/NeMo/tests/collections/asr/test_asr_filterbankfeatures_seq_len.py new file mode 100644 index 0000000000..24ba584985 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_filterbankfeatures_seq_len.py @@ -0,0 +1,99 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 librosa +import numpy as np +import pytest +import torch + +from nemo.collections.asr.parts.preprocessing.features import FilterbankFeatures + + +class TestFilterbankFeatures: + @pytest.mark.unit + def test_seq_len(self): + fb_module = FilterbankFeatures(exact_pad=False, pad_to=1) + test_1 = torch.randn(1, 800) + test_1_len = torch.tensor([800]) + fb_spec, fb_len = fb_module(test_1, test_1_len) + assert fb_spec.shape[2] == fb_len[0], f"{fb_spec.shape} != {fb_len}" + librosa_spec = librosa.stft(test_1.cpu().detach().numpy().squeeze(), n_fft=512, hop_length=160, win_length=320) + + assert librosa_spec.shape[1] == fb_spec.shape[2], f"{librosa_spec.shape} != {fb_spec.shape}" + + @pytest.mark.unit + def test_random_stft_sizes(self): + for _ in range(5): + nfft = 2 ** np.random.randint(7, 12) + window_size = np.random.randint(100, nfft) + hop_size = np.random.randint(64, window_size) + fb_module = FilterbankFeatures( + exact_pad=False, + pad_to=1, + n_fft=nfft, + n_window_size=window_size, + n_window_stride=hop_size, + normalize=False, + ) + audio_length = np.random.randint(nfft, 2 ** 16) + test_1 = torch.randn(1, audio_length) + test_1_len = torch.tensor([audio_length]) + fb_spec, fb_len = fb_module(test_1, test_1_len) + assert ( + fb_spec.shape[2] == fb_len[0] + ), f"{fb_spec.shape} != {fb_len}: {nfft}, {window_size}, {hop_size}, {audio_length}" + + librosa_spec = librosa.stft( + test_1.cpu().detach().numpy().squeeze(), n_fft=nfft, hop_length=hop_size, win_length=window_size + ) + + assert ( + librosa_spec.shape[1] == fb_spec.shape[2] + ), f"{librosa_spec.shape} != {fb_spec.shape}: {nfft}, {window_size}, {hop_size}, {audio_length}" + + @pytest.mark.unit + def test_random_stft_sizes_exact_pad(self): + for _ in range(5): + nfft = 2 ** np.random.randint(7, 12) + window_size = np.random.randint(100, nfft) + hop_size = np.random.randint(64, window_size) + if hop_size % 2 == 1: + hop_size = hop_size - 1 + fb_module = FilterbankFeatures( + exact_pad=True, + pad_to=1, + n_fft=nfft, + n_window_size=window_size, + n_window_stride=hop_size, + normalize=False, + ) + audio_length = np.random.randint(nfft, 2 ** 16) + test_1 = torch.randn(1, audio_length) + test_1_len = torch.tensor([audio_length]) + fb_spec, fb_len = fb_module(test_1, test_1_len) + assert ( + fb_spec.shape[2] == fb_len[0] + ), f"{fb_spec.shape} != {fb_len}: {nfft}, {window_size}, {hop_size}, {audio_length}" + + test_2 = test_1.cpu().detach().numpy().squeeze() + test_2 = np.pad(test_2, int((nfft - hop_size) // 2), mode="reflect") + librosa_spec = librosa.stft(test_2, n_fft=nfft, hop_length=hop_size, win_length=window_size, center=False,) + + assert ( + fb_spec.shape[2] == librosa_spec.shape[1] + ), f"{fb_spec.shape} != {librosa_spec.shape}: {nfft}, {window_size}, {hop_size}, {audio_length}" + + assert ( + fb_spec.shape[2] == audio_length // hop_size + ), f"{fb_spec.shape}, {nfft}, {window_size}, {hop_size}, {audio_length}, {audio_length // hop_size}" diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_metrics.py b/models/tts/NeMo/tests/collections/asr/test_asr_metrics.py new file mode 100644 index 0000000000..7f78c1d720 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_metrics.py @@ -0,0 +1,172 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 random + +import pytest +import torch + +from nemo.collections.asr.metrics.wer import WER, word_error_rate +from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis + + +class TestWordErrorRate: + + vocabulary = [ + " ", + "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 __string_to_ctc_tensor(self, txt: str) -> torch.Tensor: + # This function emulates how CTC output could like for txt + blank_id = len(self.vocabulary) + char_to_ind = dict([(self.vocabulary[i], i) for i in range(len(self.vocabulary))]) + string_in_id_form = [char_to_ind[c] for c in txt] + ctc_list = [] + prev_id = -1 + for c in string_in_id_form: + # when character is repeated we need to insert CTC blank symbol + if c != prev_id: + ctc_list.append(c) + else: + ctc_list.append(blank_id) + ctc_list.append(c) + prev_id = c + return torch.Tensor(ctc_list).unsqueeze(0) + + def __reference_string_to_tensor(self, txt: str) -> torch.Tensor: + # Reference tensors aren't produced by CTC logic + char_to_ind = dict([(self.vocabulary[i], i) for i in range(len(self.vocabulary))]) + string_in_id_form = [char_to_ind[c] for c in txt] + return torch.Tensor(string_in_id_form).unsqueeze(0) + + def get_wer(self, wer, prediction: str, reference: str): + wer( + predictions=self.__string_to_ctc_tensor(prediction), + targets=self.__reference_string_to_tensor(reference), + target_lengths=torch.tensor([len(reference)]), + ) + res, _, _ = wer.compute() + res = res.detach().cpu() + # return res[0] / res[1] + return res.item() + + @pytest.mark.unit + def test_wer_function(self): + assert word_error_rate(hypotheses=['cat'], references=['cot']) == 1.0 + assert word_error_rate(hypotheses=['GPU'], references=['G P U']) == 1.0 + assert word_error_rate(hypotheses=['G P U'], references=['GPU']) == 3.0 + assert word_error_rate(hypotheses=['ducati motorcycle'], references=['motorcycle']) == 1.0 + assert word_error_rate(hypotheses=['ducati motorcycle'], references=['ducuti motorcycle']) == 0.5 + assert word_error_rate(hypotheses=['a B c'], references=['a b c']) == 1.0 / 3.0 + + @pytest.mark.unit + def test_wer_metric_simple(self): + wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True) + + assert self.get_wer(wer, 'cat', 'cot') == 1.0 + assert self.get_wer(wer, 'gpu', 'g p u') == 1.0 + assert self.get_wer(wer, 'g p u', 'gpu') == 3.0 + assert self.get_wer(wer, 'ducati motorcycle', 'motorcycle') == 1.0 + assert self.get_wer(wer, 'ducati motorcycle', 'ducuti motorcycle') == 0.5 + assert abs(self.get_wer(wer, 'a f c', 'a b c') - 1.0 / 3.0) < 1e-6 + + @pytest.mark.unit + def test_wer_metric_randomized(self): + """This test relies on correctness of word_error_rate function.""" + + def __randomString(N): + return ''.join(random.choice(''.join(self.vocabulary)) for i in range(N)) + + wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True) + + for test_id in range(256): + n1 = random.randint(1, 512) + n2 = random.randint(1, 512) + s1 = __randomString(n1) + s2 = __randomString(n2) + # skip empty strings as reference + if s2.strip(): + assert ( + abs( + self.get_wer(wer, prediction=s1, reference=s2) + - word_error_rate(hypotheses=[s1], references=[s2]) + ) + < 1e-6 + ) + + @pytest.mark.unit + def test_wer_metric_decode(self): + wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True) + + tokens = self.__string_to_ctc_tensor('cat')[0].int().numpy().tolist() + assert tokens == [3, 1, 20] + + tokens_decoded = wer.decode_ids_to_tokens(tokens) + assert tokens_decoded == ['c', 'a', 't'] + + str_decoded = wer.decode_tokens_to_str(tokens) + assert str_decoded == 'cat' + + @pytest.mark.unit + def test_wer_metric_return_hypothesis(self): + wer = WER(vocabulary=self.vocabulary, batch_dim_index=0, use_cer=False, ctc_decode=True) + + tensor = self.__string_to_ctc_tensor('cat').int() + + # pass batchsize 1 tensor, get back list of length 1 Hypothesis + hyp = wer.ctc_decoder_predictions_tensor(tensor, return_hypotheses=True) + hyp = hyp[0] + assert isinstance(hyp, Hypothesis) + + assert hyp.y_sequence is None + assert hyp.score == -1.0 + assert hyp.text == 'cat' + assert hyp.alignments == [3, 1, 20] + assert hyp.length == 0 + + length = torch.tensor([tensor.shape[-1]], dtype=torch.long) + + # pass batchsize 1 tensor, get back list of length 1 Hypothesis [add length info] + hyp = wer.ctc_decoder_predictions_tensor(tensor, predictions_len=length, return_hypotheses=True) + hyp = hyp[0] + assert isinstance(hyp, Hypothesis) + assert hyp.length == 3 diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_modules.py b/models/tts/NeMo/tests/collections/asr/test_asr_modules.py new file mode 100644 index 0000000000..b4dd12b3c7 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_modules.py @@ -0,0 +1,325 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest +import torch +from omegaconf import OmegaConf + +from nemo.collections.asr import modules +from nemo.collections.asr.parts.utils.rnnt_utils import Hypothesis +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ +from nemo.utils import config_utils, logging + + +class TestASRModulesBasicTests: + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor1(self): + # Test 1 that should test the pure stft implementation as much as possible + instance1 = modules.AudioToMelSpectrogramPreprocessor( + dither=0, stft_conv=False, mag_power=1.0, normalize=False, preemph=0.0, log=False, pad_to=0 + ) + instance2 = modules.AudioToMelSpectrogramPreprocessor( + dither=0, stft_conv=True, mag_power=1.0, normalize=False, preemph=0.0, log=False, pad_to=0 + ) + + # Ensure that the two functions behave similarily + for _ in range(10): + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + res1, length1 = instance1(input_signal=input_signal, length=length) + res2, length2 = instance2(input_signal=input_signal, length=length) + for len1, len2 in zip(length1, length2): + assert len1 == len2 + assert res1.shape == res2.shape + diff = torch.mean(torch.abs(res1 - res2)) + assert diff <= 1e-3 + diff = torch.max(torch.abs(res1 - res2)) + assert diff <= 1e-2 + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_config(self): + # Test that dataclass matches signature of module + result = config_utils.assert_dataclass_signature_match( + modules.AudioToMelSpectrogramPreprocessor, + modules.audio_preprocessing.AudioToMelSpectrogramPreprocessorConfig, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor_batch(self): + # Test 1 that should test the pure stft implementation as much as possible + instance1 = modules.AudioToMelSpectrogramPreprocessor(normalize="per_feature", dither=0, pad_to=0) + + # Ensure that the two functions behave similarily + for _ in range(10): + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + + with torch.no_grad(): + # batch size 1 + res_instance, length_instance = [], [] + for i in range(input_signal.size(0)): + res_ins, length_ins = instance1(input_signal=input_signal[i : i + 1], length=length[i : i + 1]) + res_instance.append(res_ins) + length_instance.append(length_ins) + + res_instance = torch.cat(res_instance, 0) + length_instance = torch.cat(length_instance, 0) + + # batch size 4 + res_batch, length_batch = instance1(input_signal=input_signal, length=length) + + assert res_instance.shape == res_batch.shape + assert length_instance.shape == length_batch.shape + diff = torch.mean(torch.abs(res_instance - res_batch)) + assert diff <= 1e-3 + diff = torch.max(torch.abs(res_instance - res_batch)) + assert diff <= 1e-3 + + @pytest.mark.unit + def test_AudioToMelSpectrogramPreprocessor2(self): + # Test 2 that should test the stft implementation as used in ASR models + instance1 = modules.AudioToMelSpectrogramPreprocessor(dither=0, stft_conv=False) + instance2 = modules.AudioToMelSpectrogramPreprocessor(dither=0, stft_conv=True) + + # Ensure that the two functions behave similarily + for _ in range(5): + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + res1, length1 = instance1(input_signal=input_signal, length=length) + res2, length2 = instance2(input_signal=input_signal, length=length) + for len1, len2 in zip(length1, length2): + assert len1 == len2 + assert res1.shape == res2.shape + diff = torch.mean(torch.abs(res1 - res2)) + assert diff <= 3e-3 + diff = torch.max(torch.abs(res1 - res2)) + assert diff <= 3 + + @pytest.mark.unit + def test_SpectrogramAugmentationr(self): + # Make sure constructor works + instance1 = modules.SpectrogramAugmentation( + freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=False + ) + assert isinstance(instance1, modules.SpectrogramAugmentation) + + # Make sure forward doesn't throw with expected input + instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0) + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + res0 = instance0(input_signal=input_signal, length=length) + res = instance1(input_spec=res0[0]) + + assert res.shape == res0[0].shape + + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + def test_SpectrogramAugmentationr_numba_kernel(self, caplog): + numba_utils.skip_numba_cuda_test_if_unsupported(__NUMBA_MINIMUM_VERSION__) + + logging._logger.propagate = True + original_verbosity = logging.get_verbosity() + logging.set_verbosity(logging.DEBUG) + caplog.set_level(logging.DEBUG) + + # Make sure constructor works + instance1 = modules.SpectrogramAugmentation( + freq_masks=10, time_masks=3, rect_masks=3, use_numba_spec_augment=True + ) + assert isinstance(instance1, modules.SpectrogramAugmentation) + + # Make sure forward doesn't throw with expected input + instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0) + input_signal = torch.randn(size=(8, 512)) + length = torch.randint(low=161, high=500, size=[8]) + res0 = instance0(input_signal=input_signal, length=length) + res = instance1(input_spec=res0[0], length=length) + + assert res.shape == res0[0].shape + + # check tha numba kernel debug message indicates that it is available for use + assert """Numba SpecAugment kernel is available""" in caplog.text + + logging._logger.propagate = False + logging.set_verbosity(original_verbosity) + + @pytest.mark.unit + def test_SpectrogramAugmentationr_config(self): + # Test that dataclass matches signature of module + result = config_utils.assert_dataclass_signature_match( + modules.SpectrogramAugmentation, modules.audio_preprocessing.SpectrogramAugmentationConfig, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_CropOrPadSpectrogramAugmentation(self): + # Make sure constructor works + audio_length = 128 + instance1 = modules.CropOrPadSpectrogramAugmentation(audio_length=audio_length) + assert isinstance(instance1, modules.CropOrPadSpectrogramAugmentation) + + # Make sure forward doesn't throw with expected input + instance0 = modules.AudioToMelSpectrogramPreprocessor(dither=0) + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + res0 = instance0(input_signal=input_signal, length=length) + res, new_length = instance1(input_signal=res0[0], length=length) + + assert res.shape == torch.Size([4, 64, audio_length]) + assert all(new_length == torch.tensor([128] * 4)) + + @pytest.mark.unit + def test_CropOrPadSpectrogramAugmentation_config(self): + # Test that dataclass matches signature of module + result = config_utils.assert_dataclass_signature_match( + modules.CropOrPadSpectrogramAugmentation, + modules.audio_preprocessing.CropOrPadSpectrogramAugmentationConfig, + ) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_RNNTDecoder(self): + vocab = list(range(10)) + vocab = [str(x) for x in vocab] + vocab_size = len(vocab) + + pred_config = OmegaConf.create( + { + '_target_': 'nemo.collections.asr.modules.RNNTDecoder', + 'prednet': {'pred_hidden': 32, 'pred_rnn_layers': 1,}, + 'vocab_size': vocab_size, + 'blank_as_pad': True, + } + ) + + prednet = modules.RNNTDecoder.from_config_dict(pred_config) + + # num params + pred_hidden = pred_config.prednet.pred_hidden + embed = (vocab_size + 1) * pred_hidden # embedding with blank + rnn = ( + 2 * 4 * (pred_hidden * pred_hidden + pred_hidden) + ) # (ih + hh) * (ifco gates) * (indim * hiddendim + bias) + assert prednet.num_weights == (embed + rnn) + + # State initialization + x_ = torch.zeros(4, dtype=torch.float32) + states = prednet.initialize_state(x_) + + for state_i in states: + assert state_i.dtype == x_.dtype + assert state_i.device == x_.device + assert state_i.shape[1] == len(x_) + + # Blank hypotheses test + blank = vocab_size + hyp = Hypothesis(score=0.0, y_sequence=[blank]) + cache = {} + pred, states, _ = prednet.score_hypothesis(hyp, cache) + + assert pred.shape == torch.Size([1, 1, pred_hidden]) + assert len(states) == 2 + for state_i in states: + assert state_i.dtype == pred.dtype + assert state_i.device == pred.device + assert state_i.shape[1] == len(pred) + + # Blank stateless predict + g, states = prednet.predict(y=None, state=None, add_sos=False, batch_size=1) + + assert g.shape == torch.Size([1, 1, pred_hidden]) + assert len(states) == 2 + for state_i in states: + assert state_i.dtype == g.dtype + assert state_i.device == g.device + assert state_i.shape[1] == len(g) + + # Blank stateful predict + g, states2 = prednet.predict(y=None, state=states, add_sos=False, batch_size=1) + + assert g.shape == torch.Size([1, 1, pred_hidden]) + assert len(states2) == 2 + for state_i, state_j in zip(states, states2): + assert (state_i - state_j).square().sum().sqrt() > 0.0 + + # Predict with token and state + token = torch.full([1, 1], fill_value=0, dtype=torch.long) + g, states = prednet.predict(y=token, state=states2, add_sos=False, batch_size=None) + + assert g.shape == torch.Size([1, 1, pred_hidden]) + assert len(states) == 2 + + # Predict with blank token and no state + token = torch.full([1, 1], fill_value=blank, dtype=torch.long) + g, states = prednet.predict(y=token, state=None, add_sos=False, batch_size=None) + + assert g.shape == torch.Size([1, 1, pred_hidden]) + assert len(states) == 2 + + @pytest.mark.unit + def test_RNNTJoint(self): + vocab = list(range(10)) + vocab = [str(x) for x in vocab] + vocab_size = len(vocab) + + batchsize = 4 + encoder_hidden = 64 + pred_hidden = 32 + joint_hidden = 16 + + joint_cfg = OmegaConf.create( + { + '_target_': 'nemo.collections.asr.modules.RNNTJoint', + 'num_classes': vocab_size, + 'vocabulary': vocab, + 'jointnet': { + 'encoder_hidden': encoder_hidden, + 'pred_hidden': pred_hidden, + 'joint_hidden': joint_hidden, + 'activation': 'relu', + }, + } + ) + + jointnet = modules.RNNTJoint.from_config_dict(joint_cfg) + + enc = torch.zeros(batchsize, encoder_hidden, 48) # [B, D1, T] + dec = torch.zeros(batchsize, pred_hidden, 24) # [B, D2, U] + + # forward call test + out = jointnet(encoder_outputs=enc, decoder_outputs=dec) + assert out.shape == torch.Size([batchsize, 48, 24, vocab_size + 1]) # [B, T, U, V + 1] + + # joint() step test + enc2 = enc.transpose(1, 2) # [B, T, D1] + dec2 = dec.transpose(1, 2) # [B, U, D2] + out2 = jointnet.joint(enc2, dec2) # [B, T, U, V + 1] + assert (out - out2).abs().sum() <= 1e-5 + + # assert vocab size + assert jointnet.num_classes_with_blank == vocab_size + 1 diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_regression_model.py b/models/tts/NeMo/tests/collections/asr/test_asr_regression_model.py new file mode 100644 index 0000000000..eacce0f543 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_regression_model.py @@ -0,0 +1,95 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +import pytest +from omegaconf import DictConfig + +from nemo.collections.asr.models.classification_models import EncDecRegressionModel + + +@pytest.fixture() +def speech_regression_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 32, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.conv_asr.ConvASRDecoderClassification', + 'params': {'feat_in': 32, 'return_logits': True, 'num_classes': 1}, + } + + modelConfig = DictConfig( + { + 'preprocessor': DictConfig(preprocessor), + 'encoder': DictConfig(encoder), + 'decoder': DictConfig(decoder), + 'labels': None, + 'is_regression_task': True, + } + ) + model = EncDecRegressionModel(cfg=modelConfig) + return model + + +class TestEncDecRegressionModel: + @pytest.mark.unit + def test_constructor(self, speech_regression_model): + asr_model = speech_regression_model.train() + + conv_cnt = (64 * 32 * 1 + 32) + (64 * 1 * 1 + 32) # separable kernel + bias + pointwise kernel + bias + bn_cnt = (4 * 32) * 2 # 2 * moving averages + dec_cnt = 32 * 1 + 1 # fc + bias + + param_count = conv_cnt + bn_cnt + dec_cnt + assert asr_model.num_weights == param_count + + # Check to/from config_dict: + confdict = asr_model.to_config_dict() + instance2 = EncDecRegressionModel.from_config_dict(confdict) + + assert isinstance(instance2, EncDecRegressionModel) + + @pytest.mark.unit + def test_transcription(self, speech_regression_model, test_data_dir): + + audio_filenames = ['an22-flrp-b.wav', 'an90-fbbh-b.wav'] + audio_paths = [os.path.join(test_data_dir, "asr", "train", "an4", "wav", fp) for fp in audio_filenames] + + model = speech_regression_model.eval() + + # Test Top 1 classification transcription + results = model.transcribe(audio_paths, batch_size=2) + assert len(results) == 2 diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_rnnt_encdec_model.py b/models/tts/NeMo/tests/collections/asr/test_asr_rnnt_encdec_model.py new file mode 100644 index 0000000000..15124157a6 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_rnnt_encdec_model.py @@ -0,0 +1,276 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 copy + +import pytest +import torch +from omegaconf import DictConfig, ListConfig + +from nemo.collections.asr.models import EncDecRNNTModel +from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode +from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ +from nemo.utils.config_utils import assert_dataclass_signature_match + +NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__) + + +@pytest.fixture() +def asr_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + + # fmt: off + labels = [' ', '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', "'", + ] + # fmt: on + + model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64} + + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': model_defaults['enc_hidden'], + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + '_target_': 'nemo.collections.asr.modules.RNNTDecoder', + 'prednet': {'pred_hidden': model_defaults['pred_hidden'], 'pred_rnn_layers': 1}, + } + + joint = { + '_target_': 'nemo.collections.asr.modules.RNNTJoint', + 'jointnet': {'joint_hidden': 32, 'activation': 'relu'}, + } + + decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}} + + modelConfig = DictConfig( + { + 'labels': ListConfig(labels), + 'preprocessor': DictConfig(preprocessor), + 'model_defaults': DictConfig(model_defaults), + 'encoder': DictConfig(encoder), + 'decoder': DictConfig(decoder), + 'joint': DictConfig(joint), + 'decoding': DictConfig(decoding), + } + ) + + model_instance = EncDecRNNTModel(cfg=modelConfig) + return model_instance + + +class TestEncDecRNNTModel: + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_constructor(self, asr_model): + asr_model.train() + # TODO: make proper config and assert correct number of weights + # Check to/from config_dict: + confdict = asr_model.to_config_dict() + instance2 = EncDecRNNTModel.from_config_dict(confdict) + assert isinstance(instance2, EncDecRNNTModel) + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_forward(self, asr_model): + asr_model = asr_model.eval() + + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + + asr_model.compute_eval_loss = False + + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + + with torch.no_grad(): + # batch size 1 + logprobs_instance = [] + for i in range(input_signal.size(0)): + logprobs_ins, _ = asr_model.forward( + input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1] + ) + logprobs_instance.append(logprobs_ins) + logprobs_instance = torch.cat(logprobs_instance, 0) + + # batch size 4 + logprobs_batch, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length) + + assert logprobs_instance.shape == logprobs_batch.shape + diff = torch.mean(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + diff = torch.max(torch.abs(logprobs_instance - logprobs_batch)) + assert diff <= 1e-6 + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_vocab_change(self, asr_model): + old_vocab = copy.deepcopy(asr_model.joint.vocabulary) + nw1 = asr_model.num_weights + asr_model.change_vocabulary(new_vocabulary=old_vocab) + # No change + assert nw1 == asr_model.num_weights + new_vocab = copy.deepcopy(old_vocab) + new_vocab.append('!') + new_vocab.append('$') + new_vocab.append('@') + asr_model.change_vocabulary(new_vocabulary=new_vocab) + # fully connected + bias + # rnn embedding + joint + bias + pred_embedding = 3 * (asr_model.decoder.pred_hidden) + joint_joint = 3 * (asr_model.joint.joint_hidden + 1) + assert asr_model.num_weights == (nw1 + (pred_embedding + joint_joint)) + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_change_conv_asr_se_context_window(self, asr_model): + old_cfg = copy.deepcopy(asr_model.cfg) + asr_model.change_conv_asr_se_context_window(context_window=32) # 32 * 0.01s context + new_config = asr_model.cfg + + assert old_cfg.encoder.jasper[0].se_context_size == -1 + assert new_config.encoder.jasper[0].se_context_size == 32 + + for name, m in asr_model.encoder.named_modules(): + if type(m).__class__.__name__ == 'SqueezeExcite': + assert m.context_window == 32 + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_change_conv_asr_se_context_window_no_config_update(self, asr_model): + old_cfg = copy.deepcopy(asr_model.cfg) + asr_model.change_conv_asr_se_context_window(context_window=32, update_config=False) # 32 * 0.01s context + new_config = asr_model.cfg + + assert old_cfg.encoder.jasper[0].se_context_size == -1 + assert new_config.encoder.jasper[0].se_context_size == -1 # no change + + for name, m in asr_model.encoder.named_modules(): + if type(m).__class__.__name__ == 'SqueezeExcite': + assert m.context_window == 32 + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_decoding_change(self, asr_model): + assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer) + + new_strategy = DictConfig({}) + new_strategy.strategy = 'greedy' + new_strategy.greedy = DictConfig({'max_symbols': 10}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer) + + new_strategy = DictConfig({}) + new_strategy.strategy = 'beam' + new_strategy.beam = DictConfig({'beam_size': 1}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "default" + + new_strategy = DictConfig({}) + new_strategy.strategy = 'beam' + new_strategy.beam = DictConfig({'beam_size': 2}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "default" + + new_strategy = DictConfig({}) + new_strategy.strategy = 'tsd' + new_strategy.beam = DictConfig({'beam_size': 2}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "tsd" + + new_strategy = DictConfig({}) + new_strategy.strategy = 'alsd' + new_strategy.beam = DictConfig({'beam_size': 2}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "alsd" + + @pytest.mark.unit + def test_GreedyRNNTInferConfig(self): + IGNORE_ARGS = ['decoder_model', 'joint_model', 'blank_index'] + + result = assert_dataclass_signature_match( + greedy_decode.GreedyRNNTInfer, greedy_decode.GreedyRNNTInferConfig, ignore_args=IGNORE_ARGS + ) + + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_GreedyBatchedRNNTInferConfig(self): + IGNORE_ARGS = ['decoder_model', 'joint_model', 'blank_index'] + + result = assert_dataclass_signature_match( + greedy_decode.GreedyBatchedRNNTInfer, greedy_decode.GreedyBatchedRNNTInferConfig, ignore_args=IGNORE_ARGS + ) + + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_BeamRNNTInferConfig(self): + IGNORE_ARGS = ['decoder_model', 'joint_model', 'blank_index'] + + result = assert_dataclass_signature_match( + beam_decode.BeamRNNTInfer, beam_decode.BeamRNNTInferConfig, ignore_args=IGNORE_ARGS + ) + + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None diff --git a/models/tts/NeMo/tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py b/models/tts/NeMo/tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py new file mode 100644 index 0000000000..294b03dd7d --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_asr_rnnt_encoder_model_bpe.py @@ -0,0 +1,246 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import shutil +import tempfile + +import pytest +import torch +from omegaconf import DictConfig + +from nemo.collections.asr.models.rnnt_bpe_models import EncDecRNNTBPEModel +from nemo.collections.asr.parts.submodules import rnnt_beam_decoding as beam_decode +from nemo.collections.asr.parts.submodules import rnnt_greedy_decoding as greedy_decode +from nemo.collections.common import tokenizers +from nemo.core.utils import numba_utils +from nemo.core.utils.numba_utils import __NUMBA_MINIMUM_VERSION__ + +NUMBA_RNNT_LOSS_AVAILABLE = numba_utils.numba_cuda_is_supported(__NUMBA_MINIMUM_VERSION__) + + +@pytest.fixture() +def asr_model(test_data_dir): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + + model_defaults = {'enc_hidden': 1024, 'pred_hidden': 64} + + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': model_defaults['enc_hidden'], + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + '_target_': 'nemo.collections.asr.modules.RNNTDecoder', + 'prednet': {'pred_hidden': model_defaults['pred_hidden'], 'pred_rnn_layers': 1,}, + } + + joint = { + '_target_': 'nemo.collections.asr.modules.RNNTJoint', + 'jointnet': {'joint_hidden': 32, 'activation': 'relu',}, + } + + decoding = {'strategy': 'greedy_batch', 'greedy': {'max_symbols': 30}} + + tokenizer = {'dir': os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128"), 'type': 'wpe'} + + modelConfig = DictConfig( + { + 'preprocessor': DictConfig(preprocessor), + 'model_defaults': DictConfig(model_defaults), + 'encoder': DictConfig(encoder), + 'decoder': DictConfig(decoder), + 'joint': DictConfig(joint), + 'tokenizer': DictConfig(tokenizer), + 'decoding': DictConfig(decoding), + } + ) + + model_instance = EncDecRNNTBPEModel(cfg=modelConfig) + return model_instance + + +class TestEncDecRNNTBPEModel: + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_constructor(self, asr_model): + asr_model.train() + # TODO: make proper config and assert correct number of weights + # Check to/from config_dict: + confdict = asr_model.to_config_dict() + instance2 = EncDecRNNTBPEModel.from_config_dict(confdict) + assert isinstance(instance2, EncDecRNNTBPEModel) + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_forward(self, asr_model): + asr_model = asr_model.eval() + + asr_model.preprocessor.featurizer.dither = 0.0 + asr_model.preprocessor.featurizer.pad_to = 0 + + asr_model.compute_eval_loss = False + + input_signal = torch.randn(size=(4, 512)) + length = torch.randint(low=161, high=500, size=[4]) + + with torch.no_grad(): + # batch size 1 + logprobs_instance = [] + for i in range(input_signal.size(0)): + logprobs_ins, _ = asr_model.forward( + input_signal=input_signal[i : i + 1], input_signal_length=length[i : i + 1] + ) + logprobs_instance.append(logprobs_ins) + logits_instance = torch.cat(logprobs_instance, 0) + + # batch size 4 + logprobs_batch, _ = asr_model.forward(input_signal=input_signal, input_signal_length=length) + + assert logits_instance.shape == logprobs_batch.shape + diff = torch.mean(torch.abs(logits_instance - logprobs_batch)) + assert diff <= 1e-6 + diff = torch.max(torch.abs(logits_instance - logprobs_batch)) + assert diff <= 1e-6 + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_save_restore_artifact(self, asr_model): + asr_model.train() + + with tempfile.TemporaryDirectory() as tmp_dir: + path = os.path.join(tmp_dir, 'rnnt_bpe.nemo') + asr_model.save_to(path) + + new_model = EncDecRNNTBPEModel.restore_from(path) + assert isinstance(new_model, type(asr_model)) + assert new_model.vocab_path.endswith('_vocab.txt') + + assert len(new_model.tokenizer.tokenizer.get_vocab()) == 128 + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_save_restore_artifact_spe(self, asr_model, test_data_dir): + asr_model.train() + + with tempfile.TemporaryDirectory() as tmpdir: + tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_spe_128") + asr_model.change_vocabulary(new_tokenizer_dir=tokenizer_dir, new_tokenizer_type='bpe') + + save_path = os.path.join(tmpdir, 'ctc_bpe.nemo') + asr_model.train() + asr_model.save_to(save_path) + + new_model = EncDecRNNTBPEModel.restore_from(save_path) + assert isinstance(new_model, type(asr_model)) + assert isinstance(new_model.tokenizer, tokenizers.SentencePieceTokenizer) + assert new_model.model_path.endswith('_tokenizer.model') + assert new_model.vocab_path.endswith('_vocab.txt') + assert new_model.spe_vocab_path.endswith('_tokenizer.vocab') + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_vocab_change(self, test_data_dir, asr_model): + with tempfile.TemporaryDirectory() as tmpdir: + old_tokenizer_dir = os.path.join(test_data_dir, "asr", "tokenizers", "an4_wpe_128", 'vocab.txt') + new_tokenizer_dir = os.path.join(tmpdir, 'tokenizer') + + os.makedirs(new_tokenizer_dir, exist_ok=True) + shutil.copy2(old_tokenizer_dir, new_tokenizer_dir) + + nw1 = asr_model.num_weights + asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe') + # No change + assert nw1 == asr_model.num_weights + + with open(os.path.join(new_tokenizer_dir, 'vocab.txt'), 'a+') as f: + f.write("!\n") + f.write('$\n') + f.write('@\n') + + asr_model.change_vocabulary(new_tokenizer_dir=new_tokenizer_dir, new_tokenizer_type='wpe') + + # rnn embedding + joint + bias + pred_embedding = 3 * (asr_model.decoder.pred_hidden) + joint_joint = 3 * (asr_model.joint.joint_hidden + 1) + assert asr_model.num_weights == (nw1 + (pred_embedding + joint_joint)) + + @pytest.mark.skipif( + not NUMBA_RNNT_LOSS_AVAILABLE, reason='RNNTLoss has not been compiled with appropriate numba version.', + ) + @pytest.mark.unit + def test_decoding_change(self, asr_model): + assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyBatchedRNNTInfer) + + new_strategy = DictConfig({}) + new_strategy.strategy = 'greedy' + new_strategy.greedy = DictConfig({'max_symbols': 10}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, greedy_decode.GreedyRNNTInfer) + + new_strategy = DictConfig({}) + new_strategy.strategy = 'beam' + new_strategy.beam = DictConfig({'beam_size': 1}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "default" + + new_strategy = DictConfig({}) + new_strategy.strategy = 'beam' + new_strategy.beam = DictConfig({'beam_size': 2}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "default" + + new_strategy = DictConfig({}) + new_strategy.strategy = 'tsd' + new_strategy.beam = DictConfig({'beam_size': 2}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "tsd" + + new_strategy = DictConfig({}) + new_strategy.strategy = 'alsd' + new_strategy.beam = DictConfig({'beam_size': 2}) + asr_model.change_decoding_strategy(decoding_cfg=new_strategy) + assert isinstance(asr_model.decoding.decoding, beam_decode.BeamRNNTInfer) + assert asr_model.decoding.decoding.search_type == "alsd" diff --git a/models/tts/NeMo/tests/collections/asr/test_jasper_block.py b/models/tts/NeMo/tests/collections/asr/test_jasper_block.py new file mode 100644 index 0000000000..2b6cfb41eb --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_jasper_block.py @@ -0,0 +1,295 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +import torch + +from nemo.collections.asr.parts.submodules import jasper + + +class TestJasperBlock: + def jasper_base_config(self, **kwargs): + base = dict( + inplanes=16, + planes=8, + kernel_size=[11], + repeat=1, + stride=[1], + dilation=[1], + activation="relu", + conv_mask=True, + separable=False, + se=False, + ) + base.update(kwargs) + return base + + def check_module_exists(self, module, cls): + global _MODULE_EXISTS + _MODULE_EXISTS = 0 + + def _traverse(m): + if isinstance(m, cls): + global _MODULE_EXISTS + _MODULE_EXISTS += 1 + + module.apply(_traverse) + assert _MODULE_EXISTS > 0 + + @pytest.mark.unit + def test_basic_block(self): + config = self.jasper_base_config(residual=False) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + + @pytest.mark.unit + def test_residual_block(self): + config = self.jasper_base_config(residual=True) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + + @pytest.mark.unit + def test_basic_block_repeat(self): + config = self.jasper_base_config(residual=False, repeat=3) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv) + + @pytest.mark.unit + def test_basic_block_repeat_stride(self): + config = self.jasper_base_config(residual=False, repeat=3, stride=[2]) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 17]) # 131 // (stride ^ repeats) + assert ylen[0] == 17 # 131 // (stride ^ repeats) + assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv) + + @pytest.mark.unit + def test_basic_block_repeat_stride_last(self): + config = self.jasper_base_config(residual=False, repeat=3, stride=[2], stride_last=True) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 66]) # 131 // stride + assert ylen[0] == 66 # 131 // stride + assert len(block.mconv) == 3 * 3 + 1 # (3 repeats x {1 conv + 1 norm + 1 dropout} + final conv) + + @pytest.mark.unit + def test_basic_block_repeat_separable(self): + config = self.jasper_base_config(residual=False, repeat=3, separable=True) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + assert len(block.mconv) == 3 * 4 + 1 # (3 repeats x {1 dconv + 1 pconv + 1 norm + 1 dropout} + final conv) + + @pytest.mark.unit + def test_basic_block_stride(self): + config = self.jasper_base_config(stride=[2], residual=False) + act = jasper.jasper_activations.get(config.pop('activation'))() + + print(config) + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 66]) + assert ylen[0] == 66 + + @pytest.mark.unit + def test_residual_block_stride(self): + config = self.jasper_base_config(stride=[2], residual=True, residual_mode='stride_add') + act = jasper.jasper_activations.get(config.pop('activation'))() + + print(config) + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 66]) + assert ylen[0] == 66 + + @pytest.mark.unit + def test_residual_block_activations(self): + for activation in jasper.jasper_activations.keys(): + config = self.jasper_base_config(activation=activation) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + self.check_module_exists(block, act.__class__) + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + + @pytest.mark.unit + def test_residual_block_normalizations(self): + NORMALIZATIONS = ["batch", "layer", "group"] + for normalization in NORMALIZATIONS: + config = self.jasper_base_config(normalization=normalization) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + + @pytest.mark.unit + def test_residual_block_se(self): + config = self.jasper_base_config(se=True, se_reduction_ratio=8) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + self.check_module_exists(block, jasper.SqueezeExcite) + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + + @pytest.mark.unit + def test_residual_block_asymmetric_pad_future_contexts(self): + # test future contexts at various values + # 0 = no future context + # 2 = limited future context + # 5 = symmetric context + # 8 = excess future context (more future context than present or past context) + future_contexts = [0, 2, 5, 8] + for future_context in future_contexts: + print(future_context) + config = self.jasper_base_config(future_context=future_context) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + self.check_module_exists(block, torch.nn.ConstantPad1d) + self.check_module_exists(block, jasper.MaskedConv1d) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + assert block.mconv[0].pad_layer is not None + assert block.mconv[0]._padding == (config['kernel_size'][0] - 1 - future_context, future_context) + + @pytest.mark.unit + def test_residual_block_asymmetric_pad_future_context_fallback(self): + # test future contexts at various values + # 15 = K < FC; fall back to symmetric context + future_context = 15 + print(future_context) + config = self.jasper_base_config(future_context=future_context) + act = jasper.jasper_activations.get(config.pop('activation'))() + + block = jasper.JasperBlock(**config, activation=act) + + x = torch.randn(1, 16, 131) + xlen = torch.tensor([131]) + y, ylen = block(([x], xlen)) + + self.check_module_exists(block, jasper.MaskedConv1d) + + assert isinstance(block, jasper.JasperBlock) + assert y[0].shape == torch.Size([1, config['planes'], 131]) + assert ylen[0] == 131 + assert block.mconv[0].pad_layer is None + assert block.mconv[0]._padding == config['kernel_size'][0] // 2 + + @pytest.mark.unit + def test_padding_size_conv1d(self): + input_channels = 1 + output_channels = 1 + kernel_sizes = [3, 7, 11] + dilation_sizes = [2, 3, 4] + stride = 1 + inp = torch.rand(2, 1, 40) + + for kernel_size in kernel_sizes: + for dilation_size in dilation_sizes: + padding = jasper.get_same_padding(kernel_size, stride, dilation_size) + + conv = torch.nn.Conv1d( + input_channels, output_channels, kernel_size=kernel_size, dilation=dilation_size, padding=padding + ) + + out = conv(inp) + assert out.shape == inp.shape diff --git a/models/tts/NeMo/tests/collections/asr/test_label_datasets.py b/models/tts/NeMo/tests/collections/asr/test_label_datasets.py new file mode 100644 index 0000000000..fbe9025284 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_label_datasets.py @@ -0,0 +1,106 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os + +import pytest +import torch + +from nemo.collections.asr.data.audio_to_label import TarredAudioToClassificationLabelDataset +from nemo.collections.asr.data.feature_to_label import FeatureToSeqSpeakerLabelDataset +from nemo.collections.asr.parts.preprocessing.feature_loader import ExternalFeatureLoader +from nemo.collections.asr.parts.preprocessing.features import WaveformFeaturizer + + +class TestASRDatasets: + labels = ["fash", "fbbh", "fclc"] + unique_labels_in_seq = ['0', '1', '2', '3', "zero", "one", "two", "three"] + + @pytest.mark.unit + def test_tarred_dataset(self, test_data_dir): + manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/tarred_audio_manifest.json')) + + # Test braceexpand loading + tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar')) + featurizer = WaveformFeaturizer(sample_rate=16000, int_values=False, augmentor=None) + ds_braceexpand = TarredAudioToClassificationLabelDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer + ) + + assert len(ds_braceexpand) == 32 + count = 0 + for _ in ds_braceexpand: + count += 1 + assert count == 32 + + # Test loading via list + tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)] + ds_list_load = TarredAudioToClassificationLabelDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer + ) + count = 0 + for _ in ds_list_load: + count += 1 + assert count == 32 + + @pytest.mark.unit + def test_tarred_dataset_duplicate_name(self, test_data_dir): + manifest_path = os.path.abspath( + os.path.join(test_data_dir, 'asr/tarred_an4/tarred_duplicate_audio_manifest.json') + ) + + # Test braceexpand loading + tarpath = os.path.abspath(os.path.join(test_data_dir, 'asr/tarred_an4/audio_{0..1}.tar')) + featurizer = WaveformFeaturizer(sample_rate=16000, int_values=False, augmentor=None) + ds_braceexpand = TarredAudioToClassificationLabelDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer + ) + + assert len(ds_braceexpand) == 6 + count = 0 + for _ in ds_braceexpand: + count += 1 + assert count == 6 + + # Test loading via list + tarpath = [os.path.abspath(os.path.join(test_data_dir, f'asr/tarred_an4/audio_{i}.tar')) for i in range(2)] + ds_list_load = TarredAudioToClassificationLabelDataset( + audio_tar_filepaths=tarpath, manifest_filepath=manifest_path, labels=self.labels, featurizer=featurizer + ) + count = 0 + for _ in ds_list_load: + count += 1 + assert count == 6 + + @pytest.mark.unit + def test_feat_seqlabel_dataset(self, test_data_dir): + manifest_path = os.path.abspath(os.path.join(test_data_dir, 'asr/feat/emb.json')) + feature_loader = ExternalFeatureLoader(file_path=manifest_path, sample_rate=16000, augmentor=None) + ds_braceexpand = FeatureToSeqSpeakerLabelDataset( + manifest_filepath=manifest_path, labels=self.unique_labels_in_seq, feature_loader=feature_loader + ) + # fmt: off + correct_label = torch.tensor( + [0.0, 1.0, 2.0, 2.0, 1.0, 2.0, 1.0, 2.0, 2.0, 1.0, 2.0, 2.0, 3.0, 1.0, 2.0, 2.0, 2.0, 0.0, 2.0, 1.0, 1.0, 2.0, 2.0, 1.0, 1.0, 2.0, 1.0, 1.0, 1.0, 1.0, 0.0, 0.0, 0.0, 2.0, 0.0, 2.0, 2.0, 2.0, 1.0, 2.0, 1.0, 1.0, 0.0, 1.0, 1.0, 0.0, 2.0, 1.0, 2.0, 1.0,] + ) + # fmt: on + correct_label_length = torch.tensor(50) + + assert ds_braceexpand[0][0].shape == (50, 32) + assert torch.equal(ds_braceexpand[0][2], correct_label) + assert torch.equal(ds_braceexpand[0][3], correct_label_length) + + count = 0 + for _ in ds_braceexpand: + count += 1 + assert count == 2 diff --git a/models/tts/NeMo/tests/collections/asr/test_speaker_label_models.py b/models/tts/NeMo/tests/collections/asr/test_speaker_label_models.py new file mode 100644 index 0000000000..f11f880969 --- /dev/null +++ b/models/tts/NeMo/tests/collections/asr/test_speaker_label_models.py @@ -0,0 +1,63 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from unittest import TestCase + +import pytest +from omegaconf import DictConfig + +from nemo.collections.asr.models import EncDecSpeakerLabelModel + + +class EncDecSpeechLabelModelTest(TestCase): + @pytest.mark.unit + def test_constructor(self): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 512, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': False, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.SpeakerDecoder', + 'params': {'feat_in': 512, 'num_classes': 2, 'pool_mode': 'xvector', 'emb_sizes': [1024]}, + } + + modelConfig = DictConfig( + {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)} + ) + speaker_model = EncDecSpeakerLabelModel(cfg=modelConfig) + speaker_model.train() + # TODO: make proper config and assert correct number of weights + + # Check to/from config_dict: + confdict = speaker_model.to_config_dict() + instance2 = EncDecSpeakerLabelModel.from_config_dict(confdict) + self.assertTrue(isinstance(instance2, EncDecSpeakerLabelModel)) diff --git a/models/tts/NeMo/tests/collections/common/__init__.py b/models/tts/NeMo/tests/collections/common/__init__.py new file mode 100644 index 0000000000..9e32500719 --- /dev/null +++ b/models/tts/NeMo/tests/collections/common/__init__.py @@ -0,0 +1,13 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. diff --git a/models/tts/NeMo/tests/collections/common/loss_inputs.py b/models/tts/NeMo/tests/collections/common/loss_inputs.py new file mode 100644 index 0000000000..cc3c384efe --- /dev/null +++ b/models/tts/NeMo/tests/collections/common/loss_inputs.py @@ -0,0 +1,69 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# Copyright 2018-2020 William Falcon +# +# 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. + +from dataclasses import dataclass + +import torch + +from .pl_utils import BATCH_SIZE, NUM_BATCHES, NUM_CLASSES + + +@dataclass(frozen=True) +class LossInput: + """ + The input for ``nemo.collections.common.metrics.GlobalAverageLossMetric`` metric tests. + + Args: + loss_sum_or_avg: a one dimensional float tensor which contains losses for averaging. Each element is either a + sum or mean of several losses depending on the parameter ``take_avg_loss`` of the + ``nemo.collections.common.metrics.GlobalAverageLossMetric`` class. + num_measurements: a one dimensional integer tensor which contains number of measurements which sums or average + values are in ``loss_sum_or_avg``. + """ + + loss_sum_or_avg: torch.Tensor + num_measurements: torch.Tensor + + +NO_ZERO_NUM_MEASUREMENTS = LossInput( + loss_sum_or_avg=torch.rand(NUM_BATCHES) * 2.0 - 1.0, num_measurements=torch.randint(1, 100, (NUM_BATCHES,)), +) + +SOME_NUM_MEASUREMENTS_ARE_ZERO = LossInput( + loss_sum_or_avg=torch.rand(NUM_BATCHES) * 2.0 - 1.0, + num_measurements=torch.cat( + ( + torch.randint(1, 100, (NUM_BATCHES // 2,), dtype=torch.int32), + torch.zeros(NUM_BATCHES - NUM_BATCHES // 2, dtype=torch.int32), + ) + ), +) + +ALL_NUM_MEASUREMENTS_ARE_ZERO = LossInput( + loss_sum_or_avg=torch.rand(NUM_BATCHES) * 2.0 - 1.0, num_measurements=torch.zeros(NUM_BATCHES, dtype=torch.int32), +) diff --git a/models/tts/NeMo/tests/collections/common/perplexity_inputs.py b/models/tts/NeMo/tests/collections/common/perplexity_inputs.py new file mode 100644 index 0000000000..6e785ba7a4 --- /dev/null +++ b/models/tts/NeMo/tests/collections/common/perplexity_inputs.py @@ -0,0 +1,49 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# Copyright 2018-2020 William Falcon +# +# 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. + +from collections import namedtuple + +import torch + +from .pl_utils import BATCH_SIZE, NUM_BATCHES, NUM_CLASSES + +Input = namedtuple('Input', ["probs", "logits"]) + + +ONLY_PROBS = Input(probs=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES), logits=None) + +ONLY_LOGITS1 = Input(probs=None, logits=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES)) + +ONLY_LOGITS100 = Input(probs=None, logits=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES) * 200 - 100) + +PROBS_AND_LOGITS = Input( + probs=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES), + logits=torch.rand(NUM_BATCHES, BATCH_SIZE, NUM_CLASSES) * 200 - 100, +) + +NO_PROBS_NO_LOGITS = Input(probs=None, logits=None) diff --git a/models/tts/NeMo/tests/collections/common/pl_utils.py b/models/tts/NeMo/tests/collections/common/pl_utils.py new file mode 100644 index 0000000000..b902e408d4 --- /dev/null +++ b/models/tts/NeMo/tests/collections/common/pl_utils.py @@ -0,0 +1,557 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +# Copyright 2018-2020 William Falcon +# +# 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 os +import pickle +import sys +from functools import partial +from typing import Callable, Optional + +import numpy as np +import pytest +import torch +from scipy.stats import entropy +from torch.distributions.utils import logits_to_probs +from torch.multiprocessing import Pool, set_start_method +from torchmetrics import Metric + +from nemo.collections.common.metrics import GlobalAverageLossMetric, Perplexity + +NUM_PROCESSES = 2 +NUM_BATCHES = 10 +BATCH_SIZE = 16 +NUM_CLASSES = 5 +EXTRA_DIM = 3 +THRESHOLD = 0.5 + + +def setup_ddp(rank, world_size): + """ Setup ddp enviroment """ + os.environ["MASTER_ADDR"] = 'localhost' + os.environ['MASTER_PORT'] = '8088' + + if torch.distributed.is_available() and sys.platform not in ['win32', 'cygwin']: + torch.distributed.init_process_group("gloo", rank=rank, world_size=world_size) + + +def _class_test( + rank: int, + worldsize: int, + preds: torch.Tensor, + target: torch.Tensor, + metric_class: Metric, + sk_metric: Callable, + dist_sync_on_step: bool, + metric_args: dict = {}, + check_dist_sync_on_step: bool = True, + check_batch: bool = True, + atol: float = 1e-8, +): + """ Utility function doing the actual comparison between lightning class metric + and reference metric. + Args: + rank: rank of current process + worldsize: number of processes + preds: torch tensor with predictions + target: torch tensor with targets + metric_class: lightning metric class that should be tested + sk_metric: callable function that is used for comparison + dist_sync_on_step: bool, if true will synchronize metric state across + processes at each ``forward()`` + metric_args: dict with additional arguments used for class initialization + check_dist_sync_on_step: bool, if true will check if the metric is also correctly + calculated per batch per device (and not just at the end) + check_batch: bool, if true will check if the metric is also correctly + calculated across devices for each batch (and not just at the end) + """ + # Instanciate lightning metric + metric = metric_class(compute_on_step=True, dist_sync_on_step=dist_sync_on_step, **metric_args) + + # verify metrics work after being loaded from pickled state + pickled_metric = pickle.dumps(metric) + metric = pickle.loads(pickled_metric) + + for i in range(rank, NUM_BATCHES, worldsize): + batch_result = metric(preds[i], target[i]) + + if metric.dist_sync_on_step: + if rank == 0: + ddp_preds = torch.stack([preds[i + r] for r in range(worldsize)]) + ddp_target = torch.stack([target[i + r] for r in range(worldsize)]) + sk_batch_result = sk_metric(ddp_preds, ddp_target) + # assert for dist_sync_on_step + if check_dist_sync_on_step: + assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol) + else: + sk_batch_result = sk_metric(preds[i], target[i]) + # assert for batch + if check_batch: + assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol) + + # check on all batches on all ranks + result = metric.compute() + assert isinstance(result, torch.Tensor) + + total_preds = torch.stack([preds[i] for i in range(NUM_BATCHES)]) + total_target = torch.stack([target[i] for i in range(NUM_BATCHES)]) + sk_result = sk_metric(total_preds, total_target) + + # assert after aggregation + assert np.allclose(result.numpy(), sk_result, atol=atol) + + +def _functional_test( + preds: torch.Tensor, + target: torch.Tensor, + metric_functional: Callable, + sk_metric: Callable, + metric_args: dict = {}, + atol: float = 1e-8, +): + """ Utility function doing the actual comparison between lightning functional metric + and reference metric. + Args: + preds: torch tensor with predictions + target: torch tensor with targets + metric_functional: lightning metric functional that should be tested + sk_metric: callable function that is used for comparison + metric_args: dict with additional arguments used for class initialization + """ + metric = partial(metric_functional, **metric_args) + + for i in range(NUM_BATCHES): + lightning_result = metric(preds[i], target[i]) + sk_result = sk_metric(preds[i], target[i]) + + # assert its the same + assert np.allclose(lightning_result.numpy(), sk_result, atol=atol) + + +class MetricTester: + """ Class used for efficiently run alot of parametrized tests in ddp mode. + Makes sure that ddp is only setup once and that pool of processes are + used for all tests. + All tests should subclass from this and implement a new method called + `test_metric_name` + where the method `self.run_metric_test` is called inside. + """ + + atol = 1e-8 + + def setup_class(self): + """ Setup the metric class. This will spawn the pool of workers that are + used for metric testing and setup_ddp + """ + try: + set_start_method('spawn') + except RuntimeError: + pass + self.poolSize = NUM_PROCESSES + self.pool = Pool(processes=self.poolSize) + self.pool.starmap(setup_ddp, [(rank, self.poolSize) for rank in range(self.poolSize)]) + + def teardown_class(self): + """ Close pool of workers """ + self.pool.close() + self.pool.join() + + def run_functional_metric_test( + self, + preds: torch.Tensor, + target: torch.Tensor, + metric_functional: Callable, + sk_metric: Callable, + metric_args: dict = {}, + ): + """ Main method that should be used for testing functions. Call this inside + testing method + Args: + preds: torch tensor with predictions + target: torch tensor with targets + metric_functional: lightning metric class that should be tested + sk_metric: callable function that is used for comparison + metric_args: dict with additional arguments used for class initialization + """ + _functional_test( + preds=preds, + target=target, + metric_functional=metric_functional, + sk_metric=sk_metric, + metric_args=metric_args, + atol=self.atol, + ) + + def run_class_metric_test( + self, + ddp: bool, + preds: torch.Tensor, + target: torch.Tensor, + metric_class: Metric, + sk_metric: Callable, + dist_sync_on_step: bool, + metric_args: dict = {}, + check_dist_sync_on_step: bool = True, + check_batch: bool = True, + ): + """ Main method that should be used for testing class. Call this inside testing + methods. + Args: + ddp: bool, if running in ddp mode or not + preds: torch tensor with predictions + target: torch tensor with targets + metric_class: lightning metric class that should be tested + sk_metric: callable function that is used for comparison + dist_sync_on_step: bool, if true will synchronize metric state across + processes at each ``forward()`` + metric_args: dict with additional arguments used for class initialization + check_dist_sync_on_step: bool, if true will check if the metric is also correctly + calculated per batch per device (and not just at the end) + check_batch: bool, if true will check if the metric is also correctly + calculated across devices for each batch (and not just at the end) + """ + if ddp: + if sys.platform == "win32": + pytest.skip("DDP not supported on windows") + + self.pool.starmap( + partial( + _class_test, + preds=preds, + target=target, + metric_class=metric_class, + sk_metric=sk_metric, + dist_sync_on_step=dist_sync_on_step, + metric_args=metric_args, + check_dist_sync_on_step=check_dist_sync_on_step, + check_batch=check_batch, + atol=self.atol, + ), + [(rank, self.poolSize) for rank in range(self.poolSize)], + ) + else: + _class_test( + 0, + 1, + preds=preds, + target=target, + metric_class=metric_class, + sk_metric=sk_metric, + dist_sync_on_step=dist_sync_on_step, + metric_args=metric_args, + check_dist_sync_on_step=check_dist_sync_on_step, + check_batch=check_batch, + atol=self.atol, + ) + + +def reference_perplexity_func(probs): + ent = entropy(probs, axis=-1) + ppl = np.exp(ent) + return ppl.mean() + + +def _perplexity_class_test( + rank: int, + worldsize: int, + probs: Optional[torch.Tensor], + logits: Optional[torch.Tensor], + dist_sync_on_step: bool, + metric_args: dict = {}, + check_dist_sync_on_step: bool = True, + check_batch: bool = True, + atol: float = 1e-8, +): + """ Utility function doing the actual comparison between lightning class metric + and reference metric. + Args: + rank: rank of current process + worldsize: number of processes + probs: torch tensor with probabilities + logits: torch tensor with logits. The function checks ``probs`` and ``logits are mutually exclusive for + ``Perplexity`` metric. + dist_sync_on_step: bool, if true will synchronize metric state across + processes at each ``forward()`` + metric_args: dict with additional arguments used for class initialization + check_dist_sync_on_step: bool, if true will check if the metric is also correctly + calculated per batch per device (and not just at the end) + check_batch: bool, if true will check if the metric is also correctly + calculated across devices for each batch (and not just at the end) + """ + # Instanciate lightning metric + perplexity = Perplexity(compute_on_step=True, dist_sync_on_step=dist_sync_on_step, **metric_args) + if (probs is None) == (logits is None): + with pytest.raises(ValueError): + perplexity(probs, logits) + return + + # verify perplexity works after being loaded from pickled state + pickled_metric = pickle.dumps(perplexity) + perplexity = pickle.loads(pickled_metric) + + for i in range(rank, NUM_BATCHES, worldsize): + batch_result = perplexity(None if probs is None else probs[i], None if logits is None else logits[i]) + + if perplexity.dist_sync_on_step: + if rank == 0: + if probs is not None: + ddp_probs = torch.stack([probs[i + r] for r in range(worldsize)]) + else: + ddp_logits = torch.stack([logits[i + r] for r in range(worldsize)]) + ddp_probs = logits_to_probs(ddp_logits, is_binary=False) + sk_batch_result = reference_perplexity_func(ddp_probs) + # assert for dist_sync_on_step + if check_dist_sync_on_step: + assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol) + else: + if probs is None: + p = logits_to_probs(logits[i], is_binary=False) + else: + p = probs[i] + sk_batch_result = reference_perplexity_func(p) + # assert for batch + if check_batch: + assert np.allclose(batch_result.numpy(), sk_batch_result, atol=atol) + + assert (probs is None) != (logits is None) + # check on all batches on all ranks + result = perplexity.compute() + assert isinstance(result, torch.Tensor) + + if probs is None: + probs = logits_to_probs(logits, is_binary=False) + sk_result = reference_perplexity_func(probs) + + # assert after aggregation + assert np.allclose(result.numpy(), sk_result, atol=atol) + + +class PerplexityTester(MetricTester): + def run_class_perplexity_test( + self, + ddp: bool, + probs: Optional[torch.Tensor], + logits: Optional[torch.Tensor], + dist_sync_on_step: bool, + metric_args: dict = {}, + check_dist_sync_on_step: bool = True, + check_batch: bool = True, + ): + """ Main method that should be used for testing class. Call this inside testing + methods. + Args: + ddp: bool, if running in ddp mode or not + probs: torch tensor with probabilities. + logits: torch tensor with logits. This test checks that probs and logits are mutually exclusive for + ``Perplexity`` metric. + dist_sync_on_step: bool, if true will synchronize metric state across + processes at each ``forward()`` + metric_args: dict with additional arguments used for class initialization + check_dist_sync_on_step: bool, if true will check if the metric is also correctly + calculated per batch per device (and not just at the end) + check_batch: bool, if true will check if the metric is also correctly + calculated across devices for each batch (and not just at the end) + """ + if ddp: + if sys.platform == "win32": + pytest.skip("DDP not supported on windows") + + self.pool.starmap( + partial( + _perplexity_class_test, + probs=probs, + logits=logits, + dist_sync_on_step=dist_sync_on_step, + metric_args=metric_args, + check_dist_sync_on_step=check_dist_sync_on_step, + check_batch=check_batch, + atol=self.atol, + ), + [(rank, self.poolSize) for rank in range(self.poolSize)], + ) + else: + _perplexity_class_test( + 0, + 1, + probs=probs, + logits=logits, + dist_sync_on_step=dist_sync_on_step, + metric_args=metric_args, + check_dist_sync_on_step=check_dist_sync_on_step, + check_batch=check_batch, + atol=self.atol, + ) + + +def reference_loss_func(loss_sum_or_avg: torch.Tensor, num_measurements: torch.Tensor, take_avg_loss: bool): + """ + Returns average loss for data from``loss_sum_or_avg``. This function sums all losses from ``loss_sum_or_avg`` and + divides the sum by the sum of ``num_measurements`` elements. + + If ``take_avg_loss`` is ``True`` then ``loss_sum_or_avg[i]`` elements are mean values of ``num_measurements[i]`` + losses. In that case before computing sum of losses each element of ``loss_sum_or_avg`` is multiplied by + corresponding element of ``num_measurements``. + + If ``num_measurements`` sum is zero then the function returns NaN tensor. + + The function is used for testing ``nemo.collections.common.metrics.GlobalAverageLossMetric`` class. + + Args: + loss_sum_or_avg: a one dimensional float ``torch.Tensor``. Sums or mean values of loss. + num_measurements: a one dimensional integer ``torch.Tensor``. Number of values on which sums of means in + ``loss_sum_or_avg`` are calculated. + take_avg_loss: if ``True`` then ``loss_sum_or_avg`` contains mean losses else ``loss_sum_or_avg`` contains + sums of losses. + """ + if take_avg_loss: + for i in range(loss_sum_or_avg.shape[0]): + loss_sum_or_avg[i] *= num_measurements[i] + nm_sum = num_measurements.sum() + if nm_sum.eq(0): + return torch.tensor(float('nan')) + return loss_sum_or_avg.sum() / nm_sum + + +def _loss_class_test( + rank: int, + worldsize: int, + loss_sum_or_avg: Optional[torch.Tensor], + num_measurements: Optional[torch.Tensor], + dist_sync_on_step: bool, + take_avg_loss: bool, + check_dist_sync_on_step: bool = True, + check_batch: bool = True, + atol: float = 1e-8, +): + """ Utility function doing the actual comparison between lightning class metric + and reference metric. + Args: + rank: rank of current process + worldsize: number of processes + loss_sum_or_avg: a one dimensional float torch tensor with loss sums or means. + num_measurements: a one dimensional integer torch tensor with number of values on which sums or means from + ``loss_sum_or_avg`` were computed. + dist_sync_on_step: bool, if true will synchronize metric state across + processes at each ``forward()`` + take_avg_loss: dict with additional arguments used for class initialization + check_dist_sync_on_step: bool, if true will check if the metric is also correctly + calculated per batch per device (and not just at the end) + check_batch: bool, if true will check if the metric is also correctly + calculated across devices for each batch (and not just at the end) + """ + # Instantiate lightning metric + loss_metric = GlobalAverageLossMetric( + compute_on_step=True, dist_sync_on_step=dist_sync_on_step, take_avg_loss=take_avg_loss + ) + + # verify perplexity works after being loaded from pickled state + pickled_metric = pickle.dumps(loss_metric) + loss_metric = pickle.loads(pickled_metric) + + for i in range(rank, NUM_BATCHES, worldsize): + batch_result = loss_metric(loss_sum_or_avg[i], num_measurements[i]) + if loss_metric.dist_sync_on_step: + if rank == 0: + ddp_loss_sum_or_avg = torch.stack([loss_sum_or_avg[i + r] for r in range(worldsize)]) + ddp_num_measurements = torch.stack([num_measurements[i + r] for r in range(worldsize)]) + sk_batch_result = reference_loss_func(ddp_loss_sum_or_avg, ddp_num_measurements, take_avg_loss) + # assert for dist_sync_on_step + if check_dist_sync_on_step: + if sk_batch_result.isnan(): + assert batch_result.isnan() + else: + assert ( + np.allclose(batch_result.numpy(), sk_batch_result, atol=atol), + f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}", + ) + else: + ls = loss_sum_or_avg[i : i + 1] + nm = num_measurements[i : i + 1] + sk_batch_result = reference_loss_func(ls, nm, take_avg_loss) + # assert for batch + if check_batch: + if sk_batch_result.isnan(): + assert batch_result.isnan() + else: + assert ( + np.allclose(batch_result.numpy(), sk_batch_result, atol=atol), + f"batch_result = {batch_result.numpy()}, sk_batch_result = {sk_batch_result}", + ) + # check on all batches on all ranks + result = loss_metric.compute() + assert isinstance(result, torch.Tensor) + sk_result = reference_loss_func(loss_sum_or_avg, num_measurements, take_avg_loss) + + # assert after aggregation + if sk_result.isnan(): + assert result.isnan() + else: + assert ( + np.allclose(result.numpy(), sk_result, atol=atol), + f"result = {result.numpy()}, sk_result = {sk_result}", + ) + + +class LossTester(MetricTester): + def run_class_loss_test( + self, + ddp: bool, + loss_sum_or_avg: torch.Tensor, + num_measurements: torch.Tensor, + dist_sync_on_step: bool, + take_avg_loss: bool, + check_dist_sync_on_step: bool = True, + check_batch: bool = True, + ): + if ddp: + if sys.platform == "win32": + pytest.skip("DDP not supported on windows") + self.pool.starmap( + partial( + _loss_class_test, + loss_sum_or_avg=loss_sum_or_avg, + num_measurements=num_measurements, + dist_sync_on_step=dist_sync_on_step, + take_avg_loss=take_avg_loss, + check_dist_sync_on_step=check_dist_sync_on_step, + check_batch=check_batch, + atol=self.atol, + ), + [(rank, self.poolSize) for rank in range(self.poolSize)], + ) + else: + _loss_class_test( + 0, + 1, + loss_sum_or_avg=loss_sum_or_avg, + num_measurements=num_measurements, + dist_sync_on_step=dist_sync_on_step, + take_avg_loss=take_avg_loss, + check_dist_sync_on_step=check_dist_sync_on_step, + check_batch=check_batch, + atol=self.atol, + ) diff --git a/models/tts/NeMo/tests/collections/common/test_metrics.py b/models/tts/NeMo/tests/collections/common/test_metrics.py new file mode 100644 index 0000000000..6729ae3339 --- /dev/null +++ b/models/tts/NeMo/tests/collections/common/test_metrics.py @@ -0,0 +1,150 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest +import torch + +from .loss_inputs import ALL_NUM_MEASUREMENTS_ARE_ZERO, NO_ZERO_NUM_MEASUREMENTS, SOME_NUM_MEASUREMENTS_ARE_ZERO +from .perplexity_inputs import NO_PROBS_NO_LOGITS, ONLY_LOGITS1, ONLY_LOGITS100, ONLY_PROBS, PROBS_AND_LOGITS +from .pl_utils import LossTester, PerplexityTester +from nemo.collections.common.metrics.classification_accuracy import TopKClassificationAccuracy + + +class TestCommonMetrics: + top_k_logits = torch.tensor([[0.1, 0.3, 0.2, 0.0], [0.9, 0.6, 0.2, 0.3], [0.2, 0.1, 0.4, 0.3]],) # 1 # 0 # 2 + + @pytest.mark.unit + def test_top_1_accuracy(self): + labels = torch.tensor([0, 0, 2], dtype=torch.long) + + accuracy = TopKClassificationAccuracy(top_k=None) + acc = accuracy(logits=self.top_k_logits, labels=labels) + + assert accuracy.correct_counts_k.shape == torch.Size([1]) + assert accuracy.total_counts_k.shape == torch.Size([1]) + assert abs(acc[0] - 0.667) < 1e-3 + + @pytest.mark.unit + def test_top_1_2_accuracy(self): + labels = torch.tensor([0, 1, 0], dtype=torch.long) + + accuracy = TopKClassificationAccuracy(top_k=[1, 2]) + top1_acc, top2_acc = accuracy(logits=self.top_k_logits, labels=labels) + + assert accuracy.correct_counts_k.shape == torch.Size([2]) + assert accuracy.total_counts_k.shape == torch.Size([2]) + + assert abs(top1_acc - 0.0) < 1e-3 + assert abs(top2_acc - 0.333) < 1e-3 + + @pytest.mark.unit + def test_top_1_accuracy_distributed(self): + # Simulate test on 2 process DDP execution + labels = torch.tensor([[0, 0, 2], [2, 0, 0]], dtype=torch.long) + + accuracy = TopKClassificationAccuracy(top_k=None) + proc1_acc = accuracy(logits=self.top_k_logits, labels=labels[0]) + correct1, total1 = accuracy.correct_counts_k, accuracy.total_counts_k + + accuracy.reset() + proc2_acc = accuracy(logits=torch.flip(self.top_k_logits, dims=[1]), labels=labels[1]) # reverse logits + correct2, total2 = accuracy.correct_counts_k, accuracy.total_counts_k + + correct = torch.stack([correct1, correct2]) + total = torch.stack([total1, total2]) + + assert correct.shape == torch.Size([2, 1]) + assert total.shape == torch.Size([2, 1]) + + assert abs(proc1_acc[0] - 0.667) < 1e-3 # 2/3 + assert abs(proc2_acc[0] - 0.333) < 1e-3 # 1/3 + + accuracy.reset() + accuracy.correct_counts_k = torch.tensor([correct.sum()]) + accuracy.total_counts_k = torch.tensor([total.sum()]) + acc_topk = accuracy.compute() + acc_top1 = acc_topk[0] + + assert abs(acc_top1 - 0.5) < 1e-3 # 3/6 + + @pytest.mark.unit + def test_top_1_accuracy_distributed_uneven_batch(self): + # Simulate test on 2 process DDP execution + accuracy = TopKClassificationAccuracy(top_k=None) + + proc1_acc = accuracy(logits=self.top_k_logits, labels=torch.tensor([0, 0, 2])) + correct1, total1 = accuracy.correct_counts_k, accuracy.total_counts_k + + proc2_acc = accuracy( + logits=torch.flip(self.top_k_logits, dims=[1])[:2, :], # reverse logits, select first 2 samples + labels=torch.tensor([2, 0]), + ) # reduce number of labels + correct2, total2 = accuracy.correct_counts_k, accuracy.total_counts_k + + correct = torch.stack([correct1, correct2]) + total = torch.stack([total1, total2]) + + assert correct.shape == torch.Size([2, 1]) + assert total.shape == torch.Size([2, 1]) + + assert abs(proc1_acc[0] - 0.667) < 1e-3 # 2/3 + assert abs(proc2_acc[0] - 0.500) < 1e-3 # 1/2 + + accuracy.correct_counts_k = torch.tensor([correct.sum()]) + accuracy.total_counts_k = torch.tensor([total.sum()]) + acc_topk = accuracy.compute() + acc_top1 = acc_topk[0] + + assert abs(acc_top1 - 0.6) < 1e-3 # 3/5 + + +@pytest.mark.parametrize("ddp", [True, False]) +@pytest.mark.parametrize("dist_sync_on_step", [True, False]) +@pytest.mark.parametrize( + "probs, logits", + [ + (ONLY_PROBS.probs, ONLY_PROBS.logits), + (ONLY_LOGITS1.probs, ONLY_LOGITS1.logits), + (ONLY_LOGITS100.probs, ONLY_LOGITS100.logits), + (PROBS_AND_LOGITS.probs, PROBS_AND_LOGITS.logits), + (NO_PROBS_NO_LOGITS.probs, NO_PROBS_NO_LOGITS.logits), + ], +) +class TestPerplexity(PerplexityTester): + def test_perplexity(self, ddp, dist_sync_on_step, probs, logits): + self.run_class_perplexity_test( + ddp=ddp, probs=probs, logits=logits, dist_sync_on_step=dist_sync_on_step, + ) + + +@pytest.mark.parametrize("ddp", [True, False]) +@pytest.mark.parametrize("dist_sync_on_step", [True, False]) +@pytest.mark.parametrize("take_avg_loss", [True, False]) +@pytest.mark.parametrize( + "loss_sum_or_avg, num_measurements", + [ + (NO_ZERO_NUM_MEASUREMENTS.loss_sum_or_avg, NO_ZERO_NUM_MEASUREMENTS.num_measurements), + (SOME_NUM_MEASUREMENTS_ARE_ZERO.loss_sum_or_avg, SOME_NUM_MEASUREMENTS_ARE_ZERO.num_measurements), + (ALL_NUM_MEASUREMENTS_ARE_ZERO.loss_sum_or_avg, ALL_NUM_MEASUREMENTS_ARE_ZERO.num_measurements), + ], +) +class TestLoss(LossTester): + def test_loss(self, ddp, dist_sync_on_step, loss_sum_or_avg, num_measurements, take_avg_loss): + self.run_class_loss_test( + ddp=ddp, + loss_sum_or_avg=loss_sum_or_avg, + num_measurements=num_measurements, + dist_sync_on_step=dist_sync_on_step, + take_avg_loss=take_avg_loss, + ) diff --git a/models/tts/NeMo/tests/collections/common/test_spc_tokenizer.py b/models/tts/NeMo/tests/collections/common/test_spc_tokenizer.py new file mode 100644 index 0000000000..da7b079234 --- /dev/null +++ b/models/tts/NeMo/tests/collections/common/test_spc_tokenizer.py @@ -0,0 +1,267 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest + +from nemo.collections.common.tokenizers.sentencepiece_tokenizer import SentencePieceTokenizer +from nemo.collections.common.tokenizers.youtokentome_tokenizer import YouTokenToMeTokenizer + +MODEL_SPECIAL_TOKENS = { + 'unk_token': '[UNK]', + 'sep_token': '[SEP]', + 'pad_token': '[PAD]', + 'bos_token': '[CLS]', + 'mask_token': '[MASK]', + 'eos_token': '[SEP]', + 'cls_token': '[CLS]', +} + + +class TestSentencePieceTokenizerLegacy: + model_name = "/m_common.model" + + @pytest.mark.unit + def test_add_special_tokens(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + special_tokens = MODEL_SPECIAL_TOKENS + tokenizer.add_special_tokens(special_tokens) + assert tokenizer.vocab_size == tokenizer.original_vocab_size + len(set(special_tokens.values())) + + @pytest.mark.unit + def test_text_to_tokens(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + special_tokens = MODEL_SPECIAL_TOKENS + tokenizer.add_special_tokens(special_tokens) + + text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]" + tokens = tokenizer.text_to_tokens(text) + + assert len(tokens) == len(text.split()) + assert tokens.count("[CLS]") == 1 + assert tokens.count("[MASK]") == 1 + assert tokens.count("[SEP]") == 2 + + @pytest.mark.unit + def test_tokens_to_text(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + + text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]" + tokens = tokenizer.text_to_tokens(text) + result = tokenizer.tokens_to_text(tokens) + + assert text == result + + @pytest.mark.unit + def test_text_to_ids(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + special_tokens = MODEL_SPECIAL_TOKENS + tokenizer.add_special_tokens(special_tokens) + + text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]" + ids = tokenizer.text_to_ids(text) + + assert len(ids) == len(text.split()) + assert ids.count(tokenizer.token_to_id("[CLS]")) == 1 + assert ids.count(tokenizer.token_to_id("[MASK]")) == 1 + assert ids.count(tokenizer.token_to_id("[SEP]")) == 2 + + @pytest.mark.unit + def test_ids_to_text(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + special_tokens = MODEL_SPECIAL_TOKENS + tokenizer.add_special_tokens(special_tokens) + + text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]" + ids = tokenizer.text_to_ids(text) + result = tokenizer.ids_to_text(ids) + + assert text == result + + @pytest.mark.unit + def test_tokens_to_ids(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + special_tokens = MODEL_SPECIAL_TOKENS + tokenizer.add_special_tokens(special_tokens) + + text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]" + tokens = tokenizer.text_to_tokens(text) + ids = tokenizer.tokens_to_ids(tokens) + + assert len(ids) == len(tokens) + assert ids.count(tokenizer.token_to_id("[CLS]")) == 1 + assert ids.count(tokenizer.token_to_id("[MASK]")) == 1 + assert ids.count(tokenizer.token_to_id("[SEP]")) == 2 + + @pytest.mark.unit + def test_ids_to_tokens(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name, legacy=True) + special_tokens = MODEL_SPECIAL_TOKENS + tokenizer.add_special_tokens(special_tokens) + + text = "[CLS] a b c [MASK] e f [SEP] g h i [SEP]" + tokens = tokenizer.text_to_tokens(text) + ids = tokenizer.tokens_to_ids(tokens) + result = tokenizer.ids_to_tokens(ids) + + assert len(result) == len(tokens) + + for i in range(len(result)): + assert result[i] == tokens[i] + + +class TestSentencePieceTokenizer: + model_name = "/m_new.model" + + @pytest.mark.unit + def test_text_to_tokens(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name) + + # is user_defined_symbol in the test tokenizer model + # , , , and are control symbols + text = " a b c e f g h i " + tokens = tokenizer.text_to_tokens(text) + + assert tokens.count("") == 1 + assert tokens.count("") == 0 + assert tokens.count("") == 0 + + @pytest.mark.unit + def test_tokens_to_text(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name) + + # is user_defined_symbol in the test tokenizer model + text = " a b c e f g h i" + tokens = tokenizer.text_to_tokens(text) + result = tokenizer.tokens_to_text(tokens) + + assert text == result + + @pytest.mark.unit + def test_text_to_ids(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name) + + # is user_defined_symbol in the test tokenizer model + # , , , and are control symbols + text = " a b c e f g h i " + tokens = tokenizer.text_to_ids(text) + + assert tokens.count(tokenizer.token_to_id("")) == 1 + assert tokens.count(tokenizer.token_to_id("")) == 0 + assert tokens.count(tokenizer.token_to_id("")) == 0 + + @pytest.mark.unit + def test_ids_to_text(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name) + + text = " a b c e f g h i " + ids = tokenizer.text_to_ids(text) + result = tokenizer.ids_to_text(ids) + + assert text == result + + @pytest.mark.unit + def test_tokens_to_ids(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name) + + tokens = ["", "a", "b", "c", "", "e", "f", "", "g", "h", "i", ""] + ids = tokenizer.tokens_to_ids(tokens) + + assert len(ids) == len(tokens) + assert ids.count(tokenizer.token_to_id("")) == 1 + assert ids.count(tokenizer.token_to_id("")) == 1 + assert ids.count(tokenizer.token_to_id("")) == 2 + + @pytest.mark.unit + def test_ids_to_tokens(self, test_data_dir): + tokenizer = SentencePieceTokenizer(test_data_dir + self.model_name) + + tokens = ["", "a", "b", "c", "", "e", "f", "", "g", "h", "i", ""] + ids = tokenizer.tokens_to_ids(tokens) + result = tokenizer.ids_to_tokens(ids) + + assert len(result) == len(tokens) + + for i in range(len(result)): + assert result[i] == tokens[i] + + +class TestYouTokenToMeTokenizer: + model_name = "/yttm.4096.en-de.model" + + @pytest.mark.unit + def test_text_to_tokens(self, test_data_dir): + tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name) + + text = " a b c e f g h i " + tokens = tokenizer.text_to_tokens(text) + + assert tokens.count("") == 0 + assert tokens.count("") == 0 + assert tokens.count("") == 0 + + @pytest.mark.unit + def test_tokens_to_text(self, test_data_dir): + tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name) + + text = "a b c e f g h i" + tokens = tokenizer.text_to_tokens(text) + result = tokenizer.tokens_to_text(tokens) + + assert text == result + + @pytest.mark.unit + def test_text_to_ids(self, test_data_dir): + tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name) + + text = " a b c e f g h i " + tokens = tokenizer.text_to_ids(text) + + assert tokens.count(tokenizer.bos_id) == 0 + assert tokens.count(tokenizer.unk_id) == 0 + assert tokens.count(tokenizer.eos_id) == 0 + + @pytest.mark.unit + def test_ids_to_text(self, test_data_dir): + tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name) + + text = "a b c e f g h i" + ids = tokenizer.text_to_ids(text) + result = tokenizer.ids_to_text(ids) + + assert text == result + + @pytest.mark.unit + def test_tokens_to_ids(self, test_data_dir): + tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name) + + tokens = ["", "a", "b", "c", "", "e", "f", "", "g", "h", "i", ""] + ids = tokenizer.tokens_to_ids(tokens) + + assert len(ids) == len(tokens) + assert ids.count(tokenizer.bos_id) == 1 + assert ids.count(tokenizer.eos_id) == 1 + assert ids.count(tokenizer.unk_id) == 2 + + @pytest.mark.unit + def test_ids_to_tokens(self, test_data_dir): + tokenizer = YouTokenToMeTokenizer(test_data_dir + self.model_name) + + tokens = ["", "a", "b", "c", "", "e", "f", "", "g", "h", "i", ""] + ids = tokenizer.tokens_to_ids(tokens) + result = tokenizer.ids_to_tokens(ids) + + assert len(result) == len(tokens) + + for i in range(len(result)): + assert result[i] == tokens[i] diff --git a/models/tts/NeMo/tests/collections/nlp/test_classification_report.py b/models/tts/NeMo/tests/collections/nlp/test_classification_report.py new file mode 100644 index 0000000000..de2c4f7a71 --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_classification_report.py @@ -0,0 +1,48 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest +import torch +from sklearn.metrics import precision_recall_fscore_support + +from nemo.collections.nlp.metrics.classification_report import ClassificationReport + + +class ClassificationReportTests: + num_classes = 3 + label_ids = {'a': 0, 'b': 1, 'c': 2} + + @pytest.mark.unit + def test_classification_report(self): + + preds = torch.Tensor([0, 1, 1, 1, 2, 2, 0]) + labels = torch.Tensor([1, 0, 0, 1, 2, 1, 0]) + + def __convert_to_tensor(sklearn_metric): + return torch.Tensor([round(sklearn_metric * 100)])[0] + + for mode in ['macro', 'micro', 'weighted']: + + classification_report_nemo = ClassificationReport( + num_classes=self.num_classes, label_ids=self.label_ids, mode=mode + ) + # pytest.set_trace() + precision, recall, f1, _ = classification_report_nemo(preds, labels) + tp, fp, fn = classification_report_nemo.tp, classification_report_nemo.fp, classification_report_nemo.fn + pr_sklearn, recall_sklearn, f1_sklearn, _ = precision_recall_fscore_support(labels, preds, average=mode) + + self.assertEqual(torch.round(precision), __convert_to_tensor(pr_sklearn), f'wrong precision for {mode}') + self.assertEqual(torch.round(recall), __convert_to_tensor(recall_sklearn), f'wrong recall for {mode}') + self.assertEqual(torch.round(f1), __convert_to_tensor(f1_sklearn), f'wrong f1 for {mode}') diff --git a/models/tts/NeMo/tests/collections/nlp/test_entity_linking_model.py b/models/tts/NeMo/tests/collections/nlp/test_entity_linking_model.py new file mode 100644 index 0000000000..0e96b33f9b --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_entity_linking_model.py @@ -0,0 +1,83 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os +import shutil +import tempfile + +import pytest +import wget +from omegaconf import OmegaConf + +from nemo.collections.nlp.models import EntityLinkingModel + + +def get_cfg(): + + language_model = OmegaConf.create( + {"pretrained_model_name": "bert-base-uncased", "config_file": None, "config": None, "lm_checkpoint": None} + ) + + tokenizer = OmegaConf.create( + {"tokenizer_name": "bert-base-uncased", "vocab_file": None, "tokenizer_model": None, "do_lower_case": True} + ) + + model = OmegaConf.create( + { + "nemo_path": "sap_entity_linking.nemo", + "max_seq_length": 128, + "language_model": language_model, + "tokenizer": tokenizer, + "train_ds": None, + "validation_ds": None, + } + ) + + cfg = OmegaConf.create({"model": model}) + + return cfg + + +class TestEntityLinkingModel: + @pytest.mark.unit + def test_creation_saving_restoring(self): + # Create a new temporary directory + with tempfile.TemporaryDirectory() as restore_dir: + with tempfile.TemporaryDirectory() as save_dir: + model = EntityLinkingModel(cfg=get_cfg().model) + assert isinstance(model, EntityLinkingModel) + + save_dir_path = save_dir + + # Where model will be saved + model_save_path = os.path.join(save_dir, f"{model.__class__.__name__}.nemo") + model.save_to(save_path=model_save_path) + + # Where model will be restored from + model_restore_path = os.path.join(restore_dir, f"{model.__class__.__name__}.nemo") + shutil.copy(model_save_path, model_restore_path) + + # at this point save_dir should not exist + assert save_dir_path is not None and not os.path.exists(save_dir_path) + assert not os.path.exists(model_save_path) + assert os.path.exists(model_restore_path) + + # attempt to restore + model_copy = model.__class__.restore_from(restore_path=model_restore_path) + assert model.num_weights == model_copy.num_weights + + +if __name__ == "__main__": + t = TestEntityLinkingModel() + t.test_creation_saving_restoring() diff --git a/models/tts/NeMo/tests/collections/nlp/test_huggingface.py b/models/tts/NeMo/tests/collections/nlp/test_huggingface.py new file mode 100644 index 0000000000..a179c6db95 --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_huggingface.py @@ -0,0 +1,86 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import tempfile + +import pytest +import torch + +import nemo.collections.nlp as nemo_nlp +from nemo.collections.common.tokenizers.huggingface.auto_tokenizer import AutoTokenizer +from nemo.collections.nlp.modules.common.tokenizer_utils import get_tokenizer + + +def do_export(model, name: str): + with tempfile.TemporaryDirectory() as tmpdir: + # Generate filename in the temporary directory. + tmp_file_name = os.path.join(tmpdir, name + '.onnx') + # Test export. + if torch.cuda.is_available(): + model = model.cuda() + model.export(tmp_file_name) + + +class TestHuggingFace: + @pytest.mark.unit + def test_list_pretrained_models(self): + pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list() + assert len(pretrained_lm_models) > 0 + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_get_pretrained_bert_model(self): + model = nemo_nlp.modules.get_lm_model(pretrained_model_name='bert-base-uncased') + assert isinstance(model, nemo_nlp.modules.BertEncoder) + do_export(model, "bert-base-uncased") + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_get_pretrained_distilbert_model(self): + model = nemo_nlp.modules.get_lm_model(pretrained_model_name='distilbert-base-uncased') + assert isinstance(model, nemo_nlp.modules.DistilBertEncoder) + do_export(model, "distilbert-base-uncased") + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_get_pretrained_roberta_model(self): + model = nemo_nlp.modules.get_lm_model(pretrained_model_name='roberta-base') + assert isinstance(model, nemo_nlp.modules.RobertaEncoder) + do_export(model, "roberta-base-uncased") + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_get_pretrained_albert_model(self): + model = nemo_nlp.modules.get_lm_model(pretrained_model_name='albert-base-v1') + assert isinstance(model, nemo_nlp.modules.AlbertEncoder) + do_export(model, "albert-base-v1") + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_get_pretrained_chinese_bert_wwm_model(self): + model_name = 'hfl/chinese-bert-wwm' + model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name) + assert isinstance(model, nemo_nlp.modules.BertModule) + tokenizer = get_tokenizer(tokenizer_name=model_name) + assert isinstance(tokenizer, AutoTokenizer) + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_get_pretrained_arabic_model(self): + model_name = 'asafaya/bert-base-arabic' + model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name) + assert isinstance(model, nemo_nlp.modules.BertModule) + tokenizer = get_tokenizer(tokenizer_name=model_name) + assert isinstance(tokenizer, AutoTokenizer) diff --git a/models/tts/NeMo/tests/collections/nlp/test_megatron.py b/models/tts/NeMo/tests/collections/nlp/test_megatron.py new file mode 100644 index 0000000000..4e7d448a78 --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_megatron.py @@ -0,0 +1,76 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +try: + import apex + + apex_available = True +except Exception: + apex_available = False + +import os +import tempfile + +import onnx +import pytest +import torch + +import nemo.collections.nlp as nemo_nlp +from nemo.core.classes import typecheck + + +def get_pretrained_bert_345m_uncased_model(): + model_name = "megatron-bert-345m-uncased" + model = nemo_nlp.modules.get_lm_model(pretrained_model_name=model_name) + if torch.cuda.is_available(): + model = model.cuda() + return model + + +class TestMegatron: + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_list_pretrained_models(self): + pretrained_lm_models = nemo_nlp.modules.get_pretrained_lm_models_list() + assert len(pretrained_lm_models) > 0 + + @pytest.mark.with_downloads() + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + @pytest.mark.skip("Only one Megatron model is allowed") + def test_get_model(self): + model = get_pretrained_bert_345m_uncased_model() + assert isinstance(model, nemo_nlp.modules.MegatronBertEncoder) + + typecheck.set_typecheck_enabled(enabled=False) + inp = model.input_example() + out = model.forward(*inp) + typecheck.set_typecheck_enabled(enabled=True) + + @pytest.mark.skipif(not os.path.exists('/home/TestData/nlp'), reason='Not a Jenkins machine') + @pytest.mark.with_downloads() + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_onnx_export(self): + model = get_pretrained_bert_345m_uncased_model() + assert model + with tempfile.TemporaryDirectory() as tmpdir: + # Generate filename in the temporary directory. + # Test export. + model.export(os.path.join(".", "megatron.onnx")) + + +if __name__ == "__main__": + t = TestMegatron() + t.test_onnx_export() diff --git a/models/tts/NeMo/tests/collections/nlp/test_nlp_exportables.py b/models/tts/NeMo/tests/collections/nlp/test_nlp_exportables.py new file mode 100644 index 0000000000..ac5d86026a --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_nlp_exportables.py @@ -0,0 +1,176 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import tempfile + +import onnx +import pytest +import pytorch_lightning as pl +import torch +import wget +from omegaconf import OmegaConf + +from nemo.collections import nlp as nemo_nlp +from nemo.collections.nlp.models import IntentSlotClassificationModel +from nemo.collections.nlp.modules.common import ( + SequenceClassifier, + SequenceRegression, + SequenceTokenClassifier, + TokenClassifier, +) + + +def classifier_export(obj): + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, obj.__class__.__name__ + '.onnx') + obj = obj.cuda() + obj.export(output=filename) + + +class TestExportableClassifiers: + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_token_classifier_export_to_onnx(self): + for num_layers in [1, 2, 4]: + classifier_export(TokenClassifier(hidden_size=256, num_layers=num_layers, num_classes=16)) + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_bert_pretraining_export_to_onnx(self): + for num_layers in [1, 2, 4]: + classifier_export(TokenClassifier(hidden_size=256, num_layers=num_layers, num_classes=16)) + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_sequence_token_classifier_export_to_onnx(self): + for num_layers in [1, 2, 4]: + classifier_export( + SequenceTokenClassifier(hidden_size=256, num_slots=8, num_intents=8, num_layers=num_layers) + ) + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_sequence_classifier_export_to_onnx(self): + for num_layers in [1, 2, 4]: + classifier_export(SequenceClassifier(hidden_size=256, num_classes=16, num_layers=num_layers)) + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_sequence_regression_export_to_onnx(self): + for num_layers in [1, 2, 4]: + classifier_export(SequenceRegression(hidden_size=256, num_layers=num_layers)) + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_IntentSlotClassificationModel_export_to_onnx(self, dummy_data): + with tempfile.TemporaryDirectory() as tmpdir: + wget.download( + 'https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/' + 'nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml', + tmpdir, + ) + config_file = os.path.join(tmpdir, 'intent_slot_classification_config.yaml') + config = OmegaConf.load(config_file) + config = OmegaConf.create(OmegaConf.to_container(config, resolve=True)) + config.model.data_dir = dummy_data + config.trainer.gpus = 1 + config.trainer.precision = 32 + config.trainer.accelerator = None + trainer = pl.Trainer(**config.trainer) + model = IntentSlotClassificationModel(config.model, trainer=trainer) + filename = os.path.join(tmpdir, 'isc.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'input_ids' + assert onnx_model.graph.input[1].name == 'attention_mask' + assert onnx_model.graph.input[2].name == 'token_type_ids' + assert onnx_model.graph.output[0].name == 'intent_logits' + assert onnx_model.graph.output[1].name == 'slot_logits' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + @pytest.mark.skip('Are we exporting to torchscript?') + def test_IntentSlotClassificationModel_export_to_torchscript(self, dummy_data): + with tempfile.TemporaryDirectory() as tmpdir: + wget.download( + 'https://raw.githubusercontent.com/NVIDIA/NeMo/main/examples/' + 'nlp/intent_slot_classification/conf/intent_slot_classification_config.yaml', + tmpdir, + ) + config_file = os.path.join(tmpdir, 'intent_slot_classification_config.yaml') + config = OmegaConf.load(config_file) + config = OmegaConf.create(OmegaConf.to_container(config, resolve=True)) + config.model.data_dir = dummy_data + config.trainer.gpus = 1 + config.trainer.precision = 32 + config.trainer.accelerator = None + trainer = pl.Trainer(**config.trainer) + model = IntentSlotClassificationModel(config.model, trainer=trainer) + filename = os.path.join(tmpdir, 'isc.pt') + model.export(output=filename, check_trace=True) + torchscript_model = torch.jit.load(filename) + inputs = list(torchscript_model.graph.inputs()) + assert inputs[1].debugName() == 'input_ids.1' + assert inputs[2].debugName() == 'input.1' + assert inputs[3].debugName() == 'attention_mask.1' + modules = list(torchscript_model.modules()) + assert modules[2].original_name == 'BertEmbeddings' + assert modules[46].original_name == 'BertSelfAttention' + assert modules[217].original_name == 'SequenceTokenClassifier' + assert modules[219].original_name == 'MultiLayerPerceptron' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_TokenClassificationModel_export_to_onnx(self): + model = nemo_nlp.models.TokenClassificationModel.from_pretrained(model_name="ner_en_bert") + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'ner.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'input_ids' + assert onnx_model.graph.output[0].name == 'logits' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_PunctuationCapitalizationModel_export_to_onnx(self): + model = nemo_nlp.models.PunctuationCapitalizationModel.from_pretrained(model_name="punctuation_en_distilbert") + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'puncap.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'input_ids' + # assert onnx_model.graph.input[2].name == 'token_type_ids' + assert onnx_model.graph.output[0].name == 'punct_logits' + assert onnx_model.graph.output[1].name == 'capit_logits' + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_QAModel_export_to_onnx(self): + model = nemo_nlp.models.QAModel.from_pretrained(model_name="qa_squadv2.0_bertbase") + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'qa.onnx') + model.export(output=filename) + onnx_model = onnx.load(filename) + onnx.checker.check_model(onnx_model, full_check=True) # throws when failed + assert onnx_model.graph.input[0].name == 'input_ids' + assert onnx_model.graph.input[2].name == 'token_type_ids' + assert onnx_model.graph.output[0].name == 'logits' + + +@pytest.fixture() +def dummy_data(test_data_dir): + return os.path.join(test_data_dir, 'nlp', 'dummy_data') diff --git a/models/tts/NeMo/tests/collections/nlp/test_nmt_model.py b/models/tts/NeMo/tests/collections/nlp/test_nmt_model.py new file mode 100644 index 0000000000..9f0277f463 --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_nmt_model.py @@ -0,0 +1,144 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import shutil +import tempfile + +import pytest +import torch +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.nlp.models import MTEncDecModel +from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import AAYNBaseConfig + + +def export_test(model, suffix, try_script=False): + with tempfile.TemporaryDirectory() as restore_folder: + enc_filename = os.path.join(restore_folder, 'nmt_e' + suffix) + dec_filename = os.path.join(restore_folder, 'nmt_d' + suffix) + model.encoder.export(output=enc_filename, try_script=try_script) + model.decoder.export(output=dec_filename, try_script=try_script) + assert os.path.exists(enc_filename) + assert os.path.exists(dec_filename) + + +def get_cfg(): + cfg = AAYNBaseConfig() + cfg.encoder_tokenizer.tokenizer_name = 'yttm' + cfg.encoder_tokenizer.tokenizer_model = 'tests/.data/yttm.4096.en-de.model' + cfg.decoder_tokenizer.tokenizer_name = 'yttm' + cfg.decoder_tokenizer.tokenizer_model = 'tests/.data/yttm.4096.en-de.model' + cfg.train_ds = None + cfg.validation_ds = None + cfg.test_ds = None + return cfg + + +class TestMTEncDecModel: + @pytest.mark.unit + def test_creation_saving_restoring(self): + model = MTEncDecModel(cfg=get_cfg()) + assert isinstance(model, MTEncDecModel) + # Create a new temporary directory + with tempfile.TemporaryDirectory() as restore_folder: + with tempfile.TemporaryDirectory() as save_folder: + save_folder_path = save_folder + # Where model will be saved + model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo") + model.save_to(save_path=model_save_path) + # Where model will be restored from + model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo") + shutil.copy(model_save_path, model_restore_path) + # at this point save_folder should not exist + assert save_folder_path is not None and not os.path.exists(save_folder_path) + assert not os.path.exists(model_save_path) + assert os.path.exists(model_restore_path) + # attempt to restore + model_copy = model.__class__.restore_from(restore_path=model_restore_path) + assert model.num_weights == model_copy.num_weights + + @pytest.mark.unit + def test_no_artifact_name_collision(self): + model = MTEncDecModel(cfg=get_cfg()) + assert isinstance(model, MTEncDecModel) + with tempfile.TemporaryDirectory() as tmpdir1: + model.save_to("nmt_model.nemo") + with tempfile.TemporaryDirectory() as tmpdir: + model._unpack_nemo_file(path2file="nmt_model.nemo", out_folder=tmpdir) + conf = OmegaConf.load(os.path.join(tmpdir, "model_config.yaml")) + # Make sure names now differ in saved config + assert conf.encoder_tokenizer.tokenizer_model != conf.decoder_tokenizer.tokenizer_model + # Make sure names in config start with "nemo:" prefix + assert conf.encoder_tokenizer.tokenizer_model.startswith("nemo:") + assert conf.decoder_tokenizer.tokenizer_model.startswith("nemo:") + # Check if both tokenizers were included + assert os.path.exists(os.path.join(tmpdir, conf.encoder_tokenizer.tokenizer_model[5:])) + assert os.path.exists(os.path.join(tmpdir, conf.decoder_tokenizer.tokenizer_model[5:])) + + @pytest.mark.unit + def test_train_eval_loss(self): + cfg = get_cfg() + cfg.label_smoothing = 0.5 + model = MTEncDecModel(cfg=cfg) + assert isinstance(model, MTEncDecModel) + batch_size = 10 + time = 32 + vocab_size = 32000 + torch.manual_seed(42) + tgt_ids = torch.LongTensor(batch_size, time).random_(1, model.decoder_tokenizer.vocab_size) + logits = torch.FloatTensor(batch_size, time, vocab_size).random_(-1, 1) + log_probs = torch.nn.functional.log_softmax(logits, dim=-1) + train_loss = model.loss_fn(log_probs=log_probs, labels=tgt_ids) + eval_loss = model.eval_loss_fn(log_probs=log_probs, labels=tgt_ids) + assert not torch.allclose(train_loss, eval_loss) # , (train_loss, eval_loss) + + cfg.label_smoothing = 0 + model = MTEncDecModel(cfg=cfg) + # Train loss == val loss when label smoothing = 0 + train_loss = model.loss_fn(log_probs=log_probs, labels=tgt_ids) + eval_loss = model.eval_loss_fn(log_probs=log_probs, labels=tgt_ids) + assert torch.allclose(train_loss, eval_loss) + + @pytest.mark.unit + def test_cpu_export_onnx(self): + model = MTEncDecModel(cfg=get_cfg()) + assert isinstance(model, MTEncDecModel) + export_test(model, ".ts") + + @pytest.mark.unit + def test_cpu_export_ts(self): + model = MTEncDecModel(cfg=get_cfg()) + assert isinstance(model, MTEncDecModel) + export_test(model, ".onnx") + + @pytest.mark.skipif(not os.path.exists('/home/TestData/nlp'), reason='Not a Jenkins machine') + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_gpu_export_ts(self): + model = MTEncDecModel(cfg=get_cfg()).cuda() + assert isinstance(model, MTEncDecModel) + export_test(model, ".ts") + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_gpu_export_onnx(self): + model = MTEncDecModel(cfg=get_cfg()).cuda() + assert isinstance(model, MTEncDecModel) + export_test(model, ".onnx") + + +if __name__ == "__main__": + t = TestMTEncDecModel() + # t.test_gpu_export_ts() + t.test_train_eval_loss() diff --git a/models/tts/NeMo/tests/collections/nlp/test_pretrained_models_performance.py b/models/tts/NeMo/tests/collections/nlp/test_pretrained_models_performance.py new file mode 100644 index 0000000000..47f4f7640b --- /dev/null +++ b/models/tts/NeMo/tests/collections/nlp/test_pretrained_models_performance.py @@ -0,0 +1,110 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +from pathlib import Path +from shutil import rmtree +from unittest import TestCase + +import pytest +import pytorch_lightning as pl +from omegaconf import OmegaConf + +import nemo.collections.nlp.models as models + + +def get_metrics(data_dir, model): + trainer = pl.Trainer(gpus=[0]) + + model.set_trainer(trainer) + model.update_data_dir(data_dir) + + test_ds = OmegaConf.create( + { + 'text_file': 'text_dev.txt', + 'labels_file': 'labels_dev.txt', + 'shuffle': False, + 'num_samples': -1, + 'batch_size': 8, + } + ) + + model._cfg.dataset.use_cache = False + trainer.test(model) + model.setup_test_data(test_data_config=test_ds) + metrics = trainer.test(model)[0] + + return metrics + + +def data_exists(data_dir): + return os.path.exists(data_dir) + + +class TestPretrainedModelPerformance: + @pytest.mark.with_downloads() + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + @pytest.mark.skipif( + not data_exists('/home/TestData/nlp/token_classification_punctuation/fisher'), reason='Not a Jenkins machine' + ) + def test_punct_capit_with_bert(self): + data_dir = '/home/TestData/nlp/token_classification_punctuation/fisher' + model = models.PunctuationCapitalizationModel.from_pretrained("punctuation_en_bert") + metrics = get_metrics(data_dir, model) + + assert abs(metrics['punct_precision'] - 52.3024) < 0.001 + assert abs(metrics['punct_recall'] - 58.9220) < 0.001 + assert abs(metrics['punct_f1'] - 53.2976) < 0.001 + assert abs(metrics['capit_precision'] - 87.0707) < 0.001 + assert abs(metrics['capit_recall'] - 87.0707) < 0.001 + assert abs(metrics['capit_f1'] - 87.0707) < 0.001 + assert int(model.punct_class_report.total_examples) == 128 + + preds_512 = model.add_punctuation_capitalization(['what can i do for you today'], max_seq_length=512)[0] + assert preds_512 == 'What can I do for you today?' + preds_5 = model.add_punctuation_capitalization(['what can i do for you today'], max_seq_length=5)[0] + assert preds_5 == 'What can i' + + @pytest.mark.with_downloads() + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + @pytest.mark.skipif( + not data_exists('/home/TestData/nlp/token_classification_punctuation/fisher'), reason='Not a Jenkins machine' + ) + def test_punct_capit_with_distilbert(self): + data_dir = '/home/TestData/nlp/token_classification_punctuation/fisher' + model = models.PunctuationCapitalizationModel.from_pretrained("punctuation_en_distilbert") + metrics = get_metrics(data_dir, model) + + assert abs(metrics['punct_precision'] - 53.0826) < 0.001 + assert abs(metrics['punct_recall'] - 56.2905) < 0.001 + assert abs(metrics['punct_f1'] - 52.4225) < 0.001 + assert int(model.punct_class_report.total_examples) == 128 + + @pytest.mark.with_downloads() + @pytest.mark.unit + @pytest.mark.run_only_on('GPU') + @pytest.mark.skipif( + not data_exists('/home/TestData/nlp/token_classification_punctuation/gmb'), reason='Not a Jenkins machine' + ) + def test_ner_model(self): + data_dir = '/home/TestData/nlp/token_classification_punctuation/gmb' + model = models.TokenClassificationModel.from_pretrained("ner_en_bert") + metrics = get_metrics(data_dir, model) + + assert abs(metrics['precision'] - 96.0937) < 0.001 + assert abs(metrics['recall'] - 96.0146) < 0.001 + assert abs(metrics['f1'] - 95.6076) < 0.001 + assert int(model.classification_report.total_examples) == 202 diff --git a/models/tts/NeMo/tests/collections/tts/test_tts_datasets.py b/models/tts/NeMo/tests/collections/tts/test_tts_datasets.py new file mode 100644 index 0000000000..04d557ef4b --- /dev/null +++ b/models/tts/NeMo/tests/collections/tts/test_tts_datasets.py @@ -0,0 +1,43 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +import pytest +import torch + +from nemo.collections.tts.data.datalayers import FastSpeech2Dataset + + +class TestTTSDatasets: + @pytest.mark.unit + def test_fs2_dataset(self, test_data_dir): + manifest_path = os.path.join(test_data_dir, 'tts/mini_ljspeech/manifest.json') + mappings_file = os.path.join(test_data_dir, 'tts/mini_ljspeech/mappings.json') + ignore_file = os.path.join(test_data_dir, 'tts/mini_ljspeech/wavs_to_ignore.pkl') + + # Test loading data (including supplementary data) with ignore file + ds = FastSpeech2Dataset( + manifest_filepath=manifest_path, + mappings_filepath=mappings_file, + sample_rate=22050, + ignore_file=ignore_file, + load_supplementary_values=True, + ) + + assert len(ds) == 4 + count = 0 + for _ in ds: + count += 1 + assert count == 4 diff --git a/models/tts/NeMo/tests/collections/tts/test_tts_exportables.py b/models/tts/NeMo/tests/collections/tts/test_tts_exportables.py new file mode 100644 index 0000000000..59075d9e8f --- /dev/null +++ b/models/tts/NeMo/tests/collections/tts/test_tts_exportables.py @@ -0,0 +1,62 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import tempfile + +import onnx +import pytest +from omegaconf import DictConfig, ListConfig, OmegaConf + +from nemo.collections.tts.models import FastPitchModel, HifiGanModel, TalkNetSpectModel, WaveGlowModel + + +@pytest.fixture() +def fastpitch_model(): + test_root = os.path.dirname(os.path.abspath(__file__)) + conf = OmegaConf.load(os.path.join(test_root, '../../../examples/tts/conf/fastpitch.yaml')) + conf.train_dataset = conf.validation_datasets = '.' + conf.model.train_ds = conf.model.test_ds = conf.model.validation_ds = None + model = FastPitchModel(cfg=conf.model) + return model + + +@pytest.fixture() +def hifigan_model(): + test_root = os.path.dirname(os.path.abspath(__file__)) + model = HifiGanModel.from_pretrained(model_name="tts_hifigan") + return model + + +class TestExportable: + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_FastPitchModel_export_to_onnx(self, fastpitch_model): + model = fastpitch_model.cuda() + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'fp.onnx') + model.export(output=filename, verbose=True, check_trace=True) + + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_HifiGanModel_export_to_onnx(self, hifigan_model): + model = hifigan_model.cuda() + assert hifigan_model.generator is not None + with tempfile.TemporaryDirectory() as tmpdir: + filename = os.path.join(tmpdir, 'hfg.pt') + model.export(output=filename, verbose=True, check_trace=True) + + +if __name__ == "__main__": + t = TestExportable() + t.test_FastPitchModel_export_to_onnx(fastpitch_model()) diff --git a/models/tts/NeMo/tests/collections/tts/test_waveglow.py b/models/tts/NeMo/tests/collections/tts/test_waveglow.py new file mode 100644 index 0000000000..7bfd3e5b0e --- /dev/null +++ b/models/tts/NeMo/tests/collections/tts/test_waveglow.py @@ -0,0 +1,110 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import tempfile +from unittest import TestCase + +import onnx +import pytest +import torch +from omegaconf import DictConfig + +from nemo.collections.tts.models import WaveGlowModel +from nemo.collections.tts.modules import WaveGlowModule +from nemo.core.classes import typecheck + +mcfg = DictConfig( + { + "_target_": "nemo.collections.tts.modules.waveglow.WaveGlowModule", + "n_flows": 12, + "n_group": 8, + "n_mel_channels": 80, + "n_early_every": 4, + "n_early_size": 2, + "n_wn_channels": 512, + "n_wn_layers": 8, + "wn_kernel_size": 3, + } +) + +pcfg = DictConfig( + { + "_target_": "nemo.collections.asr.parts.preprocessing.features.FilterbankFeatures", + "dither": 0.0, + "nfilt": 80, + "stft_conv": True, + } +) + +wcfg = DictConfig({"waveglow": mcfg, "sigma": 1.0, "preprocessor": pcfg,}) + + +def input_example(sz): + mel = torch.randn(1, 1, 80, sz).cuda().half() + z = torch.randn(1, 8, sz * 256 // 8, 1).cuda().half() + return ( + mel, + z, + ) + + +def taco2wg(spec, z): + spec = spec.permute(0, 3, 2, 1).contiguous() + return spec.view(spec.size(0), spec.size(1), -1), z.view(z.size(0), z.size(1), -1) + + +# Wrapper method to convert Jasper's Taco2 output to WG input and call inference +def forward_wrapper(self, spec, z=None): + spec, z = taco2wg(spec, z) + audio = self.waveglow.norm_dist_to_audio(spec=spec, sigma=1.0, z=z) + return audio + + +class TestWaveGlow: + @pytest.mark.run_only_on('GPU') + @pytest.mark.unit + def test_export_to_onnx(self): + model = WaveGlowModel(wcfg) + model = model.cuda().half() + typecheck.set_typecheck_enabled(enabled=False) + with tempfile.TemporaryDirectory() as tmpdir, model.nemo_infer(): + # Generate filename in the temporary directory. + # TODO: Change `waveglow.ts` to `waveglow.onnx` for > 21.05 + tmp_file_name = os.path.join("waveglow.ts") + + n_mels = 80 + # Test export. + inp = input_example(n_mels) + inp1 = taco2wg(*inp) + inp2 = inp1 + res1 = model.waveglow(*inp1) + res2 = model.waveglow(*inp2) + assert torch.allclose(res1, res2, rtol=0.01, atol=0.1) + WaveGlowModel.forward_for_export = forward_wrapper + model.export( + tmp_file_name, + verbose=True, + input_example=inp, + output_example=res1, + try_script=False, + check_trace=False, + do_constant_folding=True, + dynamic_axes={"spec": [0], "z": [0], "audio": [0]}, + ) + + +if __name__ == "__main__": + t = TestWaveGlow() + t.test_export_to_onnx() diff --git a/models/tts/NeMo/tests/conftest.py b/models/tts/NeMo/tests/conftest.py new file mode 100644 index 0000000000..4a947a2a04 --- /dev/null +++ b/models/tts/NeMo/tests/conftest.py @@ -0,0 +1,188 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 tarfile +import urllib.request +from os import mkdir +from os.path import dirname, exists, getsize, join +from pathlib import Path +from shutil import rmtree + +import pytest + +# Those variables probably should go to main NeMo configuration file (config.yaml). +__TEST_DATA_FILENAME = "test_data.tar.gz" +__TEST_DATA_URL = "https://github.com/NVIDIA/NeMo/releases/download/v1.0.0rc1/" +__TEST_DATA_SUBDIR = ".data" + + +def pytest_addoption(parser): + """ + Additional command-line arguments passed to pytest. + For now: + --cpu: use CPU during testing (DEFAULT: GPU) + --use_local_test_data: use local test data/skip downloading from URL/GitHub (DEFAULT: False) + """ + parser.addoption( + '--cpu', action='store_true', help="pass that argument to use CPU during testing (DEFAULT: False = GPU)" + ) + parser.addoption( + '--use_local_test_data', + action='store_true', + help="pass that argument to use local test data/skip downloading from URL/GitHub (DEFAULT: False)", + ) + parser.addoption( + '--with_downloads', + action='store_true', + help="pass this argument to active tests which download models from the cloud.", + ) + parser.addoption( + '--relax_numba_compat', + action='store_false', + help="numba compatibility checks will be relaxed to just availability of cuda, " + "without cuda compatibility matrix check", + ) + + +@pytest.fixture +def device(request): + """ Simple fixture returning string denoting the device [CPU | GPU] """ + if request.config.getoption("--cpu"): + return "CPU" + else: + return "GPU" + + +@pytest.fixture(autouse=True) +def run_only_on_device_fixture(request, device): + if request.node.get_closest_marker('run_only_on'): + if request.node.get_closest_marker('run_only_on').args[0] != device: + pytest.skip('skipped on this device: {}'.format(device)) + + +@pytest.fixture(autouse=True) +def downloads_weights(request, device): + if request.node.get_closest_marker('with_downloads'): + if not request.config.getoption("--with_downloads"): + pytest.skip( + 'To run this test, pass --with_downloads option. It will download (and cache) models from cloud.' + ) + + +@pytest.fixture(autouse=True) +def cleanup_local_folder(): + # Asserts in fixture are not recommended, but I'd rather stop users from deleting expensive training runs + assert not Path("./lightning_logs").exists() + assert not Path("./NeMo_experiments").exists() + assert not Path("./nemo_experiments").exists() + + yield + + if Path("./lightning_logs").exists(): + rmtree('./lightning_logs', ignore_errors=True) + if Path("./NeMo_experiments").exists(): + rmtree('./NeMo_experiments', ignore_errors=True) + if Path("./nemo_experiments").exists(): + rmtree('./nemo_experiments', ignore_errors=True) + + +@pytest.fixture +def test_data_dir(): + """ Fixture returns test_data_dir. """ + # Test dir. + test_data_dir_ = join(dirname(__file__), __TEST_DATA_SUBDIR) + return test_data_dir_ + + +def pytest_configure(config): + """ + Initial configuration of conftest. + The function checks if test_data.tar.gz is present in tests/.data. + If so, compares its size with github's test_data.tar.gz. + If file absent or sizes not equal, function downloads the archive from github and unpacks it. + """ + config.addinivalue_line( + "markers", "run_only_on(device): runs the test only on a given device [CPU | GPU]", + ) + # Test dir and archive filepath. + test_dir = join(dirname(__file__), __TEST_DATA_SUBDIR) + test_data_archive = join(dirname(__file__), __TEST_DATA_SUBDIR, __TEST_DATA_FILENAME) + + # Get size of local test_data archive. + try: + test_data_local_size = getsize(test_data_archive) + except: + # File does not exist. + test_data_local_size = -1 + + if config.option.use_local_test_data: + if test_data_local_size == -1: + pytest.exit("Test data `{}` is not present in the system".format(test_data_archive)) + else: + print( + "Using the local `{}` test archive ({}B) found in the `{}` folder.".format( + __TEST_DATA_FILENAME, test_data_local_size, test_dir + ) + ) + return + + # Get size of remote test_data archive. + try: + url = __TEST_DATA_URL + __TEST_DATA_FILENAME + u = urllib.request.urlopen(url) + except: + # Couldn't access remote archive. + if test_data_local_size == -1: + pytest.exit("Test data not present in the system and cannot access the '{}' URL".format(url)) + else: + print( + "Cannot access the '{}' URL, using the test data ({}B) found in the `{}` folder.".format( + url, test_data_local_size, test_dir + ) + ) + return + + # Get metadata. + meta = u.info() + test_data_remote_size = int(meta["Content-Length"]) + + # Compare sizes. + if test_data_local_size != test_data_remote_size: + print( + "Downloading the `{}` test archive from `{}`, please wait...".format(__TEST_DATA_FILENAME, __TEST_DATA_URL) + ) + # Remove .data folder. + if exists(test_dir): + rmtree(test_dir) + # Create one .data folder. + mkdir(test_dir) + # Download + urllib.request.urlretrieve(url, test_data_archive) + # Extract tar + print("Extracting the `{}` test archive, please wait...".format(test_data_archive)) + tar = tarfile.open(test_data_archive) + tar.extractall(path=test_dir) + tar.close() + else: + print( + "A valid `{}` test archive ({}B) found in the `{}` folder.".format( + __TEST_DATA_FILENAME, test_data_local_size, test_dir + ) + ) + + if config.option.relax_numba_compat is not None: + from nemo.core.utils import numba_utils + + print("Setting numba compat :", config.option.relax_numba_compat) + numba_utils.set_numba_compat_strictness(strict=config.option.relax_numba_compat) diff --git a/models/tts/NeMo/tests/core/test_config_utils.py b/models/tts/NeMo/tests/core/test_config_utils.py new file mode 100644 index 0000000000..4b0d8bb67a --- /dev/null +++ b/models/tts/NeMo/tests/core/test_config_utils.py @@ -0,0 +1,125 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +from dataclasses import dataclass +from typing import Any + +import pytest +import pytorch_lightning as ptl + +from nemo.core.config.pytorch_lightning import TrainerConfig +from nemo.utils import config_utils + + +@pytest.fixture() +def cls(): + class DummyClass: + def __init__(self, a, b=5, c: int = 0, d: 'ABC' = None): + pass + + return DummyClass + + +class TestConfigUtils: + @pytest.mark.unit + def test_all_args_exist(self, cls): + @dataclass + class DummyDataClass: + a: int = -1 + b: int = 5 + c: int = 0 + d: Any = None + + result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_all_args_dont_exist(self, cls): + @dataclass + class DummyDataClass: + a: int = -1 + b: int = 5 + c: int = 0 + + result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass) + signatures_match, cls_subset, dataclass_subset = result + + assert not signatures_match + assert len(cls_subset) > 0 + assert len(dataclass_subset) == 0 + + @pytest.mark.unit + def test_extra_args_exist(self, cls): + @dataclass + class DummyDataClass: + a: int = -1 + b: int = 5 + c: int = 0 + d: Any = None + e: float = 0.0 + + result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass) + signatures_match, cls_subset, dataclass_subset = result + + assert not signatures_match + assert len(cls_subset) == 0 + assert len(dataclass_subset) > 0 + + @pytest.mark.unit + def test_extra_args_exist_but_is_ignored(self, cls): + @dataclass + class DummyDataClass: + a: int = -1 + b: int = 5 + c: int = 0 + d: Any = None + e: float = 0.0 # Assume ignored + + result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass, ignore_args=['e']) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_args_exist_but_is_remapped(self, cls): + @dataclass + class DummyDataClass: + a: int = -1 + b: int = 5 + c: int = 0 + e: Any = None # Assume remapped + + result = config_utils.assert_dataclass_signature_match(cls, DummyDataClass, remap_args={'e': 'd'}) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None + + @pytest.mark.unit + def test_ptl_config(self): + PTL_DEPRECATED = ['distributed_backend', 'automatic_optimization'] + + result = config_utils.assert_dataclass_signature_match(ptl.Trainer, TrainerConfig, ignore_args=PTL_DEPRECATED) + signatures_match, cls_subset, dataclass_subset = result + + assert signatures_match + assert cls_subset is None + assert dataclass_subset is None diff --git a/models/tts/NeMo/tests/core/test_exp_manager.py b/models/tts/NeMo/tests/core/test_exp_manager.py new file mode 100644 index 0000000000..e1fa6da95c --- /dev/null +++ b/models/tts/NeMo/tests/core/test_exp_manager.py @@ -0,0 +1,357 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 re +from pathlib import Path + +import pytest +import pytorch_lightning as pl +import torch +from omegaconf import OmegaConf +from omegaconf.errors import OmegaConfBaseException + +from nemo.constants import NEMO_ENV_VARNAME_VERSION +from nemo.core.classes import ModelPT +from nemo.utils.exp_manager import ( + CheckpointMisconfigurationError, + LoggerMisconfigurationError, + NotFoundError, + exp_manager, +) + + +class MyTestOptimizer(torch.optim.Optimizer): + def __init__(self, params): + self._step = 0 + super().__init__(params, {}) + + def step(self, *args, **kwargs): + for group in self.param_groups: + for p in group['params']: + if self._step == 0: + p.data = 0.1 * torch.ones(p.shape) + elif self._step == 1: + p.data = 0.0 * torch.ones(p.shape) + else: + p.data = 0.01 * torch.ones(p.shape) + self._step += 1 + return None + + +class OnesDataset(torch.utils.data.Dataset): + def __init__(self, dataset_len): + super().__init__() + self.__dataset_len = dataset_len + + def __getitem__(self, *args): + return torch.ones(2) + + def __len__(self): + return self.__dataset_len + + +class ExampleModel(ModelPT): + def __init__(self, *args, **kwargs): + cfg = OmegaConf.structured({}) + super().__init__(cfg) + self.l1 = torch.nn.modules.Linear(in_features=2, out_features=1) + + def train_dataloader(self): + dataset = OnesDataset(2) + return torch.utils.data.DataLoader(dataset, batch_size=2) + + def val_dataloader(self): + dataset = OnesDataset(10) + return torch.utils.data.DataLoader(dataset, batch_size=2) + + def forward(self, batch): + output = self.l1(batch) + output = torch.nn.functional.l1_loss(output, torch.zeros(output.size()).to(output.device)) + return output + + def validation_step(self, batch, batch_idx): + return self(batch) + + def training_step(self, batch, batch_idx): + return self(batch) + + def configure_optimizers(self): + return MyTestOptimizer(self.parameters()) + # return torch.optim.Adam(self.parameters(), lr=0.1) + + def list_available_models(self): + pass + + def setup_training_data(self): + pass + + def setup_validation_data(self): + pass + + def validation_epoch_end(self, loss): + self.log("val_loss", torch.stack(loss).mean()) + + +class TestExpManager: + @pytest.mark.unit + def test_omegaconf(self): + """Ensure omegaconf raises an error when an unexcepted argument is passed""" + with pytest.raises(OmegaConfBaseException): + exp_manager(pl.Trainer(), {"unused": 1}) + + @pytest.mark.unit + def test_trainer_loggers(self, tmp_path): + """ Test that a trainer with logger errors out with a number of arguments. Test that it works with + create_tensorboard_logger set to False + """ + test_trainer = pl.Trainer() # Should create logger and modelcheckpoint + + with pytest.raises(LoggerMisconfigurationError): # Fails because exp_manager defaults to trainer + exp_manager(test_trainer, {"exp_dir": str(tmp_path)}) + with pytest.raises(LoggerMisconfigurationError): # Fails because exp_manager defaults to trainer + exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path)}) + with pytest.raises(LoggerMisconfigurationError): # Fails because exp_manager defaults to trainer + exp_manager(test_trainer, {"resume_if_exists": True}) + + # Check that exp_manager uses trainer.logger, it's exp_dir, name, and version + log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False}) + assert log_dir.resolve() == Path("./lightning_logs/version_0").resolve() + assert Path("./lightning_logs").exists() + assert Path("./lightning_logs/version_0").exists() + + # Check that a trainer without a logger gets a logger attached to it + test_trainer = pl.Trainer(logger=False) + log_dir = exp_manager( + test_trainer, + {"create_tensorboard_logger": True, "create_checkpoint_callback": False, "exp_dir": str(tmp_path)}, + ) + assert isinstance(test_trainer.logger, pl.loggers.TensorBoardLogger) + + test_trainer = pl.Trainer(logger=False) + # Check that a create_wandb_logger=True errors out unless wandb_logger_kwargs is passed. + with pytest.raises(ValueError): + log_dir = exp_manager( + test_trainer, + { + "create_tensorboard_logger": False, + "create_checkpoint_callback": False, + "exp_dir": str(tmp_path), + "create_wandb_logger": True, + }, + ) + # Check that a WandbLogger is attached to logger if create_wandb_logger=True and wandb_logger_kwargs has name + # and project + log_dir = exp_manager( + test_trainer, + { + "create_tensorboard_logger": False, + "create_checkpoint_callback": False, + "exp_dir": str(tmp_path), + "create_wandb_logger": True, + "wandb_logger_kwargs": {"name": "", "project": ""}, + }, + ) + assert isinstance(test_trainer.logger, pl.loggers.WandbLogger) + + @pytest.mark.unit + def test_checkpoint_configurations(self): + """ Test that trainer creating modelcheckpoint and asking exp_manager to do it too results in errors, but + is error free if only one is asked to do so. + """ + disable_tb_logger = {"create_tensorboard_logger": False} + test_trainer = pl.Trainer() # Should create logger and modelcheckpoint + with pytest.raises(CheckpointMisconfigurationError): # Fails because both try to create modelcheckpoint + exp_manager(test_trainer, disable_tb_logger) + + # Should succeed without error + exp_manager(test_trainer, {"create_checkpoint_callback": False, "create_tensorboard_logger": False}) + + test_trainer_2 = pl.Trainer(checkpoint_callback=False) + exp_manager(test_trainer_2, disable_tb_logger) # Should succeed without error + + @pytest.mark.unit + def test_default_log_dir(self): + """Check the default of ./nemo_experiments/default/datetime works as intended""" + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + + log_dir = exp_manager(test_trainer, {"create_tensorboard_logger": False, "create_checkpoint_callback": False}) + assert (log_dir / "..").resolve() == Path("./nemo_experiments/default/").resolve() + assert Path("./nemo_experiments").exists() + assert Path("./nemo_experiments/default/").exists() + sub_dirs = [x for x in Path("./nemo_experiments/default/").iterdir() if x.is_dir()] + assert len(sub_dirs) == 1 + assert re.match(r"[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}", sub_dirs[0].name) + + @pytest.mark.unit + def test_log_dir_overrides(self, monkeypatch, tmp_path): + """Check a variety of trainer options with exp_manager""" + # Checks that explicit_log_dir ignores exp_dir, name, and version + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + log_dir = exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path / "test_log_dir_overrides")}) + assert log_dir.resolve() == (tmp_path / "test_log_dir_overrides").resolve() + assert Path(tmp_path).exists() + assert Path(tmp_path / "test_log_dir_overrides").exists() + + # Checks that exp_manager uses exp_dir, default name, and explicit version + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "version": 957}) + assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "957").resolve() + assert Path(tmp_path).exists() + assert Path(tmp_path / "test_no_name" / "default" / "957").exists() + + monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION) + # Checks that use_datetime_version False toggle works + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False}) + assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_0").resolve() + assert Path(tmp_path).exists() + assert Path(tmp_path / "test_no_name" / "default" / "version_0").exists() + + monkeypatch.delenv(NEMO_ENV_VARNAME_VERSION) + # Checks that use_datetime_version False toggle works and version increments + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + log_dir = exp_manager(test_trainer, {"exp_dir": str(tmp_path / "test_no_name"), "use_datetime_version": False}) + assert log_dir.resolve() == (tmp_path / "test_no_name" / "default" / "version_1").resolve() + assert Path(tmp_path).exists() + assert Path(tmp_path / "test_no_name" / "default" / "version_1").exists() + + @pytest.mark.unit + def test_resume(self, tmp_path): + """ Tests the resume capabilities of exp_manager""" + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + + # Error because explicit_log_dir does not exist + with pytest.raises(NotFoundError): + exp_manager( + test_trainer, + { + "exp_dir": str(tmp_path / "test_resume"), + "resume_if_exists": True, + "explicit_log_dir": "Does_not_exist", + }, + ) + + # Error because checkpoints folder does not exist + with pytest.raises(NotFoundError): + exp_manager(test_trainer, {"resume_if_exists": True, "exp_dir": str(tmp_path / "test_resume")}) + + # No error because we tell exp_manager to ignore notfounderror + exp_manager( + test_trainer, + { + "resume_if_exists": True, + "exp_dir": str(tmp_path / "test_resume_2"), + "resume_ignore_no_checkpoint": True, + }, + ) + + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints").mkdir(parents=True) + # Error because checkpoints do not exist in folder + with pytest.raises(NotFoundError): + exp_manager( + test_trainer, + { + "resume_if_exists": True, + "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"), + }, + ) + + Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").touch() + # Error because *end.ckpt is in folder indicating that training has already finished + with pytest.raises(ValueError): + exp_manager( + test_trainer, + { + "resume_if_exists": True, + "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"), + }, + ) + + Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--end.ckpt").unlink() + Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt").touch() + Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").touch() + # Error because multiple *last.ckpt is in folder. If more than one, don't know which to restore + with pytest.raises(ValueError): + exp_manager( + test_trainer, + { + "resume_if_exists": True, + "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0"), + }, + ) + + # Finally succeed + Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel2--last.ckpt").unlink() + log_dir = exp_manager( + test_trainer, + {"resume_if_exists": True, "explicit_log_dir": str(tmp_path / "test_resume" / "default" / "version_0")}, + ) + checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt") + assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve() + + # Succeed again and make sure that run_0 exists and previous log files were moved + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False) + exp_manager(test_trainer, {"resume_if_exists": True, "explicit_log_dir": str(log_dir)}) + checkpoint = Path(tmp_path / "test_resume" / "default" / "version_0" / "checkpoints" / "mymodel--last.ckpt") + assert Path(test_trainer.resume_from_checkpoint).resolve() == checkpoint.resolve() + prev_run_dir = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0") + assert prev_run_dir.exists() + prev_log = Path(tmp_path / "test_resume" / "default" / "version_0" / "run_0" / "lightning_logs.txt") + assert prev_log.exists() + + @pytest.mark.unit + def test_nemo_checkpoint_save_best_model_1(self, tmp_path): + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False, max_epochs=4) + log_dir = exp_manager( + test_trainer, + {"checkpoint_callback_params": {"save_best_model": True}, "explicit_log_dir": str(tmp_path / "test")}, + ) + model = ExampleModel() + test_trainer.fit(model) + + assert Path(str(tmp_path / "test" / "checkpoints" / "default.nemo")).exists() + + model = ExampleModel.restore_from(str(tmp_path / "test" / "checkpoints" / "default.nemo")) + assert float(model(torch.tensor([1.0, 1.0], device=model.device))) == 0.0 + + @pytest.mark.unit + def test_nemo_checkpoint_save_best_model_2(self, tmp_path): + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False, max_epochs=4) + log_dir = exp_manager(test_trainer, {"explicit_log_dir": str(tmp_path / "test")},) + model = ExampleModel() + test_trainer.fit(model) + + assert Path(str(tmp_path / "test" / "checkpoints" / "default.nemo")).exists() + + model = ExampleModel.restore_from(str(tmp_path / "test" / "checkpoints" / "default.nemo")) + assert math.fabs(float(model(torch.tensor([1.0, 1.0], device=model.device))) - 0.03) < 1e-5 + + @pytest.mark.unit + def test_nemo_checkpoint_always_save_nemo(self, tmp_path): + test_trainer = pl.Trainer(checkpoint_callback=False, logger=False, max_epochs=4) + log_dir = exp_manager( + test_trainer, + { + "checkpoint_callback_params": {"save_best_model": True, "always_save_nemo": True}, + "explicit_log_dir": str(tmp_path / "test"), + }, + ) + model = ExampleModel() + test_trainer.fit(model) + + assert Path(str(tmp_path / "test" / "checkpoints" / "default.nemo")).exists() + + model = ExampleModel.restore_from(str(tmp_path / "test" / "checkpoints" / "default.nemo")) + assert float(model(torch.tensor([1.0, 1.0], device=model.device))) == 0.0 diff --git a/models/tts/NeMo/tests/core/test_fileio.py b/models/tts/NeMo/tests/core/test_fileio.py new file mode 100644 index 0000000000..f737b23216 --- /dev/null +++ b/models/tts/NeMo/tests/core/test_fileio.py @@ -0,0 +1,299 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 os +import tempfile + +import numpy as np +import pytest +import torch +from omegaconf import DictConfig, OmegaConf + +from nemo.collections.asr.models import EncDecCTCModel + +try: + from eff.cookbooks import NeMoCookbook + + _EFF_PRESENT_ = True +except ImportError: + _EFF_PRESENT_ = False + +# A decorator marking the EFF requirement. +requires_eff = pytest.mark.skipif(not _EFF_PRESENT_, reason="Export File Format library required to run test") + + +@pytest.fixture() +def asr_model(): + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 1024, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.ConvASRDecoder', + 'params': { + 'feat_in': 1024, + 'num_classes': 28, + 'vocabulary': [ + ' ', + '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', + "'", + ], + }, + } + modelConfig = DictConfig( + {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)} + ) + + model_instance = EncDecCTCModel(cfg=modelConfig) + return model_instance + + +class TestFileIO: + @pytest.mark.unit + def test_to_from_config_file(self, asr_model): + """" Test makes sure that the second instance created with the same configuration (BUT NOT checkpoint) + has different weights. """ + + with tempfile.NamedTemporaryFile() as fp: + yaml_filename = fp.name + asr_model.to_config_file(path2yaml_file=yaml_filename) + next_instance = EncDecCTCModel.from_config_file(path2yaml_file=yaml_filename) + + assert isinstance(next_instance, EncDecCTCModel) + + assert len(next_instance.decoder.vocabulary) == 28 + assert asr_model.num_weights == next_instance.num_weights + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = next_instance.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert not np.array_equal(w1, w2) + + @pytest.mark.unit + def test_save_restore_from_nemo_file(self, asr_model): + """" Test makes sure that the second instance created from the same configuration AND checkpoint + has the same weights. """ + + with tempfile.NamedTemporaryFile() as fp: + filename = fp.name + + # Save model (with random artifact). + with tempfile.NamedTemporaryFile() as artifact: + asr_model.register_artifact(config_path="abc", src=artifact.name) + asr_model.save_to(save_path=filename) + + # Restore the model. + asr_model2 = EncDecCTCModel.restore_from(restore_path=filename) + + assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary) + assert asr_model.num_weights == asr_model2.num_weights + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert np.array_equal(w1, w2) + + @requires_eff + @pytest.mark.unit + def test_eff_save_restore_from_nemo_file_encrypted(self, asr_model): + """" Test makes sure that after encrypted save-restore the model has the same weights. """ + + with tempfile.NamedTemporaryFile() as fp: + filename = fp.name + + # Set key - use checkpoint encryption. + NeMoCookbook.set_encryption_key("test_key") + + # Save model (with random artifact). + with tempfile.NamedTemporaryFile() as artifact: + asr_model.register_artifact(config_path="abc", src=artifact.name) + asr_model.save_to(save_path=filename) + + # Try to restore the encrypted archive (weights) without the encryption key. + NeMoCookbook.set_encryption_key(None) + with pytest.raises(PermissionError): + # Restore the model. + asr_model2 = EncDecCTCModel.restore_from(restore_path=filename) + + # Restore the model. + NeMoCookbook.set_encryption_key("test_key") + asr_model3 = EncDecCTCModel.restore_from(restore_path=filename) + # Reset encryption so it won't mess up with other save/restore. + NeMoCookbook.set_encryption_key(None) + + assert asr_model.num_weights == asr_model3.num_weights + + @pytest.mark.unit + def test_save_restore_from_nemo_file_with_override(self, asr_model, tmpdir): + """" Test makes sure that the second instance created from the same configuration AND checkpoint + has the same weights. + + Args: + tmpdir: fixture providing a temporary directory unique to the test invocation. + """ + # Name of the archive in tmp folder. + filename = os.path.join(tmpdir, "eff.nemo") + + # Get path where the command is executed - the artifacts will be "retrieved" there. + # (original .nemo behavior) + cwd = os.getcwd() + + with tempfile.NamedTemporaryFile(mode='a+') as conf_fp: + + # Create a "random artifact". + with tempfile.NamedTemporaryFile(mode="w", delete=False) as artifact: + artifact.write("magic content 42") + # Remember the filename of the artifact. + _, artifact_filename = os.path.split(artifact.name) + # Add artifact to model. + asr_model.register_artifact(config_path="abc", src=artifact.name) + # Save model (with "random artifact"). + asr_model.save_to(save_path=filename) + + # Modify config slightly + cfg = asr_model.cfg + cfg.encoder.activation = 'swish' + yaml_cfg = OmegaConf.to_yaml(cfg) + conf_fp.write(yaml_cfg) + conf_fp.seek(0) + + # Restore the model. + asr_model2 = EncDecCTCModel.restore_from(restore_path=filename, override_config_path=conf_fp.name) + + assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary) + assert asr_model.num_weights == asr_model2.num_weights + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert np.array_equal(w1, w2) + + assert asr_model2.cfg.encoder.activation == 'swish' + + @pytest.mark.unit + def test_save_model_level_pt_ckpt(self, asr_model): + with tempfile.TemporaryDirectory() as ckpt_dir: + nemo_file = os.path.join(ckpt_dir, 'asr.nemo') + asr_model.save_to(nemo_file) + + # Save model level PT checkpoint + asr_model.extract_state_dict_from(nemo_file, ckpt_dir) + ckpt_path = os.path.join(ckpt_dir, 'model_weights.ckpt') + + assert os.path.exists(ckpt_path) + + # Restore the model. + asr_model2 = EncDecCTCModel.restore_from(restore_path=nemo_file) + + assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary) + assert asr_model.num_weights == asr_model2.num_weights + + # Change weights values + asr_model2.encoder.encoder[0].mconv[0].conv.weight.data += 1.0 + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert not np.array_equal(w1, w2) + + # Restore from checkpoint + asr_model2.load_state_dict(torch.load(ckpt_path)) + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert np.array_equal(w1, w2) + + @pytest.mark.unit + def test_save_module_level_pt_ckpt(self, asr_model): + with tempfile.TemporaryDirectory() as ckpt_dir: + nemo_file = os.path.join(ckpt_dir, 'asr.nemo') + asr_model.save_to(nemo_file) + + # Save model level PT checkpoint + asr_model.extract_state_dict_from(nemo_file, ckpt_dir, split_by_module=True) + encoder_path = os.path.join(ckpt_dir, 'encoder.ckpt') + decoder_path = os.path.join(ckpt_dir, 'decoder.ckpt') + preprocessor_path = os.path.join(ckpt_dir, 'preprocessor.ckpt') + + assert os.path.exists(encoder_path) + assert os.path.exists(decoder_path) + assert os.path.exists(preprocessor_path) + + # Restore the model. + asr_model2 = EncDecCTCModel.restore_from(restore_path=nemo_file) + + assert len(asr_model.decoder.vocabulary) == len(asr_model2.decoder.vocabulary) + assert asr_model.num_weights == asr_model2.num_weights + + # Change weights values + asr_model2.encoder.encoder[0].mconv[0].conv.weight.data += 1.0 + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert not np.array_equal(w1, w2) + + # Restore from checkpoint + asr_model2.encoder.load_state_dict(torch.load(encoder_path)) + + w1 = asr_model.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + w2 = asr_model2.encoder.encoder[0].mconv[0].conv.weight.data.detach().cpu().numpy() + + assert np.array_equal(w1, w2) diff --git a/models/tts/NeMo/tests/core/test_neural_types.py b/models/tts/NeMo/tests/core/test_neural_types.py new file mode 100644 index 0000000000..e40ddd93b9 --- /dev/null +++ b/models/tts/NeMo/tests/core/test_neural_types.py @@ -0,0 +1,224 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest + +from nemo.core.neural_types import ( + AcousticEncodedRepresentation, + AudioSignal, + AxisKind, + AxisKindAbstract, + AxisType, + ChannelType, + ElementType, + MelSpectrogramType, + MFCCSpectrogramType, + NeuralType, + NeuralTypeComparisonResult, + SpectrogramType, + VoidType, +) + + +class TestNeuralTypeSystem: + @pytest.mark.unit + def test_short_vs_long_version(self): + long_version = NeuralType( + axes=(AxisType(AxisKind.Batch, None), AxisType(AxisKind.Dimension, None), AxisType(AxisKind.Time, None)), + elements_type=AcousticEncodedRepresentation(), + ) + short_version = NeuralType(('B', 'D', 'T'), AcousticEncodedRepresentation()) + assert long_version.compare(short_version) == NeuralTypeComparisonResult.SAME + assert short_version.compare(long_version) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_parameterized_type_audio_sampling_frequency(self): + audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000)) + audio8K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(8000)) + another16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000)) + + assert audio8K.compare(audio16K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS + assert audio16K.compare(audio8K) == NeuralTypeComparisonResult.SAME_TYPE_INCOMPATIBLE_PARAMS + assert another16K.compare(audio16K) == NeuralTypeComparisonResult.SAME + assert audio16K.compare(another16K) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_transpose_same_1(self): + type1 = NeuralType(axes=('B', 'T', 'C')) + type2 = NeuralType(axes=('T', 'B', 'C')) + + assert type1.compare(type2) == NeuralTypeComparisonResult.TRANSPOSE_SAME + assert type2.compare(type1) == NeuralTypeComparisonResult.TRANSPOSE_SAME + + @pytest.mark.unit + def test_transpose_same_2(self): + audio16K = NeuralType(axes=('B', 'T'), elements_type=AudioSignal(16000)) + audio16K_t = NeuralType(axes=('T', 'B'), elements_type=AudioSignal(16000)) + + assert audio16K.compare(audio16K_t) == NeuralTypeComparisonResult.TRANSPOSE_SAME + + @pytest.mark.unit + def test_inheritance_spec_augment_example(self): + input = NeuralType(('B', 'D', 'T'), SpectrogramType()) + out1 = NeuralType(('B', 'D', 'T'), MelSpectrogramType()) + out2 = NeuralType(('B', 'D', 'T'), MFCCSpectrogramType()) + + assert out1.compare(out2) == NeuralTypeComparisonResult.INCOMPATIBLE + assert out2.compare(out1) == NeuralTypeComparisonResult.INCOMPATIBLE + assert input.compare(out1) == NeuralTypeComparisonResult.GREATER + assert input.compare(out2) == NeuralTypeComparisonResult.GREATER + assert out1.compare(input) == NeuralTypeComparisonResult.LESS + assert out2.compare(input) == NeuralTypeComparisonResult.LESS + + @pytest.mark.unit + def test_singletone(self): + loss_output1 = NeuralType(axes=None) + loss_output2 = NeuralType(axes=None) + + assert loss_output1.compare(loss_output2) == NeuralTypeComparisonResult.SAME + assert loss_output2.compare(loss_output1) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_list_of_lists(self): + T1 = NeuralType( + axes=( + AxisType(kind=AxisKind.Batch, size=None, is_list=True), + AxisType(kind=AxisKind.Time, size=None, is_list=True), + AxisType(kind=AxisKind.Dimension, size=32, is_list=False), + AxisType(kind=AxisKind.Dimension, size=128, is_list=False), + AxisType(kind=AxisKind.Dimension, size=256, is_list=False), + ), + elements_type=ChannelType(), + ) + T2 = NeuralType( + axes=( + AxisType(kind=AxisKind.Batch, size=None, is_list=False), + AxisType(kind=AxisKind.Time, size=None, is_list=False), + AxisType(kind=AxisKind.Dimension, size=32, is_list=False), + AxisType(kind=AxisKind.Dimension, size=128, is_list=False), + AxisType(kind=AxisKind.Dimension, size=256, is_list=False), + ), + elements_type=ChannelType(), + ) + # TODO: should this be incompatible instead??? + assert T1.compare(T2), NeuralTypeComparisonResult.TRANSPOSE_SAME + + @pytest.mark.unit + def test_void(self): + btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType()) + btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType()) + btc_void = NeuralType(('B', 'T', 'C'), VoidType()) + + assert btc_void.compare(btc_spctr) == NeuralTypeComparisonResult.SAME + assert btc_spctr.compare(btc_void) == NeuralTypeComparisonResult.INCOMPATIBLE + assert btc_void.compare(btc_spct_bad) == NeuralTypeComparisonResult.INCOMPATIBLE + + @pytest.mark.unit + def test_big_void(self): + big_void_1 = NeuralType(elements_type=VoidType()) + big_void_2 = NeuralType() + + btc_spctr = NeuralType(('B', 'T', 'C'), SpectrogramType()) + btc_spct_bad = NeuralType(('B', 'T'), SpectrogramType()) + t1 = NeuralType( + axes=( + AxisType(kind=AxisKind.Batch, size=None, is_list=True), + AxisType(kind=AxisKind.Time, size=None, is_list=True), + AxisType(kind=AxisKind.Dimension, size=32, is_list=False), + AxisType(kind=AxisKind.Dimension, size=128, is_list=False), + AxisType(kind=AxisKind.Dimension, size=256, is_list=False), + ), + elements_type=ChannelType(), + ) + t2 = NeuralType( + axes=( + AxisType(kind=AxisKind.Batch, size=None, is_list=False), + AxisType(kind=AxisKind.Time, size=None, is_list=False), + AxisType(kind=AxisKind.Dimension, size=32, is_list=False), + AxisType(kind=AxisKind.Dimension, size=128, is_list=False), + AxisType(kind=AxisKind.Dimension, size=256, is_list=False), + ), + elements_type=ChannelType(), + ) + + assert big_void_1.compare(btc_spctr) == NeuralTypeComparisonResult.SAME + assert big_void_1.compare(btc_spct_bad) == NeuralTypeComparisonResult.SAME + assert big_void_1.compare(t1) == NeuralTypeComparisonResult.SAME + assert big_void_1.compare(t2) == NeuralTypeComparisonResult.SAME + + assert big_void_2.compare(btc_spctr) == NeuralTypeComparisonResult.SAME + assert big_void_2.compare(btc_spct_bad) == NeuralTypeComparisonResult.SAME + assert big_void_2.compare(t1) == NeuralTypeComparisonResult.SAME + assert big_void_2.compare(t2) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_unspecified_dimensions(self): + t0 = NeuralType( + (AxisType(AxisKind.Batch, 64), AxisType(AxisKind.Time, 10), AxisType(AxisKind.Dimension, 128)), + SpectrogramType(), + ) + t1 = NeuralType(('B', 'T', 'C'), SpectrogramType()) + + assert t1.compare(t0), NeuralTypeComparisonResult.SAME + assert t0.compare(t1), NeuralTypeComparisonResult.DIM_INCOMPATIBLE + + @pytest.mark.unit + def test_any_axis(self): + t0 = NeuralType(('B', 'Any', 'Any'), VoidType()) + t1 = NeuralType(('B', 'Any', 'Any'), SpectrogramType()) + t2 = NeuralType(('B', 'T', 'C'), SpectrogramType()) + + assert t0.compare(t1) == NeuralTypeComparisonResult.SAME + assert t0.compare(t2) == NeuralTypeComparisonResult.SAME + assert t1.compare(t2) == NeuralTypeComparisonResult.SAME + assert t2.compare(t1) == NeuralTypeComparisonResult.INCOMPATIBLE + assert t1.compare(t0) == NeuralTypeComparisonResult.INCOMPATIBLE + + @pytest.mark.unit + def test_struct(self): + class BoundingBox(ElementType): + def __str__(self): + return "bounding box from detection model" + + def fields(self): + return ("X", "Y", "W", "H") + + # ALSO ADD new, user-defined, axis kind + class AxisKind2(AxisKindAbstract): + Image = 0 + + T1 = NeuralType( + elements_type=BoundingBox(), + axes=( + AxisType(kind=AxisKind.Batch, size=None, is_list=True), + AxisType(kind=AxisKind2.Image, size=None, is_list=True), + ), + ) + + class BadBoundingBox(ElementType): + def __str__(self): + return "bad bounding box from detection model" + + def fields(self): + return ("X", "Y", "H") + + T2 = NeuralType( + elements_type=BadBoundingBox(), + axes=( + AxisType(kind=AxisKind.Batch, size=None, is_list=True), + AxisType(kind=AxisKind2.Image, size=None, is_list=True), + ), + ) + assert T2.compare(T1) == NeuralTypeComparisonResult.INCOMPATIBLE diff --git a/models/tts/NeMo/tests/core/test_optimizers_schedulers.py b/models/tts/NeMo/tests/core/test_optimizers_schedulers.py new file mode 100644 index 0000000000..84f7d14a7c --- /dev/null +++ b/models/tts/NeMo/tests/core/test_optimizers_schedulers.py @@ -0,0 +1,809 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 random + +import omegaconf +import pytest +import pytorch_lightning as pl +import torch +import torch.optim + +from nemo.core import config, optim +from nemo.core.optim.lr_scheduler import AVAILABLE_SCHEDULERS +from nemo.core.optim.optimizers import AVAILABLE_OPTIMIZERS +from nemo.utils import logging + + +class TempModel(torch.nn.Module): + def __init__(self): + super(TempModel, self).__init__() + self.layer = torch.nn.Linear(5, 1) + + def forward(self, x): + x = self.layer(x) + return x + + +class OptCounter(torch.optim.SGD): + def __init__(self, *args, **kwargs): + self.count = 0 + super().__init__(*args, **kwargs) + + def step(self, closure=None): + try: + self.count += 1 + except AttributeError: + self.count = 1 + super().step(closure) + + +class RandomDataset(torch.utils.data.Dataset): + def __init__(self, dataset_len): + super().__init__() + self.__dataset_len = dataset_len + + def __getitem__(self, *args): + return torch.randn(2) + + def __len__(self): + return self.__dataset_len + + +class ExampleModel(pl.LightningModule): + def __init__(self, batch_size, dataset_len, drop_last, max_steps): + super().__init__() + self.l1 = torch.nn.modules.Linear(in_features=2, out_features=1) + self.batch_size = batch_size + self.dataset_len = dataset_len + self.drop_last = drop_last + self.max_steps = max_steps + + def train_dataloader(self): + dataset = RandomDataset(self.dataset_len) + return torch.utils.data.DataLoader(dataset, batch_size=self.batch_size, drop_last=self.drop_last) + + def training_step(self, batch, batch_idx): + output = self.l1(batch) + output = torch.nn.functional.l1_loss(output, torch.ones(output.size()).to(output.device)) + return {"loss": output} + + def configure_optimizers(self): + self.my_opt = OptCounter(self.parameters(), lr=0.02) + return self.my_opt + + +class Callback(pl.callbacks.Callback): + @pl.utilities.distributed.rank_zero_only + def on_train_end(self, trainer, module): + if trainer.global_step != module.my_opt.count or trainer.global_step != module.max_steps: + logging.debug(f"max_epochs: {trainer.max_epochs}") + logging.debug(f"accumulate_grad_batches: {trainer.accumulate_grad_batches}") + logging.debug(f"limit_train_batches: {trainer.limit_train_batches}") + logging.debug(f"num_processes: {trainer.num_processes}") + logging.debug(f"batch_size: {module.batch_size}") + logging.debug(f"dataset_len: {module.dataset_len}") + logging.debug(f"drop_last: {module.drop_last}") + logging.debug(f"{len(trainer.train_dataloader)}") + logging.debug(f"{trainer.num_training_batches }") + assert ( + trainer.global_step == module.my_opt.count + ), f"{trainer.global_step} != {module.my_opt.count} != {module.max_steps}" + assert ( + trainer.global_step == module.max_steps + ), f"{trainer.global_step} != {module.my_opt.count} != {module.max_steps}" + + +class TestOptimizersSchedulers: + INITIAL_LR = 0.1 + MIN_LR = 1e-3 + MAX_STEPS = 10 + + @pytest.mark.unit + def test_get_optimizer(self): + model = TempModel() + + for opt_name in AVAILABLE_OPTIMIZERS.keys(): + opt_cls = optim.get_optimizer(opt_name) + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + assert isinstance(opt, AVAILABLE_OPTIMIZERS[opt_name]) + + @pytest.mark.unit + def test_register_optimizer(self): + class TempOpt(torch.optim.SGD): + pass + + class TempOptParams(config.optimizers.SGDParams): + pass + + optim.register_optimizer('TempOpt', TempOpt, TempOptParams) + + model = TempModel() + opt_cls = optim.get_optimizer('TempOpt') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + assert isinstance(opt, TempOpt) + + @pytest.mark.unit + def test_optim_config_parse_bypass(self): + basic_optim_config = {'weight_decay': 0.001, 'betas': [0.8, 0.5]} + parsed_params = optim.parse_optimizer_args('novograd', basic_optim_config) + assert parsed_params['weight_decay'] == basic_optim_config['weight_decay'] + assert parsed_params['betas'][0] == basic_optim_config['betas'][0] + assert parsed_params['betas'][1] == basic_optim_config['betas'][1] + + dict_config = omegaconf.OmegaConf.create(basic_optim_config) + parsed_params = optim.parse_optimizer_args('novograd', dict_config) + assert parsed_params['weight_decay'] == dict_config['weight_decay'] + assert parsed_params['betas'][0] == dict_config['betas'][0] + assert parsed_params['betas'][1] == dict_config['betas'][1] + + @pytest.mark.unit + def test_optim_config_parse_arg_by_name(self): + basic_optim_config = {'name': 'auto', 'weight_decay': 0.001, 'betas': [0.8, 0.5]} + parsed_params = optim.parse_optimizer_args('novograd', basic_optim_config) + assert parsed_params['weight_decay'] == basic_optim_config['weight_decay'] + assert parsed_params['betas'][0] == basic_optim_config['betas'][0] + assert parsed_params['betas'][1] == basic_optim_config['betas'][1] + + dict_config = omegaconf.OmegaConf.create(basic_optim_config) + parsed_params = optim.parse_optimizer_args('novograd', dict_config) + assert parsed_params['weight_decay'] == dict_config['weight_decay'] + assert parsed_params['betas'][0] == dict_config['betas'][0] + assert parsed_params['betas'][1] == dict_config['betas'][1] + + with pytest.raises(omegaconf.errors.ConfigKeyError): + optim.parse_optimizer_args('sgd', dict_config) + + @pytest.mark.unit + def test_optim_config_parse_arg_by_target(self): + basic_optim_config = { + '_target_': 'nemo.core.config.NovogradParams', + 'params': {'weight_decay': 0.001, 'betas': [0.8, 0.5]}, + } + basic_optim_config = omegaconf.OmegaConf.create(basic_optim_config) + parsed_params = optim.parse_optimizer_args('novograd', basic_optim_config) + assert parsed_params['weight_decay'] == basic_optim_config['params']['weight_decay'] + assert parsed_params['betas'][0] == basic_optim_config['params']['betas'][0] + assert parsed_params['betas'][1] == basic_optim_config['params']['betas'][1] + + dict_config = omegaconf.OmegaConf.create(basic_optim_config) + parsed_params = optim.parse_optimizer_args('novograd', dict_config) + assert parsed_params['weight_decay'] == dict_config['params']['weight_decay'] + assert parsed_params['betas'][0] == dict_config['params']['betas'][0] + assert parsed_params['betas'][1] == dict_config['params']['betas'][1] + + # Names are ignored when passing class path + # This will be captured during optimizer instantiation + output_config = optim.parse_optimizer_args('sgd', dict_config) + sgd_config = vars(config.SGDParams()) + novograd_config = vars(config.NovogradParams()) + + assert set(output_config.keys()) != set(sgd_config.keys()) + assert set(output_config.keys()) == set(novograd_config) + + @pytest.mark.unit + def test_get_scheduler(self): + model = TempModel() + optimizer = optim.Novograd(model.parameters(), lr=self.INITIAL_LR) + + for sched_name in AVAILABLE_SCHEDULERS.keys(): + sched_cls = optim.lr_scheduler.get_scheduler(sched_name) + + try: + sched = sched_cls(optimizer) + assert isinstance(sched, AVAILABLE_SCHEDULERS[sched_name]) + continue + except Exception: + pass + + try: + sched = sched_cls(optimizer, max_steps=self.MAX_STEPS) + assert isinstance(sched, AVAILABLE_SCHEDULERS[sched_name]) + continue + except Exception: + pass + + @pytest.mark.unit + def test_register_scheduler(self): + class TempSched(optim.lr_scheduler.CosineAnnealing): + pass + + class TempSchedParams(config.schedulers.CosineAnnealingParams): + pass + + optim.lr_scheduler.register_scheduler('TempSched', TempSched, TempSchedParams) + + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + sched_cls = optim.lr_scheduler.get_scheduler('TempSched') + sched = sched_cls(opt, max_steps=self.MAX_STEPS) + + assert isinstance(sched, TempSched) + + @pytest.mark.unit + def test_sched_config_parse_simple(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + basic_sched_config = {'name': 'CosineAnnealing', 'max_steps': 10} + scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, basic_sched_config) + assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing) + + dict_config = omegaconf.OmegaConf.create(basic_sched_config) + scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, dict_config) + assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing) + + @pytest.mark.unit + def test_sched_config_parse_from_cls(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + basic_sched_config = { + '_target_': 'nemo.core.config.CosineAnnealingParams', + 'params': {'min_lr': 0.1}, + 'max_steps': self.MAX_STEPS, + } + scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, basic_sched_config) + assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing) + + dict_config = omegaconf.OmegaConf.create(basic_sched_config) + scheduler_setup = optim.lr_scheduler.prepare_lr_scheduler(opt, dict_config) + assert isinstance(scheduler_setup['scheduler'], optim.lr_scheduler.CosineAnnealing) + + @pytest.mark.unit + def test_WarmupPolicy(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.WarmupPolicy(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] == self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.WarmupPolicy(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 4: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] == self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_WarmupHoldPolicy(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.WarmupHoldPolicy(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] == self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.WarmupHoldPolicy(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 4: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] == self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup + Hold steps available + policy = optim.lr_scheduler.WarmupHoldPolicy( + opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 4: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] == self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_WarmupAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.WarmupAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.WarmupAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup + Hold steps available + policy = optim.lr_scheduler.WarmupHoldPolicy( + opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 4: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] == self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_SquareAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.SquareAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.SquareAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_SquareRootAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.SquareRootAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.SquareRootAnnealing( + opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_CosineAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.CosineAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.CosineAnnealing(opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_PolynomialDecayAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.PolynomialDecayAnnealing( + opt, power=2, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.PolynomialDecayAnnealing( + opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_PolynomialHoldDecayAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing( + opt, power=2, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing( + opt, power=2, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup + Hold steps available + policy = optim.lr_scheduler.PolynomialHoldDecayAnnealing( + opt, warmup_steps=5, hold_steps=3, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR, power=2 + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 4: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + elif i <= 8: + assert policy.get_last_lr()[0] == self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + def test_InverseSquareRootAnnealing(self): + model = TempModel() + opt_cls = optim.get_optimizer('novograd') + opt = opt_cls(model.parameters(), lr=self.INITIAL_LR) + + # No warmup case + policy = optim.lr_scheduler.InverseSquareRootAnnealing(opt, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr == self.INITIAL_LR + + for i in range(self.MAX_STEPS): + assert policy.get_last_lr()[0] <= self.INITIAL_LR + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + # Warmup steps available + policy = optim.lr_scheduler.InverseSquareRootAnnealing( + opt, warmup_steps=5, max_steps=self.MAX_STEPS, min_lr=self.MIN_LR + ) + initial_lr = policy.get_last_lr()[0] + + assert initial_lr < self.INITIAL_LR + + for i in range(self.MAX_STEPS): + if i <= 5: + assert policy.get_last_lr()[0] <= self.INITIAL_LR + else: + assert policy.get_last_lr()[0] < self.INITIAL_LR + + opt.step() + policy.step() + + policy.step() + final_lr = policy.get_last_lr()[0] + + assert final_lr == self.MIN_LR + + @pytest.mark.unit + @pytest.mark.run_only_on('CPU') + def test_max_step_computation(self): + def train( + max_epochs, accumulate_grad_batches, limit_train_batches, num_processes, batch_size, dataset_len, drop_last + ): + trainer = pl.Trainer( + max_epochs=max_epochs, + accelerator="ddp_cpu", + num_processes=num_processes, + accumulate_grad_batches=accumulate_grad_batches, + limit_train_batches=limit_train_batches, + checkpoint_callback=False, + progress_bar_refresh_rate=0, + weights_summary=None, + ) + max_steps = optim.lr_scheduler.compute_max_steps( + max_epochs, + accumulate_grad_batches, + limit_train_batches, + num_processes, + dataset_len, + batch_size, + drop_last, + ) + model = ExampleModel(batch_size, dataset_len, drop_last, max_steps) + trainer.callbacks.append(Callback()) + trainer.fit(model) + + # This test will break once we and lightning upgrade to pytorch 1.7.0 due to a bug fix in pytorch 1.7.0 + train( + 31, + accumulate_grad_batches=1, + limit_train_batches=1.0, + num_processes=9, + batch_size=60, + dataset_len=1613, + drop_last=True, + ) + train( + 5, + accumulate_grad_batches=1, + limit_train_batches=0.17382691901706027, + num_processes=4, + batch_size=97, + dataset_len=498, + drop_last=False, + ) + train( + 5, + accumulate_grad_batches=8, + limit_train_batches=0.1663306588594945, + num_processes=4, + batch_size=54, + dataset_len=629, + drop_last=True, + ) + train( + 5, + accumulate_grad_batches=1, + limit_train_batches=0.2121376533631948, + num_processes=1, + batch_size=68, + dataset_len=488, + drop_last=False, + ) + for _ in range(5): + drop_last = bool(random.randint(0, 1)) + accumulate_grad_batches = random.randint(1, 10) + + limit_train_batches_int = random.randint(1, 10) + limit_train_batches_float = random.uniform(0, 1) + limit_train_batches = random.choice([limit_train_batches_int, limit_train_batches_float]) + max_epochs = random.randint(4, 20) + num_processes = random.randint(1, 5) + dataset_len = random.randint(20, num_processes * 500) + batch_size = random.randint(math.ceil(5.0 / num_processes), min(dataset_len // num_processes, 128)) + train( + max_epochs, + accumulate_grad_batches, + limit_train_batches, + num_processes, + batch_size, + dataset_len, + drop_last, + ) diff --git a/models/tts/NeMo/tests/core/test_save_restore.py b/models/tts/NeMo/tests/core/test_save_restore.py new file mode 100644 index 0000000000..4c59fa0cdb --- /dev/null +++ b/models/tts/NeMo/tests/core/test_save_restore.py @@ -0,0 +1,443 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 filecmp +import os +import shutil +import tempfile +from typing import Dict, Optional, Set, Union + +import pytest +import torch +from omegaconf import DictConfig, OmegaConf, open_dict + +from nemo.collections.asr.models import EncDecCTCModel, EncDecCTCModelBPE +from nemo.collections.nlp.models import PunctuationCapitalizationModel, TransformerLMModel +from nemo.core.classes import ModelPT +from nemo.utils.app_state import AppState + + +def getattr2(object, attr): + if not '.' in attr: + return getattr(object, attr) + else: + arr = attr.split('.') + return getattr2(getattr(object, arr[0]), '.'.join(arr[1:])) + + +class MockModel(ModelPT): + def __init__(self, cfg, trainer=None): + super(MockModel, self).__init__(cfg=cfg, trainer=trainer) + self.w = torch.nn.Linear(10, 1) + # mock temp file + if 'temp_file' in self.cfg and self.cfg.temp_file is not None: + self.temp_file = self.register_artifact('temp_file', self.cfg.temp_file) + with open(self.temp_file, 'r') as f: + self.temp_data = f.readlines() + else: + self.temp_file = None + self.temp_data = None + + def forward(self, x): + y = self.w(x) + return y, self.cfg.temp_file + + def setup_training_data(self, train_data_config: Union[DictConfig, Dict]): + self._train_dl = None + + def setup_validation_data(self, val_data_config: Union[DictConfig, Dict]): + self._validation_dl = None + + def setup_test_data(self, test_data_config: Union[DictConfig, Dict]): + self._test_dl = None + + def list_available_models(cls): + return [] + + +def _mock_model_config(): + conf = {'temp_file': None} + conf = OmegaConf.create({'model': conf}) + OmegaConf.set_struct(conf, True) + return conf + + +class TestSaveRestore: + def __test_restore_elsewhere( + self, + model: ModelPT, + attr_for_eq_check: Set[str] = None, + override_config_path: Optional[Union[str, DictConfig]] = None, + map_location: Optional[torch.device] = None, + strict: bool = False, + return_config: bool = False, + ): + """Test's logic: + 1. Save model into temporary folder (save_folder) + 2. Copy .nemo file from save_folder to restore_folder + 3. Delete save_folder + 4. Attempt to restore from .nemo file in restore_folder and compare to original instance + """ + # Create a new temporary directory + with tempfile.TemporaryDirectory() as restore_folder: + with tempfile.TemporaryDirectory() as save_folder: + save_folder_path = save_folder + # Where model will be saved + model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo") + model.save_to(save_path=model_save_path) + # Where model will be restored from + model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo") + shutil.copy(model_save_path, model_restore_path) + # at this point save_folder should not exist + assert save_folder_path is not None and not os.path.exists(save_folder_path) + assert not os.path.exists(model_save_path) + assert os.path.exists(model_restore_path) + # attempt to restore + model_copy = model.__class__.restore_from( + restore_path=model_restore_path, + map_location=map_location, + strict=strict, + return_config=return_config, + override_config_path=override_config_path, + ) + + if return_config: + return model_copy + + assert model.num_weights == model_copy.num_weights + if attr_for_eq_check is not None and len(attr_for_eq_check) > 0: + for attr in attr_for_eq_check: + assert getattr2(model, attr) == getattr2(model_copy, attr) + + return model_copy + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_EncDecCTCModel(self): + # TODO: Switch to using named configs because here we don't really care about weights + qn = EncDecCTCModel.from_pretrained(model_name="QuartzNet15x5Base-En") + self.__test_restore_elsewhere(model=qn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"])) + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_EncDecCTCModelBPE(self): + # TODO: Switch to using named configs because here we don't really care about weights + cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_citrinet_256") + self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"])) + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_EncDecCTCModelBPE_v2(self): + # TODO: Switch to using named configs because here we don't really care about weights + cn = EncDecCTCModelBPE.from_pretrained(model_name="stt_en_conformer_ctc_small") + self.__test_restore_elsewhere(model=cn, attr_for_eq_check=set(["decoder._feat_in", "decoder._num_classes"])) + + @pytest.mark.with_downloads() + @pytest.mark.unit + def test_PunctuationCapitalization(self): + # TODO: Switch to using named configs because here we don't really care about weights + pn = PunctuationCapitalizationModel.from_pretrained(model_name='punctuation_en_distilbert') + self.__test_restore_elsewhere( + model=pn, attr_for_eq_check=set(["punct_classifier.log_softmax", "punct_classifier.log_softmax"]) + ) + + @pytest.mark.unit + def test_mock_save_to_restore_from(self): + with tempfile.NamedTemporaryFile('w') as empty_file: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + + # Create model + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + + assert model.temp_file == empty_file.name + + # Save test + model_copy = self.__test_restore_elsewhere(model, map_location='cpu') + + # Restore test + diff = model.w.weight - model_copy.w.weight + # because of caching - cache gets prepended + assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file)) + assert diff.mean() <= 1e-9 + # assert os.path.basename(model.temp_file) == model_copy.temp_file + assert model_copy.temp_data == ["*****\n"] + + @pytest.mark.unit + def test_mock_restore_from_config_only(self): + with tempfile.NamedTemporaryFile('w') as empty_file: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + cfg = _mock_model_config() + cfg.model.temp_file = os.path.abspath(empty_file.name) + + # Inject arbitrary config arguments (after creating model) + with open_dict(cfg.model): + cfg.model.xyz = "abc" + + # Create model + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + + assert model.temp_file == empty_file.name + model_copy = self.__test_restore_elsewhere(model, map_location='cpu', return_config=False) + # because of caching - cache gets prepended + assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file)) + # assert filecmp.cmp(model.temp_file, model_copy._cfg.temp_file) + assert model.cfg.xyz == model_copy.cfg.xyz + + @pytest.mark.unit + def test_mock_restore_from_config_override_with_OmegaConf(self): + with tempfile.NamedTemporaryFile('w') as empty_file: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + + # Create model + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + + assert model.temp_file == empty_file.name + + # Inject arbitrary config arguments (after creating model) + with open_dict(cfg.model): + cfg.model.xyz = "abc" + + # Save test (with overriden config as OmegaConf object) + model_copy = self.__test_restore_elsewhere(model, map_location='cpu', override_config_path=cfg) + + # Restore test + diff = model.w.weight - model_copy.w.weight + assert diff.mean() <= 1e-9 + assert model_copy.temp_data == ["*****\n"] + + # Test that new config has arbitrary content + assert model_copy.cfg.xyz == "abc" + + @pytest.mark.unit + def test_mock_restore_from_config_override_with_yaml(self): + with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as config_file: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + + # Create model + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + + assert model.temp_file == empty_file.name + + # Inject arbitrary config arguments (after creating model) + with open_dict(cfg.model): + cfg.model.xyz = "abc" + + # Write new config into file + OmegaConf.save(cfg, config_file) + + # Save test (with overriden config as OmegaConf object) + model_copy = self.__test_restore_elsewhere( + model, map_location='cpu', override_config_path=config_file.name + ) + + # Restore test + diff = model.w.weight - model_copy.w.weight + assert diff.mean() <= 1e-9 + assert filecmp.cmp(model.temp_file, model_copy.temp_file) + assert model_copy.temp_data == ["*****\n"] + + # Test that new config has arbitrary content + assert model_copy.cfg.xyz == "abc" + + @pytest.mark.unit + def test_mock_save_to_restore_from_with_target_class(self): + with tempfile.NamedTemporaryFile('w') as empty_file: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + + # Create model + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') # type: MockModel + + assert model.temp_file == empty_file.name + + # Save file using MockModel + with tempfile.TemporaryDirectory() as save_folder: + save_path = os.path.join(save_folder, "temp.nemo") + model.save_to(save_path) + + # Restore test (using ModelPT as restorer) + # This forces the target class = MockModel to be used as resolver + model_copy = ModelPT.restore_from(save_path, map_location='cpu') + # because of caching - cache gets prepended + assert os.path.basename(model_copy.temp_file).endswith(os.path.basename(model.temp_file)) + # assert filecmp.cmp(model.temp_file, model_copy.temp_file) + # Restore test + diff = model.w.weight - model_copy.w.weight + assert diff.mean() <= 1e-9 + assert isinstance(model_copy, MockModel) + assert model_copy.temp_data == ["*****\n"] + + @pytest.mark.unit + def test_mock_save_to_restore_from_multiple_models(self): + with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + empty_file2.writelines(["+++++\n"]) + empty_file2.flush() + + # Update config + create ,pde;s + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + cfg2 = _mock_model_config() + cfg2.model.temp_file = empty_file2.name + + # Create models + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + model2 = MockModel(cfg=cfg2.model, trainer=None) + model2 = model2.to('cpu') + + assert model.temp_file == empty_file.name + assert model2.temp_file == empty_file2.name + + # Save test + model_copy = self.__test_restore_elsewhere(model, map_location='cpu') + model2_copy = self.__test_restore_elsewhere(model2, map_location='cpu') + + # Restore test + assert model_copy.temp_data == ["*****\n"] + assert model2_copy.temp_data == ["+++++\n"] + + @pytest.mark.unit + def test_mock_save_to_restore_from_multiple_models_inverted_order(self): + with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + empty_file2.writelines(["+++++\n"]) + empty_file2.flush() + + # Update config + create ,pde;s + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + cfg2 = _mock_model_config() + cfg2.model.temp_file = empty_file2.name + + # Create models + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + model2 = MockModel(cfg=cfg2.model, trainer=None) + model2 = model2.to('cpu') + + assert model.temp_file == empty_file.name + assert model2.temp_file == empty_file2.name + + # Save test (inverted order) + model2_copy = self.__test_restore_elsewhere(model2, map_location='cpu') + model_copy = self.__test_restore_elsewhere(model, map_location='cpu') + + # Restore test + assert model_copy.temp_data == ["*****\n"] + assert model2_copy.temp_data == ["+++++\n"] + + @pytest.mark.unit + def test_mock_save_to_restore_chained(self): + with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.NamedTemporaryFile('w') as empty_file2: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + create ,pde;s + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + + # Create models + model = MockModel(cfg=cfg.model, trainer=None) + model = model.to('cpu') + + assert model.temp_file == empty_file.name + + def save_copy(model, save_folder, restore_folder): + # Where model will be saved + model_save_path = os.path.join(save_folder, f"{model.__class__.__name__}.nemo") + model.save_to(save_path=model_save_path) + # Where model will be restored from + model_restore_path = os.path.join(restore_folder, f"{model.__class__.__name__}.nemo") + shutil.copy(model_save_path, model_restore_path) + return model_restore_path + + # Save test + with tempfile.TemporaryDirectory() as level4: + with tempfile.TemporaryDirectory() as level3: + with tempfile.TemporaryDirectory() as level2: + with tempfile.TemporaryDirectory() as level1: + path = save_copy(model, level1, level2) + model_copy2 = model.__class__.restore_from(path) + path = save_copy(model_copy2, level2, level3) + model_copy3 = model.__class__.restore_from(path) + path = save_copy(model_copy3, level3, level4) + model_copy = model.__class__.restore_from(path) + + # Restore test + assert model_copy.temp_data == ["*****\n"] + + # AppState test + appstate = AppState() + metadata = appstate.get_model_metadata_from_guid(model_copy.model_guid) + assert metadata.guid != model.model_guid + assert metadata.restoration_path == path + + @pytest.mark.unit + def test_mock_save_to_multiple_times(self): + with tempfile.NamedTemporaryFile('w') as empty_file, tempfile.TemporaryDirectory() as tmpdir: + # Write some data + empty_file.writelines(["*****\n"]) + empty_file.flush() + + # Update config + cfg = _mock_model_config() + cfg.model.temp_file = empty_file.name + + # Create model + model = MockModel(cfg=cfg.model, trainer=None) # type: MockModel + model = model.to('cpu') + + assert model.temp_file == empty_file.name + + # Save test + model.save_to(os.path.join(tmpdir, 'save_0.nemo')) + model.save_to(os.path.join(tmpdir, 'save_1.nemo')) + model.save_to(os.path.join(tmpdir, 'save_2.nemo')) diff --git a/models/tts/NeMo/tests/core/test_serialization.py b/models/tts/NeMo/tests/core/test_serialization.py new file mode 100644 index 0000000000..42b1e22b71 --- /dev/null +++ b/models/tts/NeMo/tests/core/test_serialization.py @@ -0,0 +1,165 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest +from omegaconf import DictConfig + +from nemo.collections.asr.models import EncDecCTCModel +from nemo.collections.asr.modules import SpectrogramAugmentation +from nemo.core.classes.common import Serialization + + +def get_class_path(cls): + return f"{cls.__module__}.{cls.__name__}" + + +class MockSerializationImpl(Serialization): + def __init__(self, cfg: DictConfig): + self.cfg = cfg + self.value = self.__class__.__name__ + + +class MockSerializationImplV2(MockSerializationImpl): + pass + + +class TestSerialization: + @pytest.mark.unit + def test_from_config_dict_with_cls(self): + """Here we test that instantiation works for configs with cls class path in them. + Note that just Serialization.from_config_dict can be used to create an object""" + config = DictConfig( + { + 'cls': 'nemo.collections.asr.modules.SpectrogramAugmentation', + 'params': {'rect_freq': 50, 'rect_masks': 5, 'rect_time': 120,}, + } + ) + obj = Serialization.from_config_dict(config=config) + assert isinstance(obj, SpectrogramAugmentation) + + @pytest.mark.unit + def test_from_config_dict_without_cls(self): + """Here we test that instantiation works for configs without cls class path in them. + IMPORTANT: in this case, correct class type should call from_config_dict. This should work for Models.""" + preprocessor = {'cls': 'nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor', 'params': dict({})} + encoder = { + 'cls': 'nemo.collections.asr.modules.ConvASREncoder', + 'params': { + 'feat_in': 64, + 'activation': 'relu', + 'conv_mask': True, + 'jasper': [ + { + 'filters': 1024, + 'repeat': 1, + 'kernel': [1], + 'stride': [1], + 'dilation': [1], + 'dropout': 0.0, + 'residual': False, + 'separable': True, + 'se': True, + 'se_context_size': -1, + } + ], + }, + } + + decoder = { + 'cls': 'nemo.collections.asr.modules.ConvASRDecoder', + 'params': { + 'feat_in': 1024, + 'num_classes': 28, + 'vocabulary': [ + ' ', + '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', + "'", + ], + }, + } + modelConfig = DictConfig( + {'preprocessor': DictConfig(preprocessor), 'encoder': DictConfig(encoder), 'decoder': DictConfig(decoder)} + ) + obj = EncDecCTCModel.from_config_dict(config=modelConfig) + assert isinstance(obj, EncDecCTCModel) + + @pytest.mark.unit + def test_config_updated(self): + config = DictConfig( + { + 'cls': 'nemo.collections.asr.modules.SpectrogramAugmentation', + 'params': {'rect_freq': 50, 'rect_masks': 5, 'rect_time': 120,}, + } + ) + obj = Serialization.from_config_dict(config=config) + new_config = obj.to_config_dict() + assert config != new_config + assert 'params' not in new_config + assert 'cls' not in new_config + assert '_target_' in new_config + + @pytest.mark.unit + def test_base_class_instantiation(self): + # Target class is V2 impl, calling class is Serialization (base class) + config = DictConfig({'target': get_class_path(MockSerializationImplV2)}) + obj = Serialization.from_config_dict(config=config) + new_config = obj.to_config_dict() + assert config == new_config + assert isinstance(obj, MockSerializationImplV2) + assert obj.value == "MockSerializationImplV2" + + @pytest.mark.unit + def test_self_class_instantiation(self): + # Target class is V1 impl, calling class is V1 (same class) + config = DictConfig({'target': get_class_path(MockSerializationImpl)}) + obj = MockSerializationImpl.from_config_dict(config=config) # Serialization is base class + new_config = obj.to_config_dict() + assert config == new_config + assert isinstance(obj, MockSerializationImpl) + assert obj.value == "MockSerializationImpl" + + @pytest.mark.unit + def test_sub_class_instantiation(self): + # Target class is V1 impl, calling class is V2 (sub class) + config = DictConfig({'target': get_class_path(MockSerializationImpl)}) + obj = MockSerializationImplV2.from_config_dict(config=config) # Serialization is base class + new_config = obj.to_config_dict() + assert config == new_config + assert isinstance(obj, MockSerializationImplV2) + assert obj.value == "MockSerializationImplV2" diff --git a/models/tts/NeMo/tests/core/test_typecheck.py b/models/tts/NeMo/tests/core/test_typecheck.py new file mode 100644 index 0000000000..c547ab334a --- /dev/null +++ b/models/tts/NeMo/tests/core/test_typecheck.py @@ -0,0 +1,1075 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 pytest +import torch + +from nemo.core import Typing, typecheck +from nemo.core.neural_types import * + + +# Perform recursive shape assert +def recursive_assert_shape(x, shape): + if isinstance(x, list) or isinstance(x, tuple): + for xi in x: + recursive_assert_shape(xi, shape) + return + + assert x.shape == shape + + +# Perform recursive type assert +def recursive_assert_homogeneous_type(x, type_val): + if isinstance(x, list) or isinstance(x, tuple): + for xi in x: + recursive_assert_homogeneous_type(xi, type_val) + return + + assert x.neural_type.compare(type_val) == NeuralTypeComparisonResult.SAME + + +class TestNeuralTypeCheckSystem: + @pytest.mark.unit + def test_no_types_passthrough(self): + class NoTypes(Typing): + @typecheck() + def __call__(self, x): + return torch.tensor(1.0) + + obj = NoTypes() + result = obj(torch.tensor(1.0)) + + assert result == torch.tensor(1.0) + assert not hasattr(result, 'neural_type') + + @pytest.mark.unit + def test_input_output_types(self): + class InputOutputTypes(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + obj = InputOutputTypes() + result = obj(x=torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert result.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + # Test passing wrong key for input + with pytest.raises(TypeError): + _ = obj(a=torch.zeros(10)) + + # Test using positional args + with pytest.raises(TypeError): + _ = obj(torch.zeros(10)) + + @pytest.mark.unit + def test_input_types_only(self): + class InputTypes(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + obj = InputTypes() + result = obj(x=torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert hasattr(result, 'neural_type') is False + + @pytest.mark.unit + def test_multiple_input_types_only(self): + class InputTypes(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType()), "y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x, y): + x += y + return x + + obj = InputTypes() + result = obj(x=torch.zeros(10), y=torch.ones(10)) + + assert result.sum() == torch.tensor(10.0) + assert hasattr(result, 'neural_type') is False + + @pytest.mark.unit + def test_output_types_only(self): + class OutputTypes(Typing): + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + obj = OutputTypes() + result = obj(x=torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert result.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + # Test passing positional args + # Positional args allowed if input types is not set ! + result = obj(torch.zeros(10)) + assert result.sum() == torch.tensor(10.0) + + @pytest.mark.unit + def test_multiple_output_types_only(self): + class MultipleOutputTypes(Typing): + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType()), "z": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + y = x + 1 + z = x + 2 + return y, z + + obj = MultipleOutputTypes() + result_y, result_z = obj(x=torch.zeros(10)) + + assert result_y.sum() == torch.tensor(10.0) + assert result_y.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + assert result_z.sum() == torch.tensor(20.0) + assert result_z.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_multiple_mixed_output_types_only(self): + class MultipleMixedOutputTypes(Typing): + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType()), "z": [NeuralType(('B',), ElementType())]} + + @typecheck() + def __call__(self, x): + y = x + 1 + z = x + 2 + return y, [z, z] + + obj = MultipleMixedOutputTypes() + result_y, result_z = obj(x=torch.zeros(10)) + + assert result_y.sum() == torch.tensor(10.0) + assert result_y.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + assert result_z[0].sum() == torch.tensor(20.0) + assert result_z[0].neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + assert result_z[1].sum() == torch.tensor(20.0) + assert result_z[1].neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_multiple_mixed_output_types_only_mismatched(self): + class MultipleMixedOutputTypes(Typing): + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType()), "z": [NeuralType(('B',), ElementType())]} + + @typecheck() + def __call__(self, x): + # Use list of y, single z, contrary to signature + y = x + 1 + z = x + 2 + return [y, y], z + + obj = MultipleMixedOutputTypes() + with pytest.raises(TypeError): + result_y, result_z = obj(x=torch.zeros(10)) + + @pytest.mark.unit + def test_incorrect_inheritance(self): + class IncorrectInheritance(object): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + obj = IncorrectInheritance() + + with pytest.raises(RuntimeError): + _ = obj(x=torch.zeros(10)) + + @pytest.mark.unit + def test_port_definition_rejection(self): + class InputPortDefinitionRejection(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"w": NeuralType(('B',), ElementType()), "u": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x, y): + x += 1 + y -= 1 + return x, y + + # Test input port mismatch + obj = InputPortDefinitionRejection() + + with pytest.raises(TypeError): + _ = obj(x=torch.zeros(10), y=torch.zeros(10)) + + class OutputPortDefinitionRejection(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return { + "w": NeuralType(('B',), ElementType()), + } + + @typecheck() + def __call__(self, x): + return x + 1, x - 1 + + obj = OutputPortDefinitionRejection() + + with pytest.raises(TypeError): + _ = obj(x=torch.zeros(10)) + + @pytest.mark.unit + def test_port_shape_rejection(self): + class InputPortShapeRejection(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B', 'T'), ElementType())} # expect rank 2 matrix + + @property + def output_types(self): + return {"w": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + # Test input port mismatch + obj = InputPortShapeRejection() + + with pytest.raises(TypeError): + _ = obj(x=torch.zeros(10)) + + class OutputPortShapeRejection(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return { + "w": NeuralType(('B', 'T', 'D'), ElementType()), # expect rank 3 matrix + } + + @typecheck() + def __call__(self, x): + return x + 1 + + obj = OutputPortShapeRejection() + + with pytest.raises(TypeError): + _ = obj(x=torch.zeros(10)) + + @pytest.mark.unit + def test_positional_args(self): + # Test positional check on input type + class InputPositional(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + obj = InputPositional() + + with pytest.raises(TypeError): + _ = obj(torch.zeros(10)) + + # Test positional pass-through for only output ports defined + # NOTE: This is required behaviour to support type checking of NeMo Dataset class + # during collate_fn() call. + class OutputPositionalPassthrough(Typing): + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + obj = OutputPositionalPassthrough() + result = obj(torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + + @pytest.mark.unit + def test_optional_types(self): + class InputOptionalTypes(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType()), "y": NeuralType(('B',), ElementType(), optional=True)} + + @typecheck() + def __call__(self, x, y=None): + if y is None: + x += 1 + else: + x += y + return x + + obj = InputOptionalTypes() + result = obj(x=torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert hasattr(result, 'neural_type') is False + + result2 = obj(x=torch.zeros(10), y=torch.full([10], fill_value=5, dtype=torch.int32)) + + assert result2.sum() == torch.tensor(10 * 5) + assert hasattr(result, 'neural_type') is False + + @pytest.mark.unit + def test_input_output_neural_types(self): + class NodeA(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": NeuralType(('B', 'D'), LogitsType())} + + @typecheck() + def __call__(self, x): + y = torch.randn(x.shape[0], 4) + return y + + class NodeB(Typing): + @property + def input_types(self): + return {"w": NeuralType(('B', 'D'), LogitsType())} + + @property + def output_types(self): + return {"u": NeuralType(('B',), LabelsType())} + + @typecheck() + def __call__(self, w): + _, u = w.max(-1) + return u + + nodeA = NodeA() + nodeB = NodeB() + + outA = nodeA(x=torch.zeros(10)) + outB = nodeB(w=outA) + + assert outB.shape == torch.Size([10]) + assert outB.neural_type.compare(NeuralType(('B',), LabelsType())) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_nested_input_output_neural_types(self): + class NestedNodeA(Typing): + @property + def input_types(self): + return {"x": [[NeuralType(('B',), ElementType())]]} + + @property + def output_types(self): + return { + "y0": [NeuralType(('B', 'D'), LogitsType())], + "y1": [NeuralType(('B', 'D'), LogitsType())], + } + + @typecheck(ignore_collections=False) + def __call__(self, x): + # input x = [[x1, x2], [x3]] + x0 = x[0][0] + y = torch.randn(x0.shape[0], 4) + # Output is same as + # 1) return ([y, y], [y]) + # 2) return [y, y], [y] + return [[y, y], [y]] + + # Non-homogeneous output types + class NestedNodeB(Typing): + @property + def input_types(self): + return {"w": [[NeuralType(('B', 'D'), LogitsType())]]} + + @property + def output_types(self): + return { + "u0": [NeuralType(('B',), LogprobsType())], # check non homogeneous type + "u1": [NeuralType(('B',), LabelsType())], + } + + @typecheck(ignore_collections=False) + def __call__(self, w): + # input x = [[x1, x2], [x3]] + _, u00 = w[0][0].max(-1) + _, u01 = w[0][1].max(-1) + _, u10 = w[1][0].max(-1) + # Output is same as + # 1) return ([u00, u01], [u10]) + # 2) return [u00, u01], [u10] + return [[u00, u01], [u10]] + + nodeA = NestedNodeA() + nodeB = NestedNodeB() + + input_nest = [[torch.zeros(10), torch.zeros(10)], [torch.zeros(10)]] + outA = nodeA(x=input_nest) + outB = nodeB(w=outA) + + # Perform recursive shape assert + recursive_assert_shape(outB, torch.Size([10])) + + # Assert non-homogeneous type assertions + assert outB[0][0].neural_type.compare(NeuralType(('B',), LogprobsType())) == NeuralTypeComparisonResult.SAME + assert outB[0][1].neural_type.compare(NeuralType(('B',), LogprobsType())) == NeuralTypeComparisonResult.SAME + assert outB[1][0].neural_type.compare(NeuralType(('B',), LabelsType())) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_nested_input_output_neural_types_ignore_collections(self): + class NestedNodeA(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return { + "y0": NeuralType(('B', 'D'), LogitsType()), + "y1": NeuralType(('B', 'D'), LogitsType()), + } + + @typecheck(ignore_collections=True) + def __call__(self, x): + # input x = [[x1, x2], [x3]] + x0 = x[0][0] + y = torch.randn(x0.shape[0], 4) + return [[y, y], [y]] + + # Non-homogeneous output types + class NestedNodeB(Typing): + @property + def input_types(self): + return {"w": NeuralType(('B', 'D'), LogitsType())} + + @property + def output_types(self): + return { + "u0": NeuralType(('B',), LogprobsType()), # check non homogeneous type + "u1": NeuralType(('B',), LabelsType()), + } + + @typecheck(ignore_collections=True) + def __call__(self, w): + # input x = [[x1, x2], [x3]] + _, u00 = w[0][0].max(-1) + _, u01 = w[0][1].max(-1) + _, u10 = w[1][0].max(-1) + return [[u00, u01], [u10]] + + nodeA = NestedNodeA() + nodeB = NestedNodeB() + + input_nest = [[torch.zeros(10), torch.zeros(10)], [torch.zeros(10)]] + outA = nodeA(x=input_nest) + outB = nodeB(w=outA) + + # Perform recursive shape assert + recursive_assert_shape(outB, torch.Size([10])) + + # Assert non-homogeneous type assertions + assert outB[0][0].neural_type.compare(NeuralType(('B',), LogprobsType())) == NeuralTypeComparisonResult.SAME + assert outB[0][1].neural_type.compare(NeuralType(('B',), LogprobsType())) == NeuralTypeComparisonResult.SAME + assert outB[1][0].neural_type.compare(NeuralType(('B',), LabelsType())) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_nested_mixed_input_output_neural_types(self): + class NestedMixedNodeA(Typing): + @property + def input_types(self): + return {"x1": NeuralType(('B',), ElementType()), "x2": [[NeuralType(('B',), ElementType())]]} + + @property + def output_types(self): + return { + "y0": NeuralType(('B', 'D'), LogprobsType()), + "y1": [[NeuralType(('B', 'D'), LogitsType())]], + } + + @typecheck(ignore_collections=False) + def __call__(self, x1, x2): + # input x = [[x1, x2], [x3]] + x0 = x2[0][0] + y = torch.randn(x0.shape[0], 4) + return y, [[y, y], [y]] + + # Non-homogeneous output types + class NestedMixedNodeB(Typing): + @property + def input_types(self): + return {"w": [[NeuralType(('B', 'D'), LogitsType())]]} + + @property + def output_types(self): + return { + "u0": [NeuralType(('B',), LogprobsType())], # check non homogeneous type + "u1": NeuralType(('B',), LabelsType()), + } + + @typecheck(ignore_collections=False) + def __call__(self, w): + # input x = [[x1, x2], [x3]] + _, u00 = w[0][0].max(-1) + _, u01 = w[0][1].max(-1) + _, u10 = w[1][0].max(-1) + return [u00, u01], u10 + + nodeA = NestedMixedNodeA() + nodeB = NestedMixedNodeB() + + input_nest = [[torch.zeros(10), torch.zeros(10)], [torch.zeros(10)]] + out_y, outA = nodeA(x1=torch.zeros(10), x2=input_nest) + outB, out_u = nodeB(w=outA) + + # Perform recursive shape assert + assert out_y.neural_type.compare(NeuralType(('B', 'D'), LogprobsType())) + recursive_assert_shape(outB, torch.Size([10])) + + # Assert non-homogeneous type assertions + assert outB[0].neural_type.compare(NeuralType(('B',), LogprobsType())) == NeuralTypeComparisonResult.SAME + assert outB[1].neural_type.compare(NeuralType(('B',), LogprobsType())) == NeuralTypeComparisonResult.SAME + assert out_u.neural_type.compare(NeuralType(('B',), LabelsType())) == NeuralTypeComparisonResult.SAME + + @pytest.mark.unit + def test_multi_forward_type(self): + class AdaptiveTypeCheck(Typing): + @property + def input_types(self): + if self.mode == 'train': + return {"x": NeuralType(('B',), ElementType())} + + elif self.mode == 'infer': + return {"y": NeuralType(('B',), ChannelType())} + + elif self.mode == 'eval': + return {"x": NeuralType(('B',), ElementType()), "y": NeuralType(('B',), ChannelType())} + else: + raise ValueError("Wrong mode of operation") + + @property + def output_types(self): + if self.mode == 'train': + return {"u": NeuralType(('B',), ElementType())} + + elif self.mode == 'infer': + return {"v": NeuralType(('B',), ChannelType())} + + elif self.mode == 'eval': + return {"u": NeuralType(('B',), ElementType()), "v": NeuralType(('B',), ChannelType())} + else: + raise ValueError("Wrong mode of operation") + + def __init__(self): + self.mode = 'train' + + def __call__(self, **kwargs): + # Call should call and forward appropriate method in its own mode + if self.mode == 'train': + return self.train_forward(x=kwargs['x']) + + elif self.mode == 'eval': + return self.eval_forward(x=kwargs['x'], y=kwargs['y']) + + elif self.mode == 'infer': + return self.infer_forward(y=kwargs['y']) + + @typecheck() + def train_forward(self, x): + return x + 10 + + @typecheck() + def eval_forward(self, x, y): + return x - 1, y - 1 + + @typecheck() + def infer_forward(self, y): + return y - 10 + + @property + def mode(self): + return self._mode + + @mode.setter + def mode(self, val): + if val not in ['train', 'infer', 'eval']: + raise ValueError('mode must be either train infer or eval') + self._mode = val + + obj = AdaptiveTypeCheck() + + x = torch.zeros(10) + y = torch.full([10], fill_value=5, dtype=torch.int32) + + obj.mode = 'train' + x = obj(x=x) + + assert torch.all(x == 10) + assert x.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + obj.mode = 'eval' + x, y = obj(x=x, y=y) + + assert torch.all(x == 9) + assert torch.all(y == 4) + assert x.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + assert y.neural_type.compare(NeuralType(('B',), ChannelType())) == NeuralTypeComparisonResult.SAME + + obj.mode = 'infer' + y = obj(y=y) + + assert torch.all(y == -6) + assert y.neural_type.compare(NeuralType(('B',), ChannelType())) == NeuralTypeComparisonResult.SAME + + # Now perform assertions of wrong mode with wrong input combinations + obj.mode = 'train' + + # In train mode, call infer + with pytest.raises(TypeError): + _ = obj.eval_forward(x=x, y=y) + + with pytest.raises(TypeError): + # wrong input + wrong mode + _ = obj.infer_forward(y=x) + + @pytest.mark.unit + def test_input_type_override(self): + class InputTypesOverride(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + @typecheck(input_types={"y": NeuralType(('B',), CategoricalValuesType())}) + def forward(self, y): + y -= 1 + return y + + obj = InputTypesOverride() + result = obj(x=torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert hasattr(result, 'neural_type') is False + + # Test override + result2 = obj.forward(y=torch.zeros(10)) + + assert result2.sum() == torch.tensor(-10.0) + assert hasattr(result2, 'neural_type') is False + + @pytest.mark.unit + def test_output_type_override(self): + class OutputTypes(Typing): + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x): + x += 1 + return x + + @typecheck(output_types={"z": NeuralType(('B',), CategoricalValuesType())}) + def forward(self, z): + z -= 1 + return z + + obj = OutputTypes() + result = obj(x=torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert result.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + # Test passing positional args + # Positional args allowed if input types is not set ! + result = obj(torch.zeros(10)) + assert result.sum() == torch.tensor(10.0) + + # Test override + result2 = obj.forward(z=torch.zeros(10)) + + assert result2.sum() == torch.tensor(-10.0) + assert hasattr(result2, 'neural_type') + assert ( + result2.neural_type.compare(NeuralType(('B',), CategoricalValuesType())) == NeuralTypeComparisonResult.SAME + ) + + @pytest.mark.unit + def test_multi_type_override(self): + class AdaptiveTypeCheck(Typing): + @property + def input_types(self): + # __call__ assumed to be for inference only, + # therefore infer types checked at class scope + return {"y": NeuralType(('B',), ChannelType())} + + @property + def output_types(self): + # __call__ assumed to be for inference only, + # therefore infer types checked at class scope + return {"v": NeuralType(('B',), ChannelType())} + + def __call__(self, **kwargs): + # Call should call and forward appropriate method in its own mode + # Let default "forward" call be the infer mode (this is upto developer) + # Therefore default class level types == infer types + return self.infer_forward(y=kwargs['y']) + + @typecheck( + input_types={"x": NeuralType(('B',), ElementType())}, + output_types={"u": NeuralType(('B',), ElementType())}, + ) + def train_forward(self, x): + return x + 10 + + @typecheck( + input_types={"x": NeuralType(('B',), ElementType()), "y": NeuralType(('B',), ChannelType())}, + output_types={"u": NeuralType(('B',), ElementType()), "v": NeuralType(('B',), ChannelType())}, + ) + def eval_forward(self, x, y): + return x - 1, y - 1 + + @typecheck( + input_types={"y": NeuralType(('B',), ChannelType())}, + output_types={"v": NeuralType(('B',), ChannelType())}, + ) + def infer_forward(self, y): + return y - 10 + + obj = AdaptiveTypeCheck() + + x = torch.zeros(10) + y = torch.full([10], fill_value=5, dtype=torch.int32) + + # infer mode + y = obj(y=y) + + assert torch.all(y == -5) + assert y.neural_type.compare(NeuralType(('B',), ChannelType())) == NeuralTypeComparisonResult.SAME + + x, y = obj.eval_forward(x=x, y=y) + + assert torch.all(x == -1) + assert torch.all(y == -6) + assert x.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + assert y.neural_type.compare(NeuralType(('B',), ChannelType())) == NeuralTypeComparisonResult.SAME + + x = obj.train_forward(x=x) + + assert torch.all(x == 9) + assert x.neural_type.compare(NeuralType(('B',), ElementType())) == NeuralTypeComparisonResult.SAME + + # In train func, call eval signature + with pytest.raises(TypeError): + _ = obj.train_forward(x=x, y=y) + + with pytest.raises(TypeError): + # wrong input + wrong mode + _ = obj.infer_forward(x=x) + + @pytest.mark.unit + def test_disable_typecheck(self): + class InputOutputTypes(Typing): + @property + def input_types(self): + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": NeuralType(('B',), ElementType())} + + @typecheck() + def __call__(self, x, **kwargs): + x += 1 + return x + + # Disable typecheck tests + with typecheck.disable_checks(): + obj = InputOutputTypes() + + # Execute function without kwarg + result = obj(torch.zeros(10)) + + assert result.sum() == torch.tensor(10.0) + assert hasattr(result, 'neural_type') is False + + # Test passing wrong key for input + _ = obj(a=torch.zeros(10), x=torch.zeros(5)) + + @pytest.mark.unit + def test_nested_shape_mismatch(self): + class NestedShapeMismatch(Typing): + @property + def input_types(self): + return {"x": [[NeuralType(('D',), ElementType())]]} # Each element of nest will have 4 values + + @property + def output_types(self): + return {"y": [[NeuralType(('D',), ElementType())]]} # Each element of nest will have 4 values + + @typecheck() + def __call__(self, x): + # v-- this is to satisfy 1 output constraint, python will otherwise interpret x as a 3 output value + return x + + def bb(dim=4): + return torch.zeros(dim) + + obj = NestedShapeMismatch() + + # Arbitrary nest 1 (should pass) + data = [[bb(), bb(), bb()], [bb()], [bb(), bb()]] + result = obj(x=data) + + recursive_assert_shape(result, torch.Size([4])) + recursive_assert_homogeneous_type(result, NeuralType(('D',), ElementType())) + + # Arbitrary nest 2 (should pass) + def bb(dim=4): + return torch.zeros(dim, dim) + + data = [[bb(), bb(), bb()], [bb()], [bb(), bb()]] + # Fails since input shape is incorrect + with pytest.raises(TypeError): + _ = obj(x=data) + + # Arbitrary nest 3 + def bb(dim=4): + return torch.zeros(dim) + + data = [[[bb(), bb(), bb()]], [[bb()], [bb(), bb()]]] + # Check should fail since nest level is 3! + with pytest.raises(TypeError): + result = obj(x=data) + + @pytest.mark.unit + def test_nested_mixed_shape_mismatch(self): + class NestedMixedShapeMismatch(Typing): + @property + def input_types(self): + return {"x": [[NeuralType(('D',), ElementType())]]} # Each element of nest will have 4 values + + @property + def output_types(self): + return {"y": [NeuralType(('D',), ElementType())]} # Each element of nest will have 4 values + + @typecheck() + def __call__(self, x): + # v-- this is to satisfy 1 output constraint, python will otherwise interpret x as a 3 output value + x = x[0] + return x + + def bb(dim=4): + return torch.zeros(dim) + + obj = NestedMixedShapeMismatch() + + # Arbitrary nest 1 (should pass) + data = [[bb(), bb(), bb()], [bb()], [bb(), bb()]] + result = obj(x=data) + + recursive_assert_shape(result, torch.Size([4])) + recursive_assert_homogeneous_type(result, NeuralType(('D',), ElementType())) + + # Arbitrary nest 2 (should pass) + def bb(dim=4): + return torch.zeros(dim, dim) + + data = [[bb(), bb(), bb()], [bb()], [bb(), bb()]] + # Fails since input shape is incorrect + with pytest.raises(TypeError): + _ = obj(x=data) + + # Arbitrary nest 3 + def bb(dim=4): + return torch.zeros(dim) + + data = [[[bb(), bb(), bb()]], [[bb()], [bb(), bb()]]] + # Check should fail since nest level is 3! + with pytest.raises(TypeError): + result = obj(x=data) + + @pytest.mark.unit + def test_input_container_neural_types(self): + class NodeA(Typing): + @property + def input_types(self): + return {"x": [NeuralType(('B',), ElementType())]} + + @property + def output_types(self): + return {"y": NeuralType(('B', 'D'), LogitsType())} + + @typecheck() + def __call__(self, x: list()): + x1, x2, x3 = x # unpack x + y = torch.randn(x1.shape[0], 4) + return y + + nodeA = NodeA() + outA = nodeA(x=[torch.zeros(10), torch.zeros(10), torch.zeros(10)]) + + assert outA.neural_type.compare(NeuralType(('B', 'D'), LogitsType())) + + @pytest.mark.unit + def test_input_container_neural_types_incorrect(self): + class NodeA(Typing): + @property + def input_types(self): + # Nest depth level of 2 + return {"x": [[NeuralType(('B',), ElementType())]]} + + @property + def output_types(self): + return {"y": NeuralType(('B', 'D'), LogitsType())} + + @typecheck() + def __call__(self, x: list()): + x1, x2, x3 = x # unpack x + y = torch.randn(x1.shape[0], 4) + return y + + nodeA = NodeA() + # Input nest level of 1 + with pytest.raises(TypeError): + outA = nodeA(x=[torch.zeros(10), torch.zeros(10), torch.zeros(10)]) + + @pytest.mark.unit + def test_output_container_neural_types_incorrect(self): + class NodeA(Typing): + @property + def input_types(self): + # Nest depth level of 2 + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": [[NeuralType(('B', 'D'), LogitsType())]]} + + @typecheck() + def __call__(self, x): + y = torch.randn(x.shape[0], 4) + return y, y, y + + nodeA = NodeA() + # Input nest level of 1 + with pytest.raises(TypeError): + outA = nodeA(x=torch.zeros(10)) + + @pytest.mark.unit + def test_output_container_neural_types_no_tuple_wrap(self): + class NodeA(Typing): + @property + def input_types(self): + # Nest depth level of 2 + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": [NeuralType(('B', 'D'), LogitsType())]} + + @typecheck() + def __call__(self, x): + y = torch.randn(x.shape[0], 4) + y = [y, y, y] + return y + + nodeA = NodeA() + # Input nest level of 1 + outA = nodeA(x=torch.zeros(10)) + + assert len(outA) == 3 + for i in range(len(outA)): + assert outA[i].neural_type.compare(NeuralType(('B', 'D'), LogitsType())) + + @pytest.mark.unit + def test_output_container_neural_types_explicit_tuple_wrap(self): + class NodeA(Typing): + @property + def input_types(self): + # Nest depth level of 2 + return {"x": NeuralType(('B',), ElementType())} + + @property + def output_types(self): + return {"y": [NeuralType(('B', 'D'), LogitsType())]} + + @typecheck() + def __call__(self, x): + y = torch.randn(x.shape[0], 4) + y = [y, y, y] + return (y,) + + nodeA = NodeA() + # Input nest level of 1 + outA = nodeA(x=torch.zeros(10)) + + assert len(outA) == 1 + assert len(outA[0]) == 3 + for i in range(len(outA)): + assert outA[0][i].neural_type.compare(NeuralType(('B', 'D'), LogitsType())) diff --git a/models/tts/NeMo/tests/core_ptl/check_for_ranks.py b/models/tts/NeMo/tests/core_ptl/check_for_ranks.py new file mode 100644 index 0000000000..b936e2da1d --- /dev/null +++ b/models/tts/NeMo/tests/core_ptl/check_for_ranks.py @@ -0,0 +1,155 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os +import shutil + +import torch +from omegaconf import OmegaConf +from pytorch_lightning import Trainer +from pytorch_lightning.utilities.distributed import rank_zero_only + +from nemo.core import ModelPT +from nemo.utils import logging +from nemo.utils.exp_manager import ExpManagerConfig, exp_manager + + +class OnesDataset(torch.utils.data.Dataset): + def __init__(self, dataset_len): + super().__init__() + self.__dataset_len = dataset_len + + def __getitem__(self, *args): + return torch.ones(2) + + def __len__(self): + return self.__dataset_len + + +class ExampleModel(ModelPT): + def __init__(self, *args, **kwargs): + cfg = OmegaConf.structured({}) + super().__init__(cfg, trainer=kwargs.get('trainer', None)) + # dummy parameter in order to allow DDP to execute + self.l1 = torch.nn.modules.Linear(in_features=2, out_features=1) + + def train_dataloader(self): + return None + + def val_dataloader(self): + return None + + def predict_dataloader(self): + dataset = OnesDataset(2) + return torch.utils.data.DataLoader(dataset, batch_size=2) + + def forward(self, batch): + return batch.mean() + + def validation_step(self, batch, batch_idx): + return self(batch) + + def training_step(self, batch, batch_idx): + return self(batch) + + def list_available_models(self): + pass + + def setup_training_data(self): + pass + + def setup_validation_data(self): + pass + + def validation_epoch_end(self, loss): + self.log("val_loss", torch.stack(loss).mean()) + + +def instantiate_multinode_ddp_if_possible(): + num_gpus = torch.cuda.device_count() + trainer = Trainer(gpus=num_gpus, accelerator='ddp', logger=None, checkpoint_callback=None) + + exp_manager_cfg = ExpManagerConfig(exp_dir='./ddp_check/', use_datetime_version=False, version="") + exp_manager(trainer, cfg=OmegaConf.structured(exp_manager_cfg)) + return trainer + + +def setup_model(trainer: Trainer): + model = ExampleModel(trainer=trainer) + + logging.info(f"M.Global Rank:{model.global_rank}") + logging.info(f"M.Local Rank:{model.local_rank}") + logging.info(f"M.World Size:{model.trainer.world_size}") + + trainer.predict(model) + return model + + +def get_rank_info(texts: list, rank_key: str) -> int: + for line in texts: + if rank_key in line: + rank_value = line.split(":")[-1] + rank_value = int(rank_value) + return rank_value + + print("Could not find the correct rank key !") + exit(1) + + +@rank_zero_only +def check_model_ranks(model: ExampleModel): + basedir = os.path.join('./ddp_check/', 'default', 'version_0') + file_template = "nemo_log_globalrank-{rank}_localrank-{rank}.txt" + + world_size = torch.cuda.device_count() + for rank in range(world_size): + filename = file_template.format(rank=rank) + filepath = os.path.join(basedir, filename) + + with open(filepath, 'r') as f: + texts = f.readlines() + texts = [t.replace("\n", "") for t in texts] + + log_global_rank = get_rank_info(texts, rank_key='M.Global Rank') + log_world_size = get_rank_info(texts, rank_key='M.World Size') + + if log_global_rank != rank: + print("Logged global rank is not equal to trainer.global_rank !") + exit(1) + + if log_world_size != world_size: + print("Logged world size if not equal to trainer.world_size !") + exit(1) + + +@rank_zero_only +def cleanup(): + if os.path.exists('./ddp_check'): + shutil.rmtree('./ddp_check', ignore_errors=True) + + +def run_checks(): + cleanup() + + trainer = instantiate_multinode_ddp_if_possible() + model = setup_model(trainer) + check_model_ranks(model) + + print("DDP checks passed !") + + cleanup() + + +if __name__ == '__main__': + run_checks() diff --git a/models/tts/NeMo/tests/hydra/config1.yaml b/models/tts/NeMo/tests/hydra/config1.yaml new file mode 100644 index 0000000000..71425c0e8c --- /dev/null +++ b/models/tts/NeMo/tests/hydra/config1.yaml @@ -0,0 +1 @@ +dataset_name: dark diff --git a/models/tts/NeMo/tests/hydra/config1_invalid.yaml b/models/tts/NeMo/tests/hydra/config1_invalid.yaml new file mode 100644 index 0000000000..44ea0bdf04 --- /dev/null +++ b/models/tts/NeMo/tests/hydra/config1_invalid.yaml @@ -0,0 +1,2 @@ +dataset_name: magic +password: secret \ No newline at end of file diff --git a/models/tts/NeMo/tests/hydra/config_subdir/config2.yaml b/models/tts/NeMo/tests/hydra/config_subdir/config2.yaml new file mode 100644 index 0000000000..e773e40209 --- /dev/null +++ b/models/tts/NeMo/tests/hydra/config_subdir/config2.yaml @@ -0,0 +1 @@ +dataset_name: ala_config2 diff --git a/models/tts/NeMo/tests/hydra/config_subdir/config2_invalid.yaml b/models/tts/NeMo/tests/hydra/config_subdir/config2_invalid.yaml new file mode 100644 index 0000000000..f06f3eee20 --- /dev/null +++ b/models/tts/NeMo/tests/hydra/config_subdir/config2_invalid.yaml @@ -0,0 +1,2 @@ +dataset_name: ala_config2 +password: ala_secret2 \ No newline at end of file diff --git a/models/tts/NeMo/tests/hydra/my_app.py b/models/tts/NeMo/tests/hydra/my_app.py new file mode 100644 index 0000000000..dca19dc0bd --- /dev/null +++ b/models/tts/NeMo/tests/hydra/my_app.py @@ -0,0 +1,44 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +from dataclasses import dataclass + +import hydra +from omegaconf import MISSING, OmegaConf + +from nemo.core.config import hydra_runner + + +@dataclass +class DefaultConfig: + """ + This is structured config for this application. + + It provides the schema used for validation of user-written spec file + as well as default values of the selected parameters. + """ + + # Dataset. Available options: [imdb, sst2] + dataset_name: str = MISSING + + +@hydra_runner(config_name="DefaultConfig", schema=DefaultConfig) +def my_app(cfg): + print(OmegaConf.to_yaml(cfg)) + # Get dataset_name. + dataset_name = cfg.dataset_name + + +if __name__ == "__main__": + my_app() diff --git a/models/tts/NeMo/tests/hydra/test_hydra_runner.py b/models/tts/NeMo/tests/hydra/test_hydra_runner.py new file mode 100644 index 0000000000..1da0a914cf --- /dev/null +++ b/models/tts/NeMo/tests/hydra/test_hydra_runner.py @@ -0,0 +1,95 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 subprocess +import sys +from os import path + +import pytest + + +class TestHydraRunner: + @pytest.mark.integration + def test_no_config(self): + """"Test app without config - fields missing causes error. + """ + # Create system call. + call = "python tests/hydra/my_app.py" + + with pytest.raises(subprocess.CalledProcessError): + # Run the call as subprocess. + subprocess.check_call(call, shell=True, stdout=sys.stdout, stderr=sys.stdout) + + @pytest.mark.integration + def test_config1(self): + """"Test injection of valid config1. + """ + # Create system call. + call = "python tests/hydra/my_app.py --config-name config1.yaml" + + # Run the call as subprocess. + subprocess.check_call(call, shell=True, stdout=sys.stdout, stderr=sys.stdout) + + # Make sure that .hydra dir is not present. + assert not path.exists(f".hydra") + # Make sure that default hydra log file is not present. + assert not path.exists(f"my_app.log") + + @pytest.mark.integration + def test_config1_invalid(self): + """"Test injection of invalid config1. + """ + # Create system call. + call = "python tests/hydra/my_app.py --config-name config1_invalid.yaml" + + with pytest.raises(subprocess.CalledProcessError): + # Run the call as subprocess. + subprocess.check_call(call, shell=True, stdout=sys.stdout, stderr=sys.stdout) + + @pytest.mark.integration + def test_config2(self): + """"Test injection of valid config2 from a different folder. + """ + # Create system call. + call = "python tests/hydra/my_app.py --config-path config_subdir --config-name config2.yaml" + + # Run the call as subprocess. + subprocess.check_call(call, shell=True, stdout=sys.stdout, stderr=sys.stdout) + + # Make sure that .hydra dir is not present. + assert not path.exists(f".hydra") + # Make sure that default hydra log file is not present. + assert not path.exists(f"my_app.log") + + @pytest.mark.integration + def test_config2_invalid(self): + """"Test injection of invalid config2 from a different folder. + """ + # Create system call. + call = "python tests/hydra/my_app.py --config-path config_subdir --config-name config2_invalid.yaml" + + with pytest.raises(subprocess.CalledProcessError): + # Run the call as subprocess. + subprocess.check_call(call, shell=True, stdout=sys.stdout, stderr=sys.stdout) + + @pytest.mark.integration + def test_config2_filepath_schema(self): + """"Test injection of valid config2 - using namepath with schema is prohibited. + """ + # Create system call. + call = "python tests/hydra/my_app.py --config-name config_subdir/config2_invalid.yaml" + + with pytest.raises(subprocess.CalledProcessError): + # Run the call as subprocess. + subprocess.check_call(call, shell=True, stdout=sys.stdout, stderr=sys.stdout) diff --git a/models/tts/NeMo/tests/manualtest_model_downloads.py b/models/tts/NeMo/tests/manualtest_model_downloads.py new file mode 100644 index 0000000000..a3ed1734a8 --- /dev/null +++ b/models/tts/NeMo/tests/manualtest_model_downloads.py @@ -0,0 +1,92 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +""" +WARNING: Running this test will download ALL pre-trained NeMo models. +This is bandwidth and disk space consuming. +""" + +import nemo.collections.asr as nemo_asr +import nemo.collections.nlp as nemo_nlp +import nemo.collections.tts as nemo_tts + + +def testclass_downloads(cls, refresh_cache, model_names=None): + for model_info in cls.list_available_models(): + model = cls.from_pretrained(model_name=model_info.pretrained_model_name, refresh_cache=refresh_cache) + assert isinstance(model, cls) + if model_names is not None: + assert set(model_names) == set([m.pretrained_model_name for m in cls.list_available_models()]) + + +for refresh_cache in [True, False]: + # Test ASR collection + testclass_downloads( + nemo_asr.models.EncDecCTCModel, + refresh_cache, + [ + 'QuartzNet15x5Base-En', + 'QuartzNet15x5Base-Zh', + 'QuartzNet5x5LS-En', + 'QuartzNet15x5NR-En', + 'Jasper10x5Dr-En', + ], + ) + testclass_downloads(nemo_asr.models.EncDecCTCModelBPE, refresh_cache, ['ContextNet-192-WPE-1024-8x-Stride']) + testclass_downloads( + nemo_asr.models.EncDecClassificationModel, + refresh_cache, + [ + 'MatchboxNet-3x1x64-v1', + 'MatchboxNet-3x2x64-v1', + 'MatchboxNet-3x1x64-v2', + 'MatchboxNet-3x1x64-v2', + 'MatchboxNet-3x1x64-v2-subset-task', + 'MatchboxNet-3x2x64-v2-subset-task', + 'MatchboxNet-VAD-3x2', + ], + ) + testclass_downloads( + nemo_asr.models.EncDecSpeakerLabelModel, + refresh_cache, + ['speakerrecognition_speakernet', 'speakerverification_speakernet', 'speakerdiarization_speakernet'], + ) + + # Test NLP collection + testclass_downloads(nemo_nlp.models.TokenClassificationModel, refresh_cache, ['NERModel']) + testclass_downloads( + nemo_nlp.models.PunctuationCapitalizationModel, + refresh_cache, + ['Punctuation_Capitalization_with_BERT', 'Punctuation_Capitalization_with_DistilBERT'], + ) + testclass_downloads( + nemo_nlp.models.QAModel, + refresh_cache, + [ + 'BERTBaseUncasedSQuADv1.1', + 'BERTBaseUncasedSQuADv2.0', + 'BERTLargeUncasedSQuADv1.1', + 'BERTLargeUncasedSQuADv2.0', + ], + ) + # testclass_downloads(nemo_nlp.models.IntentSlotClassificationModel, refresh_cache, ['Joint_Intent_Slot_Assistant']) + + # Test TTS collection + testclass_downloads(nemo_tts.models.Tacotron2Model, refresh_cache, ['Tacotron2-22050Hz']) + testclass_downloads(nemo_tts.models.WaveGlowModel, refresh_cache, ['WaveGlow-22050Hz']) + testclass_downloads(nemo_tts.models.SqueezeWaveModel, refresh_cache, ['SqueezeWave-22050Hz']) + testclass_downloads(nemo_tts.models.GlowTTSModel, refresh_cache, ['GlowTTS-22050Hz']) + + +print("############ THAT'S ALL FOLKS! ############") diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_cardinal.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_cardinal.txt new file mode 100644 index 0000000000..fb668c03ac --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_cardinal.txt @@ -0,0 +1,24 @@ +nine trillion seven hundred eighty nine billion three hundred eighty two million five hundred thirty six thousand one hundred thirty~9789382536130 +two hundred and fifty four~254 +one hundred forty seven thousand four hundred fifty one~147451 +one million one hundred fifty six thousand one hundred seventy three~1156173 +one billion five hundred ninety three million seventy two thousand nine hundred sixty one~1593072961 +ninety seven quadrillion eight hundred eight trillion two hundred sixty four billion seven hundred seventy two million seven hundred ninety two thousand five~97808264772792005 +seventeen sextillion eight hundred fifty five quintillion thirty six quadrillion six hundred fifty seven trillion seven billion five hundred ninety six million one hundred ten thousand nine hundred forty nine~17855036657007596110949 +ten quadrillion ten trillion ten million one hundred thousand ten~10010000010100010 +minus twenty five thousand thirty seven~-25037 +one quadrillion two hundred sixty four trillion three hundred one billion nine hundred thirty eight million one hundred four~1264301938000104 +minus sixty~-60 +forty six thousand six hundred sixty four~46664 +sixty~60 +zero~zero +two million three~2000003 +one thousand thirteen~1013 +one thousand one~1001 +one thousand one hundred~1100 +one thousand twenty six~1026 +one thousand one hundred twenty six~1126 +eighteen million four hundred fifty thousand nine hundred ninety~18450990 +eighteen million nine hundred forty thousand seven hundred twenty two~18940722 +eighteen million six hundred ninety thousand nine hundred sixteen~18690916 +eighteen thousand eight hundred eighty~18880 diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_date.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_date.txt new file mode 100644 index 0000000000..3ec4268144 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_date.txt @@ -0,0 +1,34 @@ +july twenty fifth two thousand twelve~july 25 2012 +nineteen eighties~1980s +two thousand and twenty~2020 +two thousand and nine~2009 +the twenty fifth of july twenty twelve~25 july 2012 +the twenty fifth of july two thousand twelve~25 july 2012 +the twenty second of july twenty twelve~22 july 2012 +the fifteenth of january~15 january +the seventeenth of may twenty ten~17 may 2010 +january first~january 1 +july twenty second two thousand eight~july 22 2008 +june thirty~june 30 +july twenty fifth twenty twelve~july 25 2012 +nineteen seventeen~1917 +twenty twelve~2012 +march sixteen sixty five~march 1665 +sixteen sixty five~1665 +july two thousand twelve~july 2012 +october nineteen oh five~october 1905 +july fifteen o six~july 1506 +the twenty fifth of july twenty twelve~25 july 2012 +july twenty fifth twenty twelve~july 25 2012 +july twenty fifth two thousand twelve~july 25 2012 +july one thousand eight hundred seventy six~july 1876 +february twenty fifth twenty sixteen~february 25 2016 +november twenty fourth twenty fourteen~november 24 2014 +nineteen ninety four~1994 +two thousand three~2003 +one thousand eight~1008 +nineteen seventy six~1976 +june twentieth twenty fourteen~june 20 2014 +nineteen seventy three~1973 +nineteen seventy five~1975 +eleven fifty five~1155 \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_decimal.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_decimal.txt new file mode 100644 index 0000000000..2cd2272837 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_decimal.txt @@ -0,0 +1,63 @@ +five point two million~5.2 million +one hundred sixty four point five eight thousand~164.58 thousand +four hundred million~400 million +fifty billion~50 billion +four hundred five billion~405 billion +four point eight five billion~4.85 billion +one hundred billion~100 billion +one hundred ten billion~110 billion +one hundred thirty two billion~132 billion +one point eight four billion~1.84 billion +one point eight one billion~1.81 billion +one point five nine billion~1.59 billion +one point four five three billion~1.453 billion +one point seven two billion~1.72 billion +one point two five billion~1.25 billion +thirteen billion~13 billion +thirty billion~30 billion +two thousand eight hundred five point eight seven three billion~2805.873 billion +seventy trillion~70 trillion +thirteen million~13 million +eighteen billion~18 billion +four hundred fifty million~450 million +one hundred thirty million~130 million +ten million~10 million +four hundred million~400 million +five million~5 million +five hundred million~500 million +twelve million~12 million +thirteen million~13 million +four million~4 million +forty five million~45 million +fifteen million~15 million +fifteen trillion~15 trillion +fifteen billion~15 billion +two million~2 million +eight million~8 million +point one two o five~.1205 +minus sixty point two four zero zero~-60.2400 +zero point two six~0.26 +point zero two~.02 +sixty point two~60.2 +eighteen~18 +eighteen point eight five~18.85 +eighteen point five o~18.50 +eighteen point five six~18.56 +eighteen point nine~18.9 +eighteen point o five~18.05 +eighteen point one two~18.12 +eighteen point o one~18.01 +eighteen point o o o~18.000 +eighteen point six~18.6 +eighteen point three o o~18.300 +eighteen point three six~18.36 +eighteen point two five~18.25 +eighteen point two two~18.22 +eight hundred eighteen point three o three~818.303 +eight hundred eight point eight~808.8 +eight hundred eight point zero~808.0 +eight hundred eighty eight point one~888.1 +eight hundred eighty four point three~884.3 +eight hundred eighty two point eight~882.8 +eight hundred eighty two point zero~882.0 +eight hundred forty five point nine four~845.94 \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_electronic.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_electronic.txt new file mode 100644 index 0000000000..6267f87ab4 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_electronic.txt @@ -0,0 +1,7 @@ +a dot b c at g mail dot com~a.bc@gmail.com +c d f at a b c dot e d u~cdf@abc.edu +a b c at g mail dot a b c~abc@gmail.abc +a b c at a b c dot com~abc@abc.com +a s d f one two three at a b c dot com~asdf123@abc.com +a one b two at a b c dot com~a1b2@abc.com +a b three dot s d d dot three at g mail dot com~ab3.sdd.3@gmail.com \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_measure.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_measure.txt new file mode 100644 index 0000000000..4f6f540e6d --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_measure.txt @@ -0,0 +1,112 @@ +two hundred meters~200 m +fifty six point three per square kilometer~56.3 /km² +two hundred kilometers per hour~200 km/h +twenty eight kilograms force per square centimeter~28 kgf/cm² +forty two thousand two hundred fifty nine per square meter~42259 /m² +minus two thousand twelve kilo liters~-2012 kl +minus sixty six kilograms~-66 kg +two kilo watt hours~2 kwh +one point o o o o two eight cubic deci meters~1.000028 dm³ +seven point five peta bytes~7.5 pb +three hours~3 h +one milli volt~1 mv +two cubic meters~2 m³ +ninety grams~90 g +one hundred twenty four point three lumens~124.3 lm +four hundred forty milliliters~440 ml +thirty one thousand four hundred eighty square feet~31480 sq ft +one thousand six hundred hours~1600 h +thirty one thousand four hundred eighty square feet~31480 sq ft +two square miles~2 sq mi +zero point one nine square miles~0.19 sq mi +one thousand five hundred thirty one c c~1531 cc +three hundred micrometers~300 μm +sixty five thousand square kilometers~65000 km² +two miles per hour~2 mph +two hundred forty five miles per hour~245 mph +one hundred fifty c c~150 cc +sixty point two four zero zero kilograms~60.2400 kg +zero feet~0 ft +zero foot~zero foot +two feet~2 ft +twenty foot~20 foot +point two meters~.2 m +two square meters~2 m² +eighteen feet~18 ft +eighteen mega siemens~18 ms +eighteen ounces~18 oz +eighteen point five kilometers~18.5 km +eighteen point five two square kilometers~18.52 km² +eighteen point nine one square kilometers~18.91 km² +eighteen point one four percent~18.14 % +eighteen point one six percent~18.16 % +eighteen point one square kilometers~18.1 km² +eighteen point six percent~18.6 % +eighteen point two two kilometers~18.22 km +eighteen point zero kilometers~18.0 km +eighteen point zero percent~18.0 % +eighteen square kilometers~18 km² +eighteen thousand eight hundred giga watt hours~18800 gwh +eighteen thousand seven hundred hectares~18700 ha +eight hectares~8 ha +eight hundred eighty five astronomical units~885 au +eight hundred eighty hectares~880 ha +eight hundred eighty kilobytes~880 kb +eight hundred eighty kilometers~880 km +eight hundred eighty nine feet~889 ft +eight hundred eighty six kilometers~886 km +eight hundred eighty two megawatts~882 mw +eight hundred feet~800 ft +eight hundred fifty five square kilometers~855 km² +eight hundred fifty megahertz~850 mhz +eight hundred fifty meters~850 m +eight hundred fifty nanometers~850 nm +eight hundred fifty one meters~851 m +eight hundred fifty seven square kilometers~857 km² +eight hundred fifty three meters~853 m +eight hundred fifty three point six meters~853.6 m +eight hundred five point four six square kilometers~805.46 km² +eight hundred forty two point nine meters~842.9 m +eight hundred forty two square kilometers~842 km² +eight hundred gigabytes~800 gb +eight hundred horsepower~800 hp +eight hundred kilograms~800 kg +eight hundred kilo watt hours~800 kwh +eight hundred kilowatts~800 kw +eight hundred megahertz~800 mhz +eight hundred ninety four c c~894 cc +eight hundred ninety kilowatts~890 kw +eight hundred ninety millimeters~890 mm +eight hundred ninety two square kilometers~892 km² +eight hundred seventy horsepower~870 hp +eight hundred seventy meters~870 m +eight hundred sixty kilograms~860 kg +eight hundred sixty kilometers~860 km +eight hundred sixty miles~860 mi +eight hundred sixty six feet~866 ft +eight hundred ten hectares~810 ha +eight hundred ten kilohertz~810 khz +eight hundred thirty eight point two millimeters~838.2 mm +eight hundred thirty five kilometers~835 km +eight hundred thirty kilohertz~830 khz +eight hundred thirty megawatts~830 mw +eight hundred thirty nine kilometers~839 km +eight hundred thirty six meters~836 m +eight hundred twenty feet~820 ft +eight hundred twenty kilometers~820 km +eight hundred twenty meters~820 m +eight hundred twenty one point zero feet~821.0 ft +eight hundred two point eight nine kilometers~802.89 km +eight hundred volts~800 v +eight kilobits~8 kb +eight kilograms~8 kg +eight million two hundred thousand feet~8200000 ft +eight point eight kilometers~8.8 km +eight point eight meters~8.8 m +eight point eight miles~8.8 mi +eight point five centimeters~8.5 cm +eight point five five percent~8.55 % +eight point five megawatts~8.5 mw +eight point five meters~8.5 m +eight point five two percent~8.52 % +eight point four four percent~8.44 % diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_money.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_money.txt new file mode 100644 index 0000000000..24bc64b025 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_money.txt @@ -0,0 +1,48 @@ +two dollars~$2 +one cent~$0.01 +four united states dollars and sixty nine cents~$4.69 +seventy five dollars sixty three~$75.63 +twenty nine dollars fifty cents~$29.50 +eleven dollars and fifty one cents~$11.51 +nine hundred ninety three dollars and ninety two cents~$993.92 +four hundred sixty billion won~₩460 billion +thirty billion yen~¥30 billion +two point five billion dollars~$2.5 billion +forty five billion dollars~$45 billion +fifty million dollars~$50 million +fifty billion dollars~$50 billion +zero point two million dollars~$0.2 million +fifteen point two billion dollars~$15.2 billion +one point six nine billion yuan~1.69 billion yuan +one point four three six billion yuan~1.436 billion yuan +four million yuan~4 million yuan +one dollar~$1 +fifteen thousand dollars~$15000 +one dollars~one dollars +twenty dollar~20 dollar +twenty point five o six dollars~$20.506 +point five o six dollars~$.506 +eighteen dollars~$18 +eighteen million nine hundred twenty five thousand dollars~$18925000 +eighteen thousand eight hundred fifty four dollars~$18854 +eighteen thousand eight hundred one dollars~$18801 +eighteen thousand eight hundred seventy five dollars~$18875 +eighteen thousand eighty one dollars~$18081 +eighteen thousand fifty two dollars~$18052 +eighteen thousand five hundred forty two dollars~$18542 +eighteen thousand five hundred nineteen dollars~$18519 +eighteen thousand five hundred seventy dollars~$18570 +eighteen thousand five hundred seventy eight dollars~$18578 +eighteen thousand five hundred sixteen dollars~$18516 +eighteen thousand four hundred eighty two dollars~$18482 +eighteen thousand four hundred seventy eight dollars~$18478 +eighteen thousand four hundred sixty eight dollars~$18468 +eighteen thousand nine hundred three dollars~$18903 +eighteen thousand nine hundred twenty nine dollars~$18929 +eighteen thousand ninety five dollars~$18095 +eighteen thousand one hundred seventeen dollars~$18117 +eighteen thousand one hundred twenty eight dollars~$18128 +eighteen thousand one hundred twenty five dollars~$18125 +eighteen thousand one hundred twenty four dollars~$18124 +eighteen thousand one hundred twenty nine dollars~$18129 +eighteen thousand pounds~£18000 \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_ordinal.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_ordinal.txt new file mode 100644 index 0000000000..1ceaaf0acf --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_ordinal.txt @@ -0,0 +1,33 @@ +one hundredth~100th +twenty five thousand one hundred eleventh~25111th +second~2nd +zeroth~0th +first~1st +second~2nd +third~3rd +fourth~4th +eleventh~11th +twelfth~12th +thirteenth~13th +twenty first~21st +twenty third~23rd +one hundred eleventh~111th +one thousandth~1000th +one hundred twenty first~121st +second~2nd +tenth~10th +sixth~6th +third~3rd +nineteenth~19th +third~3rd +twelfth~12th +forty eighth~48th +seventy first~71st +third~3rd +forty second~42nd +seventeenth~17th +twentieth~20th +twenty first~21st +seventh~7th +second~2nd +fifth~5th \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_telephone.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_telephone.txt new file mode 100644 index 0000000000..c689d9e200 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_telephone.txt @@ -0,0 +1,4 @@ +one two three one two three five six seven eight~123-123-5678 +zero two three one two three five six seven eight~023-123-5678 +o two three one two three five six seven eight~023-123-5678 +oh two three one two three five six seven eight~023-123-5678 \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_time.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_time.txt new file mode 100644 index 0000000000..ecdb1d0d1c --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_time.txt @@ -0,0 +1,27 @@ +eight oclock g m t~08:00 gmt +seven a m e s t~07:00 a.m. est +two p m~02:00 p.m. +two thirty~02:30 +three o'clock~03:00 +quarter past one~01:15 +half past three~03:30 +eight fifty one~08:51 +eight fifty two~08:52 +eight forty~08:40 +eight nineteen~08:19 +eight o six~08:06 +eight thirty eight~08:38 +eight thirty two~08:32 +eight twenty nine~08:29 +eleven fifty five p m~11:55 p.m. +eleven fifty three p m~11:53 p.m. +eleven forty a m~11:40 a.m. +eleven forty five a m~11:45 a.m. +eleven forty p m~11:40 p.m. +eleven forty six a m~11:46 a.m. +eleven o six p m~11:06 p.m. +eleven thirteen a m~11:13 a.m. +half past twelve~12:30 +quarter past one~01:15 +quarter to one~12:45 +quarter to twelve~11:45 \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_whitelist.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_whitelist.txt new file mode 100644 index 0000000000..4f69671fd3 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_whitelist.txt @@ -0,0 +1,8 @@ +doctor dao~dr. dao +miss smith~ms. smith +misses smith~mrs. smith +mister dao~mr. dao +saint george~st. george +i like for example ice cream~i like e.g. ice cream +s and p five hundred~s&p 500 +seven eleven stores~7-eleven stores diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_word.txt b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_word.txt new file mode 100644 index 0000000000..4d12a42be8 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_inverse_text_normalization/test_cases_word.txt @@ -0,0 +1,50 @@ +~ +e s three~es3 +yahoo!~yahoo! +twenty!~20 ! +x ~x +—~— +aaa~aaa +aabach~aabach +aabenraa~aabenraa +aabye~aabye +aaccessed~aaccessed +aach~aach +aachen's~aachen's +aadri~aadri +aafia~aafia +aagaard~aagaard +aagadu~aagadu +aagard~aagard +aagathadi~aagathadi +aaghart's~aaghart's +aagnes~aagnes +aagomoni~aagomoni +aagon~aagon +aagoo~aagoo +aagot~aagot +aahar~aahar +aahh~aahh +aahperd~aahperd +aaibinterstate~aaibinterstate +aajab~aajab +aakasa~aakasa +aakervik~aakervik +aakirkeby~aakirkeby +aalam~aalam +aalbaek~aalbaek +aaldiu~aaldiu +aalem~aalem +a'ali~a'ali +aalilaassamthey~aalilaassamthey +aalin~aalin +aaliyan~aaliyan +aaliyan's~aaliyan's +aamadu~aamadu +aamara~aamara +aambala~aambala +aamera~aamera +aamer's~aamer's +aamina~aamina +aaminah~aaminah +aamjiwnaang~aamjiwnaang diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_boundary.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_boundary.txt new file mode 100644 index 0000000000..3b9d66f214 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_boundary.txt @@ -0,0 +1,7 @@ + 1~one +1~one +1, ~one , +1!!!!~one ! ! ! ! +(1)Hello~(1)Hello +!1~! one +1!hello~1!hello \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_cardinal.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_cardinal.txt new file mode 100644 index 0000000000..3d505b2ce4 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_cardinal.txt @@ -0,0 +1,18 @@ +1~one +2~two +-2~minus two +3~three +123~one hundred twenty three +13,000~thirteen thousand +13,000~thirteen thousand +123,123,000~one hundred twenty three million one hundred twenty three thousand +123,000,012~one hundred twenty three million twelve +t-0t25d12-f~t oh t two five d one two f +133-ABC~one hundred thirty three ABC +B2A23C~B two A two three C +25d08A~two five d oh eight A +C24~C two four +W2s~W two s +1-4-a-b-1-5~one four a b one five +b-c-1-5-b-s-b~b c one five b s b +1/f-4s~one f four s diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_date.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_date.txt new file mode 100644 index 0000000000..b5f95f4669 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_date.txt @@ -0,0 +1,32 @@ +july 25 2012~july twenty fifth twenty twelve +jul 25 2012~july twenty fifth twenty twelve +1980s~nineteen eighties +1980 s~nineteen eighties +25 july 2012~the twenty fifth of july twenty twelve +25 jul 2012~the twenty fifth of july twenty twelve +22 july 2012~the twenty second of july twenty twelve +15 january~the fifteenth of january +17 may 2010~the seventeenth of may twenty ten +january 1~january first +june 30~june thirtieth +2012~twenty twelve +july 1506~july fifteen o six +july 1876~july eighteen seventy six +1994~nineteen ninety four +1976~nineteen seventy six +june 20 2014~june twentieth twenty fourteen +1973~nineteen seventy three +1975~nineteen seventy five +1155~eleven fifty five +2006-08-05~august fifth two thousand six +august 23, 2002~august twenty third two thousand two +20 january~the twentieth of january +2016-11-03~november third twenty sixteen +2016/07/03~july third twenty sixteen +1980s~nineteen eighties +01.10.2010~january tenth twenty ten +10/06/2005~october sixth two thousand five +1910s~nineteen tens +25 sept.~the twenty fifth of september +1000~one thousand +SEPT. 15TH~september fifteenth diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_decimal.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_decimal.txt new file mode 100644 index 0000000000..1b839abc20 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_decimal.txt @@ -0,0 +1,8 @@ +.1665~point one six six five +2.0~two point zero +2.00~two point o o +2.050~two point o five o +100 million~one hundred million +0.1 billion~zero point one billion +.1 trillion~point one trillion +-0.1~minus zero point one diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_electronic.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_electronic.txt new file mode 100644 index 0000000000..1c5f5a71b1 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_electronic.txt @@ -0,0 +1,10 @@ +a.bc@gmail.com~a dot b c at g mail dot com +cdf@abc.edu~c d f at a b c dot e d u +abc@gmail.abc~a b c at g mail dot a b c +abc@abc.com~a b c at a b c dot com +asdf123@abc.com~a s d f one two three at a b c dot com +a1b2@abc.com~a one b two at a b c dot com +ab3.sdd.3@gmail.com~a b three dot s d d dot three at g mail dot com +email is abc1@gmail.com~email is a b c one at g mail dot com +test .b@gmail.com~test . b at g mail dot com +email is a-1.b3_c&-d@gma4i-l.com test~email is a dash one dot b three underscore c ampersand dash d at g m a four i dash l dot com test \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_fraction.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_fraction.txt new file mode 100644 index 0000000000..15d32caa9b --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_fraction.txt @@ -0,0 +1,11 @@ +1/2007~one two thousand seventh +12639/12640~twelve thousand six hundred thirty nine twelve thousand six hundred fortieths +2/4~two quarters +1/4~one quarter +31/32~thirty one thirty seconds +22/3~twenty two thirds +1/3~one third +142/1~one hundred forty two over one +1/2~one half +2 1/2~two and a half +1795 / 1805~one thousand seven hundred ninety five one thousand eight hundred fifths \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_measure.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_measure.txt new file mode 100644 index 0000000000..3279e6863c --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_measure.txt @@ -0,0 +1,15 @@ +12kg~twelve kilograms +12/kg~twelve per kilogram +12kg/kg~twelve kilograms per kilogram +1 mbps~one megabit per second +3 mbps~three megabits per second +3 cc/s~three c c per second +100 million kg~one hundred million kilograms +a 4-kilogram bag~a four kilogram bag +55-millimeters long~fifty five millimeters long +a 123.2-millimeters long~a one hundred twenty three point two millimeters long +covid-19.5~covid nineteen point five +covid-19~covid nineteen +a 4-kilogram bag~a four kilogram bag +7.2-millimeter bullet~seven point two millimeter bullet +4 1/2 lbs~four and a half pounds diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_money.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_money.txt new file mode 100644 index 0000000000..22f0b58acb --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_money.txt @@ -0,0 +1,38 @@ +$2~two dollars +₩460 billion~four hundred sixty billion won +¥30 billion~thirty billion yen +$45 billion~forty five billion dollars +$0.2 million~zero point two million dollars +$15.2 billion~fifteen point two billion dollars +1.436 billion yuan~one point four three six billion yuan +4 million yuan~four million yuan +$15000~fifteen thousand dollars +$1~one dollar +$20~twenty dollars +$20.506~twenty point five o six dollars +$.506~point five o six dollars +$18~eighteen dollars +$18925000~eighteen million nine hundred twenty five thousand dollars +$18,925,000~eighteen million nine hundred twenty five thousand dollars +$18854~eighteen thousand eight hundred fifty four dollars +$18801~eighteen thousand eight hundred one dollars +$18875~eighteen thousand eight hundred seventy five dollars +$18081~eighteen thousand eighty one dollars +$18052~eighteen thousand fifty two dollars +$18542~eighteen thousand five hundred forty two dollars +$18519~eighteen thousand five hundred nineteen dollars +$18570~eighteen thousand five hundred seventy dollars +$18578~eighteen thousand five hundred seventy eight dollars +$18516~eighteen thousand five hundred sixteen dollars +$18482~eighteen thousand four hundred eighty two dollars +$18478~eighteen thousand four hundred seventy eight dollars +$18468~eighteen thousand four hundred sixty eight dollars +$18903~eighteen thousand nine hundred three dollars +$18929~eighteen thousand nine hundred twenty nine dollars +$18095~eighteen thousand ninety five dollars +$18117~eighteen thousand one hundred seventeen dollars +$18128~eighteen thousand one hundred twenty eight dollars +$18125~eighteen thousand one hundred twenty five dollars +$18124~eighteen thousand one hundred twenty four dollars +$18129~eighteen thousand one hundred twenty nine dollars +£18000~eighteen thousand pounds \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_normalize_with_audio.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_normalize_with_audio.txt new file mode 100644 index 0000000000..436bff1920 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_normalize_with_audio.txt @@ -0,0 +1,107 @@ +~123 +one twenty three +one two three +one hundred twenty three +one hundred and twenty three +~$123 +one twenty three dollars +one two three dollars +one hundred twenty three dollars +one hundred and twenty three dollars +~$123.2 +one twenty three dollars and two cents +one two three dollars two cents +one hundred twenty three dollars and two cents +one hundred and twenty three point two dollars +~1.24 +one point two four +one two four +one point twenty four +one twenty four +~$1.21 +one dollar and twenty one cents +one dollar twenty one cents +one point two one dollars +one twenty one dollars +~£1.00 +one pound +one point o o pounds +~t-0t25d12-f +t zero t twenty five d one two f +t zero t twenty five d twelve f +t zero t two five d one two f +t zero t two five d twelve f +t oh t two five d one two f +t oh t twenty five d one two f +t oh t twenty five d twelve f +t oh t two five d twelve f +~133-A +one hundred thirty three A +~B2A23C +B two A twenty three C +~25d08A +twenty five d zero eight A +~C24 +C twenty four +~It seemed to her that the jacket Oswald wore was darker than Commission Exhibit No. 162. +It seemed to her that the jacket Oswald wore was darker than Commission Exhibit number one sixty two. +It seemed to her that the jacket Oswald wore was darker than Commission Exhibit number one hundred sixty two. +It seemed to her that the jacket Oswald wore was darker than Commission Exhibit number one hundred and sixty two. +It seemed to her that the jacket Oswald wore was darker than Commission Exhibit number one six two. +~"Father, let this cup pass." He prayed--was heard. What cup was it that passed away from him? Sure not the death-cup, now filled to the brim! There was no quailing in the awful word; He still was king of kings, of lords the lord:-- He feared lest, in the suffering waste and grim, His faith might grow too faint and sickly dim." +"Father, let this cup pass." He prayed -- was heard. What cup was it that passed away from him? Sure not the death-cup, now filled to the brim! There was no quailing in the awful word; He still was king of kings, of lords the lord: -- He feared lest, in the suffering waste and grim, His faith might grow too faint and sickly dim." +~1970-2010 +nineteen seventy to twenty ten +one thousand nine seventy to two thousand ten +one thousand nine hundred and seventy to two thousand ten +one thousand nine hundred and seventy to twenty ten +nineteen seventy twenty ten +one thousand nine seventy two thousand ten +one thousand nine hundred and seventy two thousand ten +one thousand nine hundred and seventy twenty ten +from nineteen seventy to twenty ten +from one thousand nine seventy to two thousand ten +from one thousand nine hundred and seventy to two thousand ten +from one thousand nine hundred and seventy to twenty ten +from nineteen seventy twenty ten +from one thousand nine seventy two thousand ten +from one thousand nine hundred and seventy two thousand ten +from one thousand nine hundred and seventy twenty ten +~W26s +W twenty sixes +W twenty six s +~5-3-Ws +five three W s +five three Ws +~401-ks +four oh one kays +four oh one k s +four hundred one k s +~The box was 25 x 7 m. +The box was twenty five by seven meters. +The box was two five by seven meters. +The box was two five by seven m. +The box was twenty five by seven m. +~4567 +four thousand five hundred sixty seven +four five six seven +four thousand five sixty seven +forty five sixty seven +four thousand five hundred and sixty seven +~This example number 15,000 can be a very long one, and can fail to produce valid normalization for a such an easy number like 10,125 or dollar value $5349.01, and can fail to terminate, and can fail to terminate, and can fail to terminate, and can fail to terminate, and can fail to terminate, 452. +This example number fifteen thousand can be a very long one, and can fail to produce valid normalization for a such an easy number like ten thousand one hundred twenty five or dollar value five thousand three hundred forty nine dollars and one cent, and can fail to terminate, and can fail to terminate, and can fail to terminate, and can fail to terminate, and can fail to terminate, four five two. +~$1.01 +one dollar and one cent +one dollar one cent +one dollar and zero one cents +~$17.31 +seventeen dollars and thirty one cent +seventeen dollars and thirty one cents +seventeen point three one dollars +~25.] +two five.] +twenty five.] +~Francis I--test +Francis the first -- test +Francis one -- test +Francis first -- test \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_ordinal.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_ordinal.txt new file mode 100644 index 0000000000..dce8934082 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_ordinal.txt @@ -0,0 +1,6 @@ +0th~zeroth +11th~eleventh +23rd~twenty third +1000th~one thousandth +1,000th~one thousandth +1000000th~one millionth \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_telephone.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_telephone.txt new file mode 100644 index 0000000000..2e586f60fe --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_telephone.txt @@ -0,0 +1,9 @@ ++1 123-123-5678~one, one two three, one two three, five six seven eight +123-123-5678~one two three, one two three, five six seven eight ++1-123-123-5678~one, one two three, one two three, five six seven eight ++1 (123)-123-5678~one, one two three, one two three, five six seven eight ++1(123) 123-5678~one, one two three, one two three, five six seven eight +(123)-123-5678~one two three, one two three, five six seven eight +(123) 123-5678~one two three, one two three, five six seven eight +123-123-5678-1111~one two three, one two three, five six seven eight, one one one one +123-123-5678-1~one two three, one two three, five six seven eight, one \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_time.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_time.txt new file mode 100644 index 0000000000..5388be9d4d --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_time.txt @@ -0,0 +1,14 @@ +01:00~one o'clock +01:00 am~one a m +01:00 a.m.~one a m +01:00 a.m~one a m +1 a.m.~one a m +23:00~twenty three o'clock +1:00 a.m.~one a m +1:00a.m.~one a m +1.59 p.m.~one fifty nine p m +01:59 p.m.~one fifty nine p m +1:59 p.m.~one fifty nine p m +1:59 p.m. est~one fifty nine p m e s t +1:59 p.m.est~one fifty nine p m e s t +1:59 p.m. e.s.t~one fifty nine p m e s t \ No newline at end of file diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_whitelist.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_whitelist.txt new file mode 100644 index 0000000000..7eaa6376a1 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_whitelist.txt @@ -0,0 +1,3 @@ +dr. Evil~doctor Evil +mrs. Norris~misses Norris +Dr. Evil~Dr. Evil diff --git a/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_word.txt b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_word.txt new file mode 100644 index 0000000000..e10df68195 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/data_text_normalization/test_cases_word.txt @@ -0,0 +1,23 @@ +~ +es3~e s three +x ~x +X!~X! +2!~two ! +—~— +aabach~aabach +aabenraa~aabenraa +aabye~aabye +aaccessed~aaccessed +aach~aach +aachen's~aachen's +aadri~aadri +aafia~aafia +aagaard~aagaard +aagadu~aagadu +aagard~aagard +aagathadi~aagathadi +aaghart's~aaghart's +aalem~aalem +a'ali~a'ali +aaliyan's~aaliyan's +mother-in-law~mother-in-law diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_boundary.py b/models/tts/NeMo/tests/nemo_text_processing/test_boundary.py new file mode 100644 index 0000000000..9074230050 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_boundary.py @@ -0,0 +1,39 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestBoundary: + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_boundary.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize( + test_input, n_tagged=100, punct_pre_process=False, punct_post_process=False + ) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_cardinal.py b/models/tts/NeMo/tests/nemo_text_processing/test_cardinal.py new file mode 100644 index 0000000000..7c19ea3d42 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_cardinal.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestCardinal: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_cardinal.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_cardinal.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_date.py b/models/tts/NeMo/tests/nemo_text_processing/test_date.py new file mode 100644 index 0000000000..42fce25b4d --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_date.py @@ -0,0 +1,64 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestDate: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_date.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_date.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm_uncased(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic + + normalizer_uppercased = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + cases_uppercased = {"Aug. 8": "august eighth", "8 Aug.": "the eighth of august", "aug. 8": "august eighth"} + + @parameterized.expand(cases_uppercased.items()) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm_cased(self, test_input, expected): + pred = self.normalizer_uppercased.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_decimal.py b/models/tts/NeMo/tests/nemo_text_processing/test_decimal.py new file mode 100644 index 0000000000..267fcceb5e --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_decimal.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestDecimal: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_decimal.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case="cased") if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_decimal.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_electronic.py b/models/tts/NeMo/tests/nemo_text_processing/test_electronic.py new file mode 100644 index 0000000000..a96ecc696a --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_electronic.py @@ -0,0 +1,51 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestElectronic: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_electronic.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case="cased") if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_electronic.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize( + test_input, n_tagged=100, punct_post_process=False + ) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_fraction.py b/models/tts/NeMo/tests/nemo_text_processing/test_fraction.py new file mode 100644 index 0000000000..24006c43cc --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_fraction.py @@ -0,0 +1,38 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestFraction: + normalizer = Normalizer(input_case="cased") if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_fraction.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_measure.py b/models/tts/NeMo/tests/nemo_text_processing/test_measure.py new file mode 100644 index 0000000000..12f4137f67 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_measure.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestMeasure: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_measure.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case="cased") if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_measure.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_money.py b/models/tts/NeMo/tests/nemo_text_processing/test_money.py new file mode 100644 index 0000000000..e132866b6d --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_money.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestMoney: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_money.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_money.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_normalization_with_audio.py b/models/tts/NeMo/tests/nemo_text_processing/test_normalization_with_audio.py new file mode 100644 index 0000000000..00345837b3 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_normalization_with_audio.py @@ -0,0 +1,33 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, get_test_cases_multiple + + +class TestNormalizeWithAudio: + + normalizer = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(get_test_cases_multiple('data_text_normalization/test_cases_normalize_with_audio.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, n_tagged=700) + assert len(set(pred).intersection(set(expected))) == len(expected), f'pred: {pred}' diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_ordinal.py b/models/tts/NeMo/tests/nemo_text_processing/test_ordinal.py new file mode 100644 index 0000000000..47ae51e830 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_ordinal.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestOrdinal: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_ordinal.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_ordinal.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_sparrowhawk_inverse_text_normalization.sh b/models/tts/NeMo/tests/nemo_text_processing/test_sparrowhawk_inverse_text_normalization.sh new file mode 100644 index 0000000000..268f88b057 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_sparrowhawk_inverse_text_normalization.sh @@ -0,0 +1,64 @@ +#! /bin/sh + +PROJECT_DIR=/workspace/tests + +runtest () { + input=$1 + cd /workspace/sparrowhawk/documentation/grammars + + # read test file + while read testcase; do + IFS='~' read spoken written <<< $testcase + denorm_pred=$(echo $spoken | normalizer_main --config=sparrowhawk_configuration.ascii_proto 2>&1 | tail -n 1) + + # trim white space + written="$(echo -e "${written}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')" + denorm_pred="$(echo -e "${denorm_pred}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')" + + # input expected actual + assertEquals "$spoken" "$written" "$denorm_pred" + done < "$input" +} + +testITNCardinal() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_cardinal.txt + runtest $input +} + +testITNDate() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_date.txt + runtest $input +} + +testITNDecimal() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_decimal.txt + runtest $input +} + +testITNOrdinal() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_ordinal.txt + runtest $input +} + +testITNTime() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_time.txt + runtest $input +} + +testITNMeasure() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_measure.txt + runtest $input +} + +testITNMoney() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_money.txt + runtest $input +} + +testITNWhitelist() { + input=$PROJECT_DIR/data_inverse_text_normalization/test_cases_whitelist.txt + runtest $input +} + +# Load shUnit2 +. $PROJECT_DIR/../shunit2/shunit2 diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_sparrowhawk_normalization.sh b/models/tts/NeMo/tests/nemo_text_processing/test_sparrowhawk_normalization.sh new file mode 100644 index 0000000000..c06dc05aff --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_sparrowhawk_normalization.sh @@ -0,0 +1,64 @@ +#! /bin/sh + +PROJECT_DIR=/workspace/tests + +runtest () { + input=$1 + cd /workspace/sparrowhawk/documentation/grammars + + # read test file + while read testcase; do + IFS='~' read written spoken <<< $testcase + denorm_pred=$(echo $written | normalizer_main --config=sparrowhawk_configuration.ascii_proto 2>&1 | tail -n 1) + + # trim white space + spoken="$(echo -e "${spoken}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')" + denorm_pred="$(echo -e "${denorm_pred}" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')" + + # input expected actual + assertEquals "$written" "$spoken" "$denorm_pred" + done < "$input" +} + +testTNCardinal() { + input=$PROJECT_DIR/data_text_normalization/test_cases_cardinal.txt + runtest $input +} + +testTNDate() { + input=$PROJECT_DIR/data_text_normalization/test_cases_date.txt + runtest $input +} + +testTNDecimal() { + input=$PROJECT_DIR/data_text_normalization/test_cases_decimal.txt + runtest $input +} + +testTNOrdinal() { + input=$PROJECT_DIR/data_text_normalization/test_cases_ordinal.txt + runtest $input +} + +testTNTime() { + input=$PROJECT_DIR/data_text_normalization/test_cases_time.txt + runtest $input +} + +testTNMeasure() { + input=$PROJECT_DIR/data_text_normalization/test_cases_measure.txt + runtest $input +} + +testTNMoney() { + input=$PROJECT_DIR/data_text_normalization/test_cases_money.txt + runtest $input +} + +testTNWhitelist() { + input=$PROJECT_DIR/data_text_normalization/test_cases_whitelist.txt + runtest $input +} + +# Load shUnit2 +. $PROJECT_DIR/../shunit2/shunit2 diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_telephone.py b/models/tts/NeMo/tests/nemo_text_processing/test_telephone.py new file mode 100644 index 0000000000..5ea68d6f0c --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_telephone.py @@ -0,0 +1,50 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestTelephone: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_telephone.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_telephone.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_time.py b/models/tts/NeMo/tests/nemo_text_processing/test_time.py new file mode 100644 index 0000000000..84e05a83b7 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_time.py @@ -0,0 +1,49 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestTime: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_time.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_time.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_whitelist.py b/models/tts/NeMo/tests/nemo_text_processing/test_whitelist.py new file mode 100644 index 0000000000..8b3daf2b67 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_whitelist.py @@ -0,0 +1,65 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestWhitelist: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_whitelist.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='lower_cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_whitelist.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic + + normalizer_uppercased = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + cases_uppercased = {"Dr. Evil": "doctor Evil", "No. 4": "number four", "dr. Evil": "dr. Evil", "no. 4": "no. four"} + + @parameterized.expand(cases_uppercased.items()) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm_cased(self, test_input, expected): + pred = self.normalizer_uppercased.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize(test_input, n_tagged=100) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/test_word.py b/models/tts/NeMo/tests/nemo_text_processing/test_word.py new file mode 100644 index 0000000000..ae60100c41 --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/test_word.py @@ -0,0 +1,52 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 pytest +from nemo_text_processing.inverse_text_normalization.inverse_normalize import InverseNormalizer +from nemo_text_processing.text_normalization.normalize import Normalizer +from nemo_text_processing.text_normalization.normalize_with_audio import NormalizerWithAudio +from parameterized import parameterized +from utils import PYNINI_AVAILABLE, parse_test_case_file + + +class TestWord: + inverse_normalizer = InverseNormalizer() if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_inverse_text_normalization/test_cases_word.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_denorm(self, test_input, expected): + pred = self.inverse_normalizer.inverse_normalize(test_input, verbose=False) + assert pred == expected + + normalizer = Normalizer(input_case='cased') if PYNINI_AVAILABLE else None + normalizer_with_audio = NormalizerWithAudio(input_case='cased') if PYNINI_AVAILABLE else None + + @parameterized.expand(parse_test_case_file('data_text_normalization/test_cases_word.txt')) + @pytest.mark.skipif( + not PYNINI_AVAILABLE, reason="`pynini` not installed, please install via nemo_text_processing/setup.sh" + ) + @pytest.mark.run_only_on('CPU') + @pytest.mark.unit + def test_norm(self, test_input, expected): + pred = self.normalizer.normalize(test_input, verbose=False) + assert pred == expected + pred_non_deterministic = self.normalizer_with_audio.normalize( + test_input, n_tagged=100, punct_post_process=False + ) + assert expected in pred_non_deterministic diff --git a/models/tts/NeMo/tests/nemo_text_processing/utils.py b/models/tts/NeMo/tests/nemo_text_processing/utils.py new file mode 100644 index 0000000000..d2bb3554ef --- /dev/null +++ b/models/tts/NeMo/tests/nemo_text_processing/utils.py @@ -0,0 +1,54 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 os + +try: + import pynini + + PYNINI_AVAILABLE = True +except (ImportError, ModuleNotFoundError): + PYNINI_AVAILABLE = False + + +def parse_test_case_file(file_name: str): + """ + Prepares tests pairs for ITN and TN tests + """ + test_pairs = [] + with open(os.path.dirname(os.path.abspath(__file__)) + os.path.sep + file_name, 'r') as f: + for line in f: + spoken, written = line.split('~') + test_pairs.append((spoken, written.strip("\n"))) + return test_pairs + + +def get_test_cases_multiple(file_name: str = 'data_text_normalization/test_cases_normalize_with_audio.txt'): + """ + Prepares tests pairs for audio based TN tests + """ + test_pairs = [] + with open(os.path.dirname(os.path.abspath(__file__)) + os.path.sep + file_name, 'r') as f: + written = None + normalized_options = [] + for line in f: + if line.startswith('~'): + if written: + test_pairs.append((written, normalized_options)) + normalized_options = [] + written = line.strip().replace('~', '') + else: + normalized_options.append(line.strip()) + test_pairs.append((written, normalized_options)) + return test_pairs diff --git a/models/tts/NeMo/tests/test_data_dir.py b/models/tts/NeMo/tests/test_data_dir.py new file mode 100644 index 0000000000..931e02b73c --- /dev/null +++ b/models/tts/NeMo/tests/test_data_dir.py @@ -0,0 +1,27 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + + +from os.path import exists, join + +import pytest + + +class TestDataDir: + @pytest.mark.unit + def test_test_data_dir(self, test_data_dir): + """" Just a dummy tests showing how to use the test_data_dir fixture. """ + # test_data_dir contains the absolute path to nemo -> tests/.data + assert exists(test_data_dir) + assert exists(join(test_data_dir, "test_data.tar.gz")) diff --git a/models/tts/NeMo/tools/ctc_segmentation/README.md b/models/tts/NeMo/tools/ctc_segmentation/README.md new file mode 100644 index 0000000000..bcfe5e7bd9 --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/README.md @@ -0,0 +1,31 @@ +Dataset creation tool based on CTC-Segmentation +----------------------------------------------- + +This tool provides functionality to align long audio files and the corresponding transcripts into shorter fragments +that are suitable for an Automatic Speech Recognition (ASR) model training. + +More details could be found in [this tutorial](https://github.com/NVIDIA/NeMo/blob/main/tutorials/tools/CTC_Segmentation_Tutorial.ipynb). + +The tool is based on the [CTC Segmentation](https://github.com/lumaku/ctc-segmentation): +**CTC-Segmentation of Large Corpora for German End-to-end Speech Recognition** +https://doi.org/10.1007/978-3-030-60276-5_27 or pre-print https://arxiv.org/abs/2007.09127 + +``` +@InProceedings{ctcsegmentation, +author="K{\"u}rzinger, Ludwig +and Winkelbauer, Dominik +and Li, Lujun +and Watzel, Tobias +and Rigoll, Gerhard", +editor="Karpov, Alexey +and Potapova, Rodmonga", +title="CTC-Segmentation of Large Corpora for German End-to-End Speech Recognition", +booktitle="Speech and Computer", +year="2020", +publisher="Springer International Publishing", +address="Cham", +pages="267--278", +abstract="Recent end-to-end Automatic Speech Recognition (ASR) systems demonstrated the ability to outperform conventional hybrid DNN/HMM ASR. Aside from architectural improvements in those systems, those models grew in terms of depth, parameters and model capacity. However, these models also require more training data to achieve comparable performance.", +isbn="978-3-030-60276-5" +} +``` diff --git a/models/tts/NeMo/tools/ctc_segmentation/requirements.txt b/models/tts/NeMo/tools/ctc_segmentation/requirements.txt new file mode 100644 index 0000000000..b5ef916816 --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/requirements.txt @@ -0,0 +1,2 @@ +ctc_segmentation==1.1.0 +num2words \ No newline at end of file diff --git a/models/tts/NeMo/tools/ctc_segmentation/run_sample.sh b/models/tts/NeMo/tools/ctc_segmentation/run_sample.sh new file mode 100644 index 0000000000..d865a57d0c --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/run_sample.sh @@ -0,0 +1,124 @@ +#!/bin/bash + +# default values for optional arguments +MIN_SCORE=-5 +CUT_PREFIX=0 +SCRIPTS_DIR="scripts" +OFFSET=0 +LANGUAGE='eng' # 'eng', 'ru', 'other' +MIN_SEGMENT_LEN=20 +MAX_SEGMENT_LEN=100 +ADDITIONAL_SPLIT_SYMBOLS='' +AUDIO_FORMAT='.mp3' +USE_NEMO_NORMALIZATION='True' + +for ARG in "$@" +do + key=$(echo $ARG | cut -f1 -d=) + value=$(echo $ARG | cut -f2 -d=) + + if [[ $key == *"--"* ]]; then + v="${key/--/}" + declare $v="${value}" + fi +done + +echo "MODEL_NAME_OR_PATH = $MODEL_NAME_OR_PATH" +echo "DATA_DIR = $DATA_DIR" +echo "OUTPUT_DIR = $OUTPUT_DIR" +echo "MIN_SCORE = $MIN_SCORE" +echo "CUT_PREFIX = $CUT_PREFIX" +echo "SCRIPTS_DIR = $SCRIPTS_DIR" +echo "OFFSET = $OFFSET" +echo "LANGUAGE = $LANGUAGE" +echo "MIN_SEGMENT_LEN = $MIN_SEGMENT_LEN" +echo "MAX_SEGMENT_LEN = $MAX_SEGMENT_LEN" +echo "ADDITIONAL_SPLIT_SYMBOLS = $ADDITIONAL_SPLIT_SYMBOLS" +echo "AUDIO_FORMAT = $AUDIO_FORMAT" +echo "USE_NEMO_NORMALIZATION = $USE_NEMO_NORMALIZATION" + +if [[ -z $MODEL_NAME_OR_PATH ]] || [[ -z $DATA_DIR ]] || [[ -z $OUTPUT_DIR ]]; then + echo "Usage: $(basename "$0") + --MODEL_NAME_OR_PATH=[model_name_or_path] + --DATA_DIR=[data_dir] + --OUTPUT_DIR=[output_dir] + --LANGUAGE=[language (Optional)] + --OFFSET=[offset value (Optional)] + --CUT_PREFIX=[cut prefix in sec (Optional)] + --SCRIPTS_DIR=[scripts_dir_path (Optional)] + --MIN_SEGMENT_LEN=[min number of characters of the text segment for alignment (Optional)] + --MAX_SEGMENT_LEN=[max number of characters of the text segment for alignment (Optional)] + --ADDITIONAL_SPLIT_SYMBOLS=[Additional symbols to use for + sentence split if eos sentence split resulted in sequence longer than --max_length. + Use '|' as a separator between symbols, for example: ';|:|' (Optional)] + --AUDIO_FORMAT=[choose from ['.mp3', '.wav'], input audio files format + --USE_NEMO_NORMALIZATION Set to 'True' to use NeMo Normalization tool to convert + numbers from written to spoken format. By default num2words package will be used. (Optional)" + exit 1 +fi + +NEMO_NORMALIZATION="" + if [[ ${USE_NEMO_NORMALIZATION,,} == "true" ]]; then + NEMO_NORMALIZATION="--use_nemo_normalization " + fi + +# STEP #1 +# Prepare text and audio data for segmentation +python $SCRIPTS_DIR/prepare_data.py \ +--in_text=$DATA_DIR/text \ +--audio_dir=$DATA_DIR/audio \ +--audio_format=$AUDIO_FORMAT \ +--output_dir=$OUTPUT_DIR/processed/ \ +--language=$LANGUAGE \ +--cut_prefix=$CUT_PREFIX \ +--model=$MODEL_NAME_OR_PATH \ +--min_length=$MIN_SEGMENT_LEN \ +--max_length=$MAX_SEGMENT_LEN \ +--additional_split_symbols=$ADDITIONAL_SPLIT_SYMBOLS $NEMO_NORMALIZATION || exit + +# STEP #2 +# Run CTC-segmenatation +# one might want to perform alignment with various window sizes +# note if the alignment with the initial window size isn't found, the window size will be double to re-attempt +# alignment +for WINDOW in 8000 12000 +do + python $SCRIPTS_DIR/run_ctc_segmentation.py \ + --output_dir=$OUTPUT_DIR \ + --data=$OUTPUT_DIR/processed/ \ + --model=$MODEL_NAME_OR_PATH \ + --window_len $WINDOW || exit +done + +# STEP #3 (Optional) +# Verify aligned segments only if multiple WINDOWs used in the Step #2) +python $SCRIPTS_DIR/verify_segments.py \ +--base_dir=$OUTPUT_DIR || exit + +# STEP #4 +# Cut the original audio files based on the alignments +# (use --alignment=$OUTPUT_DIR/segments if only 1 WINDOW size was used in the Step #2) +# Three manifests and corresponding clips folders will be created: +# - high scored clips +# - low scored clips +# - deleted segments +python $SCRIPTS_DIR/cut_audio.py \ +--output_dir=$OUTPUT_DIR \ +--model=$MODEL_NAME_OR_PATH \ +--alignment=$OUTPUT_DIR/verified_segments \ +--threshold=$MIN_SCORE \ +--offset=$OFFSET || exit + +# STEP #5 (Optional) +# If multiple audio files were segmented in the step #2, this step will aggregate manifests with high scored segments +# for all audio files into all_manifest.json +# Also a separate manifest with samples from across all high scored segments will be credated if --num_samples > 0 +# --num_samples samples will be taken from the beginning, end and the middle of the each audio file high score manifest +# and will be stored at sample_manifest.json +python $SCRIPTS_DIR/process_manifests.py \ +--output_dir=$OUTPUT_DIR \ +--manifests_dir=$OUTPUT_DIR/manifests/ \ +--num_samples 0 + +exit + diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/cut_audio.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/cut_audio.py new file mode 100644 index 0000000000..50bd931be9 --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/cut_audio.py @@ -0,0 +1,313 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import json +import os +import time +from pathlib import Path +from typing import List + +import torch +from scipy.io import wavfile + +from nemo.collections import asr as nemo_asr +from nemo.collections.asr.metrics.wer import WER, word_error_rate + +parser = argparse.ArgumentParser(description="Cut audio on the segments based on segments") +parser.add_argument("--output_dir", default='output', type=str, help='Path to output directory') +parser.add_argument( + "--alignment", + type=str, + required=True, + help='Path to a data directory with alignments or a single .txt file with timestamps - result of the ctc-segmentation', +) +parser.add_argument("--threshold", type=float, default=-5, help='Minimum score value accepted') +parser.add_argument( + '--model', + type=str, + default='QuartzNet15x5Base-En', + help='Path to model checkpoint or pre-trained CTC-based ASR model name', +) +parser.add_argument('--offset', type=int, default=0, help='Offset in seconds') +parser.add_argument("--batch_size", type=int, default=64, help='Batch size for inference') + + +def add_transcript_to_manifest( + manifest_original: str, manifest_updated: str, asr_model: nemo_asr.models.EncDecCTCModel, batch_size: int +) -> None: + """ + Adds transcripts generated by the asr_model to the manifest_original. + + Args: + manifest_original: path to the manifest + manifest_updated: path to the updated manifest with transcript included + asr_model: CTC-based ASR model, for example, QuartzNet15x5Base-En + batch_size: Batch size for asr_model inference + """ + transcripts = get_transcript(manifest_original, asr_model, batch_size) + with open(manifest_original, 'r', encoding='utf8') as f: + with open(manifest_updated, 'w', encoding='utf8') as f_updated: + for i, line in enumerate(f): + info = json.loads(line) + info['transcript'] = transcripts[i].strip() + info['WER'] = round(word_error_rate([info['transcript']], [info['text']]) * 100, 2) + info['CER'] = round(word_error_rate([info['transcript']], [info['text']], use_cer=True) * 100, 2) + json.dump(info, f_updated, ensure_ascii=False) + f_updated.write('\n') + + +def get_transcript(manifest_path: str, asr_model: nemo_asr.models.EncDecCTCModel, batch_size: int) -> List[str]: + """ + Returns transcripts for audio segments in the batch + + Args: + manifest_path: path to the manifest for inference + asr_model: CTC-based ASR model, for example, QuartzNet15x5Base-En + batch_size: batch size + + Returns: hypotheses: transcripts for the audio segments + """ + # batch inference + try: + from torch.cuda.amp import autocast + except ImportError: + from contextlib import contextmanager + + @contextmanager + def autocast(enabled=None): + yield + + torch.set_grad_enabled(False) + asr_model.setup_test_data( + test_data_config={ + 'sample_rate': 16000, + 'manifest_filepath': manifest_path, + 'labels': asr_model.decoder.vocabulary, + 'batch_size': batch_size, + 'normalize_transcripts': False, + } + ) + asr_model.eval() + wer = WER(vocabulary=asr_model.decoder.vocabulary) + hypotheses = [] + for test_batch in asr_model.test_dataloader(): + if torch.cuda.is_available(): + test_batch = [x.cuda() for x in test_batch] + with autocast(): + log_probs, encoded_len, greedy_predictions = asr_model( + input_signal=test_batch[0], input_signal_length=test_batch[1] + ) + hypotheses += wer.ctc_decoder_predictions_tensor(greedy_predictions) + del test_batch + torch.cuda.empty_cache() + return hypotheses + + +def process_alignment(alignment_file: str, args): + """ Cut original audio file into audio segments based on alignment_file + + Args: + alignment_file: path to the file with segmented text and corresponding time stamps. + The first line of the file contains the path to the original audio file + args: main script args + """ + if not os.path.exists(alignment_file): + raise ValueError(f'{alignment_file} not found') + + # read the segments, note the first line contains the path to the original audio + segments = [] + ref_text_processed = [] + ref_text_no_preprocessing = [] + with open(alignment_file, 'r') as f: + for line in f: + line = line.split('|') + # read audio file name from the first line + if len(line) == 1: + audio_file = line[0].strip() + continue + ref_text_processed.append(line[1].strip()) + ref_text_no_preprocessing.append(line[2].strip()) + line = line[0].split() + segments.append((float(line[0]) + args.offset / 1000, float(line[1]) + args.offset / 1000, float(line[2]))) + + # cut the audio into segments + # create manifest in /tmp directory first, then populate transcript values with the batch inference results + # and save the final manifests at output_dir + sampling_rate, signal = wavfile.read(audio_file) + original_duration = len(signal) / sampling_rate + print(f'Cutting {audio_file} based on {alignment_file}') + print(f'Original duration: {round(original_duration)}s or ~{round(original_duration / 60)}min') + + # create directories to store high score .wav audio fragments, low scored once, and deleted + args.output_dir = os.path.abspath(args.output_dir) + os.makedirs(args.output_dir, exist_ok=True) + fragments_dir = os.path.join(args.output_dir, "high_score_clips") + del_fragments = os.path.join(args.output_dir, 'deleted_clips') + low_score_segments_dir = os.path.join(args.output_dir, "low_score_clips") + + os.makedirs(fragments_dir, exist_ok=True) + os.makedirs(low_score_segments_dir, exist_ok=True) + os.makedirs(del_fragments, exist_ok=True) + + base_name = os.path.basename(alignment_file).replace('_segments.txt', '') + high_score_manifest = f'{base_name}_high_score_manifest.json' + low_score_manifest = f'{base_name}_low_score_manifest.json' + del_manifest = f'{base_name}_del_manifest.json' + manifests_dir = os.path.join(args.output_dir, 'manifests') + os.makedirs(manifests_dir, exist_ok=True) + tmp_dir = '/tmp' + + low_score_dur = 0 + high_score_dur = 0 + with open(os.path.join(tmp_dir, high_score_manifest), 'w', encoding='utf8') as f: + with open(os.path.join(tmp_dir, low_score_manifest), 'w', encoding='utf8') as low_score_f: + for i, (st, end, score) in enumerate(segments): + segment = signal[round(st * sampling_rate) : round(end * sampling_rate)] + duration = len(segment) / sampling_rate + if duration > 0: + text_processed = ref_text_processed[i].strip() + text_no_preprocessing = ref_text_no_preprocessing[i].strip() + if score > args.threshold: + high_score_dur += duration + audio_filepath = os.path.join(fragments_dir, f'{base_name}_{i:04}.wav') + file_to_write = f + else: + low_score_dur += duration + audio_filepath = os.path.join(low_score_segments_dir, f'{base_name}_{i:04}.wav') + file_to_write = low_score_f + + wavfile.write(audio_filepath, sampling_rate, segment) + + transcript = 'n/a' # asr_model.transcribe(paths2audio_files=[audio_filepath], batch_size=1)[0] + info = { + 'audio_filepath': audio_filepath, + 'duration': duration, + 'text': text_processed, + 'text_no_preprocessing': text_no_preprocessing, + 'score': round(score, 2), + 'transcript': transcript.strip(), + } + json.dump(info, file_to_write, ensure_ascii=False) + file_to_write.write('\n') + + add_transcript_to_manifest( + os.path.join(tmp_dir, high_score_manifest), + os.path.join(manifests_dir, high_score_manifest), + asr_model, + args.batch_size, + ) + add_transcript_to_manifest( + os.path.join(tmp_dir, low_score_manifest), + os.path.join(manifests_dir, low_score_manifest), + asr_model, + args.batch_size, + ) + print(f'High score files duration: {round(high_score_dur)}s or ~{round(high_score_dur/60)}min at {manifests_dir}') + print( + f'Low score files duration: {round(low_score_dur)}s or ~{round(low_score_dur/60)}min saved at {manifests_dir}' + ) + + # save deleted segments along with manifest + deleted = [] + del_duration = 0 + begin = 0 + i = 0 + with open(os.path.join(manifests_dir, del_manifest), 'w', encoding='utf8') as f: + for i, (st, end, _) in enumerate(segments): + if st - begin > 0.01: + segment = signal[int(begin * sampling_rate) : int(st * sampling_rate)] + audio_filepath = os.path.join(del_fragments, f'del_{base_name}_{i:04}.wav') + wavfile.write(audio_filepath, sampling_rate, segment) + duration = len(segment) / sampling_rate + del_duration += duration + deleted.append((begin, st)) + transcript = 'n/a' # asr_model.transcribe(paths2audio_files=[audio_filepath], batch_size=1)[0] + info = {'audio_filepath': audio_filepath, 'duration': duration, 'text': transcript} + json.dump(info, f, ensure_ascii=False) + f.write('\n') + begin = end + + segment = signal[int(begin * sampling_rate) :] + audio_filepath = os.path.join(del_fragments, f'del_{base_name}_{i+1:04}.wav') + wavfile.write(audio_filepath, sampling_rate, segment) + duration = len(segment) / sampling_rate + del_duration += duration + deleted.append((begin, original_duration)) + + info = {'audio_filepath': audio_filepath, 'duration': duration, 'text': 'n/a'} + json.dump(info, f) + f.write('\n') + + print(f'Saved DEL files duration: {round(del_duration)}s or ~ {round(del_duration/60)}min at {del_fragments}') + missing_audio = original_duration - high_score_dur - del_duration - low_score_dur + if missing_audio > 15: + raise ValueError(f'{round(missing_audio)}s or ~ {round(missing_audio/60)}min is missing. Check the args') + + stats = ( + args.output_dir, + base_name, + round(original_duration), + round(high_score_dur), + round(low_score_dur), + round(del_duration), + ) + return stats + + +if __name__ == '__main__': + start_time = time.time() + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + if os.path.exists(args.model): + asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model, strict=False) + elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names(): + asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model, strict=False) + else: + raise ValueError( + f'Provide path to the pretrained checkpoint or choose from {nemo_asr.models.EncDecCTCModel.list_available_models()}' + ) + + alignment_files = Path(args.alignment) + if os.path.isdir(args.alignment): + alignment_files = alignment_files.glob("*.txt") + else: + alignment_files = [Path(alignment_files)] + + stats_file = os.path.join(args.output_dir, 'stats.tsv') + with open(stats_file, 'w') as f: + f.write('Folder\tSegment\tOriginal dur (s)\tHigh quality dur (s)\tLow quality dur (s)\tDeleted dur (s)\n') + + high_score_dur = 0 + low_score_dur = 0 + del_duration = 0 + + for alignment_file in alignment_files: + stats = process_alignment(alignment_file, args) + high_score_dur += stats[-3] + low_score_dur += stats[-2] + del_duration += stats[-1] + stats = '\t'.join([str(t) for t in stats]) + '\n' + f.write(stats) + + f.write(f'Total\t\t{round(high_score_dur)}\t{round(low_score_dur)}\t{del_duration}') + + total_time = time.time() - start_time + print(f'High score segments duration: {round(high_score_dur)}') + print(f'Low score segments duration: {round(low_score_dur)}') + print(f'Deleted segments duration: {round(del_duration)}') + print(f'Stats saved at {stats_file}') + print(f'Total execution time: ~{round(total_time / 60)}min') diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/normalization_helpers.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/normalization_helpers.py new file mode 100644 index 0000000000..7e971e7c80 --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/normalization_helpers.py @@ -0,0 +1,69 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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. + +__all__ = ['LATIN_TO_RU', 'RU_ABBREVIATIONS', 'NUMBERS_TO_ENG', 'NUMBERS_TO_RU'] + +LATIN_TO_RU = { + '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': 'з', + 'à': 'а', + 'è': 'е', + 'é': 'е', + 'ß': 'в', + 'ä': 'а', + 'ö': 'о', + 'ü': 'у', + 'є': 'е', + 'ç': 'с', + 'ê': 'е', + 'ó': 'о', +} + +RU_ABBREVIATIONS = { + ' р.': ' рублей', + ' к.': ' копеек', + ' коп.': ' копеек', + ' копек.': ' копеек', + ' т.д.': ' так далее', + ' т. д.': ' так далее', + ' т.п.': ' тому подобное', + ' т. п.': ' тому подобное', + ' т.е.': ' то есть', + ' т. е.': ' то есть', + ' стр. ': ' страница ', +} diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/prepare_data.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/prepare_data.py new file mode 100644 index 0000000000..47373fb357 --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/prepare_data.py @@ -0,0 +1,412 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import multiprocessing +import os +import re +from pathlib import Path +from typing import List + +import regex +import scipy.io.wavfile as wav +from normalization_helpers import LATIN_TO_RU, RU_ABBREVIATIONS +from num2words import num2words + +from nemo.collections import asr as nemo_asr + +try: + from nemo_text_processing.text_normalization.normalize import Normalizer + + NEMO_NORMALIZATION_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + NEMO_NORMALIZATION_AVAILABLE = False + + +parser = argparse.ArgumentParser(description="Prepares text and audio files for segmentation") +parser.add_argument("--in_text", type=str, default=None, help='Path to a text file or a directory with .txt files') +parser.add_argument("--output_dir", type=str, required=True, help='Path to output directory') +parser.add_argument("--audio_dir", type=str, help='Path to folder with .mp3 or .wav audio files') +parser.add_argument( + "--audio_format", type=str, default='.mp3', choices=['.mp3', '.wav'], help='Audio files format in --audio_dir' +) +parser.add_argument('--sample_rate', type=int, default=16000, help='Sampling rate used during ASR model training') +parser.add_argument( + '--language', type=str, default='eng', choices=['eng', 'ru', 'add other languages supported by num2words.'] +) +parser.add_argument( + '--cut_prefix', type=int, default=0, help='Number of seconds to cut from the beginning of the audio files.', +) +parser.add_argument( + '--model', type=str, default='QuartzNet15x5Base-En', help='Pre-trained model name or path to model checkpoint' +) +parser.add_argument('--min_length', type=int, default=0, help='Min number of chars of the text segment for alignment.') +parser.add_argument( + '--max_length', type=int, default=100, help='Max number of chars of the text segment for alignment.' +) +parser.add_argument( + '--additional_split_symbols', + type=str, + default='', + help='Additional symbols to use for \ + sentence split if eos sentence split resulted in sequence longer than --max_length. ' + 'Use "|" as a separator between symbols, for example: ";|:|" ', +) +parser.add_argument( + '--use_nemo_normalization', + action='store_true', + help='Set to True to use NeMo Normalization tool to convert numbers from written to spoken format.', +) + + +def convert_audio(in_file: str, wav_file: str = None, sample_rate: int = 16000) -> str: + """ + Convert .mp3 to .wav and/or change sample rate if needed + + Args: + in_file: Path to .mp3 or .wav file + sample_rate: Desired sample rate + + Returns: + path to .wav file + """ + print(f"Converting {in_file} to .wav format with sample rate {sample_rate}") + if not os.path.exists(in_file): + raise ValueError(f'{in_file} not found') + if wav_file is None: + wav_file = in_file.replace(os.path.splitext(in_file)[-1], f"_{sample_rate}.wav") + + os.system( + f'ffmpeg -i {in_file} -acodec pcm_s16le -ac 1 -af aresample=resampler=soxr -ar {sample_rate} {wav_file} -y' + ) + return wav_file + + +def process_audio(in_file: str, wav_file: str = None, cut_prefix: int = 0, sample_rate: int = 16000): + """Process audio file: .mp3 to .wav conversion and cut a few seconds from the beginning of the audio + + Args: + in_file: path to the .mp3 or .wav file for processing + wav_file: path to the output .wav file + cut_prefix: number of seconds to cut from the beginning of the audio file + sample_rate: target sampling rate + """ + wav_audio = convert_audio(str(in_file), wav_file, sample_rate) + + if cut_prefix > 0: + # cut a few seconds of audio from the beginning + sample_rate, signal = wav.read(wav_audio) + wav.write(wav_audio, data=signal[cut_prefix * sample_rate :], rate=sample_rate) + + +def split_text( + in_file: str, + out_file: str, + vocabulary: List[str] = None, + language='eng', + remove_brackets=True, + do_lower_case=True, + min_length=0, + max_length=100, + additional_split_symbols=None, + use_nemo_normalization=False, +): + """ + Breaks down the in_file roughly into sentences. Each sentence will be on a separate line. + Written form of the numbers will be converted to its spoken equivalent, OOV punctuation will be removed. + + Args: + in_file: path to original transcript + out_file: path to the output file + vocabulary: ASR model vocabulary + language: text language + remove_brackets: Set to True if square [] and curly {} brackets should be removed from text. + Text in square/curly brackets often contains inaudible fragments like notes or translations + do_lower_case: flag that determines whether to apply lower case to the in_file text + min_length: Min number of chars of the text segment for alignment. Short segments will be combined to be + at least min_length (not recommended for multi speaker data). + max_length: Max number of chars of the text segment for alignment + additional_split_symbols: Additional symbols to use for sentence split if eos sentence split resulted in + segments longer than --max_length + use_nemo_normalization: Set to True to use NeMo normalization tool to convert numbers from written to spoken + format. Normalization using num2words will be applied afterwards to make sure there are no numbers present + in the text, otherwise they will be replaced with a space and that could deteriorate segmentation results. + """ + + print(f'Splitting text in {in_file} into sentences.') + with open(in_file, "r") as f: + transcript = f.read() + + # remove some symbols for better split into sentences + transcript = ( + transcript.replace("\n", " ") + .replace("\t", " ") + .replace("…", "...") + .replace("\\", " ") + .replace("--", " -- ") + .replace(". . .", "...") + .replace("‘", "’") + ) + # remove extra space + transcript = re.sub(r' +', ' ', transcript) + transcript = re.sub(r'(\.+)', '. ', transcript) + + if remove_brackets: + transcript = re.sub(r'(\[.*?\])', ' ', transcript) + # remove text in curly brackets + transcript = re.sub(r'(\{.*?\})', ' ', transcript) + + lower_case_unicode = '' + upper_case_unicode = '' + if language == 'ru': + lower_case_unicode = '\u0430-\u04FF' + upper_case_unicode = '\u0410-\u042F' + elif language not in ['ru', 'eng']: + print(f'Consider using {language} unicode letters for better sentence split.') + + # remove space in the middle of the lower case abbreviation to avoid splitting into separate sentences + matches = re.findall(r'[a-z' + lower_case_unicode + ']\.\s[a-z' + lower_case_unicode + ']\.', transcript) + for match in matches: + transcript = transcript.replace(match, match.replace('. ', '.')) + + # find phrases in quotes + with_quotes = re.finditer(r'“[A-Za-z ?]+.*?”', transcript) + sentences = [] + last_idx = 0 + for m in with_quotes: + match = m.group() + match_idx = m.start() + if last_idx < match_idx: + sentences.append(transcript[last_idx:match_idx]) + sentences.append(match) + last_idx = m.end() + sentences.append(transcript[last_idx:]) + sentences = [s.strip() for s in sentences if s.strip()] + + # Read and split transcript by utterance (roughly, sentences) + split_pattern = f"(? 0] + + if min_length > 0: + sentences_comb = [] + sentences_comb.append(sentences[0]) + # combines short sentence + for i in range(1, len(sentences)): + if len(sentences_comb[-1]) < min_length or len(sentences[i]) < min_length: + sentences_comb[-1] += ' ' + sentences[i].strip() + else: + sentences_comb.append(sentences[i].strip()) + sentences = sentences_comb + + sentences = [s.strip() for s in sentences if s.strip()] + + # save split text with original punctuation and case + out_dir, out_file_name = os.path.split(out_file) + with open(os.path.join(out_dir, out_file_name[:-4] + '_with_punct.txt'), "w") as f: + f.write("\n".join(sentences)) + + # substitute common abbreviations before applying lower case + if language == 'ru': + for k, v in RU_ABBREVIATIONS.items(): + sentences = [s.replace(k, v) for s in sentences] + + if language == 'ru': + # replace Latin characters with Russian + for k, v in LATIN_TO_RU.items(): + sentences = [s.replace(k, v) for s in sentences] + + if language == 'eng' and use_nemo_normalization: + if not NEMO_NORMALIZATION_AVAILABLE: + raise ValueError(f'NeMo normalization tool is not installed.') + + print('Using NeMo normalization tool...') + normalizer = Normalizer(input_case='cased') + sentences_norm = normalizer.normalize_list(sentences, verbose=False) + if len(sentences_norm) != len(sentences): + raise ValueError(f'Normalization failed, number of sentences does not match.') + + sentences = '\n'.join(sentences) + + # replace numbers with num2words + try: + p = re.compile("\d+") + new_text = '' + match_end = 0 + for i, m in enumerate(p.finditer(sentences)): + match = m.group() + match_start = m.start() + if i == 0: + new_text = sentences[:match_start] + else: + new_text += sentences[match_end:match_start] + match_end = m.end() + new_text += sentences[match_start:match_end].replace(match, num2words(match, lang=language)) + new_text += sentences[match_end:] + sentences = new_text + except NotImplementedError: + print( + f'{language} might be missing in "num2words" package. Add required language to the choices for the' + f'--language argument.' + ) + raise + + sentences = ( + sentences.replace("’", "'") + .replace("»", '"') + .replace("«", '"') + .replace("\\", "") + .replace("”", '"') + .replace("„", '"') + .replace("´", "'") + .replace("-- --", "--") + .replace("--", " -- ") + .replace("’", "'") + .replace('“', '"') + .replace('“', '"') + .replace("‘", "'") + .replace('—', '-') + .replace("- -", "--") + .replace('`', "'") + .replace(' !', '!') + .replace(' ?', '?') + .replace(' ,', ',') + .replace(' .', '.') + .replace(' ;', ';') + .replace(' :', ':') + .replace('!!', '!') + .replace('--', '-') + .replace('“', '"') + .replace(', , ', ', ') + .replace('=', '') + ) + + allowed_punct = [',', '.', '?', '!', ':', ';', '-', '"', '(', ')'] + # clean up normalized text and keep only allowed_punct and ASR vocabulary (lower and upper case) + symbols_to_remove = ''.join( + set(sentences).difference(set(vocabulary + [s.upper() for s in vocabulary] + ['\n'] + allowed_punct)) + ) + sentences_norm = sentences.translate(''.maketrans(symbols_to_remove, len(symbols_to_remove) * ' ')) + + with open(os.path.join(out_dir, out_file_name[:-4] + '_with_punct_normalized.txt'), "w") as f: + f.write(sentences_norm) + + if do_lower_case: + sentences = sentences.lower() + + # remove all OOV symbols + symbols_to_remove = ''.join(set(sentences).difference(set(vocabulary + ['\n']))) + sentences = sentences.translate(''.maketrans(symbols_to_remove, len(symbols_to_remove) * ' ')) + + # remove extra space + sentences = re.sub(r' +', ' ', sentences) + with open(out_file, "w") as f: + f.write(sentences) + + +if __name__ == '__main__': + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + + text_files = [] + if args.in_text: + vocabulary = None + if args.model is None: + print(f"No model provided, vocabulary won't be used") + elif os.path.exists(args.model): + asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model) + vocabulary = asr_model.cfg.decoder.vocabulary + elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names(): + asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model) + vocabulary = asr_model.cfg.decoder.vocabulary + else: + raise ValueError( + f'Provide path to the pretrained checkpoint or choose from {nemo_asr.models.EncDecCTCModel.get_available_model_names()}' + ) + + if os.path.isdir(args.in_text): + text_files = Path(args.in_text).glob(("*.txt")) + else: + text_files.append(Path(args.in_text)) + for text in text_files: + base_name = os.path.basename(text)[:-4] + out_text_file = os.path.join(args.output_dir, base_name + '.txt') + + split_text( + text, + out_text_file, + vocabulary=vocabulary, + language=args.language, + min_length=args.min_length, + max_length=args.max_length, + additional_split_symbols=args.additional_split_symbols, + use_nemo_normalization=args.use_nemo_normalization, + ) + print(f'Processed text saved at {args.output_dir}') + + if args.audio_dir: + if not os.path.exists(args.audio_dir): + raise ValueError(f'{args.audio_dir} not found. "--audio_dir" should contain .mp3 or .wav files.') + + audio_paths = list(Path(args.audio_dir).glob(f"*{args.audio_format}")) + + workers = [] + for i in range(len(audio_paths)): + wav_file = os.path.join(args.output_dir, audio_paths[i].name.replace(args.audio_format, ".wav")) + worker = multiprocessing.Process( + target=process_audio, args=(audio_paths[i], wav_file, args.cut_prefix, args.sample_rate), + ) + workers.append(worker) + worker.start() + for w in workers: + w.join() + + print('Done.') diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/process_manifests.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/process_manifests.py new file mode 100644 index 0000000000..d3ef9870fc --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/process_manifests.py @@ -0,0 +1,111 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 argparse +import json +import os +from pathlib import Path +from random import randint +from typing import List + +parser = argparse.ArgumentParser(description="Combine manifests and extract samples") +parser.add_argument( + "--manifests_dir", + default='_high_score_manifest.json', + type=str, + required=True, + help='Path to directory with *_high_score_manifest.json files', +) +parser.add_argument( + "--output_dir", default='output', required=True, help='Path to the directory to store final manifest' +) +parser.add_argument( + "--num_samples", + default=0, + type=int, + help='The number of samples to get from the beginning, end and randomly from the middle of highly scored samples', +) + + +def get_samples(manifest: str, num_samples: int) -> List[str]: + """ + Samples num_samples from the beginning, end and randomly from the middle of the manifest + + Args: + manifest: path to manifest + num_samples: number of samples to extract from the beginning, end and randomly from the middle of the manifest + + Returns: + num_samples: manifest lines + """ + samples = [] + with open(manifest, 'r', encoding='utf8') as f: + lines = f.readlines() + + # first n samples + samples.extend(lines[:num_samples]) + # random in the middle + for i in range(num_samples): + samples.append(lines[randint(num_samples, len(lines) - num_samples)]) + # last n samples + samples.extend(lines[-num_samples:]) + return samples + + +if __name__ == '__main__': + args = parser.parse_args() + manifest_dir = Path(args.manifests_dir) + manifests = sorted(list(manifest_dir.glob('*_high_score_manifest.json'))) + sample_manifest = os.path.join(args.output_dir, 'sample_manifest.json') + all_manifest = os.path.join(args.output_dir, 'all_manifest.json') + + if args.num_samples > 0: + with open(sample_manifest, 'w', encoding='utf8') as f: + for manifest in manifests: + samples = get_samples(manifest, args.num_samples) + for sample in samples: + f.write(sample) + + print(f'Sample manifest is saved at {sample_manifest}') + + total_duration = 0 + with open(all_manifest, 'w') as out_f: + for manifest in manifests: + with open(manifest, 'r', encoding='utf8') as f: + for line in f: + info = json.loads(line) + total_duration += info['duration'] + out_f.write(line) + + print(f'Aggregated manifest is saved at {all_manifest}') + print(f'Total files duration: ~{round(total_duration/60)} min or ~{round(total_duration/60/60)} hr') + + sets = {'_high_score_manifest.json': 0, '_low_score_manifest.json': 0, '_del_manifest.json': 0} + + for k in sets: + files = manifest_dir.glob(f'*{k}') + for manifest_f in files: + with open(manifest_f, 'r') as f: + for line in f: + line = json.loads(line) + sets[k] += line['duration'] + + total = 0 + for k in sets: + sets[k] = round(sets[k] / 3600, 1) + total += sets[k] + name = k.split('_')[1].upper() + print(f'Total {name} score duration: ~{sets[k]} hr') + print(f'Total duration: ~{total} hr') + print('Done.') diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/run_ctc_segmentation.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/run_ctc_segmentation.py new file mode 100644 index 0000000000..06b4f3669b --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/run_ctc_segmentation.py @@ -0,0 +1,180 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import logging +import multiprocessing +import os +import sys +import time +from pathlib import Path + +import numpy as np +import scipy.io.wavfile as wav +import torch +from utils import get_segments, listener_configurer, listener_process, worker_configurer, worker_process + +import nemo.collections.asr as nemo_asr + +parser = argparse.ArgumentParser(description="CTC Segmentation") +parser.add_argument("--output_dir", default='output', type=str, help='Path to output directory') +parser.add_argument( + "--data", + type=str, + required=True, + help='Path to directory with audio files and associated transcripts (same respective names only formats are ' + 'different or path to wav file (transcript should have the same base name and be located in the same folder' + 'as the wav file.', +) +parser.add_argument('--window_len', type=int, default=8000, help='Window size for ctc segmentation algorithm') +parser.add_argument('--no_parallel', action='store_true', help='Flag to disable parallel segmentation') +parser.add_argument('--sample_rate', type=int, default=16000, help='Sampling rate') +parser.add_argument( + '--model', type=str, default='QuartzNet15x5Base-En', help='Path to model checkpoint or pre-trained model name', +) +parser.add_argument('--debug', action='store_true', help='Flag to enable debugging messages') + +logger = logging.getLogger('ctc_segmentation') # use module name + +if __name__ == '__main__': + + args = parser.parse_args() + + # setup logger + log_dir = os.path.join(args.output_dir, 'logs') + os.makedirs(log_dir, exist_ok=True) + log_file = os.path.join(log_dir, f'ctc_segmentation_{args.window_len}.log') + level = 'DEBUG' if args.debug else 'INFO' + + logger = logging.getLogger('CTC') + file_handler = logging.FileHandler(filename=log_file) + stdout_handler = logging.StreamHandler(sys.stdout) + handlers = [file_handler, stdout_handler] + logging.basicConfig(handlers=handlers, level=level) + + if os.path.exists(args.model): + asr_model = nemo_asr.models.EncDecCTCModel.restore_from(args.model) + elif args.model in nemo_asr.models.EncDecCTCModel.get_available_model_names(): + asr_model = nemo_asr.models.EncDecCTCModel.from_pretrained(args.model, strict=False) + else: + raise ValueError( + f'{args.model} not a valid model name or path. Provide path to the pre-trained checkpoint ' + f'or choose from {nemo_asr.models.EncDecCTCModel.list_available_models()}' + ) + + # extract ASR vocabulary and add blank symbol + vocabulary = asr_model.cfg.decoder.vocabulary + odim = len(vocabulary) + 1 + logging.debug(f'ASR Model vocabulary: {vocabulary}') + + # add blank to vocab + vocabulary = ["ε"] + list(vocabulary) + data = Path(args.data) + output_dir = Path(args.output_dir) + + if os.path.isdir(data): + audio_paths = data.glob("*.wav") + data_dir = data + else: + audio_paths = [Path(data)] + data_dir = Path(os.path.dirname(data)) + + all_log_probs = [] + all_transcript_file = [] + all_segment_file = [] + all_wav_paths = [] + segments_dir = os.path.join(args.output_dir, 'segments') + os.makedirs(segments_dir, exist_ok=True) + for path_audio in audio_paths: + transcript_file = os.path.join(data_dir, path_audio.name.replace(".wav", ".txt")) + segment_file = os.path.join( + segments_dir, f"{args.window_len}_" + path_audio.name.replace(".wav", "_segments.txt") + ) + try: + sample_rate, signal = wav.read(path_audio) + if sample_rate != args.sample_rate: + raise ValueError( + f'Sampling rate of the audio file {path_audio} doesn\'t match --sample_rate={args.sample_rate}' + ) + except ValueError: + logging.error( + f"{path_audio} should be a .wav mono file with the sampling rate used for the ASR model training" + f"specified with {args.sample_rate}." + ) + raise + + original_duration = len(signal) / sample_rate + logging.debug(f'Duration: {original_duration}s, file_name: {path_audio}') + log_probs = asr_model.transcribe(paths2audio_files=[str(path_audio)], batch_size=1, logprobs=True)[0] + + # move blank values to the first column + blank_col = log_probs[:, -1].reshape((log_probs.shape[0], 1)) + log_probs = np.concatenate((blank_col, log_probs[:, :-1]), axis=1) + + all_log_probs.append(log_probs) + all_segment_file.append(str(segment_file)) + all_transcript_file.append(str(transcript_file)) + all_wav_paths.append(path_audio) + + del asr_model + torch.cuda.empty_cache() + + if len(all_log_probs) == 0: + raise ValueError(f'No valid audio files found at {args.data}') + start_time = time.time() + if args.no_parallel: + for i in range(len(all_log_probs)): + get_segments( + all_log_probs[i], + all_wav_paths[i], + all_transcript_file[i], + all_segment_file[i], + vocabulary, + args.window_len, + ) + else: + queue = multiprocessing.Queue(-1) + + listener = multiprocessing.Process(target=listener_process, args=(queue, listener_configurer, log_file, level)) + listener.start() + workers = [] + for i in range(len(all_log_probs)): + worker = multiprocessing.Process( + target=worker_process, + args=( + queue, + worker_configurer, + level, + all_log_probs[i], + all_wav_paths[i], + all_transcript_file[i], + all_segment_file[i], + vocabulary, + args.window_len, + ), + ) + workers.append(worker) + worker.start() + for w in workers: + w.join() + queue.put_nowait(None) + listener.join() + + total_time = time.time() - start_time + logger.info(f'Total execution time: ~{round(total_time/60)}min') + logger.info(f'Saving logs to {log_file}') + + if os.path.exists(log_file): + with open(log_file, 'r') as f: + lines = f.readlines() diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/utils.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/utils.py new file mode 100644 index 0000000000..713f8d4c5f --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/utils.py @@ -0,0 +1,173 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 logging +import logging.handlers +import multiprocessing +import os +from pathlib import PosixPath +from typing import List, Tuple, Union + +import ctc_segmentation as cs +import numpy as np + + +def get_segments( + log_probs: np.ndarray, + path_wav: Union[PosixPath, str], + transcript_file: Union[PosixPath, str], + output_file: str, + vocabulary: List[str], + window_size: int = 8000, + frame_duration_ms: int = 20, +) -> None: + """ + Segments the audio into segments and saves segments timings to a file + + Args: + log_probs: Log probabilities for the original audio from an ASR model, shape T * |vocabulary|. + values for blank should be at position 0 + path_wav: path to the audio .wav file + transcript_file: path to + output_file: path to the file to save timings for segments + vocabulary: vocabulary used to train the ASR model, note blank is at position 0 + window_size: the length of each utterance (in terms of frames of the CTC outputs) fits into that window. + frame_duration_ms: frame duration in ms + """ + config = cs.CtcSegmentationParameters() + config.char_list = vocabulary + config.min_window_size = window_size + config.frame_duration_ms = frame_duration_ms + config.blank = config.space + config.subsampling_factor = 2 + + with open(transcript_file, "r") as f: + text = f.readlines() + text = [t.strip() for t in text if t.strip()] + + # add corresponding original text without pre-processing + transcript_file_no_preprocessing = transcript_file.replace('.txt', '_with_punct.txt') + if not os.path.exists(transcript_file_no_preprocessing): + raise ValueError(f'{transcript_file_no_preprocessing} not found.') + + with open(transcript_file_no_preprocessing, "r") as f: + text_no_preprocessing = f.readlines() + text_no_preprocessing = [t.strip() for t in text_no_preprocessing if t.strip()] + + # add corresponding normalized original text + transcript_file_normalized = transcript_file.replace('.txt', '_with_punct_normalized.txt') + if not os.path.exists(transcript_file_normalized): + raise ValueError(f'{transcript_file_normalized} not found.') + + with open(transcript_file_normalized, "r") as f: + text_normalized = f.readlines() + text_normalized = [t.strip() for t in text_normalized if t.strip()] + + if len(text_no_preprocessing) != len(text): + raise ValueError(f'{transcript_file} and {transcript_file_no_preprocessing} do not match') + + if len(text_normalized) != len(text): + raise ValueError(f'{transcript_file} and {transcript_file_normalized} do not match') + + ground_truth_mat, utt_begin_indices = cs.prepare_text(config, text) + logging.debug(f"Syncing {transcript_file}") + logging.debug( + f"Audio length {os.path.basename(path_wav)}: {log_probs.shape[0]}. " + f"Text length {os.path.basename(transcript_file)}: {len(ground_truth_mat)}" + ) + + timings, char_probs, char_list = cs.ctc_segmentation(config, log_probs, ground_truth_mat) + segments = cs.determine_utterance_segments(config, utt_begin_indices, char_probs, timings, text) + write_output(output_file, path_wav, segments, text, text_no_preprocessing, text_normalized) + + +def write_output( + out_path: str, + path_wav: str, + segments: List[Tuple[float]], + text: str, + text_no_preprocessing: str, + text_normalized: str, + stride: int = 2, +): + """ + Write the segmentation output to a file + + out_path: Path to output file + path_wav: Path to the original audio file + segments: Segments include start, end and alignment score + text: Text used for alignment + text_no_preprocessing: Reference txt without any pre-processing + text_normalized: Reference text normalized + stride: Stride applied to an ASR input + """ + # Uses char-wise alignments to get utterance-wise alignments and writes them into the given file + with open(str(out_path), "w") as outfile: + outfile.write(str(path_wav) + "\n") + + for i, (start, end, score) in enumerate(segments): + outfile.write( + f'{start/stride} {end/stride} {score} | {text[i]} | {text_no_preprocessing[i]} | {text_normalized[i]}\n' + ) + + +##################### +# logging utils +##################### +def listener_configurer(log_file, level): + root = logging.getLogger() + h = logging.handlers.RotatingFileHandler(log_file, 'w') + f = logging.Formatter('%(asctime)s %(processName)-10s %(name)s %(levelname)-8s %(message)s') + h.setFormatter(f) + ch = logging.StreamHandler() + root.addHandler(h) + root.setLevel(level) + root.addHandler(ch) + + +def listener_process(queue, configurer, log_file, level): + configurer(log_file, level) + while True: + try: + record = queue.get() + if record is None: # We send this as a sentinel to tell the listener to quit. + break + logger = logging.getLogger(record.name) + logger.setLevel(logging.INFO) + logger.handle(record) # No level or filter logic applied - just do it! + + except Exception: + import sys + import traceback + + print('Problem:', file=sys.stderr) + traceback.print_exc(file=sys.stderr) + + +def worker_configurer(queue, level): + h = logging.handlers.QueueHandler(queue) # Just the one handler needed + root = logging.getLogger() + root.addHandler(h) + root.setLevel(level) + + +def worker_process( + queue, configurer, level, log_probs, path_wav, transcript_file, output_file, vocabulary, window_len +): + configurer(queue, level) + name = multiprocessing.current_process().name + innerlogger = logging.getLogger('worker') + innerlogger.info(f'{name} is processing {path_wav}, window_len={window_len}') + get_segments(log_probs, path_wav, transcript_file, output_file, vocabulary, window_len) + innerlogger.info(f'{name} completed segmentation of {path_wav}, segments saved to {output_file}') diff --git a/models/tts/NeMo/tools/ctc_segmentation/scripts/verify_segments.py b/models/tts/NeMo/tools/ctc_segmentation/scripts/verify_segments.py new file mode 100644 index 0000000000..9d648df912 --- /dev/null +++ b/models/tts/NeMo/tools/ctc_segmentation/scripts/verify_segments.py @@ -0,0 +1,143 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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 argparse +import os +import re +from pathlib import Path + +import pandas as pd + +parser = argparse.ArgumentParser(description="Compare alignment segments generated with different window sizes") +parser.add_argument( + "--base_dir", + default='output', + type=str, + required=True, + help='Path to directory with "logs" and "segments" folders generated during the segmentation step', +) + + +def _process_log_file(log_file: str): + processes = {} + with open(log_file, 'r') as f: + for line in f: + if line[0].isdigit() and 'Process-' in line: + process_id = re.search(r'Process-\d+', line)[0] + window_size = re.findall(r'\d+', line)[-1] + path = re.search(r'\/.*?\.[\w:]+', line) + if 'is processing' in line: + if process_id not in processes: + processes[process_id] = {} + processes[process_id]['Window Increased'] = False + processes[process_id]['Start Window'] = window_size + processes[process_id]['Final Window'] = window_size + base_dir, base_name = os.path.split(path[0]) + processes[process_id]['Audio Dir'] = base_dir + processes[process_id]['Audio File'] = base_name + elif 'Increasing' in line: + processes[process_id]['Final Window'] = window_size + processes[process_id]['Window Increased'] = True + elif 'completed' in line: + processes[process_id]['Completed'] = True + + df = pd.DataFrame.from_dict( + processes, + columns=['Audio Dir', 'Audio File', 'Start Window', 'Final Window', 'Window Increased', 'Completed'], + orient='index', + ).reset_index() + return df + + +if __name__ == '__main__': + args = parser.parse_args() + log_dir = os.path.join(args.base_dir, 'logs') + if not os.path.exists(log_dir): + raise ValueError(f'"log" directory was not found at {args.base_dir}.') + + log_files = Path(log_dir).glob('*.log') + dfs = [] + for log in log_files: + dfs.append(_process_log_file(log)) + + dfs = pd.concat(dfs) + summary_file = os.path.join(log_dir, 'log_summary.csv') + dfs.to_csv(summary_file, index=False) + print(f'Log summary saved to {summary_file}') + + segments_dir = os.path.join(args.base_dir, 'segments') + if not os.path.exists(segments_dir): + raise ValueError(f'"segments" directory was not found at {args.base_dir}.') + + all_files = Path(segments_dir).glob('*_segments.txt') + all_alignment_files = {} + for file in all_files: + base_name = re.sub(r'^\d+_', '', file.name) + if base_name not in all_alignment_files: + all_alignment_files[base_name] = [] + all_alignment_files[base_name].append(file) + + verified_dir = os.path.join(args.base_dir, 'verified_segments') + os.makedirs(verified_dir, exist_ok=True) + + def readlines(file): + with open(file, 'r') as f: + lines = f.readlines() + return lines + + stats = {} + for part, alignment_files in all_alignment_files.items(): + stats[part] = {} + num_alignment_files = len(alignment_files) + all_alignments = [] + for alignment in alignment_files: + all_alignments.append(readlines(alignment)) + + with open(os.path.join(verified_dir, part), 'w') as f: + num_segments = len(all_alignments[0]) + stats[part]['Original number of segments'] = num_segments + stats[part]['Verified segments'] = 0 + stats[part]['Original Duration, min'] = 0 + stats[part]['Verified Duration, min'] = 0 + + for i in range(num_segments): + line = all_alignments[0][i] + valid_line = True + if i == 0: + duration = 0 + else: + info = line.split('|')[0].split() + duration = (float(info[1]) - float(info[0])) / 60 + stats[part]['Original Duration, min'] += duration + for alignment in all_alignments: + if line != alignment[i]: + valid_line = False + if valid_line: + f.write(line) + stats[part]['Verified segments'] += 1 + stats[part]['Verified Duration, min'] += duration + + stats = pd.DataFrame.from_dict(stats, orient='index').reset_index() + stats['Number dropped'] = stats['Original number of segments'] - stats['Verified segments'] + stats['Duration of dropped, min'] = round(stats['Original Duration, min'] - stats['Verified Duration, min']) + stats['% dropped, min'] = round(stats['Duration of dropped, min'] / stats['Original number of segments'] * 100) + stats['Misalignment present'] = stats['Number dropped'] > 0 + stats['Original Duration, min'] = round(stats['Original Duration, min']) + stats['Verified Duration, min'] = round(stats['Verified Duration, min']) + stats.loc["Total"] = stats.sum() + + stats_file = os.path.join(args.base_dir, 'alignment_summary.csv') + stats.to_csv(stats_file, index=False) + print(stats) + print(f'Alignment summary saved to {stats_file}') diff --git a/models/tts/NeMo/tools/speech_data_explorer/README.md b/models/tts/NeMo/tools/speech_data_explorer/README.md new file mode 100644 index 0000000000..9d431f20fd --- /dev/null +++ b/models/tts/NeMo/tools/speech_data_explorer/README.md @@ -0,0 +1,26 @@ +Speech Data Explorer +-------------------- + +[Dash](https://plotly.com/dash/)-based tool for interactive exploration of ASR/TTS datasets. + +Features: +- dataset's statistics (alphabet, vocabulary, duration-based histograms) +- navigation across dataset (sorting, filtering) +- inspection of individual utterances (waveform, spectrogram, audio player) +- errors' analysis (Word Error Rate, Character Error Rate, Word Match Rate, Mean Word Accuracy, diff) + +Please make sure that requirements are installed. Then run: +``` +python data_explorer.py path_to_manifest.json +``` + +JSON manifest file should contain the following fields: +- "audio_filepath" (path to audio file) +- "duration" (duration of the audio file in seconds) +- "text" (reference transcript) + +Errors' analysis requires "pred_text" (ASR transcript) for all utterances. + +Any additional field will be parsed and displayed in 'Samples' tab. + +![Speech Data Explorer](screenshot.png) diff --git a/models/tts/NeMo/tools/speech_data_explorer/data_explorer.py b/models/tts/NeMo/tools/speech_data_explorer/data_explorer.py new file mode 100755 index 0000000000..e760915347 --- /dev/null +++ b/models/tts/NeMo/tools/speech_data_explorer/data_explorer.py @@ -0,0 +1,717 @@ +# Copyright (c) 2020, NVIDIA CORPORATION. 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 argparse +import base64 +import datetime +import difflib +import io +import json +import math +import operator +import os +import pickle +from collections import defaultdict + +import dash +import dash_bootstrap_components as dbc +import dash_core_components as dcc +import dash_html_components as html +import dash_table +import diff_match_patch +import editdistance +import librosa +import numpy as np +import soundfile as sf +import tqdm +from dash.dependencies import Input, Output +from dash.exceptions import PreventUpdate +from plotly import express as px +from plotly import graph_objects as go +from plotly.subplots import make_subplots + +# number of items in a table per page +DATA_PAGE_SIZE = 10 + +# operators for filtering items +filter_operators = [ + ['ge ', '>='], + ['le ', '<='], + ['lt ', '<'], + ['gt ', '>'], + ['ne ', '!='], + ['eq ', '='], + ['contains '], +] + +# parse table filter queries +def split_filter_part(filter_part): + for operator_type in filter_operators: + for op in operator_type: + if op in filter_part: + name_part, value_part = filter_part.split(op, 1) + name = name_part[name_part.find('{') + 1 : name_part.rfind('}')] + value_part = value_part.strip() + v0 = value_part[0] + if v0 == value_part[-1] and v0 in ("'", '"', '`'): + value = value_part[1:-1].replace('\\' + v0, v0) + else: + try: + value = float(value_part) + except ValueError: + value = value_part + return name, operator_type[0].strip(), value + return [None] * 3 + + +# standard command-line arguments parser +def parse_args(): + parser = argparse.ArgumentParser(description='Speech Data Explorer') + parser.add_argument( + 'manifest', help='path to JSON manifest file', + ) + parser.add_argument('--vocab', help='optional vocabulary to highlight OOV words') + parser.add_argument('--port', default='8050', help='serving port for establishing connection') + parser.add_argument( + '--disable-caching-metrics', action='store_true', help='disable caching metrics for errors analysis' + ) + parser.add_argument('--debug', '-d', action='store_true', help='enable debug mode') + args = parser.parse_args() + print(args) + return args + + +# load data from JSON manifest file +def load_data(data_filename, disable_caching, vocab=None): + if not disable_caching: + pickle_filename = data_filename.split('.json')[0] + json_mtime = datetime.datetime.fromtimestamp(os.path.getmtime(data_filename)) + timestamp = json_mtime.strftime('%Y%m%d_%H%M') + pickle_filename += '_' + timestamp + '.pkl' + if os.path.exists(pickle_filename): + with open(pickle_filename, 'rb') as f: + data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available = pickle.load(f) + return data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available + + if vocab is not None: + # load external vocab + vocabulary_ext = {} + with open(vocab, 'r') as f: + for line in f: + if '\t' in line: + # parse word from TSV file + word = line.split('\t')[0] + else: + # assume each line contains just a single word + word = line.strip() + vocabulary_ext[word] = 1 + + data = [] + wer_dist = 0.0 + wer_count = 0 + cer_dist = 0.0 + cer_count = 0 + wmr_count = 0 + wer = 0 + cer = 0 + wmr = 0 + mwa = 0 + num_hours = 0 + vocabulary = defaultdict(lambda: 0) + alphabet = set() + match_vocab = defaultdict(lambda: 0) + + sm = difflib.SequenceMatcher() + metrics_available = False + with open(data_filename, 'r', encoding='utf8') as f: + for line in tqdm.tqdm(f): + item = json.loads(line) + num_words = len(item['text'].split()) + num_chars = len(item['text']) + orig = item['text'].split() + for word in orig: + vocabulary[word] += 1 + for char in item['text']: + alphabet.add(char) + num_hours += item['duration'] + + if 'pred_text' in item: + metrics_available = True + pred = item['pred_text'].split() + word_dist = editdistance.eval(orig, pred) + char_dist = editdistance.eval(item['text'], item['pred_text']) + wer_dist += word_dist + cer_dist += char_dist + wer_count += num_words + cer_count += num_chars + + sm.set_seqs(orig, pred) + num_matches = 0 + for m in sm.get_matching_blocks(): + for word_idx in range(m[0], m[0] + m[2]): + match_vocab[orig[word_idx]] += 1 + num_matches += 1 + wmr_count += num_matches + + data.append( + { + 'audio_filepath': item['audio_filepath'], + 'duration': round(item['duration'], 2), + 'num_words': num_words, + 'num_chars': num_chars, + 'word_rate': round(num_words / item['duration'], 2), + 'char_rate': round(num_chars / item['duration'], 2), + 'text': item['text'], + } + ) + if metrics_available: + data[-1]['pred_text'] = item['pred_text'] + data[-1]['WER'] = round(word_dist / num_words * 100.0, 2) + data[-1]['CER'] = round(char_dist / num_chars * 100.0, 2) + data[-1]['WMR'] = round(num_matches / len(orig) * 100.0, 2) + for k in item: + if k not in data[-1]: + data[-1][k] = item[k] + + vocabulary_data = [{'word': word, 'count': vocabulary[word]} for word in vocabulary] + if vocab is not None: + for item in vocabulary_data: + item['OOV'] = item['word'] not in vocabulary_ext + + if metrics_available: + wer = wer_dist / wer_count * 100.0 + cer = cer_dist / cer_count * 100.0 + wmr = wmr_count / wer_count * 100.0 + + acc_sum = 0 + for item in vocabulary_data: + w = item['word'] + word_accuracy = match_vocab[w] / vocabulary[w] * 100.0 + acc_sum += word_accuracy + item['accuracy'] = round(word_accuracy, 1) + mwa = acc_sum / len(vocabulary_data) + + num_hours /= 3600.0 + + if not disable_caching: + with open(pickle_filename, 'wb') as f: + pickle.dump( + [data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available], + f, + pickle.HIGHEST_PROTOCOL, + ) + + return data, wer, cer, wmr, mwa, num_hours, vocabulary_data, alphabet, metrics_available + + +# plot histogram of specified field in data list +def plot_histogram(data, key, label): + fig = px.histogram( + data_frame=[item[key] for item in data], + nbins=50, + log_y=True, + labels={'value': label}, + opacity=0.5, + color_discrete_sequence=['green'], + height=200, + ) + fig.update_layout(showlegend=False, margin=dict(l=0, r=0, t=0, b=0, pad=0)) + return fig + + +def plot_word_accuracy(vocabulary_data): + labels = ['Unrecognized', 'Sometimes recognized', 'Always recognized'] + counts = [0, 0, 0] + for word in vocabulary_data: + if word['accuracy'] == 0: + counts[0] += 1 + elif word['accuracy'] < 100: + counts[1] += 1 + else: + counts[2] += 1 + colors = ['red', 'orange', 'green'] + + fig = go.Figure( + data=[ + go.Bar( + x=labels, + y=counts, + marker_color=colors, + text=['{:.2%}'.format(count / sum(counts)) for count in counts], + textposition='auto', + ) + ] + ) + fig.update_layout( + showlegend=False, margin=dict(l=0, r=0, t=0, b=0, pad=0), height=200, yaxis={'title_text': '#words'} + ) + + return fig + + +args = parse_args() +print('Loading data...') +data, wer, cer, wmr, mwa, num_hours, vocabulary, alphabet, metrics_available = load_data( + args.manifest, args.disable_caching_metrics, args.vocab +) +print('Starting server...') +app = dash.Dash( + __name__, + suppress_callback_exceptions=True, + external_stylesheets=[dbc.themes.BOOTSTRAP], + title=os.path.basename(args.manifest), +) + + +figure_duration = plot_histogram(data, 'duration', 'Duration (sec)') +figure_num_words = plot_histogram(data, 'num_words', '#words') +figure_num_chars = plot_histogram(data, 'num_chars', '#chars') +figure_word_rate = plot_histogram(data, 'word_rate', '#words/sec') +figure_char_rate = plot_histogram(data, 'char_rate', '#chars/sec') + +if metrics_available: + figure_wer = plot_histogram(data, 'WER', 'WER, %') + figure_cer = plot_histogram(data, 'CER', 'CER, %') + figure_wmr = plot_histogram(data, 'WMR', 'WMR, %') + figure_word_acc = plot_word_accuracy(vocabulary) + +stats_layout = [ + dbc.Row(dbc.Col(html.H5(children='Global Statistics'), className='text-secondary'), className='mt-3'), + dbc.Row( + [ + dbc.Col(html.Div('Number of hours', className='text-secondary'), width=3, className='border-right'), + dbc.Col(html.Div('Number of utterances', className='text-secondary'), width=3, className='border-right'), + dbc.Col(html.Div('Vocabulary size', className='text-secondary'), width=3, className='border-right'), + dbc.Col(html.Div('Alphabet size', className='text-secondary'), width=3), + ], + className='bg-light mt-2 rounded-top border-top border-left border-right', + ), + dbc.Row( + [ + dbc.Col( + html.H5( + '{:.2f} hours'.format(num_hours), + className='text-center p-1', + style={'color': 'green', 'opacity': 0.7}, + ), + width=3, + className='border-right', + ), + dbc.Col( + html.H5(len(data), className='text-center p-1', style={'color': 'green', 'opacity': 0.7}), + width=3, + className='border-right', + ), + dbc.Col( + html.H5( + '{} words'.format(len(vocabulary)), + className='text-center p-1', + style={'color': 'green', 'opacity': 0.7}, + ), + width=3, + className='border-right', + ), + dbc.Col( + html.H5( + '{} chars'.format(len(alphabet)), + className='text-center p-1', + style={'color': 'green', 'opacity': 0.7}, + ), + width=3, + ), + ], + className='bg-light rounded-bottom border-bottom border-left border-right', + ), +] +if metrics_available: + stats_layout += [ + dbc.Row( + [ + dbc.Col( + html.Div('Word Error Rate (WER), %', className='text-secondary'), width=3, className='border-right' + ), + dbc.Col( + html.Div('Character Error Rate (CER), %', className='text-secondary'), + width=3, + className='border-right', + ), + dbc.Col( + html.Div('Word Match Rate (WMR), %', className='text-secondary'), + width=3, + className='border-right', + ), + dbc.Col(html.Div('Mean Word Accuracy, %', className='text-secondary'), width=3), + ], + className='bg-light mt-2 rounded-top border-top border-left border-right', + ), + dbc.Row( + [ + dbc.Col( + html.H5( + '{:.2f}'.format(wer), className='text-center p-1', style={'color': 'green', 'opacity': 0.7}, + ), + width=3, + className='border-right', + ), + dbc.Col( + html.H5( + '{:.2f}'.format(cer), className='text-center p-1', style={'color': 'green', 'opacity': 0.7} + ), + width=3, + className='border-right', + ), + dbc.Col( + html.H5( + '{:.2f}'.format(wmr), className='text-center p-1', style={'color': 'green', 'opacity': 0.7}, + ), + width=3, + className='border-right', + ), + dbc.Col( + html.H5( + '{:.2f}'.format(mwa), className='text-center p-1', style={'color': 'green', 'opacity': 0.7}, + ), + width=3, + ), + ], + className='bg-light rounded-bottom border-bottom border-left border-right', + ), + ] +stats_layout += [ + dbc.Row(dbc.Col(html.H5(children='Alphabet'), className='text-secondary'), className='mt-3'), + dbc.Row( + dbc.Col(html.Div('{}'.format(sorted(alphabet))),), className='mt-2 bg-light text-monospace rounded border' + ), + dbc.Row(dbc.Col(html.H5('Duration (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='duration-graph', figure=figure_duration),),), + dbc.Row(dbc.Col(html.H5('Number of words (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='num-words-graph', figure=figure_num_words),),), + dbc.Row(dbc.Col(html.H5('Number of characters (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='num-chars-graph', figure=figure_num_chars),),), + dbc.Row(dbc.Col(html.H5('Word rate (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='word-rate-graph', figure=figure_word_rate),),), + dbc.Row(dbc.Col(html.H5('Character rate (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='char-rate-graph', figure=figure_char_rate),),), +] + +if metrics_available: + stats_layout += [ + dbc.Row(dbc.Col(html.H5('WER (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='wer-graph', figure=figure_wer),),), + dbc.Row(dbc.Col(html.H5('CER (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='cer-graph', figure=figure_cer),),), + dbc.Row(dbc.Col(html.H5('WMR (per utterance)'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='wmr-graph', figure=figure_wmr),),), + dbc.Row(dbc.Col(html.H5('Word accuracy distribution'), className='text-secondary'), className='mt-3'), + dbc.Row(dbc.Col(dcc.Graph(id='word-acc-graph', figure=figure_word_acc),),), + ] + +wordstable_columns = [{'name': 'Word', 'id': 'word'}, {'name': 'Count', 'id': 'count'}] +if args.vocab is not None: + wordstable_columns.append({'name': 'OOV', 'id': 'OOV'}) +if metrics_available: + wordstable_columns.append({'name': 'Accuracy, %', 'id': 'accuracy'}) + +stats_layout += [ + dbc.Row(dbc.Col(html.H5('Vocabulary'), className='text-secondary'), className='mt-3'), + dbc.Row( + dbc.Col( + dash_table.DataTable( + id='wordstable', + columns=wordstable_columns, + filter_action='custom', + filter_query='', + sort_action='custom', + sort_mode='single', + page_action='custom', + page_current=0, + page_size=DATA_PAGE_SIZE, + cell_selectable=False, + page_count=math.ceil(len(vocabulary) / DATA_PAGE_SIZE), + sort_by=[{'column_id': 'word', 'direction': 'asc'}], + style_cell={'maxWidth': 0, 'textAlign': 'left'}, + style_header={'color': 'text-primary'}, + ), + ), + className='m-2', + ), +] + + +@app.callback( + [Output('wordstable', 'data'), Output('wordstable', 'page_count')], + [Input('wordstable', 'page_current'), Input('wordstable', 'sort_by'), Input('wordstable', 'filter_query')], +) +def update_wordstable(page_current, sort_by, filter_query): + vocabulary_view = vocabulary + filtering_expressions = filter_query.split(' && ') + for filter_part in filtering_expressions: + col_name, op, filter_value = split_filter_part(filter_part) + + if op in ('eq', 'ne', 'lt', 'le', 'gt', 'ge'): + vocabulary_view = [x for x in vocabulary_view if getattr(operator, op)(x[col_name], filter_value)] + print(len(vocabulary_view), len(vocabulary)) + elif op == 'contains': + vocabulary_view = [x for x in vocabulary_view if filter_value in str(x[col_name])] + print(len(vocabulary_view), len(vocabulary)) + + if len(sort_by): + col = sort_by[0]['column_id'] + descending = sort_by[0]['direction'] == 'desc' + vocabulary_view = sorted(vocabulary_view, key=lambda x: x[col], reverse=descending) + if page_current * DATA_PAGE_SIZE >= len(vocabulary_view): + page_current = len(vocabulary_view) // DATA_PAGE_SIZE + return [ + vocabulary_view[page_current * DATA_PAGE_SIZE : (page_current + 1) * DATA_PAGE_SIZE], + math.ceil(len(vocabulary_view) / DATA_PAGE_SIZE), + ] + + +samples_layout = [ + dbc.Row(dbc.Col(html.H5('Data'), className='text-secondary'), className='mt-3'), + dbc.Row( + dbc.Col( + dash_table.DataTable( + id='datatable', + columns=[{'name': k.replace('_', ' '), 'id': k, 'hideable': True} for k in data[0]], + filter_action='custom', + filter_query='', + sort_action='custom', + sort_mode='single', + sort_by=[], + row_selectable='single', + selected_rows=[0], + page_action='custom', + page_current=0, + page_size=DATA_PAGE_SIZE, + page_count=math.ceil(len(data) / DATA_PAGE_SIZE), + style_cell={'overflow': 'hidden', 'textOverflow': 'ellipsis', 'maxWidth': 0, 'textAlign': 'center'}, + style_header={ + 'color': 'text-primary', + 'text_align': 'center', + 'height': 'auto', + 'whiteSpace': 'normal', + }, + css=[{'selector': '.dash-spreadsheet-menu', 'rule': 'position:absolute; bottom:-36px'}], + ), + ) + ), +] + [ + dbc.Row( + [ + dbc.Col( + html.Div(children=k.replace('_', ' ')), + width=2, + className='mt-1 bg-light text-monospace text-break small rounded border', + ), + dbc.Col(html.Div(id='_' + k), className='mt-1 bg-light text-monospace text-break small rounded border'), + ] + ) + for k in data[0] +] + +if metrics_available: + samples_layout += [ + dbc.Row( + [ + dbc.Col( + html.Div(children='diff'), + width=2, + className='mt-1 bg-light text-monospace text-break small rounded border', + ), + dbc.Col( + html.Iframe( + id='_diff', + sandbox='', + srcDoc='', + style={'border': 'none', 'width': '100%', 'height': '100%'}, + className='bg-light text-monospace text-break small', + ), + className='mt-1 bg-light text-monospace text-break small rounded border', + ), + ] + ) + ] +samples_layout += [ + dbc.Row(dbc.Col(html.Audio(id='player', controls=True),), className='mt-3 '), + dbc.Row(dbc.Col(dcc.Graph(id='signal-graph')), className='mt-3'), +] + + +@app.callback( + [Output('datatable', 'data'), Output('datatable', 'page_count')], + [Input('datatable', 'page_current'), Input('datatable', 'sort_by'), Input('datatable', 'filter_query')], +) +def update_datatable(page_current, sort_by, filter_query): + data_view = data + filtering_expressions = filter_query.split(' && ') + for filter_part in filtering_expressions: + col_name, op, filter_value = split_filter_part(filter_part) + + if op in ('eq', 'ne', 'lt', 'le', 'gt', 'ge'): + data_view = [x for x in data_view if getattr(operator, op)(x[col_name], filter_value)] + elif op == 'contains': + data_view = [x for x in data_view if filter_value in str(x[col_name])] + + if len(sort_by): + col = sort_by[0]['column_id'] + descending = sort_by[0]['direction'] == 'desc' + data_view = sorted(data_view, key=lambda x: x[col], reverse=descending) + if page_current * DATA_PAGE_SIZE >= len(data_view): + page_current = len(data_view) // DATA_PAGE_SIZE + return [ + data_view[page_current * DATA_PAGE_SIZE : (page_current + 1) * DATA_PAGE_SIZE], + math.ceil(len(data_view) / DATA_PAGE_SIZE), + ] + + +app.layout = html.Div( + [ + dcc.Location(id='url', refresh=False), + dbc.NavbarSimple( + children=[ + dbc.NavItem(dbc.NavLink('Statistics', id='stats_link', href='/', active=True)), + dbc.NavItem(dbc.NavLink('Samples', id='samples_link', href='/samples')), + ], + brand='Speech Data Explorer', + sticky='top', + color='green', + dark=True, + ), + dbc.Container(id='page-content'), + ] +) + + +@app.callback( + [Output('page-content', 'children'), Output('stats_link', 'active'), Output('samples_link', 'active')], + [Input('url', 'pathname')], +) +def nav_click(url): + if url == '/samples': + return [samples_layout, False, True] + else: + return [stats_layout, True, False] + + +@app.callback( + [Output('_' + k, 'children') for k in data[0]], [Input('datatable', 'selected_rows'), Input('datatable', 'data')] +) +def show_item(idx, data): + if len(idx) == 0: + raise PreventUpdate + return [data[idx[0]][k] for k in data[0]] + + +@app.callback(Output('_diff', 'srcDoc'), [Input('datatable', 'selected_rows'), Input('datatable', 'data')]) +def show_diff(idx, data): + if len(idx) == 0: + raise PreventUpdate + + orig_words = data[idx[0]]['text'] + orig_words = '\n'.join(orig_words.split()) + '\n' + + pred_words = data[idx[0]]['pred_text'] + pred_words = '\n'.join(pred_words.split()) + '\n' + + diff = diff_match_patch.diff_match_patch() + diff.Diff_Timeout = 0 + orig_enc, pred_enc, enc = diff.diff_linesToChars(orig_words, pred_words) + diffs = diff.diff_main(orig_enc, pred_enc, False) + diff.diff_charsToLines(diffs, enc) + diffs_post = [] + for d in diffs: + diffs_post.append((d[0], d[1].replace('\n', ' '))) + + diff_html = diff.diff_prettyHtml(diffs_post) + + return diff_html + + +@app.callback(Output('signal-graph', 'figure'), [Input('datatable', 'selected_rows'), Input('datatable', 'data')]) +def plot_signal(idx, data): + if len(idx) == 0: + raise PreventUpdate + figs = make_subplots(rows=2, cols=1, subplot_titles=('Waveform', 'Spectrogram')) + try: + filename = data[idx[0]]['audio_filepath'] + audio, fs = librosa.load(filename, sr=None) + if 'offset' in data[idx[0]]: + audio = audio[ + int(data[idx[0]]['offset'] * fs) : int((data[idx[0]]['offset'] + data[idx[0]]['duration']) * fs) + ] + time_stride = 0.01 + hop_length = int(fs * time_stride) + n_fft = 512 + # linear scale spectrogram + s = librosa.stft(y=audio, n_fft=n_fft, hop_length=hop_length) + s_db = librosa.power_to_db(np.abs(s) ** 2, ref=np.max, top_db=100) + figs.add_trace( + go.Scatter( + x=np.arange(audio.shape[0]) / fs, + y=audio, + line={'color': 'green'}, + name='Waveform', + hovertemplate='Time: %{x:.2f} s
Amplitude: %{y:.2f}
', + ), + row=1, + col=1, + ) + figs.add_trace( + go.Heatmap( + z=s_db, + colorscale=[[0, 'rgb(30,62,62)'], [0.5, 'rgb(30,128,128)'], [1, 'rgb(30,255,30)'],], + colorbar=dict(yanchor='middle', lenmode='fraction', y=0.2, len=0.5, ticksuffix=' dB'), + dx=time_stride, + dy=fs / n_fft / 1000, + name='Spectrogram', + hovertemplate='Time: %{x:.2f} s
Frequency: %{y:.2f} kHz
Magnitude: %{z:.2f} dB', + ), + row=2, + col=1, + ) + figs.update_layout({'margin': dict(l=0, r=0, t=20, b=0, pad=0), 'height': 500}) + figs.update_xaxes(title_text='Time, s', row=1, col=1) + figs.update_yaxes(title_text='Amplitude', row=1, col=1) + figs.update_xaxes(title_text='Time, s', row=2, col=1) + figs.update_yaxes(title_text='Frequency, kHz', row=2, col=1) + except Exception: + pass + + return figs + + +@app.callback(Output('player', 'src'), [Input('datatable', 'selected_rows'), Input('datatable', 'data')]) +def update_player(idx, data): + if len(idx) == 0: + raise PreventUpdate + try: + filename = data[idx[0]]['audio_filepath'] + signal, sr = librosa.load(filename, sr=None) + if 'offset' in data[idx[0]]: + signal = signal[ + int(data[idx[0]]['offset'] * sr) : int((data[idx[0]]['offset'] + data[idx[0]]['duration']) * sr) + ] + with io.BytesIO() as buf: + # convert to PCM .wav + sf.write(buf, signal, sr, format='WAV') + buf.seek(0) + encoded = base64.b64encode(buf.read()) + return 'data:audio/wav;base64,{}'.format(encoded.decode()) + except Exception: + return '' + + +if __name__ == '__main__': + app.run_server(host='0.0.0.0', port=args.port, debug=args.debug) diff --git a/models/tts/NeMo/tools/speech_data_explorer/requirements.txt b/models/tts/NeMo/tools/speech_data_explorer/requirements.txt new file mode 100644 index 0000000000..54d15352c5 --- /dev/null +++ b/models/tts/NeMo/tools/speech_data_explorer/requirements.txt @@ -0,0 +1,10 @@ +librosa +dash +dash_html_components +dash_bootstrap_components +plotly +dash_core_components +dash_table +numpy +diff_match_patch +editdistance diff --git a/models/tts/NeMo/tools/speech_data_explorer/screenshot.png b/models/tts/NeMo/tools/speech_data_explorer/screenshot.png new file mode 100755 index 0000000000..c892347bc3 Binary files /dev/null and b/models/tts/NeMo/tools/speech_data_explorer/screenshot.png differ diff --git a/models/tts/NeMo/tools/text_processing_deployment/Dockerfile b/models/tts/NeMo/tools/text_processing_deployment/Dockerfile new file mode 100644 index 0000000000..9b2a4ae3e8 --- /dev/null +++ b/models/tts/NeMo/tools/text_processing_deployment/Dockerfile @@ -0,0 +1,38 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +# Dockerfile for C++ (inverse) text normalization backend Sparrowhawk https://github.com/google/sparrowhawk + +# set base image (host OS) +FROM conda/miniconda3 + +# set the working directory in the container +WORKDIR /workspace + +# install dependencies +RUN conda update -n base -c defaults conda -y +RUN conda install conda-build -y +RUN apt-get update && apt-get install -y --reinstall build-essential && apt-get upgrade -y && apt-get install -y git && apt-get install make +RUN git clone https://github.com/google/re2 +RUN cd re2 && make && make install +RUN apt-get install build-essential -y && apt-get install wget -y +RUN wget https://github.com/protocolbuffers/protobuf/releases/download/v2.5.0/protobuf-2.5.0.tar.gz +RUN tar xzvf protobuf-2.5.0.tar.gz +RUN cd protobuf-2.5.0 && ./configure && make && make install && ldconfig +RUN conda install -c conda-forge thrax=1.3.4 -y +RUN git clone https://github.com/yzhang123/sparrowhawk.git +RUN cd sparrowhawk && git checkout test && apt-get install -y autoconf && bash autoreconf && ./configure && make && make install && ldconfig +RUN git clone https://github.com/kward/shunit2.git +RUN echo "DONE" diff --git a/models/tts/NeMo/tools/text_processing_deployment/README.rst b/models/tts/NeMo/tools/text_processing_deployment/README.rst new file mode 100644 index 0000000000..5c1df016d0 --- /dev/null +++ b/models/tts/NeMo/tools/text_processing_deployment/README.rst @@ -0,0 +1,10 @@ +**Text Processing Deployment** +========================================= + +Introduction +------------ + +This folder provides scripts to deploy WFST-based grammars in `NeMo Text Processing `_ for +for production. + +See `documentation `_ for details. \ No newline at end of file diff --git a/models/tts/NeMo/tools/text_processing_deployment/docker/build.sh b/models/tts/NeMo/tools/text_processing_deployment/docker/build.sh new file mode 100644 index 0000000000..fbca6722b5 --- /dev/null +++ b/models/tts/NeMo/tools/text_processing_deployment/docker/build.sh @@ -0,0 +1,18 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + + +#!/bin/bash + +docker build . --rm -t sparrowhawk \ No newline at end of file diff --git a/models/tts/NeMo/tools/text_processing_deployment/docker/launch.sh b/models/tts/NeMo/tools/text_processing_deployment/docker/launch.sh new file mode 100644 index 0000000000..ddc35cacf7 --- /dev/null +++ b/models/tts/NeMo/tools/text_processing_deployment/docker/launch.sh @@ -0,0 +1,44 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. 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. + +#!/bin/bash + +MODE=${1:-"export"} +SCRIPT_DIR=$(cd $(dirname $0); pwd) +: ${CLASSIFY_DIR:="$SCRIPT_DIR/../classify"} +: ${VERBALIZE_DIR:="$SCRIPT_DIR/../verbalize"} +: ${CMD:=${2:-"/bin/bash"}} + +MOUNTS="" +MOUNTS+=" -v $CLASSIFY_DIR:/workspace/sparrowhawk/documentation/grammars/en_toy/classify" +MOUNTS+=" -v $VERBALIZE_DIR:/workspace/sparrowhawk/documentation/grammars/en_toy/verbalize" + +WORK_DIR="/workspace/sparrowhawk/documentation/grammars" +if [[ $MODE == "test_tn_grammars" ]]; then + CMD="bash test_sparrowhawk_normalization.sh" + WORK_DIR="/workspace/tests" +elif [[ $MODE == "test_itn_grammars" ]]; then + CMD="bash test_sparrowhawk_inverse_text_normalization.sh" + WORK_DIR="/workspace/tests" +fi + +echo $MOUNTS +docker run -it --rm \ + --shm-size=4g \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + $MOUNTS \ + -v $SCRIPT_DIR/../../../tests/nemo_text_processing/:/workspace/tests \ + -w $WORK_DIR \ + sparrowhawk:latest $CMD \ No newline at end of file diff --git a/models/tts/NeMo/tools/text_processing_deployment/export_grammars.sh b/models/tts/NeMo/tools/text_processing_deployment/export_grammars.sh new file mode 100644 index 0000000000..393fca3d5e --- /dev/null +++ b/models/tts/NeMo/tools/text_processing_deployment/export_grammars.sh @@ -0,0 +1,59 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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. + +#!/bin/bash + +# This script compiles and exports WFST-grammars from nemo_text_processing, builds C++ production backend Sparrowhawk (https://github.com/google/sparrowhawk) in docker, +# plugs grammars into Sparrowhawk and returns prompt inside docker. +# For inverse text normalization run: +# bash export_grammars.sh --GRAMMARS=itn_grammars +# echo "two dollars fifty" | ../../src/bin/normalizer_main --config=sparrowhawk_configuration.ascii_proto +# For text normalization run: +# bash export_grammars.sh --GRAMMARS=tn_grammars +# echo "\$2.5" | ../../src/bin/normalizer_main --config=sparrowhawk_configuration.ascii_proto +# +# To test TN grammars, run: +# bash export_grammars.sh --GRAMMARS=tn_grammars --MODE=test +# +# To test ITN grammars, run: +# bash export_grammars.sh --GRAMMARS=itn_grammars --MODE=test + +GRAMMARS="itn_grammars" # tn_grammars +INPUT_CASE="cased" # lower_cased, only for tn_grammars +MODE="" + +for ARG in "$@" +do + key=$(echo $ARG | cut -f1 -d=) + value=$(echo $ARG | cut -f2 -d=) + + if [[ $key == *"--"* ]]; then + v="${key/--/}" + declare $v="${value}" + fi +done + +echo "GRAMMARS = $GRAMMARS" +echo "MODE = $MODE" + +python pynini_export.py --output_dir=. --grammars=${GRAMMARS} --input_case=${INPUT_CASE} || exit 1 +find . -name "Makefile" -type f -delete +bash docker/build.sh + +if [[ $MODE == "test" ]]; then + MODE=${MODE}_${GRAMMARS} +fi + +bash docker/launch.sh $MODE \ No newline at end of file diff --git a/models/tts/NeMo/tools/text_processing_deployment/pynini_export.py b/models/tts/NeMo/tools/text_processing_deployment/pynini_export.py new file mode 100644 index 0000000000..80ccace9f0 --- /dev/null +++ b/models/tts/NeMo/tools/text_processing_deployment/pynini_export.py @@ -0,0 +1,111 @@ +# Copyright (c) 2021, NVIDIA CORPORATION. All rights reserved. +# Copyright 2015 and onwards Google, Inc. +# +# 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 os +import time +from argparse import ArgumentParser +from typing import Dict + +from nemo_text_processing.inverse_text_normalization.taggers.tokenize_and_classify import ClassifyFst as ITNClassifyFst +from nemo_text_processing.inverse_text_normalization.verbalizers.verbalize import VerbalizeFst as ITNVerbalizeFst +from nemo_text_processing.text_normalization.taggers.tokenize_and_classify import ClassifyFst as TNClassifyFst +from nemo_text_processing.text_normalization.verbalizers.verbalize import VerbalizeFst as TNVerbalizeFst + +from nemo.utils import logging + +try: + import pynini + from pynini.export import export + + PYNINI_AVAILABLE = True +except (ModuleNotFoundError, ImportError): + + logging.warning( + "`pynini` is not installed ! \n" + "Please run the `nemo_text_processing/setup.sh` script" + "prior to usage of this toolkit." + ) + + PYNINI_AVAILABLE = False + + +# This script exports compiled grammars inside nemo_text_processing into OpenFst finite state archive files tokenize_and_classify.far and verbalize.far for production purposes + + +def _generator_main(file_name: str, graphs: Dict[str, pynini.FstLike]): + """ + Exports graph as OpenFst finite state archive (FAR) file with given file name and rule name. + + Args: + file_name: exported file name + graph: Pynini WFST graph to be exported + rule_name: rule name for graph in created FAR file + + """ + exporter = export.Exporter(file_name) + for rule, graph in graphs.items(): + exporter[rule] = graph.optimize() + exporter.close() + print(f'Created {file_name}') + + +def itn_grammars(**kwargs): + d = {} + d['classify'] = {'TOKENIZE_AND_CLASSIFY': ITNClassifyFst().fst} + d['verbalize'] = {'ALL': ITNVerbalizeFst().fst, 'REDUP': pynini.accep("REDUP")} + return d + + +def tn_grammars(**kwargs): + d = {} + d['classify'] = {'TOKENIZE_AND_CLASSIFY': TNClassifyFst(input_case=kwargs["input_case"], deterministic=True).fst} + d['verbalize'] = {'ALL': TNVerbalizeFst(deterministic=True).fst, 'REDUP': pynini.accep("REDUP")} + return d + + +def export_grammars(output_dir, grammars): + """ + Exports tokenizer_and_classify and verbalize Fsts as OpenFst finite state archive (FAR) files. + + Args: + output_dir: directory to export FAR files to. Subdirectories will be created for tagger and verbalizer respectively. + """ + + for category, graphs in grammars.items(): + out_dir = os.path.join(output_dir, category) + if not os.path.exists(out_dir): + os.makedirs(out_dir) + time.sleep(1) + if category == "classify": + category = "tokenize_and_classify" + _generator_main(f"{out_dir}/{category}.far", graphs) + + +def parse_args(): + parser = ArgumentParser() + parser.add_argument("--output_dir", help="output directory for grammars", required=True, type=str) + parser.add_argument( + "--grammars", help="grammars to be exported", choices=["tn_grammars", "itn_grammars"], type=str, required=True + ) + parser.add_argument( + "--input_case", help="input capitalization", choices=["lower_cased", "cased"], default="cased", type=str + ) + return parser.parse_args() + + +if __name__ == '__main__': + args = parse_args() + export_grammars(output_dir=args.output_dir, grammars=locals()[args.grammars](input_case=args.input_case)) diff --git a/models/tts/NeMo/tutorials/00_NeMo_Primer.ipynb b/models/tts/NeMo/tutorials/00_NeMo_Primer.ipynb new file mode 100644 index 0000000000..db3bdfa814 --- /dev/null +++ b/models/tts/NeMo/tutorials/00_NeMo_Primer.ipynb @@ -0,0 +1,1281 @@ +{ + "nbformat": 4, + "nbformat_minor": 0, + "metadata": { + "colab": { + "name": "00_NeMo_Primer.ipynb", + "provenance": [], + "collapsed_sections": [], + "toc_visible": true + }, + "kernelspec": { + "display_name": "Python 3", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.8.5" + } + }, + "cells": [ + { + "cell_type": "markdown", + "metadata": { + "id": "7LfkL2r2Q1tr" + }, + "source": [ + "# Getting Started: Exploring Nemo Fundamentals\n", + "\n", + "NeMo is a toolkit for creating [Conversational AI](https://developer.nvidia.com/conversational-ai#started) applications.\n", + "\n", + "NeMo toolkit makes it possible for researchers to easily compose complex neural network architectures for conversational AI using reusable components - Neural Modules. Neural Modules are conceptual blocks of neural networks that take typed inputs and produce typed outputs. Such modules typically represent data layers, encoders, decoders, language models, loss functions, or methods of combining activations.\n", + "\n", + "The toolkit comes with extendable collections of pre-built modules and ready-to-use models for automatic speech recognition (ASR), natural language processing (NLP) and text synthesis (TTS). Built for speed, NeMo can utilize NVIDIA's Tensor Cores and scale out training to multiple GPUs and multiple nodes.\n", + "\n", + "For more information, please visit https://docs.nvidia.com/deeplearning/nemo/user-guide/docs/en/main/#" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "zLSy94NEQi-e" + }, + "source": [ + "\"\"\"\n", + "You can run either this notebook locally (if you have all the dependencies and a GPU) or on Google Colab.\n", + "\n", + "Instructions for setting up Colab are as follows:\n", + "1. Open a new Python 3 notebook.\n", + "2. Import this notebook from GitHub (File -> Upload Notebook -> \"GITHUB\" tab -> copy/paste GitHub URL)\n", + "3. Connect to an instance with a GPU (Runtime -> Change runtime type -> select \"GPU\" for hardware accelerator)\n", + "4. Run this cell to set up dependencies.\n", + "\"\"\"\n", + "# If you're using Google Colab and not running locally, run this cell.\n", + "\n", + "## Install dependencies\n", + "!pip install wget\n", + "!apt-get install sox libsndfile1 ffmpeg\n", + "!pip install unidecode\n", + "\n", + "# ## Install NeMo\n", + "BRANCH = 'v1.0.2'\n", + "!python -m pip install git+https://github.com/NVIDIA/NeMo.git@$BRANCH#egg=nemo_toolkit[all]\n", + "\n", + "## Install TorchAudio\n", + "!pip install torchaudio>=0.6.0 -f https://download.pytorch.org/whl/torch_stable.html\n", + "\n", + "## Grab the config we'll use in this example\n", + "!mkdir configs" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6G2TZkaxcM0e" + }, + "source": [ + "## Foundations of NeMo\n", + "---------\n", + "\n", + "NeMo models leverage [PyTorch Lightning](https://github.com/PyTorchLightning/pytorch-lightning) Module, and are compatible with the entire PyTorch ecosystem. This means that users have the full flexibility of using the higher level APIs provided by PyTorch Lightning (via Trainer), or write their own training and evaluation loops in PyTorch directly (by simply calling the model and the individual components of the model).\n", + "\n", + "For NeMo developers, a \"Model\" is the neural network(s) as well as all the infrastructure supporting those network(s), wrapped into a singular, cohesive unit. As such, all NeMo models are constructed to contain the following out of the box (at the bare minimum, some models support additional functionality too!) - \n", + "\n", + " - Neural Network architecture - all of the modules that are required for the model.\n", + "\n", + " - Dataset + Data Loaders - all of the components that prepare the data for consumption during training or evaluation.\n", + "\n", + " - Preprocessing + Postprocessing - all of the components that process the datasets so they can easily be consumed by the modules.\n", + "\n", + " - Optimizer + Schedulers - basic defaults that work out of the box, and allow further experimentation with ease.\n", + "\n", + " - Any other supporting infrastructure - tokenizers, language model configuration, data augmentation etc.\n", + "\n", + "\n" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "XxAwtqWBQrNk" + }, + "source": [ + "import nemo\n", + "nemo.__version__" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "H01SHfKQh-gV" + }, + "source": [ + "## NeMo Collections\n", + "\n", + "NeMo is sub-divided into a few fundamental collections based on their domains - `asr`, `nlp`, `tts`. When you performed the `import nemo` statement above, none of the above collections were imported. This is because you might not need all of the collections at once, so NeMo allows partial imports of just one or more collection, as and when you require them.\n", + "\n", + "-------\n", + "Let's import the above three collections - " + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "J09NNa8fhth7" + }, + "source": [ + "import nemo.collections.asr as nemo_asr\n", + "import nemo.collections.nlp as nemo_nlp\n", + "import nemo.collections.tts as nemo_tts" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "bSvYoeBrjPza" + }, + "source": [ + "## NeMo Models in Collections\n", + "\n", + "NeMo contains several models for each of its collections, pertaining to certain common tasks involved in conversational AI. At a brief glance, let's look at all the Models that NeMo offers for the above 3 collections." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "9LbbC_92i41f" + }, + "source": [ + "asr_models = [model for model in dir(nemo_asr.models) if model.endswith(\"Model\")]\n", + "asr_models" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "t5_ax9Z8j9FC" + }, + "source": [ + "nlp_models = [model for model in dir(nemo_nlp.models) if model.endswith(\"Model\")]\n", + "nlp_models" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "bQdR6RJdkezq" + }, + "source": [ + "tts_models = [model for model in dir(nemo_tts.models) if model.endswith(\"Model\")]\n", + "tts_models" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "iWKxKQnSkj9Z" + }, + "source": [ + "## The NeMo Model\n", + "\n", + "Let's dive deeper into what a NeMo model really is. There are many ways we can create these models - we can use the constructor and pass in a config, we can instantiate the model from a pre-trained checkpoint, or simply pass a pre-trained model name and instantiate a model directly from the cloud !\n", + "\n", + "---------\n", + "For now, let's try to work with an ASR model - [Citrinet](https://arxiv.org/abs/2104.01721)" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "n-XOQaW1kh3v" + }, + "source": [ + "citrinet = nemo_asr.models.EncDecCTCModelBPE.from_pretrained('stt_en_citrinet_512')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "YP4X7KVPli6g" + }, + "source": [ + "citrinet.summarize();" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "MB91Swu0pIKr" + }, + "source": [ + "## Model Configuration using OmegaConf\n", + "--------\n", + "\n", + "So we could download, instantiate and analyse the high level structure of the `Citrinet` model in a few lines! Now let's delve deeper into the configuration file that makes the model work.\n", + "\n", + "First, we import [OmegaConf](https://omegaconf.readthedocs.io/en/latest/). OmegaConf is an excellent library that is used throughout NeMo in order to enable us to perform yaml configuration management more easily. Additionally, it plays well with another library, [Hydra](https://hydra.cc/docs/intro/), that is used by NeMo to perform on the fly config edits from the command line, dramatically boosting ease of use of our config files !" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "RkgrDJvumFER" + }, + "source": [ + "from omegaconf import OmegaConf" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "CktakfBluA56" + }, + "source": [ + "All NeMo models come packaged with their model configuration inside the `cfg` attribute. While technically it is meant to be config declaration of the model as it has been currently constructed, `cfg` is an essential tool to modify the behaviour of the Model after it has been constructed. It can be safely used to make it easier to perform many essential tasks inside Models. \n", + "\n", + "To be doubly sure, we generally work on a copy of the config until we are ready to edit it inside the model" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "ISd6z7sXt9Mm" + }, + "source": [ + "import copy" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "N2_SiLHRve8A" + }, + "source": [ + "cfg = copy.deepcopy(citrinet.cfg)\n", + "print(OmegaConf.to_yaml(cfg))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W_V3e3W7vqOb" + }, + "source": [ + "## Analysing the contents of the Model config\n", + "----------\n", + "\n", + "Above we see a configuration for the Citrinet model. As discussed in the beginning, NeMo models contain the entire definition of the neural network(s) as well as most of the surrounding infrastructure to support that model within themselves. Here, we see a perfect example of this behaviour.\n", + "\n", + "Citrinet contains within its config - \n", + "\n", + "- `preprocessor` - MelSpectrogram preprocessing layer\n", + "- `encoder` - The acoustic encoder model.\n", + "- `decoder` - The CTC decoder layer.\n", + "- `optim` (and potentially `sched`) - Optimizer configuration. Can optionally include Scheduler information.\n", + "- `spec_augment` - Spectrogram Augmentation support.\n", + "- `train_ds`, `validation_ds` and `test_ds` - Dataset and data loader construction information." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "sIwhdXkwxn6R" + }, + "source": [ + "## Modifying the contents of the Model config\n", + "----------\n", + "\n", + "Say we want to experiment with a different preprocessor (we want MelSpectrogram, but with different configuration than was provided in the original configuration). Or say we want to add a scheduler to this model during training. \n", + "\n", + "OmegaConf makes this a very simple task for us!" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "WlSZ8EA4yGKo" + }, + "source": [ + "# OmegaConf won't allow you to add new config items, so we temporarily disable this safeguard.\n", + "OmegaConf.set_struct(cfg, False)\n", + "\n", + "# Let's see the old optim config\n", + "print(\"Old Config: \")\n", + "print(OmegaConf.to_yaml(cfg.optim))\n", + "\n", + "sched = {'name': 'CosineAnnealing', 'warmup_steps': 1000, 'min_lr': 1e-6}\n", + "sched = OmegaConf.create(sched) # Convert it into a DictConfig\n", + "\n", + "# Assign it to cfg.optim.sched namespace\n", + "cfg.optim.sched = sched\n", + "\n", + "# Let's see the new optim config\n", + "print(\"New Config: \")\n", + "print(OmegaConf.to_yaml(cfg.optim))\n", + "\n", + "# Here, we restore the safeguards so no more additions can be made to the config\n", + "OmegaConf.set_struct(cfg, True)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "-nMDN66502kn" + }, + "source": [ + "## Updating the model from config\n", + "----------\n", + "\n", + "NeMo Models can be updated in a few ways, but we follow similar patterns within each collection so as to maintain consistency.\n", + "\n", + "Here, we will show the two most common ways to modify core components of the model - using the `from_config_dict` method, and updating a few special parts of the model.\n", + "\n", + "Remember, all NeMo models are PyTorch Lightning modules, which themselves are PyTorch modules, so we have a lot of flexibility here!" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "qrKzFYkZ20aa" + }, + "source": [ + "### Update model using `from_config_dict`\n", + "\n", + "In certain config files, you will notice the following pattern : \n", + "\n", + "```yaml\n", + "preprocessor:\n", + " _target_: nemo.collections.asr.modules.AudioToMelSpectrogramPreprocessor\n", + " normalize: per_feature\n", + " window_size: 0.02\n", + " sample_rate: 16000\n", + " window_stride: 0.01\n", + " window: hann\n", + " features: 64\n", + " n_fft: 512\n", + " frame_splicing: 1\n", + " dither: 1.0e-05\n", + " stft_conv: false\n", + "```\n", + "\n", + "You might ask, why are we using `_target_`? Well, it is generally rare for the preprocessor, encoder, decoder and perhaps a few other details to be changed often from the command line when experimenting. In order to stabilize these settings, we enforce that our preprocessor will always be of type `AudioToMelSpectrogramPreprocessor` for this model by setting its `_target_` attribute in the config. In order to provide its parameters in the class constructor, we simply add them after `_target_`.\n", + "\n", + "---------\n", + "Note, we can still change all of the parameters of this `AudioToMelSpectrogramPreprocessor` class from the command line using hydra, so we don't lose any flexibility once we decide what type of preprocessing class we want !\n", + "\n", + "This also gives us a flexible way to instantiate parts of the model from just the config object !" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "1Be08R4szkT3" + }, + "source": [ + "new_preprocessor_config = copy.deepcopy(cfg.preprocessor)\n", + "new_preprocessor = citrinet.from_config_dict(new_preprocessor_config)\n", + "print(new_preprocessor)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "UzJQ7Y8H4S_U" + }, + "source": [ + "So how do we actually update our model's internal preprocessor with something new? Well, since NeMo Model's are just pytorch Modules, we can just replace their attribute !" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "WdtnPKX84OJ-" + }, + "source": [ + "citrinet.preprocessor = new_preprocessor" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "OMz2KR-24xTO" + }, + "source": [ + "citrinet.summarize();" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gPb_BdPN40Ro" + }, + "source": [ + "--------\n", + "This might look like nothing changed - because we didn't actually modify the config for the preprocessor at all ! But as we showed above, we can easily modify the config for the preprocessor, instantiate it from config, and then just set it to the model." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "IV8WKJkD5E_Q" + }, + "source": [ + "-------\n", + "**NOTE**: Preprocessors don't generally have weights, so this was easy, but say we want to replace a part of the model which actually has trained parameters?\n", + "\n", + "Well, the above approach will still work, just remember the fact that the new module you inserted into `citrinet.encoder` or `citrinet.decoder` actually won't have pretrained weights. You can easily rectify that by loading the state dict for the module *before* you set it to the Model though!" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "YplQcgfG6S1U" + }, + "source": [ + "### Preserving the new config\n", + "\n", + "So we went ahead and updated the preprocessor of the model. We however also need to perform a crucial step - **preserving the updated config**!\n", + "\n", + "Why do we want to do this? NeMo has many ways of saving and restoring its models, which we will discuss a bit later. All of them depend on having an updated config that defines the model in its entirety, so if we modify anything, we should also update the corresponding part of the config to safely save and restore models." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "dsxQHBV86R4a" + }, + "source": [ + "# Update the config copy\n", + "cfg.preprocessor = new_preprocessor_config\n", + "# Update the model config\n", + "citrinet.cfg = cfg" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "eXRRBnJk5tCv" + }, + "source": [ + "## Update a few special components of the Model\n", + "---------\n", + "\n", + "While the above approach is good for most major components of the model, NeMo has special utilities for a few components.\n", + "\n", + "They are - \n", + "\n", + " - `setup_training_data`\n", + " - `setup_validation_data` and `setup_multi_validation_data`\n", + " - `setup_test_data` and `setup_multi_test_data`\n", + " - `setup_optimization`\n", + "\n", + "These special utilities are meant to help you easily setup training, validation, testing once you restore a model from a checkpoint.\n", + "\n", + "------\n", + "One of the major tasks of all conversational AI models is fine-tuning onto new datasets - new languages, new corpus of text, new voices etc. It is often insufficient to have just a pre-trained model. So these setup methods are provided to enable users to adapt models *after* they have been already trained or provided to you.\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "B7Y7wt2x9goJ" + }, + "source": [ + "You might remember having seen a few warning messages the moment you tried to instantiate the pre-trained model. Those warnings are in fact reminders to call the appropriate setup methods for the task you want to perform. \n", + "\n", + "Those warnings are simply displaying the old config that was used to train that model, and are a basic template that you can easily modify. You have the ability to modify the `train_ds`, `validation_ds` and `test_ds` sub-configs in their entirety in order to evaluate, fine-tune or train from scratch the model, or any further purpose as you require it.\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1hXXdaup-QmG" + }, + "source": [ + "Let's discuss how to add the scheduler to the model below (which initially had just an optimizer in its config)" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "cveKWvMZ4zBo" + }, + "source": [ + "# Let's print out the current optimizer\n", + "print(OmegaConf.to_yaml(citrinet.cfg.optim))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "XVguw3k0-f6b" + }, + "source": [ + "# Now let's update the config\n", + "citrinet.setup_optimization(cfg.optim);" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "1JZBCQeW-21X" + }, + "source": [ + "-------\n", + "We see a warning - \n", + "\n", + "```\n", + "Neither `max_steps` nor `iters_per_batch` were provided to `optim.sched`, cannot compute effective `max_steps` !\n", + " Scheduler will not be instantiated !\n", + "```\n", + "\n", + "We don't have a train dataset setup, nor do we have max_steps in the config. Most NeMo schedulers cannot be instantiated without computing how many train steps actually exist!\n", + "\n", + "Here, we can temporarily allow the scheduler construction by explicitly passing a max_steps value to be 100" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "mqC89hfE-tqf" + }, + "source": [ + "OmegaConf.set_struct(cfg.optim.sched, False)\n", + "\n", + "cfg.optim.sched.max_steps = 100\n", + "\n", + "OmegaConf.set_struct(cfg.optim.sched, True)" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "r22IqOBK_q6l" + }, + "source": [ + "# Now let's update the config and try again\n", + "citrinet.setup_optimization(cfg.optim);" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "U7Eezf_sAVS0" + }, + "source": [ + "You might wonder why we didnt explicitly set `citrinet.cfg.optim = cfg.optim`. \n", + "\n", + "This is because the `setup_optimization()` method does it for you! You can still update the config manually." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "THqhXy_lQ7i8" + }, + "source": [ + "### Optimizer & Scheduler Config\n", + "\n", + "Optimizers and schedulers are common components of models, and are essential to train the model from scratch.\n", + "\n", + "They are grouped together under a unified `optim` namespace, as schedulers often operate on a given optimizer.\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "6HY51nuoSJs5" + }, + "source": [ + "### Let's breakdown the general `optim` structure\n", + "```yaml\n", + "optim:\n", + " name: novograd\n", + " lr: 0.01\n", + "\n", + " # optimizer arguments\n", + " betas: [0.8, 0.25]\n", + " weight_decay: 0.001\n", + "\n", + " # scheduler setup\n", + " sched:\n", + " name: CosineAnnealing\n", + "\n", + " # Optional arguments\n", + " max_steps: null # computed at runtime or explicitly set here\n", + " monitor: val_loss\n", + " reduce_on_plateau: false\n", + "\n", + " # scheduler config override\n", + " warmup_steps: 1000\n", + " warmup_ratio: null\n", + " min_lr: 1e-9\n", + "```\n", + "\n", + "Essential Optimizer components - \n", + "\n", + " - `name`: String name of the optimizer. Generally a lower case of the class name.\n", + " - `lr`: Learning rate is a required argument to all optimizers.\n", + "\n", + "Optional Optimizer components - after the above two arguments are provided, any additional arguments added under `optim` will be passed to the constructor of that optimizer as keyword arguments\n", + "\n", + " - `betas`: List of beta values to pass to the optimizer\n", + " - `weight_decay`: Optional weight decay passed to the optimizer.\n", + "\n", + "Optional Scheduler components - `sched` is an optional setup of the scheduler for the given optimizer.\n", + "\n", + "If `sched` is provided, only one essential argument needs to be provided : \n", + "\n", + " - `name`: The name of the scheduler. Generally, it is the full class name.\n", + "\n", + "Optional Scheduler components - \n", + "\n", + " - `max_steps`: Max steps as an override from the user. If one provides `trainer.max_steps` inside the trainer configuration, that value is used instead. If neither value is set, the scheduler will attempt to compute the `effective max_steps` using the size of the train data loader. If that too fails, then the scheduler will not be created at all.\n", + "\n", + " - `monitor`: Used if you are using an adaptive scheduler such as ReduceLROnPlateau. Otherwise ignored. Defaults to `loss` - indicating train loss as monitor.\n", + "\n", + " - `reduce_on_plateau`: Required to be set to true if using an adaptive scheduler.\n", + "\n", + "Any additional arguments under `sched` will be supplied as keyword arguments to the constructor of the scheduler.\n", + "\n", + "\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "V3pQM2aj_6WX" + }, + "source": [ + "## Difference between the data loader setup methods\n", + "----------\n", + "\n", + "You might notice, we have multiple setup methods for validation and test data sets. We also don't have an equivalent `setup_multi_train_data`. \n", + "\n", + "In general, the `multi` methods refer to multiple data sets / data loaders. \n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "g33nMx9WCJdj" + }, + "source": [ + "### Where's `setup_multi_train_data`?\n", + "With the above in mind, let's tackle why we don't have `setup_multi_train_data`. \n", + "\n", + "NeMo is concerned with multiple domains - `asr`, `nlp` and `tts`. The way datasets are setup and used in these domains is dramatically different. It is often unclear what it means to have multiple train datasets - do we concatenate them? Do we randomly sample (with same or different probability) from each of them? \n", + "\n", + "Therefore we leave such support for multiple datasets up to the model itself. For example, in ASR, you can concatenate multiple train manifest files by using commas when providing the `manifest_filepath` value!" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "BjI2Q5LECJib" + }, + "source": [ + "### What are multi methods?\n", + "\n", + "In many cases, especially true for ASR and NLP, we may have multiple validation and test datasets. The most common example for this in ASR is `Librispeech`, which has `dev_clean`, `dev_other`, `test_clean`, `test_other`.\n", + "\n", + "NeMo standardizes how to handle multiple data loaders for validation and testing, so that all of our collections have a similar look and feel, as well as ease development of our models. During evaluation, these datasets are treated independently and prepended with resolved names so that logs are separate!\n", + "\n", + "The `multi` methods are therefore generalizations of the single validation and single test data setup methods, with some additional functionality. If you provide multiple datasets, you still have to write code for just one dataset and NeMo will automatically attach the appropriate names to your logs so you can differentiate between them!\n", + "\n", + "Furthermore, they also automatically preserve the config the user passes to them when updating the validation or test data loaders.\n", + "\n", + "**In general, it is preferred to call the `setup_multi_validation_data` and `setup_multi_test_data` methods, even if you are only using single datasets, simply for the automated management they provide.**" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ZKURHn0jH_52" + }, + "source": [ + "## Creating Model from constructor vs restoring a model\n", + "---------\n", + "\n", + "You might notice, we discuss all of the above setup methods in the context of model after it is restored. However, NeMo scripts do not call them inside any of the example train scripts themselves.\n", + "\n", + "This is because these methods are automatically called by the constructor when the Model is created for the first time, but these methods are skipped during restoration (either from a PyTorch Lightning checkpoint using `load_from_checkpoint`, or via `restore_from` method inside NeMo Models).\n", + "\n", + "This is done as most datasets are stored on a user's local directory, and the path to these datasets is set in the config (either set by default, or set by Hydra overrides). On the other hand, the models are meant to be portable. On another user's system, the data might not be placed at exactly the same location, or even on the same drive as specified in the model's config!\n", + "\n", + "Therefore we allow the constructor some brevity and automate such dataset setup, whereas restoration warns that data loaders were not set up and provides the user with ways to set up their own datasets.\n", + "\n", + "------\n", + "\n", + "Why are optimizers not restored automatically? Well, optimizers themselves don't face an issue, but as we saw before, schedulers depend on the number of train steps in order to calculate their schedule.\n", + "\n", + "However, if you don't wish to modify the optimizer and scheduler, and prefer to leave them to their default values, that's perfectly alright. The `setup_optimization()` method is automatically called by PyTorch Lightning for you when you begin training your model!" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "g91FE8mlMcnh" + }, + "source": [ + "## Saving and restoring models\n", + "----------\n", + "\n", + "NeMo provides a few ways to save and restore models. If you utilize the Experiment Manager that is part of all NeMo train scripts, PyTorch Lightning will automatically save checkpoints for you in the experiment directory.\n", + "\n", + "We can also use packaged files using the specialized `save_to` and `restore_from` methods." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "NzMxga7QNYn8" + }, + "source": [ + "### Saving and Restoring from PTL Checkpoints\n", + "----------\n", + "\n", + "The PyTorch Lightning Trainer object will periodically save checkpoints when the experiment manager is being used during training.\n", + "\n", + "PyTorch Lightning checkpoints can then be loaded and evaluated / fine-tuned just as always using the class method `load_from_checkpoint`.\n", + "\n", + "For example, restore a Citrinet model from a checkpoint - \n", + "\n", + "```python\n", + "citrinet = nemo_asr.models.EncDecCTCModelBPE.load_from_checkpoint()\n", + "```" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "W4YzAG-KOBkZ" + }, + "source": [ + "### Saving and Restoring from .nemo files\n", + "----------\n", + "\n", + "There are a few models which might require external dependencies to be packaged with them in order to restore them properly.\n", + "\n", + "One such example is an ASR model with an external BPE tokenizer. It is preferred if the model includes all of the components required to restore it, but a binary file for a tokenizer cannot be serialized into a PyTorch Lightning checkpoint.\n", + "\n", + "In such cases, we can use the `save_to` and `restore_from` method to package the entire model + its components (here, the tokenizer file(s)) into a tarfile. This can then be easily imported by the user and used to restore the model." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "P6_vMSwXNJ74" + }, + "source": [ + "# Save the model\n", + "citrinet.save_to('citrinet_512.nemo')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "HrBhgaqyP4rU" + }, + "source": [ + "!ls -d -- *.nemo " + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "Tyht1E0DQGb_" + }, + "source": [ + "# Restore the model\n", + "temp_cn = nemo_asr.models.EncDecCTCModelBPE.restore_from('citrinet_512.nemo')" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "dqNpmYYJQS2H" + }, + "source": [ + "temp_cn.summarize();" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "A5e42EoiZYjf" + }, + "source": [ + "# Note that the preprocessor + optimizer config have been preserved after the changes we made !\n", + "print(OmegaConf.to_yaml(temp_cn.cfg))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "OI3RxwpcV-UF" + }, + "source": [ + "Note, that .nemo file is a simple .tar.gz with checkpoint, configuration and, potentially, other artifacts such as tokenizer configs being used by the model" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "jFBAGcaDWLiu" + }, + "source": [ + "!cp citrinet_512.nemo citrinet_512.tar.gz\n", + "!tar -xvf citrinet_512.tar.gz" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "mkau4Q9jZo1l" + }, + "source": [ + "### Extracting PyTorch checkpoints from NeMo tarfiles (Model level)\n", + "-----------\n", + "\n", + "While the .nemo tarfile is an excellent way to have a portable model, sometimes it is necessary for researchers to have access to the basic PyTorch save format. NeMo aims to be entirely compatible with PyTorch, and therefore offers a simple method to extract just the PyTorch checkpoint from the .nemo tarfile." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "qccPANeycCoq" + }, + "source": [ + "import torch" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "metadata": { + "id": "A4zswOKHar9q" + }, + "source": [ + "state_dict = temp_cn.extract_state_dict_from('citrinet_512.nemo', save_dir='./pt_ckpt/')\n", + "!ls ./pt_ckpt/" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ACB-0dfnbFG3" + }, + "source": [ + "As we can see below, there is now a single basic PyTorch checkpoint available inside the `pt_ckpt` directory, which we can use to load the weights of the entire model as below" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "4ZAF_A0uc5bB" + }, + "source": [ + "temp_cn.load_state_dict(torch.load('./pt_ckpt/model_weights.ckpt'))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "Hkq6EM99cS6y" + }, + "source": [ + "### Extracting PyTorch checkpoints from NeMo tarfiles (Module level)\n", + "----------\n", + "\n", + "While the above method is exceptional when extracting the checkpoint of the entire model, sometimes there may be a necessity to load and save the individual modules that comprise the Model.\n", + "\n", + "The same extraction method offers a flag to extract the individual model level checkpoints into their individual files, so that users have access to per-module level checkpoints." + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "LW6wve2zbT9D" + }, + "source": [ + "state_dict = temp_cn.extract_state_dict_from('citrinet_512.nemo', save_dir='./pt_module_ckpt/', split_by_module=True)\n", + "!ls ./pt_module_ckpt/" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DtV5vpb5d1ni" + }, + "source": [ + "Now, we can load and assign the weights of the individual modules of the above Citrinet Model !" + ] + }, + { + "cell_type": "code", + "metadata": { + "id": "rVHylSKFdywn" + }, + "source": [ + "temp_cn.preprocessor.load_state_dict(torch.load('./pt_module_ckpt/preprocessor.ckpt'))\n", + "temp_cn.encoder.load_state_dict(torch.load('./pt_module_ckpt/encoder.ckpt'))\n", + "temp_cn.decoder.load_state_dict(torch.load('./pt_module_ckpt/decoder.ckpt'))" + ], + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "88vOGV7VYcuu" + }, + "source": [ + "# NeMo with Hydra\n", + "\n", + "[Hydra](https://hydra.cc/docs/intro/) is used throughout NeMo as a way to enable rapid prototyping using predefined config files. Hydra and OmegaConf offer great compatibility with each other, and below we show a few general helpful tips to improve productivity with Hydra when using NeMo." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "DfY6Ha3qYcxG" + }, + "source": [ + "## Hydra Help\n", + "--------\n", + "\n", + "Since our scripts are written with hydra in mind, you might notice that using `python --help` returns you a config rather than the usual help format from argparse. \n", + "\n", + "Using `--help` you can see the default config attached to the script - every NeMo script has at least one default config file attached to it. This gives you a guide on how you can modify values for an experiment.\n", + "\n", + "Hydra also has a special `--hydra-help` flag, which will offer you more help with respect to hydra itself as it is set up in the script.\n" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "gEsZlnfaYc3X" + }, + "source": [ + "## Changing config paths and files\n", + "---------\n", + "\n", + "While all NeMo models come with at least 1 default config file, one might want to switch configs without changing code. This is easily achieved by the following commands : \n", + "\n", + "- `--config-path`: Path to the directory which contains the config files\n", + "- `--config-name`: Name of the config file we wish to load.\n", + "\n", + "Note that these two arguments need to be at the very beginning of your execution statement, before you provide any command line overrides to your config file." + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "ZyNHlArpYc9A" + }, + "source": [ + "## Overriding config from the command line\n", + "----------\n", + "\n", + "Hydra allows users to provide command line overrides to any part of the config. There are three cases to consider - \n", + "\n", + " - Override existing value in config\n", + " - Add new value in config\n", + " - Remove old value in config" + ] + }, + { + "cell_type": "markdown", + "metadata": { + "id": "96CKbvn6Yc7f" + }, + "source": [ + "### Overriding existing values in config\n", + "\n", + "Let's take the case where we want to change the optimizer from `novograd` to `adam`. Let's also change the beta values to default adam values.\n", + "\n", + "Hydra overrides are based on the `.` syntax - each `.` representing a level in the config itself.\n", + "\n", + "```sh\n", + "$ python + + +
+ + +
+ + + +
+ + + +
+ + +
+ + +''' + + +#class Converter(): +# def __init__(self, model_dir, model_name): +# self.num_features = 24 +# self.sampling_rate = 16000 +# self.frame_period = 5.0 +# +# self.model = CycleGAN(num_features = self.num_features, mode = 'test') +# +# self.model.load(filepath = os.path.join(model_dir, model_name)) +# +# """ +# # NB: Save the graph +# definition = self.model.sess.graph_def +# directory = 'saved_model_2' +# tf.train.write_graph(definition, directory, 'saved_model_2.pb', as_text=True) +# # https://github.com/tensorflow/models/issues/3530#issuecomment-395968881 +# output_dir = './saved_model/' +# builder = tf.saved_model.builder.SavedModelBuilder(output_dir) +# builder.add_meta_graph_and_variables( +# self.model.sess, +# [tf.saved_model.tag_constants.SERVING], +# main_op=tf.tables_initializer(), +# ) +# builder.save() +# """ +# +# """ +# builder.add_meta_graph_and_variables( +# self.model.sess, +# [tf.saved_model.tag_constants.SERVING], +# signature_def_map={ +# 'predict_images': +# prediction_signature, +# signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: +# classification_signature, +# }, +# main_op=tf.tables_initializer()) +# """ +# +# self.mcep_normalization_params = np.load(os.path.join(model_dir, 'mcep_normalization.npz')) +# self.mcep_mean_A = self.mcep_normalization_params['mean_A'] +# self.mcep_std_A = self.mcep_normalization_params['std_A'] +# self.mcep_mean_B = self.mcep_normalization_params['mean_B'] +# self.mcep_std_B = self.mcep_normalization_params['std_B'] +# +# self.logf0s_normalization_params = np.load(os.path.join(model_dir, 'logf0s_normalization.npz')) +# self.logf0s_mean_A = self.logf0s_normalization_params['mean_A'] +# self.logf0s_std_A = self.logf0s_normalization_params['std_A'] +# self.logf0s_mean_B = self.logf0s_normalization_params['mean_B'] +# self.logf0s_std_B = self.logf0s_normalization_params['std_B'] +# +# def convert(self, wav, conversion_direction='A2B'): +# wav = wav_padding(wav = wav, sr = self.sampling_rate, frame_period = self.frame_period, multiple = 4) +# f0, timeaxis, sp, ap = world_decompose(wav = wav, fs = self.sampling_rate, frame_period = self.frame_period) +# coded_sp = world_encode_spectral_envelop(sp = sp, fs = self.sampling_rate, dim = self.num_features) +# coded_sp_transposed = coded_sp.T +# +# if conversion_direction == 'A2B': +# f0_converted = pitch_conversion(f0 = f0, mean_log_src = self.logf0s_mean_A, std_log_src = self.logf0s_std_A, mean_log_target = self.logf0s_mean_B, std_log_target = self.logf0s_std_B) +# coded_sp_norm = (coded_sp_transposed - self.mcep_mean_A) / self.mcep_std_A +# coded_sp_converted_norm = self.model.test(inputs = np.array([coded_sp_norm]), direction = conversion_direction)[0] +# coded_sp_converted = coded_sp_converted_norm * self.mcep_std_B + self.mcep_mean_B +# else: +# f0_converted = pitch_conversion(f0 = f0, mean_log_src = self.logf0s_mean_B, std_log_src = self.logf0s_std_B, mean_log_target = self.logf0s_mean_A, std_log_target = self.logf0s_std_A) +# coded_sp_norm = (coded_sp_transposed - self.mcep_mean_B) / self.mcep_std_B +# coded_sp_converted_norm = self.model.test(inputs = np.array([coded_sp_norm]), direction = conversion_direction)[0] +# coded_sp_converted = coded_sp_converted_norm * self.mcep_std_A + self.mcep_mean_A +# +# coded_sp_converted = coded_sp_converted.T +# coded_sp_converted = np.ascontiguousarray(coded_sp_converted) +# decoded_sp_converted = world_decode_spectral_envelop(coded_sp = coded_sp_converted, fs = self.sampling_rate) +# wav_transformed = world_speech_synthesis(f0 = f0_converted, decoded_sp = decoded_sp_converted, ap = ap, fs = self.sampling_rate, frame_period = self.frame_period) +# +# # For debugging model output, uncomment the following line: +# # librosa.output.write_wav('model_output.wav', wav_transformed, self.sampling_rate) +# +# # TODO: Perhaps ditch this. It's probably unnecessary work. +# upsampled = librosa.resample(wav_transformed, self.sampling_rate, 48000) +# pcm_data = upsampled.astype(np.float64) +# stereo_pcm_data = np.tile(pcm_data, (2,1)).T +# +# buf = io.BytesIO() +# scipy.io.wavfile.write(buf, 48000, stereo_pcm_data.astype(np.float32)) +# return buf + +## Set up model +## This should live long in memory, so we do it up front. +#model_dir_default = './model/sf1_tm1' +#model_name_default = 'sf1_tm1.ckpt' +#converter = Converter(model_dir_default, model_name_default) + +#tacotron = load_tacotron_model('/home/bt/models/tacotron2/tacotron2_uberduck_Noire.pt') +#waveglow = load_waveglow_model('/home/bt/models/waveglow/waveglow_256channels_universal_v5.pt') +pipeline = TacotronWaveglowPipeline() + +class IndexHandler(): + def on_get(self, request, response): + response.content_type = 'text/html' + response.body = INDEX_HTML + +class ApiHandler(): + def on_post(self, request, response): + raw_data = json.load(request.bounded_stream) + + # Request parameters + synthesizer_checkpoint_path = raw_data.get('synthesizer_checkpoint_path') + vocoder_type = raw_data.get('vocoder_type') + waveglow_vocoder_checkpoint_path = raw_data.get('waveglow_vocoder_checkpoint_path') + hifigan_vocoder_checkpoint_path = raw_data.get('hifigan_vocoder_checkpoint_path') + hifigan_superres_vocoder_checkpoint_path = raw_data.get('hifigan_superres_vocoder_checkpoint_path') + inference_text = raw_data.get('inference_text') + output_audio_filename = raw_data.get('output_audio_filename') + output_spectrogram_filename = raw_data.get('output_spectrogram_filename') + output_metadata_filename = raw_data.get('output_metadata_filename') + + # Optional request parameters + maybe_clear_synthesizer_checkpoint_path = \ + raw_data.get('maybe_clear_synthesizer_checkpoint_path') + + print('synthesizer_checkpoint_path: {}'.format(synthesizer_checkpoint_path)) + print('vocoder_type: {}'.format(vocoder_type)) + print('waveglow_vocoder_checkpoint_path: {}'.format(waveglow_vocoder_checkpoint_path)) + print('hifigan_vocoder_checkpoint_path: {}'.format(hifigan_vocoder_checkpoint_path)) + print('hifigan_superres_vocoder_checkpoint_path: {}'.format(hifigan_superres_vocoder_checkpoint_path)) + print('inference_text: {}'.format(inference_text)) + print('output_audio_filename: {}'.format(output_audio_filename)) + print('output_spectrogram_filename: {}'.format(output_spectrogram_filename)) + print('output_metadata_filename: {}'.format(output_metadata_filename)) + print('maybe_clear_synthesizer_checkpoint_path: {}'.format( + maybe_clear_synthesizer_checkpoint_path)) + + pipeline.maybe_load_waveglow_model(waveglow_vocoder_checkpoint_path) + pipeline.maybe_load_hifigan(hifigan_vocoder_checkpoint_path) + pipeline.maybe_load_hifigan_super_resolution(hifigan_superres_vocoder_checkpoint_path) + + if maybe_clear_synthesizer_checkpoint_path: + pipeline.uncache_tacotron_model(maybe_clear_synthesizer_checkpoint_path) + + pipeline.maybe_load_tacotron_model_to_cache(synthesizer_checkpoint_path) + + inference_args = { + 'raw_text': inference_text, + 'synthesizer_checkpoint_path': synthesizer_checkpoint_path, + 'vocoder_type': vocoder_type, + 'output_audio_filename': output_audio_filename, + 'output_spectrogram_filename': output_spectrogram_filename, + 'output_metadata_filename': output_metadata_filename, + } + + pipeline.infer(inference_args) + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--port', type=int, default=8000) + args = parser.parse_args() + + api = falcon.API(middleware=[MultipartMiddleware()]) + api.add_route('/', IndexHandler()) + api.add_route('/infer', ApiHandler()) + + print('Serving on 0.0.0.0:%d' % args.port) + simple_server.make_server('0.0.0.0', args.port, api).serve_forever() + +if __name__ == '__main__': + main() diff --git a/models/tts/vocodes_server_talknet.py b/models/tts/vocodes_server_talknet.py new file mode 100755 index 0000000000..32cc859623 --- /dev/null +++ b/models/tts/vocodes_server_talknet.py @@ -0,0 +1,283 @@ +#!/usr/bin/env python3 + +""" +Starts a demo HTTP server to capture and transform audio +as a live demonstration of the trained model. +Brandon Thomas 2019-07-29 +""" + +import argparse +import falcon +import io +import librosa +import numpy as np +import os +import scipy +import soundfile +import tensorflow as tf +import subprocess +import tempfile +import json + +from falcon_multipart.middleware import MultipartMiddleware +#from model import CycleGAN +#from preprocess import * +from wsgiref import simple_server + +from vocodes_common_talknet import TacotronWaveglowPipeline +from vocodes_common_talknet import print_gpu_info +from vocodes_common_talknet import load_tacotron_model +from vocodes_common_talknet import load_waveglow_model +from vocodes_common_talknet import load_hifigan_model + +print("TensorFlow version: {}".format(tf.version.VERSION)) + +print_gpu_info() + +INDEX_HTML = ''' + + + + + + + +

TTS Inference

+ + + +
+ + +
+ + + +
+ + + +
+ + +
+ + +''' + + +#class Converter(): +# def __init__(self, model_dir, model_name): +# self.num_features = 24 +# self.sampling_rate = 16000 +# self.frame_period = 5.0 +# +# self.model = CycleGAN(num_features = self.num_features, mode = 'test') +# +# self.model.load(filepath = os.path.join(model_dir, model_name)) +# +# """ +# # NB: Save the graph +# definition = self.model.sess.graph_def +# directory = 'saved_model_2' +# tf.train.write_graph(definition, directory, 'saved_model_2.pb', as_text=True) +# # https://github.com/tensorflow/models/issues/3530#issuecomment-395968881 +# output_dir = './saved_model/' +# builder = tf.saved_model.builder.SavedModelBuilder(output_dir) +# builder.add_meta_graph_and_variables( +# self.model.sess, +# [tf.saved_model.tag_constants.SERVING], +# main_op=tf.tables_initializer(), +# ) +# builder.save() +# """ +# +# """ +# builder.add_meta_graph_and_variables( +# self.model.sess, +# [tf.saved_model.tag_constants.SERVING], +# signature_def_map={ +# 'predict_images': +# prediction_signature, +# signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY: +# classification_signature, +# }, +# main_op=tf.tables_initializer()) +# """ +# +# self.mcep_normalization_params = np.load(os.path.join(model_dir, 'mcep_normalization.npz')) +# self.mcep_mean_A = self.mcep_normalization_params['mean_A'] +# self.mcep_std_A = self.mcep_normalization_params['std_A'] +# self.mcep_mean_B = self.mcep_normalization_params['mean_B'] +# self.mcep_std_B = self.mcep_normalization_params['std_B'] +# +# self.logf0s_normalization_params = np.load(os.path.join(model_dir, 'logf0s_normalization.npz')) +# self.logf0s_mean_A = self.logf0s_normalization_params['mean_A'] +# self.logf0s_std_A = self.logf0s_normalization_params['std_A'] +# self.logf0s_mean_B = self.logf0s_normalization_params['mean_B'] +# self.logf0s_std_B = self.logf0s_normalization_params['std_B'] +# +# def convert(self, wav, conversion_direction='A2B'): +# wav = wav_padding(wav = wav, sr = self.sampling_rate, frame_period = self.frame_period, multiple = 4) +# f0, timeaxis, sp, ap = world_decompose(wav = wav, fs = self.sampling_rate, frame_period = self.frame_period) +# coded_sp = world_encode_spectral_envelop(sp = sp, fs = self.sampling_rate, dim = self.num_features) +# coded_sp_transposed = coded_sp.T +# +# if conversion_direction == 'A2B': +# f0_converted = pitch_conversion(f0 = f0, mean_log_src = self.logf0s_mean_A, std_log_src = self.logf0s_std_A, mean_log_target = self.logf0s_mean_B, std_log_target = self.logf0s_std_B) +# coded_sp_norm = (coded_sp_transposed - self.mcep_mean_A) / self.mcep_std_A +# coded_sp_converted_norm = self.model.test(inputs = np.array([coded_sp_norm]), direction = conversion_direction)[0] +# coded_sp_converted = coded_sp_converted_norm * self.mcep_std_B + self.mcep_mean_B +# else: +# f0_converted = pitch_conversion(f0 = f0, mean_log_src = self.logf0s_mean_B, std_log_src = self.logf0s_std_B, mean_log_target = self.logf0s_mean_A, std_log_target = self.logf0s_std_A) +# coded_sp_norm = (coded_sp_transposed - self.mcep_mean_B) / self.mcep_std_B +# coded_sp_converted_norm = self.model.test(inputs = np.array([coded_sp_norm]), direction = conversion_direction)[0] +# coded_sp_converted = coded_sp_converted_norm * self.mcep_std_A + self.mcep_mean_A +# +# coded_sp_converted = coded_sp_converted.T +# coded_sp_converted = np.ascontiguousarray(coded_sp_converted) +# decoded_sp_converted = world_decode_spectral_envelop(coded_sp = coded_sp_converted, fs = self.sampling_rate) +# wav_transformed = world_speech_synthesis(f0 = f0_converted, decoded_sp = decoded_sp_converted, ap = ap, fs = self.sampling_rate, frame_period = self.frame_period) +# +# # For debugging model output, uncomment the following line: +# # librosa.output.write_wav('model_output.wav', wav_transformed, self.sampling_rate) +# +# # TODO: Perhaps ditch this. It's probably unnecessary work. +# upsampled = librosa.resample(wav_transformed, self.sampling_rate, 48000) +# pcm_data = upsampled.astype(np.float64) +# stereo_pcm_data = np.tile(pcm_data, (2,1)).T +# +# buf = io.BytesIO() +# scipy.io.wavfile.write(buf, 48000, stereo_pcm_data.astype(np.float32)) +# return buf + +## Set up model +## This should live long in memory, so we do it up front. +#model_dir_default = './model/sf1_tm1' +#model_name_default = 'sf1_tm1.ckpt' +#converter = Converter(model_dir_default, model_name_default) + +#tacotron = load_tacotron_model('/home/bt/models/tacotron2/tacotron2_uberduck_Noire.pt') +#waveglow = load_waveglow_model('/home/bt/models/waveglow/waveglow_256channels_universal_v5.pt') +pipeline = TacotronWaveglowPipeline() + +class IndexHandler(): + def on_get(self, request, response): + response.content_type = 'text/html' + response.body = INDEX_HTML + +class ApiHandler(): + def on_post(self, request, response): + raw_data = json.load(request.bounded_stream) + + # Request parameters + synthesizer_checkpoint_path = raw_data.get('synthesizer_checkpoint_path') + vocoder_type = raw_data.get('vocoder_type') + waveglow_vocoder_checkpoint_path = raw_data.get('waveglow_vocoder_checkpoint_path') + hifigan_vocoder_checkpoint_path = raw_data.get('hifigan_vocoder_checkpoint_path') + hifigan_superres_vocoder_checkpoint_path = raw_data.get('hifigan_superres_vocoder_checkpoint_path') + inference_text = raw_data.get('inference_text') + output_audio_filename = raw_data.get('output_audio_filename') + output_spectrogram_filename = raw_data.get('output_spectrogram_filename') + output_metadata_filename = raw_data.get('output_metadata_filename') + + # Optional request parameters + maybe_clear_synthesizer_checkpoint_path = \ + raw_data.get('maybe_clear_synthesizer_checkpoint_path') + + print('synthesizer_checkpoint_path: {}'.format(synthesizer_checkpoint_path)) + print('vocoder_type: {}'.format(vocoder_type)) + print('waveglow_vocoder_checkpoint_path: {}'.format(waveglow_vocoder_checkpoint_path)) + print('hifigan_vocoder_checkpoint_path: {}'.format(hifigan_vocoder_checkpoint_path)) + print('hifigan_superres_vocoder_checkpoint_path: {}'.format(hifigan_superres_vocoder_checkpoint_path)) + print('inference_text: {}'.format(inference_text)) + print('output_audio_filename: {}'.format(output_audio_filename)) + print('output_spectrogram_filename: {}'.format(output_spectrogram_filename)) + print('output_metadata_filename: {}'.format(output_metadata_filename)) + print('maybe_clear_synthesizer_checkpoint_path: {}'.format( + maybe_clear_synthesizer_checkpoint_path)) + + #pipeline.maybe_load_waveglow_model(waveglow_vocoder_checkpoint_path) + #pipeline.maybe_load_hifigan(hifigan_vocoder_checkpoint_path) + #pipeline.maybe_load_hifigan_super_resolution(hifigan_superres_vocoder_checkpoint_path) + + #if maybe_clear_synthesizer_checkpoint_path: + # pipeline.uncache_tacotron_model(maybe_clear_synthesizer_checkpoint_path) + + #pipeline.maybe_load_tacotron_model_to_cache(synthesizer_checkpoint_path) + + inference_args = { + 'raw_text': inference_text, + 'synthesizer_checkpoint_path': synthesizer_checkpoint_path, + 'vocoder_type': vocoder_type, + 'output_audio_filename': output_audio_filename, + 'output_spectrogram_filename': output_spectrogram_filename, + 'output_metadata_filename': output_metadata_filename, + } + + pipeline.infer(inference_args) + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument('--port', type=int, default=8000) + args = parser.parse_args() + + api = falcon.API(middleware=[MultipartMiddleware()]) + api.add_route('/', IndexHandler()) + api.add_route('/infer', ApiHandler()) + + print('Serving on 0.0.0.0:%d' % args.port) + simple_server.make_server('0.0.0.0', args.port, api).serve_forever() + +if __name__ == '__main__': + """ + Launching a TalkNet server: + + from ControllableTalkNet.controllable_talknet import * + + app.run_server( + mode="inline", + threaded=True, + #dev_tools_ui=True, + #dev_tools_hot_reload=True, + ) + """ + main() diff --git a/models/tts/vocodes_wav_images.py b/models/tts/vocodes_wav_images.py new file mode 100644 index 0000000000..ec5949805d --- /dev/null +++ b/models/tts/vocodes_wav_images.py @@ -0,0 +1,49 @@ +from PIL import Image, ImageDraw + +def render_histogram(wav, filename='histograms/output.png', row_index=None, col_index=None): + dimensions = wav.shape + print('>>> Tensor Image input size: {}'.format(dimensions)) + + wav_signal = None + rows = -1 + cols = -1 + if row_index is not None and col_index is not None: + rows = dimensions[row_index] + cols = dimensions[col_index] + wav_signal = wav + else: + # Remove 1-sized dimensions + wav_signal = wav.squeeze() + rows = wav_signal.shape[0] + cols = wav_signal.shape[1] + + print('>>> Generating image (rows: {} cols: {})'.format(rows, cols)) + image = Image.new('RGB', (rows, cols)) + pixels = image.load() + + minimum = wav_signal[0,0] + maximum = wav_signal[0,0] + + for x in range(0, rows): + for y in range(0, cols): + if wav_signal[x,y] > maximum: + maximum = wav_signal[x,y] + if wav_signal[x,y] < minimum: + minimum = wav_signal[x,y] + + print('Minimum: ' + str(minimum)) + print('Maximum: ' + str(maximum)) + + if maximum == minimum: + maximum = 1 # nb: to prevent division by zero below + + for x in range(0, rows): + for y in range(0, cols): + v = wav_signal[x,y] + scaled = int((v - minimum) / (maximum - minimum) * 255) + pixels[x, y] = (scaled, scaled, scaled) + + # Image Show only works on local X server: + #image.show() + image.save(filename) + diff --git a/models/tts/waveglow/.gitmodules b/models/tts/waveglow/.gitmodules new file mode 100644 index 0000000000..e3bf8bf069 --- /dev/null +++ b/models/tts/waveglow/.gitmodules @@ -0,0 +1,3 @@ +[submodule "tacotron2"] + path = tacotron2 + url = http://github.com/NVIDIA/tacotron2 diff --git a/models/tts/waveglow/LICENSE b/models/tts/waveglow/LICENSE new file mode 100644 index 0000000000..2a718d63da --- /dev/null +++ b/models/tts/waveglow/LICENSE @@ -0,0 +1,29 @@ +BSD 3-Clause License + +Copyright (c) 2018, NVIDIA Corporation +All rights reserved. + +Redistribution and use in source and binary forms, with or without +modification, are permitted provided that the following conditions are met: + +* Redistributions of source code must retain the above copyright notice, this + list of conditions and the following disclaimer. + +* Redistributions in binary form must reproduce the above copyright notice, + this list of conditions and the following disclaimer in the documentation + and/or other materials provided with the distribution. + +* Neither the name of the copyright holder nor the names of its + contributors may be used to endorse or promote products derived from + this software without specific prior written permission. + +THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE +FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL +DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR +SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER +CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, +OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE +OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. diff --git a/models/tts/waveglow/README.md b/models/tts/waveglow/README.md new file mode 100755 index 0000000000..89886f00a6 --- /dev/null +++ b/models/tts/waveglow/README.md @@ -0,0 +1,91 @@ +![WaveGlow](waveglow_logo.png "WaveGLow") + +## WaveGlow: a Flow-based Generative Network for Speech Synthesis + +### Ryan Prenger, Rafael Valle, and Bryan Catanzaro + +In our recent [paper], we propose WaveGlow: a flow-based network capable of +generating high quality speech from mel-spectrograms. WaveGlow combines insights +from [Glow] and [WaveNet] in order to provide fast, efficient and high-quality +audio synthesis, without the need for auto-regression. WaveGlow is implemented +using only a single network, trained using only a single cost function: +maximizing the likelihood of the training data, which makes the training +procedure simple and stable. + +Our [PyTorch] implementation produces audio samples at a rate of 4850 +kHz on an NVIDIA V100 GPU. Mean Opinion Scores show that it delivers audio +quality as good as the best publicly available WaveNet implementation. + +Visit our [website] for audio samples. + +## Setup + +1. Clone our repo and initialize submodule + + ```command + git clone https://github.com/NVIDIA/waveglow.git + cd waveglow + git submodule init + git submodule update + ``` + +2. Install requirements `pip3 install -r requirements.txt` + +3. Install [Apex] + + +## Generate audio with our pre-existing model + +1. Download our [published model] +2. Download [mel-spectrograms] +3. Generate audio `python3 inference.py -f <(ls mel_spectrograms/*.pt) -w waveglow_256channels.pt -o . --is_fp16 -s 0.6` + +N.b. use `convert_model.py` to convert your older models to the current model +with fused residual and skip connections. + +## Train your own model + +1. Download [LJ Speech Data]. In this example it's in `data/` + +2. Make a list of the file names to use for training/testing + + ```command + ls data/*.wav | tail -n+10 > train_files.txt + ls data/*.wav | head -n10 > test_files.txt + ``` + +3. Train your WaveGlow networks + + ```command + mkdir checkpoints + python train.py -c config.json + ``` + + For multi-GPU training replace `train.py` with `distributed.py`. Only tested with single node and NCCL. + + For mixed precision training set `"fp16_run": true` on `config.json`. + +4. Make test set mel-spectrograms + + `python mel2samp.py -f test_files.txt -o . -c config.json` + +5. Do inference with your network + + ```command + ls *.pt > mel_files.txt + python3 inference.py -f mel_files.txt -w checkpoints/waveglow_10000 -o . --is_fp16 -s 0.6 + ``` + +[//]: # (TODO) +[//]: # (PROVIDE INSTRUCTIONS FOR DOWNLOADING LJS) +[pytorch 1.0]: https://github.com/pytorch/pytorch#installation +[website]: https://nv-adlr.github.io/WaveGlow +[paper]: https://arxiv.org/abs/1811.00002 +[WaveNet implementation]: https://github.com/r9y9/wavenet_vocoder +[Glow]: https://blog.openai.com/glow/ +[WaveNet]: https://deepmind.com/blog/wavenet-generative-model-raw-audio/ +[PyTorch]: http://pytorch.org +[published model]: https://ngc.nvidia.com/catalog/models/nvidia:waveglow_ljs_256channels +[mel-spectrograms]: https://drive.google.com/file/d/1g_VXK2lpP9J25dQFhQwx7doWl_p20fXA/view?usp=sharing +[LJ Speech Data]: https://keithito.com/LJ-Speech-Dataset +[Apex]: https://github.com/nvidia/apex diff --git a/models/tts/waveglow/config.json b/models/tts/waveglow/config.json new file mode 100755 index 0000000000..7160474748 --- /dev/null +++ b/models/tts/waveglow/config.json @@ -0,0 +1,41 @@ +{ + "train_config": { + "fp16_run": true, + "output_directory": "checkpoints", + "epochs": 100000, + "learning_rate": 1e-4, + "sigma": 1.0, + "iters_per_checkpoint": 2000, + "batch_size": 12, + "seed": 1234, + "checkpoint_path": "", + "with_tensorboard": false + }, + "data_config": { + "training_files": "train_files.txt", + "segment_length": 16000, + "sampling_rate": 22050, + "filter_length": 1024, + "hop_length": 256, + "win_length": 1024, + "mel_fmin": 0.0, + "mel_fmax": 8000.0 + }, + "dist_config": { + "dist_backend": "nccl", + "dist_url": "tcp://localhost:54321" + }, + + "waveglow_config": { + "n_mel_channels": 80, + "n_flows": 12, + "n_group": 8, + "n_early_every": 4, + "n_early_size": 2, + "WN_config": { + "n_layers": 8, + "n_channels": 256, + "kernel_size": 3 + } + } +} diff --git a/models/tts/waveglow/convert_model.py b/models/tts/waveglow/convert_model.py new file mode 100644 index 0000000000..32e77e2b20 --- /dev/null +++ b/models/tts/waveglow/convert_model.py @@ -0,0 +1,72 @@ +import sys +import copy +import torch + +def _check_model_old_version(model): + if hasattr(model.WN[0], 'res_layers') or hasattr(model.WN[0], 'cond_layers'): + return True + else: + return False + + +def _update_model_res_skip(old_model, new_model): + for idx in range(0, len(new_model.WN)): + wavenet = new_model.WN[idx] + n_channels = wavenet.n_channels + n_layers = wavenet.n_layers + wavenet.res_skip_layers = torch.nn.ModuleList() + for i in range(0, n_layers): + if i < n_layers - 1: + res_skip_channels = 2*n_channels + else: + res_skip_channels = n_channels + res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1) + skip_layer = torch.nn.utils.remove_weight_norm(wavenet.skip_layers[i]) + if i < n_layers - 1: + res_layer = torch.nn.utils.remove_weight_norm(wavenet.res_layers[i]) + res_skip_layer.weight = torch.nn.Parameter(torch.cat([res_layer.weight, skip_layer.weight])) + res_skip_layer.bias = torch.nn.Parameter(torch.cat([res_layer.bias, skip_layer.bias])) + else: + res_skip_layer.weight = torch.nn.Parameter(skip_layer.weight) + res_skip_layer.bias = torch.nn.Parameter(skip_layer.bias) + res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight') + wavenet.res_skip_layers.append(res_skip_layer) + del wavenet.res_layers + del wavenet.skip_layers + +def _update_model_cond(old_model, new_model): + for idx in range(0, len(new_model.WN)): + wavenet = new_model.WN[idx] + n_channels = wavenet.n_channels + n_layers = wavenet.n_layers + n_mel_channels = wavenet.cond_layers[0].weight.shape[1] + cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels*n_layers, 1) + cond_layer_weight = [] + cond_layer_bias = [] + for i in range(0, n_layers): + _cond_layer = torch.nn.utils.remove_weight_norm(wavenet.cond_layers[i]) + cond_layer_weight.append(_cond_layer.weight) + cond_layer_bias.append(_cond_layer.bias) + cond_layer.weight = torch.nn.Parameter(torch.cat(cond_layer_weight)) + cond_layer.bias = torch.nn.Parameter(torch.cat(cond_layer_bias)) + cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight') + wavenet.cond_layer = cond_layer + del wavenet.cond_layers + +def update_model(old_model): + if not _check_model_old_version(old_model): + return old_model + new_model = copy.deepcopy(old_model) + if hasattr(old_model.WN[0], 'res_layers'): + _update_model_res_skip(old_model, new_model) + if hasattr(old_model.WN[0], 'cond_layers'): + _update_model_cond(old_model, new_model) + return new_model + +if __name__ == '__main__': + old_model_path = sys.argv[1] + new_model_path = sys.argv[2] + model = torch.load(old_model_path) + model['model'] = update_model(model['model']) + torch.save(model, new_model_path) + diff --git a/models/tts/waveglow/denoiser.py b/models/tts/waveglow/denoiser.py new file mode 100644 index 0000000000..0e73584d54 --- /dev/null +++ b/models/tts/waveglow/denoiser.py @@ -0,0 +1,40 @@ +import sys +sys.path.append('tacotron2') +import torch +from layers import STFT + + +class Denoiser(torch.nn.Module): + """ Removes model bias from audio produced with waveglow """ + + def __init__(self, waveglow, filter_length=1024, n_overlap=4, + win_length=1024, mode='zeros'): + super(Denoiser, self).__init__() + self.stft = STFT(filter_length=filter_length, + hop_length=int(filter_length/n_overlap), + win_length=win_length).cuda() + if mode == 'zeros': + mel_input = torch.zeros( + (1, 80, 88), + dtype=waveglow.upsample.weight.dtype, + device=waveglow.upsample.weight.device) + elif mode == 'normal': + mel_input = torch.randn( + (1, 80, 88), + dtype=waveglow.upsample.weight.dtype, + device=waveglow.upsample.weight.device) + else: + raise Exception("Mode {} if not supported".format(mode)) + + with torch.no_grad(): + bias_audio = waveglow.infer(mel_input, sigma=0.0).float() + bias_spec, _ = self.stft.transform(bias_audio) + + self.register_buffer('bias_spec', bias_spec[:, :, 0][:, :, None]) + + def forward(self, audio, strength=0.1): + audio_spec, audio_angles = self.stft.transform(audio.cuda().float()) + audio_spec_denoised = audio_spec - self.bias_spec * strength + audio_spec_denoised = torch.clamp(audio_spec_denoised, 0.0) + audio_denoised = self.stft.inverse(audio_spec_denoised, audio_angles) + return audio_denoised diff --git a/models/tts/waveglow/distributed.py b/models/tts/waveglow/distributed.py new file mode 100644 index 0000000000..19cbfd2cca --- /dev/null +++ b/models/tts/waveglow/distributed.py @@ -0,0 +1,184 @@ +# ***************************************************************************** +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the +# names of its contributors may be used to endorse or promote products +# derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY +# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# ***************************************************************************** +import os +import sys +import time +import subprocess +import argparse + +import torch +import torch.distributed as dist +from torch.autograd import Variable + +def reduce_tensor(tensor, num_gpus): + rt = tensor.clone() + dist.all_reduce(rt, op=dist.reduce_op.SUM) + rt /= num_gpus + return rt + +def init_distributed(rank, num_gpus, group_name, dist_backend, dist_url): + assert torch.cuda.is_available(), "Distributed mode requires CUDA." + print("Initializing Distributed") + + # Set cuda device so everything is done on the right GPU. + torch.cuda.set_device(rank % torch.cuda.device_count()) + + # Initialize distributed communication + dist.init_process_group(dist_backend, init_method=dist_url, + world_size=num_gpus, rank=rank, + group_name=group_name) + +def _flatten_dense_tensors(tensors): + """Flatten dense tensors into a contiguous 1D buffer. Assume tensors are of + same dense type. + Since inputs are dense, the resulting tensor will be a concatenated 1D + buffer. Element-wise operation on this buffer will be equivalent to + operating individually. + Arguments: + tensors (Iterable[Tensor]): dense tensors to flatten. + Returns: + A contiguous 1D buffer containing input tensors. + """ + if len(tensors) == 1: + return tensors[0].contiguous().view(-1) + flat = torch.cat([t.contiguous().view(-1) for t in tensors], dim=0) + return flat + +def _unflatten_dense_tensors(flat, tensors): + """View a flat buffer using the sizes of tensors. Assume that tensors are of + same dense type, and that flat is given by _flatten_dense_tensors. + Arguments: + flat (Tensor): flattened dense tensors to unflatten. + tensors (Iterable[Tensor]): dense tensors whose sizes will be used to + unflatten flat. + Returns: + Unflattened dense tensors with sizes same as tensors and values from + flat. + """ + outputs = [] + offset = 0 + for tensor in tensors: + numel = tensor.numel() + outputs.append(flat.narrow(0, offset, numel).view_as(tensor)) + offset += numel + return tuple(outputs) + +def apply_gradient_allreduce(module): + """ + Modifies existing model to do gradient allreduce, but doesn't change class + so you don't need "module" + """ + if not hasattr(dist, '_backend'): + module.warn_on_half = True + else: + module.warn_on_half = True if dist._backend == dist.dist_backend.GLOO else False + + for p in module.state_dict().values(): + if not torch.is_tensor(p): + continue + dist.broadcast(p, 0) + + def allreduce_params(): + if(module.needs_reduction): + module.needs_reduction = False + buckets = {} + for param in module.parameters(): + if param.requires_grad and param.grad is not None: + tp = type(param.data) + if tp not in buckets: + buckets[tp] = [] + buckets[tp].append(param) + if module.warn_on_half: + if torch.cuda.HalfTensor in buckets: + print("WARNING: gloo dist backend for half parameters may be extremely slow." + + " It is recommended to use the NCCL backend in this case. This currently requires" + + "PyTorch built from top of tree master.") + module.warn_on_half = False + + for tp in buckets: + bucket = buckets[tp] + grads = [param.grad.data for param in bucket] + coalesced = _flatten_dense_tensors(grads) + dist.all_reduce(coalesced) + coalesced /= dist.get_world_size() + for buf, synced in zip(grads, _unflatten_dense_tensors(coalesced, grads)): + buf.copy_(synced) + + for param in list(module.parameters()): + def allreduce_hook(*unused): + Variable._execution_engine.queue_callback(allreduce_params) + if param.requires_grad: + param.register_hook(allreduce_hook) + dir(param) + + def set_needs_reduction(self, input, output): + self.needs_reduction = True + + module.register_forward_hook(set_needs_reduction) + return module + + +def main(config, stdout_dir, args_str): + args_list = ['train.py'] + args_list += args_str.split(' ') if len(args_str) > 0 else [] + + args_list.append('--config={}'.format(config)) + + num_gpus = torch.cuda.device_count() + args_list.append('--num_gpus={}'.format(num_gpus)) + args_list.append("--group_name=group_{}".format(time.strftime("%Y_%m_%d-%H%M%S"))) + + if not os.path.isdir(stdout_dir): + os.makedirs(stdout_dir) + os.chmod(stdout_dir, 0o775) + + workers = [] + + for i in range(num_gpus): + args_list[-2] = '--rank={}'.format(i) + stdout = None if i == 0 else open( + os.path.join(stdout_dir, "GPU_{}.log".format(i)), "w") + print(args_list) + p = subprocess.Popen([str(sys.executable)]+args_list, stdout=stdout) + workers.append(p) + + for p in workers: + p.wait() + + +if __name__ == '__main__': + parser = argparse.ArgumentParser() + parser.add_argument('-c', '--config', type=str, required=True, + help='JSON file for configuration') + parser.add_argument('-s', '--stdout_dir', type=str, default=".", + help='directory to save stoud logs') + parser.add_argument( + '-a', '--args_str', type=str, default='', + help='double quoted string with space separated key value pairs') + + args = parser.parse_args() + main(args.config, args.stdout_dir, args.args_str) diff --git a/models/tts/waveglow/glow.py b/models/tts/waveglow/glow.py new file mode 100644 index 0000000000..7a7696403d --- /dev/null +++ b/models/tts/waveglow/glow.py @@ -0,0 +1,311 @@ +# ***************************************************************************** +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the +# names of its contributors may be used to endorse or promote products +# derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY +# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# ***************************************************************************** +import copy +import torch +from torch.autograd import Variable +import torch.nn.functional as F + + +@torch.jit.script +def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): + n_channels_int = n_channels[0] + in_act = input_a+input_b + t_act = torch.tanh(in_act[:, :n_channels_int, :]) + s_act = torch.sigmoid(in_act[:, n_channels_int:, :]) + acts = t_act * s_act + return acts + + +class WaveGlowLoss(torch.nn.Module): + def __init__(self, sigma=1.0): + super(WaveGlowLoss, self).__init__() + self.sigma = sigma + + def forward(self, model_output): + z, log_s_list, log_det_W_list = model_output + for i, log_s in enumerate(log_s_list): + if i == 0: + log_s_total = torch.sum(log_s) + log_det_W_total = log_det_W_list[i] + else: + log_s_total = log_s_total + torch.sum(log_s) + log_det_W_total += log_det_W_list[i] + + loss = torch.sum(z*z)/(2*self.sigma*self.sigma) - log_s_total - log_det_W_total + return loss/(z.size(0)*z.size(1)*z.size(2)) + + +class Invertible1x1Conv(torch.nn.Module): + """ + The layer outputs both the convolution, and the log determinant + of its weight matrix. If reverse=True it does convolution with + inverse + """ + def __init__(self, c): + super(Invertible1x1Conv, self).__init__() + self.conv = torch.nn.Conv1d(c, c, kernel_size=1, stride=1, padding=0, + bias=False) + + # Sample a random orthonormal matrix to initialize weights + W = torch.qr(torch.FloatTensor(c, c).normal_())[0] + + # Ensure determinant is 1.0 not -1.0 + if torch.det(W) < 0: + W[:,0] = -1*W[:,0] + W = W.view(c, c, 1) + self.conv.weight.data = W + + def forward(self, z, reverse=False): + # shape + batch_size, group_size, n_of_groups = z.size() + + W = self.conv.weight.squeeze() + + if reverse: + if not hasattr(self, 'W_inverse'): + # Reverse computation + W_inverse = W.float().inverse() + W_inverse = Variable(W_inverse[..., None]) + if z.type() == 'torch.cuda.HalfTensor': + W_inverse = W_inverse.half() + self.W_inverse = W_inverse + z = F.conv1d(z, self.W_inverse, bias=None, stride=1, padding=0) + return z + else: + # Forward computation + log_det_W = batch_size * n_of_groups * torch.logdet(W) + z = self.conv(z) + return z, log_det_W + + +class WN(torch.nn.Module): + """ + This is the WaveNet like layer for the affine coupling. The primary difference + from WaveNet is the convolutions need not be causal. There is also no dilation + size reset. The dilation only doubles on each layer + """ + def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, + kernel_size): + super(WN, self).__init__() + assert(kernel_size % 2 == 1) + assert(n_channels % 2 == 0) + self.n_layers = n_layers + self.n_channels = n_channels + self.in_layers = torch.nn.ModuleList() + self.res_skip_layers = torch.nn.ModuleList() + + start = torch.nn.Conv1d(n_in_channels, n_channels, 1) + start = torch.nn.utils.weight_norm(start, name='weight') + self.start = start + + # Initializing last layer to 0 makes the affine coupling layers + # do nothing at first. This helps with training stability + end = torch.nn.Conv1d(n_channels, 2*n_in_channels, 1) + end.weight.data.zero_() + end.bias.data.zero_() + self.end = end + + cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels*n_layers, 1) + self.cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight') + + for i in range(n_layers): + dilation = 2 ** i + padding = int((kernel_size*dilation - dilation)/2) + in_layer = torch.nn.Conv1d(n_channels, 2*n_channels, kernel_size, + dilation=dilation, padding=padding) + in_layer = torch.nn.utils.weight_norm(in_layer, name='weight') + self.in_layers.append(in_layer) + + + # last one is not necessary + if i < n_layers - 1: + res_skip_channels = 2*n_channels + else: + res_skip_channels = n_channels + res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1) + res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight') + self.res_skip_layers.append(res_skip_layer) + + def forward(self, forward_input): + audio, spect = forward_input + audio = self.start(audio) + output = torch.zeros_like(audio) + n_channels_tensor = torch.IntTensor([self.n_channels]) + + spect = self.cond_layer(spect) + + for i in range(self.n_layers): + spect_offset = i*2*self.n_channels + acts = fused_add_tanh_sigmoid_multiply( + self.in_layers[i](audio), + spect[:,spect_offset:spect_offset+2*self.n_channels,:], + n_channels_tensor) + + res_skip_acts = self.res_skip_layers[i](acts) + if i < self.n_layers - 1: + audio = audio + res_skip_acts[:,:self.n_channels,:] + output = output + res_skip_acts[:,self.n_channels:,:] + else: + output = output + res_skip_acts + + return self.end(output) + + +class WaveGlow(torch.nn.Module): + def __init__(self, n_mel_channels, n_flows, n_group, n_early_every, + n_early_size, WN_config): + super(WaveGlow, self).__init__() + + self.upsample = torch.nn.ConvTranspose1d(n_mel_channels, + n_mel_channels, + 1024, stride=256) + assert(n_group % 2 == 0) + self.n_flows = n_flows + self.n_group = n_group + self.n_early_every = n_early_every + self.n_early_size = n_early_size + self.WN = torch.nn.ModuleList() + self.convinv = torch.nn.ModuleList() + + n_half = int(n_group/2) + + # Set up layers with the right sizes based on how many dimensions + # have been output already + n_remaining_channels = n_group + for k in range(n_flows): + if k % self.n_early_every == 0 and k > 0: + n_half = n_half - int(self.n_early_size/2) + n_remaining_channels = n_remaining_channels - self.n_early_size + self.convinv.append(Invertible1x1Conv(n_remaining_channels)) + self.WN.append(WN(n_half, n_mel_channels*n_group, **WN_config)) + self.n_remaining_channels = n_remaining_channels # Useful during inference + + def forward(self, forward_input): + """ + forward_input[0] = mel_spectrogram: batch x n_mel_channels x frames + forward_input[1] = audio: batch x time + """ + spect, audio = forward_input + + # Upsample spectrogram to size of audio + spect = self.upsample(spect) + assert(spect.size(2) >= audio.size(1)) + if spect.size(2) > audio.size(1): + spect = spect[:, :, :audio.size(1)] + + spect = spect.unfold(2, self.n_group, self.n_group).permute(0, 2, 1, 3) + spect = spect.contiguous().view(spect.size(0), spect.size(1), -1).permute(0, 2, 1) + + audio = audio.unfold(1, self.n_group, self.n_group).permute(0, 2, 1) + output_audio = [] + log_s_list = [] + log_det_W_list = [] + + for k in range(self.n_flows): + if k % self.n_early_every == 0 and k > 0: + output_audio.append(audio[:,:self.n_early_size,:]) + audio = audio[:,self.n_early_size:,:] + + audio, log_det_W = self.convinv[k](audio) + log_det_W_list.append(log_det_W) + + n_half = int(audio.size(1)/2) + audio_0 = audio[:,:n_half,:] + audio_1 = audio[:,n_half:,:] + + output = self.WN[k]((audio_0, spect)) + log_s = output[:, n_half:, :] + b = output[:, :n_half, :] + audio_1 = torch.exp(log_s)*audio_1 + b + log_s_list.append(log_s) + + audio = torch.cat([audio_0, audio_1],1) + + output_audio.append(audio) + return torch.cat(output_audio,1), log_s_list, log_det_W_list + + def infer(self, spect, sigma=1.0): + spect = self.upsample(spect) + # trim conv artifacts. maybe pad spec to kernel multiple + time_cutoff = self.upsample.kernel_size[0] - self.upsample.stride[0] + spect = spect[:, :, :-time_cutoff] + + spect = spect.unfold(2, self.n_group, self.n_group).permute(0, 2, 1, 3) + spect = spect.contiguous().view(spect.size(0), spect.size(1), -1).permute(0, 2, 1) + + if spect.type() == 'torch.cuda.HalfTensor': + audio = torch.cuda.HalfTensor(spect.size(0), + self.n_remaining_channels, + spect.size(2)).normal_() + else: + audio = torch.cuda.FloatTensor(spect.size(0), + self.n_remaining_channels, + spect.size(2)).normal_() + + audio = torch.autograd.Variable(sigma*audio) + + for k in reversed(range(self.n_flows)): + n_half = int(audio.size(1)/2) + audio_0 = audio[:,:n_half,:] + audio_1 = audio[:,n_half:,:] + + output = self.WN[k]((audio_0, spect)) + + s = output[:, n_half:, :] + b = output[:, :n_half, :] + audio_1 = (audio_1 - b)/torch.exp(s) + audio = torch.cat([audio_0, audio_1],1) + + audio = self.convinv[k](audio, reverse=True) + + if k % self.n_early_every == 0 and k > 0: + if spect.type() == 'torch.cuda.HalfTensor': + z = torch.cuda.HalfTensor(spect.size(0), self.n_early_size, spect.size(2)).normal_() + else: + z = torch.cuda.FloatTensor(spect.size(0), self.n_early_size, spect.size(2)).normal_() + audio = torch.cat((sigma*z, audio),1) + + audio = audio.permute(0,2,1).contiguous().view(audio.size(0), -1).data + return audio + + @staticmethod + def remove_weightnorm(model): + waveglow = model + for WN in waveglow.WN: + WN.start = torch.nn.utils.remove_weight_norm(WN.start) + WN.in_layers = remove(WN.in_layers) + WN.cond_layer = torch.nn.utils.remove_weight_norm(WN.cond_layer) + WN.res_skip_layers = remove(WN.res_skip_layers) + return waveglow + + +def remove(conv_list): + new_conv_list = torch.nn.ModuleList() + for old_conv in conv_list: + old_conv = torch.nn.utils.remove_weight_norm(old_conv) + new_conv_list.append(old_conv) + return new_conv_list diff --git a/models/tts/waveglow/glow_old.py b/models/tts/waveglow/glow_old.py new file mode 100644 index 0000000000..0de2375cc5 --- /dev/null +++ b/models/tts/waveglow/glow_old.py @@ -0,0 +1,233 @@ +import copy +import torch +from glow import Invertible1x1Conv, remove + + +@torch.jit.script +def fused_add_tanh_sigmoid_multiply(input_a, input_b, n_channels): + n_channels_int = n_channels[0] + in_act = input_a+input_b + t_act = torch.tanh(in_act[:, :n_channels_int, :]) + s_act = torch.sigmoid(in_act[:, n_channels_int:, :]) + acts = t_act * s_act + return acts + + +class WN(torch.nn.Module): + """ + This is the WaveNet like layer for the affine coupling. The primary difference + from WaveNet is the convolutions need not be causal. There is also no dilation + size reset. The dilation only doubles on each layer + """ + def __init__(self, n_in_channels, n_mel_channels, n_layers, n_channels, + kernel_size): + super(WN, self).__init__() + assert(kernel_size % 2 == 1) + assert(n_channels % 2 == 0) + self.n_layers = n_layers + self.n_channels = n_channels + self.in_layers = torch.nn.ModuleList() + self.res_skip_layers = torch.nn.ModuleList() + self.cond_layers = torch.nn.ModuleList() + + start = torch.nn.Conv1d(n_in_channels, n_channels, 1) + start = torch.nn.utils.weight_norm(start, name='weight') + self.start = start + + # Initializing last layer to 0 makes the affine coupling layers + # do nothing at first. This helps with training stability + end = torch.nn.Conv1d(n_channels, 2*n_in_channels, 1) + end.weight.data.zero_() + end.bias.data.zero_() + self.end = end + + for i in range(n_layers): + dilation = 2 ** i + padding = int((kernel_size*dilation - dilation)/2) + in_layer = torch.nn.Conv1d(n_channels, 2*n_channels, kernel_size, + dilation=dilation, padding=padding) + in_layer = torch.nn.utils.weight_norm(in_layer, name='weight') + self.in_layers.append(in_layer) + + cond_layer = torch.nn.Conv1d(n_mel_channels, 2*n_channels, 1) + cond_layer = torch.nn.utils.weight_norm(cond_layer, name='weight') + self.cond_layers.append(cond_layer) + + # last one is not necessary + if i < n_layers - 1: + res_skip_channels = 2*n_channels + else: + res_skip_channels = n_channels + res_skip_layer = torch.nn.Conv1d(n_channels, res_skip_channels, 1) + res_skip_layer = torch.nn.utils.weight_norm(res_skip_layer, name='weight') + self.res_skip_layers.append(res_skip_layer) + + def forward(self, forward_input): + audio, spect = forward_input + audio = self.start(audio) + + for i in range(self.n_layers): + acts = fused_add_tanh_sigmoid_multiply( + self.in_layers[i](audio), + self.cond_layers[i](spect), + torch.IntTensor([self.n_channels])) + + res_skip_acts = self.res_skip_layers[i](acts) + if i < self.n_layers - 1: + audio = res_skip_acts[:,:self.n_channels,:] + audio + skip_acts = res_skip_acts[:,self.n_channels:,:] + else: + skip_acts = res_skip_acts + + if i == 0: + output = skip_acts + else: + output = skip_acts + output + return self.end(output) + + +class WaveGlow(torch.nn.Module): + def __init__(self, n_mel_channels, n_flows, n_group, n_early_every, + n_early_size, WN_config): + super(WaveGlow, self).__init__() + + self.upsample = torch.nn.ConvTranspose1d(n_mel_channels, + n_mel_channels, + 1024, stride=256) + assert(n_group % 2 == 0) + self.n_flows = n_flows + self.n_group = n_group + self.n_early_every = n_early_every + self.n_early_size = n_early_size + self.WN = torch.nn.ModuleList() + self.convinv = torch.nn.ModuleList() + + n_half = int(n_group/2) + + # Set up layers with the right sizes based on how many dimensions + # have been output already + n_remaining_channels = n_group + for k in range(n_flows): + if k % self.n_early_every == 0 and k > 0: + n_half = n_half - int(self.n_early_size/2) + n_remaining_channels = n_remaining_channels - self.n_early_size + self.convinv.append(Invertible1x1Conv(n_remaining_channels)) + self.WN.append(WN(n_half, n_mel_channels*n_group, **WN_config)) + self.n_remaining_channels = n_remaining_channels # Useful during inference + + def forward(self, forward_input): + return None + """ + forward_input[0] = audio: batch x time + forward_input[1] = upsamp_spectrogram: batch x n_cond_channels x time + """ + """ + spect, audio = forward_input + + # Upsample spectrogram to size of audio + spect = self.upsample(spect) + assert(spect.size(2) >= audio.size(1)) + if spect.size(2) > audio.size(1): + spect = spect[:, :, :audio.size(1)] + + spect = spect.unfold(2, self.n_group, self.n_group).permute(0, 2, 1, 3) + spect = spect.contiguous().view(spect.size(0), spect.size(1), -1).permute(0, 2, 1) + + audio = audio.unfold(1, self.n_group, self.n_group).permute(0, 2, 1) + output_audio = [] + s_list = [] + s_conv_list = [] + + for k in range(self.n_flows): + if k%4 == 0 and k > 0: + output_audio.append(audio[:,:self.n_multi,:]) + audio = audio[:,self.n_multi:,:] + + # project to new basis + audio, s = self.convinv[k](audio) + s_conv_list.append(s) + + n_half = int(audio.size(1)/2) + if k%2 == 0: + audio_0 = audio[:,:n_half,:] + audio_1 = audio[:,n_half:,:] + else: + audio_1 = audio[:,:n_half,:] + audio_0 = audio[:,n_half:,:] + + output = self.nn[k]((audio_0, spect)) + s = output[:, n_half:, :] + b = output[:, :n_half, :] + audio_1 = torch.exp(s)*audio_1 + b + s_list.append(s) + + if k%2 == 0: + audio = torch.cat([audio[:,:n_half,:], audio_1],1) + else: + audio = torch.cat([audio_1, audio[:,n_half:,:]], 1) + output_audio.append(audio) + return torch.cat(output_audio,1), s_list, s_conv_list + """ + + def infer(self, spect, sigma=1.0): + spect = self.upsample(spect) + # trim conv artifacts. maybe pad spec to kernel multiple + time_cutoff = self.upsample.kernel_size[0] - self.upsample.stride[0] + spect = spect[:, :, :-time_cutoff] + + spect = spect.unfold(2, self.n_group, self.n_group).permute(0, 2, 1, 3) + spect = spect.contiguous().view(spect.size(0), spect.size(1), -1).permute(0, 2, 1) + + if spect.type() == 'torch.cuda.HalfTensor': + audio = torch.cuda.HalfTensor(spect.size(0), + self.n_remaining_channels, + spect.size(2)).normal_() + else: + audio = torch.cuda.FloatTensor(spect.size(0), + self.n_remaining_channels, + spect.size(2)).normal_() + + audio = torch.autograd.Variable(sigma*audio) + + for k in reversed(range(self.n_flows)): + n_half = int(audio.size(1)/2) + if k%2 == 0: + audio_0 = audio[:,:n_half,:] + audio_1 = audio[:,n_half:,:] + else: + audio_1 = audio[:,:n_half,:] + audio_0 = audio[:,n_half:,:] + + output = self.WN[k]((audio_0, spect)) + s = output[:, n_half:, :] + b = output[:, :n_half, :] + audio_1 = (audio_1 - b)/torch.exp(s) + if k%2 == 0: + audio = torch.cat([audio[:,:n_half,:], audio_1],1) + else: + audio = torch.cat([audio_1, audio[:,n_half:,:]], 1) + + audio = self.convinv[k](audio, reverse=True) + + if k%4 == 0 and k > 0: + if spect.type() == 'torch.cuda.HalfTensor': + z = torch.cuda.HalfTensor(spect.size(0), + self.n_early_size, + spect.size(2)).normal_() + else: + z = torch.cuda.FloatTensor(spect.size(0), + self.n_early_size, + spect.size(2)).normal_() + audio = torch.cat((sigma*z, audio),1) + + return audio.permute(0,2,1).contiguous().view(audio.size(0), -1).data + + @staticmethod + def remove_weightnorm(model): + waveglow = model + for WN in waveglow.WN: + WN.start = torch.nn.utils.remove_weight_norm(WN.start) + WN.in_layers = remove(WN.in_layers) + WN.cond_layers = remove(WN.cond_layers) + WN.res_skip_layers = remove(WN.res_skip_layers) + return waveglow diff --git a/models/tts/waveglow/inference.py b/models/tts/waveglow/inference.py new file mode 100644 index 0000000000..39e7e222d3 --- /dev/null +++ b/models/tts/waveglow/inference.py @@ -0,0 +1,84 @@ +# ***************************************************************************** +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the +# names of its contributors may be used to endorse or promote products +# derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" +# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE +# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE +# ARE DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY +# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# ***************************************************************************** +import os +from scipy.io.wavfile import write +import torch +from mel2samp import files_to_list, MAX_WAV_VALUE +from denoiser import Denoiser + + +def main(mel_files, waveglow_path, sigma, output_dir, sampling_rate, is_fp16, + denoiser_strength): + mel_files = files_to_list(mel_files) + waveglow = torch.load(waveglow_path)['model'] + waveglow = waveglow.remove_weightnorm(waveglow) + waveglow.cuda().eval() + if is_fp16: + from apex import amp + waveglow, _ = amp.initialize(waveglow, [], opt_level="O3") + + if denoiser_strength > 0: + denoiser = Denoiser(waveglow).cuda() + + for i, file_path in enumerate(mel_files): + file_name = os.path.splitext(os.path.basename(file_path))[0] + mel = torch.load(file_path) + mel = torch.autograd.Variable(mel.cuda()) + mel = torch.unsqueeze(mel, 0) + mel = mel.half() if is_fp16 else mel + with torch.no_grad(): + audio = waveglow.infer(mel, sigma=sigma) + if denoiser_strength > 0: + audio = denoiser(audio, denoiser_strength) + audio = audio * MAX_WAV_VALUE + audio = audio.squeeze() + audio = audio.cpu().numpy() + audio = audio.astype('int16') + audio_path = os.path.join( + output_dir, "{}_synthesis.wav".format(file_name)) + write(audio_path, sampling_rate, audio) + print(audio_path) + + +if __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument('-f', "--filelist_path", required=True) + parser.add_argument('-w', '--waveglow_path', + help='Path to waveglow decoder checkpoint with model') + parser.add_argument('-o', "--output_dir", required=True) + parser.add_argument("-s", "--sigma", default=1.0, type=float) + parser.add_argument("--sampling_rate", default=22050, type=int) + parser.add_argument("--is_fp16", action="store_true") + parser.add_argument("-d", "--denoiser_strength", default=0.0, type=float, + help='Removes model bias. Start with 0.1 and adjust') + + args = parser.parse_args() + + main(args.filelist_path, args.waveglow_path, args.sigma, args.output_dir, + args.sampling_rate, args.is_fp16, args.denoiser_strength) diff --git a/models/tts/waveglow/mel2samp.py b/models/tts/waveglow/mel2samp.py new file mode 100644 index 0000000000..f13f4af8a7 --- /dev/null +++ b/models/tts/waveglow/mel2samp.py @@ -0,0 +1,142 @@ +# ***************************************************************************** +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the +# names of its contributors may be used to endorse or promote products +# derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY +# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# *****************************************************************************\ +import os +import random +import argparse +import json +import torch +import torch.utils.data +import sys +from scipy.io.wavfile import read + +# We're using the audio processing from TacoTron2 to make sure it matches +sys.path.insert(0, 'tacotron2') +from tacotron2.layers import TacotronSTFT + +MAX_WAV_VALUE = 32768.0 + +def files_to_list(filename): + """ + Takes a text file of filenames and makes a list of filenames + """ + with open(filename, encoding='utf-8') as f: + files = f.readlines() + + files = [f.rstrip() for f in files] + return files + +def load_wav_to_torch(full_path): + """ + Loads wavdata into torch array + """ + sampling_rate, data = read(full_path) + return torch.from_numpy(data).float(), sampling_rate + + +class Mel2Samp(torch.utils.data.Dataset): + """ + This is the main class that calculates the spectrogram and returns the + spectrogram, audio pair. + """ + def __init__(self, training_files, segment_length, filter_length, + hop_length, win_length, sampling_rate, mel_fmin, mel_fmax): + self.audio_files = files_to_list(training_files) + random.seed(1234) + random.shuffle(self.audio_files) + self.stft = TacotronSTFT(filter_length=filter_length, + hop_length=hop_length, + win_length=win_length, + sampling_rate=sampling_rate, + mel_fmin=mel_fmin, mel_fmax=mel_fmax) + self.segment_length = segment_length + self.sampling_rate = sampling_rate + + def get_mel(self, audio): + audio_norm = audio / MAX_WAV_VALUE + audio_norm = audio_norm.unsqueeze(0) + audio_norm = torch.autograd.Variable(audio_norm, requires_grad=False) + melspec = self.stft.mel_spectrogram(audio_norm) + melspec = torch.squeeze(melspec, 0) + return melspec + + def __getitem__(self, index): + # Read audio + filename = self.audio_files[index] + audio, sampling_rate = load_wav_to_torch(filename) + if sampling_rate != self.sampling_rate: + raise ValueError("{} SR doesn't match target {} SR".format( + sampling_rate, self.sampling_rate)) + + # Take segment + if audio.size(0) >= self.segment_length: + max_audio_start = audio.size(0) - self.segment_length + audio_start = random.randint(0, max_audio_start) + audio = audio[audio_start:audio_start+self.segment_length] + else: + audio = torch.nn.functional.pad(audio, (0, self.segment_length - audio.size(0)), 'constant').data + + mel = self.get_mel(audio) + audio = audio / MAX_WAV_VALUE + + return (mel, audio) + + def __len__(self): + return len(self.audio_files) + +# =================================================================== +# Takes directory of clean audio and makes directory of spectrograms +# Useful for making test sets +# =================================================================== +if __name__ == "__main__": + # Get defaults so it can work with no Sacred + parser = argparse.ArgumentParser() + parser.add_argument('-f', "--filelist_path", required=True) + parser.add_argument('-c', '--config', type=str, + help='JSON file for configuration') + parser.add_argument('-o', '--output_dir', type=str, + help='Output directory') + args = parser.parse_args() + + with open(args.config) as f: + data = f.read() + data_config = json.loads(data)["data_config"] + mel2samp = Mel2Samp(**data_config) + + filepaths = files_to_list(args.filelist_path) + + # Make directory if it doesn't exist + if not os.path.isdir(args.output_dir): + os.makedirs(args.output_dir) + os.chmod(args.output_dir, 0o775) + + for filepath in filepaths: + audio, sr = load_wav_to_torch(filepath) + melspectrogram = mel2samp.get_mel(audio) + filename = os.path.basename(filepath) + new_filepath = args.output_dir + '/' + filename + '.pt' + print(new_filepath) + torch.save(melspectrogram, new_filepath) diff --git a/models/tts/waveglow/requirements.txt b/models/tts/waveglow/requirements.txt new file mode 100644 index 0000000000..20c8b3eefb --- /dev/null +++ b/models/tts/waveglow/requirements.txt @@ -0,0 +1,10 @@ +torch==1.0 +matplotlib==2.1.0 +tensorflow +numpy==1.13.3 +inflect==0.2.5 +librosa==0.6.0 +scipy==1.0.0 +tensorboardX==1.1 +Unidecode==1.0.22 +pillow diff --git a/models/tts/waveglow/train.py b/models/tts/waveglow/train.py new file mode 100644 index 0000000000..e8035c1528 --- /dev/null +++ b/models/tts/waveglow/train.py @@ -0,0 +1,188 @@ +# ***************************************************************************** +# Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved. +# +# Redistribution and use in source and binary forms, with or without +# modification, are permitted provided that the following conditions are met: +# * Redistributions of source code must retain the above copyright +# notice, this list of conditions and the following disclaimer. +# * Redistributions in binary form must reproduce the above copyright +# notice, this list of conditions and the following disclaimer in the +# documentation and/or other materials provided with the distribution. +# * Neither the name of the NVIDIA CORPORATION nor the +# names of its contributors may be used to endorse or promote products +# derived from this software without specific prior written permission. +# +# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND +# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED +# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE +# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY +# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES +# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; +# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND +# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT +# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS +# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE. +# +# ***************************************************************************** +import argparse +import json +import os +import torch + +#=====START: ADDED FOR DISTRIBUTED====== +from distributed import init_distributed, apply_gradient_allreduce, reduce_tensor +from torch.utils.data.distributed import DistributedSampler +#=====END: ADDED FOR DISTRIBUTED====== + +from torch.utils.data import DataLoader +from glow import WaveGlow, WaveGlowLoss +from mel2samp import Mel2Samp + +def load_checkpoint(checkpoint_path, model, optimizer): + assert os.path.isfile(checkpoint_path) + checkpoint_dict = torch.load(checkpoint_path, map_location='cpu') + iteration = checkpoint_dict['iteration'] + optimizer.load_state_dict(checkpoint_dict['optimizer']) + model_for_loading = checkpoint_dict['model'] + model.load_state_dict(model_for_loading.state_dict()) + print("Loaded checkpoint '{}' (iteration {})" .format( + checkpoint_path, iteration)) + return model, optimizer, iteration + +def save_checkpoint(model, optimizer, learning_rate, iteration, filepath): + print("Saving model and optimizer state at iteration {} to {}".format( + iteration, filepath)) + model_for_saving = WaveGlow(**waveglow_config).cuda() + model_for_saving.load_state_dict(model.state_dict()) + torch.save({'model': model_for_saving, + 'iteration': iteration, + 'optimizer': optimizer.state_dict(), + 'learning_rate': learning_rate}, filepath) + +def train(num_gpus, rank, group_name, output_directory, epochs, learning_rate, + sigma, iters_per_checkpoint, batch_size, seed, fp16_run, + checkpoint_path, with_tensorboard): + torch.manual_seed(seed) + torch.cuda.manual_seed(seed) + #=====START: ADDED FOR DISTRIBUTED====== + if num_gpus > 1: + init_distributed(rank, num_gpus, group_name, **dist_config) + #=====END: ADDED FOR DISTRIBUTED====== + + criterion = WaveGlowLoss(sigma) + model = WaveGlow(**waveglow_config).cuda() + + #=====START: ADDED FOR DISTRIBUTED====== + if num_gpus > 1: + model = apply_gradient_allreduce(model) + #=====END: ADDED FOR DISTRIBUTED====== + + optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate) + + if fp16_run: + from apex import amp + model, optimizer = amp.initialize(model, optimizer, opt_level='O1') + + # Load checkpoint if one exists + iteration = 0 + if checkpoint_path != "": + model, optimizer, iteration = load_checkpoint(checkpoint_path, model, + optimizer) + iteration += 1 # next iteration is iteration + 1 + + trainset = Mel2Samp(**data_config) + # =====START: ADDED FOR DISTRIBUTED====== + train_sampler = DistributedSampler(trainset) if num_gpus > 1 else None + # =====END: ADDED FOR DISTRIBUTED====== + train_loader = DataLoader(trainset, num_workers=1, shuffle=False, + sampler=train_sampler, + batch_size=batch_size, + pin_memory=False, + drop_last=True) + + # Get shared output_directory ready + if rank == 0: + if not os.path.isdir(output_directory): + os.makedirs(output_directory) + os.chmod(output_directory, 0o775) + print("output directory", output_directory) + + if with_tensorboard and rank == 0: + from tensorboardX import SummaryWriter + logger = SummaryWriter(os.path.join(output_directory, 'logs')) + + model.train() + epoch_offset = max(0, int(iteration / len(train_loader))) + # ================ MAIN TRAINNIG LOOP! =================== + for epoch in range(epoch_offset, epochs): + print("Epoch: {}".format(epoch)) + for i, batch in enumerate(train_loader): + model.zero_grad() + + mel, audio = batch + mel = torch.autograd.Variable(mel.cuda()) + audio = torch.autograd.Variable(audio.cuda()) + outputs = model((mel, audio)) + + loss = criterion(outputs) + if num_gpus > 1: + reduced_loss = reduce_tensor(loss.data, num_gpus).item() + else: + reduced_loss = loss.item() + + if fp16_run: + with amp.scale_loss(loss, optimizer) as scaled_loss: + scaled_loss.backward() + else: + loss.backward() + + optimizer.step() + + print("{}:\t{:.9f}".format(iteration, reduced_loss)) + if with_tensorboard and rank == 0: + logger.add_scalar('training_loss', reduced_loss, i + len(train_loader) * epoch) + + if (iteration % iters_per_checkpoint == 0): + if rank == 0: + checkpoint_path = "{}/waveglow_{}".format( + output_directory, iteration) + save_checkpoint(model, optimizer, learning_rate, iteration, + checkpoint_path) + + iteration += 1 + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument('-c', '--config', type=str, + help='JSON file for configuration') + parser.add_argument('-r', '--rank', type=int, default=0, + help='rank of process for distributed') + parser.add_argument('-g', '--group_name', type=str, default='', + help='name of group for distributed') + args = parser.parse_args() + + # Parse configs. Globals nicer in this case + with open(args.config) as f: + data = f.read() + config = json.loads(data) + train_config = config["train_config"] + global data_config + data_config = config["data_config"] + global dist_config + dist_config = config["dist_config"] + global waveglow_config + waveglow_config = config["waveglow_config"] + + num_gpus = torch.cuda.device_count() + if num_gpus > 1: + if args.group_name == '': + print("WARNING: Multiple GPUs detected but no distributed group set") + print("Only running 1 GPU. Use distributed.py for multiple GPUs") + num_gpus = 1 + + if num_gpus == 1 and args.rank != 0: + raise Exception("Doing single GPU training on rank > 0") + + torch.backends.cudnn.enabled = True + torch.backends.cudnn.benchmark = False + train(num_gpus, args.rank, args.group_name, **train_config) diff --git a/models/tts/waveglow/waveglow_logo.png b/models/tts/waveglow/waveglow_logo.png new file mode 100755 index 0000000000..d15bed6048 Binary files /dev/null and b/models/tts/waveglow/waveglow_logo.png differ diff --git a/rustfmt.toml b/rustfmt.toml new file mode 100644 index 0000000000..e4a54388f3 --- /dev/null +++ b/rustfmt.toml @@ -0,0 +1,42 @@ +# Storyteller rustfmt rules +# (Not yet happy with these.) + +# Four spaces is excessive in any language. +# Conserve horizontal space for small laptop screens. +tab_spaces = 2 + +# Only Unix +newline_style = "unix" + +# We won't change any vendor code. +# Note: This appears not to work in stable, so it's necessary to run `cargo +nightly fmt`. +ignore = [ + "vendor" +] + +# Item-level imports makes for better diffs and fewer merge conflicts. +# Another suitable option is "Module", but that's not the default either. +# (This is commented out for now to reduce the first round of rustfmt diffs, but will be enabled.) +imports_granularity = "Item" + +# It's better to keep stdlib, crate (self), and external crates separate. +# (This is commented out for now to reduce the first round of rustfmt diffs, but will be enabled.) +group_imports = "StdExternalCrate" + +# The default is wrong +match_block_trailing_comma = true + +# 'try!' shouldn't exist in our codebase. +use_try_shorthand = true + +# MAYBE +#struct_lit_single_line = true + +# This fixes method chaining. +# (rustfmt ordinarily puts everything on one line.) +use_small_heuristics = "Off" + +# TEMPORARY +#reorder_imports = false +#max_width = 180 # default 100 +#chain_width = 10 diff --git a/scripts/download_internet_file.py b/scripts/download_internet_file.py new file mode 100755 index 0000000000..7485701842 --- /dev/null +++ b/scripts/download_internet_file.py @@ -0,0 +1,92 @@ +#!/usr/bin/env python3 + +# For youtube-dl +from __future__ import unicode_literals + +import argparse +import gdown +import glob +import os +import re +import urllib +#import youtube_dl # NB: No longer maintained +import yt_dlp + +parser = argparse.ArgumentParser(description='Download Google Drive files, Youtube videos, or web urls') + +parser.add_argument('--url', type=str, help='the web url, google drive link, or youtube url', required=True) +parser.add_argument('--output_filename', type=str, help='the output filename', required=True) + +# This URL format gets an OAuth screen!! UGH!! +# 'https://drive.google.com/file/d/15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe/view?usp=sharing' +# Fixed URL: +# 'https://drive.google.com/uc?id=15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe' + +def _normalize_google_drive_url(url): + """ + URLs must be in a certain form, or they present the user with an OAuth wall. + This attempts to fix URLs in the wrong format. + """ + matches = re.match(r"^https:\/\/drive\.google\.com\/file\/d\/([^\/]+)\/", url) + + if not matches: + return url + + match_groups = matches.groups() + + if not match_groups or not len(match_groups) == 1: + return url + + document_token = match_groups[0] + + # NB(2022-02-17): Google appears to have broken downloads + # Using this fix with URL structure: + # https://github.com/wkentaro/gdown/issues/146#issuecomment-1042382215 + # OLD FORMAT: return 'https://drive.google.com/uc?id={}'.format(document_token) + return 'https://drive.google.com/uc?export=download&confirm=pbef&id={}'.format(document_token) + + +def download_url_or_gdrive(url, output_filename): + url = _normalize_google_drive_url(url) + print('Normalized URL: {}'.format(url)) + + gdown.download(url, output_filename, quiet=False) + +def download_youtube(youtube_url, output_filename): + + filename_expected = "ytdl_download/filename" + + ydl_opts = { + 'outtmpl': filename_expected, + 'noplaylist': True, + } + + with yt_dlp.YoutubeDL(ydl_opts) as ydl: + ydl.download([youtube_url]) + + # ytdl forcefully appends the extension + filename_search = '{}*'.format(filename_expected) + + found_filename = None + for file in glob.glob(filename_search): + found_filename = file + break + print('Downloaded as: {}'.format(found_filename)) + + os.rename(found_filename, output_filename) + +args = parser.parse_args() +url = args.url +output_filename = args.output_filename + +print('URL: {}'.format(url)) +print('Output filename: {}'.format(output_filename)) + +parsed_url = urllib.parse.urlparse(url) +if 'youtube' in parsed_url.hostname: + print('Downloading from YouTube...') + download_youtube(url, output_filename) +else: + print('Downloading from Gdrive or Web...') + download_url_or_gdrive(url, output_filename) + diff --git a/seed/bootstrap_inserts_roles_etc.sh b/seed/bootstrap_inserts_roles_etc.sh new file mode 100755 index 0000000000..fdf2fc31ed --- /dev/null +++ b/seed/bootstrap_inserts_roles_etc.sh @@ -0,0 +1,24 @@ +#!/bin/bash +# ONLY FOR LOCAL DEV + +USERNAME='storyteller' +PASSWORD='password' +HOST='localhost' + +# --default-character-set=UTF8mb4 \ + +echo 'Inserting System User Roles...' +mysql -u "${USERNAME}" \ + -p${PASSWORD} \ + -h $HOST \ + -D storyteller \ + -e "source ./seed/sql/system_roles.sql" + +echo 'Inserting Badges...' +mysql -u $USERNAME \ + -p${PASSWORD} \ + -h $HOST \ + -D storyteller \ + -e "source ./seed/sql/user_badges.sql" + + diff --git a/seed/sql/system_roles.sql b/seed/sql/system_roles.sql new file mode 100644 index 0000000000..d9e1c42c11 --- /dev/null +++ b/seed/sql/system_roles.sql @@ -0,0 +1,83 @@ + + +INSERT INTO user_roles +SET + slug = 'user', + name = 'User', + -- Usage + can_use_tts = TRUE, + can_use_w2l = TRUE, + can_delete_own_tts_results = TRUE, + can_delete_own_w2l_results = TRUE, + can_delete_own_account = TRUE, + -- Contribution + can_upload_tts_models = TRUE, + can_upload_w2l_templates = TRUE, + can_delete_own_tts_models = TRUE, + can_delete_own_w2l_templates = TRUE, + -- Moderation + can_approve_w2l_templates = FALSE, + can_edit_other_users_profiles = FALSE, + can_edit_other_users_tts_models = FALSE, + can_edit_other_users_w2l_templates = FALSE, + can_delete_other_users_tts_models = FALSE, + can_delete_other_users_tts_results = FALSE, + can_delete_other_users_w2l_templates = FALSE, + can_delete_other_users_w2l_results = FALSE, + can_ban_users = FALSE, + can_delete_users = FALSE; + +INSERT INTO user_roles +SET + slug = 'mod', + name = 'Moderator', + -- Usage + can_use_tts = TRUE, + can_use_w2l = TRUE, + can_delete_own_tts_results = TRUE, + can_delete_own_w2l_results = TRUE, + can_delete_own_account = TRUE, + -- Contribution + can_upload_tts_models = TRUE, + can_upload_w2l_templates = TRUE, + can_delete_own_tts_models = TRUE, + can_delete_own_w2l_templates = TRUE, + -- Moderation + can_approve_w2l_templates = TRUE, + can_edit_other_users_profiles = TRUE, + can_edit_other_users_tts_models = TRUE, + can_edit_other_users_w2l_templates = TRUE, + can_delete_other_users_tts_models = TRUE, + can_delete_other_users_tts_results = TRUE, + can_delete_other_users_w2l_templates = TRUE, + can_delete_other_users_w2l_results = TRUE, + can_ban_users = TRUE, + can_delete_users = TRUE; + +INSERT INTO user_roles +SET + slug = 'admin', + name = 'Admin', + -- Usage + can_use_tts = TRUE, + can_use_w2l = TRUE, + can_delete_own_tts_results = TRUE, + can_delete_own_w2l_results = TRUE, + can_delete_own_account = TRUE, + -- Contribution + can_upload_tts_models = TRUE, + can_upload_w2l_templates = TRUE, + can_delete_own_tts_models = TRUE, + can_delete_own_w2l_templates = TRUE, + -- Moderation + can_approve_w2l_templates = TRUE, + can_edit_other_users_profiles = TRUE, + can_edit_other_users_tts_models = TRUE, + can_edit_other_users_w2l_templates = TRUE, + can_delete_other_users_tts_models = TRUE, + can_delete_other_users_tts_results = TRUE, + can_delete_other_users_w2l_templates = TRUE, + can_delete_other_users_w2l_results = TRUE, + can_ban_users = TRUE, + can_delete_users = TRUE; + diff --git a/seed/sql/user_badges.sql b/seed/sql/user_badges.sql new file mode 100644 index 0000000000..5e05455939 --- /dev/null +++ b/seed/sql/user_badges.sql @@ -0,0 +1,146 @@ +-- USER SIGNUP / TIME LIMITED + +INSERT INTO badges +SET + slug = 'early_user', + title = 'Early Adopter of Vocodes', + description = 'Early Adopter of Vocodes', + image_url = ''; + +-- TTS MODEL UPLOADS + +INSERT INTO badges +SET + slug = 'tts_model_uploader_1', + title = 'Uploaded 1 TTS Model', + description = 'Uploaded 1 TTS model to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_5', + title = 'Uploaded 5 TTS Models', + description = 'Uploaded 5 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_10', + title = 'Uploaded 10 TTS Models', + description = 'Uploaded 10 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_20', + title = 'Uploaded 20 TTS Models', + description = 'Uploaded 20 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_50', + title = 'Uploaded 50 TTS Models', + description = 'Uploaded 50 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_100', + title = 'Uploaded 100 TTS Models', + description = 'Uploaded 100 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_150', + title = 'Uploaded 150 TTS Models', + description = 'Uploaded 150 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_200', + title = 'Uploaded 200 TTS Models', + description = 'Uploaded 200 TTS models to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'tts_model_uploader_250', + title = 'Uploaded 250 TTS Models', + description = 'Uploaded 250 TTS models to Vocodes', + image_url = ''; + +-- W2L TEMPLATE UPLOADS + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_1', + title = 'Uploaded 1 W2L Template', + description = 'Uploaded 1 W2L Template to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_10', + title = 'Uploaded 10 W2L Templates', + description = 'Uploaded 10 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_50', + title = 'Uploaded 50 W2L Templates', + description = 'Uploaded 50 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_100', + title = 'Uploaded 100 W2L Templates', + description = 'Uploaded 100 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_200', + title = 'Uploaded 200 W2L Templates', + description = 'Uploaded 200 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_500', + title = 'Uploaded 500 W2L Templates', + description = 'Uploaded 500 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_1000', + title = 'Uploaded 1000 W2L Templates', + description = 'Uploaded 1000 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_2000', + title = 'Uploaded 2000 W2L Templates', + description = 'Uploaded 2000 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_5000', + title = 'Uploaded 5000 W2L Templates', + description = 'Uploaded 5000 W2L Templates to Vocodes', + image_url = ''; + +INSERT INTO badges +SET + slug = 'w2l_template_uploader_10000', + title = 'Uploaded 10000 W2L Templates', + description = 'Uploaded 10000 W2L Templates to Vocodes', + image_url = ''; + diff --git a/sqlx-data.json b/sqlx-data.json new file mode 100644 index 0000000000..4c2b18f872 --- /dev/null +++ b/sqlx-data.json @@ -0,0 +1,11252 @@ +{ + "db": "MySQL", + "0033c76be0ea1c99069f1a8df1ce400c7064526feaf5dc88930da0718a885b19": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.is_banned as creator_is_banned,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n w2l.title,\n w2l.description_markdown,\n w2l.description_rendered_html,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.creator_set_visibility,\n w2l.is_public_listing_approved,\n w2l.is_locked_from_use,\n w2l.is_locked_from_user_modification,\n w2l.creator_ip_address_creation,\n w2l.creator_ip_address_last_update,\n w2l.created_at,\n w2l.updated_at,\n w2l.user_deleted_at,\n w2l.mod_deleted_at\nFROM w2l_templates as w2l\nJOIN users\n ON users.token = w2l.creator_user_token\nWHERE w2l.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 4, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 8, + "name": "description_markdown", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 9, + "name": "description_rendered_html", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 10, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 11, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 12, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 13, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 14, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 15, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 16, + "name": "is_public_listing_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 17, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "is_locked_from_user_modification", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 19, + "name": "creator_ip_address_creation", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 20, + "name": "creator_ip_address_last_update", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 21, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 23, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 24, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + true, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "01f91b5f3c3357aa3397102f52524b4221124caecda2a62f4a4bb31855547b68": { + "query": "\nUPDATE api_tokens\nSET\n deleted_at = CURRENT_TIMESTAMP,\n ip_address_last_update = ?\nWHERE\n user_token = ?\n AND api_token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "047d52d9fe0d14843a5ade550e3fea114e1e369d330d2239f49194d34486c2de": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n tts.title,\n tts.ietf_language_tag,\n tts.ietf_primary_language_subtag,\n tts.is_locked_from_use,\n tts.is_front_page_featured,\n tts.is_twitch_featured,\n tts.maybe_suggested_unique_bot_command,\n tts.created_at,\n tts.updated_at\nFROM tts_models as tts\nJOIN users\nON\n users.token = tts.creator_user_token\nWHERE\n users.username = ?\n AND tts.user_deleted_at IS NULL\n AND tts.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 7, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 9, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "is_front_page_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "is_twitch_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "maybe_suggested_unique_bot_command", + "type_info": { + "type": "VarString", + "flags": { + "bits": 132 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 13, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 14, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false + ] + } + }, + "05c3ebf54cd8a5374d6f4bbe21b3c6027ce6314917af8acd0ed136d375d53342": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n w2l.title,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.is_public_listing_approved,\n w2l.created_at,\n w2l.updated_at\nFROM w2l_templates as w2l\nJOIN users\nON\n users.token = w2l.creator_user_token\nWHERE\n users.username = ?\n AND w2l.user_deleted_at IS NULL\n AND w2l.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 7, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 8, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 9, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 10, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 11, + "name": "is_public_listing_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 13, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + true, + false, + false + ] + } + }, + "089986414149d2d189f4262f8aa7488d35f6a162a4e56f06c1b897092d37ea20": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.is_banned as creator_is_banned,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n w2l.title,\n w2l.description_markdown,\n w2l.description_rendered_html,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.creator_set_visibility,\n w2l.is_public_listing_approved,\n w2l.is_locked_from_use,\n w2l.is_locked_from_user_modification,\n w2l.creator_ip_address_creation,\n w2l.creator_ip_address_last_update,\n w2l.created_at,\n w2l.updated_at,\n w2l.user_deleted_at,\n w2l.mod_deleted_at\nFROM w2l_templates as w2l\nJOIN users\n ON users.token = w2l.creator_user_token\nWHERE\n w2l.token = ?\n AND w2l.user_deleted_at IS NULL\n AND w2l.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 4, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 8, + "name": "description_markdown", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 9, + "name": "description_rendered_html", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 10, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 11, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 12, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 13, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 14, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 15, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 16, + "name": "is_public_listing_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 17, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "is_locked_from_user_modification", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 19, + "name": "creator_ip_address_creation", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 20, + "name": "creator_ip_address_last_update", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 21, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 23, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 24, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + true, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "09e655ea4b3504cbb3bc18b92c6c25861c941ba073e3aead414e560b05bd7638": { + "query": "\nSELECT count(*) as count\nFROM tts_models\nWHERE\n creator_user_token = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false + ] + } + }, + "10d93ef1350115742d8dddb951f98dfebabf7df3e44577b97b52343f0650eba1": { + "query": "\nSELECT\n badges.slug,\n badges.title,\n badges.description,\n badges.image_url,\n user_badges.created_at as user_created_at\n\nFROM badges\nJOIN user_badges\nON\n badges.slug = user_badges.badge_slug\nWHERE\n user_badges.user_token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "slug", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 2, + "name": "description", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 3, + "name": "image_url", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 4, + "name": "user_created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + }, + "12a176252df63d1e3a958362d754c828991d64ade427ab11a2025098d1581976": { + "query": "\nSELECT\n assignments.model_token AS tts_model_token,\n assignments.category_token\nFROM\n tts_category_assignments AS assignments\nJOIN\n tts_models AS tts\n ON assignments.model_token = tts.token\nWHERE\n tts.is_locked_from_use IS FALSE\n AND tts.user_deleted_at IS NULL\n AND tts.mod_deleted_at IS NULL\n AND assignments.deleted_at IS NULL\n\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "tts_model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "category_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false + ] + } + }, + "150c98c00f4cd6954428736583eae913be578c4e69dac317d4445bc0e97743c9": { + "query": "\nSELECT\n next_id\nFROM\n tts_result_synthetic_ids\nWHERE\n user_token = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "next_id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 20 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false + ] + } + }, + "28cf5843d68d6cb0e67eb316ef9e7e5fb05a6be2433f86c081576b279167508d": { + "query": "\nINSERT INTO tts_result_synthetic_ids\nSET\n user_token = ?,\n next_id = 1\nON DUPLICATE KEY UPDATE\n user_token = ?,\n next_id = next_id + 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "2b80162b8957ac044f245af93d678aee976535d58fe0de0b39087323aa720792": { + "query": "\nUPDATE tts_models\nSET\n maybe_mod_user_token = ?,\n mod_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "314904656a333caa75a8644e73f4d338b7bd6f68c7ea9e37b5bc7f7b257bf34d": { + "query": "\nSELECT 1 as does_exist\nFROM user_badges\nWHERE\n user_token = ?\nAND\n badge_slug = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "does_exist", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 2 + } + } + ], + "parameters": { + "Right": 2 + }, + "nullable": [ + false + ] + } + }, + "345cbabbe69652c89ccd384ed7db44f70c53791488cf63721f9563291b45db3e": { + "query": "\nINSERT INTO user_sessions (\n token,\n user_token,\n ip_address_creation,\n expires_at\n)\nVALUES ( ?, ?, ?, NOW() + interval 1 year )\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "376e36b1357952599866262c7356ba35f15ccdfb338d8ccdbdc06ccd69331fc3": { + "query": "\nUPDATE tts_results\nSET\n mod_deleted_at = CURRENT_TIMESTAMP,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "38d9e9860cc87398a27053a90d46183bae6589555c7f4424929d9f900ae1809d": { + "query": "\nSELECT\n id,\n token AS inference_job_token,\n uuid_idempotency_token,\n\n model_token,\n raw_inference_text,\n\n creator_ip_address,\n maybe_creator_user_token,\n creator_set_visibility as `creator_set_visibility: crate::column_types::record_visibility::RecordVisibility`,\n\n is_from_api,\n is_for_twitch,\n\n status as `status: crate::column_types::job_status::JobStatus`,\n priority_level,\n attempt_count,\n failure_reason,\n\n created_at,\n updated_at,\n retry_at\nFROM tts_inference_jobs\nWHERE\n (\n status IN (\"pending\", \"attempt_failed\")\n )\n AND\n (\n retry_at IS NULL\n OR\n retry_at < CURRENT_TIMESTAMP\n )\n ORDER BY priority_level DESC, id ASC\n LIMIT ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "inference_job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "uuid_idempotency_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 3, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "raw_inference_text", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 5, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 6, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "creator_set_visibility: crate::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 8, + "name": "is_from_api", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "is_for_twitch", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "status: crate::column_types::job_status::JobStatus", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 11, + "name": "priority_level", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 33 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 12, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 13, + "name": "failure_reason", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 14, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 15, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 16, + "name": "retry_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + true, + false, + false, + true + ] + } + }, + "40a3301e9a3201e74d49b6b339f67291e2ae22d315188401708609ac2d28a68c": { + "query": "\nUPDATE tts_results\nSET\n mod_deleted_at = NULL,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "417dd070beeaf6ccd86852b7e99bf7fda5c5b15da4a3bcc066a0e2bb90ca7df1": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n w2l.title,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.is_public_listing_approved,\n w2l.created_at,\n w2l.updated_at\nFROM w2l_templates as w2l\nJOIN users\n ON users.token = w2l.creator_user_token\nWHERE\n w2l.is_public_listing_approved IS TRUE\n AND w2l.user_deleted_at IS NULL\n AND w2l.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 7, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 8, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 9, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 10, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 11, + "name": "is_public_listing_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 13, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + true, + false, + false + ] + } + }, + "48dd7546f75e3c997d1b9c8f1ef88adcd5a2d37dae517aa7047f285c8c760c59": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.text_preprocessing_algorithm,\n tts.maybe_default_pretrained_vocoder,\n\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n users.is_banned as creator_is_banned,\n\n tts.title,\n tts.description_markdown,\n tts.description_rendered_html,\n\n tts.ietf_language_tag,\n tts.ietf_primary_language_subtag,\n\n tts.is_front_page_featured,\n tts.is_twitch_featured,\n\n tts.maybe_suggested_unique_bot_command,\n\n tts.creator_set_visibility,\n\n tts.is_locked_from_use,\n tts.is_locked_from_user_modification,\n\n tts.created_at,\n tts.updated_at,\n\n tts.creator_ip_address_creation,\n tts.creator_ip_address_last_update,\n tts.mod_deleted_at,\n tts.user_deleted_at\n\nFROM tts_models as tts\nJOIN users\n ON users.token = tts.creator_user_token\nWHERE\n tts.token = ?\n AND tts.user_deleted_at IS NULL\n AND tts.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "text_preprocessing_algorithm", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 3, + "name": "maybe_default_pretrained_vocoder", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 4, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 5, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 7, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 8, + "name": "creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 10, + "name": "description_markdown", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 11, + "name": "description_rendered_html", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 12, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 13, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 14, + "name": "is_front_page_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "is_twitch_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "maybe_suggested_unique_bot_command", + "type_info": { + "type": "VarString", + "flags": { + "bits": 132 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 17, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 18, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 19, + "name": "is_locked_from_user_modification", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "creator_ip_address_creation", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 23, + "name": "creator_ip_address_last_update", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 24, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 25, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "4c7a92b6ac43e271096138f3594ca27b0ab6622b1e29a28f4a034c7732be9312": { + "query": "\nSELECT\n users.username,\n users.display_name,\n users.email_gravatar_hash as gravatar_hash,\n creator_user_token,\n count(*) as uploaded_count\nFROM w2l_templates\nJOIN users\nON\n users.token = creator_user_token\nWHERE\n is_public_listing_approved IS TRUE\n AND users.is_banned IS FALSE\n AND w2l_templates.user_deleted_at IS NULL\n AND w2l_templates.mod_deleted_at IS NULL\nGROUP BY creator_user_token\nORDER BY uploaded_count desc\nLIMIT 25;\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 1, + "name": "display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "uploaded_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + }, + "4cea10944b8869419ecdfea30c58f49a87c090b9083078043177c33737ecb949": { + "query": "\nINSERT INTO voice_clone_requests\nSET\n token = ?,\n uuid_idempotency_token = ?,\n\n maybe_user_token = ?,\n email_address = ?,\n discord_username = ?,\n\n is_for_private_use = ?,\n is_for_public_use = ?,\n\n is_for_studio = ?,\n is_for_twitch_tts = ?,\n is_for_api_use = ?,\n is_for_music = ?,\n is_for_games = ?,\n is_for_other = ?,\n optional_notes_on_use = ?,\n\n is_own_voice = ?,\n is_third_party_voice = ?,\n optional_notes_on_subject = ?,\n\n has_clean_audio_recordings = ?,\n has_good_microphone = ?,\n\n optional_questions = ?,\n optional_extra_comments = ?,\n\n ip_address_creation = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 22 + }, + "nullable": [] + } + }, + "4df4765816d8c431a241b5ea2e7e9b292bf8e603937b095fe256cbc8a7a165f6": { + "query": "\nSELECT\n internal_token,\n api_token,\n maybe_short_description,\n created_at,\n updated_at\nFROM api_tokens\nWHERE\n user_token = ?\n AND deleted_at IS NULL\nORDER BY created_at DESC\nLIMIT 5\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "internal_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "api_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "maybe_short_description", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1161 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 4, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + true, + false, + false + ] + } + }, + "4f61f621682ba255814f8a55455ea2b7fba9b6c85affaf87abfe2f81e11322f0": { + "query": "\nSELECT\n id,\n status\nFROM tts_inference_jobs\nWHERE id = ?\nFOR UPDATE\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false + ] + } + }, + "4ff1b1e0cc36f6c2f1a9b099f38751ae7f35666be939b4074bdfc98e031cd6ac": { + "query": "\nSELECT\n w2l_results.token as w2l_result_token,\n w2l_results.maybe_tts_inference_result_token,\n\n w2l_results.public_bucket_video_path,\n\n w2l_templates.token as maybe_w2l_template_token,\n w2l_templates.template_type,\n w2l_templates.title as template_title,\n\n users.token as maybe_creator_user_token,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n users.is_banned as maybe_creator_is_banned,\n\n template_users.token as maybe_template_creator_user_token,\n template_users.username as maybe_template_creator_username,\n template_users.display_name as maybe_template_creator_display_name,\n template_users.email_gravatar_hash as maybe_template_creator_gravatar_hash,\n template_users.is_banned as maybe_template_creator_is_banned,\n\n w2l_results.creator_set_visibility,\n\n w2l_results.file_size_bytes,\n w2l_results.frame_width,\n w2l_results.frame_height,\n w2l_results.duration_millis,\n w2l_results.created_at,\n w2l_results.updated_at,\n\n w2l_results.creator_ip_address,\n w2l_results.user_deleted_at,\n w2l_results.mod_deleted_at\n\nFROM w2l_results\nLEFT OUTER JOIN w2l_templates\n ON w2l_results.maybe_w2l_template_token = w2l_templates.token\nLEFT OUTER JOIN users\n ON w2l_results.maybe_creator_user_token = users.token\nLEFT OUTER JOIN users as template_users\n ON w2l_templates.creator_user_token = template_users.token\nWHERE\n w2l_results.token = ?\n AND w2l_results.user_deleted_at IS NULL\n AND w2l_results.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "w2l_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "maybe_tts_inference_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "public_bucket_video_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "maybe_w2l_template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 392 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 5, + "name": "template_title", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "maybe_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 8, + "name": "maybe_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 9, + "name": "maybe_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 10, + "name": "maybe_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "maybe_template_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 12, + "name": "maybe_template_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 13, + "name": "maybe_template_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 14, + "name": "maybe_template_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 15, + "name": "maybe_template_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 17, + "name": "file_size_bytes", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 18, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 19, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 20, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 21, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 23, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 24, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 25, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + true, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "53f0792791317fcd4038308a6630f541d5508ae913cbfa9bc3d4081fba9a048a": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.maybe_default_pretrained_vocoder,\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n tts.title,\n tts.private_bucket_hash,\n tts.created_at,\n tts.updated_at,\n tts.user_deleted_at,\n tts.mod_deleted_at\nFROM tts_models as tts\nJOIN users\nON users.token = tts.creator_user_token\nWHERE tts.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "maybe_default_pretrained_vocoder", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 3, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 7, + "name": "private_bucket_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 9, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 10, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 11, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "551749bc153f9c8683a3a1ca5bb50744322aac046db34327501cf991485c4dee": { + "query": "\nUPDATE tts_models\nSET\n creator_ip_address_last_update = ?,\n user_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "57cf798a0fe32a1d6d0349ee61baac6e536fe3e1a3aa7bd93b34c519812bf70a": { + "query": "\nSELECT\n id,\n token AS inference_job_token,\n uuid_idempotency_token,\n\n model_token,\n raw_inference_text,\n\n creator_ip_address,\n maybe_creator_user_token,\n creator_set_visibility as `creator_set_visibility: crate::column_types::record_visibility::RecordVisibility`,\n\n is_from_api,\n is_for_twitch,\n\n status as `status: crate::column_types::job_status::JobStatus`,\n priority_level,\n attempt_count,\n failure_reason,\n\n created_at,\n updated_at,\n retry_at\nFROM tts_inference_jobs\nWHERE\n (\n status IN (\"pending\", \"attempt_failed\")\n )\n AND\n (\n retry_at IS NULL\n OR\n retry_at < CURRENT_TIMESTAMP\n )\n LIMIT ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "inference_job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "uuid_idempotency_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 3, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "raw_inference_text", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 5, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 6, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "creator_set_visibility: crate::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 8, + "name": "is_from_api", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "is_for_twitch", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "status: crate::column_types::job_status::JobStatus", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 11, + "name": "priority_level", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 33 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 12, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 13, + "name": "failure_reason", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 14, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 15, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 16, + "name": "retry_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + true, + false, + false, + true + ] + } + }, + "5e4a00c339ce4d7fac19aff728449156b1e02b9e1d9fd534ac37d4ac79f44820": { + "query": "\nINSERT INTO twitch_bits_events\nSET\n sender_twitch_user_id = ?,\n sender_twitch_username = ?,\n sender_twitch_username_lowercase = ?,\n destination_channel_id = ?,\n destination_channel_name = ?,\n is_anonymous= ?,\n bits_used = ?,\n total_bits_used = ?,\n chat_message = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "5e54fd8569b2d6d3a0d11acee6bd03b75400b6325952eb16c8dac68a7933b13a": { + "query": "\nSELECT\n users.token as user_token,\n username,\n display_name,\n email_gravatar_hash,\n profile_markdown,\n profile_rendered_html,\n user_role_slug,\n is_banned,\n disable_gravatar,\n discord_username,\n twitch_username,\n twitter_username,\n patreon_username,\n github_username,\n created_at\nFROM users\nWHERE\n users.username = ?\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 3, + "name": "email_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "profile_markdown", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 5, + "name": "profile_rendered_html", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 6, + "name": "user_role_slug", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 7, + "name": "is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 8, + "name": "disable_gravatar", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "discord_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 10, + "name": "twitch_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 11, + "name": "twitter_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 12, + "name": "patreon_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 13, + "name": "github_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 14, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + true, + true, + true, + false + ] + } + }, + "5eacdf9c65dc13820077be2930cca7047aed07804ed2e7b2ff38e12a617b4858": { + "query": "\nINSERT INTO twitch_channel_point_events\nSET\n sender_twitch_user_id = ?,\n sender_twitch_username = ?,\n sender_twitch_username_lowercase = ?,\n destination_channel_id = ?,\n destination_channel_name = ?,\n title = ?,\n prompt = ?,\n user_text_input = ?,\n redemption_id = ?,\n reward_id = ?,\n reward_cost = ?,\n is_sub_only = ?,\n max_per_stream = ?,\n max_per_user_per_stream = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 14 + }, + "nullable": [] + } + }, + "65ce50a3aafe979e77a541d576228d7fcfe76e33400427d01382da36c1585edb": { + "query": "\nUPDATE tts_inference_jobs\nSET\n status = ?,\n on_success_result_token = ?,\n failure_reason = NULL,\n retry_at = NULL\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "663506988f146de6f559e5df4f4a2b607e9421a69bb57db1f656774bf8b375db": { + "query": "\nUPDATE w2l_results\nSET\n mod_deleted_at = CURRENT_TIMESTAMP,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "6aaa7c04e6984434bd2feb6f42a9014cd45e995b94e4d5e29f2d502f09136691": { + "query": "\nUPDATE tts_models\nSET\n maybe_mod_user_token = ?,\n mod_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "78006a7a9388557d53c9a9cc8aeec84496e87355501c5b412f788a2291077473": { + "query": "\nSELECT\n slug,\n name,\n\n can_use_tts,\n can_use_w2l,\n can_delete_own_tts_results,\n can_delete_own_w2l_results,\n can_delete_own_account,\n\n can_upload_tts_models,\n can_upload_w2l_templates,\n can_delete_own_tts_models,\n can_delete_own_w2l_templates,\n\n can_approve_w2l_templates,\n can_edit_other_users_profiles,\n can_edit_other_users_tts_models,\n can_edit_other_users_w2l_templates,\n can_delete_other_users_tts_models,\n can_delete_other_users_tts_results,\n can_delete_other_users_w2l_templates,\n can_delete_other_users_w2l_results,\n can_ban_users,\n can_delete_users,\n\n created_at,\n updated_at\n\nFROM user_roles\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "slug", + "type_info": { + "type": "String", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 1, + "name": "name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 2, + "name": "can_use_tts", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 3, + "name": "can_use_w2l", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 4, + "name": "can_delete_own_tts_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 5, + "name": "can_delete_own_w2l_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 6, + "name": "can_delete_own_account", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 7, + "name": "can_upload_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 8, + "name": "can_upload_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "can_delete_own_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "can_delete_own_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "can_approve_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "can_edit_other_users_profiles", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 13, + "name": "can_edit_other_users_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 14, + "name": "can_edit_other_users_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "can_delete_other_users_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "can_delete_other_users_tts_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 17, + "name": "can_delete_other_users_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "can_delete_other_users_w2l_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 19, + "name": "can_ban_users", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 20, + "name": "can_delete_users", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 21, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ] + } + }, + "7cd0ad5d7b9a2c393f799beadc1e0a6bc85d952e07a1b82a52115b99be40ff6d": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n w2l.title,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.is_public_listing_approved,\n w2l.created_at,\n w2l.updated_at\nFROM w2l_templates as w2l\nJOIN users\n ON users.token = w2l.creator_user_token\nWHERE\n w2l.user_deleted_at IS NULL\n AND w2l.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 7, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 8, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 9, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 10, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 11, + "name": "is_public_listing_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 13, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + true, + false, + false + ] + } + }, + "80e8ad515ab8fbfadc12463b87059a1eb6554be32527729d6075fc46c966630c": { + "query": "\nINSERT INTO twitch_oauth_tokens\nSET\n internal_token = ?,\n oauth_refresh_grouping_token = ?,\n maybe_user_token = ?,\n twitch_user_id = ?,\n twitch_username = ?,\n twitch_username_lowercase = ?,\n access_token = ?,\n maybe_refresh_token = ?,\n refresh_count = ?,\n token_type = ?,\n expires_in_seconds = ?,\n has_bits_read = ?,\n has_channel_read_redemptions= ?,\n has_channel_read_subscriptions = ?,\n has_chat_edit = ?,\n has_chat_read = ?,\n ip_address_creation = ?,\n expires_at = DATE_ADD(NOW(), INTERVAL ? SECOND)\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 18 + }, + "nullable": [] + } + }, + "865ce236444a36d597939660b14eb4b0cee05c70fdb73f31c2fb8cddddfef42c": { + "query": "\nINSERT INTO firehose_entries\nSET\n token = ?,\n event_type = ?,\n maybe_target_user_token = ?,\n maybe_target_entity_token = ?,\n maybe_created_entity_token = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 5 + }, + "nullable": [] + } + }, + "87d7261dd2a589220c5539f51dde42ef2407677f7ece7968761426c6fd0e0360": { + "query": "\nSELECT\n category.token as category_token,\n category.model_type,\n category.maybe_super_category_token,\n category.can_directly_have_models,\n category.can_have_subcategories,\n category.can_only_mods_apply,\n category.name,\n category.maybe_dropdown_name,\n category.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash AS creator_gravatar_hash,\n category.creator_ip_address_creation,\n category.creator_ip_address_last_update,\n category.is_mod_approved,\n category.maybe_mod_comments,\n category.maybe_mod_user_token,\n category.created_at,\n category.updated_at,\n category.deleted_at\nFROM model_categories AS category\nJOIN users\nON category.creator_user_token = users.token\nWHERE\n category.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "category_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "maybe_super_category_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "can_directly_have_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 4, + "name": "can_have_subcategories", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 5, + "name": "can_only_mods_apply", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 6, + "name": "name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 7, + "name": "maybe_dropdown_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 8, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 9, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 10, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 11, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 12, + "name": "creator_ip_address_creation", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 13, + "name": "creator_ip_address_last_update", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 14, + "name": "is_mod_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "maybe_mod_comments", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 16, + "name": "maybe_mod_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 17, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 18, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 19, + "name": "deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + true, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + true, + true, + true, + false, + false, + true + ] + } + }, + "89b66d91d1f2cd5c9871671b21a90e77d9041572a90cc6be6dd49845424677fa": { + "query": "\nUPDATE tts_inference_jobs\nSET status='complete_failure'\nWHERE\n status IN ('pending');\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 0 + }, + "nullable": [] + } + }, + "8ac132004408188cc55710185539ff6cd2011de3e0f80b4c91e2bf8b3212135e": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n tts.title,\n tts.ietf_language_tag,\n tts.ietf_primary_language_subtag,\n tts.is_locked_from_use,\n tts.is_front_page_featured,\n tts.is_twitch_featured,\n tts.maybe_suggested_unique_bot_command,\n tts.created_at,\n tts.updated_at\nFROM tts_models as tts\nJOIN users\nON\n users.token = tts.creator_user_token\nWHERE\n users.username = ?\n AND tts.is_locked_from_use IS FALSE\n AND tts.user_deleted_at IS NULL\n AND tts.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 7, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 9, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "is_front_page_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "is_twitch_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "maybe_suggested_unique_bot_command", + "type_info": { + "type": "VarString", + "flags": { + "bits": 132 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 13, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 14, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false + ] + } + }, + "8c3f19f67e8871ada450832b5342efd870c87975f065f3d6b3789d537cab01ed": { + "query": "\nUPDATE tts_inference_jobs\nSET status='complete_failure'\nWHERE\n status IN ('pending', 'attempt_failed');\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 0 + }, + "nullable": [] + } + }, + "8c49202950a8d6cf87bbc60517168fba4c9b0c52c71a25d766be49fedbd28827": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n tts.title,\n tts.ietf_language_tag,\n tts.ietf_primary_language_subtag,\n tts.is_locked_from_use,\n tts.is_front_page_featured,\n tts.is_twitch_featured,\n tts.maybe_suggested_unique_bot_command,\n tts.created_at,\n tts.updated_at\nFROM tts_models as tts\nJOIN users\n ON users.token = tts.creator_user_token\nWHERE\n tts.user_deleted_at IS NULL\n AND tts.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 7, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 9, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "is_front_page_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "is_twitch_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "maybe_suggested_unique_bot_command", + "type_info": { + "type": "VarString", + "flags": { + "bits": 132 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 13, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 14, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false + ] + } + }, + "8c6a56bfa985324a0dd138c9ac9570d72195ee721d5e2d6c50c4266d1590ed20": { + "query": "\nUPDATE twitch_event_rules\nSET\n event_match_predicate = ?,\n event_response = ?,\n rule_is_disabled = ?,\n ip_address_last_update = ?\nWHERE\n token = ?\n AND deleted_at IS NULL\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 5 + }, + "nullable": [] + } + }, + "8d4fd8365f8f2843c8ade46b5bfd4add943f7dbdf02c582edf8d5f690ba608cd": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.text_preprocessing_algorithm,\n tts.maybe_default_pretrained_vocoder,\n\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n users.is_banned as creator_is_banned,\n\n tts.title,\n tts.description_markdown,\n tts.description_rendered_html,\n\n tts.ietf_language_tag,\n tts.ietf_primary_language_subtag,\n\n tts.is_front_page_featured,\n tts.is_twitch_featured,\n\n tts.maybe_suggested_unique_bot_command,\n\n tts.creator_set_visibility,\n\n tts.is_locked_from_use,\n tts.is_locked_from_user_modification,\n\n tts.created_at,\n tts.updated_at,\n\n tts.creator_ip_address_creation,\n tts.creator_ip_address_last_update,\n tts.mod_deleted_at,\n tts.user_deleted_at\n\nFROM tts_models as tts\nJOIN users\n ON users.token = tts.creator_user_token\nWHERE tts.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "text_preprocessing_algorithm", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 3, + "name": "maybe_default_pretrained_vocoder", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 4, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 5, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 7, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 8, + "name": "creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 10, + "name": "description_markdown", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 11, + "name": "description_rendered_html", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 12, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 13, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 14, + "name": "is_front_page_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "is_twitch_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "maybe_suggested_unique_bot_command", + "type_info": { + "type": "VarString", + "flags": { + "bits": 132 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 17, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 18, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 19, + "name": "is_locked_from_user_modification", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "creator_ip_address_creation", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 23, + "name": "creator_ip_address_last_update", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 24, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 25, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "8ff27fc5583db97f157f1b808f593a72c77dc5dcf5a633cd9c1162af6b97c011": { + "query": "\nUPDATE w2l_templates\nSET\n creator_ip_address_last_update = ?,\n user_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "9134347abc2b89b7336394894ed420db725be4efebd1c36e959652a10daf8551": { + "query": "\nUPDATE twitch_event_rules\nSET\n deleted_at = CURRENT_TIMESTAMP,\n ip_address_last_update = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "91e88859f2975447b7232264b6fbf98e35064531b7bc1ffe66552a8d7fee8bed": { + "query": "\nSELECT\n token,\n event_category as `event_category: crate::column_types::twitch_event_category::TwitchEventCategory`,\n event_match_predicate,\n event_response,\n user_specified_rule_order,\n rule_is_disabled,\n created_at,\n updated_at,\n deleted_at\nFROM twitch_event_rules\nWHERE\n token = ?\n AND user_token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "event_category: crate::column_types::twitch_event_category::TwitchEventCategory", + "type_info": { + "type": "String", + "flags": { + "bits": 4489 + }, + "char_set": 224, + "max_size": 56 + } + }, + { + "ordinal": 2, + "name": "event_match_predicate", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 3, + "name": "event_response", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 4, + "name": "user_specified_rule_order", + "type_info": { + "type": "Long", + "flags": { + "bits": 33 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 5, + "name": "rule_is_disabled", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 6, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 7, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 8, + "name": "deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 2 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + true + ] + } + }, + "922a7dbcfa24a2c5901d0453c229b066b3b4eea2c73fa533651bb27365f17be0": { + "query": "\nSELECT\n token,\n\n maybe_user_token,\n email_address,\n discord_username,\n\n is_for_private_use,\n is_for_public_use,\n\n is_for_studio,\n is_for_twitch_tts,\n is_for_api_use,\n is_for_music,\n is_for_games,\n is_for_other,\n optional_notes_on_use,\n\n is_own_voice,\n is_third_party_voice,\n optional_notes_on_subject,\n\n has_clean_audio_recordings,\n has_good_microphone,\n\n optional_questions,\n optional_extra_comments,\n\n created_at,\n updated_at\nFROM voice_clone_requests\nWHERE\n token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "maybe_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "email_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "discord_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 4, + "name": "is_for_private_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 5, + "name": "is_for_public_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 6, + "name": "is_for_studio", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 7, + "name": "is_for_twitch_tts", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 8, + "name": "is_for_api_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "is_for_music", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "is_for_games", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "is_for_other", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "optional_notes_on_use", + "type_info": { + "type": "Blob", + "flags": { + "bits": 144 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 13, + "name": "is_own_voice", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 14, + "name": "is_third_party_voice", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "optional_notes_on_subject", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 16, + "name": "has_clean_audio_recordings", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 17, + "name": "has_good_microphone", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "optional_questions", + "type_info": { + "type": "Blob", + "flags": { + "bits": 144 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 19, + "name": "optional_extra_comments", + "type_info": { + "type": "Blob", + "flags": { + "bits": 144 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + true, + false, + false, + true, + true, + false, + false + ] + } + }, + "961f792c24bee6009e3e8a215561db932871a53c48839b8cb80292de00e7b2e1": { + "query": "\nUPDATE tts_results\nSET\n user_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "9b63758edee6341cb417073775eea2782137c680c1827189d6f7a93c40d70e44": { + "query": "\nSELECT\n tts.token as model_token,\n tts.tts_model_type,\n tts.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n tts.title,\n tts.ietf_language_tag,\n tts.ietf_primary_language_subtag,\n tts.is_locked_from_use,\n tts.is_front_page_featured,\n tts.is_twitch_featured,\n tts.maybe_suggested_unique_bot_command,\n tts.created_at,\n tts.updated_at\nFROM tts_models as tts\nJOIN users\n ON users.token = tts.creator_user_token\nWHERE\n tts.is_locked_from_use IS FALSE\n AND tts.user_deleted_at IS NULL\n AND tts.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 7, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": { + "bits": 129 + }, + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 9, + "name": "is_locked_from_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 9 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "is_front_page_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "is_twitch_featured", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "maybe_suggested_unique_bot_command", + "type_info": { + "type": "VarString", + "flags": { + "bits": 132 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 13, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 14, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false + ] + } + }, + "a60b239dff945c0089ded216dc6b4ea1bd7acb2125671bebdd34e5e1312270ce": { + "query": "\nSELECT\n users.username,\n users.display_name,\n users.email_gravatar_hash as gravatar_hash,\n creator_user_token,\n count(*) as uploaded_count\nFROM tts_models\nJOIN users\nON\n users.token = creator_user_token\nWHERE\n users.is_banned IS FALSE\n AND tts_models.user_deleted_at IS NULL\n AND tts_models.mod_deleted_at IS NULL\nGROUP BY creator_user_token\nORDER BY uploaded_count desc\nLIMIT 25;\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 1, + "name": "display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "uploaded_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + }, + "ab6f25eb0f6340d49aa705835ed428434d8b6ab14b0740f6a5ebe7841b9ece20": { + "query": "\nUPDATE tts_inference_jobs\nSET\n status = 'started',\n attempt_count = attempt_count + 1,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "adf2d5eb967fa2bc00fad1c8fcc2e1bf4384d39dc54e8240d78bb4ca6fcc3409": { + "query": "\nUPDATE w2l_results\nSET\n mod_deleted_at = NULL,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "aeeed0077e567e4b67ffa30e32f5475e168c6c68d2f206cc7b6c37168384089c": { + "query": "\nINSERT INTO twitch_oauth_tokens\nSET\n internal_token = ?,\n oauth_refresh_grouping_token = ?,\n maybe_user_token = ?,\n twitch_user_id = ?,\n twitch_username = ?,\n twitch_username_lowercase = ?,\n access_token = ?,\n maybe_refresh_token = ?,\n refresh_count = ?,\n token_type = ?,\n expires_in_seconds = NULL,\n has_bits_read = ?,\n has_channel_read_redemptions= ?,\n has_channel_read_subscriptions = ?,\n has_chat_edit = ?,\n has_chat_read = ?,\n ip_address_creation = ?,\n expires_at = NULL\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 16 + }, + "nullable": [] + } + }, + "af032ec861cb91f15ec763d9452e03c23797af15b5714565aa330f6cc420e8d6": { + "query": "\nINSERT INTO api_tokens\nSET\n internal_token = ?,\n api_token = ?,\n user_token = ?,\n uuid_idempotency_token = ?,\n ip_address_creation = ?,\n ip_address_last_update = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 6 + }, + "nullable": [] + } + }, + "b483867c3f767a01cc80e4c0e0f1950c65154e06c13fc2fcc6e3ec96e8e2c956": { + "query": "\nUPDATE api_tokens\nSET\n maybe_short_description = ?,\n ip_address_last_update = ?\nWHERE\n user_token = ?\n AND api_token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 4 + }, + "nullable": [] + } + }, + "b58e215a61914d06cfdadc92130e8423e1b2db24c03850697873cd21a75dae61": { + "query": "\nSELECT\n tts_results.token as tts_result_token,\n\n tts_results.raw_inference_text,\n\n tts_results.model_token as tts_model_token,\n tts_models.title as tts_model_title,\n\n tts_results.maybe_pretrained_vocoder_used,\n\n users.token as maybe_creator_user_token,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n users.is_banned as maybe_creator_is_banned,\n\n model_users.token as maybe_model_creator_user_token,\n model_users.username as maybe_model_creator_username,\n model_users.display_name as maybe_model_creator_display_name,\n model_users.email_gravatar_hash as maybe_model_creator_gravatar_hash,\n model_users.is_banned as maybe_model_creator_is_banned,\n\n tts_results.public_bucket_wav_audio_path,\n tts_results.public_bucket_spectrogram_path,\n\n tts_results.creator_set_visibility,\n\n tts_results.file_size_bytes,\n tts_results.duration_millis,\n tts_results.created_at,\n tts_results.updated_at,\n\n tts_results.creator_ip_address,\n tts_results.user_deleted_at,\n tts_results.mod_deleted_at\n\nFROM tts_results\nLEFT OUTER JOIN tts_models\n ON tts_results.model_token = tts_models.token\nLEFT OUTER JOIN users\n ON tts_results.maybe_creator_user_token = users.token\nLEFT OUTER JOIN users as model_users\n ON tts_models.creator_user_token = model_users.token\nWHERE\n tts_results.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "tts_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "raw_inference_text", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 2, + "name": "tts_model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "tts_model_title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 4, + "name": "maybe_pretrained_vocoder_used", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 5, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "maybe_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 7, + "name": "maybe_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 8, + "name": "maybe_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 9, + "name": "maybe_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "maybe_model_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 11, + "name": "maybe_model_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 12, + "name": "maybe_model_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 13, + "name": "maybe_model_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 14, + "name": "maybe_model_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "public_bucket_wav_audio_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 16, + "name": "public_bucket_spectrogram_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 17, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 18, + "name": "file_size_bytes", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 19, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 23, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 24, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "b8746ae4ff2d4dcd54563eff64877f924bb119681ad78f1300146e00b762ca4f": { + "query": "\nINSERT INTO tts_results\nSET\n token = ?,\n\n model_token = ?,\n maybe_pretrained_vocoder_used = ?,\n raw_inference_text = ?,\n raw_inference_text_hash_sha2 = ?,\n normalized_inference_text = ?,\n\n maybe_creator_user_token = ?,\n maybe_creator_synthetic_id = ?,\n\n creator_ip_address = ?,\n creator_set_visibility = ?,\n\n public_bucket_wav_audio_path = ?,\n public_bucket_spectrogram_path = ?,\n\n file_size_bytes = ?,\n duration_millis = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 14 + }, + "nullable": [] + } + }, + "ba31fefd3e96fd238fc2bcdfab64f390a17293f2c0996ae891ed8821e57512cd": { + "query": "\nSELECT\n w2l_results.token as w2l_result_token,\n w2l_results.maybe_tts_inference_result_token,\n\n w2l_results.public_bucket_video_path,\n\n w2l_templates.token as maybe_w2l_template_token,\n w2l_templates.template_type,\n w2l_templates.title as template_title,\n\n users.token as maybe_creator_user_token,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n users.is_banned as maybe_creator_is_banned,\n\n template_users.token as maybe_template_creator_user_token,\n template_users.username as maybe_template_creator_username,\n template_users.display_name as maybe_template_creator_display_name,\n template_users.email_gravatar_hash as maybe_template_creator_gravatar_hash,\n template_users.is_banned as maybe_template_creator_is_banned,\n\n w2l_results.creator_set_visibility,\n\n w2l_results.file_size_bytes,\n w2l_results.frame_width,\n w2l_results.frame_height,\n w2l_results.duration_millis,\n w2l_results.created_at,\n w2l_results.updated_at,\n\n w2l_results.creator_ip_address,\n w2l_results.user_deleted_at,\n w2l_results.mod_deleted_at\n\nFROM w2l_results\nLEFT OUTER JOIN w2l_templates\n ON w2l_results.maybe_w2l_template_token = w2l_templates.token\nLEFT OUTER JOIN users\n ON w2l_results.maybe_creator_user_token = users.token\nLEFT OUTER JOIN users as template_users\n ON w2l_templates.creator_user_token = template_users.token\nWHERE\n w2l_results.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "w2l_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "maybe_tts_inference_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "public_bucket_video_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "maybe_w2l_template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 392 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 5, + "name": "template_title", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "maybe_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 8, + "name": "maybe_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 9, + "name": "maybe_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 10, + "name": "maybe_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "maybe_template_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 12, + "name": "maybe_template_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 13, + "name": "maybe_template_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 14, + "name": "maybe_template_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 15, + "name": "maybe_template_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 17, + "name": "file_size_bytes", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 18, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 19, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 20, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 21, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 23, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 24, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 25, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + true, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "c0f56e2656b566d3a8c869d39b0b56d1fbb75d7d9e817a24b10add44d348092d": { + "query": "\nSELECT count(*) as count\nFROM w2l_templates\nWHERE\n creator_user_token = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false + ] + } + }, + "c4640a618e1d179fa95adce67a329818f8b1626ac88fa76dbb871b8788136191": { + "query": "\nUPDATE w2l_templates\nSET\n maybe_mod_user_token = ?,\n mod_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "c7bff6a198838ec7209f2b48ef1d8fe850e66aa532ed562aac75dd3adfaf48a9": { + "query": "\nSELECT\n tts_results.token as tts_result_token,\n\n tts_results.raw_inference_text,\n\n tts_results.model_token as tts_model_token,\n tts_models.title as tts_model_title,\n\n tts_results.maybe_pretrained_vocoder_used,\n\n users.token as maybe_creator_user_token,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n users.is_banned as maybe_creator_is_banned,\n\n model_users.token as maybe_model_creator_user_token,\n model_users.username as maybe_model_creator_username,\n model_users.display_name as maybe_model_creator_display_name,\n model_users.email_gravatar_hash as maybe_model_creator_gravatar_hash,\n model_users.is_banned as maybe_model_creator_is_banned,\n\n tts_results.public_bucket_wav_audio_path,\n tts_results.public_bucket_spectrogram_path,\n\n tts_results.creator_set_visibility,\n\n tts_results.file_size_bytes,\n tts_results.duration_millis,\n tts_results.created_at,\n tts_results.updated_at,\n\n tts_results.creator_ip_address,\n tts_results.user_deleted_at,\n tts_results.mod_deleted_at\n\nFROM tts_results\nLEFT OUTER JOIN tts_models\n ON tts_results.model_token = tts_models.token\nLEFT OUTER JOIN users\n ON tts_results.maybe_creator_user_token = users.token\nLEFT OUTER JOIN users as model_users\n ON tts_models.creator_user_token = model_users.token\nWHERE\n tts_results.token = ?\n AND tts_results.user_deleted_at IS NULL\n AND tts_results.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "tts_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "raw_inference_text", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 2, + "name": "tts_model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "tts_model_title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 4, + "name": "maybe_pretrained_vocoder_used", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 5, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "maybe_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 7, + "name": "maybe_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 8, + "name": "maybe_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 9, + "name": "maybe_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "maybe_model_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 11, + "name": "maybe_model_creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 12, + "name": "maybe_model_creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 13, + "name": "maybe_model_creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 14, + "name": "maybe_model_creator_is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "public_bucket_wav_audio_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 16, + "name": "public_bucket_spectrogram_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 17, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 18, + "name": "file_size_bytes", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 19, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 22, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 23, + "name": "user_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 24, + "name": "mod_deleted_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 136 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true + ] + } + }, + "cb732b1732a36c826317f835b87c17d1b6d061c9363e9477f2603b435f9bf28f": { + "query": "\nUPDATE tts_inference_jobs\nSET status='complete_failure'\nWHERE\n status IN ('pending')\n AND priority_level = 0;\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 0 + }, + "nullable": [] + } + }, + "d51e6e901cc8641b7b076f42b0fc507dad2168617c33c615ef97417f3a78ccc1": { + "query": "\nUPDATE api_tokens\nSET\n deleted_at = CURRENT_TIMESTAMP,\n ip_address_last_update = ?\nWHERE\n user_token = ?\nAND api_token NOT IN (?, ?, ?, ?, ?)\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 7 + }, + "nullable": [] + } + }, + "d71ef2cfed337f9e3f9e1c2255bdec6347763910fd589482f0b0e56b794d7bb7": { + "query": "\nUPDATE tts_models\nSET\n creator_ip_address_last_update = ?,\n user_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "da3197d279b94817c6c8a6277aff7a0cf44d66e6777f46680b1b07de48875dd5": { + "query": "\nINSERT INTO user_badges\nSET\n user_token = ?,\n badge_slug = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "dbce51eb1c80dded902c9509e50713483c7131f115e8ec10646e99e88110278b": { + "query": "\nUPDATE tts_results\nSET\n user_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "def4202f63874e7b04d8f8a5da028a9aacc4dfe19b8898d3d6820d393a633cf3": { + "query": "\nINSERT INTO tts_inference_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n model_token = ?,\n raw_inference_text = ?,\n maybe_creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n is_from_api = ?,\n is_for_twitch = ?,\n priority_level = ?,\n status = \"pending\"\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 10 + }, + "nullable": [] + } + }, + "dfe8dedd1838627a7bebbc2facabc3c51a7f3122e19ece6259629696643ae877": { + "query": "\nSELECT NOW() as present_time\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "present_time", + "type_info": { + "type": "Datetime", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false + ] + } + }, + "e0ad2a7fe4ae795e4f097f604d5cfcf9605298b58cde5ca23bdf9c8ec7383b00": { + "query": "\nUPDATE w2l_templates\nSET\n creator_ip_address_last_update = ?,\n user_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "e0d442313b7861debaef984c5809647bd5df85106c194b66e864674b6af95330": { + "query": "\nUPDATE tts_inference_jobs\nSET\n status = \"dead\",\n failure_reason = ?,\n retry_at = NULL\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "e3a0bd8073c3c0d84b9ab0e2e1ad8371f14b10801b7e4658b73c42ce1b27fa0b": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n w2l.title,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.is_public_listing_approved,\n w2l.created_at,\n w2l.updated_at\nFROM w2l_templates as w2l\nJOIN users\nON\n users.token = w2l.creator_user_token\nWHERE\n users.username = ?\n AND w2l.is_public_listing_approved IS TRUE\n AND w2l.user_deleted_at IS NULL\n AND w2l.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 7, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 8, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 9, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 10, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 11, + "name": "is_public_listing_approved", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 8 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 13, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + true, + false, + false + ] + } + }, + "e3a4f23245c8cfd9eb288548174cb4c5d550ceeff3ee4ec8cd76a94c5afa3eda": { + "query": "\nINSERT INTO twitch_event_rules\nSET\n uuid_idempotency_token = ?,\n token = ?,\n user_token = ?,\n event_category = ?,\n event_match_predicate = ?,\n event_response = ?,\n user_specified_rule_order = ?,\n rule_is_disabled = ?,\n ip_address_creation = ?,\n ip_address_last_update = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 10 + }, + "nullable": [] + } + }, + "e4e5be790f087b9aa6bee87e2315140bf621e4919362ebe72ab070163f18e9c8": { + "query": "\nUPDATE w2l_templates\nSET\n maybe_mod_user_token = ?,\n mod_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "e5ef2bda28221a570da6dcfa73e76cae9910a0b4ee53005a87176ecf1be73dd7": { + "query": "\nSELECT\n token,\n\n maybe_user_token,\n email_address,\n discord_username,\n\n is_for_private_use,\n is_for_public_use,\n\n is_for_studio,\n is_for_twitch_tts,\n is_for_api_use,\n is_for_music,\n is_for_games,\n is_for_other,\n optional_notes_on_use,\n\n is_own_voice,\n is_third_party_voice,\n optional_notes_on_subject,\n\n has_clean_audio_recordings,\n has_good_microphone,\n\n optional_questions,\n optional_extra_comments,\n\n created_at,\n updated_at\nFROM voice_clone_requests\nWHERE\n maybe_user_token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "maybe_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "email_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "discord_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 4, + "name": "is_for_private_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 5, + "name": "is_for_public_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 6, + "name": "is_for_studio", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 7, + "name": "is_for_twitch_tts", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 8, + "name": "is_for_api_use", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "is_for_music", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 10, + "name": "is_for_games", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 11, + "name": "is_for_other", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 12, + "name": "optional_notes_on_use", + "type_info": { + "type": "Blob", + "flags": { + "bits": 144 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 13, + "name": "is_own_voice", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 14, + "name": "is_third_party_voice", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "optional_notes_on_subject", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 16, + "name": "has_clean_audio_recordings", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 17, + "name": "has_good_microphone", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "optional_questions", + "type_info": { + "type": "Blob", + "flags": { + "bits": 144 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 19, + "name": "optional_extra_comments", + "type_info": { + "type": "Blob", + "flags": { + "bits": 144 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + true, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + false, + false, + true, + false, + false, + true, + true, + false, + false + ] + } + }, + "e9ec861d6eb135370f9baf3a866282ec5686d4b299fc4ec10900ab32ea8577d8": { + "query": "\nUPDATE w2l_results\nSET\n user_deleted_at = NULL\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "eacc0e022cd188eb930bbf4ffbf2fc8fc9081fa7a44119924b114422ceae7ad3": { + "query": "\nSELECT\n token,\n event_category as `event_category: crate::column_types::twitch_event_category::TwitchEventCategory`,\n event_match_predicate,\n event_response,\n user_specified_rule_order,\n rule_is_disabled,\n created_at,\n updated_at\nFROM twitch_event_rules\nWHERE\n user_token = ?\n AND deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "event_category: crate::column_types::twitch_event_category::TwitchEventCategory", + "type_info": { + "type": "String", + "flags": { + "bits": 4489 + }, + "char_set": 224, + "max_size": 56 + } + }, + { + "ordinal": 2, + "name": "event_match_predicate", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 3, + "name": "event_response", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 67108860 + } + }, + { + "ordinal": 4, + "name": "user_specified_rule_order", + "type_info": { + "type": "Long", + "flags": { + "bits": 33 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 5, + "name": "rule_is_disabled", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 6, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 7, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false + ] + } + }, + "eadbe83c69c09f3e049ec7c760d1af0313cec8f3d82c0993e526738aa4d0a6dc": { + "query": "\nSELECT\n ip_bans.ip_address,\n\n ip_bans.maybe_target_user_token,\n banned_users.username as maybe_target_username,\n\n ip_bans.mod_user_token,\n mod_users.username as mod_username,\n mod_users.display_name as mod_display_name,\n\n ip_bans.mod_notes,\n ip_bans.created_at,\n ip_bans.updated_at\nFROM\n ip_address_bans AS ip_bans\nLEFT OUTER JOIN users as banned_users\n ON ip_bans.maybe_target_user_token = banned_users.token\nJOIN users as mod_users\n ON ip_bans.mod_user_token = mod_users.token\nWHERE\n ip_bans.deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 1, + "name": "maybe_target_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "maybe_target_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 3, + "name": "mod_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "mod_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "mod_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "mod_notes", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 7, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 8, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + true, + true, + false, + false, + false, + false, + false, + false + ] + } + }, + "f19300f4e8434c2669dac062e5a79c5ca71767c7380aab0ab40b51be8ff449d8": { + "query": "\nUPDATE tts_inference_jobs\nSET\n status = ?,\n failure_reason = ?,\n retry_at = NOW() + interval 1 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "ffa85f7c65101a43653c724479ee080e0a74577d2ae37863b039f3a452493d1e": { + "query": "\nUPDATE w2l_results\nSET\n user_deleted_at = CURRENT_TIMESTAMP\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "02060ab9b8ca472a2bd4d4fa3120743e90148ccd6a4cb13bc491d00f36ed7a8f": { + "query": "\nSELECT\n id,\n status\nFROM w2l_inference_jobs\nWHERE id = ?\nFOR UPDATE\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false + ] + } + }, + "0caed749627ff22ae1cea52621a375156259a731e8070459593082b80d463491": { + "query": "\nSELECT\n id,\n status\nFROM tts_model_upload_jobs\nWHERE id = ?\nFOR UPDATE\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false + ] + } + }, + "0ff13b84956c2625050211a5bb2352a489ef8435571c364c8b2d1d21d936cbd4": { + "query": "\nUPDATE tts_results\nSET\n creator_set_visibility = ?,\n maybe_mod_user_token = ?\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "1236640a30d0d5177e13c31b61b791874f7c70026003a58ba6e07e7af46c30c1": { + "query": "\nSELECT\n next_id\nFROM\n w2l_result_synthetic_ids\nWHERE\n user_token = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "next_id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 20 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false + ] + } + }, + "17c8d210d93ee9c0d8e8422eed605c9150f19f9bf72affcf111b15a67a5b9e70": { + "query": "\nUPDATE w2l_templates\nSET\n is_public_listing_approved = ?,\n is_locked_from_user_modification = ?,\n is_locked_from_use = ?,\n maybe_mod_comments = ?,\n maybe_mod_user_token = ?,\n version = version + 1\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 6 + }, + "nullable": [] + } + }, + "1cc62e5cdc8de18969d7e8018d91c38794e5a7d23f117e5b0a4126845f300e60": { + "query": "\nUPDATE user_sessions\nSET deleted_at = CURRENT_TIMESTAMP()\nWHERE\n token = ?\n AND deleted_at IS NULL\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "1d0f302c0b7c4301000c024d6bbae4d3b7370ed7c944e4c557c64b0bd2f31c14": { + "query": "\nSELECT public_count, all_count FROM (\n (SELECT count(*) as public_count \n FROM tts_models \n WHERE user_deleted_at IS NULL \n AND mod_deleted_at IS NULL) as pc,\n (SELECT count(*) as all_count \n FROM tts_models) as ac\n)\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "public_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 21 + } + }, + { + "ordinal": 1, + "name": "all_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false + ] + } + }, + "218b3fe667b34bfa158bb7d4f600a3b91a22460adf566dcdf6ce540006e6c684": { + "query": "\nUPDATE model_categories\nSET\n name = ?,\n maybe_dropdown_name = ?,\n\n can_directly_have_models = ?,\n can_have_subcategories = ?,\n can_only_mods_apply = ?,\n\n maybe_super_category_token = ?,\n\n is_mod_approved = ?,\n maybe_mod_user_token = ?,\n maybe_mod_comments = ?,\n\n version = version + 1\n\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 10 + }, + "nullable": [] + } + }, + "25e7858d370ac1359deb42ef5292e85ccd1fe7debf3a1ba73303e5882523a4d0": { + "query": "\nINSERT INTO w2l_template_upload_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n title = ?,\n template_type = ?,\n download_url = ?,\n status = \"pending\"\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 8 + }, + "nullable": [] + } + }, + "294c089b17fdd87d7e579c05a9d90e9bfa5b7465cafa9a95af08d58709212440": { + "query": "\nUPDATE w2l_template_upload_jobs\nSET\n status = ?,\n failure_reason = ?,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "2a7837b626439198db1d9d63b22f6f6cac602d5e6b09de45f17b70272d1679d4": { + "query": "\nUPDATE tts_models\nSET\n maybe_default_pretrained_vocoder = ?,\n title = ?,\n description_markdown = ?,\n description_rendered_html = ?,\n ietf_language_tag = ?,\n ietf_primary_language_subtag = ?,\n creator_set_visibility = ?,\n maybe_mod_user_token = ?,\n version = version + 1\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "2b04dd7db52ef3a97a6f420a22a4dbee7416ae58697438df69e87ffd47bbf350": { + "query": "\nUPDATE w2l_inference_jobs\nSET\n status = ?,\n failure_reason = ?,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "33cd0c0b8eaaf494ef97232ac89775b8eb6a22d310824b7aa7db310dd4fd41b1": { + "query": "\nUPDATE w2l_results\nSET\n creator_set_visibility = ?,\n maybe_mod_user_token = ?\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "34364f9bc851971f8291f5f9942dd23d10718ccb210a93a236d96d0591f65f35": { + "query": "\nINSERT INTO users (\n token,\n username,\n display_name,\n email_address,\n email_gravatar_hash,\n profile_markdown,\n profile_rendered_html,\n user_role_slug,\n password_hash,\n ip_address_creation,\n ip_address_last_login,\n ip_address_last_update\n)\nVALUES ( ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? )\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 12 + }, + "nullable": [] + } + }, + "34642399c3e91960768cb659b6006db291574d9cd07a5bace3db573eff20c3fb": { + "query": "\nINSERT INTO w2l_result_synthetic_ids\nSET\n user_token = ?,\n next_id = 1\nON DUPLICATE KEY UPDATE\n user_token = ?,\n next_id = next_id + 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "346c02b215cb35efaa35fe96d866a8c9056ac46e987b4e01f64222259dccdcbf": { + "query": "\nUPDATE\n ip_address_bans\nSET\n mod_user_token = ?,\n deleted_at = NULL\nWHERE\n ip_address = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "369a7576edc128cf2b310648d6fc45a2191c1263d433fe6d1461d15dbf26690e": { + "query": "\nINSERT INTO w2l_results\nSET\n token = ?,\n maybe_w2l_template_token = ?,\n maybe_creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n\n maybe_creator_synthetic_id = ?,\n\n public_bucket_video_path = ?,\n\n file_size_bytes = ?,\n mime_type = ?,\n frame_width = ?,\n frame_height = ?,\n duration_millis = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 12 + }, + "nullable": [] + } + }, + "3a2dc0a7e41dcbdf21f91d1a1e54e813f74b0bac7aea052302089cae195eeb53": { + "query": "\nUPDATE users\nSET\n profile_markdown = ?,\n profile_rendered_html = ?,\n discord_username = ?,\n twitter_username = ?,\n twitch_username = ?,\n github_username = ?,\n cashapp_username = ?,\n website_url = ?,\n version = version + 1\n\nWHERE users.token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "4014c934d3705ce4293cef2ca3ebf14d36b6d569a2af6743a4e0bf94cdda07f9": { + "query": "\nSELECT\n id,\n token AS inference_job_token,\n uuid_idempotency_token,\n\n maybe_w2l_template_token,\n maybe_public_image_bucket_location,\n maybe_tts_inference_result_token,\n maybe_public_audio_bucket_hash,\n maybe_public_audio_bucket_location,\n\n maybe_original_audio_filename,\n maybe_original_audio_download_url,\n maybe_audio_mime_type,\n\n creator_ip_address,\n maybe_creator_user_token,\n\n creator_set_visibility as `creator_set_visibility: database_queries::column_types::record_visibility::RecordVisibility`,\n disable_end_bump,\n disable_watermark,\n\n status,\n attempt_count,\n failure_reason,\n created_at,\n updated_at,\n retry_at\nFROM w2l_inference_jobs\nWHERE\n (\n status IN (\"pending\", \"attempt_failed\")\n )\n AND\n (\n retry_at IS NULL\n OR\n retry_at < CURRENT_TIMESTAMP\n )\n LIMIT ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "inference_job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "uuid_idempotency_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 3, + "name": "maybe_w2l_template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "maybe_public_image_bucket_location", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 5, + "name": "maybe_tts_inference_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "maybe_public_audio_bucket_hash", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 7, + "name": "maybe_public_audio_bucket_location", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 8, + "name": "maybe_original_audio_filename", + "type_info": { + "type": "String", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 9, + "name": "maybe_original_audio_download_url", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 10, + "name": "maybe_audio_mime_type", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 11, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 12, + "name": "maybe_creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 13, + "name": "creator_set_visibility: database_queries::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 14, + "name": "disable_end_bump", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "disable_watermark", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 17, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 18, + "name": "failure_reason", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 19, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 20, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 21, + "name": "retry_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + true, + false, + true, + false, + false, + false, + false, + false, + true, + false, + false, + true + ] + } + }, + "5473d3857478d001688998056377fefe4bcc792477f303b41fa3cf64f79170f3": { + "query": "\n INSERT INTO tts_category_assignments\n SET\n model_token = ?,\n category_token = ?,\n category_removal_user_token = ?,\n ip_address_creation = ?,\n ip_address_last_update = ?,\n deleted_at = CURRENT_TIMESTAMP\n ON DUPLICATE KEY UPDATE\n model_token = ?,\n category_token = ?,\n category_removal_user_token = ?,\n ip_address_last_update = ?,\n deleted_at = CURRENT_TIMESTAMP\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "553034e2300bb7cce54964b1cd93d6f4826eccef9b1aee3bc9abd69433ee0d5c": { + "query": "\nUPDATE tts_models\nSET\n is_public_listing_approved = ?,\n is_locked_from_user_modification = ?,\n is_locked_from_use = ?,\n maybe_mod_comments = ?,\n maybe_mod_user_token = ?,\n is_front_page_featured = ?,\n is_twitch_featured = ?,\n maybe_suggested_unique_bot_command = ?,\n version = version + 1\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "56a0a2bfc3c73d6845ca9cf1d285b175a40b50a250b7effb9ce6474e58113d6e": { + "query": "\nSELECT\n jobs.token as job_token,\n\n jobs.status,\n jobs.attempt_count,\n jobs.on_success_result_token as maybe_model_token,\n\n jobs.created_at,\n jobs.updated_at\n\nFROM tts_model_upload_jobs as jobs\n\nWHERE jobs.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 2, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 3, + "name": "maybe_model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 5, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + false, + false + ] + } + }, + "57928b1001a10a88ac5a07f43b1dd30f6dee6edbc0808be99fe1233f9e276c56": { + "query": "\nSELECT\n jobs.token as job_token,\n\n jobs.status,\n jobs.attempt_count,\n jobs.on_success_result_token as maybe_result_token,\n results.public_bucket_video_path as maybe_public_bucket_video_path,\n\n jobs.maybe_w2l_template_token,\n w2l.template_type as w2l_template_type,\n w2l.title,\n\n jobs.created_at,\n jobs.updated_at\n\nFROM w2l_inference_jobs as jobs\nJOIN w2l_templates as w2l\n ON w2l.token = jobs.maybe_w2l_template_token\nLEFT OUTER JOIN w2l_results as results\n ON jobs.on_success_result_token = results.token\n\nWHERE jobs.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 2, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 3, + "name": "maybe_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "maybe_public_bucket_video_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 5, + "name": "maybe_w2l_template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "w2l_template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 7, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 8, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 9, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + true, + false, + false, + false, + false + ] + } + }, + "5ccee09700cfc6d511451c9734d6dc1a8a0cc855ad45c3a73128fc3cf092c0a6": { + "query": "\nUPDATE tts_models\nSET\n maybe_default_pretrained_vocoder = ?,\n title = ?,\n description_markdown = ?,\n description_rendered_html = ?,\n ietf_language_tag = ?,\n ietf_primary_language_subtag = ?,\n creator_set_visibility = ?,\n creator_ip_address_last_update = ?,\n version = version + 1\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "60c06a507349919addb5d26b1c4a283a04f66fce86e84ed0a51db6001c9fc2b5": { + "query": "\nUPDATE w2l_templates\nSET\n title = ?,\n description_markdown = ?,\n description_rendered_html = ?,\n creator_set_visibility = ?,\n maybe_mod_user_token = ?,\n version = version + 1\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 6 + }, + "nullable": [] + } + }, + "619d18c87075a86c01f1e0ac26744774293ff8ca6b0011ce611b3f47d8a97a0e": { + "query": "\nSELECT\n w2l_templates.token as template_token,\n w2l_templates.title,\n w2l_templates.template_type,\n w2l_templates.duration_millis,\n w2l_templates.frame_width,\n w2l_templates.frame_height,\n w2l_templates.creator_user_token,\n users.username AS creator_username,\n users.display_name AS creator_display_name,\n users.email_gravatar_hash AS creator_gravatar_hash,\n w2l_templates.created_at\nFROM\n w2l_templates\nJOIN\n users\nON\n users.token = w2l_templates.creator_user_token\nWHERE\n w2l_templates.is_public_listing_approved IS NULL\n AND w2l_templates.user_deleted_at IS NULL\n AND w2l_templates.mod_deleted_at IS NULL\n AND w2l_templates.is_locked_from_use IS FALSE\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 2, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 3, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 4, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 5, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 6, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 8, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 9, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 10, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + false + ] + } + }, + "641253d4c621f75037877651bba044a21033f50a5e3f8fed9b06ec967a508853": { + "query": "\nUPDATE w2l_template_upload_jobs\nSET\n status = 'started',\n attempt_count = attempt_count + 1,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "68376f750b8421c10d319fa1045a2f3ce8fb1d35a5fe6ec92b2022babafbde80": { + "query": "\nUPDATE users\nSET\n user_role_slug = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "6b66db376952ff4cfcea6ffe1f9ecd4462722843a3d19fcd4df16f0fb993decc": { + "query": "\nUPDATE tts_model_upload_jobs\nSET\n status = 'started',\n attempt_count = attempt_count + 1,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "720d7a26fb8ce75a8a1b3a62ef509d8ca99d419172e9cc5d8e2c803e0eebcb7f": { + "query": "\n UPDATE model_categories\n SET\n deleted_at = NULL\n WHERE\n token = ?\n LIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "73558104936aab6add2d9c789b368b5f155f1fe01d32e7314f2f120422866f54": { + "query": "\nUPDATE w2l_template_upload_jobs\nSET\n status = ?,\n on_success_result_token = ?,\n failure_reason = NULL,\n retry_at = NULL\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "7490dcb8de10b1eda2c857e53ebc5c0465cdb589d3df4c5549859f7919acac77": { + "query": "\nSELECT token, username, email_address, password_hash, is_banned\nFROM users\nWHERE email_address = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "email_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "password_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 63, + "max_size": 60 + } + }, + { + "ordinal": 4, + "name": "is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + }, + "7622292c05fe1b5832136762042780a12a02048166bb1b4492dfd627c48c8968": { + "query": "\nSELECT\n token as template_token\nFROM w2l_templates\nWHERE\n token = ?\n AND user_deleted_at IS NULL\n AND mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false + ] + } + }, + "7ab54466ca3021d90e0feb234acba56bf53d415dde69bc1c35120ed0c8dd1858": { + "query": "\nSELECT *\nFROM tts_model_upload_jobs\nWHERE\n (\n status IN (\"pending\", \"attempt_failed\")\n )\n AND\n (\n retry_at IS NULL\n OR\n retry_at < CURRENT_TIMESTAMP\n )\n LIMIT ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "uuid_idempotency_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 3, + "name": "on_success_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 5, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 6, + "name": "download_url", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 7, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 8, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 9, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 10, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 11, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 12, + "name": "failure_reason", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 13, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 14, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 15, + "name": "retry_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + false, + false, + true, + false, + false, + false, + false, + false, + true, + false, + false, + true + ] + } + }, + "88328e04ce947d3f541951aada9cee25135aa950ab478cf0dc6de453bb079e17": { + "query": "\nSELECT\n events.token as event_token,\n events.event_type,\n events.maybe_target_user_token,\n users.username as maybe_target_username,\n users.display_name as maybe_target_display_name,\n users.email_gravatar_hash as maybe_target_user_gravatar_hash,\n events.maybe_target_entity_token,\n events.created_at,\n events.updated_at\nFROM firehose_entries as events\nLEFT OUTER JOIN users\nON events.maybe_target_user_token = users.token\nORDER BY events.id DESC\nLIMIT 25\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "event_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "event_type", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "maybe_target_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "maybe_target_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "maybe_target_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "maybe_target_user_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "maybe_target_entity_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 8, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + true, + true, + true, + true, + true, + false, + false + ] + } + }, + "886861e91c7e8af53eb856c949389a172b8d9e079f1c2070d115423907c47c50": { + "query": "\nINSERT INTO tts_model_upload_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n title = ?,\n tts_model_type = ?,\n download_url = ?,\n status = \"pending\"\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 8 + }, + "nullable": [] + } + }, + "8b1b9cf11b80da18d79237f2579a96b4a85fa7201118e01fffc01c1eba64cbaf": { + "query": "\n UPDATE model_categories\n SET\n deleted_at = CURRENT_TIMESTAMP\n WHERE\n token = ?\n LIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "9a879f76345ef5682b995f99026fac8b2f4c64c8ec4bb078187d1e53effcfbaf": { + "query": "\nINSERT INTO w2l_templates\nSET\n token = ?,\n template_type = ?,\n title = ?,\n description_markdown = '',\n description_rendered_html = '',\n creator_user_token = ?,\n creator_ip_address_creation = ?,\n creator_ip_address_last_update = ?,\n original_download_url = ?,\n private_bucket_hash = ?,\n private_bucket_object_name = ?,\n private_bucket_cached_faces_object_name = ?,\n maybe_public_bucket_preview_image_object_name = ?,\n maybe_public_bucket_preview_video_object_name = ?,\n file_size_bytes = ?,\n mime_type = ?,\n frame_width = ?,\n frame_height = ?,\n frame_count = ?,\n fps = ?,\n duration_millis = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 19 + }, + "nullable": [] + } + }, + "9b607b814cd7c445c2c46ec5c13dc23f3a6a771a3469a786ed46ec099577d95e": { + "query": "\nUPDATE\n ip_address_bans\nSET\n mod_user_token = ?,\n deleted_at = CURRENT_TIMESTAMP\nWHERE\n ip_address = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "9cc5f697350814199344d3d9a0ff026e0d1e09e9097648400d4b63cfcbd31178": { + "query": "\nUPDATE w2l_results\nSET\n creator_set_visibility = ?,\n creator_ip_address = ?\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "9cef4989671fb55d8e732f6f3d5c16aea7a9b3e01fc6a3bc7a56d0b066f07730": { + "query": "\nINSERT INTO tts_models\nSET\n token = ?,\n tts_model_type = \"tacotron2\",\n title = ?,\n description_markdown = '',\n description_rendered_html = '',\n creator_user_token = ?,\n creator_ip_address_creation = ?,\n creator_ip_address_last_update = ?,\n original_download_url = ?,\n private_bucket_hash = ?,\n private_bucket_object_name = ?,\n file_size_bytes = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "9d4932baa4023accc2d28af98b58a47149d37c4937cd0171e6b955c09855c5f8": { + "query": "\nINSERT INTO w2l_inference_jobs\nSET\n maybe_w2l_template_token = ?,\n maybe_tts_inference_result_token = ?,\n maybe_public_audio_bucket_location = NULL,\n maybe_public_image_bucket_location = NULL,\n maybe_creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n status = \"pending\"\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 5 + }, + "nullable": [] + } + }, + "a125f183f7ae5627fc2118e7302cea73d967200ba1cbb12550f5880187b7db5e": { + "query": "\nUPDATE w2l_inference_jobs\nSET\n status = ?,\n on_success_result_token = ?,\n failure_reason = NULL,\n retry_at = NULL\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "a2b586dd2120f29dd26281626f7381e292c821f8e826b9bc584b92ffa45035b2": { + "query": "\nSELECT\n jobs.token as job_token,\n\n jobs.status,\n jobs.attempt_count,\n jobs.on_success_result_token as maybe_result_token,\n results.public_bucket_wav_audio_path as maybe_public_bucket_wav_audio_path,\n\n jobs.model_token,\n tts.tts_model_type,\n tts.title,\n\n jobs.raw_inference_text,\n\n jobs.created_at,\n jobs.updated_at\n\nFROM tts_inference_jobs as jobs\nJOIN tts_models as tts\n ON tts.token = jobs.model_token\nLEFT OUTER JOIN tts_results as results\n ON jobs.on_success_result_token = results.token\n\nWHERE jobs.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 2, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 3, + "name": "maybe_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "maybe_public_bucket_wav_audio_path", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4224 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 5, + "name": "model_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "tts_model_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 60 + } + }, + { + "ordinal": 7, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 8, + "name": "raw_inference_text", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 9, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 10, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + false, + false, + false, + false, + false, + false + ] + } + }, + "ab981965f518a8986aaaba41ed821fcb438db4c6156cbb26f639921225a88d59": { + "query": "\nUPDATE w2l_template_upload_jobs\nSET\n status = \"dead\",\n failure_reason = ?,\n retry_at = NULL\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + } + }, + "ae04054c4d32e581dd183fd7ebb8e48fc7df856d2ffd8036d9ce6006bb7a296a": { + "query": "\nSELECT\n jobs.token as job_token,\n\n jobs.status,\n jobs.attempt_count,\n jobs.on_success_result_token as maybe_template_token,\n\n jobs.failure_reason as maybe_failure_reason,\n\n jobs.created_at,\n jobs.updated_at\n\nFROM w2l_template_upload_jobs as jobs\n\nWHERE jobs.token = ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "job_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 2, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 3, + "name": "maybe_template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "maybe_failure_reason", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 5, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 6, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + true, + false, + false + ] + } + }, + "b1ad806aa93f66cab358d7cdb405d57c643e2ff131cfec9f1ced55c2e6ca9511": { + "query": "\nUPDATE tts_model_upload_jobs\nSET\n status = ?,\n on_success_result_token = ?,\n failure_reason = NULL,\n retry_at = NULL\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "b257c267dfe063995374786d3004f215e77e8a32843c024645da6fa435507dce": { + "query": "\nINSERT INTO\n ip_address_bans\nSET\n ip_address = ?,\n maybe_target_user_token = ?,\n mod_user_token = ?,\n mod_notes = ?\nON DUPLICATE KEY UPDATE\n expires_at = NULL,\n deleted_at = NULL,\n ip_address = ?,\n maybe_target_user_token = ?,\n mod_user_token = ?,\n mod_notes = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 8 + }, + "nullable": [] + } + }, + "b7535ecf2e63b7f02b5b96f833fd8cd28285b24473df6f89795ff360c1723215": { + "query": "\nSELECT *\nFROM w2l_template_upload_jobs\nWHERE\n (\n status IN (\"pending\", \"attempt_failed\")\n )\n AND\n (\n retry_at IS NULL\n OR\n retry_at < CURRENT_TIMESTAMP\n )\n LIMIT ?\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "uuid_idempotency_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 3, + "name": "on_success_result_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "title", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 5, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 6, + "name": "download_url", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 7, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 8, + "name": "creator_ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 9, + "name": "creator_set_visibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 10, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 11, + "name": "attempt_count", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 3 + } + }, + { + "ordinal": 12, + "name": "failure_reason", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 2048 + } + }, + { + "ordinal": 13, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 14, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 15, + "name": "retry_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + true, + false, + false, + true, + false, + false, + false, + false, + false, + true, + false, + false, + true + ] + } + }, + "bc9c6b57b2578b8163956d11a5f6a9a77ddab6098d2c4207e1a70502db30c9d1": { + "query": "\nSELECT\n w2l.token as template_token,\n w2l.template_type,\n w2l.creator_user_token,\n users.username as creator_username,\n users.display_name as creator_display_name,\n w2l.title,\n w2l.frame_width,\n w2l.frame_height,\n w2l.duration_millis,\n w2l.private_bucket_hash,\n w2l.maybe_public_bucket_preview_image_object_name,\n w2l.maybe_public_bucket_preview_video_object_name,\n w2l.created_at,\n w2l.updated_at\nFROM w2l_templates as w2l\nJOIN users\nON users.token = w2l.creator_user_token\nWHERE w2l.token = ?\nAND w2l.user_deleted_at IS NULL\nAND w2l.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "template_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "template_type", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 2, + "name": "creator_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 4, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "title", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 6, + "name": "frame_width", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 7, + "name": "frame_height", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 8, + "name": "duration_millis", + "type_info": { + "type": "Long", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 10 + } + }, + { + "ordinal": 9, + "name": "private_bucket_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 10, + "name": "maybe_public_bucket_preview_image_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 11, + "name": "maybe_public_bucket_preview_video_object_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 12, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 13, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + false + ] + } + }, + "bee13160674414391587f56e8fb0aa3629e3dfcead0afa5759bff1e161f1bb6e": { + "query": "\nUPDATE w2l_templates\nSET\n is_public_listing_approved = ?,\n maybe_mod_user_token = ?\nWHERE\n token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "bf489cf419d342d144f072c310d6d70c9073b092b5721f489eaf2ffc8ad3915e": { + "query": "\nUPDATE tts_model_upload_jobs\nSET\n status = ?,\n failure_reason = ?,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "cc4372cdd1a65999cbaf008f200683eedc93ab83e17256e47e011667ba8384d8": { + "query": "\nUPDATE tts_results\nSET\n creator_set_visibility = ?,\n creator_ip_address = ?\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 3 + }, + "nullable": [] + } + }, + "ce629cea6c387590759cc89a060c471ecdcb5d5f46a53576897a963154ed4d74": { + "query": "\n INSERT INTO tts_category_assignments\n SET\n model_token = ?,\n category_token = ?,\n category_addition_user_token = ?,\n ip_address_creation = ?,\n ip_address_last_update = ?,\n deleted_at = NULL\n ON DUPLICATE KEY UPDATE\n model_token = ?,\n category_token = ?,\n category_addition_user_token = ?,\n ip_address_last_update = ?,\n deleted_at = NULL\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 9 + }, + "nullable": [] + } + }, + "d12757abef811cd6951e0e1427fdb496a8931b547519d320f5ade56007d3eb5d": { + "query": "\nSELECT\n id,\n status\nFROM w2l_template_upload_jobs\nWHERE id = ?\nFOR UPDATE\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "id", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 515 + }, + "char_set": 63, + "max_size": 20 + } + }, + { + "ordinal": 1, + "name": "status", + "type_info": { + "type": "String", + "flags": { + "bits": 393 + }, + "char_set": 224, + "max_size": 64 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false + ] + } + }, + "d2ef33027abc15b6ec00260f7d30a8c66a3190229ea103ecb571d4957af63525": { + "query": "\nUPDATE w2l_templates\nSET\n title = ?,\n description_markdown = ?,\n description_rendered_html = ?,\n creator_set_visibility = ?,\n creator_ip_address_last_update = ?,\n version = version + 1\nWHERE token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 6 + }, + "nullable": [] + } + }, + "d850784ca167e686b91fdd4d3cce20c7a00e4e38e4da629733114a1c55ba4ab3": { + "query": "\nINSERT INTO w2l_inference_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n\n maybe_w2l_template_token = ?,\n maybe_public_audio_bucket_hash = ?,\n maybe_public_audio_bucket_location = ?,\n\n maybe_original_audio_filename = ?,\n maybe_audio_mime_type = ?,\n\n maybe_creator_user_token = ?,\n creator_ip_address = ?,\n disable_end_bump = false,\n creator_set_visibility = ?,\n status = \"pending\"\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 10 + }, + "nullable": [] + } + }, + "de2d38b505ecff7b8b73996888ff6a904c5a0fab38e1c0b6156929c812b94302": { + "query": "\nSELECT\n users.token as user_token,\n users.username,\n users.display_name,\n user_roles.slug as user_role_slug,\n user_roles.name as user_role_name\nFROM\n users\nJOIN user_roles\n ON users.user_role_slug = user_roles.slug\nWHERE\n user_roles.slug != 'user'\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 3, + "name": "user_role_slug", + "type_info": { + "type": "String", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 4, + "name": "user_role_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 1020 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + }, + "dfe818ab849c5f6e2a5931142aab640d249aa96cb85dc06bfffae4d7feef541e": { + "query": "\nSELECT\n NOW() - t2.created_at AS seconds_since_first,\n (\n SELECT\n count(t1.id) as pending_count\n FROM w2l_inference_jobs AS t1\n WHERE t1.status = \"pending\"\n ) as pending_count\nFROM w2l_inference_jobs AS t2\nWHERE t2.status = \"pending\"\nORDER BY t2.id ASC\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "seconds_since_first", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 16 + } + }, + { + "ordinal": 1, + "name": "pending_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 128 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + true + ] + } + }, + "e5d2e0881e3003256fbf54f494c356187c331c2311a6d8b53f1b4f0e72def332": { + "query": "\nINSERT INTO tts_inference_jobs\nSET\n token = ?,\n uuid_idempotency_token = ?,\n\n model_token = ?,\n raw_inference_text = ?,\n maybe_creator_user_token = ?,\n creator_ip_address = ?,\n creator_set_visibility = ?,\n priority_level = ?,\n status = \"pending\"\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 8 + }, + "nullable": [] + } + }, + "e73192a3e763a287d5737821443c5a81649f17e39978497b99c52e1e896a75f6": { + "query": "\nSELECT\n token as session_token,\n user_token\nFROM user_sessions\nWHERE token = ?\nAND deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "session_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false + ] + } + }, + "eab69a8d679b2dcfb06a9ccc79851f6c693d0c372142ff0e7ded89c95e4bcf44": { + "query": "\nUPDATE users\nSET\n is_banned = ?,\n maybe_mod_comments = ?,\n maybe_mod_user_token = ?,\n version = version + 1\n\nWHERE users.token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 4 + }, + "nullable": [] + } + }, + "ed8d32dd94719ad6479dd638cdc7a8cfbf725785f51678533988ddba94aac112": { + "query": "\nSELECT\n IFNULL((\n SELECT\n NOW() - t1.created_at AS seconds_since_first\n FROM tts_inference_jobs AS t1\n WHERE t1.status = \"pending\"\n ORDER BY t1.id ASC\n LIMIT 1\n ), 0) as seconds_since_first,\n sub2.pending_count,\n sub3.pending_priority_nonzero_count,\n sub4.pending_priority_gt_one_count,\n sub5.attempt_failed_count\nFROM\n (\n SELECT\n count(t2.id) as pending_count\n FROM tts_inference_jobs AS t2\n WHERE t2.status = \"pending\"\n ) as sub2,\n (\n SELECT\n count(t3.id) as pending_priority_nonzero_count\n FROM tts_inference_jobs AS t3\n WHERE t3.status = \"pending\"\n AND t3.priority_level > 0\n ) as sub3,\n (\n SELECT\n count(t4.id) as pending_priority_gt_one_count\n FROM tts_inference_jobs AS t4\n WHERE t4.status = \"pending\"\n AND t4.priority_level > 1\n ) as sub4,\n (\n SELECT\n count(t5.id) as attempt_failed_count\n FROM tts_inference_jobs AS t5\n WHERE t5.status = \"attempt_failed\"\n ) as sub5\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "seconds_since_first", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 129 + }, + "char_set": 63, + "max_size": 16 + } + }, + { + "ordinal": 1, + "name": "pending_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 21 + } + }, + { + "ordinal": 2, + "name": "pending_priority_nonzero_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 21 + } + }, + { + "ordinal": 3, + "name": "pending_priority_gt_one_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 21 + } + }, + { + "ordinal": 4, + "name": "attempt_failed_count", + "type_info": { + "type": "LongLong", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 21 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + }, + "ee05f95e45f7aaa5811d70b1f00d3f40a361536dfc6451affcad97586dc3338f": { + "query": "\nUPDATE w2l_inference_jobs\nSET\n status = 'started',\n attempt_count = attempt_count + 1,\n retry_at = NOW() + interval 2 minute\nWHERE id = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 1 + }, + "nullable": [] + } + }, + "ef828549a221fe6fad3105fede0b38a439a685d6ea7f83a0657df9cdbb80c589": { + "query": "\nSELECT\n ip_bans.ip_address,\n\n ip_bans.maybe_target_user_token,\n banned_users.username as maybe_target_username,\n\n ip_bans.mod_user_token,\n mod_users.username as mod_username,\n mod_users.display_name as mod_display_name,\n\n ip_bans.mod_notes,\n ip_bans.created_at,\n ip_bans.updated_at\nFROM\n ip_address_bans AS ip_bans\nLEFT OUTER JOIN users as banned_users\n ON ip_bans.maybe_target_user_token = banned_users.token\nJOIN users as mod_users\n ON ip_bans.mod_user_token = mod_users.token\nWHERE\n ip_bans.ip_address = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "ip_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 160 + } + }, + { + "ordinal": 1, + "name": "maybe_target_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 136 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "maybe_target_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4228 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 3, + "name": "mod_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "mod_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "mod_display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "mod_notes", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 7, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 8, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 9345 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + true, + true, + false, + false, + false, + false, + false, + false + ] + } + }, + "f70bd775718829cd10f508dfe09286020467cc34fb54daf0758dba2545b164f8": { + "query": "\nSELECT\n users.token as user_token,\n users.username,\n users.display_name,\n\n users.email_address,\n users.email_confirmed,\n users.email_gravatar_hash,\n\n users.disable_gravatar,\n users.auto_play_audio_preference,\n users.auto_play_video_preference,\n users.preferred_tts_result_visibility as `preferred_tts_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`,\n users.preferred_w2l_result_visibility as `preferred_w2l_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`,\n\n users.user_role_slug,\n users.is_banned,\n\n user_roles.can_use_tts,\n user_roles.can_use_w2l,\n user_roles.can_delete_own_tts_results,\n user_roles.can_delete_own_w2l_results,\n user_roles.can_delete_own_account,\n\n user_roles.can_upload_tts_models,\n user_roles.can_upload_w2l_templates,\n user_roles.can_delete_own_tts_models,\n user_roles.can_delete_own_w2l_templates,\n\n user_roles.can_approve_w2l_templates,\n user_roles.can_edit_other_users_profiles,\n user_roles.can_edit_other_users_tts_models,\n user_roles.can_edit_other_users_w2l_templates,\n user_roles.can_delete_other_users_tts_models,\n user_roles.can_delete_other_users_tts_results,\n user_roles.can_delete_other_users_w2l_templates,\n user_roles.can_delete_other_users_w2l_results,\n user_roles.can_ban_users,\n user_roles.can_delete_users\n\nFROM users\nLEFT OUTER JOIN user_sessions\n ON users.token = user_sessions.user_token\nLEFT OUTER JOIN user_roles\n ON users.user_role_slug = user_roles.slug\nWHERE user_sessions.token = ?\n AND user_sessions.deleted_at IS NULL\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 3, + "name": "email_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 4, + "name": "email_confirmed", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 5, + "name": "email_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 6, + "name": "disable_gravatar", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 7, + "name": "auto_play_audio_preference", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 8, + "name": "auto_play_video_preference", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 9, + "name": "preferred_tts_result_visibility: database_queries::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 10, + "name": "preferred_w2l_result_visibility: database_queries::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 11, + "name": "user_role_slug", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 12, + "name": "is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 13, + "name": "can_use_tts", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 14, + "name": "can_use_w2l", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 15, + "name": "can_delete_own_tts_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 16, + "name": "can_delete_own_w2l_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 17, + "name": "can_delete_own_account", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "can_upload_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 19, + "name": "can_upload_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 20, + "name": "can_delete_own_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 21, + "name": "can_delete_own_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 22, + "name": "can_approve_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 23, + "name": "can_edit_other_users_profiles", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 24, + "name": "can_edit_other_users_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 25, + "name": "can_edit_other_users_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 26, + "name": "can_delete_other_users_tts_models", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 27, + "name": "can_delete_other_users_tts_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 28, + "name": "can_delete_other_users_w2l_templates", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 29, + "name": "can_delete_other_users_w2l_results", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 30, + "name": "can_ban_users", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 31, + "name": "can_delete_users", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 0 + }, + "char_set": 63, + "max_size": 1 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true, + true + ] + } + }, + "f70dbc04f76d5ab9ea790482dffe06ebd2a6685b825b522f3bd16c3d14330453": { + "query": "\nSELECT\n users.token as user_token,\n username,\n display_name,\n email_gravatar_hash,\n profile_markdown,\n profile_rendered_html,\n user_role_slug,\n disable_gravatar,\n preferred_tts_result_visibility as `preferred_tts_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`,\n preferred_w2l_result_visibility as `preferred_w2l_result_visibility: database_queries::column_types::record_visibility::RecordVisibility`,\n discord_username,\n twitch_username,\n twitter_username,\n patreon_username,\n github_username,\n cashapp_username,\n website_url,\n is_banned,\n maybe_mod_comments,\n maybe_mod_user_token,\n created_at\nFROM users\nWHERE\n users.username = ?\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "display_name", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 3, + "name": "email_gravatar_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "profile_markdown", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 5, + "name": "profile_rendered_html", + "type_info": { + "type": "Blob", + "flags": { + "bits": 4241 + }, + "char_set": 224, + "max_size": 262140 + } + }, + { + "ordinal": 6, + "name": "user_role_slug", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4233 + }, + "char_set": 224, + "max_size": 64 + } + }, + { + "ordinal": 7, + "name": "disable_gravatar", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 8, + "name": "preferred_tts_result_visibility: database_queries::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 9, + "name": "preferred_w2l_result_visibility: database_queries::column_types::record_visibility::RecordVisibility", + "type_info": { + "type": "String", + "flags": { + "bits": 385 + }, + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 10, + "name": "discord_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 11, + "name": "twitch_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 12, + "name": "twitter_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 13, + "name": "patreon_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 14, + "name": "github_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 15, + "name": "cashapp_username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 144 + } + }, + { + "ordinal": 16, + "name": "website_url", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 17, + "name": "is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + }, + { + "ordinal": 18, + "name": "maybe_mod_comments", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 19, + "name": "maybe_mod_user_token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 128 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 20, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": { + "bits": 1153 + }, + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + false, + false, + false, + true, + true, + true, + true, + true, + true, + true, + false, + true, + true, + false + ] + } + }, + "f75b713d166398ebea2bdbde4e416dc8eb2c48dba4aa9bd164bcf5bc77237ce1": { + "query": "\nUPDATE users\nSET\n profile_markdown = ?,\n profile_rendered_html = ?,\n preferred_tts_result_visibility = ?,\n preferred_w2l_result_visibility = ?,\n discord_username = ?,\n twitter_username = ?,\n twitch_username = ?,\n github_username = ?,\n cashapp_username = ?,\n website_url = ?,\n ip_address_last_update = ?,\n version = version + 1\n\nWHERE users.token = ?\nLIMIT 1\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 12 + }, + "nullable": [] + } + }, + "fd31f7abe386147c07f862d668c9e56eb3d809c4f0a61b7debf9b2e3785a4c83": { + "query": "\nINSERT INTO model_categories\nSET\n token = ?,\n uuid_idempotency_token = ?,\n model_type = ?,\n name = ?,\n\n creator_user_token = ?,\n creator_ip_address_creation = ?,\n creator_ip_address_last_update = ?,\n\n is_mod_approved = ?,\n maybe_mod_user_token = ?,\n can_directly_have_models = ?,\n can_have_subcategories = ?,\n can_only_mods_apply = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 12 + }, + "nullable": [] + } + }, + "feb0dae816a5701707a76e15011ea6cb262b10c663abf2cf1fb46bd9b1cebe09": { + "query": "\nSELECT token, username, email_address, password_hash, is_banned\nFROM users\nWHERE username = ?\nLIMIT 1\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "username", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 2, + "name": "email_address", + "type_info": { + "type": "VarString", + "flags": { + "bits": 4229 + }, + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "password_hash", + "type_info": { + "type": "String", + "flags": { + "bits": 4225 + }, + "char_set": 63, + "max_size": 60 + } + }, + { + "ordinal": 4, + "name": "is_banned", + "type_info": { + "type": "Tiny", + "flags": { + "bits": 1 + }, + "char_set": 63, + "max_size": 1 + } + } + ], + "parameters": { + "Right": 1 + }, + "nullable": [ + false, + false, + false, + false, + false + ] + } + } +} diff --git a/static/test.js b/static/test.js new file mode 100644 index 0000000000..c477d0f0a6 --- /dev/null +++ b/static/test.js @@ -0,0 +1,13 @@ +//var sock = new WebSocket("ws://api.jungle.horse:80/twitch"); +var sock = new WebSocket("ws://localhost:12345/twitch"); + +sock.onopen = function (event) { + console.log('on open event', event); + sock.send('on open message from client'); +}; + +sock.onmessage = function (event) { + console.log('on message event', event.data); +} + +sock.send('this is a message from the client'); diff --git a/static/websocket_test.html b/static/websocket_test.html new file mode 100644 index 0000000000..5af8d02f58 --- /dev/null +++ b/static/websocket_test.html @@ -0,0 +1,12 @@ + + + + Websocket Test + + +

Websocket Test

+

+

Use the URL hash to set the username, eg. #testytest

+ + + diff --git a/static/websocket_test.js b/static/websocket_test.js new file mode 100644 index 0000000000..674881d3c6 --- /dev/null +++ b/static/websocket_test.js @@ -0,0 +1,45 @@ + +const WS_URL = 'ws://localhost:54321/obs'; + +const DEFAULT_USERNAME = 'testytest512'; + +(function() { + console.log('installing script'); + + const maybeUsername = window.location.hash.replace('#', '').trim(); + + const username = (!!maybeUsername) ? maybeUsername : DEFAULT_USERNAME; + + console.log('username', username); + document.getElementById('username').innerHTML = username; + + const url = `${WS_URL}/${username}`; + + let sock = new WebSocket(url); + + sock.onopen = function (event) { + console.log('on open event', event); + sock.send('on open message from browser'); + }; + + sock.onmessage = function (event) { + console.log('on message event', event.data); + } + + sock.onerror = function(event) { + console.log('on error event', event.data); + } + + setInterval(() => { + console.log('sending browser message to server on interval'); + sock.send('interval trigger from browser'); + }, 2000); + + //setTimeout(() => { + // console.log('disconnect') + // const normal_close = 1000; + // sock.close(normal_close, 'close'); + // + //}, 10000) + +})(); \ No newline at end of file