Improved examples [skip ci]

Author: ankane

examples/hybrid-search/example.js

@@ -12,22 +12,21 @@ await client.query('DROP TABLE IF EXISTS documents');
 await client.query('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))');
 await client.query("CREATE INDEX ON documents USING GIN (to_tsvector('english', content))");
 
-const input = [
-  'The dog is barking',
-  'The cat is purring',
-  'The bear is growling'
-];
-
 const extractor = await pipeline('feature-extraction', 'Xenova/multi-qa-MiniLM-L6-cos-v1', {dtype: 'fp32'});
 
 async function embed(content) {
   const output = await extractor(content, {pooling: 'mean', normalize: true});
-  return Array.from(output.data);
+  return output.tolist();
 }
 
-for (let content of input) {
-  const embedding = await embed(content);
-  await client.query('INSERT INTO documents (content, embedding) VALUES ($1, $2)', [content, pgvector.toSql(embedding)]);
+const input = [
+  'The dog is barking',
+  'The cat is purring',
+  'The bear is growling'
+];
+const embeddings = await embed(input);
+for (let [i, content] of input.entries()) {
+  await client.query('INSERT INTO documents (content, embedding) VALUES ($1, $2)', [content, pgvector.toSql(embeddings[i])]);
 }
 
 const sql = `
@@ -54,9 +53,9 @@ ORDER BY score DESC
 LIMIT 5
 `;
 const query = 'growling bear'
-const embedding = await embed(query);
+const queryEmbedding = (await embed([query]))[0];
 const k = 60
-const { rows } = await client.query(sql, [query, pgvector.toSql(embedding), k]);
+const { rows } = await client.query(sql, [query, pgvector.toSql(queryEmbedding), k]);
 for (let row of rows) {
   console.log(row);
 }

examples/transformers/example.js

@@ -1,4 +1,4 @@
-import { pipeline } from '@xenova/transformers';
+import { pipeline } from '@huggingface/transformers';
 import pg from 'pg';
 import pgvector from 'pgvector/pg';
 
@@ -11,22 +11,26 @@ await pgvector.registerTypes(client);
 await client.query('DROP TABLE IF EXISTS documents');
 await client.query('CREATE TABLE documents (id bigserial PRIMARY KEY, content text, embedding vector(384))');
 
+const extractor = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2', {dtype: 'fp32'});
+
+async function embed(input) {
+  const output = await extractor(input, {pooling: 'mean', normalize: true});
+  return output.tolist();
+}
+
 const input = [
   'The dog is barking',
   'The cat is purring',
   'The bear is growling'
 ];
-
-const extractor = await pipeline('feature-extraction', 'Xenova/all-MiniLM-L6-v2');
-
-for (let content of input) {
-  const output = await extractor(content, {pooling: 'mean', normalize: true});
-  const embedding = Array.from(output.data);
-  await client.query('INSERT INTO documents (content, embedding) VALUES ($1, $2)', [content, pgvector.toSql(embedding)]);
+const embeddings = await embed(input);
+for (let [i, content] of input.entries()) {
+  await client.query('INSERT INTO documents (content, embedding) VALUES ($1, $2)', [content, pgvector.toSql(embeddings[i])]);
 }
 
-const documentId = 2;
-const { rows } = await client.query('SELECT * FROM documents WHERE id != $1 ORDER BY embedding <=> (SELECT embedding FROM documents WHERE id = $1) LIMIT 5', [documentId]);
+const query = 'forest';
+const queryEmbedding = (await embed([query]))[0];
+const { rows } = await client.query('SELECT content FROM documents ORDER BY embedding <=> $1 LIMIT 5', [pgvector.toSql(queryEmbedding)]);
 for (let row of rows) {
   console.log(row.content);
 }

examples/transformers/package.json

@@ -2,7 +2,7 @@
     "private": true,
     "type": "module",
     "dependencies": {
-        "@xenova/transformers": "^2.6.0",
+        "@huggingface/transformers": "^3.3.3",
         "pg": "^8.11.3",
         "pgvector": "file:../.."
     }