orbitals 0.0.2__tar.gz → 0.0.3__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. orbitals-0.0.3/.dockerignore +5 -0
  2. orbitals-0.0.3/CONTRIBUTING.md +45 -0
  3. {orbitals-0.0.2 → orbitals-0.0.3}/Dockerfile +4 -2
  4. orbitals-0.0.3/PKG-INFO +161 -0
  5. orbitals-0.0.3/README.md +109 -0
  6. {orbitals-0.0.2 → orbitals-0.0.3}/README.scope-guard.md +90 -39
  7. orbitals-0.0.3/examples/scope-guard/test.py +27 -0
  8. {orbitals-0.0.2 → orbitals-0.0.3}/pyproject.toml +8 -3
  9. orbitals-0.0.3/src/orbitals/utils.py +53 -0
  10. orbitals-0.0.2/.dockerignore +0 -2
  11. orbitals-0.0.2/PKG-INFO +0 -124
  12. orbitals-0.0.2/README.md +0 -77
  13. orbitals-0.0.2/src/orbitals/utils.py +0 -47
  14. {orbitals-0.0.2 → orbitals-0.0.3}/.gitignore +0 -0
  15. {orbitals-0.0.2 → orbitals-0.0.3}/AGENTS.md +0 -0
  16. {orbitals-0.0.2 → orbitals-0.0.3}/LICENSE +0 -0
  17. {orbitals-0.0.2 → orbitals-0.0.3}/assets/orbitals-banner.png +0 -0
  18. {orbitals-0.0.2 → orbitals-0.0.3}/assets/orbitals.svg +0 -0
  19. {orbitals-0.0.2 → orbitals-0.0.3}/assets/scope-guard.svg +0 -0
  20. {orbitals-0.0.2 → orbitals-0.0.3}/examples/scope-guard/api.py +0 -0
  21. {orbitals-0.0.2 → orbitals-0.0.3}/examples/scope-guard/async_api.py +0 -0
  22. {orbitals-0.0.2 → orbitals-0.0.3}/examples/scope-guard/local.py +0 -0
  23. {orbitals-0.0.2 → orbitals-0.0.3}/examples/scope-guard/vllm_serve.py +0 -0
  24. {orbitals-0.0.2 → orbitals-0.0.3}/src/hf_pipeline/scope_guard.py +0 -0
  25. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/__init__.py +0 -0
  26. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/cli/__init__.py +0 -0
  27. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/cli/main.py +0 -0
  28. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/__init__.py +0 -0
  29. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/cli/__init__.py +0 -0
  30. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/cli/convert_default_model_name.py +0 -0
  31. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/cli/main.py +0 -0
  32. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/cli/serve.py +0 -0
  33. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/guards/__init__.py +0 -0
  34. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/guards/api.py +0 -0
  35. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/guards/base.py +0 -0
  36. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/guards/hf.py +0 -0
  37. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/guards/vllm.py +0 -0
  38. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/modeling.py +0 -0
  39. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/prompting.py +0 -0
  40. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/serving/__init__.py +0 -0
  41. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/serving/main.py +0 -0
  42. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/scope_guard/serving/vllm_logging_config.json +0 -0
  43. {orbitals-0.0.2 → orbitals-0.0.3}/src/orbitals/types.py +0 -0
  44. {orbitals-0.0.2 → orbitals-0.0.3}/src/scripts/playground.ipynb +0 -0
  45. {orbitals-0.0.2 → orbitals-0.0.3}/src/scripts/push_hf_pipeline.py +0 -0
  46. {orbitals-0.0.2 → orbitals-0.0.3}/src/scripts/push_model.py +0 -0
  47. {orbitals-0.0.2 → orbitals-0.0.3}/uv.lock +0 -0
@@ -0,0 +1,5 @@
1
+ .venv
2
+ dist
3
+ __pycache__/
4
+ .vscode/
5
+ examples/
@@ -0,0 +1,45 @@
1
+ # Contributing to Orbitals
2
+
3
+ Thank you for your interest in contributing to `orbitals`! We welcome contributions from the community.
4
+
5
+ ## Getting Started
6
+
7
+ 1. **Fork the repository** and clone it locally
8
+ 2. **Install dependencies** using `uv`:
9
+ ```bash
10
+ uv sync
11
+ ```
12
+
13
+ ## Development Workflow
14
+
15
+ 1. Create a new branch for your feature or fix:
16
+ ```bash
17
+ git checkout -b feature/your-feature-name
18
+ ```
19
+
20
+ 2. Make your changes and ensure they follow the project's coding conventions
21
+
22
+ 3. Test your changes thoroughly
23
+
24
+ 4. Commit your changes with a clear, descriptive message
25
+
26
+ 5. Push to your fork and submit a pull request
27
+
28
+ ## Coding Conventions
29
+
30
+ - Use **Google-style docstrings** for all documentation
31
+ - Follow PEP 8 style guidelines
32
+ - Keep code modular and well-organized
33
+
34
+ ## Reporting Issues
35
+
36
+ When reporting issues, please include:
37
+
38
+ - A clear description of the problem
39
+ - Steps to reproduce the issue
40
+ - Expected vs actual behavior
41
+ - Your environment details (Python version, OS, etc.)
42
+
43
+ ## Questions?
44
+
45
+ Feel free to open an issue for any questions or discussions about contributing.
@@ -25,13 +25,15 @@ RUN --mount=type=cache,target=/root/.cache/uv \
25
25
  COPY src /app/src/
26
26
 
27
27
  RUN --mount=type=cache,target=/root/.cache/uv \
28
- uv sync --locked --extra serve --no-editable --no-dev
28
+ uv sync --locked --extra scope-guard-serve --no-editable --no-dev
29
29
 
30
30
  # TODO remove next line
31
31
  ENTRYPOINT ["/bin/bash", "-c"]
32
32
 
33
33
  FROM vllm/vllm-openai:v0.14.1 as runner
34
34
 
35
+ ARG MODEL
36
+
35
37
  WORKDIR /app
36
38
  ENV PATH="/app/.venv/bin:$PATH"
37
39
 
@@ -43,4 +45,4 @@ COPY --from=builder /root/.cache/huggingface/hub /root/.cache/huggingface/hub
43
45
  EXPOSE 8000
44
46
 
45
47
  ENTRYPOINT ["/bin/bash", "-c"]
46
- CMD [ "orbitals", "scope-guard", "serve", "${MODEL}", "--port", "8000", "--host", "0.0.0.0" ]
48
+ CMD [ "orbitals", "scope-guard", "serve", "principled-intelligence/${MODEL}", "--port", "8000", "--host", "0.0.0.0" ]
@@ -0,0 +1,161 @@
1
+ Metadata-Version: 2.4
2
+ Name: orbitals
3
+ Version: 0.0.3
4
+ Summary: LLM Guardrails tailored to your Principles
5
+ Author-email: Luigi Procopio <luigi@principled-intelligence.com>, Edoardo Barba <edoardo@principled-intelligence.com>
6
+ License: Apache-2.0
7
+ License-File: LICENSE
8
+ Classifier: Programming Language :: Python :: 3 :: Only
9
+ Classifier: Programming Language :: Python :: 3.10
10
+ Classifier: Programming Language :: Python :: 3.11
11
+ Classifier: Programming Language :: Python :: 3.12
12
+ Classifier: Programming Language :: Python :: 3.13
13
+ Classifier: Programming Language :: Python :: 3.14
14
+ Requires-Python: >=3.10
15
+ Requires-Dist: aiohttp
16
+ Requires-Dist: pydantic>=2.0.0
17
+ Requires-Dist: requests
18
+ Requires-Dist: typer>=0.12.3
19
+ Provides-Extra: all
20
+ Requires-Dist: accelerate>=1.11.0; extra == 'all'
21
+ Requires-Dist: fastapi[standard]>=0.119.1; extra == 'all'
22
+ Requires-Dist: pynvml; extra == 'all'
23
+ Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'all'
24
+ Requires-Dist: uvicorn>=0.29.0; extra == 'all'
25
+ Requires-Dist: vllm>=0.11.0; extra == 'all'
26
+ Requires-Dist: xgrammar; extra == 'all'
27
+ Provides-Extra: scope-guard-all
28
+ Requires-Dist: accelerate>=1.11.0; extra == 'scope-guard-all'
29
+ Requires-Dist: fastapi[standard]>=0.119.1; extra == 'scope-guard-all'
30
+ Requires-Dist: pynvml; extra == 'scope-guard-all'
31
+ Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-all'
32
+ Requires-Dist: uvicorn>=0.29.0; extra == 'scope-guard-all'
33
+ Requires-Dist: vllm>=0.11.0; extra == 'scope-guard-all'
34
+ Requires-Dist: xgrammar; extra == 'scope-guard-all'
35
+ Provides-Extra: scope-guard-hf
36
+ Requires-Dist: accelerate>=1.11.0; extra == 'scope-guard-hf'
37
+ Requires-Dist: pynvml; extra == 'scope-guard-hf'
38
+ Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-hf'
39
+ Provides-Extra: scope-guard-serve
40
+ Requires-Dist: fastapi[standard]>=0.119.1; extra == 'scope-guard-serve'
41
+ Requires-Dist: pynvml; extra == 'scope-guard-serve'
42
+ Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-serve'
43
+ Requires-Dist: uvicorn>=0.29.0; extra == 'scope-guard-serve'
44
+ Requires-Dist: vllm>=0.11.0; extra == 'scope-guard-serve'
45
+ Requires-Dist: xgrammar; extra == 'scope-guard-serve'
46
+ Provides-Extra: scope-guard-vllm
47
+ Requires-Dist: pynvml; extra == 'scope-guard-vllm'
48
+ Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-vllm'
49
+ Requires-Dist: vllm>=0.11.0; extra == 'scope-guard-vllm'
50
+ Requires-Dist: xgrammar; extra == 'scope-guard-vllm'
51
+ Description-Content-Type: text/markdown
52
+
53
+ <div align="center">
54
+ <img src="assets/orbitals-banner.png" width="70%" />
55
+ <h3 align="center">
56
+ <p>
57
+ <b>LLM Guardrails tailored to your Principles</b>
58
+ </p>
59
+ </h4>
60
+ </div>
61
+
62
+ <p align="center">
63
+ <img src="https://img.shields.io/pypi/v/orbitals?color=green" alt="PyPI Version">
64
+ <!-- <img src="https://img.shields.io/badge/type%20checked-ty-blue.svg?color=green" alt="Type Checked with ty"> -->
65
+ <img src="https://img.shields.io/pypi/pyversions/orbitals" alt="Python Versions">
66
+ <img src="https://img.shields.io/github/license/principled-intelligence/orbitals" alt="GitHub License">
67
+ </p>
68
+
69
+ ## Overview
70
+
71
+ **Orbitals** is a lightweight Python library for adding LLM guardrails in just a few lines of code. With Orbitals, you can add a governance layer tailored to **user-specific principles**. Rather than enforcing generic notions of safety, compliance, and correctness, Orbitals validates inputs (e.g., user requests) and outputs (e.g., assistant responses) against user-defined specifications and custom policies. This makes guardrails explicit, auditable, and aligned with the user's philosophy.
72
+
73
+ ### Key Features
74
+
75
+ - **User-defined specifications** — Guardrails that match your use case and your custom policies, not generic safety rules
76
+ - **Simple integration** — Add guardrails with minimal code changes
77
+ - **Open framework, open models** — Orbitals is open-source and is a simple interface for our open models
78
+
79
+ ## Getting started
80
+
81
+ ### Installation
82
+
83
+ You can install Orbitals via pip:
84
+
85
+ ```bash
86
+ pip install orbitals[all]
87
+ ```
88
+
89
+ ### Basic Usage
90
+
91
+ Here's a quick example to get you started with Orbitals, in which we use the ScopeGuard module to guard an AI service (for example, a customer support chatbot) from user requests that violate specified principles or fall outside of the scope of the core task of the assistant.
92
+
93
+ ```python
94
+ from orbitals.scope_guard import ScopeGuard
95
+
96
+ ai_service_description = "You are a helpful assistant for ..."
97
+ user_message = "Can I buy ..."
98
+
99
+ guardrail = ScopeGuard()
100
+ result = guardrail.validate(user_message, ai_service_description)
101
+ ```
102
+
103
+ The result of a guardrail validation will indicate whether the input or output passed the guardrail checks, along with details on any violations. You can then handle violations as needed, such as by rejecting the input or modifying the output. For example:
104
+
105
+ ```python
106
+ if result.scope_class.value == "Restricted" or result.scope_class.value == "Out of Scope":
107
+ print("Request violates guardrail:", result.evidences)
108
+ else:
109
+ # The user request is safe!
110
+ # We can now pass it to the AI assistant for processing.
111
+ ...
112
+ ```
113
+
114
+ ### Available Guardrails
115
+
116
+ Orbitals currently provides the following guardrail modules:
117
+
118
+ | Guardrail | Description | Hosting Options |
119
+ |:----------|:------------|:----------------|
120
+ | **[ScopeGuard](README.scope-guard.md)** | Classifies user queries against AI assistant specifications to detect out-of-scope requests, policy violations, and chit-chat | Self-hosted / Cloud hosting |
121
+ | 🚀 *Coming Soon* | More guardrails are on the way — stay tuned for updates! | — |
122
+
123
+ #### Hosting Options
124
+
125
+ - **Self-hosted**: Use open-weight models that you can deploy on your own infrastructure, ensuring data privacy and control.
126
+ - **Cloud hosting**: (Coming soon) Managed hosting options for ease of use and scalability
127
+
128
+ ### Documentation
129
+
130
+ For detailed documentation, including installation instructions, usage guides, and API references, please visit the Orbitals Documentation.
131
+
132
+ - [ScopeGuard Documentation](README.scope-guard.md)
133
+
134
+ ### FAQ
135
+
136
+ - **Can I use Orbitals for commercial applications?**
137
+ Yes, Orbitals is designed to be used in both research and commercial applications. It is licensed under the Apache 2.0 License, which allows for commercial use.
138
+ - **Other questions?**
139
+ Feel free to reach out to us at [orbitals@principled-intelligence.com](mailto:orbitals@principled-intelligence.com)!
140
+
141
+ ### Contributing
142
+
143
+ We welcome contributions from the community! If you'd like to contribute to Orbitals, please check out our [Contributing Guide](CONTRIBUTING.md) for guidelines on how to get started.
144
+
145
+ ### License
146
+
147
+ This project is licensed under the Apache 2.0 License. See the [LICENSE](LICENSE) file for details.
148
+
149
+ ### Contact
150
+
151
+ For questions, feedback, or support, please reach out to us at [orbitals@principled-intelligence.com](mailto:orbitals@principled-intelligence.com).
152
+
153
+ ---
154
+
155
+ <div align="center">
156
+ <p>
157
+ <b>Built with ❤️ by <a href="https://principled-intelligence.com">Principled Intelligence</a></b>
158
+ <br />
159
+ Follow us on <a href="https://www.linkedin.com/company/principled-ai/">LinkedIn</a> for the latest updates.
160
+ </p>
161
+ </div>
@@ -0,0 +1,109 @@
1
+ <div align="center">
2
+ <img src="assets/orbitals-banner.png" width="70%" />
3
+ <h3 align="center">
4
+ <p>
5
+ <b>LLM Guardrails tailored to your Principles</b>
6
+ </p>
7
+ </h4>
8
+ </div>
9
+
10
+ <p align="center">
11
+ <img src="https://img.shields.io/pypi/v/orbitals?color=green" alt="PyPI Version">
12
+ <!-- <img src="https://img.shields.io/badge/type%20checked-ty-blue.svg?color=green" alt="Type Checked with ty"> -->
13
+ <img src="https://img.shields.io/pypi/pyversions/orbitals" alt="Python Versions">
14
+ <img src="https://img.shields.io/github/license/principled-intelligence/orbitals" alt="GitHub License">
15
+ </p>
16
+
17
+ ## Overview
18
+
19
+ **Orbitals** is a lightweight Python library for adding LLM guardrails in just a few lines of code. With Orbitals, you can add a governance layer tailored to **user-specific principles**. Rather than enforcing generic notions of safety, compliance, and correctness, Orbitals validates inputs (e.g., user requests) and outputs (e.g., assistant responses) against user-defined specifications and custom policies. This makes guardrails explicit, auditable, and aligned with the user's philosophy.
20
+
21
+ ### Key Features
22
+
23
+ - **User-defined specifications** — Guardrails that match your use case and your custom policies, not generic safety rules
24
+ - **Simple integration** — Add guardrails with minimal code changes
25
+ - **Open framework, open models** — Orbitals is open-source and is a simple interface for our open models
26
+
27
+ ## Getting started
28
+
29
+ ### Installation
30
+
31
+ You can install Orbitals via pip:
32
+
33
+ ```bash
34
+ pip install orbitals[all]
35
+ ```
36
+
37
+ ### Basic Usage
38
+
39
+ Here's a quick example to get you started with Orbitals, in which we use the ScopeGuard module to guard an AI service (for example, a customer support chatbot) from user requests that violate specified principles or fall outside of the scope of the core task of the assistant.
40
+
41
+ ```python
42
+ from orbitals.scope_guard import ScopeGuard
43
+
44
+ ai_service_description = "You are a helpful assistant for ..."
45
+ user_message = "Can I buy ..."
46
+
47
+ guardrail = ScopeGuard()
48
+ result = guardrail.validate(user_message, ai_service_description)
49
+ ```
50
+
51
+ The result of a guardrail validation will indicate whether the input or output passed the guardrail checks, along with details on any violations. You can then handle violations as needed, such as by rejecting the input or modifying the output. For example:
52
+
53
+ ```python
54
+ if result.scope_class.value == "Restricted" or result.scope_class.value == "Out of Scope":
55
+ print("Request violates guardrail:", result.evidences)
56
+ else:
57
+ # The user request is safe!
58
+ # We can now pass it to the AI assistant for processing.
59
+ ...
60
+ ```
61
+
62
+ ### Available Guardrails
63
+
64
+ Orbitals currently provides the following guardrail modules:
65
+
66
+ | Guardrail | Description | Hosting Options |
67
+ |:----------|:------------|:----------------|
68
+ | **[ScopeGuard](README.scope-guard.md)** | Classifies user queries against AI assistant specifications to detect out-of-scope requests, policy violations, and chit-chat | Self-hosted / Cloud hosting |
69
+ | 🚀 *Coming Soon* | More guardrails are on the way — stay tuned for updates! | — |
70
+
71
+ #### Hosting Options
72
+
73
+ - **Self-hosted**: Use open-weight models that you can deploy on your own infrastructure, ensuring data privacy and control.
74
+ - **Cloud hosting**: (Coming soon) Managed hosting options for ease of use and scalability
75
+
76
+ ### Documentation
77
+
78
+ For detailed documentation, including installation instructions, usage guides, and API references, please visit the Orbitals Documentation.
79
+
80
+ - [ScopeGuard Documentation](README.scope-guard.md)
81
+
82
+ ### FAQ
83
+
84
+ - **Can I use Orbitals for commercial applications?**
85
+ Yes, Orbitals is designed to be used in both research and commercial applications. It is licensed under the Apache 2.0 License, which allows for commercial use.
86
+ - **Other questions?**
87
+ Feel free to reach out to us at [orbitals@principled-intelligence.com](mailto:orbitals@principled-intelligence.com)!
88
+
89
+ ### Contributing
90
+
91
+ We welcome contributions from the community! If you'd like to contribute to Orbitals, please check out our [Contributing Guide](CONTRIBUTING.md) for guidelines on how to get started.
92
+
93
+ ### License
94
+
95
+ This project is licensed under the Apache 2.0 License. See the [LICENSE](LICENSE) file for details.
96
+
97
+ ### Contact
98
+
99
+ For questions, feedback, or support, please reach out to us at [orbitals@principled-intelligence.com](mailto:orbitals@principled-intelligence.com).
100
+
101
+ ---
102
+
103
+ <div align="center">
104
+ <p>
105
+ <b>Built with ❤️ by <a href="https://principled-intelligence.com">Principled Intelligence</a></b>
106
+ <br />
107
+ Follow us on <a href="https://www.linkedin.com/company/principled-ai/">LinkedIn</a> for the latest updates.
108
+ </p>
109
+ </div>
@@ -16,29 +16,28 @@
16
16
  </div>
17
17
 
18
18
  Given the specifications of an AI assistant and a user query, `scope-guard` maps the user query to one of the following five classes:
19
- * **Directly Supported**: The query is clearly within the assistant's capabilities.
20
- * **Potentially Supported**: The query could plausibly be handled by the assistant.
21
- * **Out of Scope**: The query is outside the assistant's defined role.
22
- * **Restricted**: The query cannot be handled due to a specific constraint.
23
- * **Chit Chat**: The query is a social interaction not related to the assistant's function.
19
+
20
+ * **Directly Supported**: The query is clearly within the assistant's capabilities.
21
+ * **Potentially Supported**: The query could plausibly be handled by the assistant.
22
+ * **Out of Scope**: The query is outside the assistant's defined role.
23
+ * **Restricted**: The query cannot be handled due to a specific constraint.
24
+ * **Chit Chat**: The query is a social interaction not related to the assistant's function.
24
25
 
25
26
  <p align="center">
26
27
  <img src="assets/scope-guard.svg" width="80%" />
27
28
  </p>
28
29
 
29
- To do this, `scope-guard` leverages specialized language models trained to evaluate whether a user query falls within an assistant’s intended scope. Different models are available, released with different modality flavors:
30
+ To do this, `scope-guard` leverages specialized language models trained to evaluate whether a user query falls within an assistant’s intended scope. Different models are available, released with different deployment options:
30
31
 
31
- | Flavor | Model | Parameters | Scope Classification | Custom Safety | Vanilla Safety | Tested GPUs |
32
+ | Model | Parameters | Hosting Options | Scope Classification | Custom Safety | Vanilla Safety | Tested GPUs |
32
33
  | :--- | :--- | :--- | :--- | :--- | :--- | :--- |
33
- | Open | scope-guard-4B-q-2601 | 4B | 89.1 | 79.1 | 90,3 | T4, L4, L40S, 4090, 5090 |
34
- | Open | scope-guard-4B-g-2601 | 4B | 90.1 | 78.0 | 88.0 | T4, L4, L40S, 4090, 5090 |
35
- | Closed | scope-guard-pro | ~ | 91.9 | 81.8 | 92.0 | ~ |
34
+ | scope-guard-4B-q-2601 | 4B | Self-hosted | 89.1 | 79.1 | 90.3 | T4, L4, L40S, 4090, 5090 |
35
+ | scope-guard-4B-g-2601 | 4B | Self-hosted | 90.1 | 78.0 | 88.0 | T4, L4, L40S, 4090, 5090 |
36
+ | scope-guard-pro | ~ | Cloud-only | 91.9 | 81.8 | 92.0 | ~ |
36
37
 
37
- ## Quickstart
38
+ ## Quickstart with our open models
38
39
 
39
- <details open>
40
- <summary>Open Flavor</summary>
41
- <br>
40
+ The easiest way to get started with ScopeGuard is to use our open models, which you can self-host on consumer-grade GPUs (e.g., Nvidia RTX series) and use via the `vllm` or `huggingface` backends. `scope-guard-4B-q-2601` and `scope-guard-4B-g-2601` are both 4B-parameter models so they can run even on a free instance on Google Colab with a T4 GPU.
42
41
 
43
42
  First, we need to install `orbitals` and `scope-guard`:
44
43
 
@@ -53,7 +52,7 @@ Then:
53
52
  ```python
54
53
  from orbitals.scope_guard import ScopeGuard
55
54
 
56
- scope_guard = ScopeGuard(
55
+ sg = ScopeGuard(
57
56
  backend="vllm",
58
57
  model="scope-guard-q", # for the Qwen-family model
59
58
  # model="scope-guard-g", # for the Gemma-family model
@@ -66,7 +65,7 @@ Never respond to requests for refunds.
66
65
  """
67
66
 
68
67
  user_query = "If the package hasn't arrived by tomorrow, can I get my money back?"
69
- result = scope_guard.validate(user_query, ai_service_description)
68
+ result = sg.validate(user_query, ai_service_description)
70
69
 
71
70
  print(f"Scope: {result.scope_class.value}")
72
71
  if result.evidences:
@@ -79,11 +78,9 @@ if result.evidences:
79
78
  # - Never respond to requests for refunds.
80
79
  ```
81
80
 
82
- </details>
81
+ ## Quickstart with our hosted models
83
82
 
84
- <details>
85
- <summary>Hosted Flavor</summary>
86
- <br>
83
+ ScopeGuard Pro is our most advanced model, available via API through our managed cloud hosting (get in touch with us if you are interested in on-premise deployment). It achieves best-in-class performance on scope classification tasks, including custom safety evaluations based on user-defined policies.
87
84
 
88
85
  > Open access to the hosted models is coming soon.
89
86
  > Need access earlier? Contact us at [orbitals@principled-intelligence.com](mailto:orbitals@principled-intelligence.com).
@@ -99,9 +96,9 @@ Then:
99
96
  ```python
100
97
  from orbitals.scope_guard import ScopeGuard
101
98
 
102
- scope_guard = ScopeGuard(
99
+ sg = ScopeGuard(
103
100
  backend="api",
104
- api_key="principled_1234",
101
+ api_key="principled_1234", # replace with your actual API key
105
102
  )
106
103
 
107
104
  ai_service_description = """
@@ -111,7 +108,7 @@ Never respond to requests for refunds.
111
108
  """
112
109
 
113
110
  user_query = "If the package hasn't arrived by tomorrow, can I get my money back?"
114
- result = scope_guard.validate(user_query, ai_service_description)
111
+ result = sg.validate(user_query, ai_service_description)
115
112
 
116
113
  print(f"Scope: {result.scope_class.value}")
117
114
  if result.evidences:
@@ -124,10 +121,62 @@ if result.evidences:
124
121
  # - Never respond to requests for refunds.
125
122
  ```
126
123
 
127
- </details>
128
-
129
124
  ## Usage
130
125
 
126
+ ### Initialization
127
+
128
+ Initialize the `ScopeGuard` object by specifying the backend and model you want to use.
129
+
130
+ If you are using the self-hosted models, you can choose between the `vllm` and `huggingface` backends:
131
+
132
+ ```python
133
+ from orbitals.scope_guard import ScopeGuard
134
+
135
+ sg = ScopeGuard(
136
+ model="scope-guard-q", # for the Qwen-family model
137
+ # model="scope-guard-g", # for the Gemma-family model
138
+ backend="vllm", # or "huggingface"
139
+ )
140
+ ```
141
+
142
+ If you are using the hosted models, use the `api` backend and provide your API key:
143
+
144
+ ```python
145
+ from orbitals.scope_guard import ScopeGuard
146
+
147
+ sg = ScopeGuard(
148
+ backend="api",
149
+ api_key="principled_1234", # replace with your actual API key
150
+ )
151
+ ```
152
+
153
+ ### Scope Classes
154
+
155
+ The possible scope classes returned by `scope-guard` are:
156
+
157
+ ```python
158
+ from orbitals.scope_guard import ScopeClass
159
+ print(ScopeClass.DIRECTLY_SUPPORTED.value) # "Directly Supported"
160
+ print(ScopeClass.POTENTIALLY_SUPPORTED.value) # "Potentially Supported"
161
+ print(ScopeClass.OUT_OF_SCOPE.value) # "Out of Scope"
162
+ print(ScopeClass.RESTRICTED.value) # "Restricted"
163
+ print(ScopeClass.CHIT_CHAT.value) # "Chit Chat"
164
+ ```
165
+
166
+ For example, you can check the scope class of a validation result as follows:
167
+
168
+ ```python
169
+ result = sg.validate(user_query, ai_service_description)
170
+
171
+ # Using the Enum member:
172
+ if result.scope_class == ScopeClass.RESTRICTED:
173
+ print("The user query is restricted.")
174
+
175
+ # Or using the string value:
176
+ if result.scope_class.value == "Restricted":
177
+ print("The user query is restricted.")
178
+ ```
179
+
131
180
  ### Input Formats
132
181
 
133
182
  The `validate` method is flexible and accepts various input formats for the conversation.
@@ -135,7 +184,7 @@ The `validate` method is flexible and accepts various input formats for the conv
135
184
  #### User query as a string
136
185
 
137
186
  ```python
138
- result = scope_guard.validate(
187
+ result = sg.validate(
139
188
  "When is my package scheduled to arrive?",
140
189
  ai_service_description
141
190
  )
@@ -143,9 +192,8 @@ result = scope_guard.validate(
143
192
 
144
193
  #### User query as a dictionary (OpenAI's API Message)
145
194
 
146
-
147
195
  ```python
148
- result = scope_guard.validate(
196
+ result = sg.validate(
149
197
  {
150
198
  "role": "user",
151
199
  "content": "When is my package scheduled to arrive?"
@@ -154,10 +202,10 @@ result = scope_guard.validate(
154
202
  )
155
203
  ```
156
204
 
157
- #### Conversation as a list of dictionaries
205
+ #### Conversation as a list of dictionaries
158
206
 
159
207
  ```python
160
- result = scope_guard.validate(
208
+ result = sg.validate(
161
209
  [
162
210
  {
163
211
  "role": "user",
@@ -197,7 +245,7 @@ queries = [
197
245
  "When is the package expected to be delivered?"
198
246
  ]
199
247
 
200
- result = scope_guard.batch_validate(
248
+ result = sg.batch_validate(
201
249
  queries,
202
250
  ai_service_description=ai_service_description
203
251
  )
@@ -211,19 +259,21 @@ ai_service_descriptions = [
211
259
  "You are a virtual assistant for a Courier. You answer questions about package tracking. Never respond to refund requests."
212
260
  ]
213
261
 
214
- result = scope_guard.batch_validate(
262
+ result = sg.batch_validate(
215
263
  queries,
216
264
  ai_service_descriptions=ai_service_descriptions
217
265
  )
218
266
  ```
219
267
 
220
- ## Serving
268
+ ## Serving ScopeGuard on-premise or on your infrastructure
221
269
 
222
270
  `scope-guard` comes with built-in support for serving. For better performance, it consists of two components:
271
+
223
272
  1. A **vLLM serving engine** that runs the model
224
273
  2. A **FastAPI server** that provides the end-to-end API interface, mapping input data to prompts, invoking the vLLM serving engine and returning the response to the user
225
274
 
226
275
  All of this is configured via the `orbitals scope-guard serve` command:
276
+
227
277
  ```bash
228
278
  # install the necessary packages
229
279
  pip install orbitals[scope-guard-serve]
@@ -233,6 +283,7 @@ orbitals scope-guard serve scope-guard --port 8000
233
283
  ```
234
284
 
235
285
  Alternatively, we also release a pre-built Docker image:
286
+
236
287
  ```bash
237
288
  docker run --runtime nvidia --gpus all \
238
289
  -p 8000:8000 \
@@ -277,14 +328,14 @@ Response:
277
328
  **Synchronous API client:**
278
329
 
279
330
  ```python
280
- from orbitals.guard import ScopeGuard
331
+ from orbitals.scope_guard import ScopeGuard
281
332
 
282
- scope_guard = ScopeGuard(
333
+ sg = ScopeGuard(
283
334
  backend="api",
284
335
  api_url="http://localhost:8000"
285
336
  )
286
337
 
287
- result = scope_guard.validate(
338
+ result = sg.validate(
288
339
  "If the package doesn't arrive by tomorrow, can I get my money back?",
289
340
  "You are a virtual assistant for a parcel delivery service. You can only answer questions about package tracking. Never respond to requests for refunds."
290
341
  )
@@ -293,14 +344,14 @@ result = scope_guard.validate(
293
344
  **Asynchronous API client:**
294
345
 
295
346
  ```python
296
- from orbitals.guard import AsyncScopeGuard
347
+ from orbitals.scope_guard import AsyncScopeGuard
297
348
 
298
- scope_guard = AsyncScopeGuard(
349
+ sg = AsyncScopeGuard(
299
350
  backend="api",
300
351
  api_url="http://localhost:8000"
301
352
  )
302
353
 
303
- result = await scope_guard.validate(
354
+ result = await sg.validate(
304
355
  "If the package doesn't arrive by tomorrow, can I get my money back?",
305
356
  "You are a virtual assistant for a parcel delivery service. You can only answer questions about package tracking. Never respond to requests for refunds."
306
357
  )
@@ -0,0 +1,27 @@
1
+ def main():
2
+ from orbitals.scope_guard import ScopeGuard
3
+
4
+ sg = ScopeGuard(
5
+ backend="vllm",
6
+ model="scope-guard-q", # for the Qwen-family model
7
+ # model="scope-guard-g", # for the Gemma-family model
8
+ )
9
+
10
+ ai_service_description = """
11
+ You are a virtual assistant for a parcel delivery service.
12
+ You can only answer questions about package tracking.
13
+ Never respond to requests for refunds.
14
+ """
15
+
16
+ user_query = "If the package hasn't arrived by tomorrow, can I get my money back?"
17
+ result = sg.validate(user_query, ai_service_description)
18
+
19
+ print(f"Scope: {result.scope_class.value}")
20
+ if result.evidences:
21
+ print("Evidences:")
22
+ for evidence in result.evidences:
23
+ print(f" - {evidence}")
24
+
25
+
26
+ if __name__ == "__main__":
27
+ main()
@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
4
4
 
5
5
  [project]
6
6
  name = "orbitals"
7
- version = "0.0.2"
7
+ version = "0.0.3"
8
8
  description = "LLM Guardrails tailored to your Principles"
9
9
  readme = "README.md"
10
10
 
@@ -29,8 +29,13 @@ classifiers = [
29
29
  dependencies = ["pydantic>=2.0.0", "typer>=0.12.3", "requests", "aiohttp"]
30
30
 
31
31
  [project.optional-dependencies]
32
- scope-guard-hf = ["transformers>=4.47.0,<5.0.0", "accelerate>=1.11.0"]
33
- scope-guard-vllm = ["transformers>=4.47.0,<5.0.0", "vllm>=0.11.0", "xgrammar"]
32
+ scope-guard-hf = ["transformers>=4.47.0,<5.0.0", "accelerate>=1.11.0", "pynvml"]
33
+ scope-guard-vllm = [
34
+ "transformers>=4.47.0,<5.0.0",
35
+ "vllm>=0.11.0",
36
+ "xgrammar",
37
+ "pynvml",
38
+ ]
34
39
  scope-guard-serve = [
35
40
  "orbitals[scope-guard-vllm]",
36
41
  "fastapi[standard]>=0.119.1",
@@ -0,0 +1,53 @@
1
+ import logging
2
+ import os
3
+
4
+
5
+ def maybe_configure_gpu_usage():
6
+ """
7
+ If the user hasn't explicitly set CUDA_VISIBLE_DEVICES, auto-configure it for
8
+ optimal usage: search for the gpu with the most free memory, and
9
+ set CUDA_VISIBLE_DEVICES to that GPU only.
10
+
11
+ Uses pynvml to avoid triggering CUDA initialization from torch.
12
+ """
13
+ if "CUDA_VISIBLE_DEVICES" in os.environ:
14
+ logging.info(
15
+ "CUDA_VISIBLE_DEVICES is already set, not auto-configuring GPU usage"
16
+ )
17
+ return
18
+
19
+ try:
20
+ import pynvml # ty: ignore[unresolved-import]
21
+ except ModuleNotFoundError:
22
+ logging.debug("pynvml not available, skipping GPU auto-configuration")
23
+ return
24
+
25
+ try:
26
+ pynvml.nvmlInit()
27
+ except pynvml.NVMLError:
28
+ logging.debug("NVML initialization failed, skipping GPU auto-configuration")
29
+ return
30
+
31
+ try:
32
+ device_count = pynvml.nvmlDeviceGetCount()
33
+ if device_count == 0:
34
+ return
35
+
36
+ best_idx = None
37
+ best_free = -1
38
+
39
+ for i in range(device_count):
40
+ handle = pynvml.nvmlDeviceGetHandleByIndex(i)
41
+ mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
42
+ if mem_info.free > best_free:
43
+ best_idx = i
44
+ best_free = mem_info.free
45
+
46
+ if best_idx is not None:
47
+ if device_count > 1:
48
+ logging.warning(
49
+ f"Auto-configuring to use GPU {best_idx} with {best_free / 1024**3:.2f} GB free"
50
+ )
51
+ os.environ["CUDA_VISIBLE_DEVICES"] = str(best_idx)
52
+ finally:
53
+ pynvml.nvmlShutdown()
@@ -1,2 +0,0 @@
1
- .venv
2
- examples/
orbitals-0.0.2/PKG-INFO DELETED
@@ -1,124 +0,0 @@
1
- Metadata-Version: 2.4
2
- Name: orbitals
3
- Version: 0.0.2
4
- Summary: LLM Guardrails tailored to your Principles
5
- Author-email: Luigi Procopio <luigi@principled-intelligence.com>, Edoardo Barba <edoardo@principled-intelligence.com>
6
- License: Apache-2.0
7
- License-File: LICENSE
8
- Classifier: Programming Language :: Python :: 3 :: Only
9
- Classifier: Programming Language :: Python :: 3.10
10
- Classifier: Programming Language :: Python :: 3.11
11
- Classifier: Programming Language :: Python :: 3.12
12
- Classifier: Programming Language :: Python :: 3.13
13
- Classifier: Programming Language :: Python :: 3.14
14
- Requires-Python: >=3.10
15
- Requires-Dist: aiohttp
16
- Requires-Dist: pydantic>=2.0.0
17
- Requires-Dist: requests
18
- Requires-Dist: typer>=0.12.3
19
- Provides-Extra: all
20
- Requires-Dist: accelerate>=1.11.0; extra == 'all'
21
- Requires-Dist: fastapi[standard]>=0.119.1; extra == 'all'
22
- Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'all'
23
- Requires-Dist: uvicorn>=0.29.0; extra == 'all'
24
- Requires-Dist: vllm>=0.11.0; extra == 'all'
25
- Requires-Dist: xgrammar; extra == 'all'
26
- Provides-Extra: scope-guard-all
27
- Requires-Dist: accelerate>=1.11.0; extra == 'scope-guard-all'
28
- Requires-Dist: fastapi[standard]>=0.119.1; extra == 'scope-guard-all'
29
- Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-all'
30
- Requires-Dist: uvicorn>=0.29.0; extra == 'scope-guard-all'
31
- Requires-Dist: vllm>=0.11.0; extra == 'scope-guard-all'
32
- Requires-Dist: xgrammar; extra == 'scope-guard-all'
33
- Provides-Extra: scope-guard-hf
34
- Requires-Dist: accelerate>=1.11.0; extra == 'scope-guard-hf'
35
- Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-hf'
36
- Provides-Extra: scope-guard-serve
37
- Requires-Dist: fastapi[standard]>=0.119.1; extra == 'scope-guard-serve'
38
- Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-serve'
39
- Requires-Dist: uvicorn>=0.29.0; extra == 'scope-guard-serve'
40
- Requires-Dist: vllm>=0.11.0; extra == 'scope-guard-serve'
41
- Requires-Dist: xgrammar; extra == 'scope-guard-serve'
42
- Provides-Extra: scope-guard-vllm
43
- Requires-Dist: transformers<5.0.0,>=4.47.0; extra == 'scope-guard-vllm'
44
- Requires-Dist: vllm>=0.11.0; extra == 'scope-guard-vllm'
45
- Requires-Dist: xgrammar; extra == 'scope-guard-vllm'
46
- Description-Content-Type: text/markdown
47
-
48
- <div align="center">
49
- <img src="assets/orbitals-banner.png" width="70%" />
50
- <h3 align="center">
51
- <p>
52
- <b>LLM Guardrails tailored to your Principles</b>
53
- </p>
54
- </h4>
55
- <hr/>
56
- </div>
57
-
58
- <p align="center">
59
- <img src="https://img.shields.io/badge/type%20checked-ty-blue.svg?color=green" alt="Type Checked with ty">
60
- <img src="https://img.shields.io/pypi/v/orbitals?color=green" alt="PyPI Version">
61
- <img src="https://img.shields.io/github/license/principled-intelligence/orbitals" alt="GitHub License">
62
- <img src="https://img.shields.io/pypi/pyversions/orbitals" alt="Python Versions">
63
- <img src="https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Fprincipled-intelligence%2Forbitals%2Fmain%2Fpyproject.toml" alt="Required Python Version">
64
- </p>
65
-
66
- `orbitals` is an ecosystem of LLM guardrails, designed to provide a governance layer tailored to **user-specific principles, requirements and use cases**. Rather than enforcing generic notions of safety, correctness, etc., Orbitals evaluates inputs and outputs against *user-defined specifications*. This makes guardrails explicit, auditable, and aligned with the user's philosophy.
67
-
68
- Orbitals guardrails fall into two typologies:
69
- - **Guards** operate on the *input* of a guardrailed LLM, assessing whether a user request is legitimate under the provided specifications.
70
- - **Supervisors** operate on the *output* of a guardrailed LLM, evaluating the assistant’s response before it is returned.
71
-
72
- Guardrails may be released under different modality flavors:
73
- - **Open** (open-source and open-weight), allowing users to run guardrails and underlying models on their own infrastructure.
74
- - **Hosted**, accessible via simple HTTP calls (API key required).
75
-
76
- ## Available Guardrails
77
-
78
- | Name | Flavor | Description |
79
- |-------------|-------------------------|-----------------------------------------------------------------------------|
80
- | [ScopeGuard](README.scope-guard.md) | Open / Hosted | Validates whether a user request falls within the intended use of an AI service. |
81
- | RagSupervisor | Coming soon | Ensures LLM responses remain grounded in retrieved context for RAG setups. |
82
-
83
- <details>
84
- <summary>ScopeGuard</summary>
85
- <br>
86
-
87
- First, we need to install `orbitals` and `scope-guard`:
88
-
89
- ```bash
90
- pip install orbitals[scope-guard-vllm]
91
- ```
92
-
93
- Then:
94
-
95
- ```python
96
- from orbitals.scope_guard import ScopeGuard
97
-
98
- scope_guard = ScopeGuard(
99
- backend="vllm",
100
- model="scope-guard-q", # for the Qwen-family model
101
- # model="scope-guard-g", # for the Gemma-family model
102
- )
103
-
104
- ai_service_description = """
105
- You are a virtual assistant for a parcel delivery service.
106
- You can only answer questions about package tracking.
107
- Never respond to requests for refunds.
108
- """
109
-
110
- user_query = "If the package hasn't arrived by tomorrow, can I get my money back?"
111
- result = scope_guard.validate(user_query, ai_service_description)
112
-
113
- print(f"Scope: {result.scope_class.value}")
114
- if result.evidences:
115
- print("Evidences:")
116
- for evidence in result.evidences:
117
- print(f" - {evidence}")
118
-
119
- # Scope: Restricted
120
- # Evidences:
121
- # - Never respond to requests for refunds.
122
- ```
123
-
124
- </details>
orbitals-0.0.2/README.md DELETED
@@ -1,77 +0,0 @@
1
- <div align="center">
2
- <img src="assets/orbitals-banner.png" width="70%" />
3
- <h3 align="center">
4
- <p>
5
- <b>LLM Guardrails tailored to your Principles</b>
6
- </p>
7
- </h4>
8
- <hr/>
9
- </div>
10
-
11
- <p align="center">
12
- <img src="https://img.shields.io/badge/type%20checked-ty-blue.svg?color=green" alt="Type Checked with ty">
13
- <img src="https://img.shields.io/pypi/v/orbitals?color=green" alt="PyPI Version">
14
- <img src="https://img.shields.io/github/license/principled-intelligence/orbitals" alt="GitHub License">
15
- <img src="https://img.shields.io/pypi/pyversions/orbitals" alt="Python Versions">
16
- <img src="https://img.shields.io/python/required-version-toml?tomlFilePath=https%3A%2F%2Fraw.githubusercontent.com%2Fprincipled-intelligence%2Forbitals%2Fmain%2Fpyproject.toml" alt="Required Python Version">
17
- </p>
18
-
19
- `orbitals` is an ecosystem of LLM guardrails, designed to provide a governance layer tailored to **user-specific principles, requirements and use cases**. Rather than enforcing generic notions of safety, correctness, etc., Orbitals evaluates inputs and outputs against *user-defined specifications*. This makes guardrails explicit, auditable, and aligned with the user's philosophy.
20
-
21
- Orbitals guardrails fall into two typologies:
22
- - **Guards** operate on the *input* of a guardrailed LLM, assessing whether a user request is legitimate under the provided specifications.
23
- - **Supervisors** operate on the *output* of a guardrailed LLM, evaluating the assistant’s response before it is returned.
24
-
25
- Guardrails may be released under different modality flavors:
26
- - **Open** (open-source and open-weight), allowing users to run guardrails and underlying models on their own infrastructure.
27
- - **Hosted**, accessible via simple HTTP calls (API key required).
28
-
29
- ## Available Guardrails
30
-
31
- | Name | Flavor | Description |
32
- |-------------|-------------------------|-----------------------------------------------------------------------------|
33
- | [ScopeGuard](README.scope-guard.md) | Open / Hosted | Validates whether a user request falls within the intended use of an AI service. |
34
- | RagSupervisor | Coming soon | Ensures LLM responses remain grounded in retrieved context for RAG setups. |
35
-
36
- <details>
37
- <summary>ScopeGuard</summary>
38
- <br>
39
-
40
- First, we need to install `orbitals` and `scope-guard`:
41
-
42
- ```bash
43
- pip install orbitals[scope-guard-vllm]
44
- ```
45
-
46
- Then:
47
-
48
- ```python
49
- from orbitals.scope_guard import ScopeGuard
50
-
51
- scope_guard = ScopeGuard(
52
- backend="vllm",
53
- model="scope-guard-q", # for the Qwen-family model
54
- # model="scope-guard-g", # for the Gemma-family model
55
- )
56
-
57
- ai_service_description = """
58
- You are a virtual assistant for a parcel delivery service.
59
- You can only answer questions about package tracking.
60
- Never respond to requests for refunds.
61
- """
62
-
63
- user_query = "If the package hasn't arrived by tomorrow, can I get my money back?"
64
- result = scope_guard.validate(user_query, ai_service_description)
65
-
66
- print(f"Scope: {result.scope_class.value}")
67
- if result.evidences:
68
- print("Evidences:")
69
- for evidence in result.evidences:
70
- print(f" - {evidence}")
71
-
72
- # Scope: Restricted
73
- # Evidences:
74
- # - Never respond to requests for refunds.
75
- ```
76
-
77
- </details>
@@ -1,47 +0,0 @@
1
- import logging
2
- import os
3
- from typing import TYPE_CHECKING
4
-
5
- if TYPE_CHECKING:
6
- import torch
7
- else:
8
- try:
9
- import torch # ty: ignore[unresolved-import]
10
- except ModuleNotFoundError:
11
- torch = None # ty: ignore[invalid-assignment]
12
-
13
-
14
- def maybe_configure_gpu_usage():
15
- """
16
- If the user hasn't explicitly set CUDA_VISIBLE_DEVICES, auto-configure it for
17
- optimal usage: search for the gpu with the most free memory, and
18
- set CUDA_VISIBLE_DEVICES to that GPU only.
19
- """
20
- if torch is None:
21
- return
22
-
23
- if "CUDA_VISIBLE_DEVICES" in os.environ:
24
- logging.info(
25
- "CUDA_VISIBLE_DEVICES is already set, not auto-configuring GPU usage"
26
- )
27
- return
28
-
29
- if not torch.cuda.is_available():
30
- return
31
-
32
- best_idx = None
33
- best_free = -1
34
- is_multi_gpu = torch.cuda.device_count() > 1
35
-
36
- for i in range(torch.cuda.device_count()):
37
- free_bytes, _ = torch.cuda.mem_get_info(i) # (free, total)
38
- if free_bytes > best_free:
39
- best_idx = i
40
- best_free = free_bytes
41
-
42
- if best_idx is not None:
43
- if is_multi_gpu:
44
- logging.warning(
45
- f"Auto-configuring to use GPU {best_idx} with {best_free / 1024**3:.2f} GB free"
46
- )
47
- os.environ["CUDA_VISIBLE_DEVICES"] = str(best_idx)
File without changes
File without changes
File without changes
File without changes
File without changes
File without changes