CoolFace
Apppublic

pennaburry/parallel-constrained-decoding

sourceHugging Faceapache-2.0updated 6d agoView on Hugging Face
1likes
engine.py44 linesDownload Raw Back to core
1"""2Unified Engine Router for Parallel Constrained Decoding.3Automatically selects MLX backend on Apple Silicon macOS,4or PyTorch / CUDA backend on Linux, Docker, and Hugging Face Spaces.5"""6 7import os8import platform9 10USE_MLX = False11if platform.system() == "Darwin" and os.environ.get("BACKEND", "").lower() != "torch":12    try:13        import mlx.core as mx14        import mlx_lm15        USE_MLX = True16    except Exception:17        USE_MLX = False18 19if USE_MLX:20    from core.engine_mlx import (21        get_engine,22        run_parallel_generation,23        run_naive_generation,24        stream_naive_generation,25        run_rlcd_generation,26    )27else:28    from core.engine_torch import (29        get_torch_engine as get_engine,30        run_parallel_generation_torch as run_parallel_generation,31        run_naive_generation_torch as run_naive_generation,32        stream_naive_generation_torch as stream_naive_generation,33    )34    run_rlcd_generation = run_parallel_generation35 36__all__ = [37    "get_engine",38    "run_parallel_generation",39    "run_naive_generation",40    "stream_naive_generation",41    "run_rlcd_generation",42    "USE_MLX",43]44