BOUND/setup.py at main · RUCAIBox/BOUND · 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
from pathlib import Path

from setuptools import find_packages, setup


ROOT = Path(__file__).resolve().parent


setup(
    name="search-control-preferences",
    version="0.1.0",
    description=(
        "Reference implementation for search-control preference "
        "construction and training"
    ),
    long_description=(ROOT / "README.md").read_text(encoding="utf-8"),
    long_description_content_type="text/markdown",
    license="MIT",
    python_requires=">=3.10",

    package_dir={"": "src"},
    packages=find_packages("src"),

    py_modules=[
        "context",
        "jsonl_io",
        "preferences",
        "prompts",
        "questions",
        "schema",
        "teacher",
    ],

    # Lightweight dependencies used by the core preference-construction
    # pipeline. Training, local inference, and benchmark-specific utilities
    # are exposed through optional extras below.
    install_requires=[
        "openai>=1.30",
    ],

    extras_require={
        # DPO training dependencies.
        "train": [
            "accelerate>=0.30",
            "datasets>=4.2,<4.3",
            "torch>=2.1",
            "transformers>=4.57,<4.58",
            "trl>=0.24,<0.25",
        ],

        # Local policy inference and rerouting passage scoring.
        "infer": [
            "accelerate>=0.30",
            "torch>=2.1",
            "transformers>=4.57,<4.58",
            "vllm",
        ],

        # Benchmark and training-question preparation.
        "data": [
            "datasets>=2.19",
            "huggingface-hub>=0.23",
        ],

        # GAIA retriever service.
        "gaia": [
            "requests>=2.31",
        ],

        # Development and testing.
        "test": [
            "pytest>=8",
        ],

        # Full installation for reproducing the released pipeline.
        "all": [
            "accelerate>=0.30",
            "datasets>=4.2,<4.3",
            "huggingface-hub>=0.23",
            "pytest>=8",
            "requests>=2.31",
            "torch>=2.1",
            "transformers>=4.57,<4.58",
            "trl>=0.24,<0.25",
            "vllm",
        ],
    },

    entry_points={
        "console_scripts": [
            "search-build-preferences=cli.build_preferences:main",
            "search-train=cli.train:main",
            "search-infer=cli.infer:main",
            "search-evaluate=cli.evaluate:main",
            "search-prepare-benchmark=cli.prepare_benchmark:main",
            (
                "search-prepare-training-questions="
                "cli.prepare_training_questions:main"
            ),
            (
                "search-serve-gaia-retriever="
                "cli.serve_gaia_retriever:main"
            ),
        ],
    },
)