Update app.py

app.py

@@ -35,9 +35,10 @@ def prompt_model(model_id, system_prompt, user_prompt, schema):
     return result
 
 def fine_tune_model(base_model_id, dataset):
-    tokenizer = download_model(base_model_id)
-    #download_dataset(dataset)
-    fine_tuned_model_id = upload_model(base_model_id, tokenizer)
+    #tokenizer = download_model(base_model_id)
+    download_dataset(dataset)
+    #fine_tuned_model_id = upload_model(base_model_id, tokenizer)
+    fine_tuned_model_id = base_model_id # DELETE
     return fine_tuned_model_id
         
 def download_model(base_model_id):
@@ -46,9 +47,32 @@ def download_model(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):
-    ds = load_dataset(dataset)
-    return ""
+    dataset = load_dataset("b-mc2/sql-create-context", 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")
+    ###
+    dataset = load_dataset("json", data_files="train_dataset.json", split="train")
+    return "Done"
 
 def upload_model(base_model_id, tokenizer):
     fine_tuned_model_id = replace_hf_profile(base_model_id)