Skip to content

Commit cfdbab0

Browse files
feat(medcat): CU-869cw9zmj Improve inference speed (#410)
* CU-869cw9zmj: Use faster way to calculate unit vector * CU-869cw9zmj: Speed up context vector obtaining * CU-869ctq789: Avoid leaking normalized vectors * Revert "CU-869ctq789: Avoid leaking normalized vectors" This reverts commit 7856c12. * Revert "CU-869cw9zmj: Use faster way to calculate unit vector" This reverts commit 12b7c52. * CU-869cw9zmj: Fix usage of tokens in the correct order * CU-869cw9zmj: Add small comment on left token slicing * CU-869cw9zmj: Separate centre context vectors calculation outside the loop * CU-869cw9zmj: Get centre tokens separately only if they're required * CU-869cw9zmj: Fix linting issue --------- Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
1 parent 939d015 commit cfdbab0

1 file changed

Lines changed: 58 additions & 24 deletions

File tree

medcat-v2/medcat/components/linking/vector_context_model.py

Lines changed: 58 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@ def __init__(self, cui2info: dict[str, CUIInfo],
5858
def get_context_tokens(self, entity: MutableEntity, doc: MutableDocument,
5959
size: int,
6060
per_doc_valid_token_cache: 'PerDocumentTokenCache',
61+
fill_centre_tokens: bool = True,
6162
) -> tuple[list[MutableToken],
6263
list[MutableToken],
6364
list[MutableToken]]:
@@ -83,17 +84,21 @@ def get_context_tokens(self, entity: MutableEntity, doc: MutableDocument,
8384
per_doc_valid_token_cache[tkn]]
8485
# Reverse because the first token should be the one closest to center
8586
tokens_left.reverse()
86-
tokens_center: list[MutableToken] = list(
87-
cast(Iterable[MutableToken], entity))
87+
if fill_centre_tokens:
88+
tokens_center: list[MutableToken] = list(
89+
cast(Iterable[MutableToken], entity))
90+
else:
91+
tokens_center = []
8892
_right_tokens = doc[end_ind + 1:end_ind + 1 + size]
8993
tokens_right = [tkn for tkn in _right_tokens if
9094
per_doc_valid_token_cache[tkn]]
9195

9296
return tokens_left, tokens_center, tokens_right
9397

94-
def _tokens2vecs(self, tokens: Sequence[Union[MutableToken, str]]
95-
) -> Iterable[np.ndarray]:
96-
for step, tkn in enumerate(tokens):
98+
def _tokens2vecs(self, tokens: Sequence[Union[MutableToken, str]],
99+
step_start: int = 0
100+
) -> Iterable[np.ndarray]:
101+
for step, tkn in enumerate(tokens, start=step_start):
97102
lower = tkn.lower() if isinstance(tkn, str) else tkn.base.lower
98103
if lower not in self.vocab:
99104
continue
@@ -137,26 +142,55 @@ def get_context_vectors(self, entity: MutableEntity,
137142
"""
138143
vectors: dict[str, np.ndarray] = {}
139144

140-
context_vector_sizes = self.config.context_vector_sizes
141-
for context_type, window_size in context_vector_sizes.items():
142-
tokens_left, tokens_center, tokens_right = self.get_context_tokens(
143-
entity, doc, window_size, per_doc_valid_token_cache)
144-
145-
values: list[np.ndarray] = []
146-
# Add left
147-
values.extend(self._tokens2vecs(tokens_left))
148-
149-
if not self.config.context_ignore_center_tokens:
150-
# Add center
151-
values.extend(
152-
self._preprocess_center_tokens(cui, tokens_center))
153-
154-
# Add right
155-
values.extend(self._tokens2vecs(tokens_right))
156-
145+
# Sort ascending so each iteration is a superset of the previous
146+
sorted_contexts = sorted(
147+
self.config.context_vector_sizes.items(), key=lambda x: x[1])
148+
149+
prev_left: list[MutableToken] = []
150+
prev_right: list[MutableToken] = []
151+
# Accumulated weighted vecs from previous (smaller) windows,
152+
# excluding center (center is the same for all window sizes)
153+
prev_left_vecs: list[np.ndarray] = []
154+
prev_right_vecs: list[np.ndarray] = []
155+
156+
# Center is identical for all window sizes, only compute once
157+
if not self.config.context_ignore_center_tokens:
158+
tokens_center = list(
159+
cast(Iterable[MutableToken], entity))
160+
center_vecs = list(
161+
self._preprocess_center_tokens(cui, tokens_center))
162+
else:
163+
center_vecs = []
164+
165+
for context_type, window_size in sorted_contexts:
166+
tokens_left, _, tokens_right = self.get_context_tokens(
167+
entity, doc, window_size, per_doc_valid_token_cache,
168+
fill_centre_tokens=False)
169+
170+
# New outer tokens only — the inner ones were already processed
171+
# NOTE: left hand tokens are in order of closest first, which is why
172+
# we're slicing from the start of the list
173+
new_left = tokens_left[len(prev_left):]
174+
new_right = tokens_right[len(prev_right):]
175+
176+
# step_start for new left tokens: they are further from centre
177+
# so their step index is
178+
# len(tokens_left) - len(new_left) ... len(tokens_left)-1
179+
# i.e. the new tokens are the outermost, highest-step ones
180+
new_left_vecs = list(self._tokens2vecs(
181+
new_left, step_start=len(prev_left)))
182+
new_right_vecs = list(self._tokens2vecs(
183+
new_right, step_start=len(prev_right)))
184+
185+
prev_left_vecs = new_left_vecs + prev_left_vecs
186+
prev_right_vecs = prev_right_vecs + new_right_vecs
187+
prev_left = tokens_left
188+
prev_right = tokens_right
189+
190+
values = prev_left_vecs + center_vecs + prev_right_vecs
157191
if values:
158-
value = np.average(values, axis=0)
159-
vectors[context_type] = value
192+
vectors[context_type] = np.average(values, axis=0)
193+
160194
return vectors
161195

162196
def similarity(self, cui: str, entity: MutableEntity, doc: MutableDocument,

0 commit comments

Comments
 (0)