add parquet file output to the data command

megacamelus · Jul 30, 2024 · 2c3303a · 2c3303a
1 parent faf59db
commit 2c3303a
Show file tree

Hide file tree

Showing 6 changed files with 101 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -132,6 +132,8 @@ java -jar target/camel-jbang-plugin-explain-4.7.0-SNAPSHOT-jar-with-dependencies
 
 You can generate LLM training datasets from the catalog information.
 
+JSON and Parquet files are generated in the `dataset` directory.
+
 Generate training data using the component information:
 ```shell
 java -jar target/camel-jbang-plugin-explain-4.7.0-SNAPSHOT-jar-with-dependencies.jar data --model-name --data-type components mistral:latest

diff --git a/pom.xml b/pom.xml
@@ -21,6 +21,10 @@
         <main.class>org.apache.camel.standalone.Main</main.class>
         <velocity-version>2.3</velocity-version>
         <commonmark-version>0.22.0</commonmark-version>
+
+        <avro-version>1.11.3</avro-version>
+        <parquet-avro-version>1.14.1</parquet-avro-version>
+        <hadoop-version>3.4.0</hadoop-version>
     </properties>
 
     <dependencyManagement>
@@ -77,6 +81,28 @@
             <version>${commonmark-version}</version>
         </dependency>
 
+        <dependency>
+            <groupId>org.apache.parquet</groupId>
+            <artifactId>parquet-avro</artifactId>
+            <version>${parquet-avro-version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.avro</groupId>
+            <artifactId>avro</artifactId>
+            <version>${avro-version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.hadoop</groupId>
+            <artifactId>hadoop-common</artifactId>
+            <version>${hadoop-version}</version>
+        </dependency>
+        <dependency>
+            <groupId>org.apache.hadoop</groupId>
+            <artifactId>hadoop-mapreduce-client-core</artifactId>
+            <version>${hadoop-version}</version>  
+        </dependency>
+
+
     </dependencies>
 
     <profiles>

diff --git a/src/main/java/org/apache/camel/jbang/ai/data/CatalogProcessor.java b/src/main/java/org/apache/camel/jbang/ai/data/CatalogProcessor.java
@@ -60,7 +60,7 @@ protected void processRecords(int startFrom, List<String> componentNames, int to
     protected void processOption(
             List<AlpacaRecord> alpacaRecords, String componentName,
             List<? extends BaseOptionModel> optionModels, String type) {
-        int componentOptionCount = 0;
+        int componentOptionCount = 1;
         final int componentOptionTotal = optionModels.size();
         for (BaseOptionModel optionModel : optionModels) {
             StopWatch watch = new StopWatch();

diff --git a/src/main/java/org/apache/camel/jbang/ai/util/CatalogUtil.java b/src/main/java/org/apache/camel/jbang/ai/util/CatalogUtil.java
@@ -2,6 +2,7 @@
 
 import java.io.File;
 import java.io.IOException;
+import java.nio.file.Path;
 import java.time.Instant;
 import java.time.ZoneId;
 import java.time.format.DateTimeFormatter;
@@ -13,22 +14,30 @@
 import org.apache.camel.jbang.ai.types.AlpacaRecord;
 import org.apache.camel.tooling.model.BaseOptionModel;
 
-public class CatalogUtil {
+public final class CatalogUtil {
+    private static final String DATASET_DIR = "dataset";
     private static final String PATTERN_FORMAT = "HH:mm:ss";
 
+    private CatalogUtil() {
+        throw new IllegalStateException("Util final class should be instantiated.");
+    }
+
     public static void saveRecords(List<AlpacaRecord> alpacaRecords, String componentName) {
         if (!alpacaRecords.isEmpty()) {
             ObjectMapper mapper = new ObjectMapper();
             mapper.enable(SerializationFeature.INDENT_OUTPUT);
 
-            final File file = new File("dataset", String.format("camel-%s.json", componentName));
+            final File file = new File(DATASET_DIR, String.format("camel-%s.json", componentName));
             file.getParentFile().mkdirs();
 
             try {
                 mapper.writeValue(file, alpacaRecords);
             } catch (IOException e) {
-                throw new RuntimeException(e);
+                throw new RuntimeException("Failed to write json file", e);
             }
+
+            Path parquetOutput = Path.of(DATASET_DIR, String.format("camel-%s.parquet", componentName));
+            ParquetUtil.saveParquet(alpacaRecords, parquetOutput);
         }
     }
 

diff --git a/src/main/java/org/apache/camel/jbang/ai/util/ParquetUtil.java b/src/main/java/org/apache/camel/jbang/ai/util/ParquetUtil.java
@@ -0,0 +1,50 @@
+package org.apache.camel.jbang.ai.util;
+
+import java.io.IOException;
+import java.nio.file.Path;
+import java.util.List;
+
+import org.apache.avro.Schema;
+import org.apache.avro.generic.GenericData;
+import org.apache.avro.generic.GenericRecord;
+import org.apache.camel.jbang.ai.types.AlpacaRecord;
+import org.apache.parquet.avro.AvroParquetWriter;
+import org.apache.parquet.io.LocalOutputFile;
+
+public final class ParquetUtil {
+
+    private static final String SCHEMA_FILE = "/training-set-schema.avsc";
+
+    private ParquetUtil() {
+        throw new IllegalStateException("Util final class should be instantiated.");
+    }
+
+    public static void saveParquet(List<AlpacaRecord> alpacaRecords, Path outputPath) {
+
+        Schema schema;
+        try (var schemaContent = ParquetUtil.class.getResourceAsStream(SCHEMA_FILE)) {
+            schema = new Schema.Parser().parse(schemaContent);
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to load Parquet schema", e);
+        }
+
+        var outputFile = new LocalOutputFile(outputPath);
+        var builder = AvroParquetWriter
+                .<GenericRecord>builder(outputFile)
+                .withSchema(schema);
+
+        try (var writer = builder.build()) {
+            for (AlpacaRecord ar : alpacaRecords) {
+                GenericRecord r = new GenericData.Record(schema);
+                r.put("input", ar.getInput());
+                r.put("instruction", ar.getInstruction());
+                r.put("output", ar.getOutput());
+                writer.write(r);
+            }
+        } catch (IOException e) {
+            throw new RuntimeException("Failed to write Parquet file", e);
+        }
+
+    }
+
+}
diff --git a/src/main/resources/training-set-schema.avsc b/src/main/resources/training-set-schema.avsc
@@ -0,0 +1,10 @@
+{
+    "namespace": "org.apache.camel",
+    "type": "record",
+    "name": "AlpacaRecord",
+    "fields": [
+        {"name": "instruction", "type": "string"},
+        {"name": "input", "type": "string"},
+        {"name": "output", "type": "string"}
+    ]
+}