async function embedBatchPaLM()

in firestore-semantic-search/functions/src/common/palm_embeddings.ts [46:81]


async function embedBatchPaLM(batch: string[]): Promise<number[][]> {
  const instances = batch.map(text =>
    helpers.toValue({content: text})
  ) as protobuf.common.IValue[];

  const parameters = helpers.toValue({});

  const [response] = await client.predict({
    endpoint,
    instances,
    parameters,
  });

  if (!response || !response.predictions || response.predictions.length === 0)
    throw new Error('Error with embedding');

  const predictionValues = response.predictions as protobuf.common.IValue[];

  const predictions = predictionValues.map(helpers.fromValue) as {
    embeddings: {values: number[]};
  }[];

  if (
    predictions.some(
      prediction => !prediction.embeddings || !prediction.embeddings.values
    )
  ) {
    throw new Error('Error with embedding');
  }

  const embeddings = predictions.map(
    prediction => prediction.embeddings.values
  );

  return embeddings;
}