import dataclasses
from typing import Optional


@dataclasses.dataclass
class BelleParam:
    num_heads: int = 32
    size_per_head: int = 128
    inter_size: int = 16384
    num_layers: int = 30
    vocab_size: int = 250880
    start_id: Optional[int] = 1
    end_id: Optional[int] = 2
    tensor_para_size: int = 1
    pipeline_para_size: int = 1
    remove_padding: bool = True
    shared_contexts_ratio: float = 1.0
    weights_data_type: str = "fp16"

    def __post_init__(self):
        if not 0.0 <= self.shared_contexts_ratio <= 1.0:
            raise ValueError(
                f'Got an invalid value of shared_context_ratio '
                f'{self.shared_contexts_ratio} - range: [0.0, 1.0]')

    def asdict(self):
        return dataclasses.asdict(self)


BELLE_PARAM = BelleParam()
import os
current_dir = os.path.dirname(os.path.abspath(__file__))
LIB_SO_PATH = os.path.join(current_dir, 'libth_transformer.so')