OpenAI runtime (#944)

levkk · SilasMarvin · web-flow · commit 56cb69b69af4 · 2023-08-23T14:44:07.000-07:00
Co-authored-by: SilasMarvin &lt;19626586+SilasMarvin@users.noreply.github.com&gt;
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(),
         }
     }
 }

Original file line number	Diff line number	Diff line change
`@@ -6,6 +6,7 @@ use serde::Deserialize;`
`6`	`6`	`pub enum Runtime {`
`7`	`7`	`python,`
`8`	`8`	`rust,`
	`9`	`+ openai,`
`9`	`10`	`}`
`10`	`11`
`11`	`12`	`impl std::str::FromStr for Runtime {`
`@@ -15,6 +16,7 @@ impl std::str::FromStr for Runtime {`
`15`	`16`	`match input {`
`16`	`17`	`"python" => Ok(Runtime::python),`
`17`	`18`	`"rust" => Ok(Runtime::rust),`
	`19`	`+ "openai" => Ok(Runtime::openai),`
`18`	`20`	`_ => Err(()),`
`19`	`21`	`}`
`20`	`22`	`}`
`@@ -25,6 +27,7 @@ impl std::string::ToString for Runtime {`
`25`	`27`	`match *self {`
`26`	`28`	`Runtime::python => "python".to_string(),`
`27`	`29`	`Runtime::rust => "rust".to_string(),`
	`30`	`+ Runtime::openai => "openai".to_string(),`
`28`	`31`	`}`
`29`	`32`	`}`
`30`	`33`	`}`
Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,7 @@ pub enum Task {`
`13`	`13`	`text_generation,`
`14`	`14`	`text2text,`
`15`	`15`	`cluster,`
	`16`	`+ embedding,`
`16`	`17`	`}`
`17`	`18`
`18`	`19`	`// unfortunately the pgrx macro expands the enum names to underscore, but huggingface uses dash`
`@@ -28,6 +29,7 @@ impl Task {`
`28`	`29`	`Task::text_generation => "text_generation".to_string(),`
`29`	`30`	`Task::text2text => "text2text".to_string(),`
`30`	`31`	`Task::cluster => "cluster".to_string(),`
	`32`	`+ Task::embedding => "embedding".to_string(),`
`31`	`33`	`}`
`32`	`34`	`}`
`33`	`35`
`@@ -46,6 +48,7 @@ impl Task {`
`46`	`48`	`Task::text_generation => "perplexity",`
`47`	`49`	`Task::text2text => "perplexity",`
`48`	`50`	`Task::cluster => "silhouette",`
	`51`	`+ Task::embedding => error!("No default target metric for embedding task")`
`49`	`52`	`}`
`50`	`53`	`.to_string()`
`51`	`54`	`}`
`@@ -61,6 +64,7 @@ impl Task {`
`61`	`64`	`Task::text_generation => false,`
`62`	`65`	`Task::text2text => false,`
`63`	`66`	`Task::cluster => true,`
	`67`	`+ Task::embedding => error!("No default target metric positive for embedding task")`
`64`	`68`	`}`
`65`	`69`	`}`
`66`	`70`
`@@ -117,6 +121,7 @@ impl std::string::ToString for Task {`
`117`	`121`	`Task::text_generation => "text-generation".to_string(),`
`118`	`122`	`Task::text2text => "text2text".to_string(),`
`119`	`123`	`Task::cluster => "cluster".to_string(),`
	`124`	`+ Task::embedding => "embedding".to_string(),`
`120`	`125`	`}`
`121`	`126`	`}`
`122`	`127`	`}`