41 lines
1020 B
Python
Executable File
41 lines
1020 B
Python
Executable File
"""Emit the install command for the Axolotl cut-cross-entropy fork."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import shutil
|
|
import sys
|
|
|
|
try:
|
|
import torch
|
|
except ImportError as exc: # pragma: no cover - defensive
|
|
raise ImportError("Install torch via `pip install torch`") from exc
|
|
|
|
from packaging.version import Version as V
|
|
|
|
USE_UV_FLAG = "--uv" in sys.argv[1:]
|
|
USE_PIP_FLAG = "--pip" in sys.argv[1:]
|
|
|
|
if USE_UV_FLAG and USE_PIP_FLAG:
|
|
raise SystemExit("Specify only one of --uv or --pip")
|
|
|
|
if USE_PIP_FLAG:
|
|
use_uv = False
|
|
elif USE_UV_FLAG:
|
|
use_uv = True
|
|
else:
|
|
use_uv = shutil.which("uv") is not None
|
|
|
|
if V(torch.__version__) < V("2.4.0"):
|
|
print("")
|
|
sys.exit(0)
|
|
|
|
# No need to uninstall in CI runs; the environment is fresh. Just emit the install command.
|
|
installer = "uv pip install --system" if use_uv else "pip install"
|
|
command = (
|
|
f"{installer} "
|
|
'"cut-cross-entropy[transformers] '
|
|
'@ git+https://github.com/axolotl-ai-cloud/ml-cross-entropy.git@147ea28"'
|
|
)
|
|
|
|
print(command)
|