pcatt/setup.py at main · PreferredAI/pcatt · GitHub

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
from sysconfig import get_path
from setuptools import setup, Extension
PATH_PREFIXES = [get_path(p) for p in ["data", "platlib"]]

include_dirs = [
    path
    for prefix in PATH_PREFIXES
    for path in (
        f"{prefix}/include/",
        f"{prefix}/include/tbb",
        f"{prefix}/pybind11/include",
    )
]

modules = []

for code in ["greedy_encoder", "pco_tokenizer"]:
    modules.append(
        Extension(
            name=f"pcatt.{code}",
            sources=[f"pcatt/{code}.cpp"],
            include_dirs=include_dirs,
            library_dirs=[f"{prefix}/lib/" for prefix in PATH_PREFIXES],
            libraries=["tbb"],
            extra_compile_args=["-O3", "-std=c++23"],
            define_macros=[
                ("MAJOR_VERSION", "0"),
                ("MINOR_VERSION", "15"),
            ],
            extra_link_args=[
                "-Wl,-rpath,$ORIGIN/.libs",
                "-Wl,--enable-new-dtags",
            ],
        )
)


setup(ext_modules=modules)