Working non streaming open source ai replacement

SilasMarvin · SilasMarvin · commit 23ef2a336aeb · 2023-11-29T14:59:21.000-08:00
diff --git a/pgml-extension/src/bindings/transformers/transform.rs b/pgml-extension/src/bindings/transformers/transform.rs
@@ -32,7 +32,9 @@ impl Iterator for TransformStreamIterator {
             if res.is_none() {
                 Ok(None)
             } else {
+                eprintln!("\nHERE WE ARE!\n");
                 let res: Vec<String> = res.extract()?;
+                eprintln!("\nYUP WE DIDNT GET HERE\n");
                 Ok(Some(JsonB(serde_json::to_value(res).unwrap())))
             }
         })
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
@@ -126,6 +126,7 @@ def __iter__(self):
 
     def __next__(self):
         value = self.text_queue.get(timeout=self.timeout)
+        print("\n\n", value, "\n\n", file=sys.stderr)
         if value != self.stop_signal:
             return value
 
diff --git a/pgml-sdks/pgml/src/languages/javascript.rs b/pgml-sdks/pgml/src/languages/javascript.rs
@@ -95,8 +95,10 @@ fn transform_stream_iterate_next(mut cx: FunctionContext) -> JsResult<JsPromise>
             .try_settle_with(&channel, move |mut cx| {
                 let o = cx.empty_object();
                 if let Some(v) = v {
-                    let v: String = v.expect("Error calling next on TransformerStream");
-                    let v = cx.string(v);
+                    let v: Json = v.expect("Error calling next on TransformerStream");
+                    let v = v
+                        .into_js_result(&mut cx)
+                        .expect("Error converting rust Json to JavaScript Object");
                     let d = cx.boolean(false);
                     o.set(&mut cx, "value", v)
                         .expect("Error setting object value in transformer_sream_iterate_next");
diff --git a/pgml-sdks/pgml/src/languages/python.rs b/pgml-sdks/pgml/src/languages/python.rs
@@ -72,7 +72,7 @@ impl TransformerStreamPython {
             if let Some(o) = ts.next().await {
                 Ok(Some(Python::with_gil(|py| {
                     o.expect("Error calling next on TransformerStream")
-                        .to_object(py)
+                        .into_py(py)
                 })))
             } else {
                 Err(pyo3::exceptions::PyStopAsyncIteration::new_err(
diff --git a/pgml-sdks/pgml/src/open_source_ai.rs b/pgml-sdks/pgml/src/open_source_ai.rs
@@ -49,21 +49,17 @@ impl OpenSourceAI {
         Self { database_url }
     }
 
-    pub async fn chat_completions_create_async(
+    fn create_pipeline_model_name_parameters(
         &self,
         mut model: Json,
-        messages: Json,
-        max_tokens: Option<i32>,
-        temperature: Option<f64>,
-        n: Option<i32>,
-    ) -> anyhow::Result<Json> {
-        let (transformer_pipeline, model_name, model_parameters) = if model.is_object() {
+    ) -> anyhow::Result<(TransformerPipeline, String, Json)> {
+        if model.is_object() {
             let args = model.as_object_mut().unwrap();
             let model_name = args
                 .remove("model")
                 .context("`model` is a required key in the model object")?;
             let model_name = model_name.as_str().context("`model` must be a string")?;
-            (
+            Ok((
                 TransformerPipeline::new(
                     "conversational",
                     Some(model_name.to_string()),
@@ -72,7 +68,7 @@ impl OpenSourceAI {
                 ),
                 model_name.to_string(),
                 model,
-            )
+            ))
         } else {
             let model_name = model
                 .as_str()
@@ -83,7 +79,7 @@ impl OpenSourceAI {
 mistralai/Mistral-7B-v0.1
 "#,
                 )?;
-            (
+            Ok((
                 TransformerPipeline::new(
                     "conversational",
                     Some(real_model_name.to_string()),
@@ -92,36 +88,52 @@ mistralai/Mistral-7B-v0.1
                 ),
                 model_name.to_string(),
                 parameters,
-            )
-        };
+            ))
+        }
+    }
+
+    pub async fn chat_completions_create_stream_async(
+        &self,
+        model: Json,
+        messages: Vec<Json>,
+        max_tokens: Option<i32>,
+        temperature: Option<f64>,
+        n: Option<i32>,
+    ) -> anyhow::Result<()> {
+        Ok(())
+    }
+
+    pub async fn chat_completions_create_async(
+        &self,
+        model: Json,
+        messages: Vec<Json>,
+        max_tokens: Option<i32>,
+        temperature: Option<f64>,
+        n: Option<i32>,
+    ) -> anyhow::Result<Json> {
+        let (transformer_pipeline, model_name, model_parameters) =
+            self.create_pipeline_model_name_parameters(model)?;
 
         let max_tokens = max_tokens.unwrap_or(1000);
         let temperature = temperature.unwrap_or(0.8);
         let n = n.unwrap_or(1) as usize;
-        let to_hash = format!(
-            "{}{}{}{}",
-            model_parameters.to_string(),
-            max_tokens,
-            temperature,
-            n
-        );
+        let to_hash = format!("{}{}{}{}", *model_parameters, max_tokens, temperature, n);
         let md5_digest = md5::compute(to_hash.as_bytes());
         let fingerprint = uuid::Uuid::from_slice(&md5_digest.0)?;
 
-        let messages: Vec<Json> = std::iter::repeat(messages).take(n).collect();
         let choices = transformer_pipeline
             .transform(
                 messages,
                 Some(
-                    serde_json::json!({ "max_length": max_tokens, "temperature": temperature })
+                    serde_json::json!({ "max_length": max_tokens, "temperature": temperature, "do_sample": true, "num_return_sequences": n })
                         .into(),
                 ),
             )
             .await?;
         let choices: Vec<Json> = choices
             .as_array()
             .context("Error parsing return from TransformerPipeline")?
-            .into_iter()
+            .iter()
             .enumerate()
             .map(|(i, c)| {
                 serde_json::json!({
@@ -157,7 +169,7 @@ mistralai/Mistral-7B-v0.1
     pub fn chat_completions_create(
         &self,
         model: Json,
-        messages: Json,
+        messages: Vec<Json>,
         max_tokens: Option<i32>,
         temperature: Option<f64>,
         n: Option<i32>,
@@ -177,14 +189,14 @@ mistralai/Mistral-7B-v0.1
 mod tests {
     use super::*;
 
-    #[sqlx::test]
-    async fn can_open_source_ai_create() -> anyhow::Result<()> {
+    #[test]
+    fn can_open_source_ai_create() -> anyhow::Result<()> {
         let client = OpenSourceAI::new(None);
-        let results = client.chat_completions_create(Json::from_serializable("mistralai/Mistral-7B-v0.1"), serde_json::json!([
-          {"role": "system", "content": "You are a friendly chatbot who always responds in the style of a pirate"},
-          {"role": "user", "content": "How many helicopters can a human eat in one sitting?"}
-        ]).into(), Some(1000), None, None)?;
-        assert!(results.as_array().is_some());
+        let results = client.chat_completions_create(Json::from_serializable("mistralai/Mistral-7B-v0.1"), vec![
+          serde_json::json!({"role": "system", "content": "You are a friendly chatbot who always responds in the style of a pirate"}).into(),
+          serde_json::json!({"role": "user", "content": "How many helicopters can a human eat in one sitting?"}).into(),
+        ], Some(1000), None, Some(3))?;
+        assert!(results["choices"].as_array().is_some());
         Ok(())
     }
 }
diff --git a/pgml-sdks/pgml/src/transformer_pipeline.rs b/pgml-sdks/pgml/src/transformer_pipeline.rs
@@ -55,7 +55,7 @@ impl TransformerStream {
 }
 
 impl Stream for TransformerStream {
-    type Item = Result<String, sqlx::Error>;
+    type Item = Result<Json, sqlx::Error>;
 
     fn poll_next(
         mut self: Pin<&mut Self>,
@@ -106,7 +106,7 @@ impl Stream for TransformerStream {
 
         if !self.results.is_empty() {
             let r = self.results.pop_front().unwrap();
-            Poll::Ready(Some(Ok(r.get::<String, _>(0))))
+            Poll::Ready(Some(Ok(r.get::<Json, _>(0))))
         } else if self.done {
             Poll::Ready(None)
         } else {
@@ -251,10 +251,10 @@ mod tests {
         internal_init_logger(None, None).ok();
         let t = TransformerPipeline::new(
             "text-generation",
-            Some("TheBloke/zephyr-7B-beta-GGUF".to_string()),
+            Some("TheBloke/zephyr-7B-beta-GPTQ".to_string()),
             Some(
                 serde_json::json!({
-                  "model_file": "zephyr-7b-beta.Q5_K_M.gguf", "model_type": "mistral"
+                  "model_type": "mistral", "revision": "main", "device_map": "auto"
                 })
                 .into(),
             ),
