From 14beb4f9429b5932b0375b536786cd2882b8fad9 Mon Sep 17 00:00:00 2001 From: makaveli10 Date: Thu, 14 Dec 2023 23:52:25 +0530 Subject: [PATCH] update chrome plugin with model size dropdown --- Audio-Transcription-Chrome/background.js | 3 +- Audio-Transcription-Chrome/options.js | 3 +- Audio-Transcription-Chrome/popup.html | 13 +- Audio-Transcription-Chrome/popup.js | 20 +- Audio-Transcription-Chrome/test.py | 243 +++++++++++++++++++++++ 5 files changed, 277 insertions(+), 5 deletions(-) create mode 100644 Audio-Transcription-Chrome/test.py diff --git a/Audio-Transcription-Chrome/background.js b/Audio-Transcription-Chrome/background.js index f860209..a5028a7 100644 --- a/Audio-Transcription-Chrome/background.js +++ b/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 { diff --git a/Audio-Transcription-Chrome/options.js b/Audio-Transcription-Chrome/options.js index af77b38..9717856 100644 --- a/Audio-Transcription-Chrome/options.js +++ b/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 }) ); }; diff --git a/Audio-Transcription-Chrome/popup.html b/Audio-Transcription-Chrome/popup.html index 4f2b6ec..79f18ca 100644 --- a/Audio-Transcription-Chrome/popup.html +++ b/Audio-Transcription-Chrome/popup.html @@ -125,11 +125,22 @@
-
+
+ + +
diff --git a/Audio-Transcription-Chrome/popup.js b/Audio-Transcription-Chrome/popup.js index f67f8b7..a1dd210 100644 --- a/Audio-Transcription-Chrome/popup.js +++ b/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; diff --git a/Audio-Transcription-Chrome/test.py b/Audio-Transcription-Chrome/test.py new file mode 100644 index 0000000..ddd69c9 --- /dev/null +++ b/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( \ No newline at end of file