Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
"""Declares specification of the Transformer model."""
import numpy as np
from ctranslate2.specs import attention_spec
from ctranslate2.specs import common_spec
from ctranslate2.specs import model_spec
class TransformerSpec(model_spec.ModelSpec):
"""Describes a Transformer model.
The specification is invariant to hidden dimensions but requires to
explicitly set the number of layers and attention heads.
"""
def __init__(self, num_layers, num_heads):
self.num_heads = np.dtype("int8").type(num_heads)
self.encoder = TransformerEncoderSpec(num_layers)
self.decoder = TransformerDecoderSpec(num_layers)
@property
def revision(self):
return 3
@property
def source_vocabulary_size(self):