From 730d61cfe38a5b7b17a6909977e7b3c4b265c234 Mon Sep 17 00:00:00 2001
From: leduc <leduc@MBP14.local>
Date: Sat, 15 Mar 2025 11:27:46 +0100
Subject: [PATCH] ollama automatic pull

---
 .DS_Store | Bin 6148 -> 6148 bytes
 README.md |   8 ++++++-
 main.py   |  64 +++++++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 70 insertions(+), 2 deletions(-)

diff --git a/.DS_Store b/.DS_Store
index 844a1869da8f4dd6ca6f4d6caba5ece61e3e4835..099c148775b0e005cda257c85038e6ab4d9347d3 100644
GIT binary patch
delta 112
zcmZoMXffEJ$`a2h8qL7Kz`~%%kj{|FP?DSP;*yk;p9B=+xc^DNr18)(M^yO~xO~rz
h4Qc4|1sR6H$@#ejKs^i$JPr(#^H|k3uVRT31^|Cx9fbe@

delta 112
zcmZoMXffEJ$`a3bwx5B4frUYjA)O(Up(Hoo#U&{xKM5$tq3^7v9Jc$IBdUA~UipFy
a!{Frn+ybB;ps5WElk-^BHm_od5(WVI!5Yl~

diff --git a/README.md b/README.md
index d07f751..e879411 100644
--- a/README.md
+++ b/README.md
@@ -1,3 +1,5 @@
+![Model Comparison](benchmark_results/model_comparison.png)
+
 # Codebench - Ollama Model Benchmark Tool
 
 A Python-based benchmarking tool for testing and comparing different Ollama models on coding tasks. This tool allows you to benchmark multiple Ollama models against common coding problems, measure their performance, and visualize the results.
@@ -13,6 +15,7 @@ A Python-based benchmarking tool for testing and comparing different Ollama mode
 - Measure performance metrics (tokens/sec, response time)
 - Track success rates across different coding challenges
 - Support for local and remote Ollama servers
+- Automatic model download if not available locally
 - Detailed test results and leaderboard generation
 - CPU information tracking for benchmarks
 
@@ -61,7 +64,7 @@ python main.py --server [local|z60] --model [model_name] --number [count|all] --
 ## Arguments:
 
 - --server : Choose Ollama server (default: local)
-- --model : Test specific model only
+- --model : Test specific model only (will be automatically downloaded if not available locally)
 - --number : Number of models to test
 - --verbose : Enable detailed output
 - --plot-only : Skip benchmarking and just generate graphs from existing results
@@ -125,6 +128,9 @@ By default, main.py will now automatically generate graphs after benchmarking. Y
 # Run benchmarks and generate graphs (default behavior)
 python3 main.py
 
+# Test a specific model (will be downloaded automatically if not available locally)
+python3 main.py --model llama3
+
 # Skip benchmarking and just generate graphs from the latest results
 python3 main.py --plot-only
 
diff --git a/main.py b/main.py
index a357ce0..3a0e30a 100755
--- a/main.py
+++ b/main.py
@@ -591,6 +591,53 @@ def get_available_models(server_url: str) -> List[str]:
         print(f"{RED}Error getting model list from {server_url}: {e}{ENDC}")
         return []
 
+def check_model_exists_locally(model_name: str, server_url: str) -> bool:
+    """Check if a model exists locally on the specified Ollama server."""
+    available_models = get_available_models(server_url)
+    return model_name in available_models
+
+def download_model(model_name: str) -> bool:
+    """Download a model using ollama pull command.
+    
+    Args:
+        model_name: Name of the model to download
+        
+    Returns:
+        bool: True if download was successful, False otherwise
+    """
+    print(f"\n{INFO}Model '{model_name}' not found locally. Downloading...{ENDC}")
+    try:
+        # Run ollama pull command and capture output
+        process = subprocess.Popen(
+            ["ollama", "pull", model_name],
+            stdout=subprocess.PIPE,
+            stderr=subprocess.PIPE,
+            text=True,
+            bufsize=1
+        )
+        
+        # Print output in real-time
+        print(f"{INFO}Download progress:{ENDC}")
+        while True:
+            output = process.stdout.readline()
+            if output == '' and process.poll() is not None:
+                break
+            if output:
+                print(f"{MUTED}{output.strip()}{ENDC}")
+                
+        # Check if download was successful
+        return_code = process.poll()
+        if return_code == 0:
+            print(f"\n{SUCCESS}Successfully downloaded model '{model_name}'.{ENDC}")
+            return True
+        else:
+            error = process.stderr.read()
+            print(f"\n{ERROR}Failed to download model '{model_name}': {error}{ENDC}")
+            return False
+    except Exception as e:
+        print(f"\n{ERROR}Error downloading model '{model_name}': {e}{ENDC}")
+        return False
+
 def get_model_details(model_name):
     try:
         result = subprocess.run(
@@ -961,7 +1008,22 @@ def main():
 
     # Get available models or use specified model
     if args.model:
-        models = [args.model]
+        # Check if the specified model exists locally
+        if not check_model_exists_locally(args.model, server_url):
+            # If not, try to download it
+            if download_model(args.model):
+                # Verify the model is now available
+                if check_model_exists_locally(args.model, server_url):
+                    models = [args.model]
+                else:
+                    print(f"{WARNING}Model '{args.model}' was downloaded but not found on server. Please check manually.{ENDC}")
+                    return
+            else:
+                print(f"{RED}Could not download model '{args.model}'. Exiting.{ENDC}")
+                return
+        else:
+            print(f"{SUCCESS}Using locally available model: {args.model}{ENDC}")
+            models = [args.model]
     else:
         models = get_available_models(server_url)