Update model.py

model.py

@@ -6,448 +6,190 @@ from urllib.parse import urlparse
 import numpy as np
 import onnxruntime as ort
 
-import config
 from utils import preprocess_text
 
-# Global variables to hold the loaded models and sessions
+# --------- DEFAULT CONFIG ---------
+
+MODEL_DIR = "models"
+
+SPAM_THRESHOLD = 0.5
+
+SHORT_TEXT_WORD_COUNT = 6
+SHORT_TEXT_THRESHOLD = 0.65
+
+VERY_SHORT_TEXT_WORD_COUNT = 3
+VERY_SHORT_TEXT_THRESHOLD = 0.75
+
+LONG_TEXT_WORD_THRESHOLD = 80
+
+CHUNK_MAX_WORDS = 40
+MAX_CHUNKS = 24
+
+BLOCKED_URL_DOMAINS = set()
+
+METADATA_PATH = os.path.join("/", "metadata.json")
+
+# --------- GLOBALS ---------
+
 _binary_session = None
 _category_session = None
 _metadata = None
 
-# Signals that are strong spam indicators even in short messages.
+# --------- REGEX ---------
+
 SPAM_HINT_PATTERN = re.compile(
     r"(http|www|win|winner|claim|click|offer|bonus|urgent|verify|password|"
     r"account|bank|deposit|earn|investment|crypto|btc|telegram|airdrop|giveaway|jackpot|prize)",
     re.IGNORECASE,
 )
 
-BENIGN_WIN_CONTEXT_PATTERN = re.compile(
-    r"\b(won|win|winner)\b.*\b(match|game|tournament|league|race|finals|team|football|cricket|basketball)\b",
-    re.IGNORECASE,
-)
-
 SCAM_ACTION_PATTERN = re.compile(
     r"(claim|click|prize|reward|link|http|www|money|cash|gift|airdrop|crypto|account|verify|urgent)",
     re.IGNORECASE,
 )
 
 GIVEAWAY_OVERRIDE_PATTERN = re.compile(
-    r"(\b(won|winner|jackpot|lucky draw)\b.*\b(lambo|lamboo|prize|reward|gift|voucher|tesla|iphone|cash)\b)|"
-    r"(\b(claim|redeem)\b.*\b(prize|reward|gift|voucher)\b)",
+    r"(\b(won|winner|jackpot)\b.*\b(prize|reward|gift|voucher|iphone|cash)\b)",
     re.IGNORECASE,
 )
 
-SENTENCE_SPLIT_PATTERN = re.compile(r"(?<=[.!?])\s+|\n+")
-URL_TOKEN_PATTERN = re.compile(r"^(https?://\S+|www\.\S+)$", re.IGNORECASE)
 URL_ANY_PATTERN = re.compile(r"(https?://\S+|www\.\S+)", re.IGNORECASE)
-LINK_SPAM_CUE_PATTERN = re.compile(
-    r"(claim|verify|password|bank|urgent|winner|prize|reward|bonus|airdrop|crypto|"
-    r"deposit|investment|gift card|limited time|act now|suspended|login|otp|kyc)",
-    re.IGNORECASE,
-)
-
-BENIGN_ADULT_CONTEXT_PATTERN = re.compile(
-    r"(older than 18|over 18|under 18|age requirement|adult supervision|age limit|"
-    r"content rating|parental guidance|legal age|years old)",
-    re.IGNORECASE,
-)
-
-BENIGN_WORK_CONTEXT_PATTERN = re.compile(
-    r"(pull request|code review|deployment|sprint|bug fix|qa|release note|"
-    r"project update|meeting notes|standup|ticket|merge request|ci pipeline)",
-    re.IGNORECASE,
-)
-
-SHORT_BRAND_ALERT_PATTERN = re.compile(
-    r"^[A-Za-z0-9&'._+-]{2,32}\s*:\s*[^:]{2,90}[.!?]?$"
-)
 
-MONEY_JOB_SCAM_PATTERN = re.compile(
-    r"(\$\s?\d[\d,]*(?:\.\d+)?\s*/?\s*(day|week|month))|"
-    r"(earn\s+\$?\s?\d[\d,]*)|"
-    r"(get rich quick)|"
-    r"(no experience needed)",
+LINK_SPAM_CUE_PATTERN = re.compile(
+    r"(claim|verify|password|bank|urgent|winner|prize|reward|bonus|airdrop|crypto|deposit)",
     re.IGNORECASE,
 )
 
+# --------- LOADERS ---------
 
 def _load_metadata():
-    if os.path.exists(config.METADATA_PATH):
-        with open(config.METADATA_PATH, "r", encoding="utf-8") as f:
+    if os.path.exists(METADATA_PATH):
+        with open(METADATA_PATH, "r", encoding="utf-8") as f:
             return json.load(f)
 
     return {
-        "spam_threshold": config.SPAM_THRESHOLD,
-        "short_text_word_count": config.SHORT_TEXT_WORD_COUNT,
-        "short_text_threshold": config.SHORT_TEXT_THRESHOLD,
-        "very_short_text_word_count": config.VERY_SHORT_TEXT_WORD_COUNT,
-        "very_short_text_threshold": config.VERY_SHORT_TEXT_THRESHOLD,
+        "spam_threshold": SPAM_THRESHOLD
     }
 
 
 def load_models():
-    """Loads the model sessions from disk only once."""
     global _binary_session, _category_session, _metadata
 
     if _binary_session is None:
-        onnx_path = os.path.join(config.MODEL_DIR, "binary_model.onnx")
-        _binary_session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
-    
+        binary_path = os.path.join(MODEL_DIR, "binary_model.onnx")
+        _binary_session = ort.InferenceSession(binary_path)
+
     if _category_session is None:
-        onnx_path = os.path.join(config.MODEL_DIR, "category_model.onnx")
-        _category_session = ort.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])
-        
+        category_path = os.path.join(MODEL_DIR, "category_model.onnx")
+        _category_session = ort.InferenceSession(category_path)
+
     if _metadata is None:
         _metadata = _load_metadata()
 
 
-def _effective_threshold(raw_text, cleaned_text):
-    threshold = float(_metadata.get("spam_threshold", config.SPAM_THRESHOLD))
-    short_word_count = int(
-        _metadata.get("short_text_word_count", config.SHORT_TEXT_WORD_COUNT)
-    )
-    short_threshold = float(
-        _metadata.get("short_text_threshold", config.SHORT_TEXT_THRESHOLD)
-    )
-    very_short_word_count = int(
-        _metadata.get("very_short_text_word_count", config.VERY_SHORT_TEXT_WORD_COUNT)
-    )
-    very_short_threshold = float(
-        _metadata.get("very_short_text_threshold", config.VERY_SHORT_TEXT_THRESHOLD)
-    )
-
-    words = [w for w in cleaned_text.split(" ") if w]
-    has_spam_hint = bool(SPAM_HINT_PATTERN.search(raw_text or ""))
-
-    if not has_spam_hint:
-        if len(words) <= very_short_word_count:
-            threshold = max(threshold, very_short_threshold)
-        elif len(words) <= short_word_count:
-            threshold = max(threshold, short_threshold)
+# --------- HELPERS ---------
+
+def _contains_url(text):
+    return bool(URL_ANY_PATTERN.search(text or ""))
+
+
+def _effective_threshold(text):
+    threshold = float(_metadata.get("spam_threshold", SPAM_THRESHOLD))
+
+    words = text.split()
+    if len(words) <= VERY_SHORT_TEXT_WORD_COUNT:
+        threshold = max(threshold, VERY_SHORT_TEXT_THRESHOLD)
+    elif len(words) <= SHORT_TEXT_WORD_COUNT:
+        threshold = max(threshold, SHORT_TEXT_THRESHOLD)
 
     return threshold
 
 
-def _is_benign_win_context(raw_text):
-    if not raw_text:
-        return False
-    return bool(BENIGN_WIN_CONTEXT_PATTERN.search(raw_text)) and not bool(
-        SCAM_ACTION_PATTERN.search(raw_text)
-    )
-
-
-def _is_benign_context(raw_text):
-    if not raw_text:
-        return False
-    if bool(SCAM_ACTION_PATTERN.search(raw_text)):
-        return False
-    return bool(BENIGN_ADULT_CONTEXT_PATTERN.search(raw_text)) or bool(
-        BENIGN_WORK_CONTEXT_PATTERN.search(raw_text)
-    )
-
-
-def _extract_url_domains(raw_text: str) -> list[str]:
-    if not raw_text:
-        return []
-    domains = []
-    for m in URL_ANY_PATTERN.finditer(raw_text):
-        url = m.group(0).strip()
-        if url.lower().startswith("www."):
-            url = "https://" + url
-        try:
-            parsed = urlparse(url)
-            host = (parsed.netloc or "").lower().strip()
-        except Exception:
-            continue
-        if not host:
-            continue
-        if host.startswith("m."):
-            host = host[2:]
-        domains.append(host)
-    return domains
-
-
-def _has_blocked_domain(raw_text: str) -> bool:
-    blocked = set(getattr(config, "BLOCKED_URL_DOMAINS", set()))
-    if not blocked:
-        return False
-    domains = _extract_url_domains(raw_text)
-    if not domains:
-        return False
-    for host in domains:
-        if host in blocked:
-            return True
-        for base in blocked:
-            if host.endswith('.' + base):
-                return True
-    return False
-
-
-def _contains_url(raw_text: str) -> bool:
-    return bool(URL_ANY_PATTERN.search(raw_text or ""))
-
-
-def _has_link_spam_cues(raw_text: str) -> bool:
-    return bool(LINK_SPAM_CUE_PATTERN.search(raw_text or ""))
-
-
-def _split_long_text(text: str) -> list[str]:
-    max_words = int(getattr(config, "CHUNK_MAX_WORDS", 40))
-    max_chunks = int(getattr(config, "MAX_CHUNKS", 24))
-
-    parts = [p.strip() for p in SENTENCE_SPLIT_PATTERN.split(text or "") if p.strip()]
+def _split_text(text):
+    words = text.split()
     chunks = []
-    current = []
-    current_words = 0
-
-    for part in parts:
-        words = part.split()
-        if not words:
-            continue
-        if len(words) > max_words:
-            for i in range(0, len(words), max_words):
-                piece = " ".join(words[i : i + max_words]).strip()
-                if piece:
-                    chunks.append(piece)
-                if len(chunks) >= max_chunks:
-                    return chunks[:max_chunks]
-            continue
-
-        if current_words + len(words) > max_words and current:
-            chunks.append(" ".join(current).strip())
-            current = [part]
-            current_words = len(words)
-        else:
-            current.append(part)
-            current_words += len(words)
-
-        if len(chunks) >= max_chunks:
-            return chunks[:max_chunks]
-
-    if current and len(chunks) < max_chunks:
-        chunks.append(" ".join(current).strip())
-
-    return chunks[:max_chunks]
-
-
-def _predict_single(raw_text: str, cleaned_text: str) -> dict:
-    if _contains_url(raw_text):
-        if _has_blocked_domain(raw_text):
-            return {
-                "is_spam": True,
-                "confidence": 0.99,
-                "category": "spam",
-                "threshold_used": float(_metadata.get("spam_threshold", config.SPAM_THRESHOLD)),
-            }
-        if not _has_link_spam_cues(raw_text):
-            return {
-                "is_spam": False,
-                "confidence": 0.05,
-                "category": "normal",
-                "threshold_used": float(_metadata.get("spam_threshold", config.SPAM_THRESHOLD)),
-            }
-
-    # Prepare input for ONNX
-    # Input name was set to 'input' in conversion script.
-    # It expects StringTensorType([None, 1])
-    onnx_input = np.array([[cleaned_text]], dtype=object)
-    
-    # Binary prediction
-    binary_inputs = {_binary_session.get_inputs()[0].name: onnx_input}
-    # Output names are usually 'label' and 'probabilities'
-    binary_outputs = _binary_session.run(None, binary_inputs)
-    
-    # binary_outputs[1] is a list of dictionaries like [{'0': 0.9, '1': 0.1}]
-    # Let's verify the actual output format.
-    # usually it's [labels, [{0: prob, 1: prob}]]
-    probs = binary_outputs[1][0]
-    spam_prob = float(probs.get(1, probs.get('1', 0.0)))
 
-    threshold = _effective_threshold(raw_text, cleaned_text)
-    is_spam = spam_prob >= threshold
+    for i in range(0, len(words), CHUNK_MAX_WORDS):
+        chunk = " ".join(words[i:i + CHUNK_MAX_WORDS])
+        chunks.append(chunk)
 
-    if is_spam and spam_prob < 0.92 and _is_benign_win_context(raw_text):
-        is_spam = False
-    if is_spam and spam_prob < 0.85 and _is_benign_context(raw_text):
-        is_spam = False
+        if len(chunks) >= MAX_CHUNKS:
+            break
 
-    has_giveaway_override = bool(GIVEAWAY_OVERRIDE_PATTERN.search(raw_text or ""))
-    if not is_spam and has_giveaway_override and not _is_benign_win_context(raw_text):
-        is_spam = True
+    return chunks
 
-    short_brand_alert = bool(SHORT_BRAND_ALERT_PATTERN.match((raw_text or "").strip()))
-    money_job_scam = bool(MONEY_JOB_SCAM_PATTERN.search(raw_text or ""))
 
-    if cleaned_text.strip().lower() == "join now":
-        is_spam = False
+# --------- CORE ---------
+
+def _predict_single(raw_text, cleaned_text):
+    onnx_input = np.array([[cleaned_text]], dtype=object)
 
+    # Binary
+    inputs = {_binary_session.get_inputs()[0].name: onnx_input}
+    outputs = _binary_session.run(None, inputs)
+
+    probs = outputs[1][0]
+    spam_prob = float(probs.get(1, probs.get('1', 0.0)))
+
+    threshold = _effective_threshold(cleaned_text)
+    is_spam = spam_prob >= threshold
+
+    # Heuristics
     if not is_spam:
-        if money_job_scam and spam_prob >= max(0.55, threshold - 0.20):
-            is_spam = True
-        elif short_brand_alert and spam_prob >= max(0.50, threshold - 0.16):
+        if GIVEAWAY_OVERRIDE_PATTERN.search(raw_text or ""):
             is_spam = True
 
+    # Category
     if is_spam:
-        if money_job_scam:
-            category = "job_scam"
-        elif short_brand_alert:
-            category = "phishing"
-        elif has_giveaway_override:
-            category = "giveaway"
-        else:
-            category_inputs = {_category_session.get_inputs()[0].name: onnx_input}
-            category_outputs = _category_session.run(None, category_inputs)
-            category = str(category_outputs[0][0])
+        cat_inputs = {_category_session.get_inputs()[0].name: onnx_input}
+        cat_outputs = _category_session.run(None, cat_inputs)
+        category = str(cat_outputs[0][0])
     else:
         category = "normal"
 
     return {
-        "is_spam": bool(is_spam),
-        "confidence": float(spam_prob),
-        "category": str(category),
-        "threshold_used": float(threshold),
+        "is_spam": is_spam,
+        "confidence": spam_prob,
+        "category": category,
+        "threshold_used": threshold
     }
 
 
-def validate_message(text: str) -> tuple[bool, str]:
-    if text is None:
-        return False, "Input is required."
-    if not isinstance(text, str):
-        return False, "Input must be a string."
-    normalized = text.strip()
-    if not normalized:
-        return False, "Input cannot be empty."
-    if not any(ch.isalnum() for ch in normalized):
-        return True, ""
-    if len(normalized) < 2:
-        return True, ""
-    return True, ""
-
+# --------- PUBLIC API ---------
 
-def predict_message(text: str) -> dict:
+def predict_message(text):
     load_models()
-    cleaned_text = preprocess_text(text)
-    word_count = len([w for w in cleaned_text.split(" ") if w])
-    long_threshold = int(getattr(config, "LONG_TEXT_WORD_THRESHOLD", 80))
-
-    if word_count <= long_threshold:
-        pred = _predict_single(text, cleaned_text)
-        return {
-            "is_spam": pred["is_spam"],
-            "confidence": round(pred["confidence"], 4),
-            "category": pred["category"],
-            "threshold_used": round(pred["threshold_used"], 4),
-            "chunked": False,
-        }
-
-    chunks = _split_long_text(text)
-    if not chunks:
-        pred = _predict_single(text, cleaned_text)
-        return {
-            "is_spam": pred["is_spam"],
-            "confidence": round(pred["confidence"], 4),
-            "category": pred["category"],
-            "threshold_used": round(pred["threshold_used"], 4),
-            "chunked": False,
-        }
-
-    chunk_predictions = []
-    for chunk in chunks:
-        cp = _predict_single(chunk, preprocess_text(chunk))
-        chunk_predictions.append(cp)
-
-    highest = max(chunk_predictions, key=lambda x: x["confidence"])
-    spam_chunks = [cp for cp in chunk_predictions if cp["is_spam"]]
-    is_spam = len(spam_chunks) > 0
 
-    if is_spam:
-        representative = max(spam_chunks, key=lambda x: x["confidence"])
-    else:
-        representative = highest
+    cleaned = preprocess_text(text)
+    word_count = len(cleaned.split())
 
-    return {
-        "is_spam": bool(is_spam),
-        "confidence": round(float(highest["confidence"]), 4),
-        "category": str(representative["category"] if is_spam else "normal"),
-        "threshold_used": round(float(representative["threshold_used"]), 4),
-        "chunked": True,
-        "chunk_count": len(chunks),
-    }
+    if word_count <= LONG_TEXT_WORD_THRESHOLD:
+        pred = _predict_single(text, cleaned)
+        pred["confidence"] = round(pred["confidence"], 4)
+        return pred
 
+    chunks = _split_text(text)
+
+    preds = [_predict_single(chunk, preprocess_text(chunk)) for chunk in chunks]
+
+    best = max(preds, key=lambda x: x["confidence"])
 
-def run_model(text: str) -> dict:
-    ok, error = validate_message(text)
-    if not ok:
-        return {
-            "ok": False,
-            "error": error,
-            "input": text,
-        }
-    prediction = predict_message(text)
     return {
-        "ok": True,
-        "input": text.strip(),
-        "result": prediction,
+        "is_spam": any(p["is_spam"] for p in preds),
+        "confidence": round(best["confidence"], 4),
+        "category": best["category"] if best["is_spam"] else "normal",
+        "threshold_used": best["threshold_used"]
     }
 
 
-def update_model(text: str, label: int, category: str):
-    if text is None:
-        return
-    os.makedirs("dataset", exist_ok=True)
-    feedback_path = os.path.join("dataset", "feedback.jsonl")
-    payload = {
-        "text": text.strip(),
-        "label": int(label),
-        "category": str(category),
-    }
-    with open(feedback_path, "a", encoding="utf-8") as f:
-        f.write(json.dumps(payload, ensure_ascii=True) + "\n")
-
-
-def get_model_specs() -> dict:
-    specs = {
-        "model_dir": config.MODEL_DIR,
-        "binary_model_path": config.BINARY_MODEL_PATH,
-        "category_model_path": config.CATEGORY_MODEL_PATH,
-        "metadata_path": config.METADATA_PATH,
-        "spam_threshold": config.SPAM_THRESHOLD,
-        "word_max_features": getattr(config, "WORD_MAX_FEATURES", None),
-        "char_max_features": getattr(config, "CHAR_MAX_FEATURES", None),
-        "files": {
-            "binary_onnx_exists": os.path.exists(os.path.join(config.MODEL_DIR, "binary_model.onnx")),
-            "category_onnx_exists": os.path.exists(os.path.join(config.MODEL_DIR, "category_model.onnx")),
-            "metadata_exists": os.path.exists(config.METADATA_PATH),
-        },
-    }
-    try:
-        load_models()
-        specs["loaded"] = True
-        specs["runtime_threshold"] = _metadata.get("spam_threshold", config.SPAM_THRESHOLD)
-    except Exception as exc:
-        specs["loaded"] = False
-        specs["load_error"] = str(exc)
-    return specs
-
-
-def print_model_specs() -> None:
-    specs = get_model_specs()
-    print("Model Specs (ONNX)")
-    print(f"- Model dir: {specs['model_dir']}")
-    print(f"- Base threshold: {specs['spam_threshold']}")
-    print(
-        "- Files exist: "
-        f"binary_onnx={specs['files']['binary_onnx_exists']}, "
-        f"category_onnx={specs['files']['category_onnx_exists']}, "
-        f"metadata={specs['files']['metadata_exists']}"
-    )
-    if specs.get("loaded"):
-        print("- Loaded: True")
-        print(f"- Runtime threshold: {specs['runtime_threshold']}")
-    else:
-        print("- Loaded: False")
-        print(f"- Load error: {specs.get('load_error', 'unknown error')}")
+def run_model(text):
+    if not text or not isinstance(text, str):
+        return {"ok": False, "error": "Invalid input"}
 
+    result = predict_message(text)
 
-if __name__ == "__main__":
-    print_model_specs()
+    return {
+        "ok": True,
+        "input": text.strip(),
+        "result": result
+    }