Use built-in AI models for in-database inference

Embed a single string

SELECT pgml.embed('thenlper/gte-small-zh', 'The self-attention mechanism is widely used in neural networks.');

Result:

                           embed
------------------------------------------------------------
 {0.011534364,-0.029397607, ... -0.00056651415,-0.05465962}

Embed multiple strings in one call

SELECT pgml.embed(
    'thenlper/gte-small-zh',
    ARRAY[
        'The self-attention mechanism is widely used in neural networks.',
        'The core idea is to allow models to consider other elements in the sequence when processing a single element'
    ]
);

Result:

                           embed
------------------------------------------------------------
 {0.011534364,-0.029397607, ... -0.00056651415,-0.05465962}
 {0.024123996,0.0360483154, ... -0.029659372,-0.0198373856}

Embed existing table data

Generate embeddings from a column and store the results in a new table.

-- Source table
CREATE TABLE dump_table (id int, content text);
-- id | content
-- ---+--------------------------------------------------------
--  1 | The self-attention mechanism generates attention scores...
--  2 | Attention weights are generated by calculating the dot product...

-- Generate embeddings and save to result_table
CREATE TABLE result_table AS
    SELECT id, content, pgml.embed('thenlper/gte-small-zh', content) AS embed
    FROM dump_table;

-- Preview the first two dimensions of each embedding
SELECT id, content, embed[1:2] FROM result_table;

Result:

 id | content                                         | embed
----+-------------------------------------------------+-------------------
  1 | The self-attention mechanism generates...       | {-0.020080239,...}
  2 | Attention weights are generated by calculating..| {-0.036033697,...}
(2 rows)

Classify a single string

SELECT pgml.transform(
    task   => '{"task": "text-classification",
                "model": "lxyuan/distilbert-base-multilingual-cased-sentiments-student"}'::JSONB,
    inputs => ARRAY['The product quality is good']
) AS positivity;

The result shows a label of "positive" with a confidence score of 0.9924:

                     positivity
-----------------------------------------------------
 [{"label": "positive", "score": 0.992497742176056}]
(1 row)

Classify existing table data

Read rows from dump_table, classify the content column, and save a boolean pos flag to result_table.

CREATE TABLE result_table AS
    SELECT
        id,
        content,
        pgml.transform(
            task   => '{"task": "text-classification",
                        "model": "lxyuan/distilbert-base-multilingual-cased-sentiments-student"}'::JSONB,
            inputs => ARRAY[content]::text[]
        )->0->>'label' = 'positive' AS pos
    FROM dump_table;

Extract entities from a single string

SELECT pgml.transform(
    task   => '{"task": "token-classification",
                "model": "Babelscape/wikineural-multilingual-ner"}'::JSONB,
    inputs => ARRAY['Li Ming arrived in a city yesterday']
);

Result:

                      transform
------------------------------------------------------------
 [[{"end": 1, "word": "Li",   "index": 1, "score": 0.9668, "start": 0, "entity": "B-PER"},
   {"end": 2, "word": "Ming", "index": 2, "score": 0.9899, "start": 1, "entity": "I-PER"},
   {"end": 8, "word": "Some", "index": 8, "score": 0.9998, "start": 7, "entity": "B-LOC"},
   {"end": 9, "word": "City", "index": 9, "score": 0.9989, "start": 8, "entity": "I-LOC"}]]
(1 row)

Each object in the result array describes one recognized token:

Extract entities from existing table data

CREATE TABLE result_table AS
    SELECT
        id,
        content,
        pgml.transform(
            task   => '{"task": "token-classification",
                        "model": "Babelscape/wikineural-multilingual-ner"}'::JSONB,
            inputs => ARRAY[content]
        ) AS result
    FROM dump_table;

SELECT * FROM result_table;

Result:

 id | content                                   | result
----+-------------------------------------------+-------------------------------------------------------
  1 | My name IS Wolfgang AND I live IN Berlin. | [[{"word": "Wolfgang", "entity": "B-PER", ...},
    |                                           |   {"word": "Berlin",   "entity": "B-LOC", ...}]]
  2 | Li Ming will go to a certain city today.  | [[{"word": "Li",       "entity": "B-PER", ...},
    |                                           |   {"word": "Ming",     "entity": "I-PER", ...}, ...]]
(2 rows)