update chrome plugin with model size dropdown
This commit is contained in:
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
})
|
||||
);
|
||||
};
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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(
|
||||
Reference in New Issue
Block a user