Update app.py

app.py

@@ -32,6 +32,7 @@ def fine_tune_model(base_model_name, dataset_name):
 
     print("### Dataset")
     print(dataset)
+    print(dataset.head())
     print("###")
     
     # Load model
@@ -44,43 +45,46 @@ def fine_tune_model(base_model_name, dataset_name):
     print("###")
     
     # Pre-process dataset
+    
     def preprocess(examples):
-        model_inputs = tokenizer(examples["sql_prompt"], text_target=examples["sql"], max_length=512, padding="max_length", truncation=True)
+        model_inputs = tokenizer(examples["sql_prompt"], text_target=examples["sql"]) #, max_length=512, padding="max_length", truncation=True)
         return model_inputs
+        
     dataset = dataset.map(preprocess, batched=True)
 
     print("### Pre-processed dataset")
     print(dataset)
+    print(dataset.head())
     print("###")
     
     # Split dataset into training and validation sets
-    train_dataset = dataset["train"].shuffle(seed=42).select(range(1000))
-    test_dataset = dataset["test"].shuffle(seed=42).select(range(100))
+    train_dataset = dataset["train"] #.shuffle(seed=42).select(range(1000))
+    test_dataset = dataset["test"] #.shuffle(seed=42).select(range(100))
 
     print("### Training dataset")
-    print(test_dataset)
+    print(train_dataset)
     print("### Validation dataset")
     print(test_dataset)
     print("###")
     
     # Configure training arguments
     training_args = Seq2SeqTrainingArguments(
-        output_dir="./results",
-        logging_dir="./logs",
+        output_dir="./output",
+        logging_dir="./logging",
         num_train_epochs=1,
-        per_device_train_batch_size=16,
-        per_device_eval_batch_size=64,
-        eval_strategy="steps",
-        save_total_limit=2,
-        save_steps=500,
-        eval_steps=500,
-        warmup_steps=500,
-        weight_decay=0.01,
-        metric_for_best_model="accuracy",
-        greater_is_better=True,
-        load_best_model_at_end=True,
-        push_to_hub=True,
-        save_on_each_node=True,
+        #per_device_train_batch_size=16,
+        #per_device_eval_batch_size=64,
+        #eval_strategy="steps",
+        #save_total_limit=2,
+        #save_steps=500,
+        #eval_steps=500,
+        #warmup_steps=500,
+        #weight_decay=0.01,
+        #metric_for_best_model="accuracy",
+        #greater_is_better=True,
+        #load_best_model_at_end=True,
+        #push_to_hub=True,
+        #save_on_each_node=True,
     )
 
     print("### Training arguments")
@@ -93,13 +97,9 @@ def fine_tune_model(base_model_name, dataset_name):
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=test_dataset,
-        compute_metrics=lambda pred: {"accuracy": torch.sum(pred.label_ids == pred.predictions.argmax(-1))},
+        #compute_metrics=lambda pred: {"accuracy": torch.sum(pred.label_ids == pred.predictions.argmax(-1))},
     )
 
-    print("### Trainer")
-    print(trainer)
-    print("###")
-    
     # Train and save model
     #trainer.train()
     #trainer.save_model()
@@ -128,8 +128,8 @@ def prompt_model(model_name, system_prompt, user_prompt, sql_schema):
 
 def load_model(model_name):
     model = AutoModelForCausalLM.from_pretrained(model_name, device_map="auto")
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    tokenizer.pad_token = tokenizer.eos_token
+    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL_NAME)
+    #tokenizer.pad_token = tokenizer.eos_token
     
     return model, tokenizer