From a23ab95e3bc1793db0fb4a5eacbd2add06b99458 Mon Sep 17 00:00:00 2001
From: ssiddhantsharma <29195354+ssiddhantsharma@users.noreply.github.com>
Date: Thu, 13 Aug 2026 12:08:58 +0530
Subject: [PATCH] Cache template features per query in TemplateHitFeaturizer

get_templates re-runs the full template search (remote mmCIF fetch + kalign)
on every call. In iterative inference (e.g. hallucination) the same fixed
target chain is featurized every optimization step, repeating identical
remote fetches and stalling the loop. Memoize by (sequence_uid, query_sequence)
so templates are computed once per unique query.
---
 opendde/data/template/template_utils.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/opendde/data/template/template_utils.py b/opendde/data/template/template_utils.py
index a67f623..f0a265b 100644
--- a/opendde/data/template/template_utils.py
+++ b/opendde/data/template/template_utils.py
@@ -794,6 +794,7 @@ def __init__(
         self._zero_center_positions = _zero_center_positions
         self._max_template_candidates_num = _max_template_candidates_num
         self._fetch_remote = fetch_remote
+        self._template_feature_cache = {}
 
         if max_template_date:
             if isinstance(max_template_date, str):
@@ -868,6 +869,10 @@ def get_templates(
         Returns:
             A tuple of (TemplateSearchResult, timing_dict).
         """
+        cache_key = (sequence_uid, query_sequence)
+        if cache_key in self._template_feature_cache:
+            return self._template_feature_cache[cache_key]
+
         cutoff = self._max_template_date
         if max_template_date:
             if isinstance(max_template_date, str):
@@ -949,4 +954,6 @@ def get_templates(
                     features.append(res.features)
                     final_hits.append(res.hit)
 
-        return TemplateSearchResult(features, final_hits, errors, warnings), last_track
+        result = TemplateSearchResult(features, final_hits, errors, warnings), last_track
+        self._template_feature_cache[cache_key] = result
+        return result