postgresml · SilasMarvin · Aug 23, 2023 · Aug 23, 2023 · Aug 23, 2023 · Aug 23, 2023
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
diff --git a/pgml-extension/Cargo.toml b/pgml-extension/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "pgml"
-version = "2.7.7"
+version = "2.7.8"
 edition = "2021"
 
 [lib]

diff --git a/pgml-extension/sql/pgml--2.7.7--2.7.8.sql b/pgml-extension/sql/pgml--2.7.7--2.7.8.sql
diff --git a/pgml-extension/sql/schema.sql b/pgml-extension/sql/schema.sql
@@ -82,7 +82,7 @@ SELECT pgml.auto_updated_at('pgml.snapshots');
 CREATE TABLE IF NOT EXISTS pgml.models(
 	id BIGSERIAL PRIMARY KEY,
 	project_id BIGINT NOT NULL,
-	snapshot_id BIGINT NOT NULL,
+	snapshot_id BIGINT,
 	num_features INT NOT NULL,
 	algorithm TEXT NOT NULL,
 	runtime pgml.runtime DEFAULT 'python'::pgml.runtime,

diff --git a/pgml-extension/src/orm/file.rs b/pgml-extension/src/orm/file.rs
@@ -115,6 +115,10 @@ pub fn find_deployed_estimator_by_model_id(model_id: i64) -> Result<Arc<Box<dyn
         Runtime::python => {
             anyhow::bail!("Python runtime not supported, recompile with `--features python`")
         }
+
+        Runtime::openai => {
+            error!("OpenAI runtime is not supported for training or inference");
+        }
     };
 
     // Cache the estimator in process memory.

diff --git a/pgml-extension/src/orm/model.rs b/pgml-extension/src/orm/model.rs
@@ -306,6 +306,10 @@ impl Model {
                 .unwrap();
 
                 let bindings: Box<dyn Bindings> = match runtime {
+                    Runtime::openai => {
+                        error!("OpenAI runtime is not supported for training or inference");
+                    }
+
                     Runtime::rust => {
                         match algorithm {
                             Algorithm::xgboost => {
@@ -396,6 +400,10 @@ impl Model {
 
     fn get_fit_function(&self) -> crate::bindings::Fit {
         match self.runtime {
+            Runtime::openai => {
+                error!("OpenAI runtime is not supported for training or inference");
+            }
+
             Runtime::rust => match self.project.task {
                 Task::regression => match self.algorithm {
                     Algorithm::xgboost => xgboost::fit_regression,

diff --git a/pgml-extension/src/orm/runtime.rs b/pgml-extension/src/orm/runtime.rs
@@ -6,6 +6,7 @@ use serde::Deserialize;
 pub enum Runtime {
     python,
     rust,
+    openai,
 }
 
 impl std::str::FromStr for Runtime {
@@ -15,6 +16,7 @@ impl std::str::FromStr for Runtime {
         match input {
             "python" => Ok(Runtime::python),
             "rust" => Ok(Runtime::rust),
+            "openai" => Ok(Runtime::openai),
             _ => Err(()),
         }
     }
@@ -25,6 +27,7 @@ impl std::string::ToString for Runtime {
         match *self {
             Runtime::python => "python".to_string(),
             Runtime::rust => "rust".to_string(),
+            Runtime::openai => "openai".to_string(),
         }
     }
 }
diff --git a/pgml-extension/src/orm/task.rs b/pgml-extension/src/orm/task.rs
@@ -13,6 +13,7 @@ pub enum Task {
     text_generation,
     text2text,
     cluster,
+    embedding,
 }
 
 // unfortunately the pgrx macro expands the enum names to underscore, but huggingface uses dash
@@ -28,6 +29,7 @@ impl Task {
             Task::text_generation => "text_generation".to_string(),
             Task::text2text => "text2text".to_string(),
             Task::cluster => "cluster".to_string(),
+            Task::embedding => "embedding".to_string(),
         }
     }
 
@@ -46,6 +48,7 @@ impl Task {
             Task::text_generation => "perplexity",
             Task::text2text => "perplexity",
             Task::cluster => "silhouette",
+            Task::embedding => error!("No default target metric for embedding task")
         }
         .to_string()
     }
@@ -61,6 +64,7 @@ impl Task {
             Task::text_generation => false,
             Task::text2text => false,
             Task::cluster => true,
+            Task::embedding => error!("No default target metric positive for embedding task")
         }
     }
 
@@ -117,6 +121,7 @@ impl std::string::ToString for Task {
             Task::text_generation => "text-generation".to_string(),
             Task::text2text => "text2text".to_string(),
             Task::cluster => "cluster".to_string(),
+            Task::embedding => "embedding".to_string(),
         }
     }
 }