diff --git a/osschat_insert.py b/osschat_insert.py
index 5e69af5..710033d 100644
--- a/osschat_insert.py
+++ b/osschat_insert.py
@@ -98,8 +98,8 @@ def osschat_insert_pipe(config):
     p = (
         pipe.input('doc', 'project_name')
             .map('doc', 'text', data_loader)
-            .map('doc', 'token_count', ops.token_counter(config.token_model))
             .flat_map('text', 'chunk', text_split_op)
+            .map('chunk', 'token_count', ops.token_counter(config.token_model))
             .map('chunk', 'embedding', sentence_embedding_op, config=sentence_embedding_config)
     )
     if config.embedding_normalize: