Merge pull request #2 from mgonzs13/main

hf_hub_download_with_shards

Author: agonzc34

README.md

@@ -6,7 +6,6 @@ With this library, C++ developers can integrate Hugging Face model downloads dir
 
 [![GitHub License](https://img.shields.io/github/license/agonzc34/huggingface-hub-cpp)](https://opensource.org/license/mit) [![GitHub release](https://img.shields.io/github/release/agonzc34/huggingface-hub-cpp.svg)](https://github.com/agonzc34/huggingface-hub-cpp/releases) [![Code Size](https://img.shields.io/github/languages/code-size/agonzc34/huggingface-hub-cpp.svg?branch=main)](https://github.com/agonzc34/huggingface-hub-cpp?branch=main) [![Last Commit](https://img.shields.io/github/last-commit/agonzc34/huggingface-hub-cpp.svg)](https://github.com/agonzc34/huggingface-hub-cpp/commits/main) [![GitHub issues](https://img.shields.io/github/issues/agonzc34/huggingface-hub-cpp)](https://github.com/agonzc34/huggingface-hub-cpp/issues) [![GitHub pull requests](https://img.shields.io/github/issues-pr/agonzc34/huggingface-hub-cpp)](https://github.com/agonzc34/huggingface-hub-cpp/pulls) [![Contributors](https://img.shields.io/github/contributors/agonzc34/huggingface-hub-cpp.svg)](https://github.com/agonzc34/huggingface-hub-cpp/graphs/contributors) [![Build](https://github.com/agonzc34/huggingface-hub-cpp/actions/workflows/cmake-build-status.yml/badge.svg)](https://github.com/agonzc34/huggingface-hub-cpp/actions/workflows/cmake-build-status.yml?branch=main) [![Doxygen Deployment](https://github.com/agonzc34/huggingface-hub-cpp/actions/workflows/doxygen-deployment.yml/badge.svg)](https://agonzc34.github.io/huggingface-hub-cpp/)
 
-
 ## Table of Contents
 
 - [huggingface-hub-cpp](#huggingface-hub-cpp)
@@ -64,6 +63,7 @@ int main() {
   } else {
     std::cout << "Error" << std::endl;
   }
+}
 ```
 
 ### Running the demo app

include/huggingface_hub.h

@@ -95,5 +95,25 @@ hf_hub_download(const std::string &repo_id, const std::string &filename,
                 const std::string &cache_dir = "~/.cache/huggingface/hub",
                 bool force_download = false);
 
+/**
+ * @brief Download a file from Hugging Face Hub.
+ *
+ * This function downloads a specified file from a given repository on the
+ * Hugging Face Hub and saves it to the specified cache directory.
+ *
+ * @param repo_id The repository ID.
+ * @param filename The name of the file to download.
+ * @param cache_dir The directory to cache the downloaded file. Default is
+ * "~/.cache/huggingface/hub".
+ * @param force_download If true, forces the download even if the file already
+ * exists in the cache.
+ * @return A DownloadResult structure containing the success status and the path
+ * of the downloaded file.
+ */
+struct DownloadResult hf_hub_download_with_shards(
+    const std::string &repo_id, const std::string &filename,
+    const std::string &cache_dir = "~/.cache/huggingface/hub",
+    bool force_download = false);
+
 #endif // HUGGINGFACE_HUB_H
 } // namespace huggingface_hub

src/huggingface_hub.cpp

@@ -21,19 +21,23 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-#include "huggingface_hub.h"
 #include <algorithm>
 #include <chrono>
 #include <csignal>
-#include <curl/curl.h>
 #include <filesystem>
 #include <fstream>
 #include <iomanip>
+#include <regex>
 #include <sstream>
+
+#include <curl/curl.h>
 #include <sys/stat.h>
 #include <sys/types.h>
 
+#include "huggingface_hub.h"
+
 namespace huggingface_hub {
+
 volatile sig_atomic_t stop_download = 0;
 void handle_sigint(int) { stop_download = 1; }
 
@@ -369,4 +373,40 @@ struct DownloadResult hf_hub_download(const std::string &repo_id,
   result.success = res == CURLE_OK;
   return result;
 }
+
+struct DownloadResult hf_hub_download_with_shards(const std::string &repo_id,
+                                                  const std::string &filename,
+                                                  const std::string &cache_dir,
+                                                  bool force_download) {
+
+  std::regex pattern(R"(-([0-9]+)-of-([0-9]+)\.gguf)");
+  std::smatch match;
+
+  if (std::regex_search(filename, match, pattern)) {
+    int total_shards = std::stoi(match[2]);
+    std::string base_name = filename.substr(0, match.position(0));
+
+    // Download shards
+    for (int i = 1; i <= total_shards; ++i) {
+      char shard_file[512];
+      snprintf(shard_file, sizeof(shard_file), "%s-%05d-of-%05d.gguf",
+               base_name.c_str(), i, total_shards);
+      auto aux_res =
+          hf_hub_download(repo_id, shard_file, cache_dir, force_download);
+
+      if (!aux_res.success) {
+        return aux_res;
+      }
+    }
+
+    // Return first shard
+    char first_shard[512];
+    snprintf(first_shard, sizeof(first_shard), "%s-00001-of-%05d.gguf",
+             base_name.c_str(), total_shards);
+    return hf_hub_download(repo_id, first_shard, cache_dir, false);
+  }
+
+  return hf_hub_download(repo_id, filename, cache_dir, force_download);
+}
+
 } // namespace huggingface_hub