%pip install azure-ai-ml azure-identity --upgrade --quiet

%env LOCATION eastus
%env SUBSCRIPTION_ID <YOUR_SUBSCRIPTION_ID>
%env RESOURCE_GROUP <YOUR_RESOURCE_GROUP>
%env AI_FOUNDRY_HUB_PROJECT <YOUR_AI_FOUNDRY_HUB_PROJECT>

import os
from uuid import uuid4

os.environ["ENDPOINT_NAME"] = f"nvidia-parakeet-{str(uuid4())[:8]}"
os.environ["DEPLOYMENT_NAME"] = f"nvidia-parakeet-{str(uuid4())[:8]}"

import os
from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential

client = MLClient(
    credential=DefaultAzureCredential(),
    subscription_id=os.getenv("SUBSCRIPTION_ID"),
    resource_group_name=os.getenv("RESOURCE_GROUP"),
    workspace_name=os.getenv("AI_FOUNDRY_HUB_PROJECT"),
)

model_id = "nvidia/parakeet-tdt-0.6b-v2"

model_uri = (
    f"azureml://registries/HuggingFace/models/{model_id.replace('/', '-').replace('_', '-').lower()}/labels/latest"
)
model_uri

from azure.ai.ml.entities import ManagedOnlineEndpoint, ManagedOnlineDeployment

endpoint = ManagedOnlineEndpoint(name=os.getenv("ENDPOINT_NAME"))

deployment = ManagedOnlineDeployment(
    name=os.getenv("DEPLOYMENT_NAME"),
    endpoint_name=os.getenv("ENDPOINT_NAME"),
    model=model_uri,
    instance_type="Standard_NC40ads_H100_v5",
    instance_count=1,
)

client.begin_create_or_update(endpoint).wait()

client.online_deployments.begin_create_or_update(deployment).wait()

api_key = client.online_endpoints.get_keys(os.getenv("ENDPOINT_NAME")).primary_key
api_url = client.online_endpoints.get(os.getenv("ENDPOINT_NAME")).scoring_uri

!wget https://dldata-public.s3.us-east-2.amazonaws.com/2086-149220-0033.wav

from io import BytesIO
import requests

audio_file = BytesIO(open("2086-149220-0033.wav", "rb").read())
audio_file.name = "2086-149220-0033.wav"

response = requests.post(
    api_url,
    headers={
        "Authorization": f"Bearer {api_key}",
        "azureml-model-deployment": os.getenv("DEPLOYMENT_NAME"),
    },
    files={"file": (audio_file.name, audio_file, "audio/wav")},
    data={"model": model_id},
)
print(response.json())
# {'text': "Well, I don't wish to see it any more, observed Phebe, turning away her eyes. It is certainly very like the old portrait."}

%pip install openai --upgrade --quiet

api_url = client.online_endpoints.get(os.getenv("ENDPOINT_NAME")).scoring_uri.replace("/audio/transcriptions", "")

import os
from openai import OpenAI

openai_client = OpenAI(
    base_url=api_url,
    api_key=api_key,
    default_headers={"azureml-model-deployment": os.getenv("DEPLOYMENT_NAME")},
)

transcription = openai_client.audio.transcriptions.create(
    model=model_id,
    file=open("2086-149220-0033.wav", "rb"),
    response_format="json",
)
print(transcription.text)
# Well, I don't wish to see it any more, observed Phebe, turning away her eyes. It is certainly very like the old portrait.

os.environ["API_URL"] = api_url
os.environ["API_KEY"] = api_key

!curl -sS $API_URL/audio/transcriptions \
    -H "Authorization: Bearer $API_KEY" \
    -H "azureml-model-deployment: $DEPLOYMENT_NAME" \
    -H "Content-Type: multipart/form-data" \
    -F file=@2086-149220-0033.wav \
    -F model=nvidia/parakeet-tdt-0.6b-v2

%pip install gradio --upgrade --quiet

import os
from pathlib import Path

import gradio as gr
from openai import OpenAI

openai_client = OpenAI(
    base_url=os.getenv("API_URL"),
    api_key=os.getenv("API_KEY"),
    default_headers={"azureml-model-deployment": os.getenv("DEPLOYMENT_NAME")},
)


def transcribe(audio: Path, temperature: float = 1.0) -> str:
    return openai_client.audio.transcriptions.create(
        model=model_id,
        file=open(audio, "rb"),
        temperature=temperature,
        response_format="text",
    )


demo = gr.Interface(
    fn=transcribe,
    inputs=[
        # https://gradio.org.cn/docs/gradio/audio
        gr.Audio(type="filepath", streaming=False, label="Upload or Record Audio"),
        gr.Slider(0, 1, value=0.0, step=0.1, label="Temperature"),
    ],
    outputs=gr.Textbox(label="Transcribed Text"),
    title="NVIDIA Parakeet on Azure AI",
    description="Upload or record audio and get the transcribed text using NVIDIA Parakeet on Azure AI via the OpenAI's Transcription API.",
)

demo.launch()

client.online_endpoints.begin_delete(name=os.getenv("ENDPOINT_NAME")).result()