Mar 18, 2024 · Mar 15, 2024
diff --git a/pgml-extension/src/bindings/mod.rs b/pgml-extension/src/bindings/mod.rs
 use anyhow::{anyhow, Result};
 #[allow(unused_imports)] // used for test macros
 use pgrx::*;
 use pyo3::{PyResult, Python};
 use pyo3::{pyfunction,PyResult, Python};

 use crate::orm::*;

 #[pyfunction]
 fn r_insert_logs(project_id: i64, model_id: i64, logs: String) -> PyResult<String> {
    let id_value = Spi::get_one_with_args::<i64>(
        "INSERT INTO pgml.logs (project_id, model_id, logs) VALUES ($1, $2, $3::JSONB) RETURNING id;",
        vec![
            (PgBuiltInOids::INT8OID.oid(), project_id.into_datum()),
            (PgBuiltInOids::INT8OID.oid(), model_id.into_datum()),
            (PgBuiltInOids::TEXTOID.oid(), logs.into_datum()),
        ],
    )
    .unwrap()
    .unwrap();
    Ok(format!("Inserted logs with id: {}", id_value))
 }

 #[pyfunction]
 fn r_print_info(info: String) -> PyResult<String> {
    info!("{}", info);
    Ok(info)
 }

 #[cfg(feature = "python")]
 #[macro_export]
 macro_rules! create_pymodule {
                pyo3::Python::with_gil(|py| -> anyhow::Result<pyo3::Py<pyo3::types::PyModule>> {
                    use $crate::bindings::TracebackError;
                    let src = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), $pyfile));
 Ok(
 pyo3::types::PyModule::from_code(py, src, "transformers.py", "__main__")
        .format_traceback(py)?
        .into(),
                    )
 let module = pyo3::types::PyModule::from_code(py, src, "transformers.py", "__main__")
 .format_traceback(py)?;
 module.add_function(wrap_pyfunction!($crate::bindings::r_insert_logs, module)?)?;
 module.add_function(wrap_pyfunction!($crate::bindings::r_print_info, module)?)?;
 Ok(module.into())
                })
            });
    };
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
 from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
 from trl.trainer import ConstantLengthDataset
 from peft import LoraConfig, get_peft_model
 from pypgrx import print_info, insert_logs
 from abc import abstractmethod

 transformers.logging.set_verbosity_info()
            logs["step"] = state.global_step
            logs["max_steps"] = state.max_steps
            logs["timestamp"] = str(datetime.now())
            print_info(json.dumps(logs, indent=4))
            insert_logs(self.project_id, self.model_id, json.dumps(logs))
            r_print_info(json.dumps(logs, indent=4))


 class FineTuningBase:
                trainable_model_params += param.numel()

        # Calculate and print the number and percentage of trainable parameters
 print_info(f"Trainable model parameters: {trainable_model_params}")
 print_info(f"All model parameters: {all_model_params}")
 print_info(
 r_print_info(f"Trainable model parameters: {trainable_model_params}")
 r_print_info(f"All model parameters: {all_model_params}")
 r_print_info(
            f"Percentage of trainable model parameters: {100 * trainable_model_params / all_model_params:.2f}%"
        )

                "bias": "none",
                "task_type": "CAUSAL_LM",
            }
 print_info(
 r_print_info(
                "LoRA configuration are not set. Using default parameters"
                + json.dumps(self.lora_config_params)
            )
            peft_config=LoraConfig(**self.lora_config_params),
            callbacks=[PGMLCallback(self.project_id, self.model_id)],
        )
 print_info("Creating Supervised Fine Tuning trainer done. Training ... ")
 r_print_info("Creating Supervised Fine Tuning trainer done. Training ... ")

        # Train
        self.trainer.train()
Original file line number	Diff line number	Diff line change
Expand Up		@@ -3,10 +3,31 @@ use std::fmt::Debug;
		use anyhow::{anyhow, Result};
		#[allow(unused_imports)] // used for test macros
		use pgrx::*;
		use pyo3::{PyResult, Python};
		use pyo3::{pyfunction,PyResult, Python};

		use crate::orm::*;

		#[pyfunction]
		fn r_insert_logs(project_id: i64, model_id: i64, logs: String) -> PyResult<String> {
		let id_value = Spi::get_one_with_args::<i64>(
		"INSERT INTO pgml.logs (project_id, model_id, logs) VALUES ($1, $2, $3::JSONB) RETURNING id;",
		vec![
		(PgBuiltInOids::INT8OID.oid(), project_id.into_datum()),
		(PgBuiltInOids::INT8OID.oid(), model_id.into_datum()),
		(PgBuiltInOids::TEXTOID.oid(), logs.into_datum()),
		],
		)
		.unwrap()
		.unwrap();
		Ok(format!("Inserted logs with id: {}", id_value))
		}

		#[pyfunction]
		fn r_print_info(info: String) -> PyResult<String> {
		info!("{}", info);
		Ok(info)
		}

		#[cfg(feature = "python")]
		#[macro_export]
		macro_rules! create_pymodule {
Expand All		@@ -16,11 +37,11 @@ macro_rules! create_pymodule {
		pyo3::Python::with_gil(\|py\| -> anyhow::Result<pyo3::Py<pyo3::types::PyModule>> {
		use $crate::bindings::TracebackError;
		let src = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), $pyfile));
		Ok(
		pyo3::types::PyModule::from_code(py, src, "transformers.py", "__main__")
		.format_traceback(py)?
		.into(),
		)
		let module = pyo3::types::PyModule::from_code(py, src, "transformers.py", "__main__")
		.format_traceback(py)?;
		module.add_function(wrap_pyfunction!($crate::bindings::r_insert_logs, module)?)?;
		module.add_function(wrap_pyfunction!($crate::bindings::r_print_info, module)?)?;
		Ok(module.into())
		})
		});
		};
Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -55,7 +55,6 @@
		from trl import SFTTrainer, DataCollatorForCompletionOnlyLM
		from trl.trainer import ConstantLengthDataset
		from peft import LoraConfig, get_peft_model
		from pypgrx import print_info, insert_logs
		from abc import abstractmethod

		transformers.logging.set_verbosity_info()
Expand DownExpand Up		@@ -1017,8 +1016,7 @@ def on_log(self, args, state, control, logs=None, **kwargs):
		logs["step"] = state.global_step
		logs["max_steps"] = state.max_steps
		logs["timestamp"] = str(datetime.now())
		print_info(json.dumps(logs, indent=4))
		insert_logs(self.project_id, self.model_id, json.dumps(logs))
		r_print_info(json.dumps(logs, indent=4))


		class FineTuningBase:
Expand DownExpand Up		@@ -1100,9 +1098,9 @@ def print_number_of_trainable_model_parameters(self, model):
		trainable_model_params += param.numel()

		# Calculate and print the number and percentage of trainable parameters
		print_info(f"Trainable model parameters: {trainable_model_params}")
		print_info(f"All model parameters: {all_model_params}")
		print_info(
		r_print_info(f"Trainable model parameters: {trainable_model_params}")
		r_print_info(f"All model parameters: {all_model_params}")
		r_print_info(
		f"Percentage of trainable model parameters: {100 * trainable_model_params / all_model_params:.2f}%"
		)

Expand DownExpand Up		@@ -1398,7 +1396,7 @@ def __init__(
		"bias": "none",
		"task_type": "CAUSAL_LM",
		}
		print_info(
		r_print_info(
		"LoRA configuration are not set. Using default parameters"
		+ json.dumps(self.lora_config_params)
		)
Expand DownExpand Up		@@ -1465,7 +1463,7 @@ def formatting_prompts_func(example):
		peft_config=LoraConfig(**self.lora_config_params),
		callbacks=[PGMLCallback(self.project_id, self.model_id)],
		)
		print_info("Creating Supervised Fine Tuning trainer done. Training ... ")
		r_print_info("Creating Supervised Fine Tuning trainer done. Training ... ")

		# Train
		self.trainer.train()
Expand Down