Update handler.py

handler.py

@@ -12,45 +12,108 @@ class EndpointHandler:
         self.model.to(self.device)
         self.model.eval()
 
+        self.bad_prefixes = [
+            "extract keyphrases:",
+            "extract keywords:",
+            "keyphrases:",
+            "keywords:",
+        ]
+
+    def _clean_keywords(self, raw_text: str, source_text: str) -> List[str]:
+        source_lower = source_text.lower().strip()
+        raw_parts = [part.strip() for part in raw_text.split(";") if part.strip()]
+
+        seen = set()
+        cleaned: List[str] = []
+
+        for kw in raw_parts:
+            kw_clean = " ".join(kw.split()).strip()
+            kw_lower = kw_clean.lower()
+
+            if not kw_clean:
+                continue
+
+            # Remove instruction leakage
+            if any(kw_lower.startswith(prefix) for prefix in self.bad_prefixes):
+                continue
+
+            # Remove exact/near-full input echoes
+            if kw_lower == source_lower:
+                continue
+            if len(kw_lower) > 30 and kw_lower in source_lower:
+                continue
+            if len(source_lower) > 30 and source_lower in kw_lower:
+                continue
+
+            # Skip very long outputs that are likely sentence fragments, not keywords
+            if len(kw_clean.split()) > 6:
+                continue
+
+            # Skip obvious clause/sentence-like phrases
+            sentence_markers = [" and ", " because ", " that ", " which ", " where ", " when "]
+            if any(marker in kw_lower for marker in sentence_markers) and len(kw_clean.split()) > 4:
+                continue
+
+            # Trim surrounding punctuation
+            kw_clean = kw_clean.strip(" ,.;:-")
+
+            if not kw_clean:
+                continue
+
+            # Dedupe case-insensitively
+            normalized = kw_clean.lower()
+            if normalized in seen:
+                continue
+
+            seen.add(normalized)
+            cleaned.append(kw_clean)
+
+        return cleaned
+
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         text = data.get("inputs")
         if text is None:
             return {"error": "Missing required field: inputs"}
 
+        if not isinstance(text, str):
+            return {"error": "The 'inputs' field must be a string"}
+
         parameters = data.get("parameters", {})
 
+        max_input_length = int(parameters.get("max_input_length", 1024))
+        max_new_tokens = int(parameters.get("max_new_tokens", 48))
+        num_beams = int(parameters.get("num_beams", 4))
+        do_sample = bool(parameters.get("do_sample", False))
+        temperature = float(parameters.get("temperature", 1.0))
+        no_repeat_ngram_size = int(parameters.get("no_repeat_ngram_size", 3))
+
         encoded = self.tokenizer(
             text,
             return_tensors="pt",
             truncation=True,
-            max_length=int(parameters.get("max_input_length", 1024)),
+            max_length=max_input_length,
         )
         encoded = {k: v.to(self.device) for k, v in encoded.items()}
 
-        with torch.inference_mode():
-            output_ids = self.model.generate(
-                **encoded,
-                max_new_tokens=int(parameters.get("max_new_tokens", 48)),
-                num_beams=int(parameters.get("num_beams", 4)),
-                do_sample=bool(parameters.get("do_sample", False)),
-                temperature=float(parameters.get("temperature", 1.0)),
-                no_repeat_ngram_size=int(parameters.get("no_repeat_ngram_size", 3)),
-                early_stopping=True,
-            )
+        generate_kwargs = {
+            **encoded,
+            "max_new_tokens": max_new_tokens,
+            "num_beams": num_beams,
+            "do_sample": do_sample,
+            "no_repeat_ngram_size": no_repeat_ngram_size,
+            "early_stopping": True,
+        }
 
-        raw_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+        if do_sample:
+            generate_kwargs["temperature"] = temperature
 
-        keywords = [x.strip() for x in raw_text.split(";") if x.strip()]
+        with torch.inference_mode():
+            output_ids = self.model.generate(**generate_kwargs)
 
-        seen = set()
-        deduped: List[str] = []
-        for kw in keywords:
-            k = kw.lower()
-            if k not in seen:
-                seen.add(k)
-                deduped.append(kw)
+        raw_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+        keywords = self._clean_keywords(raw_text, text)
 
         return {
             "generated_text": raw_text,
-            "keywords": deduped,
+            "keywords": keywords,
         }