update chrome plugin with model size dropdown

This commit is contained in:
makaveli10
2023-12-14 23:52:25 +05:30
parent 7ffcad64ba
commit 14beb4f942
5 changed files with 277 additions and 5 deletions
+2 -1
View File
@@ -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 {
+2 -1
View File
@@ -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
})
);
};
+12 -1
View File
@@ -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>
+18 -2
View File
@@ -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;
+243
View File
@@ -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(