WIP

1 files changed, 4 insertions, 4 deletions

diff --git a/finetuning.py b/finetuning.py
index a0ad1fe..935adfa 100755
--- a/finetuning.py
+++ b/finetuning.py
@@ -5,7 +5,7 @@ import json
 import os
 import torch
 
-from datasets import Dataset, Image, load_dataset
+from datasets import Dataset, load_dataset
 from functools import partial
 from glob import glob
 from peft import LoraConfig
@@ -75,10 +75,10 @@ def main():
     parser.add_argument("--lora-dropout", default=0.05, type=float)
     parser.add_argument("--lora-r", default=16, type=int)
     parser.add_argument("--bnb-4bit", action="store_true")
-    parser.add_argument("--lora-small", action="store_true")
     parser.add_argument("--max-grad-norm", default=1.0, type=float)
     parser.add_argument("--max-length", default=512, type=int)
-    parser.add_argument("--lora-config", choices=["S", "M", "L"], default="M")
+    parser.add_argument("--lora-config", choices=["S", "M", "L"], default="M", type=str)
+    parser.add_argument("--dataset", default="asdf2k/caption_translation", type=str)
     args = parser.parse_args()
 
     if args.bnb_4bit:
@@ -120,7 +120,7 @@ def main():
 
     peft_config = LoraConfig(**lora_kwargs)
 
-    dataset = load_dataset("asdf2k/caption_translation")
+    dataset = load_dataset(args.dataset)
     dev_ds = dataset["dev"].map(partial(add_chat_text, processor=processor))
     train_ds = dataset["train"].map(partial(add_chat_text, processor=processor))