NotificationsYou must be signed in to change notification settings
Fork352
Star6.6k

Commit96dd570

authored

Adding embed_array for getting the embeddings of multiple strings (#686)

1 parent7d1f4cd commit96dd570Copy full SHA for 96dd570

File tree

3 files changed

+36

-19

lines changed

pgml-extension/src
- api.rs
- bindings
  - transformers.py
  - transformers.rs

3 files changed

+36

-19

lines changed

`‎pgml-extension/src/api.rs‎`

Lines changed: 14 additions & 2 deletions

Original file line number	Diff line number	Diff line change
`@@ -563,9 +563,21 @@ fn load_dataset(`
`563`	`563`	`TableIterator::new(vec![(name, rows)].into_iter())`
`564`	`564`	`}`
`565`	`565`
`566`		`-#[pg_extern(immutable, parallel_safe)]`
	`566`	`+#[pg_extern(immutable, parallel_safe, name ="embed")]`
`567`	`567`	`pubfnembed(transformer:&str,text:&str,kwargs:default!(JsonB,"'{}'")) ->Vec<f32>{`
`568`		`-crate::bindings::transformers::embed(transformer, text,&kwargs.0)`
	`568`	`+embed_batch(transformer,Vec::from([text]), kwargs)`
	`569`	`+.first()`
	`570`	`+.unwrap()`
	`571`	`+.to_vec()`
	`572`	`+}`
	`573`	`+`
	`574`	`+#[pg_extern(immutable, parallel_safe, name ="embed")]`
	`575`	`+pubfnembed_batch(`
	`576`	`+transformer:&str,`
	`577`	`+inputs:Vec<&str>,`
	`578`	`+kwargs:default!(JsonB,"'{}'"),`
	`579`	`+) ->Vec<Vec<f32>>{`
	`580`	`+crate::bindings::transformers::embed(transformer,&inputs,&kwargs.0)`
`569`	`581`	`}`
`570`	`582`
`571`	`583`	`#[pg_extern(immutable, parallel_safe)]`

`‎pgml-extension/src/bindings/transformers.py‎`

Lines changed: 12 additions & 7 deletions

Original file line number	Diff line number	Diff line change
`@@ -110,25 +110,30 @@ def transform(task, args, inputs):`
`110`	`110`	`returnjson.dumps(pipe(inputs,**args),cls=NumpyJSONEncoder)`
`111`	`111`
`112`	`112`
`113`		`-defembed(transformer,text,kwargs):`
	`113`	`+defembed(transformer,inputs,kwargs):`
	`114`	`+`
	`115`	`+inputs=json.loads(inputs)`
`114`	`116`	`kwargs=json.loads(kwargs)`
`115`	`117`	`ensure_device(kwargs)`
`116`	`118`	`instructor=transformer.startswith("hkunlp/instructor")`
	`119`	`+`
`117`	`120`	`ifinstructor:`
`118`	`121`	`klass=INSTRUCTOR`
`119`		`-text= [[kwargs.pop("instruction"),text]]`
	`122`	`+`
	`123`	`+texts_with_instructions= []`
	`124`	`+instruction=kwargs.pop("instruction")`
	`125`	`+fortextininputs:`
	`126`	`+texts_with_instructions.append([instruction,text])`
	`127`	`+`
	`128`	`+inputs=texts_with_instructions`
`120`	`129`	`else:`
`121`	`130`	`klass=SentenceTransformer`
`122`	`131`
`123`	`132`	`iftransformernotin__cache_sentence_transformer_by_name:`
`124`	`133`	`__cache_sentence_transformer_by_name[transformer]=klass(transformer)`
`125`	`134`	`model=__cache_sentence_transformer_by_name[transformer]`
`126`	`135`
`127`		`-result=model.encode(text,**kwargs)`
`128`		`-ifinstructor:`
`129`		`-result=result[0]`
`130`		`-`
`131`		`-returnresult`
	`136`	`+returnmodel.encode(inputs,**kwargs)`
`132`	`137`
`133`	`138`
`134`	`139`	`defload_dataset(name,subset,limit:None,kwargs:"{}"):`

`‎pgml-extension/src/bindings/transformers.rs‎`

Lines changed: 10 additions & 10 deletions

Original file line number	Diff line number	Diff line change
`@@ -35,14 +35,13 @@ pub fn transform(`
`35`	`35`	`let results =Python::with_gil(\|py\| ->String{`
`36`	`36`	`let transform:Py<PyAny> =PY_MODULE.getattr(py,"transform").unwrap().into();`
`37`	`37`
`38`		`-let result = transform`
`39`		`-.call1(`
	`38`	`+let result = transform.call1(`
	`39`	`+ py,`
	`40`	`+PyTuple::new(`
`40`	`41`	`py,`
`41`		`-PyTuple::new(`
`42`		`- py,`
`43`		`-&[task.into_py(py), args.into_py(py), inputs.into_py(py)],`
`44`		`-),`
`45`		`-);`
	`42`	`+&[task.into_py(py), args.into_py(py), inputs.into_py(py)],`
	`43`	`+),`
	`44`	`+);`
`46`	`45`
`47`	`46`	`let result =match result{`
`48`	`47`	`Err(e) =>{`
`@@ -57,11 +56,12 @@ pub fn transform(`
`57`	`56`	`serde_json::from_str(&results).unwrap()`
`58`	`57`	`}`
`59`	`58`
`60`		`-pubfnembed(transformer:&str,text:&str,kwargs:&serde_json::Value) ->Vec<f32>{`
	`59`	`+pubfnembed(transformer:&str,inputs:Vec<&str>,kwargs:&serde_json::Value) ->Vec<Vec<f32>>{`
`61`	`60`	`crate::bindings::venv::activate();`
`62`	`61`
`63`	`62`	`let kwargs = serde_json::to_string(kwargs).unwrap();`
`64`		`-Python::with_gil(\|py\| ->Vec<f32>{`
	`63`	`+let inputs = serde_json::to_string(&inputs).unwrap();`
	`64`	`+Python::with_gil(\|py\| ->Vec<Vec<f32>>{`
`65`	`65`	`let embed:Py<PyAny> =PY_MODULE.getattr(py,"embed").unwrap().into();`
`66`	`66`	`embed`
`67`	`67`	`.call1(`
`@@ -70,7 +70,7 @@ pub fn embed(transformer: &str, text: &str, kwargs: &serde_json::Value) -> Vec<f`
`70`	`70`	`py,`
`71`	`71`	`&[`
`72`	`72`	`transformer.to_string().into_py(py),`
`73`		`-text.to_string().into_py(py),`
	`73`	`+inputs.into_py(py),`
`74`	`74`	`kwargs.into_py(py),`
`75`	`75`	`],`
`76`	`76`	`),`

0 commit comments

Comments

(0)

Movatterモバイル変換

Navigation Menu

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

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Commit96dd570

File tree

3 files changed

3 files changed

`‎pgml-extension/src/api.rs‎`

`‎pgml-extension/src/bindings/transformers.py‎`

`‎pgml-extension/src/bindings/transformers.rs‎`

0 commit comments