@cocoindex.flow_def(name="TextEmbedding")
def text_embedding(flow: cocoindex.FlowBuilder, data: cocoindex.DataScope):
  # Add data source.
  data["documents"] = flow.add_source(
    cocoindex.sources.LocalFile(path="sourceFiles"))
  doc_embeddings = data.add_collector()

  with data["documents"].row() as doc:
    # Split into chunks.
    doc["chunks"] = doc["content"].transform(
      cocoindex.functions.SplitRecursively(), chunk_size=512)

    with doc["chunks"].row() as chunk:
      # Embed each chunk.
      chunk["embedding"] = chunk["text"].transform(
        cocoindex.functions.SentenceTransformerEmbed(
            model="sentence-transformers/all-MiniLM-L6-v2")

      # Collect embedding and metadata for indexing.
      doc_embeddings.collect(
        filename=doc["filename"], location=chunk["location"],
        text=chunk["text"], embedding=chunk["embedding"])

  # Export to vector store.
  doc_embeddings.export(
      "doc_embeddings", cocoindex.storages.Postgres(),
      primary_key_fields=["filename", "location"],
      vector_index=["embedding"])