File size: 2,535 Bytes
5ce0678 | 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 | #!/usr/bin/env python3
"""
Setup script to prepare OpenEnv model for supply chain environment.
Usage: python setup_model.py [--model-id <model_id>] [--hf-token <token>]
"""
import argparse
import subprocess
import sys
from pathlib import Path
MODELS_DIR = Path(__file__).parent
def install_dependencies():
"""Install required packages for model loading."""
deps = [
"huggingface_hub",
"transformers",
"torch",
]
print("[*] Installing model dependencies...")
for dep in deps:
try:
__import__(dep.replace("-", "_"))
print(f" ✓ {dep} already installed")
except ImportError:
print(f" → Installing {dep}...")
subprocess.check_call([sys.executable, "-m", "pip", "install", dep, "-q"])
print(f" ✓ {dep} installed")
def download_model(model_id: str, hf_token: str = None):
"""Download model from Hugging Face."""
from model_loader import get_model_path
print(f"\n[*] Setting up model: {model_id}")
print(f" Cache dir: {MODELS_DIR / '.cache'}")
try:
path = get_model_path(model_id, hf_token)
print(f" ✓ Model ready at: {path}")
return True
except Exception as e:
print(f" ✗ Error: {e}")
return False
def main():
parser = argparse.ArgumentParser(
description="Setup OpenEnv model for supply chain environment"
)
parser.add_argument(
"--model-id",
default="openenv/supply-chain-model",
help="Hugging Face model ID (default: openenv/supply-chain-model)",
)
parser.add_argument(
"--hf-token",
help="Hugging Face API token (or set HF_TOKEN env var)",
)
parser.add_argument(
"--skip-deps",
action="store_true",
help="Skip installing dependencies",
)
args = parser.parse_args()
print("=" * 60)
print("OpenEnv Model Setup for Supply Chain Disruption Management")
print("=" * 60)
if not args.skip_deps:
install_dependencies()
success = download_model(args.model_id, args.hf_token)
print("\n" + "=" * 60)
if success:
print("✓ Setup complete! Model is ready to use.")
print("\nTo use the local model in inference.py, set:")
print(" export USE_LOCAL_MODEL=true")
print(f" export LOCAL_MODEL_ID={args.model_id}")
else:
print("✗ Setup failed. Check your Hugging Face token and model ID.")
sys.exit(1)
if __name__ == "__main__":
main()
|