forked from k2-fsa/k2
-
Notifications
You must be signed in to change notification settings - Fork 0
/
setup.py
132 lines (110 loc) · 4.12 KB
/
setup.py
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
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# Welcome to the k2 setup.py.
#
# Please follow instructions in scripts/build_pip.sh to use this file.
#
import datetime
import re
import setuptools
import sys
if sys.version_info < (3,):
print('Python 2 has reached end-of-life and is no longer supported by k2.')
sys.exit(-1)
if sys.version_info < (3, 6):
print('Python 3.5 has reached end-of-life on September 13th, 2020 '
'and is no longer supported by k2.')
sys.exit(-1)
# Refer to https://stackoverflow.com/questions/45150304/how-to-force-a-python-wheel-to-be-platform-specific-when-building-it
# for why to introduce `bdist_wheel`.
#
# With `bdist_wheel`, the final wheel name looks like `k2-0.0.1.dev20201104-cp37-cp37m-linux_x86_64.whl`
try:
from wheel.bdist_wheel import bdist_wheel as _bdist_wheel
class bdist_wheel(_bdist_wheel):
def finalize_options(self):
_bdist_wheel.finalize_options(self)
self.root_is_pure = False
except ImportError:
bdist_wheel = None
def get_long_description():
with open('README.md', 'r') as f:
long_description = f.read()
return long_description
def get_cuda_version():
import torch
from torch.utils import collect_env
running_cuda_version = collect_env.get_running_cuda_version(
collect_env.run)
cuda_version = torch.version.cuda
if running_cuda_version is not None:
assert cuda_version in running_cuda_version, \
f'PyTorch is built with CUDA version: {cuda_version}.\n' \
f'The current running CUDA version is: {running_cuda_version}'
cuda_version = cuda_version.split('.')
major, minor = int(cuda_version[0]), int(cuda_version[1])
cuda_version = major * 10 + minor
return f'{cuda_version}'
def get_package_version():
# Set a default CUDA version here so that `pip install k2`
# uses the default CUDA version.
#
# `pip install k2==x.x.x+cu100` to install k2 with CUDA 10.0
#
default_cuda_version = '101' # CUDA 10.1
cuda_version = get_cuda_version()
if default_cuda_version != cuda_version:
cuda_version = f'+cu{cuda_version}'
else:
cuda_version = ''
with open('CMakeLists.txt') as f:
content = f.read()
latest_version = re.search(r'set\(K2_VERSION (.*)\)', content).group(1)
latest_version = latest_version.strip('"')
dt = datetime.datetime.utcnow()
package_version = f'{latest_version}{cuda_version}.dev{dt.year}{dt.month:02d}{dt.day:02d}'
return package_version
def get_short_description():
return 'FSA/FST algorithms, intended to (eventually) be interoperable with PyTorch and similar'
with open('k2/python/k2/__init__.py', 'a') as f:
f.write(f"__dev_version__ = '{get_package_version()}'\n")
dev_requirements = [
'clang-format==9.0.0',
'flake8==3.8.3',
'yapf==0.27.0',
]
setuptools.setup(
python_requires='>=3.6',
name='k2',
version=get_package_version(),
author='Daniel Povey',
author_email='[email protected]',
keywords='k2, FSA, FST',
description=get_short_description(),
long_description=get_long_description(),
long_description_content_type='text/markdown',
url='https://github.com/k2-fsa/k2',
package_dir={
'k2': 'k2/python/k2',
'k2.ragged': 'k2/python/k2/ragged',
'k2.sparse': 'k2/python/k2/sparse',
},
packages=['k2', 'k2.ragged', 'k2.sparse'],
install_requires=['torch', 'graphviz'],
extras_require={'dev': dev_requirements},
data_files=[('', ['LICENSE'])],
cmdclass={'bdist_wheel': bdist_wheel},
classifiers=[
'Development Status :: 3 - Alpha',
'Programming Language :: Python :: 3',
'Programming Language :: C++',
'Programming Language :: Python :: Implementation :: CPython',
'Topic :: Scientific/Engineering :: Artificial Intelligence',
'Operating System :: OS Independent',
],
)
# remove the line __dev_version__ from k2/python/k2/__init__.py
with open('k2/python/k2/__init__.py', 'r') as f:
lines = f.readlines()
with open('k2/python/k2/__init__.py', 'w') as f:
for line in lines:
if '__dev_version__' not in line:
f.write(line)