Update app.py

app.py

@@ -6,7 +6,7 @@ from peft import LoraConfig
 from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, pipeline
 from trl import SFTTrainer, setup_chat_format
 
-# Fine-tune on NVidia A10G Large (sleep after 1 hour)
+# Fine-tune on NVidia 4xL4 (sleep after 10 hours)
 
 hf_profile = "bstraehle"
 
@@ -38,11 +38,42 @@ def prompt_model(model_id, system_prompt, user_prompt, schema):
 
 def fine_tune_model(base_model_id, dataset):
     #tokenizer = download_model(base_model_id)
-    download_dataset(dataset)
+    prepare_dataset(dataset)
     train_model(base_model_id)
     #fine_tuned_model_id = upload_model(base_model_id, tokenizer)
     return "fine_tuned_model_id"
 
+def download_model(base_model_id):
+    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
+    model = AutoModelForCausalLM.from_pretrained(base_model_id)
+    model.save_pretrained(base_model_id)
+    return tokenizer
+
+def create_conversation(sample):
+  return {
+    "messages": [
+      {"role": "system", "content": system_prompt.format(schema=sample["context"])},
+      {"role": "user", "content": sample["question"]},
+      {"role": "assistant", "content": sample["answer"]}
+    ]
+  }
+    
+def prepare_dataset(dataset):
+    dataset = load_dataset(dataset, split="train")
+    dataset = dataset.shuffle().select(range(12500))
+     
+    # Convert dataset to OAI messages
+    dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
+    # split dataset into 10,000 training samples and 2,500 test samples
+    dataset = dataset.train_test_split(test_size=2500/12500)
+     
+    print(dataset["train"][345]["messages"])
+     
+    # save datasets to disk
+    dataset["train"].to_json("train_dataset.json", orient="records")
+    dataset["test"].to_json("test_dataset.json", orient="records")
+    ###
+
 def train_model(model_id):
     print("111")
     dataset = load_dataset("json", data_files="train_dataset.json", split="train")
@@ -124,38 +155,7 @@ def train_model(model_id):
     del model
     del trainer
     torch.cuda.empty_cache()
-
-def download_model(base_model_id):
-    tokenizer = AutoTokenizer.from_pretrained(base_model_id)
-    model = AutoModelForCausalLM.from_pretrained(base_model_id)
-    model.save_pretrained(base_model_id)
-    return tokenizer
-
-def create_conversation(sample):
-  return {
-    "messages": [
-      {"role": "system", "content": system_prompt.format(schema=sample["context"])},
-      {"role": "user", "content": sample["question"]},
-      {"role": "assistant", "content": sample["answer"]}
-    ]
-  }
     
-def download_dataset(dataset):
-    dataset = load_dataset(dataset, split="train")
-    dataset = dataset.shuffle().select(range(12500))
-     
-    # Convert dataset to OAI messages
-    dataset = dataset.map(create_conversation, remove_columns=dataset.features,batched=False)
-    # split dataset into 10,000 training samples and 2,500 test samples
-    dataset = dataset.train_test_split(test_size=2500/12500)
-     
-    print(dataset["train"][345]["messages"])
-     
-    # save datasets to disk
-    dataset["train"].to_json("train_dataset.json", orient="records")
-    dataset["test"].to_json("test_dataset.json", orient="records")
-    ###
-
 def upload_model(base_model_id, tokenizer):
     fine_tuned_model_id = replace_hf_profile(base_model_id)
     login(token=os.environ["HF_TOKEN"])