Update scripts/download_model.py

scripts/download_model.py

@@ -5,35 +5,71 @@ Script to download the Phi-3.5-mini-instruct model for local development
 """
 
 import os
-from huggingface_hub import hf_hub_download
+import sys
+from huggingface_hub import hf_hub_download, snapshot_download
 import argparse
+import requests
+import hashlib
+import time
 
-def download_model(model_repo: str, model_file: str, local_dir: str = "./models"):
-    """Download the model from Hugging Face Hub"""
-    
-    print(f"📥 Downloading model from {model_repo}")
-    print(f"📄 Model file: {model_file}")
-    print(f"💾 Local directory: {local_dir}")
-    
-    # Create local directory if it doesn't exist
-    os.makedirs(local_dir, exist_ok=True)
+def calculate_file_hash(filepath):
+    """Calculate MD5 hash of a file"""
+    hash_md5 = hashlib.md5()
+    with open(filepath, "rb") as f:
+        for chunk in iter(lambda: f.read(4096), b""):
+            hash_md5.update(chunk)
+    return hash_md5.hexdigest()
+
+def download_with_progress(repo_id, filename, local_dir):
+    """Download file with progress tracking"""
+    print(f"📥 Downloading {filename} from {repo_id}")
     
     try:
-        # Download the model
+        # Download with progress
         model_path = hf_hub_download(
-            repo_id=model_repo,
-            filename=model_file,
+            repo_id=repo_id,
+            filename=filename,
             local_dir=local_dir,
-            local_dir_use_symlinks=False
+            local_dir_use_symlinks=False,
+            resume_download=True
         )
         
-        print(f"✅ Model downloaded successfully: {model_path}")
-        return model_path
-        
+        # Verify download
+        if os.path.exists(model_path):
+            file_size = os.path.getsize(model_path) / (1024 * 1024 * 1024)  # GB
+            file_hash = calculate_file_hash(model_path)
+            
+            print(f"✅ Download completed successfully!")
+            print(f"📁 File: {model_path}")
+            print(f"📊 Size: {file_size:.2f} GB")
+            print(f"🔒 Hash: {file_hash}")
+            
+            return model_path
+        else:
+            print(f"❌ Downloaded file not found: {model_path}")
+            return None
+            
     except Exception as e:
         print(f"❌ Error downloading model: {e}")
         return None
 
+def check_disk_space(required_gb=3):
+    """Check if there's enough disk space"""
+    try:
+        stat = os.statvfs('/')
+        free_gb = (stat.f_bavail * stat.f_frsize) / (1024 ** 3)
+        
+        print(f"💾 Available disk space: {free_gb:.1f} GB")
+        print(f"💾 Required disk space: {required_gb} GB")
+        
+        if free_gb < required_gb:
+            print(f"❌ Not enough disk space. Need {required_gb} GB, have {free_gb:.1f} GB")
+            return False
+        return True
+    except:
+        print("⚠️  Could not check disk space")
+        return True
+
 def main():
     parser = argparse.ArgumentParser(description="Download Phi-3.5-mini-instruct model")
     parser.add_argument("--repo", default="Thetima4/Phi-3.5-mini-instruct-Q4_K_M-GGUF",
@@ -42,13 +78,53 @@ def main():
                        help="Model filename")
     parser.add_argument("--dir", default="./models",
                        help="Local directory to save the model")
+    parser.add_argument("--force", action="store_true",
+                       help="Force download even if file exists")
     
     args = parser.parse_args()
     
     print("🚀 Saem's Tunes AI Model Downloader")
     print("=" * 50)
     
-    download_model(args.repo, args.file, args.dir)
+    # Check disk space
+    if not check_disk_space(3):
+        sys.exit(1)
+    
+    # Create local directory if it doesn't exist
+    os.makedirs(args.dir, exist_ok=True)
+    
+    # Check if file already exists
+    local_path = os.path.join(args.dir, args.file)
+    if os.path.exists(local_path) and not args.force:
+        print(f"✅ Model already exists: {local_path}")
+        file_size = os.path.getsize(local_path) / (1024 * 1024 * 1024)
+        print(f"📊 Size: {file_size:.2f} GB")
+        return local_path
+    
+    # Download the model
+    start_time = time.time()
+    model_path = download_with_progress(args.repo, args.file, args.dir)
+    download_time = time.time() - start_time
+    
+    if model_path:
+        print(f"⏱️  Download time: {download_time:.1f} seconds")
+        print(f"🎉 Model ready for use!")
+        
+        # Create model info file
+        info_file = os.path.join(args.dir, "model_info.txt")
+        with open(info_file, 'w') as f:
+            f.write(f"Model: {args.repo}\n")
+            f.write(f"File: {args.file}\n")
+            f.write(f"Downloaded: {time.ctime()}\n")
+            f.write(f"Path: {model_path}\n")
+            f.write(f"Size: {os.path.getsize(model_path) / (1024**3):.2f} GB\n")
+        
+        print(f"📄 Model info saved to: {info_file}")
+    else:
+        print("❌ Model download failed")
+        sys.exit(1)
+    
+    return model_path
 
 if __name__ == "__main__":
     main()