Update handler.py

handler.py

@@ -1,6 +1,9 @@
-from typing import Any, Dict, List
+from typing import Any, Dict, List, Tuple
+import math
+import re
+
 import torch
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 
 class EndpointHandler:
@@ -19,56 +22,199 @@ class EndpointHandler:
             "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()]
+        self.generic_phrases = {
+            "new platform",
+            "platform",
+            "company",
+            "market",
+            "markets",
+            "system",
+            "technology",
+            "solution",
+            "services",
+            "service",
+            "product",
+            "products",
+            "tool",
+            "tools",
+        }
 
-        seen = set()
-        cleaned: List[str] = []
+        self.stopwords = {
+            "a", "an", "the", "and", "or", "of", "for", "to", "in", "on", "with",
+            "by", "at", "from", "into", "over", "under", "through", "across",
+            "is", "are", "was", "were", "be", "been", "being",
+            "this", "that", "these", "those", "it", "its", "their",
+            "new", "latest"
+        }
 
-        for kw in raw_parts:
-            kw_clean = " ".join(kw.split()).strip()
-            kw_lower = kw_clean.lower()
+    def _normalize_space(self, text: str) -> str:
+        return " ".join(text.split()).strip()
 
-            if not kw_clean:
-                continue
+    def _normalize_phrase(self, text: str) -> str:
+        text = self._normalize_space(text)
+        text = text.strip(" ,.;:-_")
+        return text
 
-            # Remove instruction leakage
-            if any(kw_lower.startswith(prefix) for prefix in self.bad_prefixes):
-                continue
+    def _phrase_tokens(self, text: str) -> List[str]:
+        return re.findall(r"[A-Za-z0-9][A-Za-z0-9\-+/.]*", text.lower())
 
-            # 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
+    def _contains_instruction_leakage(self, phrase_lower: str) -> bool:
+        return any(phrase_lower.startswith(prefix) for prefix in self.bad_prefixes)
 
-            # Skip very long outputs that are likely sentence fragments, not keywords
-            if len(kw_clean.split()) > 6:
-                continue
+    def _looks_sentence_like(self, phrase: str) -> bool:
+        lower = phrase.lower()
+        markers = [" and ", " because ", " which ", " where ", " when ", " while ", " after ", " before "]
+        if any(m in lower for m in markers) and len(phrase.split()) > 4:
+            return True
+        if phrase.endswith("."):
+            return True
+        return False
 
-            # 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
+    def _is_too_generic(self, phrase: str) -> bool:
+        lower = phrase.lower()
+        if lower in self.generic_phrases:
+            return True
 
-            # Trim surrounding punctuation
-            kw_clean = kw_clean.strip(" ,.;:-")
+        tokens = self._phrase_tokens(lower)
+        if len(tokens) == 1 and tokens[0] in self.generic_phrases:
+            return True
 
-            if not kw_clean:
-                continue
+        # phrases like "new platform" or "new system"
+        if len(tokens) == 2 and tokens[0] in {"new", "latest"} and tokens[1] in self.generic_phrases:
+            return True
+
+        return False
+
+    def _jaccard(self, a: List[str], b: List[str]) -> float:
+        sa, sb = set(a), set(b)
+        if not sa or not sb:
+            return 0.0
+        return len(sa & sb) / len(sa | sb)
+
+    def _text_coverage_score(self, phrase: str, source_text: str) -> float:
+        """
+        Soft relevance score using literal presence and token overlap.
+        Keeps semantically good present phrases near the top.
+        """
+        phrase_lower = phrase.lower()
+        source_lower = source_text.lower()
+
+        score = 0.0
+
+        if phrase_lower in source_lower:
+            score += 4.0
+
+        phrase_tokens = self._phrase_tokens(phrase)
+        source_tokens = self._phrase_tokens(source_text)
+
+        if not phrase_tokens:
+            return 0.0
+
+        overlap = len(set(phrase_tokens) & set(source_tokens))
+        score += overlap * 1.25
+        score += self._jaccard(phrase_tokens, source_tokens) * 2.0
+
+        # prefer 2–3 word phrases slightly
+        wc = len(phrase.split())
+        if wc == 2:
+            score += 1.0
+        elif wc == 3:
+            score += 0.75
+        elif wc == 1:
+            score += 0.25
+        elif wc >= 5:
+            score -= 1.0
+
+        # penalize generic lead words
+        if phrase_tokens and phrase_tokens[0] in self.stopwords:
+            score -= 0.75
+
+        return score
+
+    def _parse_candidates(self, generated_texts: List[str], source_text: str, max_keyword_words: int) -> List[str]:
+        source_lower = self._normalize_space(source_text.lower())
+        candidates: List[str] = []
+
+        for raw_text in generated_texts:
+            parts = [self._normalize_phrase(p) for p in raw_text.split(";")]
+            for part in parts:
+                if not part:
+                    continue
+
+                lower = part.lower()
 
-            # Dedupe case-insensitively
-            normalized = kw_clean.lower()
-            if normalized in seen:
+                if self._contains_instruction_leakage(lower):
+                    continue
+
+                if lower == source_lower:
+                    continue
+
+                if len(lower) > 30 and lower in source_lower:
+                    # likely near-complete echo
+                    continue
+
+                if self._looks_sentence_like(part):
+                    continue
+
+                wc = len(part.split())
+                if wc == 0 or wc > max_keyword_words:
+                    continue
+
+                if self._is_too_generic(part):
+                    continue
+
+                candidates.append(part)
+
+        return candidates
+
+    def _dedupe_and_prune(self, phrases: List[str], source_text: str, top_k: int) -> List[Tuple[str, float]]:
+        # First score
+        scored: List[Tuple[str, float]] = []
+        seen_exact = set()
+
+        for phrase in phrases:
+            norm = phrase.lower()
+            if norm in seen_exact:
                 continue
+            seen_exact.add(norm)
+
+            score = self._text_coverage_score(phrase, source_text)
+            if score > 0:
+                scored.append((phrase, score))
+
+        # Sort best first
+        scored.sort(key=lambda x: x[1], reverse=True)
+
+        # Remove subsumed / near-duplicate phrases
+        final_scored: List[Tuple[str, float]] = []
+        for phrase, score in scored:
+            ptoks = self._phrase_tokens(phrase)
+            pset = set(ptoks)
+
+            should_skip = False
+            for kept_phrase, kept_score in final_scored:
+                ktoks = self._phrase_tokens(kept_phrase)
+                kset = set(ktoks)
 
-            seen.add(normalized)
-            cleaned.append(kw_clean)
+                # exact token subset of a better phrase -> drop shorter one
+                if pset and pset.issubset(kset):
+                    should_skip = True
+                    break
 
-        return cleaned
+                # heavy overlap and shorter/weaker -> drop
+                jac = self._jaccard(ptoks, ktoks)
+                if jac >= 0.6:
+                    if len(ptoks) <= len(ktoks) and score <= kept_score + 0.5:
+                        should_skip = True
+                        break
+
+            if not should_skip:
+                final_scored.append((phrase, round(score, 4)))
+
+            if len(final_scored) >= top_k:
+                break
+
+        return final_scored
 
     def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]:
         text = data.get("inputs")
@@ -81,11 +227,20 @@ class EndpointHandler:
         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))
+        max_new_tokens = int(parameters.get("max_new_tokens", 32))
+        num_beams = int(parameters.get("num_beams", 6))
+        num_return_sequences = int(parameters.get("num_return_sequences", 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))
+        temperature = float(parameters.get("temperature", 0.9))
+        top_p = float(parameters.get("top_p", 0.95))
+        no_repeat_ngram_size = int(parameters.get("no_repeat_ngram_size", 2))
+        max_keyword_words = int(parameters.get("max_keyword_words", 4))
+        top_k_keywords = int(parameters.get("top_k_keywords", 6))
+        return_scores = bool(parameters.get("return_scores", False))
+
+        if not do_sample:
+            # beam search requires return_sequences <= beams
+            num_return_sequences = min(num_return_sequences, num_beams)
 
         encoded = self.tokenizer(
             text,
@@ -99,6 +254,7 @@ class EndpointHandler:
             **encoded,
             "max_new_tokens": max_new_tokens,
             "num_beams": num_beams,
+            "num_return_sequences": num_return_sequences,
             "do_sample": do_sample,
             "no_repeat_ngram_size": no_repeat_ngram_size,
             "early_stopping": True,
@@ -106,14 +262,37 @@ class EndpointHandler:
 
         if do_sample:
             generate_kwargs["temperature"] = temperature
+            generate_kwargs["top_p"] = top_p
 
         with torch.inference_mode():
             output_ids = self.model.generate(**generate_kwargs)
 
-        raw_text = self.tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
-        keywords = self._clean_keywords(raw_text, text)
+        generated_texts = self.tokenizer.batch_decode(output_ids, skip_special_tokens=True)
+        generated_texts = [self._normalize_space(t) for t in generated_texts if self._normalize_space(t)]
+
+        candidates = self._parse_candidates(
+            generated_texts=generated_texts,
+            source_text=text,
+            max_keyword_words=max_keyword_words,
+        )
+
+        ranked = self._dedupe_and_prune(
+            phrases=candidates,
+            source_text=text,
+            top_k=top_k_keywords,
+        )
+
+        keywords = [phrase for phrase, _ in ranked]
 
-        return {
-            "generated_text": raw_text,
+        response: Dict[str, Any] = {
+            "generated_texts": generated_texts,
             "keywords": keywords,
-        }
+        }
+
+        if return_scores:
+            response["keyword_scores"] = [
+                {"keyword": phrase, "score": score}
+                for phrase, score in ranked
+            ]
+
+        return response