Upload folder using huggingface_hub

requirements-base.txt

@@ -0,0 +1,8 @@
+torch>=2.0.0
+accelerate>=0.27.0
+bitsandbytes>=0.41.0
+datasets>=2.15.0
+gradio>=5.17.0
+huggingface-hub>=0.19.0
+tensorboard>=2.15.0
+transformers>=4.36.0

requirements-flash.txt

@@ -0,0 +1,2 @@
+-r requirements-base.txt
+flash-attn==2.5.2

requirements.txt

@@ -22,4 +22,3 @@ tqdm>=4.65.0
 transformers>=4.36.0
 typing-extensions>=4.8.0
 unsloth>=2024.3
-flash-attn==2.5.2

update_space.py

@@ -97,10 +97,12 @@ def verify_configs():
             raise ValueError(f"Invalid JSON in {json_file}: {e}")
 
 def update_requirements():
-    """Update requirements.txt with necessary packages."""
+    """Update requirements.txt with necessary packages using a two-stage installation process."""
     current_dir = Path(__file__).parent
-    req_path = current_dir / "requirements.txt"
+    base_req_path = current_dir / "requirements-base.txt"
+    flash_req_path = current_dir / "requirements-flash.txt"
     
+    # First ensure base requirements exist
     required_packages = {
         "torch>=2.0.0",
         "transformers>=4.36.0",
@@ -112,34 +114,34 @@ def update_requirements():
         "datasets>=2.15.0"
     }
     
-    # Read existing requirements
+    # Read existing base requirements
     existing_requirements = set()
-    if req_path.exists():
-        with open(req_path) as f:
-            existing_requirements = {line.strip() for line in f if line.strip()}
+    if base_req_path.exists():
+        with open(base_req_path) as f:
+            existing_requirements = {line.strip() for line in f if line.strip() and not line.startswith('-r')}
     
     # Add new requirements
     updated_requirements = existing_requirements.union(required_packages)
     
-    # Write updated requirements with torch first and flash-attn last
-    with open(req_path, 'w') as f:
+    # Write updated base requirements
+    with open(base_req_path, 'w') as f:
         # Ensure torch is first
         torch_req = next((req for req in updated_requirements if req.startswith("torch")), "torch>=2.0.0")
         f.write(f"{torch_req}\n")
         
-        # Extract flash-attn to add it last
-        flash_attn_req = next((req for req in updated_requirements if req.startswith("flash-attn")), None)
-        
-        # Write all other requirements (excluding torch and flash-attn)
-        for req in sorted(r for r in updated_requirements 
-                         if not r.startswith("torch") and not r.startswith("flash-attn")):
+        # Write all other requirements (excluding torch)
+        for req in sorted(r for r in updated_requirements if not r.startswith("torch")):
             f.write(f"{req}\n")
-        
-        # Add flash-attn as the very last package
-        if flash_attn_req:
-            f.write(f"{flash_attn_req}\n")
     
-    logger.info("Updated requirements.txt with torch listed first and flash-attn listed last")
+    # Create or update flash-attn requirements
+    with open(flash_req_path, 'w') as f:
+        f.write("-r requirements-base.txt\n")
+        f.write("flash-attn==2.5.2\n")
+    
+    logger.info("Updated requirements files for two-stage installation:")
+    logger.info(f"1. Base requirements in {base_req_path}")
+    logger.info(f"2. Flash-attention requirements in {flash_req_path}")
+    logger.info("This ensures torch is installed before flash-attn")
 
 def create_space(username, space_name):
     """Create or get a Hugging Face Space."""