Update app.py

app.py

@@ -4,15 +4,7 @@ from datasets import load_dataset
 from huggingface_hub import HfApi, login
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# Run on NVidia A10G Large (sleep after 1 hour)
-
-# Model IDs:
-#
-# meta-llama/Meta-Llama-3-8B-Instruct
-
-# Datasets:
-#
-# gretelai/synthetic_text_to_sql
+# Fine-tune on NVidia A10G Large (sleep after 1 hour)
 
 profile = "bstraehle"
 
@@ -24,7 +16,6 @@ user_prompt = "What is the total trade value and average price for each trader a
 schema = "CREATE TABLE trade_history (id INT, trader_id INT, stock VARCHAR(255), price DECIMAL(5,2), quantity INT, trade_time TIMESTAMP);"
 
 base_model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
-fine_tuned_model_id = "bstraehle/Meta-Llama-3-8B-Instruct"
 dataset = "gretelai/synthetic_text_to_sql"
 
 def prompt_model(model_id, system_prompt, user_prompt, schema):
@@ -40,16 +31,16 @@ def prompt_model(model_id, system_prompt, user_prompt, schema):
     
     return output[0]["generated_text"][-1]["content"]
 
-def fine_tune_model(model_id):
-    tokenizer = download_model(model_id)
-    model_repo_name = upload_model(model_id, tokenizer)
+def fine_tune_model(base_model_id):
+    tokenizer = download_model(base_model_id)
+    fine_tuned_model_id = upload_model(base_model_id, tokenizer)
 
-    return model_repo_name
+    return fine_tuned_model_id
         
-def download_model(model_id):
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)
-    model.save_pretrained(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
 
@@ -57,29 +48,31 @@ def download_model(model_id):
 #    ds = load_dataset(dataset)
 #    return ""
 
-def upload_model(model_id, tokenizer):
-    model_name = model_id[model_id.rfind('/')+1:]
-    model_repo_name = f"{profile}/{model_name}"
+def upload_model(base_model_id, tokenizer):
+    fine_tuned_model_id = replace_profile(base_model_id)
 
     login(token=os.environ["HF_TOKEN"])
 
     api = HfApi()
-    api.create_repo(repo_id=model_repo_name)
+    api.create_repo(repo_id=fine_tuned_model_id)
     api.upload_folder(
-        folder_path=model_id,
-        repo_id=model_repo_name
+        folder_path=base_model_id,
+        repo_id=fine_tuned_model_id)
     )
 
-    tokenizer.push_to_hub(model_repo_name)
+    tokenizer.push_to_hub(fine_tuned_model_id)
+
+    return fine_tuned_model_id
 
-    return model_repo_name
+def replace_profile(base_model_id):
+    model_id = base_model_id[base_model_id.rfind('/')+1:]
+    return f"{profile}/{model_id}"
 
 def process(action, base_model_id, dataset, system_prompt, user_prompt, schema):
     if action == action_1:
         result = fine_tune_model(base_model_id)
     elif action == action_2:
-        model_id = base_model_id[base_model_id.rfind('/')+1:]
-        fine_tuned_model_id = f"{profile}/{model_id}"
+        fine_tuned_model_id = replace_profile(base_model_id)
         result = prompt_model(fine_tuned_model_id, system_prompt, user_prompt, schema)
     
     return result