From eb1b359dae48b4996dfbbf4898ae08bf326c87fa Mon Sep 17 00:00:00 2001
From: Aleksa Gordic <gordicaleksa@gmail.com>
Date: Sat, 10 Aug 2024 22:45:46 +0200
Subject: [PATCH] Replace fire with argparse

---
 dev/data/README.md          |  6 +++---
 dev/data/tinyshakespeare.py | 13 ++++++-------
 dev/data/tinystories.py     | 13 ++++++-------
 requirements.txt            |  1 -
 4 files changed, 15 insertions(+), 18 deletions(-)

diff --git a/dev/data/README.md b/dev/data/README.md
index 424dce5..7966a42 100644
--- a/dev/data/README.md
+++ b/dev/data/README.md
@@ -2,9 +2,9 @@
 
 The idea is that each dataset has a .py file here in the root of `dev/data`, and each dataset then creates a directory here, and writes and caches anything inside that directory. So for example:
 
-- running `python tinystories.py --model gpt-2` will create a directory `tinystories` with its .bin files inside it
-- running `python tinyshakespeare.py -- model gpt-2` will create a directory `tinyshakespeare` with its .bin files inside it
+- running `python tinystories.py` will create a directory `tinystories` with its .bin files inside it
+- running `python tinyshakespeare.py` will create a directory `tinyshakespeare` with its .bin files inside it
 
 And so on. This way we can nicely organize multiple datasets here, share common utilities between them, and then point the .py/.c code in the root of the project accordingly to these.
 
-Note: `--model` can be either "gpt-2" or "llama" (we assume llama 3).
+Note: we support "gpt-2" and "llama" (llama 3 in particular) models and the above scripts will tokenize gpt-2 by default.
diff --git a/dev/data/tinyshakespeare.py b/dev/data/tinyshakespeare.py
index eed862e..814704e 100644
--- a/dev/data/tinyshakespeare.py
+++ b/dev/data/tinyshakespeare.py
@@ -19,9 +19,9 @@ connection and computer. The .bin files are raw byte
 streams of uint16 (gpt-2) or uint32 (llama) numbers indicating the token ids.
 """
 
+import argparse
 import os
 
-import fire
 import tiktoken
 from transformers import AutoTokenizer
 
@@ -69,10 +69,9 @@ def tokenize(model):
     write_datafile(val_filename, val_tokens, model)
     write_datafile(train_filename, train_tokens, model)
 
-def process(model):
-    assert model in ["gpt-2", "llama"], f"unknown model {model} (choose from gpt-2, llama)"
-    download()
-    tokenize(model)
-
 if __name__ == "__main__":
-    fire.Fire(process)
\ No newline at end of file
+    parser = argparse.ArgumentParser(description="Tiny Shakespeare dataset preprocessing")
+    parser.add_argument("-m", "--model", type=str, default="gpt-2", choices=["gpt-2", "llama"], help="Model type, gpt-2|llama")
+    args = parser.parse_args()
+    download()
+    tokenize(args.model)
\ No newline at end of file
diff --git a/dev/data/tinystories.py b/dev/data/tinystories.py
index 81de5f6..5bdf699 100644
--- a/dev/data/tinystories.py
+++ b/dev/data/tinystories.py
@@ -25,13 +25,13 @@ connection and computer. The .bin files are raw byte
 streams of uint16 (gpt-2) or uint32 (llama) numbers indicating the token ids.
 """
 
+import argparse
 import os
 import glob
 import json
 import random
 from concurrent.futures import ProcessPoolExecutor, as_completed
 
-import fire
 import tiktoken
 from transformers import AutoTokenizer
 
@@ -116,10 +116,9 @@ def tokenize(model):
         split_filename = os.path.join(DATA_CACHE_DIR, f"TinyStories_{split_name}.bin")
         write_datafile(split_filename, all_tokens, model)
 
-def process(model):
-    assert model in ["gpt-2", "llama"], f"unknown model {model} (choose from gpt-2, llama)"
-    download()
-    tokenize(model)
-
 if __name__ == "__main__":
-    fire.Fire(process)
\ No newline at end of file
+    parser = argparse.ArgumentParser(description="Tiny Stories dataset preprocessing")
+    parser.add_argument("-m", "--model", type=str, default="gpt-2", choices=["gpt-2", "llama"], help="Model type, gpt-2|llama")
+    args = parser.parse_args()
+    download()
+    tokenize(args.model)
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 8971c62..ea4bc76 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,4 +1,3 @@
-fire
 tqdm
 numpy<2
 torch