fixed bug

llm_toolkit/tune_logical_reasoning.py

@@ -6,8 +6,6 @@ from trl import SFTTrainer
 from transformers import TrainingArguments
 
 from dotenv import find_dotenv, load_dotenv
-from llm_toolkit.logical_reasoning_utils import *
-from llm_toolkit.llm_utils import *
 
 found_dotenv = find_dotenv(".env")
 
@@ -20,11 +18,14 @@ path = os.path.dirname(found_dotenv)
 print(f"Adding {path} to sys.path")
 sys.path.append(path)
 
+from llm_toolkit.logical_reasoning_utils import *
+from llm_toolkit.llm_utils import *
+
 model_name = os.getenv("MODEL_NAME")
 token = os.getenv("HF_TOKEN") or None
 load_in_4bit = os.getenv("LOAD_IN_4BIT") == "true"
-local_model = os.getenv("LOCAL_MODEL")
-hub_model = os.getenv("HUB_MODEL")
+local_model = os.getenv("LOCAL_MODEL") or "gemma-2-9b-it-lora"
+hub_model = os.getenv("HUB_MODEL") or "inflaton-ai/gemma-2-9b-it-lora"
 num_train_epochs = int(os.getenv("NUM_TRAIN_EPOCHS") or 0)
 data_path = os.getenv("LOGICAL_REASONING_DATA_PATH")
 results_path = os.getenv("LOGICAL_REASONING_RESULTS_PATH")
@@ -140,6 +141,9 @@ print(f"Peak reserved memory for training = {used_memory_for_lora} GB.")
 print(f"Peak reserved memory % of max memory = {used_percentage} %.")
 print(f"Peak reserved memory for training % of max memory = {lora_percentage} %.")
 
+model.save_pretrained(local_model)  # Local saving
+tokenizer.save_pretrained(local_model)
+
 print("Evaluating fine-tuned model: " + model_name)
 FastLanguageModel.for_inference(model)  # Enable native 2x faster inference
 predictions = eval_model(model, tokenizer, datasets["test"])
@@ -160,3 +164,6 @@ save_results(
 
 metrics = calc_metrics(datasets["test"]["label"], predictions, debug=True)
 print(metrics)
+
+model.push_to_hub(hub_model, token=token)  # Online saving
+tokenizer.push_to_hub(hub_model, token=token)  # Online saving