Working e2e finetunig for two tasks

postgresml · santiatpml · Mar 26, 2024 · Jan 31, 2024 · Feb 1, 2024 · Feb 6, 2024
commit 097a8cf00116d2e1f14dd0f8046024d1c4c0b33c
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
@@ -906,25 +906,27 @@ fn tune(
         LIMIT 1;",
         vec![(PgBuiltInOids::TEXTOID.oid(), project_name.into_datum())],
     );
-
+    
     let mut deploy = true;
     match automatic_deploy {
         // Deploy only if metrics are better than previous model.
         Some(true) | None => {
             if let Ok(Some(deployed_metrics)) = deployed_metrics {
                 let deployed_metrics = deployed_metrics.0.as_object().unwrap();
-                if project.task.value_is_better(
-                    deployed_metrics
-                        .get(&project.task.default_target_metric())
-                        .unwrap()
-                        .as_f64()
-                        .unwrap(),
-                    new_metrics
-                        .get(&project.task.default_target_metric())
-                        .unwrap()
-                        .as_f64()
-                        .unwrap(),
-                ) {
+
+                let deployed_value = deployed_metrics
+                    .get(&project.task.default_target_metric())
+                    .and_then(|value| value.as_f64())
+                    .unwrap_or_default(); // Default to 0.0 if the key is not present or conversion fails
+
+                // Get the value for the default target metric from new_metrics or provide a default value
+                let new_value = new_metrics
+                    .get(&project.task.default_target_metric())
+                    .and_then(|value| value.as_f64())
+                    .unwrap_or_default(); // Default to 0.0 if the key is not present or conversion fails
+
+
+                if project.task.value_is_better(deployed_value, new_value){
                     deploy = false;
                 }
             }

diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -1052,26 +1052,14 @@ def tokenize_function(example):
     # Generate tokens
     train_tokenized_datasets = train_dataset.map(tokenize_function, batched=True)
     test_tokenized_datasets = test_dataset.map(tokenize_function, batched=True)
-    log.info("Tokenization done")
-    log.info("Train dataset")
-    log.info(train_tokenized_datasets[0:2])
-    log.info("Test dataset")
-    log.info(test_tokenized_datasets[0:2])
+
     # Data collator
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
     # Training Args
-    log.info("Training args setup started path=%s"%path)
     training_args=TrainingArguments(output_dir=path, logging_dir=path, **hyperparams["training_args"])
 
-    log.info("Trainer setup done")
     # Trainer
-    log.info(model)
-    log.info(training_args)
-    log.info(train_tokenized_datasets)
-    log.info(test_tokenized_datasets)
-    log.info(tokenizer)
-    log.info(data_collator)
     try:
         trainer = Trainer(
             model=model,
@@ -1083,15 +1071,16 @@ def tokenize_function(example):
         )
     except Exception as e:
         log.error(e)
-    log.info("Training started")
-
+
     # Train
     trainer.train()
 
     # Save model
     trainer.save_model()
 
-    metrics = {"loss" : 0.0}
+    # TODO: compute real metrics
+    metrics = {"loss" : 0.0, "f1": 1.0}
+
     return metrics
 
 def finetune_text_pair_classification(task, hyperparams, path, text1_train, text1_test, text2_train, text2_test, class_train, class_test):
@@ -1147,42 +1136,30 @@ def tokenize_function(example):
     # Generate tokens
     train_tokenized_datasets = train_dataset.map(tokenize_function, batched=True)
     test_tokenized_datasets = test_dataset.map(tokenize_function, batched=True)
-    log.info("Tokenization done")
-    log.info("Train dataset")
-    log.info(train_tokenized_datasets[0:2])
-    log.info("Test dataset")
-    log.info(test_tokenized_datasets[0:2])
+
     # Data collator
     data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
 
     # Training Args
-    log.info("Training args setup started path=%s"%path)
     training_args=TrainingArguments(output_dir=path, logging_dir=path, **hyperparams["training_args"])
-    log.info("Trainer setup done")
+
     # Trainer
-    log.info(model)
-    log.info(training_args)
-    log.info(train_tokenized_datasets)
-    log.info(test_tokenized_datasets)
-    log.info(tokenizer)
-    log.info(data_collator)
-    try:
-        trainer = Trainer(
-            model=model,
-            args=training_args,
-            train_dataset=train_tokenized_datasets,
-            eval_dataset=test_tokenized_datasets,
-            tokenizer=tokenizer,
-            data_collator=data_collator,
-        )
-    except Exception as e:
-        log.error(e)
-    log.info("Training started")
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_tokenized_datasets,
+        eval_dataset=test_tokenized_datasets,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+    )
 
     # Train
     trainer.train()
 
     # Save model
     trainer.save_model()
-    metrics = {"loss" : 0.0}
+
+    # TODO: Get real metrics
+    metrics = {"loss" : 0.0, "f1": 1.0}
+
     return metrics
diff --git a/pgml-extension/src/orm/model.rs b/pgml-extension/src/orm/model.rs
@@ -242,12 +242,6 @@ impl Model {
 
         model.metrics = Some(JsonB(json!(metrics)));
         info!("Metrics: {:?}", &metrics);
-        // let metrics = match transformers::finetune(&project.task, dataset, &model.hyperparams, &path) {
-        //     Ok(metrics) => metrics,
-        //     Err(e) => error!("{e}"),
-        // };
-        // model.metrics = Some(JsonB(json!(metrics)));
-        // info!("Metrics: {:?}", &metrics);
 
         Spi::get_one_with_args::<i64>(
             "UPDATE pgml.models SET hyperparams = $1, metrics = $2 WHERE id = $3 RETURNING id",
@@ -266,26 +260,31 @@ impl Model {
         .unwrap();
 
         // Save the bindings.
-        /*for entry in std::fs::read_dir(&path).unwrap() {
+        for entry in std::fs::read_dir(&path).unwrap() {
             let path = entry.unwrap().path();
-            let bytes = std::fs::read(&path).unwrap();
-            for (i, chunk) in bytes.chunks(100_000_000).enumerate() {
-                Spi::get_one_with_args::<i64>(
-                    "INSERT INTO pgml.files (model_id, path, part, data) VALUES($1, $2, $3, $4) RETURNING id",
-                    vec![
-                        (PgBuiltInOids::INT8OID.oid(), model.id.into_datum()),
-                        (
-                            PgBuiltInOids::TEXTOID.oid(),
-                            path.file_name().unwrap().to_str().into_datum(),
-                        ),
-                        (PgBuiltInOids::INT8OID.oid(), (i as i64).into_datum()),
-                        (PgBuiltInOids::BYTEAOID.oid(), chunk.into_datum()),
-                    ],
-                )
-                .unwrap();
-            }
-        }*/
 
+            if path.is_file() {
+
+                let bytes = std::fs::read(&path).unwrap();
+
+                for (i, chunk) in bytes.chunks(100_000_000).enumerate() {
+                    Spi::get_one_with_args::<i64>(
+                        "INSERT INTO pgml.files (model_id, path, part, data) VALUES($1, $2, $3, $4) RETURNING id",
+                        vec![
+                            (PgBuiltInOids::INT8OID.oid(), model.id.into_datum()),
+                            (
+                                PgBuiltInOids::TEXTOID.oid(),
+                                path.file_name().unwrap().to_str().into_datum(),
+                            ),
+                            (PgBuiltInOids::INT8OID.oid(), (i as i64).into_datum()),
+                            (PgBuiltInOids::BYTEAOID.oid(), chunk.into_datum()),
+                        ],
+                    )
+                    .unwrap();
+                }
+            }
+        }
+
         Spi::run_with_args(
             "UPDATE pgml.models SET status = $1::pgml.status WHERE id = $2",
             Some(vec![
@@ -297,6 +296,7 @@ impl Model {
             ]),
         )
         .unwrap();
+
         model
     }