Overloads download method to specify progress

tjake · Oct 16, 2024 · 6536a29 · 6536a29
1 parent 6d3ab99
commit 6536a29
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java b/jlama-core/src/main/java/com/github/tjake/jlama/safetensors/SafeTensorSupport.java
@@ -310,7 +310,7 @@ public void write(byte[] b, int off, int len) throws IOException {
         return qPath;
     }
 
-    public static File maybeDownloadModel(String modelDir, String fullModelName) throws IOException {
+    public static File maybeDownloadModel(String modelDir, String fullModelName, TriConsumer<String, Long, Long> progressReporter) throws IOException {
         String[] parts = fullModelName.split("/");
         if (parts.length == 0 || parts.length > 2) {
             throw new IllegalArgumentException("Model must be in the form owner/name");
@@ -327,7 +327,11 @@ public static File maybeDownloadModel(String modelDir, String fullModelName) thr
             name = parts[1];
         }
 
-        return maybeDownloadModel(modelDir, Optional.ofNullable(owner), name, true, Optional.empty(), Optional.empty(), Optional.empty());
+        return maybeDownloadModel(modelDir, Optional.ofNullable(owner), name, true, Optional.empty(), Optional.empty(), Optional.ofNullable(progressReporter));
+    }
+
+    public static File maybeDownloadModel(String modelDir, String fullModelName) throws IOException {
+        return maybeDownloadModel(modelDir, fullModelName, null);
     }
 
     public static Path constructLocalModelPath(String modelDir, String owner, String modelName) {