File size: 1,165 Bytes
b88c26d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 | """
BURT-IMMA: Bi-encoder Unified Retrieval Transformer with
Interferometric Matrix Memory Architecture
License: BSL-1.1
Contact: jessica@collectivekitty.com
"""
__version__ = "0.1.0"
# Try to import CUDA-accelerated kernels, fall back to pure PyTorch
_CUDA_AVAILABLE = False
try:
import _burt_imma_cuda
_CUDA_AVAILABLE = True
except ImportError:
_burt_imma_cuda = None
# Import kernel wrappers (these handle CUDA/CPU fallback internally)
from .kernels import (
constrained_softmax,
cifg_update,
batched_cifg_update,
sparse_moe_dispatch,
biencoder_attention,
attention_softmax,
)
from .utils import (
spectral_norm,
entropy,
check_huntington,
load_config,
set_seed,
Timer,
)
__all__ = [
# Kernel functions
"constrained_softmax",
"cifg_update",
"batched_cifg_update",
"sparse_moe_dispatch",
"biencoder_attention",
"attention_softmax",
# Utilities
"spectral_norm",
"entropy",
"check_huntington",
"load_config",
"set_seed",
"Timer",
# Meta
"__version__",
"_CUDA_AVAILABLE",
]
|