Skip to content

Commit

Permalink
Working C
Browse files Browse the repository at this point in the history
  • Loading branch information
SilasMarvin committed Apr 25, 2024
1 parent c72cb30 commit 5d276fc
Show file tree
Hide file tree
Showing 14 changed files with 327 additions and 96 deletions.
8 changes: 8 additions & 0 deletions pgml-sdks/pgml/go/test.c
Expand Up @@ -33,5 +33,13 @@ int main() {
printf("Result %u -> %s\n", i, results[i]);
}

// Test the TransformerPipeline
TransformerPipelineC * t_pipeline = TransformerPipelineC_new("text-generation", "TheBloke/zephyr-7B-beta-GPTQ", "{\"revision\": \"main\"}", "postgres://pg:ml@sql.cloud.postgresml.org:38042/pgml");
GeneralJsonAsyncIteratorC * t_pipeline_iter = TransformerPipelineC_transform_stream(t_pipeline, "\"AI is going to\"", "{\"max_new_tokens\": 100}", NULL);
while (!GeneralJsonAsyncIteratorC_done(t_pipeline_iter)) {
char * res = GeneralJsonAsyncIteratorC_next(t_pipeline_iter);
printf("Token -> %s\n", res);
}

return 0;
}
8 changes: 5 additions & 3 deletions pgml-sdks/pgml/src/builtins.rs
Expand Up @@ -3,8 +3,7 @@ use sqlx::Row;
use tracing::instrument;

/// Provides access to builtin database methods
// #[derive(alias, Debug, Clone)]
#[derive(Debug, Clone)]
#[derive(alias, Debug, Clone)]
pub struct Builtins {
database_url: Option<String>,
}
Expand All @@ -14,7 +13,10 @@ use crate::{get_or_initialize_pool, query_runner::QueryRunner, types::Json};
#[cfg(feature = "python")]
use crate::{query_runner::QueryRunnerPython, types::JsonPython};

// #[alias_methods(new, query, transform)]
#[cfg(feature = "c")]
use crate::{languages::c::JsonC, query_runner::QueryRunnerC};

#[alias_methods(new, query, transform)]
impl Builtins {
pub fn new(database_url: Option<String>) -> Self {
Self { database_url }
Expand Down
82 changes: 82 additions & 0 deletions pgml-sdks/pgml/src/languages/c.rs
@@ -1,11 +1,15 @@
use crate::types::{DateTime, GeneralJsonAsyncIterator, GeneralJsonIterator, Json};
use futures::pin_mut;
use futures::stream::Stream;
use rust_bridge::c::CustomInto;
use std::pin::Pin;

pub type JsonC = std::ffi::c_char;

unsafe impl CustomInto<Json> for *mut JsonC {
unsafe fn custom_into(self) -> Json {
let s = std::ffi::CStr::from_ptr(self).to_str().unwrap();
eprintln!("\nABOU TO DECODE: {}\n", s);
serde_json::from_str::<serde_json::Value>(s).unwrap().into()
}
}
Expand All @@ -16,3 +20,81 @@ unsafe impl CustomInto<*mut JsonC> for Json {
std::ffi::CString::new(s).unwrap().into_raw()
}
}

#[repr(C)]
pub struct GeneralJsonIteratorC {
pub wrapped:
*mut std::iter::Peekable<Box<dyn Iterator<Item = Result<Json, anyhow::Error>> + Send>>,
}

unsafe impl CustomInto<*mut GeneralJsonIteratorC> for GeneralJsonIterator {
unsafe fn custom_into(self) -> *mut GeneralJsonIteratorC {
Box::into_raw(Box::new(GeneralJsonIteratorC {
wrapped: Box::into_raw(Box::new(self.0.peekable())),
}))
}
}

#[no_mangle]
pub unsafe extern "C" fn GeneralJsonIteratorC_done(iterator: *mut GeneralJsonIteratorC) -> bool {
let mut c = Box::leak(Box::from_raw(iterator));
if let Some(_) = (*c.wrapped).peek() {
false
} else {
true
}
}

#[no_mangle]
pub unsafe extern "C" fn GeneralJsonIteratorC_next(
iterator: *mut GeneralJsonIteratorC,
) -> *mut JsonC {
let c = Box::leak(Box::from_raw(iterator));
let b = Box::leak(Box::from_raw(c.wrapped));
(*b).next().unwrap().unwrap().custom_into()
}

#[repr(C)]
pub struct GeneralJsonAsyncIteratorC {
pub wrapped: *mut futures::stream::Peekable<
Pin<Box<dyn Stream<Item = Result<Json, anyhow::Error>> + Send>>,
>,
}

unsafe impl CustomInto<*mut GeneralJsonAsyncIteratorC> for GeneralJsonAsyncIterator {
unsafe fn custom_into(self) -> *mut GeneralJsonAsyncIteratorC {
use futures::stream::StreamExt;
Box::into_raw(Box::new(GeneralJsonAsyncIteratorC {
wrapped: Box::into_raw(Box::new(self.0.peekable())),
}))
}
}

#[no_mangle]
pub unsafe extern "C" fn GeneralJsonAsyncIteratorC_done(
iterator: *mut GeneralJsonAsyncIteratorC,
) -> bool {
crate::get_or_set_runtime().block_on(async move {
use futures::stream::StreamExt;
let c = Box::leak(Box::from_raw(iterator));
let s = Box::leak(Box::from_raw(c.wrapped));
let mut s = Pin::new(s);
let res = s.as_mut().peek_mut().await;
if let Some(res) = res {
false
} else {
true
}
})
}

#[no_mangle]
pub unsafe extern "C" fn GeneralJsonAsyncIteratorC_next(
iterator: *mut GeneralJsonAsyncIteratorC,
) -> *mut JsonC {
crate::get_or_set_runtime().block_on(async move {
use futures::stream::StreamExt;
let mut c = Box::leak(Box::from_raw(iterator));
(*c.wrapped).next().await.unwrap().unwrap().custom_into()
})
}
8 changes: 5 additions & 3 deletions pgml-sdks/pgml/src/model.rs
Expand Up @@ -11,6 +11,9 @@ use crate::{
#[cfg(feature = "python")]
use crate::types::JsonPython;

#[cfg(feature = "c")]
use crate::languages::c::JsonC;

/// A few notes on the following enums:
/// - Sqlx does provide type derivation for enums, but it's not very good
/// - Queries using these enums require a number of additional queries to get their oids and
Expand Down Expand Up @@ -52,8 +55,7 @@ pub(crate) struct ModelDatabaseData {
}

/// A model used for embedding, inference, etc...
// #[derive(alias, Debug, Clone)]
#[derive(Debug, Clone)]
#[derive(alias, Debug, Clone)]
pub struct Model {
pub(crate) name: String,
pub(crate) runtime: ModelRuntime,
Expand All @@ -67,7 +69,7 @@ impl Default for Model {
}
}

// #[alias_methods(new, transform)]
#[alias_methods(new, transform)]
impl Model {
/// Creates a new [Model]
pub fn new(name: Option<String>, source: Option<String>, parameters: Option<Json>) -> Self {
Expand Down
23 changes: 14 additions & 9 deletions pgml-sdks/pgml/src/open_source_ai.rs
Expand Up @@ -13,9 +13,14 @@ use crate::{
#[cfg(feature = "python")]
use crate::types::{GeneralJsonAsyncIteratorPython, GeneralJsonIteratorPython, JsonPython};

#[cfg(feature = "c")]
use crate::{
languages::c::JsonC,
languages::c::{GeneralJsonAsyncIteratorC, GeneralJsonIteratorC},
};

/// A drop in replacement for OpenAI
// #[derive(alias, Debug, Clone)]
#[derive(Debug, Clone)]
#[derive(alias, Debug, Clone)]
pub struct OpenSourceAI {
database_url: Option<String>,
}
Expand Down Expand Up @@ -163,13 +168,13 @@ impl Iterator for AsyncToSyncJsonIterator {
}
}

// #[alias_methods(
// new,
// chat_completions_create,
// chat_completions_create_async,
// chat_completions_create_stream,
// chat_completions_create_stream_async
// )]
#[alias_methods(
new,
chat_completions_create,
chat_completions_create_async,
chat_completions_create_stream,
chat_completions_create_stream_async
)]
impl OpenSourceAI {
/// Creates a new [OpenSourceAI]
///
Expand Down
5 changes: 4 additions & 1 deletion pgml-sdks/pgml/src/query_builder.rs
Expand Up @@ -12,14 +12,17 @@ use crate::{pipeline::Pipeline, types::Json, Collection};
#[cfg(feature = "python")]
use crate::{pipeline::PipelinePython, types::JsonPython};

#[cfg(feature = "c")]
use crate::{languages::c::JsonC, pipeline::PipelineC};

#[derive(alias, Clone, Debug)]
pub struct QueryBuilder {
collection: Collection,
query: Json,
pipeline: Option<Pipeline>,
}

#[alias_methods(limit, filter, vector_recall, to_full_string, fetch_all)]
#[alias_methods(limit, filter, vector_recall, to_full_string, fetch_all(skip = "C"))]
impl QueryBuilder {
pub fn new(collection: Collection) -> Self {
let query = json!({
Expand Down
24 changes: 13 additions & 11 deletions pgml-sdks/pgml/src/query_runner.rs
Expand Up @@ -8,6 +8,9 @@ use crate::{get_or_initialize_pool, types::Json};
#[cfg(feature = "python")]
use crate::types::JsonPython;

#[cfg(feature = "c")]
use crate::languages::c::JsonC;

#[derive(Clone, Debug)]
enum BindValue {
String(String),
Expand All @@ -17,23 +20,22 @@ enum BindValue {
Json(Json),
}

// #[derive(alias, Clone, Debug)]
#[derive(Clone, Debug)]
#[derive(alias, Clone, Debug)]
pub struct QueryRunner {
query: String,
bind_values: Vec<BindValue>,
database_url: Option<String>,
}

// #[alias_methods(
// fetch_all,
// execute,
// bind_string,
// bind_int,
// bind_float,
// bind_bool,
// bind_json
// )]
#[alias_methods(
fetch_all,
execute,
bind_string,
bind_int,
bind_float,
bind_bool,
bind_json
)]
impl QueryRunner {
pub fn new(query: &str, database_url: Option<String>) -> Self {
Self {
Expand Down
8 changes: 5 additions & 3 deletions pgml-sdks/pgml/src/splitter.rs
Expand Up @@ -11,6 +11,9 @@ use crate::{
#[cfg(feature = "python")]
use crate::types::JsonPython;

#[cfg(feature = "c")]
use crate::languages::c::JsonC;

#[allow(dead_code)]
#[derive(Debug, Clone)]
pub(crate) struct SplitterDatabaseData {
Expand All @@ -19,8 +22,7 @@ pub(crate) struct SplitterDatabaseData {
}

/// A text splitter
// #[derive(alias, Debug, Clone)]
#[derive(Debug, Clone)]
#[derive(alias, Debug, Clone)]
pub struct Splitter {
pub(crate) name: String,
pub(crate) parameters: Json,
Expand All @@ -33,7 +35,7 @@ impl Default for Splitter {
}
}

// #[alias_methods(new)]
#[alias_methods(new)]
impl Splitter {
/// Creates a new [Splitter]
///
Expand Down
8 changes: 5 additions & 3 deletions pgml-sdks/pgml/src/transformer_pipeline.rs
Expand Up @@ -4,8 +4,7 @@ use sqlx::Row;
use tracing::instrument;

/// Provides access to builtin database methods
// #[derive(alias, Debug, Clone)]
#[derive(Debug, Clone)]
#[derive(alias, Debug, Clone)]
pub struct TransformerPipeline {
task: Json,
database_url: Option<String>,
Expand All @@ -17,7 +16,10 @@ use crate::{get_or_initialize_pool, types::Json};
#[cfg(feature = "python")]
use crate::types::{GeneralJsonAsyncIteratorPython, JsonPython};

// #[alias_methods(new, transform, transform_stream)]
#[cfg(feature = "c")]
use crate::{languages::c::GeneralJsonAsyncIteratorC, languages::c::JsonC};

#[alias_methods(new, transform, transform_stream)]
impl TransformerPipeline {
/// Creates a new [TransformerPipeline]
///
Expand Down

0 comments on commit 5d276fc

Please sign in to comment.