dice-embeddings/setup.py at develop · dice-group/dice-embeddings · 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
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
import re
from setuptools import setup, find_packages

# Installation options:
# - Default install (CPU-only recommended): pip install dicee --extra-index-url https://download.pytorch.org/whl/cpu
# - GPU/CUDA install: pip install dicee
# - Development: pip install dicee[dev] --extra-index-url https://download.pytorch.org/whl/cpu
# - Documentation: pip install dicee[docs]
# - All extras: pip install dicee[all] --extra-index-url https://download.pytorch.org/whl/cpu
#
# NOTE: To avoid ~2GB NVIDIA CUDA dependencies, always use --extra-index-url https://download.pytorch.org/whl/cpu

# Core dependencies (includes torch and lightning)
# Use --extra-index-url https://download.pytorch.org/whl/cpu for CPU-only installation
_core_deps = [
    "numpy==1.26.4",
    "torch>=2.5.1",
    "lightning>=2.5.0.post0",
    "pandas>=2.1.0",
    "requests>=2.32.3",
    "polars>=0.16.14",
    "tiktoken>=0.5.1",
    "psutil>=5.9.4",
]

# Optional dependencies for various features
_optional_deps = [
    "pyarrow>=11.0.0",
    "rdflib>=7.0.0",
    "tiktoken>=0.5.1",
    "pykeen>=1.10.2",
    "matplotlib>=3.8.2",
    "zstandard>=0.21.0",
    "requests>=2.32.3",
    "scikit-learn>=1.2.2",
]

# Development dependencies
_dev_deps = [
    "pytest>=7.2.2",
    "ruff>=0.0.284",
    "mypy>=1.0.0",
    "scikit-learn>=1.2.2",
]

# Documentation dependencies
_docs_deps = [
    "sphinx>=7.2.6",
    "sphinx-autoapi>=3.0.0",
    "myst-parser>=2.0.0",
    "sphinx_rtd_theme>=2.0.0",
    "sphinx-theme>=1.0",
    "sphinxcontrib-plantuml>=0.27",
    "plantuml-local-client>=1.2022.6",
]

# Combine all dependencies for regex parsing
_all_deps = _core_deps + _optional_deps + _dev_deps + _docs_deps

# Parse dependencies into a dictionary
deps = {b: a for a, b in (re.findall(r"^(([^!=<>~ ]+)(?:[!=<>~ ].*)?$)", x)[0] for x in _all_deps)}


def deps_list(*pkgs):
    return [deps[pkg] for pkg in pkgs]


# Define extras
extras = dict()

# Minimal installation - only core dependencies
extras["min"] = _core_deps

# Development extras
extras["dev"] = _core_deps + _optional_deps + _dev_deps

# Documentation extras
extras["docs"] = _docs_deps

# All extras
extras["all"] = _core_deps + _optional_deps + _dev_deps + _docs_deps

# Base installation includes core dependencies (torch and lightning included)
# Use --extra-index-url https://download.pytorch.org/whl/cpu to avoid NVIDIA CUDA dependencies
install_requires = _core_deps

with open('README.md', 'r') as fh:
    long_description = fh.read()
setup(
    name="dicee",
    description="Dice embedding is an hardware-agnostic framework for large-scale knowledge graph embedding applications",
    version="0.3.3",
    packages=find_packages(exclude=["tests", "tests.*"]),
    extras_require=extras,
    install_requires=list(install_requires),
    author='Caglar Demir',
    author_email='caglardemir8@gmail.com',
    url='https://github.com/dice-group/dice-embeddings',
    classifiers=[
        "Programming Language :: Python :: 3.11",
        "License :: OSI Approved :: MIT License"],
    python_requires='>=3.11',
    entry_points={"console_scripts":
                      ["dicee=dicee.scripts.run:main",
                       "dicee_vector_db=dicee.scripts.index_serve:main"]},
    long_description=long_description,
    long_description_content_type="text/markdown",
)