Skip to content

Commit fc31f41

Browse files
committed
fix typo in doc_pooler.collate
1 parent 8072a21 commit fc31f41

File tree

1 file changed

+3
-1
lines changed

1 file changed

+3
-1
lines changed

edsnlp/pipes/trainable/embeddings/doc_pooler/doc_pooler.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,7 +79,9 @@ def collate(self, batch: Dict[str, Any]) -> DocPoolerBatchInput:
7979
stats = batch["stats"]
8080
return {
8181
"embedding": embedding_batch,
82-
"stats": stats,
82+
"stats": {
83+
"doc_length": sum(stats["doc_length"])
84+
}, # <-- sum(...) pour aggréger les comptes par doc en un compte par batch
8385
}
8486

8587
def forward(self, batch: DocPoolerBatchInput) -> DocPoolerBatchOutput:

0 commit comments

Comments
 (0)