Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit67ba48c

Browse files
authored
Updated to support streaming (#1151)
1 parent3e8cc28 commit67ba48c

File tree

11 files changed

+328
-135
lines changed

11 files changed

+328
-135
lines changed

‎pgml-sdks/pgml/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
[package]
22
name ="pgml"
3-
version ="0.9.5"
3+
version ="0.9.6"
44
edition ="2021"
55
authors = ["PosgresML <team@postgresml.org>"]
66
homepage ="https://postgresml.org/"

‎pgml-sdks/pgml/javascript/package.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
{
22
"name":"pgml",
3-
"version":"0.9.5",
3+
"version":"0.9.6",
44
"description":"Open Source Alternative for Building End-to-End Vector Search Applications without OpenAI & Pinecone",
55
"keywords": [
66
"postgres",

‎pgml-sdks/pgml/javascript/tests/typescript-tests/test.ts

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -280,6 +280,28 @@ it("can order documents", async () => {
280280
awaitcollection.archive();
281281
});
282282

283+
///////////////////////////////////////////////////
284+
// Transformer Pipeline Tests /////////////////////
285+
///////////////////////////////////////////////////
286+
287+
it("can transformer pipeline",async()=>{
288+
constt=pgml.newTransformerPipeline("text-generation");
289+
constit=awaitt.transform(["AI is going to"],{max_new_tokens:5});
290+
expect(it.length).toBeGreaterThan(0)
291+
});
292+
293+
it("can transformer pipeline stream",async()=>{
294+
constt=pgml.newTransformerPipeline("text-generation");
295+
constit=awaitt.transform_stream("AI is going to",{max_new_tokens:5});
296+
letresult=awaitit.next();
297+
letoutput=[];
298+
while(!result.done){
299+
output.push(result.value);
300+
result=awaitit.next();
301+
}
302+
expect(output.length).toBeGreaterThan(0)
303+
});
304+
283305
///////////////////////////////////////////////////
284306
// Test migrations ////////////////////////////////
285307
///////////////////////////////////////////////////

‎pgml-sdks/pgml/pyproject.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@ build-backend = "maturin"
55
[project]
66
name ="pgml"
77
requires-python =">=3.7"
8-
version ="0.9.5"
8+
version ="0.9.6"
99
description ="Python SDK is designed to facilitate the development of scalable vector search applications on PostgreSQL databases."
1010
authors = [
1111
{name ="PostgresML",email ="team@postgresml.org"},

‎pgml-sdks/pgml/python/pgml/pgml.pyi

Lines changed: 0 additions & 96 deletions
This file was deleted.

‎pgml-sdks/pgml/python/tests/test.py

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -298,6 +298,27 @@ async def test_order_documents():
298298
awaitcollection.archive()
299299

300300

301+
###################################################
302+
## Transformer Pipeline Tests #####################
303+
###################################################
304+
305+
306+
@pytest.mark.asyncio
307+
asyncdeftest_transformer_pipeline():
308+
t=pgml.TransformerPipeline("text-generation")
309+
it=awaitt.transform(["AI is going to"], {"max_new_tokens":5})
310+
assert (len(it))>0
311+
312+
@pytest.mark.asyncio
313+
asyncdeftest_transformer_pipeline_stream():
314+
t=pgml.TransformerPipeline("text-generation")
315+
it=awaitt.transform_stream("AI is going to", {"max_new_tokens":5})
316+
total= []
317+
asyncforcinit:
318+
total.append(c)
319+
assert (len(total))>0
320+
321+
301322
###################################################
302323
## Migration tests ################################
303324
###################################################

‎pgml-sdks/pgml/src/languages/javascript.rs

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
1+
use futures::StreamExt;
12
use neon::prelude::*;
23
use rust_bridge::javascript::{FromJsType,IntoJsResult};
4+
use std::sync::Arc;
35

46
usecrate::{
57
pipeline::PipelineSyncData,
8+
transformer_pipeline::TransformerStream,
69
types::{DateTime,Json},
710
};
811

@@ -16,8 +19,9 @@ impl IntoJsResult for DateTime {
1619
self,
1720
cx:&mutC,
1821
) ->JsResult<'b,Self::Output>{
19-
let date = neon::types::JsDate::new(cx,self.0.assume_utc().unix_timestamp()asf64*1000.0)
20-
.expect("Error converting to JS Date");
22+
let date =
23+
neon::types::JsDate::new(cx,self.0.assume_utc().unix_timestamp()asf64*1000.0)
24+
.expect("Error converting to JS Date");
2125
Ok(date)
2226
}
2327
}
@@ -69,6 +73,64 @@ impl IntoJsResult for PipelineSyncData {
6973
}
7074
}
7175

76+
#[derive(Clone)]
77+
structTransformerStreamArcMutex(Arc<tokio::sync::Mutex<TransformerStream>>);
78+
79+
implFinalizeforTransformerStreamArcMutex{}
80+
81+
fntransform_stream_iterate_next(mutcx:FunctionContext) ->JsResult<JsPromise>{
82+
let this = cx.this();
83+
let s:Handle<JsBox<TransformerStreamArcMutex>> = this
84+
.get(&mut cx,"s")
85+
.expect("Error getting self in transformer_stream_iterate_next");
86+
let ts:&TransformerStreamArcMutex =&s;
87+
let ts:TransformerStreamArcMutex = ts.clone();
88+
89+
let channel = cx.channel();
90+
let(deferred, promise) = cx.promise();
91+
crate::get_or_set_runtime().spawn(asyncmove{
92+
letmut ts = ts.0.lock().await;
93+
let v = ts.next().await;
94+
deferred
95+
.try_settle_with(&channel,move |mut cx|{
96+
let o = cx.empty_object();
97+
ifletSome(v) = v{
98+
let v:String = v.expect("Error calling next on TransformerStream");
99+
let v = cx.string(v);
100+
let d = cx.boolean(false);
101+
o.set(&mut cx,"value", v)
102+
.expect("Error setting object value in transformer_sream_iterate_next");
103+
o.set(&mut cx,"done", d)
104+
.expect("Error setting object value in transformer_sream_iterate_next");
105+
}else{
106+
let d = cx.boolean(true);
107+
o.set(&mut cx,"done", d)
108+
.expect("Error setting object value in transformer_sream_iterate_next");
109+
}
110+
Ok(o)
111+
})
112+
.expect("Error sending js");
113+
});
114+
Ok(promise)
115+
}
116+
117+
implIntoJsResultforTransformerStream{
118+
typeOutput =JsObject;
119+
fninto_js_result<'a,'b,'c:'b,C:Context<'c>>(
120+
self,
121+
cx:&mutC,
122+
) ->JsResult<'b,Self::Output>{
123+
let o = cx.empty_object();
124+
let f:Handle<JsFunction> =JsFunction::new(cx, transform_stream_iterate_next)?;
125+
o.set(cx,"next", f)?;
126+
let s = cx.boxed(TransformerStreamArcMutex(Arc::new(
127+
tokio::sync::Mutex::new(self),
128+
)));
129+
o.set(cx,"s", s)?;
130+
Ok(o)
131+
}
132+
}
133+
72134
////////////////////////////////////////////////////////////////////////////////
73135
// JS To Rust //////////////////////////////////////////////////////////////////
74136
////////////////////////////////////////////////////////////////////////////////

‎pgml-sdks/pgml/src/languages/python.rs

Lines changed: 58 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,65 +1,99 @@
1+
use futures::StreamExt;
12
use pyo3::conversion::IntoPy;
23
use pyo3::types::{PyDict,PyFloat,PyInt,PyList,PyString};
34
use pyo3::{prelude::*, types::PyBool};
5+
use std::sync::Arc;
46

57
use rust_bridge::python::CustomInto;
68

7-
usecrate::{pipeline::PipelineSyncData, types::Json};
9+
usecrate::{pipeline::PipelineSyncData,transformer_pipeline::TransformerStream,types::Json};
810

911
////////////////////////////////////////////////////////////////////////////////
1012
// Rust to PY //////////////////////////////////////////////////////////////////
1113
////////////////////////////////////////////////////////////////////////////////
1214

13-
implToPyObjectforJson{
14-
fnto_object(&self,py:Python) ->PyObject{
15+
implIntoPy<PyObject>forJson{
16+
fninto_py(self,py:Python) ->PyObject{
1517
match&self.0{
16-
serde_json::Value::Bool(x) => x.to_object(py),
18+
serde_json::Value::Bool(x) => x.into_py(py),
1719
serde_json::Value::Number(x) =>{
1820
if x.is_f64(){
1921
x.as_f64()
2022
.expect("Error converting to f64 in impl ToPyObject for Json")
21-
.to_object(py)
23+
.into_py(py)
2224
}else{
2325
x.as_i64()
2426
.expect("Error converting to i64 in impl ToPyObject for Json")
25-
.to_object(py)
27+
.into_py(py)
2628
}
2729
}
28-
serde_json::Value::String(x) => x.to_object(py),
30+
serde_json::Value::String(x) => x.into_py(py),
2931
serde_json::Value::Array(x) =>{
3032
let list =PyList::empty(py);
3133
for vin x.iter(){
32-
list.append(Json(v.clone()).to_object(py)).unwrap();
34+
list.append(Json(v.clone()).into_py(py)).unwrap();
3335
}
34-
list.to_object(py)
36+
list.into_py(py)
3537
}
3638
serde_json::Value::Object(x) =>{
3739
let dict =PyDict::new(py);
3840
for(k, v)in x.iter(){
39-
dict.set_item(k,Json(v.clone()).to_object(py)).unwrap();
41+
dict.set_item(k,Json(v.clone()).into_py(py)).unwrap();
4042
}
41-
dict.to_object(py)
43+
dict.into_py(py)
4244
}
4345
serde_json::Value::Null => py.None(),
4446
}
4547
}
4648
}
4749

48-
implIntoPy<PyObject>forJson{
50+
implIntoPy<PyObject>forPipelineSyncData{
4951
fninto_py(self,py:Python) ->PyObject{
50-
self.to_object(py)
52+
Json::from(self).into_py(py)
5153
}
5254
}
5355

54-
implToPyObjectforPipelineSyncData{
55-
fnto_object(&self,py:Python) ->PyObject{
56-
Json::from(self.clone()).to_object(py)
56+
#[pyclass]
57+
#[derive(Clone)]
58+
structTransformerStreamPython{
59+
wrapped:Arc<tokio::sync::Mutex<TransformerStream>>,
60+
}
61+
62+
#[pymethods]
63+
implTransformerStreamPython{
64+
fn__aiter__(slf:PyRef<'_,Self>) ->PyRef<'_,Self>{
65+
slf
66+
}
67+
68+
fn__anext__<'p>(slf:PyRefMut<'_,Self>,py:Python<'p>) ->PyResult<Option<PyObject>>{
69+
let ts = slf.wrapped.clone();
70+
let fut = pyo3_asyncio::tokio::future_into_py(py,asyncmove{
71+
letmut ts = ts.lock().await;
72+
ifletSome(o) = ts.next().await{
73+
Ok(Some(Python::with_gil(|py|{
74+
o.expect("Error calling next on TransformerStream")
75+
.to_object(py)
76+
})))
77+
}else{
78+
Err(pyo3::exceptions::PyStopAsyncIteration::new_err(
79+
"stream exhausted",
80+
))
81+
}
82+
})?;
83+
Ok(Some(fut.into()))
5784
}
5885
}
5986

60-
implIntoPy<PyObject>forPipelineSyncData{
87+
implIntoPy<PyObject>forTransformerStream{
6188
fninto_py(self,py:Python) ->PyObject{
62-
self.to_object(py)
89+
let f:Py<TransformerStreamPython> =Py::new(
90+
py,
91+
TransformerStreamPython{
92+
wrapped:Arc::new(tokio::sync::Mutex::new(self)),
93+
},
94+
)
95+
.expect("Error converting TransformerStream to TransformerStreamPython");
96+
f.to_object(py)
6397
}
6498
}
6599

@@ -115,6 +149,12 @@ impl FromPyObject<'_> for PipelineSyncData {
115149
}
116150
}
117151

152+
implFromPyObject<'_>forTransformerStream{
153+
fnextract(_ob:&PyAny) ->PyResult<Self>{
154+
panic!("We must implement this, but this is impossible to be reached")
155+
}
156+
}
157+
118158
////////////////////////////////////////////////////////////////////////////////
119159
// Rust to Rust //////////////////////////////////////////////////////////////////
120160
////////////////////////////////////////////////////////////////////////////////

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp