update chrome plugin with model size dropdown

Audio-Transcription-Chrome/background.js

@@ -156,7 +156,8 @@ async function startCapture(options) {
           port: options.port, 
           multilingual: options.useMultilingual,
           language: options.language,
-          task: options.task
+          task: options.task,
+          modelSize: options.modelSize
         },
       });
     } else {

Audio-Transcription-Chrome/options.js

@@ -102,7 +102,8 @@ async function startRecord(option) {
           uid: uuid,
           multilingual: option.multilingual,
           language: option.language,
-          task: option.task
+          task: option.task,
+          model_size: option.modelSize
         })
       );
     };

Audio-Transcription-Chrome/popup.html

@@ -125,11 +125,22 @@
   </div>
   <div class="dropdown-container">
     <label for="taskDropdown">Select task:</label>
-    <select id="taskDropdown" disabled>
+    <select id="taskDropdown" >
       <option value="">Select Task</option>
       <option value="transcribe" selected>Transcribe</option>
       <option value="translate">Translate</option>
     </select>
   </div>
+  <div class="dropdown-container">
+    <label for="modelSizeDropdown">Select Model Size:</label>
+    <select id="modelSizeDropdown">
+      <option value="">Select Task</option>
+      <option value="tiny">Tiny</option>
+      <option value="base">Base</option>
+      <option value="small" selected>Small</option>
+      <option value="medium">Medium</option>
+      <option value="large-v2">Large-v2</option>
+    </select>
+  </div>
 </body>
 </html>

Audio-Transcription-Chrome/popup.js

@@ -7,8 +7,10 @@ document.addEventListener("DOMContentLoaded", function () {
   const useMultilingualCheckbox = document.getElementById('useMultilingualCheckbox');
   const languageDropdown = document.getElementById('languageDropdown');
   const taskDropdown = document.getElementById('taskDropdown');
+  const modelSizeDropdown = document.getElementById('modelSizeDropdown');
   let selectedLanguage = null;
   let selectedTask = taskDropdown.value;
+  let selectedModelSize = modelSizeDropdown.value;
 
   // Add click event listeners to the buttons
   startButton.addEventListener("click", startCapture);
@@ -52,6 +54,13 @@ document.addEventListener("DOMContentLoaded", function () {
     }
   });
 
+  chrome.storage.local.get("selectedModelSize", ({ selectedModelSize: storedModelSize }) => {
+    if (storedModelSize !== undefined) {
+      modelSizeDropdown.value = storedModelSize;
+      selectedModelSize = storedModelSize;
+    }
+  });
+
   // Function to handle the start capture button click event
   async function startCapture() {
     // Ignore click if the button is disabled
@@ -64,7 +73,7 @@ document.addEventListener("DOMContentLoaded", function () {
 
     // Send a message to the background script to start capturing
     let host = "localhost";
-    let port = "9090";
+    let port = "5901";
     const useCollaboraServer = useServerCheckbox.checked;
     if (useCollaboraServer){
       host = "transcription.kurg.org"
@@ -79,7 +88,8 @@ document.addEventListener("DOMContentLoaded", function () {
         port: port,
         useMultilingual: useMultilingualCheckbox.checked,
         language: selectedLanguage,
-        task: selectedTask
+        task: selectedTask,
+        modelSize: selectedModelSize
       }, () => {
         // Update capturing state in storage and toggle the buttons
         chrome.storage.local.set({ capturingState: { isCapturing: true } }, () => {
@@ -120,6 +130,7 @@ document.addEventListener("DOMContentLoaded", function () {
     stopButton.disabled = !isCapturing;
     useServerCheckbox.disabled = isCapturing; 
     useMultilingualCheckbox.disabled = isCapturing;
+    modelSizeDropdown.disabled = isCapturing;
     
     startButton.classList.toggle("disabled", isCapturing);
     stopButton.classList.toggle("disabled", !isCapturing);
@@ -157,6 +168,11 @@ document.addEventListener("DOMContentLoaded", function () {
     chrome.storage.local.set({ selectedTask });
   });
 
+  modelSizeDropdown.addEventListener('change', function() {
+    selectedModelSize = modelSizeDropdown.value;
+    chrome.storage.local.set({ selectedModelSize });
+  });
+
   chrome.runtime.onMessage.addListener(async (request, sender, sendResponse) => {
     if (request.action === "updateSelectedLanguage") {
       const detectedLanguage = request.detectedLanguage;

Audio-Transcription-Chrome/test.py

@@ -0,0 +1,243 @@
+import configparser
+
+import numpy as np
+import torch
+
+from tensorrt_llm._utils import str_dtype_to_torch, torch_to_numpy
+from tensorrt_llm.functional import LayerNormPositionType, LayerNormType
+
+layernorm_type_map = {i.name: i.value for i in LayerNormType}
+layernorm_position_map = {i.name: i.value for i in LayerNormPositionType}
+
+def parse_config(ini_file, component, args):
+    config = configparser.ConfigParser()
+    config.read(ini_file)
+    if component == 'encoder':
+        args.n_layer = config.getint(component, 'n_layer')
+        args.n_head = config.getint(component, 'n_head')
+        args.hidden_size = config.getint(component, 'n_state')
+        args.n_ctx = config.getint(component, 'n_ctx')
+        args.ffn_hidden_size = config.getint(component, 'ffn_hidden_size')
+        args.vocab_size = config.getint(component, 'vocab_size')
+        args.n_positions = config.getint(component, 'n_positions')
+        args.has_position_embedding = config.getboolean(
+            component, 'has_position_embedding', fallback=False)
+        args.has_token_type_embedding = config.getboolean(
+            component, 'has_token_type_embedding', fallback=False)
+        args.has_embedding_layernorm = config.getboolean(
+            component, 'has_embedding_layernorm', fallback=False)
+        args.has_embedding_scale = config.getboolean(component,
+                                                     'has_embedding_scale',
+                                                     fallback=False)
+        args.q_scaling = config.getfloat(component, 'q_scaling', fallback=1.0)
+        args.has_attention_qkvo_bias = config.getboolean(
+            component, 'has_attention_qkvo_bias', fallback=False)
+        args.has_mlp_bias = config.getboolean(component,
+                                              'has_mlp_bias',
+                                              fallback=False)
+        args.has_model_final_layernorm = config.getboolean(
+            component, 'has_model_final_layernorm', fallback=False)
+        args.layernorm_eps = config.getfloat(component,
+                                             'layernorm_eps',
+                                             fallback=1e-5)
+        args.layernorm_position = layernorm_position_map[config.get(
+            component, 'layernorm_position')]
+        args.layernorm_type = layernorm_type_map[config.get(
+            component, 'layernorm_type')]
+        args.hidden_act = config.get(component, 'hidden_act')
+        args.relative_attention = config.getboolean(component,
+                                                    'relative_attention',
+                                                    fallback=False)
+    return args
+
+def fuse_qkv(q, k, v):
+    qkv_weight = np.concatenate((q, k, v))
+    return qkv_weight
+
+def load_whisper_from_pytorch(tllm_model,
+                         pytorch_ckpt_path,
+                         component,
+                         model_size,
+                         multilingual=False,
+                         dtype="float32"):
+    torch_dtype = str_dtype_to_torch(dtype)
+    model_name = "whisper-" + f"{model_size}"
+    if not multilingual:
+        model_name = model_name + ".en"
+
+    pytorch_ckpt = torch.load(pytorch_ckpt_path + f"{model_name}.ckpt")
+    pytorch_model = {
+        key: torch_to_numpy(value.to(torch_dtype))
+        for key, value in pytorch_ckpt.items()
+    }
+
+    if component == "encoder":
+        # set conv1d
+        tllm_model.conv1.weight.value = pytorch_model['encoder.conv1.weight'].unsqueeze(3)
+        tllm_model.conv1.bias.value = pytorch_model['encoder.conv1.bias']
+        tllm_model.conv2.weight.value = pytorch_model['encoder.conv2.weight'].unsqueeze(3)
+        tllm_model.conv2.bias.value = pytorch_model['encoder.conv2.bias']
+        # if tllm_model.embedding.position_embedding:
+        #     tllm_model.embedding.position_embedding.weight.value = pytorch_model[
+        #         'encoder.embed_positions.weight']
+
+        for i in range(tllm_model.num_layers):
+            layer = tllm_model.encoder_layers[i]
+            layer_prefix = f'encoder.layers.{i}.'
+
+            # attention table for all layers
+            layer.attention.rel_attn_table.value = relative_attention_table
+
+            layer.attention.qkv.weight.value = fuse_qkv(
+                pytorch_model[f'{layer_prefix}self_attn.q.weight'],
+                pytorch_model[f'{layer_prefix}self_attn.k.weight'],
+                pytorch_model[f'{layer_prefix}self_attn.v.weight'])
+            layer.attention.dense.weight.value = pytorch_model[
+                f'{layer_prefix}self_attn.out_proj.weight']
+
+            if tllm_model.has_attention_qkvo_bias:
+                layer.attention.qkv.bias.value = fuse_qkv(
+                    pytorch_model[f'{layer_prefix}self_attn.q.bias'],
+                    torch.zeros(pytorch_model[f'{layer_prefix}self_attn.q.bias'].shape), # no bias for k
+                    pytorch_model[f'{layer_prefix}self_attn.v.bias']
+                )
+                layer.attention.dense.bias.value = pytorch_model[
+                    f'{layer_prefix}self_attn.out_proj.bias']
+
+            layer.attention_layernorm.weight.value = pytorch_model[
+                f'{layer_prefix}self_attn_layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                layer.attention_layernorm.bias.value = pytorch_model[
+                    f'{layer_prefix}self_attn_layer_norm.bias']
+
+            layer.mlp.fc.weight.value = pytorch_model[
+                f'{layer_prefix}fc1.weight']
+            layer.mlp.proj.weight.value = pytorch_model[
+                f'{layer_prefix}fc2.weight']
+
+            if tllm_model.has_mlp_bias:
+                layer.mlp.fc.bias.value = pytorch_model[
+                    f'{layer_prefix}fc1.bias']
+                layer.mlp.proj.bias.value = pytorch_model[
+                    f'{layer_prefix}fc2.bias']
+
+            layer.mlp_layernorm.weight.value = pytorch_model[
+                f'{layer_prefix}final_layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                layer.mlp_layernorm.bias.value = pytorch_model[
+                    f'{layer_prefix}final_layer_norm.bias']
+
+        if tllm_model.final_layernorm:
+            tllm_model.final_layernorm.weight.value = pytorch_model[
+                'encoder.layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                tllm_model.final_layernorm.bias.value = pytorch_model[
+                    'encoder.layer_norm.bias']
+
+    if component == "decoder":
+        tllm_model.embedding.vocab_embedding.weight.value = pytorch_model[
+            'shared.weight']
+        if tllm_model.embedding.position_embedding:
+            tllm_model.embedding.position_embedding.weight.value = pytorch_model[
+                'decoder.embed_positions.weight']
+        if tllm_model.embedding.token_type_embedding:
+            tllm_model.embedding.token_type_embedding.weight.value = pytorch_model[
+                'decoder.embed_token_type.weight']
+
+        # all layers use 1st layer's attn table
+        # transpose from [num_buckets, num_heads] --> [num_heads, num_buckets]
+        # ascontiguousarray is very important! otherwise TRT always receives the original layout
+        relative_attention_table = np.ascontiguousarray(pytorch_model[
+            f'decoder.block.0.layer.0.SelfAttention.relative_attention_bias.weight']
+                                                        .T)
+
+        for i in range(tllm_model.num_layers):
+            layer = tllm_model.decoder_layers[i]
+            layer_prefix = f'decoder.block.{i}.'
+
+            # attention table for all layers
+            layer.self_attention.rel_attn_table.value = relative_attention_table
+
+            # self attn
+            layer.self_attention.qkv.weight.value = fuse_qkv(
+                pytorch_model[f'{layer_prefix}layer.0.SelfAttention.q.weight'],
+                pytorch_model[f'{layer_prefix}layer.0.SelfAttention.k.weight'],
+                pytorch_model[f'{layer_prefix}layer.0.SelfAttention.v.weight'])
+            layer.self_attention.dense.weight.value = pytorch_model[
+                f'{layer_prefix}layer.0.SelfAttention.o.weight']
+
+            if tllm_model.has_attention_qkvo_bias:
+                layer.self_attention.qkv.bias.value = fuse_qkv(
+                    pytorch_model[
+                        f'{layer_prefix}layer.0.SelfAttention.q.bias'],
+                    pytorch_model[
+                        f'{layer_prefix}layer.0.SelfAttention.k.bias'],
+                    pytorch_model[f'{layer_prefix}layer.0.SelfAttention.v.bias']
+                )
+                layer.self_attention.dense.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.0.SelfAttention.o.bias']
+
+            layer.self_attention_layernorm.weight.value = pytorch_model[
+                f'{layer_prefix}layer.0.layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                layer.self_attention_layernorm.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.0.layer_norm.bias']
+
+            # cross attn
+            layer.cross_attention.qkv.weight.value = fuse_qkv(
+                pytorch_model[
+                    f'{layer_prefix}layer.1.EncDecAttention.q.weight'],
+                pytorch_model[
+                    f'{layer_prefix}layer.1.EncDecAttention.k.weight'],
+                pytorch_model[f'{layer_prefix}layer.1.EncDecAttention.v.weight']
+            )
+            layer.cross_attention.dense.weight.value = pytorch_model[
+                f'{layer_prefix}layer.1.EncDecAttention.o.weight']
+
+            if tllm_model.has_attention_qkvo_bias:
+                layer.cross_attention.qkv.bias.value = fuse_qkv(
+                    pytorch_model[
+                        f'{layer_prefix}layer.1.EncDecAttention.q.bias'],
+                    pytorch_model[
+                        f'{layer_prefix}layer.1.EncDecAttention.k.bias'],
+                    pytorch_model[
+                        f'{layer_prefix}layer.1.EncDecAttention.v.bias'])
+                layer.cross_attention.dense.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.1.EncDecAttention.o.bias']
+
+            layer.cross_attention_layernorm.weight.value = pytorch_model[
+                f'{layer_prefix}layer.1.layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                layer.cross_attention_layernorm.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.1.layer_norm.bias']
+
+            layer.mlp.fc.weight.value = pytorch_model[
+                f'{layer_prefix}layer.2.DenseReluDense.wi.weight']
+            layer.mlp.proj.weight.value = pytorch_model[
+                f'{layer_prefix}layer.2.DenseReluDense.wo.weight']
+
+            if tllm_model.has_mlp_bias:
+                layer.mlp.fc.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.2.DenseReluDense.wi.bias']
+                layer.mlp.proj.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.2.DenseReluDense.wo.bias']
+
+            layer.mlp_layernorm.weight.value = pytorch_model[
+                f'{layer_prefix}layer.2.layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                layer.mlp_layernorm.bias.value = pytorch_model[
+                    f'{layer_prefix}layer.2.layer_norm.bias']
+
+        if tllm_model.final_layernorm:
+            tllm_model.final_layernorm.weight.value = pytorch_model[
+                'decoder.final_layer_norm.weight']
+            if tllm_model.layernorm_type != LayerNormType.RmsNorm:
+                tllm_model.final_layernorm.bias.value = pytorch_model[
+                    'decoder.final_layer_norm.bias']
+
+        tllm_model.lm_head.weight.value = pytorch_model['lm_head.weight']
+
+if __name__=="__main__":
+    # print(layernorm_type_map)
+    # print(layernorm_position_map)
+    load_whisper_from_pytorch(