move embedding tests into their own file

montanalow · montanalow · commit 470f2d3d7619 · 2024-07-12T07:57:36.000-07:00
diff --git a/pgml-extension/examples/embedding.sql b/pgml-extension/examples/embedding.sql
@@ -0,0 +1,7 @@
+\timing on
+
+SELECT pgml.embed('Alibaba-NLP/gte-base-en-v1.5', 'hi mom', '{"trust_remote_code": true}');
+SELECT pgml.embed('Alibaba-NLP/gte-base-en-v1.5', 'hi mom', '{"device": "cuda", "trust_remote_code": true}');
+SELECT pgml.embed('Alibaba-NLP/gte-base-en-v1.5', 'hi mom', '{"device": "cpu", "trust_remote_code": true}');
+SELECT pgml.embed('hkunlp/instructor-xl', 'hi mom', '{"instruction": "Encode it with love"}');
+SELECT pgml.embed('mixedbread-ai/mxbai-embed-large-v1', 'test', '{"prompt": "test prompt: "}');
diff --git a/pgml-extension/examples/image_classification.sql b/pgml-extension/examples/image_classification.sql
@@ -66,7 +66,7 @@ SELECT * FROM pgml.train('Handwritten Digits', algorithm => 'xgboost', hyperpara
 
 -- runtimes
 SELECT * FROM pgml.train('Handwritten Digits', algorithm => 'linear', runtime => 'python');
-SELECT * FROM pgml.train('Handwritten Digits', algorithm => 'linear', runtime => 'rust');
+--SELECT * FROM pgml.train('Handwritten Digits', algorithm => 'linear', runtime => 'rust');
 
 --SELECT * FROM pgml.train('Handwritten Digits', algorithm => 'xgboost', runtime => 'python', hyperparams => '{"n_estimators": 10}'); -- too slow
 SELECT * FROM pgml.train('Handwritten Digits', algorithm => 'xgboost', runtime => 'rust', hyperparams => '{"n_estimators": 10}');
diff --git a/pgml-extension/examples/transformers.sql b/pgml-extension/examples/transformers.sql
@@ -18,7 +18,8 @@ SELECT pgml.transform_stream(
   }'::JSONB,
   input => 'AI is going to',
   args   => '{
-    "max_new_tokens": 100
+    "max_new_tokens": 100,
+    "token": "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"
   }'::JSONB
 );
 -- BitsAndBytes support
diff --git a/pgml-extension/tests/test.sql b/pgml-extension/tests/test.sql
@@ -31,5 +31,6 @@ SELECT pgml.load_dataset('wine');
 \i examples/vectors.sql
 \i examples/chunking.sql
 \i examples/preprocessing.sql
+\i examples/embedding.sql
 -- transformers are generally too slow to run in the test suite
 --\i examples/transformers.sql
