train: add simple loading already tokenized data from parquet dataset

Author: lexasub

CMakeLists.txt

@@ -84,6 +84,7 @@ option(LLAMA_BUILD_SERVER   "llama: build server example" ${LLAMA_STANDALONE})
 # 3rd party libs
 option(LLAMA_CURL       "llama: use libcurl to download model from an URL" ON)
 option(LLAMA_LLGUIDANCE "llama-common: include LLGuidance library for structured output in common utils" OFF)
+option(LLAMA_PARQUET "Enable Parquet dataset support via Arrow/Parquet C++" OFF)
 
 # Required for relocatable CMake package
 include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/build-info.cmake)
@@ -173,6 +174,12 @@ if (MINGW)
     add_compile_definitions(_WIN32_WINNT=${GGML_WIN_VER})
 endif()
 
+if(LLAMA_PARQUET)
+    find_package(Arrow REQUIRED)
+    find_package(Parquet REQUIRED)
+    add_definitions(-DLLAMA_PARQUET)
+endif()
+
 #
 # build the library
 #

examples/training/finetune.cpp

@@ -2,6 +2,7 @@
 #include "common.h"
 #include "log.h"
 #include "llama.h"
+#include "../../src/parquet_dataset.h"
 
 #include <cmath>
 #include <cstdio>
@@ -57,7 +58,17 @@ int main(int argc, char ** argv) {
 
     constexpr float val_split = 0.05f;
 
-    std::vector<llama_token> tokens = common_tokenize(ctx.get(), params.prompt, true);
+#ifndef LLAMA_PARQUET
+    std::vector<llama_token> tokens = common_tokenize(ctx.get(), params.prompt, true); //load from text file
+#else
+    auto tokens = load_parquet_dataset("test.parquet" /*params.parquet_path, params.tokens_column*/ ,"tokens");
+    if (tokens.empty()) {
+        //LOG_ERR("No tokens in %s, or column %s not found/invalid", params.parquet_path.c_str(), params.tokens_column.c_str());
+        return 1;
+    }
+    LOG_INF("Loaded %zu tokens from Parquet", tokens.size());
+#endif
+
     ggml_opt_dataset_t dataset = common_opt_dataset_init(ctx.get(), tokens, llama_n_ctx(ctx.get())/2);
 
     struct ggml_opt_optimizer_params optimizer_params = ggml_opt_get_default_optimizer_params(nullptr);

src/CMakeLists.txt

@@ -32,6 +32,7 @@ add_library(llama
             llama-quant.cpp
             llama-sampling.cpp
             llama-vocab.cpp
+            parquet_dataset.cpp
             unicode-data.cpp
             unicode.cpp
             unicode.h
@@ -41,7 +42,12 @@ target_include_directories(llama PRIVATE .)
 target_include_directories(llama PUBLIC ../include)
 target_compile_features   (llama PRIVATE cxx_std_17) # don't bump
 
-target_link_libraries(llama PUBLIC ggml)
+
+if(LLAMA_PARQUET)
+    target_link_libraries(llama PUBLIC ggml Arrow::arrow_shared Parquet::parquet_shared)
+else()
+    target_link_libraries(llama PUBLIC ggml)
+endif()
 
 if (BUILD_SHARED_LIBS)
     set_target_properties(llama PROPERTIES POSITION_INDEPENDENT_CODE ON)

src/parquet_dataset.cpp

@@ -0,0 +1,47 @@
+#ifdef LLAMA_PARQUET
+#include "parquet_dataset.h"
+#include <arrow/api.h>
+#include <arrow/io/file.h>
+#include <parquet/arrow/reader.h>
+#include "llama-impl.h"
+
+std::vector<llama_token> load_parquet_dataset(const std::string &path, const std::string &column) {
+    arrow::MemoryPool *pool = arrow::default_memory_pool();
+    std::shared_ptr<arrow::io::RandomAccessFile> infile;
+    PARQUET_ASSIGN_OR_THROW(infile, arrow::io::ReadableFile::Open(path));
+    arrow::Result<std::unique_ptr<parquet::arrow::FileReader>> reader_raw;
+    PARQUET_ASSIGN_OR_THROW(reader_raw, parquet::arrow::OpenFile(infile, pool));
+
+    std::unique_ptr<parquet::arrow::FileReader> reader = std::move(reader_raw.ValueUnsafe());
+    std::shared_ptr<arrow::Table> table;
+    PARQUET_THROW_NOT_OK(reader->ReadTable(&table));
+
+    auto field = table->schema()->GetFieldByName(column);
+    if (!field || !field->type()->Equals(arrow::list(arrow::int32()))) {
+        LLAMA_LOG_ERROR("Parquet column '%s' missing or not list<int32>", column.c_str());
+        return {};
+    }
+
+    auto col = table->GetColumnByName(column);
+    std::vector<llama_token> tokens;
+    for (int chunk = 0; chunk < col->num_chunks(); ++chunk) {
+        auto list_arr = std::static_pointer_cast<arrow::ListArray>(col->chunk(chunk));
+        auto values_arr = std::static_pointer_cast<arrow::Int32Array>(list_arr->values());
+        // get raw offsets (int32_t or int64_t based on ListArray template)
+        const auto *offsets = list_arr->raw_value_offsets();
+        // offsets length = list_arr->length() + 1
+        int64_t values_length = values_arr->length();
+        for (int64_t i = 0; i < list_arr->length(); ++i) {
+            int64_t start = offsets[i];
+            int64_t end   = offsets[i + 1];
+            // Clamp end
+            if (start < 0) start = 0;
+            if (end > values_length) end = values_length;
+            for (int64_t j = start; j < end; ++j) {
+                tokens.push_back(static_cast<llama_token>(values_arr->Value(j)));
+            }
+        }
+    }
+    return tokens;
+}
+#endif // LLAMA_PARQUET

src/parquet_dataset.h

@@ -0,0 +1,10 @@
+#ifndef PARQUET_DATASET_H
+#define PARQUET_DATASET_H
+#include <string>
+#include <vector>
+#include "llama.h"
+
+#ifdef LLAMA_PARQUET
+std::vector<llama_token> load_parquet_dataset(const std::string &path, const std::string &column);
+#endif
+#endif  //