From d0b846abb02dc3c95db5ca3fa928efc85bef8972 Mon Sep 17 00:00:00 2001 From: nsantacruz Date: Sun, 19 Oct 2025 13:12:43 +0300 Subject: [PATCH] feat: add type hints for span partitioning functions in app_helper.py --- app/app_helper.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/app/app_helper.py b/app/app_helper.py index edfb125..c440f91 100644 --- a/app/app_helper.py +++ b/app/app_helper.py @@ -1,5 +1,6 @@ from functools import reduce from collections import defaultdict +from ne_span import NESpan BATCH_SIZE = 150 @@ -35,7 +36,7 @@ def _get_linker_entities(text, ner_model, ref_part_model): return cit_spans, ref_parts, other_spans -def _partition_spans(spans): +def _partition_spans(spans: list[NESpan]): cit_spans, other_spans = [], [] for span in spans: if LABEL2TYPE.get(span.label) == 'citation': @@ -45,7 +46,7 @@ def _partition_spans(spans): return cit_spans, other_spans -def _bulk_partition_spans(spans_list): +def _bulk_partition_spans(spans_list: list[list[NESpan]]): cit_spans_list, other_spans_list = [], [] for spans in spans_list: inner_cit_spans, inner_other_spans = _partition_spans(spans)