Rename --n_steps to --sequence_length

clebert · Oct 19, 2023 · 6406661 · 6406661
1 parent 389afbe
commit 6406661
Show file tree

Hide file tree

Showing 8 changed files with 62 additions and 49 deletions.
diff --git a/README.md b/README.md
@@ -38,11 +38,11 @@ zig build -Doptimize=ReleaseFast run-generator -- models/llama2_7b_hf --prompt "
 Usage: llama2-generator <model_path> [options]
 
 Options:
-  --temperature   <float>  = 1.0
-  --top_p         <float>  = 0.9
-  --random_seed   <int>    = <milli_timestamp>
-  --n_steps       <int>    = <max_sequence_length>
-  --prompt        <string> = ""
+  --temperature     <float>  = 1.0
+  --top_p           <float>  = 0.9
+  --random_seed     <int>    = <milli_timestamp>
+  --sequence_length <int>    = <max_sequence_length>
+  --prompt          <string> = ""
   --verbose
   --help
 ```
@@ -53,12 +53,12 @@ Options:
 Usage: llama2-chat <model_path> [options]
 
 Options:
-  --temperature   <float>  = 1.0
-  --top_p         <float>  = 0.9
-  --random_seed   <int>    = <milli_timestamp>
-  --n_steps       <int>    = <max_sequence_length>
-  --system_prompt <string> = ""
-  --user_prompt   <string> = ""
+  --temperature     <float>  = 1.0
+  --top_p           <float>  = 0.9
+  --random_seed     <int>    = <milli_timestamp>
+  --sequence_length <int>    = <max_sequence_length>
+  --system_prompt   <string> = ""
+  --user_prompt     <string> = ""
   --help
 ```
 

diff --git a/src/chat.zig b/src/chat.zig
@@ -15,7 +15,7 @@ system_prompt: []const u8,
 user_prompt: []const u8,
 
 pub fn init(allocator: std.mem.Allocator, args: ChatArgs) !Self {
-    const transformer = try Transformer.init(allocator, args.model_path, args.n_steps);
+    const transformer = try Transformer.init(allocator, args.model_path, args.sequence_length);
 
     errdefer transformer.deinit();
 

diff --git a/src/chat_args.zig b/src/chat_args.zig
@@ -7,11 +7,18 @@ model_path: []const u8,
 temperature: f32,
 top_p: f32,
 random_seed: u64,
-n_steps: usize,
+sequence_length: usize,
 system_prompt: []const u8,
 user_prompt: []const u8,
 
-const Option = enum { temperature, top_p, random_seed, n_steps, system_prompt, user_prompt };
+const Option = enum {
+    temperature,
+    top_p,
+    random_seed,
+    sequence_length,
+    system_prompt,
+    user_prompt,
+};
 
 pub fn init(allocator: std.mem.Allocator) !Self {
     var arg_iterator = try std.process.argsWithAllocator(allocator);
@@ -26,7 +33,7 @@ pub fn init(allocator: std.mem.Allocator) !Self {
     var temperature: ?f32 = null;
     var top_p: ?f32 = null;
     var random_seed: ?u64 = null;
-    var n_steps: ?usize = null;
+    var sequence_length: ?usize = null;
     var system_prompt: ?[]const u8 = null;
     var user_prompt: ?[]const u8 = null;
 
@@ -38,8 +45,8 @@ pub fn init(allocator: std.mem.Allocator) !Self {
                 top_p = try std.fmt.parseFloat(f32, arg);
             } else if (option == .random_seed and random_seed == null) {
                 random_seed = try std.fmt.parseInt(u64, arg, 10);
-            } else if (option == .n_steps and n_steps == null) {
-                n_steps = try std.fmt.parseInt(usize, arg, 10);
+            } else if (option == .sequence_length and sequence_length == null) {
+                sequence_length = try std.fmt.parseInt(usize, arg, 10);
             } else if (option == .system_prompt and system_prompt == null) {
                 system_prompt = arg;
             } else if (option == .user_prompt and user_prompt == null) {
@@ -55,8 +62,8 @@ pub fn init(allocator: std.mem.Allocator) !Self {
             current_option = .top_p;
         } else if (std.mem.eql(u8, arg, "--random_seed")) {
             current_option = .random_seed;
-        } else if (std.mem.eql(u8, arg, "--n_steps")) {
-            current_option = .n_steps;
+        } else if (std.mem.eql(u8, arg, "--sequence_length")) {
+            current_option = .sequence_length;
         } else if (std.mem.eql(u8, arg, "--system_prompt")) {
             current_option = .system_prompt;
         } else if (std.mem.eql(u8, arg, "--user_prompt")) {
@@ -76,7 +83,7 @@ pub fn init(allocator: std.mem.Allocator) !Self {
         .temperature = @max(@min(temperature orelse 1, 1), 0),
         .top_p = @max(@min(top_p orelse 0.9, 1), 0),
         .random_seed = random_seed orelse @intCast(std.time.milliTimestamp()),
-        .n_steps = n_steps orelse 0,
+        .sequence_length = sequence_length orelse 0,
         .system_prompt = system_prompt orelse "",
         .user_prompt = user_prompt orelse "",
     };
@@ -95,12 +102,12 @@ fn help(exit_status: u8) !noreturn {
     try console.print("Usage: llama2-chat <model_path> [options]\n\n", .{});
 
     try console.print("Options:\n", .{});
-    try console.print("  --temperature   <float>  = 1.0\n", .{});
-    try console.print("  --top_p         <float>  = 0.9\n", .{});
-    try console.print("  --random_seed   <int>    = <milli_timestamp>\n", .{});
-    try console.print("  --n_steps       <int>    = <max_sequence_length>\n", .{});
-    try console.print("  --system_prompt <string> = \"\"\n", .{});
-    try console.print("  --user_prompt   <string> = \"\"\n", .{});
+    try console.print("  --temperature     <float>  = 1.0\n", .{});
+    try console.print("  --top_p           <float>  = 0.9\n", .{});
+    try console.print("  --random_seed     <int>    = <milli_timestamp>\n", .{});
+    try console.print("  --sequence_length <int>    = <max_sequence_length>\n", .{});
+    try console.print("  --system_prompt   <string> = \"\"\n", .{});
+    try console.print("  --user_prompt     <string> = \"\"\n", .{});
     try console.print("  --help\n", .{});
 
     std.process.exit(exit_status);

diff --git a/src/converter_args.zig b/src/converter_args.zig
@@ -5,8 +5,6 @@ const std = @import("std");
 arg_iterator: std.process.ArgIterator,
 model_path: []const u8,
 
-const Option = enum { temperature, top_p, random_seed, n_steps, prompt };
-
 pub fn init(allocator: std.mem.Allocator) !Self {
     var arg_iterator = try std.process.argsWithAllocator(allocator);
 

diff --git a/src/generator.zig b/src/generator.zig
@@ -15,7 +15,7 @@ prompt_tokens: []usize,
 verbose: bool,
 
 pub fn init(allocator: std.mem.Allocator, args: GeneratorArgs) !Self {
-    const transformer = try Transformer.init(allocator, args.model_path, args.n_steps);
+    const transformer = try Transformer.init(allocator, args.model_path, args.sequence_length);
 
     errdefer transformer.deinit();
 
@@ -113,7 +113,7 @@ test "generate tiny story" {
         .temperature = 1,
         .top_p = 0.9,
         .random_seed = 42,
-        .n_steps = 10,
+        .sequence_length = 10,
         .prompt = "There was",
         .verbose = false,
     };

diff --git a/src/generator_args.zig b/src/generator_args.zig
@@ -7,11 +7,11 @@ model_path: []const u8,
 temperature: f32,
 top_p: f32,
 random_seed: u64,
-n_steps: usize,
+sequence_length: usize,
 prompt: []const u8,
 verbose: bool,
 
-const Option = enum { temperature, top_p, random_seed, n_steps, prompt };
+const Option = enum { temperature, top_p, random_seed, sequence_length, prompt };
 
 pub fn init(allocator: std.mem.Allocator) !Self {
     var arg_iterator = try std.process.argsWithAllocator(allocator);
@@ -26,7 +26,7 @@ pub fn init(allocator: std.mem.Allocator) !Self {
     var temperature: ?f32 = null;
     var top_p: ?f32 = null;
     var random_seed: ?u64 = null;
-    var n_steps: ?usize = null;
+    var sequence_length: ?usize = null;
     var prompt: ?[]const u8 = null;
     var verbose: bool = false;
 
@@ -38,8 +38,8 @@ pub fn init(allocator: std.mem.Allocator) !Self {
                 top_p = try std.fmt.parseFloat(f32, arg);
             } else if (option == .random_seed and random_seed == null) {
                 random_seed = try std.fmt.parseInt(u64, arg, 10);
-            } else if (option == .n_steps and n_steps == null) {
-                n_steps = try std.fmt.parseInt(usize, arg, 10);
+            } else if (option == .sequence_length and sequence_length == null) {
+                sequence_length = try std.fmt.parseInt(usize, arg, 10);
             } else if (option == .prompt and prompt == null) {
                 prompt = arg;
             } else {
@@ -53,8 +53,8 @@ pub fn init(allocator: std.mem.Allocator) !Self {
             current_option = .top_p;
         } else if (std.mem.eql(u8, arg, "--random_seed")) {
             current_option = .random_seed;
-        } else if (std.mem.eql(u8, arg, "--n_steps")) {
-            current_option = .n_steps;
+        } else if (std.mem.eql(u8, arg, "--sequence_length")) {
+            current_option = .sequence_length;
         } else if (std.mem.eql(u8, arg, "--prompt")) {
             current_option = .prompt;
         } else if (std.mem.eql(u8, arg, "--verbose") and !verbose) {
@@ -74,7 +74,7 @@ pub fn init(allocator: std.mem.Allocator) !Self {
         .temperature = @max(@min(temperature orelse 1, 1), 0),
         .top_p = @max(@min(top_p orelse 0.9, 1), 0),
         .random_seed = random_seed orelse @intCast(std.time.milliTimestamp()),
-        .n_steps = n_steps orelse 0,
+        .sequence_length = sequence_length orelse 0,
         .prompt = prompt orelse "",
         .verbose = verbose,
     };
@@ -93,11 +93,11 @@ fn help(exit_status: u8) !noreturn {
     try console.print("Usage: llama2-generator <model_path> [options]\n\n", .{});
 
     try console.print("Options:\n", .{});
-    try console.print("  --temperature   <float>  = 1.0\n", .{});
-    try console.print("  --top_p         <float>  = 0.9\n", .{});
-    try console.print("  --random_seed   <int>    = <milli_timestamp>\n", .{});
-    try console.print("  --n_steps       <int>    = <max_sequence_length>\n", .{});
-    try console.print("  --prompt        <string> = \"\"\n", .{});
+    try console.print("  --temperature     <float>  = 1.0\n", .{});
+    try console.print("  --top_p           <float>  = 0.9\n", .{});
+    try console.print("  --random_seed     <int>    = <milli_timestamp>\n", .{});
+    try console.print("  --sequence_length <int>    = <max_sequence_length>\n", .{});
+    try console.print("  --prompt          <string> = \"\"\n", .{});
     try console.print("  --verbose\n", .{});
     try console.print("  --help\n", .{});
 

diff --git a/src/transformer.zig b/src/transformer.zig
@@ -14,12 +14,20 @@ ffn: FFN,
 hidden_buffer: Tensor(1),
 output_buffer: Tensor(1),
 
-pub fn init(allocator: std.mem.Allocator, model_path: []const u8, n_steps: usize) !Self {
+pub fn init(
+    allocator: std.mem.Allocator,
+    model_path: []const u8,
+    custom_sequence_length: usize,
+) !Self {
     const checkpoint = try Checkpoint.init(allocator, model_path);
 
     errdefer checkpoint.deinit();
 
-    const sequence_length = if (n_steps == 0) checkpoint.max_sequence_length else n_steps;
+    const sequence_length = if (custom_sequence_length == 0)
+        checkpoint.max_sequence_length
+    else
+        custom_sequence_length;
+
     const attention = try Attention.init(allocator, checkpoint, sequence_length);
 
     errdefer attention.deinit();

diff --git a/test.sh b/test.sh
@@ -9,7 +9,7 @@ zig build
 
 model_path="models/tinystories_260k"
 
-actual_output=$(./zig-out/bin/llama2-generator $model_path --temperature 0 --n_steps 200)
+actual_output=$(./zig-out/bin/llama2-generator $model_path --temperature 0 --sequence_length 200)
 
 # Generated with llama2.c (https://github.com/karpathy/llama2.c/tree/7ac65cb2c2b169050747be92011b7bebdd1b4544)
 expected_output="Once upon a time, there was a little girl named Lily. She loved to play outside in the park. One day, she saw a big, red ball. She wanted to play with it, but it was too high.
@@ -22,7 +22,7 @@ if [ "$actual_output" != "$expected_output" ]; then
     exit 1
 fi
 
-actual_output=$(./zig-out/bin/llama2-generator $model_path --top_p 1 --random_seed 42 --n_steps 200)
+actual_output=$(./zig-out/bin/llama2-generator $model_path --top_p 1 --random_seed 42 --sequence_length 200)
 
 # Generated with llama2.c (https://github.com/karpathy/llama2.c/tree/7ac65cb2c2b169050747be92011b7bebdd1b4544)
 expected_output="Once upon a time, there was a big roof. The fox was ready to look for people inside. He saw a big rock near a big tree. The roof was very small and fun! He ate the roof too. He got a shiny stool, so he sicked the roof with his friend, the girl named Mia.
@@ -35,7 +35,7 @@ if [ "$actual_output" != "$expected_output" ]; then
     exit 1
 fi
 
-actual_output=$(./zig-out/bin/llama2-generator $model_path --top_p 0.95 --random_seed 42 --n_steps 200)
+actual_output=$(./zig-out/bin/llama2-generator $model_path --top_p 0.95 --random_seed 42 --sequence_length 200)
 
 # Generated with llama2.c (https://github.com/karpathy/llama2.c/tree/7ac65cb2c2b169050747be92011b7bebdd1b4544)
 expected_output="Once upon a time, there was a little boy named Timmy. Timmy loved going to the park with his mom. One day, Lily went outside to play outside in her pocket. He was scared and didn't know where to buy some colorful animals.
@@ -46,7 +46,7 @@ if [ "$actual_output" != "$expected_output" ]; then
     exit 1
 fi
 
-actual_output=$(./zig-out/bin/llama2-generator $model_path --top_p 0.95 --random_seed 42 --n_steps 200 --prompt "There was a big")
+actual_output=$(./zig-out/bin/llama2-generator $model_path --top_p 0.95 --random_seed 42 --sequence_length 200 --prompt "There was a big")
 
 # Generated with llama2.c (https://github.com/karpathy/llama2.c/tree/7ac65cb2c2b169050747be92011b7bebdd1b4544)
 expected_output="There was a big pretty grass. It was a long elephant. The cars wanted to tell him that as they spin before the amazing doll, just like it she was always okay.