File size: 3,656 Bytes
d572bbd | 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 109 110 111 112 113 114 115 116 117 118 119 120 121 | #!/usr/bin/env python3
"""
Google Colab Environment Setup Script for EasyTranslate.
This script is designed to be run as the first cell in a Colab notebook.
It handles:
1. GPU verification and CUDA setup
2. Repository cloning from GitHub
3. Dependency installation
4. Google Drive mounting
5. Environment variable configuration
Usage in Colab:
!wget -q https://raw.githubusercontent.com/your-org/UCAS-EasyTranslate/main/scripts/setup_colab.py
%run setup_colab.py
"""
import os
import subprocess
import sys
from pathlib import Path
REPO_URL = "https://github.com/your-org/UCAS-EasyTranslate.git"
REPO_DIR = Path("/content/UCAS-EasyTranslate")
DRIVE_MOUNT_POINT = "/content/drive"
DRIVE_BASE = "/content/drive/MyDrive/EasyTranslate"
def check_gpu():
try:
import torch
if torch.cuda.is_available():
print(f"[OK] GPU detected: {torch.cuda.get_device_name(0)}")
print(f" CUDA version: {torch.version.cuda}")
print(f" GPU memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB")
return True
else:
print("[WARN] No GPU detected. Training will be very slow.")
return False
except ImportError:
print("[WARN] PyTorch not found. Installing...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "torch"])
return check_gpu()
def clone_repository():
if REPO_DIR.exists():
print(f"[INFO] Repository exists at {REPO_DIR}, pulling latest...")
os.chdir(REPO_DIR)
subprocess.check_call(["git", "pull", "origin", "main"])
else:
print(f"[INFO] Cloning repository from {REPO_URL}...")
subprocess.check_call(["git", "clone", REPO_URL, str(REPO_DIR)])
os.chdir(REPO_DIR)
sys.path.insert(0, str(REPO_DIR / "src"))
print(f"[OK] Repository ready at {REPO_DIR}")
def install_dependencies():
req_file = REPO_DIR / "requirements-colab.txt"
if req_file.exists():
print("[INFO] Installing Colab dependencies...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-r", str(req_file)])
else:
print("[INFO] Installing from requirements.txt...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-r", str(REPO_DIR / "requirements.txt")])
print("[INFO] Installing EasyTranslate package...")
subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "-e", str(REPO_DIR)])
print("[OK] All dependencies installed")
def mount_drive():
try:
from google.colab import drive
drive.mount(DRIVE_MOUNT_POINT)
if os.path.exists(DRIVE_MOUNT_POINT):
os.makedirs(DRIVE_BASE, exist_ok=True)
print(f"[OK] Google Drive mounted at {DRIVE_MOUNT_POINT}")
print(f" Base path: {DRIVE_BASE}")
return True
except Exception as e:
print(f"[WARN] Google Drive mount failed: {e}")
return False
def setup_environment():
os.environ["TOKENIZERS_PARALLELISM"] = "false"
os.environ["WANDB_MODE"] = os.environ.get("WANDB_MODE", "offline")
os.environ["PYTHONHASHSEED"] = "42"
print("[OK] Environment variables configured")
def main():
print("=" * 60)
print(" EasyTranslate — Colab Environment Setup")
print("=" * 60)
print()
check_gpu()
print()
clone_repository()
print()
install_dependencies()
print()
mount_drive()
print()
setup_environment()
print()
print("=" * 60)
print(" Setup Complete! Ready for training.")
print("=" * 60)
if __name__ == "__main__":
main() |