May 18, 2023 · May 18, 2023
diff --git a/pgml-extension/examples/transformers.sql b/pgml-extension/examples/transformers.sql
 \timingon

 SELECTpgml.embed('intfloat/e5-small','hi mom');

 SELECTpgml.embed('intfloat/e5-small','hi mom','{"device": "cuda"}');
 SELECTpgml.embed('intfloat/e5-small','hi mom','{"device": "cpu"}');

 SELECTpgml.transform(
 'translation_en_to_fr',
 SELECTpgml.transform(
 '{"model": "roberta-large-mnli"}'::JSONB,
    inputs=> ARRAY[
 'I love how amazingly simple ML has become!',
 'I love how amazingly simple ML has become!',
 'Some models are painfully slow and expensive ☹️'
    ]
 )AS result;
    ]
 );
 SELECTpgml.transform(
    task=>'{"task": "text-classification",
              "model": "finiteautomata/bertweet-base-sentiment-analysis"
             }'::JSONB,
    inputs=> ARRAY[
 'I love how amazingly simple ML has become!',
 'I hate doing mundane and thankless tasks. ☹️'
    ],
    task=>'{"task": "text-classification",
              "model": "finiteautomata/bertweet-base-sentiment-analysis"
             }'::JSONB
 )AS positivity;

 SELECTpgml.transform(
diff --git a/pgml-extension/requirements.txt b/pgml-extension/requirements.txt
 accelerate==0.16.0
 accelerate==0.19.0
 datasets==2.10.1
 deepspeed==0.8.1
 InstructorEmbedding
 torchaudio==0.13.1
 torchvision==0.14.1
 tqdm==4.64.1
 transformers==4.26.1
 transformers==4.28.1
 xgboost
diff --git a/pgml-extension/src/api.rs b/pgml-extension/src/api.rs
    task: JsonB,
    args: default!(JsonB, "'{}'"),
    inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
    cache: default!(bool, false),
 ) -> JsonB {
    JsonB(crate::bindings::transformers::transform(
        &task.0, &args.0, &inputs, cache,
        &task.0, &args.0, &inputs,
    ))
 }

    task: String,
    args: default!(JsonB, "'{}'"),
    inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
    cache: default!(bool, false),
 ) -> JsonB {
    let mut task_map = HashMap::new();
    task_map.insert("task", task);
    let task_json = json!(task_map);
    JsonB(crate::bindings::transformers::transform(
        &task_json, &args.0, &inputs, cache,
        &task_json, &args.0, &inputs,
    ))
 }

diff --git a/pgml-extension/src/bindings/transformers.py b/pgml-extension/src/bindings/transformers.py
        return super().default(obj)


 def transform(task, args, inputs, cache):
 def transform(task, args, inputs):
    task = json.loads(task)
    args = json.loads(args)
    inputs = json.loads(inputs)

 task["device"] = assign_device(task.get("device"))
 ensure_device(task)

    if cache:
        key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
        if key not in __cache_transform_pipeline_by_task:
            __cache_transform_pipeline_by_task[key] = transformers.pipeline(**task)
        pipe = __cache_transform_pipeline_by_task[key]
    else:
        pipe = transformers.pipeline(**task)
    key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
    if key not in __cache_transform_pipeline_by_task:
        __cache_transform_pipeline_by_task[key] = transformers.pipeline(**task)
    pipe = __cache_transform_pipeline_by_task[key]

    if pipe.task == "question-answering":
        inputs = [json.loads(input) for input in inputs]

 def embed(transformer, text, kwargs):
    kwargs = json.loads(kwargs)
 kwargs["device"] = assign_device(kwargs.get("device"))
 ensure_device(kwargs)
    instructor = transformer.startswith("hkunlp/instructor")
    if instructor:
        klass = INSTRUCTOR
    return all_preds


 def assign_device(device=None):
    if device is not None:
        if device == "cpu" or "cuda:" in device:
            return device
        if "cuda" in device and not torch.cuda.is_available():
            raise Exception("CUDA is not available")

    if torch.cuda.is_available():
        device = "cuda:" + str(os.getpid() % torch.cuda.device_count())
    else:
        device = "cpu"
 def ensure_device(kwargs):
    device = kwargs.get("device")
    device_map = kwargs.get("device_map")
    if device is None and device_map is None:
        if torch.cuda.is_available():
            kwargs["device"] = "cuda:" + str(os.getpid() % torch.cuda.device_count())
        else:
            kwargs["device"] = "cpu"

    return device
diff --git a/pgml-extension/src/bindings/transformers.rs b/pgml-extension/src/bindings/transformers.rs
    task: &serde_json::Value,
    args: &serde_json::Value,
    inputs: &Vec<String>,
    cache: bool,
 ) -> serde_json::Value {
    crate::bindings::venv::activate();

                        task.into_py(py),
                        args.into_py(py),
                        inputs.into_py(py),
                        cache.into_py(py),
                    ],
                ),
            )
Original file line number	Diff line number	Diff line change
Expand Up		@@ -3,7 +3,8 @@
		\timingon

		SELECTpgml.embed('intfloat/e5-small','hi mom');

		SELECTpgml.embed('intfloat/e5-small','hi mom','{"device": "cuda"}');
		SELECTpgml.embed('intfloat/e5-small','hi mom','{"device": "cpu"}');

		SELECTpgml.transform(
		'translation_en_to_fr',
Expand All		@@ -16,7 +17,7 @@ SELECT pgml.transform(
		SELECTpgml.transform(
		'{"model": "roberta-large-mnli"}'::JSONB,
		inputs=> ARRAY[
		'I love how amazingly simple ML has become!',
		'I love how amazingly simple ML has become!',
		'Some models are painfully slow and expensive ☹️'
		]
		)AS result;
Expand All		@@ -35,13 +36,13 @@ SELECT pgml.transform(
		]
		);
		SELECTpgml.transform(
		task=>'{"task": "text-classification",
		"model": "finiteautomata/bertweet-base-sentiment-analysis"
		}'::JSONB,
		inputs=> ARRAY[
		'I love how amazingly simple ML has become!',
		'I hate doing mundane and thankless tasks. ☹️'
		],
		task=>'{"task": "text-classification",
		"model": "finiteautomata/bertweet-base-sentiment-analysis"
		}'::JSONB
		)AS positivity;

		SELECTpgml.transform(
Expand Down
Original file line number	Diff line number	Diff line change
		@@ -1,4 +1,4 @@
		accelerate==0.16.0
		accelerate==0.19.0
		datasets==2.10.1
		deepspeed==0.8.1
		InstructorEmbedding
Expand All		@@ -15,5 +15,5 @@ torch==1.13.1
		torchaudio==0.13.1
		torchvision==0.14.1
		tqdm==4.64.1
		transformers==4.26.1
		transformers==4.28.1
		xgboost
Original file line number	Diff line number	Diff line change
Expand Up		@@ -574,10 +574,9 @@ pub fn transform_json(
		task: JsonB,
		args: default!(JsonB, "'{}'"),
		inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
		cache: default!(bool, false),
		) -> JsonB {
		JsonB(crate::bindings::transformers::transform(
		&task.0, &args.0, &inputs, cache,
		&task.0, &args.0, &inputs,
		))
		}

Expand All		@@ -587,13 +586,12 @@ pub fn transform_string(
		task: String,
		args: default!(JsonB, "'{}'"),
		inputs: default!(Vec<String>, "ARRAY[]::TEXT[]"),
		cache: default!(bool, false),
		) -> JsonB {
		let mut task_map = HashMap::new();
		task_map.insert("task", task);
		let task_json = json!(task_map);
		JsonB(crate::bindings::transformers::transform(
		&task_json, &args.0, &inputs, cache,
		&task_json, &args.0, &inputs,
		))
		}

Expand Down
Original file line number	Diff line number	Diff line change
Expand Up		@@ -50,20 +50,17 @@ def default(self, obj):
		return super().default(obj)


		def transform(task, args, inputs, cache):
		def transform(task, args, inputs):
		task = json.loads(task)
		args = json.loads(args)
		inputs = json.loads(inputs)

		task["device"] = assign_device(task.get("device"))
		ensure_device(task)

		if cache:
		key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
		if key not in __cache_transform_pipeline_by_task:
		__cache_transform_pipeline_by_task[key] = transformers.pipeline(**task)
		pipe = __cache_transform_pipeline_by_task[key]
		else:
		pipe = transformers.pipeline(**task)
		key = ",".join([f"{key}:{val}" for (key, val) in sorted(task.items())])
		if key not in __cache_transform_pipeline_by_task:
		__cache_transform_pipeline_by_task[key] = transformers.pipeline(**task)
		pipe = __cache_transform_pipeline_by_task[key]

		if pipe.task == "question-answering":
		inputs = [json.loads(input) for input in inputs]
Expand All		@@ -73,7 +70,7 @@ def transform(task, args, inputs, cache):

		def embed(transformer, text, kwargs):
		kwargs = json.loads(kwargs)
		kwargs["device"] = assign_device(kwargs.get("device"))
		ensure_device(kwargs)
		instructor = transformer.startswith("hkunlp/instructor")
		if instructor:
		klass = INSTRUCTOR
Expand DownExpand Up		@@ -543,16 +540,12 @@ def generate(model_id, data, config):
		return all_preds


		def assign_device(device=None):
		if device is not None:
		if device == "cpu" or "cuda:" in device:
		return device
		if "cuda" in device and not torch.cuda.is_available():
		raise Exception("CUDA is not available")

		if torch.cuda.is_available():
		device = "cuda:" + str(os.getpid() % torch.cuda.device_count())
		else:
		device = "cpu"
		def ensure_device(kwargs):
		device = kwargs.get("device")
		device_map = kwargs.get("device_map")
		if device is None and device_map is None:
		if torch.cuda.is_available():
		kwargs["device"] = "cuda:" + str(os.getpid() % torch.cuda.device_count())
		else:
		kwargs["device"] = "cpu"

		return device
Original file line number	Diff line number	Diff line change
Expand Up		@@ -25,7 +25,6 @@ pub fn transform(
		task: &serde_json::Value,
		args: &serde_json::Value,
		inputs: &Vec<String>,
		cache: bool,
		) -> serde_json::Value {
		crate::bindings::venv::activate();

Expand All		@@ -45,7 +44,6 @@ pub fn transform(
		task.into_py(py),
		args.into_py(py),
		inputs.into_py(py),
		cache.into_py(py),
		],
		),
		)
Expand Down