Robert Riachi RobertRiachi

Research Engineer @deepmind // prev SWE & ML @splunk, @bloomberg, @SAP, @Deloitte // UWaterloo Double Major CS & Stats

RobertRiachi / Whisper_ANE_export.py

Last active February 3, 2024 01:42

Export an optimized version of Whisper for ANE using coreml

	import torch
	import torch.nn.functional as F
	import coremltools as ct

	from torch import Tensor
	from torch import nn
	from typing import Dict
	from typing import Optional
	from ane_transformers.reference.layer_norm import LayerNormANE as LayerNormANEBase
	from coremltools.models.neural_network.quantization_utils import quantize_weights