pgml 1.1.1

The official pgml Rust SDK
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60

const pgml = require("pgml");
require("dotenv").config();

const main = async () => {
  // Initialize the collection
  const collection = pgml.newCollection("qa_collection");

  // Add a pipeline
  const pipeline = pgml.newPipeline("qa_pipeline", {
    text: {
      splitter: { model: "recursive_character" },
      semantic_search: {
        model: "Alibaba-NLP/gte-base-en-v1.5",
      },
    },
  });
  await collection.add_pipeline(pipeline);

  // Upsert documents, these documents are automatically split into chunks and embedded by our pipeline
  const documents = [
    {
      id: "Document One",
      text: "PostgresML is the best tool for machine learning applications!",
    },
    {
      id: "Document Two",
      text: "PostgresML is open source and available to everyone!",
    },
  ];
  await collection.upsert_documents(documents);

  // Perform vector search
  const query = "What is the best tool for building machine learning applications?";
  const queryResults = await collection.vector_search(
    {
      query: {
        fields: {
          text: { query: query }
        }
      }, limit: 1
    }, pipeline);
  console.log("The results");
  console.log(queryResults);

  const context = queryResults.map((result) => result["chunk"]).join("\n\n");

  // Query for summarization
  const builtins = pgml.newBuiltins();
  const answer = await builtins.transform(
    { task: "summarization", model: "sshleifer/distilbart-cnn-12-6" },
    [context],
  );
  console.log("The summary");
  console.log(answer);

  // Archive the collection
  await collection.archive();
};

main().then(() => console.log("Done!"));