handler.py

import json
import re
import sys
import unicodedata
from pathlib import Path
from typing import Set

import numpy as np
from sentence_transformers import SentenceTransformer


# ── Inline text_utils (avoids sys.path / import issues on HF endpoints) ──

_ARABIC_DIACRITICS_RE = re.compile(
    "[\u0610-\u061A\u064B-\u065F\u0670\u06D6-\u06DC\u06DF-\u06E4"
    "\u06E7-\u06E8\u06EA-\u06ED]+"
)
_ALEF_VARIANTS = {
    "\u0622": "\u0627", "\u0623": "\u0627",
    "\u0625": "\u0627", "\u0671": "\u0627",
}
_ALEF_RE = re.compile("[" + "".join(_ALEF_VARIANTS.keys()) + "]")
_TURKISH_UPPER_MAP = {"\u0130": "i", "I": "\u0131"}
_ARABIC_PREFIXES = ("وال", "بال", "كال", "فال", "لل", "ال")

ALL_STOPWORDS: Set[str] = {
    # Arabic
    "في", "من", "على", "الى", "إلى", "عن", "مع", "هذا", "هذه",
    "ذلك", "تلك", "هو", "هي", "هم", "هن", "نحن", "انا", "أنا",
    "كان", "كانت", "يكون", "ان", "أن", "لا", "ما", "لم", "لن",
    "قد", "او", "أو", "و", "ثم", "بعد", "قبل", "حتى", "عند",
    "كل", "بعض", "غير", "بين", "الذي", "التي", "الذين", "اللاتي",
    "اللذان", "اللتان", "اذا", "إذا", "لو", "كي", "حيث",
    # English
    "the", "a", "an", "and", "or", "but", "in", "on", "at", "to",
    "for", "of", "with", "by", "from", "is", "it", "as", "was",
    "are", "were", "be", "been", "being", "have", "has", "had",
    "do", "does", "did", "will", "would", "could", "should", "may",
    "might", "shall", "can", "not", "no", "this", "that", "these",
    "those", "i", "you", "he", "she", "we", "they", "me", "him",
    "her", "us", "them", "my", "your", "his", "its", "our", "their",
    "so", "if", "then", "than", "too", "very", "just",
    # Turkish
    "ve", "bir", "bu", "de", "da", "ile", "icin", "için", "ama",
    "fakat", "veya", "ya", "ki", "ne", "o", "su", "şu", "ben",
    "sen", "biz", "siz", "onlar", "gibi", "daha", "en", "cok",
    "çok", "var", "yok", "mi", "mu", "mı", "mü", "ise", "hem",
    "hep", "her", "kadar", "sonra", "once", "önce", "bunu", "sunu",
    "şunu", "onu",
}


def _normalize_text(text: str) -> str:
    text = unicodedata.normalize("NFC", text)
    text = _ARABIC_DIACRITICS_RE.sub("", text)
    text = _ALEF_RE.sub(lambda m: _ALEF_VARIANTS[m.group()], text)
    text = text.replace("\u0629", "\u0647")  # taa marbuta -> haa
    text = text.replace("\u0649", "\u064A")  # alef maksura -> yaa
    for src, dst in _TURKISH_UPPER_MAP.items():
        text = text.replace(src, dst)
    text = text.lower()
    text = text.replace("-", " ").replace("_", " ")
    text = re.sub(r"\s+", " ", text).strip()
    return text


def tokenize(text: str) -> Set[str]:
    normalized = _normalize_text(text)
    tokens = set(normalized.split()) - ALL_STOPWORDS
    extra: Set[str] = set()
    for token in tokens:
        for prefix in _ARABIC_PREFIXES:
            if token.startswith(prefix) and len(token) > len(prefix) + 1:
                extra.add(token[len(prefix):])
    tokens |= extra
    return {t for t in tokens if len(t) > 1}


# ── Handler ──────────────────────────────────────────────────────────────

class EndpointHandler:
    def __init__(self, path: str = ""):
        model_dir = Path(path)

        # Load from model/ subdir to avoid ST auto-detection at repo root
        model_path = model_dir / "model"
        if model_path.exists():
            self.bi_encoder = SentenceTransformer(str(model_path))
        else:
            self.bi_encoder = SentenceTransformer(str(model_dir))

        # Load precomputed category embeddings (numpy, no faiss needed)
        self.category_embeddings = np.load(
            str(model_dir / "category_embeddings.npy")
        )

        with open(model_dir / "category_metadata.json", "r", encoding="utf-8") as f:
            self.categories = json.load(f)

    def __call__(self, data):
        inputs = data.get("inputs", "")
        params = data.get("parameters", {})
        top_k = params.get("top_k", 5)
        prefer_level = params.get("prefer_level", None)

        if isinstance(inputs, str):
            inputs = [inputs]

        results = []
        for text in inputs:
            preds = self._predict(text, top_k=top_k, prefer_level=prefer_level)
            results.append({"input": text, "predictions": preds})

        return results if len(results) > 1 else results[0]

    def _predict(self, product_text, top_k=5, prefer_level=None):
        # Bi-encoder retrieval via numpy dot product (embeddings are normalized)
        query_emb = self.bi_encoder.encode(
            [product_text], normalize_embeddings=True, convert_to_numpy=True
        )
        scores = (self.category_embeddings @ query_emb.T).flatten()

        # Top 20 candidates
        top_indices = np.argsort(scores)[::-1][:20]

        candidates = []
        for idx in top_indices:
            cat = self.categories[idx]
            candidates.append({**cat, "bi_score": float(scores[idx])})

        if not candidates:
            return []

        # Keyword boosting
        query_tokens = tokenize(product_text)
        for c in candidates:
            cat_text = f"{c['alias']} {c['name_en']} {c['name_ar']} {c['name_ku']}"
            cat_tokens = tokenize(cat_text)
            path_tokens = tokenize(" ".join(c["path_en"]))
            overlap = len(query_tokens & (cat_tokens | path_tokens))
            c["keyword_boost"] = overlap * 0.02
            c["bi_score_boosted"] = c["bi_score"] + c["keyword_boost"]

        candidates.sort(key=lambda x: x["bi_score_boosted"], reverse=True)

        if prefer_level is not None:
            for c in candidates:
                if c["level"] == prefer_level:
                    c["bi_score_boosted"] *= 1.1
            candidates.sort(key=lambda x: x["bi_score_boosted"], reverse=True)

        return [
            {
                "alias": c["alias"],
                "level": c["level"],
                "name_en": c["name_en"],
                "name_ar": c["name_ar"],
                "path_en": " > ".join(c["path_en"]),
                "path_ar": " > ".join(c["path_ar"]),
                "score": round(c["bi_score_boosted"], 4),
            }
            for c in candidates[:top_k]
        ]