mirror of
https://github.com/storytold/storyteller-ml.git
synced 2026-10-23 08:26:06 +00:00
Introduce code from Rust monorepo. This is a mess
This commit is contained in:
Brandon Thomas
@@ -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"]
|
||||
@@ -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 <https://gunicorn.org/>`_, 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
|
||||
|
||||
<!-- Load .nemo checkpoint -->
|
||||
<!-- Comment this section out to remove ability to upload files -->
|
||||
<br>
|
||||
NeMo File :
|
||||
<input type = "file" name = "nemo_model" accept=".nemo" id="nemo_model_file" />
|
||||
<!-- Comment out upto here to remove ability to upload models -->
|
||||
@@ -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 = """
|
||||
<script>
|
||||
document.getElementById('nemo_model_file').value = ""
|
||||
</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)
|
||||
@@ -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 .
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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 !")
|
||||
@@ -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 = "<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
|
||||
@@ -0,0 +1,4 @@
|
||||
-e git+git://github.com/NVIDIA/NeMo.git@v1.0.2#egg=nemo_toolkit[all]
|
||||
flask
|
||||
flask-cors
|
||||
gunicorn
|
||||
@@ -0,0 +1,150 @@
|
||||
<!--
|
||||
# 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.
|
||||
-->
|
||||
<!doctype html>
|
||||
<html lang="en">
|
||||
<head>
|
||||
<link rel="stylesheet" href="http://fonts.googleapis.com/css?family=Roboto:300,400,500,700" type="text/css">
|
||||
<meta charset="utf-8">
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1">
|
||||
<title>Automatic Speech Recognition</title>
|
||||
<link rel="stylesheet" href="//code.jquery.com/ui/1.12.1/themes/base/jquery-ui.css">
|
||||
|
||||
<link rel="stylesheet" href="https://fonts.googleapis.com/icon?family=Material+Icons">
|
||||
<link rel="stylesheet" href="https://code.getmdl.io/1.3.0/material.light_green-lime.min.css">
|
||||
<script defer src="https://code.getmdl.io/1.3.0/material.min.js"></script>
|
||||
<style>
|
||||
.ui-resizable-se {
|
||||
bottom: 17px;
|
||||
}
|
||||
|
||||
</style>
|
||||
<script src="https://code.jquery.com/jquery-1.12.4.js"></script>
|
||||
<script src="https://code.jquery.com/ui/1.12.1/jquery-ui.js"></script>
|
||||
|
||||
<!-- Load from https://htmx.org/ -->
|
||||
<script src="https://unpkg.com/htmx.org@1.3.1"></script>
|
||||
</head>
|
||||
<body>
|
||||
|
||||
<!-- Always shows a header, even in smaller screens. -->
|
||||
<div class="mdl-layout mdl-js-layout mdl-layout--fixed-header">
|
||||
<header class="mdl-layout__header">
|
||||
<div class="mdl-layout__header-row">
|
||||
<!-- Title -->
|
||||
<span class="mdl-layout-title">NeMo Automatic Speech Recognition Demo</span>
|
||||
<!-- Add spacer, to align navigation to the right -->
|
||||
<div class="mdl-layout-spacer"></div>
|
||||
<!-- Navigation. We hide it in small screens. -->
|
||||
<nav class="mdl-navigation mdl-layout--large-screen-only">
|
||||
<a class="mdl-navigation__link" href="#top">Top</a>
|
||||
</nav>
|
||||
</div>
|
||||
</header>
|
||||
<div class="mdl-layout__drawer">
|
||||
<span class="mdl-layout-title">ASR Transcription</span>
|
||||
<nav class="mdl-navigation">
|
||||
<a class="mdl-navigation__link" href="https://github.com/NVIDIA/NeMo">Github</a>
|
||||
</nav>
|
||||
</div>
|
||||
<main class="mdl-layout__content">
|
||||
<div class="page-content">
|
||||
|
||||
<!--Placeholder for top of app-->
|
||||
<div id="top"></div>
|
||||
|
||||
<form id="model_names_form" action="" enctype="multipart/form-data" style="display: inline;">
|
||||
|
||||
<h5> Model Name :
|
||||
|
||||
<!-- Format selectable list of models -->
|
||||
<div class="mdl-textfield mdl-js-textfield mdl-textfield--floating-label" style="width: 25%">
|
||||
<select name="model_names_select" class="mdl-textfield__input" id="model_names_select_id">
|
||||
{% for x in model_names %}
|
||||
<option value="{{ x }}" {% if loop.first %} SELECTED{% endif %}>{{ x }}</option>
|
||||
{% endfor %}
|
||||
</select>
|
||||
</div>
|
||||
|
||||
<!-- Load Model -->
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url_for('initialize_model') }}" hx-target="this" hx-swap="afterend"
|
||||
hx-encoding="multipart/form-data">
|
||||
Load Model
|
||||
</button>
|
||||
|
||||
<!-- Load .nemo checkpoint -->
|
||||
<!-- Comment this section out to remove ability to upload files -->
|
||||
|
||||
<br>
|
||||
NeMo File :
|
||||
<input type="file" name="nemo_model" accept=".nemo" id="nemo_model_file"/>
|
||||
|
||||
<!-- Comment out upto here to remove ability to upload models -->
|
||||
|
||||
</h5>
|
||||
|
||||
<!-- GPU checkbox -->
|
||||
<label class="mdl-checkbox mdl-js-checkbox mdl-js-ripple-effect" for="use_gpu_ckbx_id">
|
||||
<input type="checkbox" id="use_gpu_ckbx_id" name="use_gpu_ckbx" class="mdl-checkbox__input" checked>
|
||||
<span class="mdl-checkbox__label">Use GPU (if available)</span>
|
||||
</label>
|
||||
|
||||
</form>
|
||||
|
||||
<hr>
|
||||
<br>
|
||||
|
||||
<!-- Form input for files -->
|
||||
<form action="" enctype="multipart/form-data">
|
||||
<input type="file" name="file" accept=".wav" multiple=""/>
|
||||
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url_for('upload_audio_files') }}"
|
||||
hx-target="this" hx-swap="outerHTML" hx-encoding="multipart/form-data">
|
||||
Upload audio file(s)
|
||||
</button>
|
||||
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url_for('remove_audio_files') }}"
|
||||
hx-target="this" hx-swap="outerHTML">
|
||||
Remove all files
|
||||
</button>
|
||||
</form>
|
||||
<hr>
|
||||
<br>
|
||||
|
||||
<!-- Transcribe Button -->
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url_for('transcribe') }}" hx-target="#transcribeResults">
|
||||
Transcribe all
|
||||
</button>
|
||||
|
||||
<div id="transcribeResults"></div>
|
||||
|
||||
|
||||
<!--Material Toast -->
|
||||
|
||||
<div id="toast" class="mdl-js-snackbar mdl-snackbar">
|
||||
<div class="mdl-snackbar__text"></div>
|
||||
<button class="mdl-snackbar__action" type="button"></button>
|
||||
</div>
|
||||
|
||||
</div>
|
||||
</main>
|
||||
</div>
|
||||
|
||||
</body>
|
||||
</html>
|
||||
@@ -0,0 +1,28 @@
|
||||
<!--
|
||||
# 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.
|
||||
-->
|
||||
|
||||
<script>
|
||||
(function() {
|
||||
'use strict';
|
||||
var snackbarContainer = document.querySelector('#toast');
|
||||
var data = {
|
||||
message: '{{ toast_message }}',
|
||||
timeout: {{ timeout }},
|
||||
};
|
||||
snackbarContainer.MaterialSnackbar.showSnackbar(data);
|
||||
}());
|
||||
|
||||
</script>
|
||||
@@ -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.
|
||||
-->
|
||||
|
||||
<table class="transcript-view-table mdl-data-table mdl-js-data-table mdl-shadow--4dp">
|
||||
<thead>
|
||||
<tr>
|
||||
<th class="mdl-data-table__cell--non-numeric">Filename</th>
|
||||
<th class="mdl-data-table__cell--non-numeric">Transcript</th>
|
||||
</tr>
|
||||
</thead>
|
||||
<tbody>
|
||||
{% for transcript in transcripts %}
|
||||
<tr>
|
||||
<td class="property-value-col mdl-data-table__cell--non-numeric" style="width: 10%">
|
||||
<div class="col-value" style="overflow-x:auto;">
|
||||
{{ transcript.filename }}
|
||||
</div>
|
||||
</td>
|
||||
<td class="property-value-col mdl-data-table__cell--non-numeric">
|
||||
<div class="col-value mdl-textfield mdl-js-textfield"
|
||||
style="width: 100%; display: block;">
|
||||
<textarea class="mdl-textfield__input" type="text" rows="2" readonly>
|
||||
{{ transcript.transcription }}
|
||||
</textarea>
|
||||
</div>
|
||||
</td>
|
||||
</tr>
|
||||
{% endfor %}
|
||||
</tbody>
|
||||
</table>
|
||||
@@ -0,0 +1,22 @@
|
||||
<!--
|
||||
# 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.
|
||||
-->
|
||||
|
||||
{{ pre_exec }}
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url }}"
|
||||
hx-target="this" hx-swap="outerHTML">
|
||||
Remove all files
|
||||
</button>
|
||||
@@ -0,0 +1,22 @@
|
||||
<!--
|
||||
# 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.
|
||||
-->
|
||||
|
||||
{{ pre_exec }}
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url }}"
|
||||
hx-target="this" hx-swap="outerHTML" hx-encoding="multipart/form-data">
|
||||
Upload audio file(s)
|
||||
</button>
|
||||
@@ -0,0 +1,22 @@
|
||||
<!--
|
||||
# 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.
|
||||
-->
|
||||
|
||||
{{ pre_exec }}
|
||||
<button class="btn mdl-button mdl-js-button mdl-button--raised mdl-button--colored mdl-js-ripple-effect"
|
||||
hx-post="{{ url }}"
|
||||
hx-target="this" hx-swap="outerHTML" hx-encoding="multipart/form-data">
|
||||
{{ msg }}
|
||||
</button>
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user