models/bert_modules/attention/__init__.py

from .multi_head import MultiHeadedAttention
from .single import Attention