Update handler.py

handler.py

@@ -173,6 +173,7 @@ class EndpointHandler:
         # Validate and format inputs for the pipeline
         pipeline_inputs = []
         valid_indices = []
+        contexts = []  # Store contexts for later use in cleaning
 
         for i, item in enumerate(inputs):
             utterance = item.get("utterance", "").strip()
@@ -194,6 +195,7 @@ class EndpointHandler:
                 'context': context  # Now includes conversation history
             })
             valid_indices.append(i)
+            contexts.append(context)  # Store for later use
 
         # Run prediction
         results = []
@@ -258,8 +260,12 @@ class EndpointHandler:
                     raw_answers = [p.get("answer", "") for p in current_preds]
                     raw_answers = [a for a in raw_answers if is_good_span(a)]
                     
-                    # Clean spans against ORIGINAL utterance (not full context)
-                    triggers = self._clean_spans(raw_answers, utterance)
+                    # Extract context text (part before [TARGET] marker)
+                    full_context = contexts[pred_idx]
+                    context_without_target = full_context.split("[TARGET]")[0].strip() if "[TARGET]" in full_context else ""
+                    
+                    # Clean spans against BOTH target utterance AND context
+                    triggers = self._clean_spans(raw_answers, utterance, context_without_target)
 
                     results.append({
                         "utterance": utterance,
@@ -280,12 +286,14 @@ class EndpointHandler:
                 "triggers": []
             } for item in inputs]
 
-    def _clean_spans(self, spans: List[str], target_text: str) -> List[str]:
+    def _clean_spans(self, spans: List[str], target_text: str, context_text: str = "") -> List[str]:
         """
         Clean and filter extracted trigger spans.
-        (Logic preserved exactly as provided)
+        Spans can come from either target_text or context_text.
         """
         target_text = target_text or ""
+        context_text = context_text or ""
+        full_text = (context_text + " " + target_text).lower()
         target_lower = target_text.lower()
 
         def _norm(s: str) -> str:
@@ -316,7 +324,8 @@ class EndpointHandler:
             s_norm = _norm(s)
             if not s_norm:
                 continue
-            if target_text and s_norm not in target_lower:
+            # Check if span exists in EITHER target OR context
+            if full_text and s_norm not in full_text:
                 continue
             tokens = s_norm.split()
             if len(tokens) > 8 or len(s_norm) > 80: