Jul 12, 2024 · Jul 12, 2024 · Jul 12, 2024 · Jul 12, 2024
diff --git a/pgml-extension/src/bindings/transformers/transformers.py b/pgml-extension/src/bindings/transformers/transformers.py
 returnrank_using(model,query,documents,kwargs)


 defcreate_embedding(transformer):
 returnSentenceTransformer(transformer)
 defcreate_embedding(transformer,kwargs):
 returnSentenceTransformer(transformer,**kwargs)


 defembed_using(model,transformer,inputs,kwargs):

 defembed(transformer,inputs,kwargs):
 kwargs=orjson.loads(kwargs)

 ensure_device(kwargs)

 init_kwarg_keys= [
 "device",
 "trust_remote_code",
 "revision",
 "model_kwargs",
 "tokenizer_kwargs",
 "config_kwargs",
 "truncate_dim",
 "token",
    ]
 init_kwargs= {
 key:valueforkey,valueinkwargs.items()ifkeyininit_kwarg_keys
    }
 encode_kwargs= {
 key:valueforkey,valueinkwargs.items()ifkeynotininit_kwarg_keys
    }

 iftransformernotin__cache_sentence_transformer_by_name:
 __cache_sentence_transformer_by_name[transformer]=create_embedding(
 transformer
 transformer,init_kwargs
        )
 model=__cache_sentence_transformer_by_name[transformer]

 returnembed_using(model,transformer,inputs,kwargs)
 returnembed_using(model,transformer,inputs,encode_kwargs)


 defclear_gpu_cache(memory_usage:None):
Original file line number	Diff line number	Diff line change
Expand Up		@@ -527,8 +527,8 @@ def rank(transformer, query, documents, kwargs):
		returnrank_using(model,query,documents,kwargs)


		defcreate_embedding(transformer):
		returnSentenceTransformer(transformer)
		defcreate_embedding(transformer,kwargs):
		returnSentenceTransformer(transformer,**kwargs)


		defembed_using(model,transformer,inputs,kwargs):
Expand All		@@ -545,16 +545,32 @@ def embed_using(model, transformer, inputs, kwargs):

		defembed(transformer,inputs,kwargs):
		kwargs=orjson.loads(kwargs)

		ensure_device(kwargs)

		init_kwarg_keys= [
		"device",
		"trust_remote_code",
		"revision",
		"model_kwargs",
		"tokenizer_kwargs",
		"config_kwargs",
		"truncate_dim",
		"token",
		]
		init_kwargs= {
		key:valueforkey,valueinkwargs.items()ifkeyininit_kwarg_keys
		}
		encode_kwargs= {
		key:valueforkey,valueinkwargs.items()ifkeynotininit_kwarg_keys
		}

		iftransformernotin__cache_sentence_transformer_by_name:
		__cache_sentence_transformer_by_name[transformer]=create_embedding(
		transformer
		transformer,init_kwargs
		)
		model=__cache_sentence_transformer_by_name[transformer]

		returnembed_using(model,transformer,inputs,kwargs)
		returnembed_using(model,transformer,inputs,encode_kwargs)


		defclear_gpu_cache(memory_usage:None):
Expand Down