Example of BlindAI deployment with Wav2vec2¶

This example shows how you can run a Wav2Vec2 model to perform Speech-To-Text with confidentiality guarantees.

By using BlindAI, people can send data for the AI to analyze their conversations without having to fear privacy leaks.

Wav2Vec2 is a state-of-the art Transformers model for speech. You can learn more about it on FAIR blog's post.

Installing dependencies¶

Install the dependencies this example needs.

In [ ]:

!pip install -q transformers[onnx] torch

!pip install -q transformers[onnx] torch

We will need librosa to load the "hello world" audio file. You might need to downgrade numpy to 1.21 to make it work. The following commands should do the trick to install librosa:

In [ ]:

!pip install -q --upgrade numpy==1.21
!pip install -q librosa

!pip install -q --upgrade numpy==1.21 !pip install -q librosa

In addition, you might need to install ffmpeg to have a backend to process the wav file.

In [ ]:

!sudo apt-get install -y ffmpeg

!sudo apt-get install -y ffmpeg

Install the latest version of BlindAI.

In [ ]:

!pip install blindai

!pip install blindai

Preparing the model¶

Here we will use a large Wav2Vec2 model. First step is to get the model and tokenizers.

In [ ]:

from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
import torch

# load model and processor
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")
model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")

from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC import torch # load model and processor processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h") model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h")

We can download an hello world audio file to be used as example. Let's download it.

In [ ]:

!wget https://github.com/mithril-security/blindai/raw/master/examples/wav2vec2/hello_world.wav

!wget https://github.com/mithril-security/blindai/raw/master/examples/wav2vec2/hello_world.wav

We can hear it here:

In [ ]:

import IPython.display as ipd
ipd.Audio('hello_world.wav')

import IPython.display as ipd ipd.Audio('hello_world.wav')

We can then see the Wav2vec2 model in action on the hello world file.

In [ ]:

import librosa

audio, rate = librosa.load("hello_world.wav", sr = 16000)

# Tokenize sampled audio to input into model
input_values = processor(audio, sampling_rate=rate, return_tensors="pt", padding="longest").input_values

# Retrieve logits
logits = model(input_values).logits

# Take argmax and decode
predicted_ids = torch.argmax(logits, dim=-1)
transcription = processor.batch_decode(predicted_ids)
print(transcription)

import librosa audio, rate = librosa.load("hello_world.wav", sr = 16000) # Tokenize sampled audio to input into model input_values = processor(audio, sampling_rate=rate, return_tensors="pt", padding="longest").input_values # Retrieve logits logits = model(input_values).logits # Take argmax and decode predicted_ids = torch.argmax(logits, dim=-1) transcription = processor.batch_decode(predicted_ids) print(transcription)

In order to facilitate the deployment, we will add the post processing directly to the full model. This way the client will not have to do the post processing.

In [ ]:

import torch.nn as nn

# Let's embed the post-processing phase with argmax inside our model
class ArgmaxLayer(nn.Module):
    def __init__(self):
        super(ArgmaxLayer, self).__init__()

    def forward(self, outputs):
        return torch.argmax(outputs.logits, dim = -1)

import torch.nn as nn # Let's embed the post-processing phase with argmax inside our model class ArgmaxLayer(nn.Module): def __init__(self): super(ArgmaxLayer, self).__init__() def forward(self, outputs): return torch.argmax(outputs.logits, dim = -1)

In [ ]:

final_layer = ArgmaxLayer()

# Finally we concatenate everything
full_model = nn.Sequential(model, final_layer)

final_layer = ArgmaxLayer() # Finally we concatenate everything full_model = nn.Sequential(model, final_layer)

We can check the results are the same.

In [ ]:

predicted_ids = full_model(input_values)
transcription = processor.batch_decode(predicted_ids)
transcription

predicted_ids = full_model(input_values) transcription = processor.batch_decode(predicted_ids) transcription

Now we can export the model in ONNX format, so that we can feed later the ONNX to our BlindAI server.

In [ ]:

torch.onnx.export(
    full_model,
    input_values,
    'wav2vec2_hello_world.onnx',
    export_params=True,
     opset_version = 11)

torch.onnx.export( full_model, input_values, 'wav2vec2_hello_world.onnx', export_params=True, opset_version = 11)

Deployment on BlindAI¶

Now we can upload the model to BlindAI Cloud. To upload of the model, make sure you have an API key.

You can get one on the Mithril Cloud.

You might get an error if the name you want to use is already taken, as models are uniquely identified by their model_id. We will implement namespace soon to avoid that. Meanwhile, you will have to choose a unique ID. We provide an example below to upload your model with a unique name:

In [ ]:

import blindai
import uuid

api_key = "YOUR_API_KEY" # Enter your API key here
model_id = "wav2vec2-" + str(uuid.uuid4())

# Upload the ONNX file to the remote enclave
with blindai.connect(api_key=api_key) as client:
    response = client.upload_model("wav2vec2_hello_world.onnx", model_id=model_id)

import blindai import uuid api_key = "YOUR_API_KEY" # Enter your API key here model_id = "wav2vec2-" + str(uuid.uuid4()) # Upload the ONNX file to the remote enclave with blindai.connect(api_key=api_key) as client: response = client.upload_model("wav2vec2_hello_world.onnx", model_id=model_id)

Sending data for confidential prediction¶

Now it's time to check it's working live!

We will just prepare some input for the model inside the secure enclave of BlindAI to process it.

First we prepare our input data, the hello world audio file.

In [ ]:

from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
import torch
import librosa

# load model and processor
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")

audio, rate = librosa.load("hello_world.wav", sr = 16000)

# Tokenize sampled audio to input into model
input_values = processor(audio, sampling_rate=rate, return_tensors="pt", padding="longest").input_values

from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC import torch import librosa # load model and processor processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h") audio, rate = librosa.load("hello_world.wav", sr = 16000) # Tokenize sampled audio to input into model input_values = processor(audio, sampling_rate=rate, return_tensors="pt", padding="longest").input_values

Now we can send the audio data to be processed confidentially!

In [ ]:

with blindai.connect() as client:
  response = client.predict(model_id, input_values)

with blindai.connect() as client: response = client.predict(model_id, input_values)

We can reconstruct the output now:

In [ ]:

# Decode the output
processor.batch_decode(response.output[0].as_torch().unsqueeze(0))

# Decode the output processor.batch_decode(response.output[0].as_torch().unsqueeze(0))

Et voila! We have been able to apply a start of the art model of speech recognition, without ever having to show the data in clear to the people operating the service!

If you have liked this example, do not hesitate to drop a star on our GitHub and chat with us on our Discord!