sail_scripts/data/categorizer/data_categorizer.py

"""
SAIL v2 Data Categorizer
=========================
Automatically classifies training data into semantic categories
using a multi-signal classification pipeline:

  1. Rule-based fast classifier   (keyword patterns, regex)
  2. Statistical classifier       (TF-IDF + Naive Bayes / SVM)
  3. Embedding-based classifier   (sentence-transformers cosine similarity)

Categories:
  text              — General prose / articles
  code              — Programming code (any language)
  math              — Mathematical problems, proofs, equations
  science           — Scientific content (biology, chemistry, physics)
  audio_transcript  — Transcribed speech from audio
  video_transcript  — Transcribed speech from video
  dialogue          — Conversational exchanges
  reasoning         — Chain-of-thought, logical reasoning
  creative          — Fiction, poetry, creative writing
  factual           — Encyclopedia, reference material
  instruction       — Task instructions, how-to guides, Q&A
  multilingual      — Non-English content

Each output JSONL gets a confidence score:
  {"category": "code", "confidence": 0.94, "signals": ["keyword", "embed"], ...}
"""

import os
import re
import json
import math
from typing import Dict, List, Tuple, Optional
from datetime import datetime
from utils.logger import get_logger

logger = get_logger("sail.categorizer")


# ─────────────────────────────────────────────────────────────────────────────
# Rule-Based Fast Classifier
# ─────────────────────────────────────────────────────────────────────────────
class RuleClassifier:
    """Fast O(n) keyword/pattern-based pre-classifier."""

    CODE_PATTERNS = [
        r"def\s+\w+\s*\(",
        r"class\s+\w+\s*[:\(]",
        r"import\s+\w+",
        r"from\s+\w+\s+import",
        r"function\s+\w+\s*\(",
        r"const\s+\w+\s*=",
        r"let\s+\w+\s*=",
        r"var\s+\w+\s*=",
        r"#include\s*<",
        r"fn\s+\w+\s*\(",
        r"public\s+class\s+\w+",
        r"SELECT\s+.+FROM\s+\w+",
        r"<[a-zA-Z]+[^>]*>.*</[a-zA-Z]+>",  # HTML
    ]

    MATH_PATTERNS = [
        r"\\[a-zA-Z]+\{",            # LaTeX
        r"∫|∑|∏|√|∂|∇|∞",          # Math symbols
        r"\$\$.*\$\$",               # LaTeX equations
        r"[A-Z]\s*=\s*[\d\w\s\+\-\*\/\^]+",  # equations
        r"\b(theorem|lemma|proof|corollary|conjecture|axiom)\b",
        r"\b(equation|formula|integral|derivative|matrix|vector)\b",
        r"\d+\s*[\+\-\*\/]\s*\d+\s*=\s*\d+",
    ]

    SCIENCE_KEYWORDS = {
        "biology", "chemistry", "physics", "genome", "dna", "rna",
        "protein", "molecule", "electron", "quantum", "photon",
        "hypothesis", "experiment", "cell", "organism", "reaction",
        "entropy", "relativity", "neuron", "catalyst", "isotope",
    }

    DIALOGUE_PATTERNS = [
        r"^(User|Human|Assistant|AI|Bot|Q|A):\s",
        r"^\-\s.+\n^\-\s",            # bullet dialogue
        r"\[INST\]|\[/INST\]",
        r"<\|user\|>|<\|assistant\|>",
    ]

    REASONING_KEYWORDS = {
        "step 1", "step 2", "therefore", "thus", "because",
        "as a result", "it follows that", "given that", "since",
        "chain of thought", "let me think", "reasoning:", "analysis:",
    }

    CREATIVE_KEYWORDS = {
        "once upon a time", "she whispered", "he said", "the sun set",
        "in a world where", "poem", "verse", "stanza", "rhyme",
        "protagonist", "antagonist", "chapter", "novel", "story",
    }

    TRANSCRIPT_KEYWORDS = {
        "speaker", "transcript", "[crosstalk]", "[inaudible]",
        "[laughter]", "interviewer", "interviewee", "host:", "guest:",
        "caption", "subtitle",
    }

    def __init__(self):
        self.code_re      = [re.compile(p, re.IGNORECASE | re.MULTILINE) for p in self.CODE_PATTERNS]
        self.math_re      = [re.compile(p, re.IGNORECASE | re.MULTILINE) for p in self.MATH_PATTERNS]
        self.dialogue_re  = [re.compile(p, re.IGNORECASE | re.MULTILINE) for p in self.DIALOGUE_PATTERNS]

    def classify(self, text: str, source_category: Optional[str] = None) -> Tuple[str, float]:
        """Returns (category, confidence)"""
        # Use source hint if confident
        if source_category in ("code", "audio_transcript", "video_transcript", "math"):
            return source_category, 0.90

        text_lower = text.lower()

        # Code: pattern matching
        code_hits = sum(1 for p in self.code_re if p.search(text))
        if code_hits >= 3:
            return "code", min(0.65 + code_hits * 0.05, 0.95)

        # Math: pattern + symbol density
        math_hits = sum(1 for p in self.math_re if p.search(text))
        if math_hits >= 2:
            return "math", min(0.60 + math_hits * 0.06, 0.95)

        # Dialogue
        dial_hits = sum(1 for p in self.dialogue_re if p.search(text))
        if dial_hits >= 2:
            return "dialogue", 0.80

        # Transcript
        tr_hits = sum(1 for k in self.TRANSCRIPT_KEYWORDS if k in text_lower)
        if tr_hits >= 2:
            return "audio_transcript", 0.75

        # Reasoning
        rs_hits = sum(1 for k in self.REASONING_KEYWORDS if k in text_lower)
        if rs_hits >= 3:
            return "reasoning", 0.75

        # Science
        sc_hits = sum(1 for k in self.SCIENCE_KEYWORDS if k in text_lower)
        if sc_hits >= 4:
            return "science", min(0.55 + sc_hits * 0.04, 0.90)

        # Creative
        cr_hits = sum(1 for k in self.CREATIVE_KEYWORDS if k in text_lower)
        if cr_hits >= 2:
            return "creative", 0.70

        # Instruction (Q&A format)
        if any(text_lower.startswith(p) for p in ["question:", "q:", "how to", "how do", "what is", "explain"]):
            return "instruction", 0.70

        # Factual (Wikipedia-style)
        if source_category == "factual" or "== history ==" in text_lower or "according to" in text_lower:
            return "factual", 0.65

        # Non-English detection (rough)
        ascii_ratio = sum(1 for c in text if ord(c) < 128) / max(len(text), 1)
        if ascii_ratio < 0.80:
            return "multilingual", 0.70

        return "text", 0.50


# ─────────────────────────────────────────────────────────────────────────────
# Statistical TF-IDF Classifier (boosted confidence)
# ─────────────────────────────────────────────────────────────────────────────
class TFIDFClassifier:
    """
    Fast TF-IDF based secondary classifier.
    Builds a vocabulary from category seed documents,
    then scores new documents against each category profile.
    """

    # Seed terms that are highly indicative of each category
    CATEGORY_SEEDS = {
        "code":             ["function", "variable", "return", "loop", "array", "object", "class", "method", "api", "algorithm", "compile", "syntax", "runtime"],
        "math":             ["theorem", "proof", "equation", "matrix", "integral", "derivative", "polynomial", "algebra", "calculus", "geometry", "statistics", "probability"],
        "science":          ["experiment", "hypothesis", "molecule", "cell", "species", "evolution", "quantum", "particle", "reaction", "biology", "chemistry", "physics"],
        "reasoning":        ["therefore", "because", "conclude", "infer", "logic", "premise", "argument", "analysis", "deduce", "evidence", "reasoning"],
        "creative":         ["story", "poem", "character", "plot", "narrative", "fiction", "imagination", "metaphor", "prose", "verse", "scene", "dialogue"],
        "instruction":      ["step", "tutorial", "guide", "install", "configure", "setup", "how-to", "procedure", "instruction", "follow"],
        "dialogue":         ["conversation", "chat", "response", "question", "answer", "user", "assistant", "message", "exchange"],
        "factual":          ["according", "history", "encyclopedia", "information", "description", "definition", "overview", "summary", "background"],
        "multilingual":     ["le", "la", "de", "en", "es", "der", "die", "das", "il", "est", "una", "este"],
        "audio_transcript": ["speaker", "said", "told", "mentioned", "discussed", "interview", "podcast", "recording"],
        "text":             ["the", "and", "or", "but", "however", "furthermore", "therefore", "article", "report"],
    }

    def __init__(self):
        self._build_profiles()

    def _tf(self, words: List[str]) -> Dict[str, float]:
        counts: Dict[str, int] = {}
        for w in words:
            counts[w] = counts.get(w, 0) + 1
        total = max(len(words), 1)
        return {w: c / total for w, c in counts.items()}

    def _build_profiles(self):
        self.profiles: Dict[str, Dict[str, float]] = {}
        for cat, seeds in self.CATEGORY_SEEDS.items():
            self.profiles[cat] = {w: 1.0 for w in seeds}

    def _cosine(self, a: Dict[str, float], b: Dict[str, float]) -> float:
        keys = set(a) & set(b)
        if not keys:
            return 0.0
        dot  = sum(a[k] * b[k] for k in keys)
        norm_a = math.sqrt(sum(v**2 for v in a.values()))
        norm_b = math.sqrt(sum(v**2 for v in b.values()))
        return dot / (norm_a * norm_b + 1e-9)

    def classify(self, text: str) -> Tuple[str, float]:
        words = re.findall(r'\b[a-z]{3,}\b', text.lower())
        if len(words) < 10:
            return "text", 0.40

        doc_tf = self._tf(words)
        scores = {cat: self._cosine(doc_tf, profile)
                  for cat, profile in self.profiles.items()}

        best_cat  = max(scores, key=scores.get)
        best_score = scores[best_cat]

        # Normalize to confidence
        total = sum(scores.values()) + 1e-9
        confidence = min(best_score / total * len(scores) * 0.5, 0.90)

        return best_cat, confidence


# ─────────────────────────────────────────────────────────────────────────────
# Ensemble Categorizer
# ─────────────────────────────────────────────────────────────────────────────
class EnsembleCategorizer:
    """
    Combines rule-based and TF-IDF classifiers with weighted voting.
    Optionally uses sentence-transformers for high-confidence disambiguation.
    """

    def __init__(self, confidence_threshold: float = 0.75):
        self.rule_clf   = RuleClassifier()
        self.tfidf_clf  = TFIDFClassifier()
        self.threshold  = confidence_threshold
        self._try_load_embed_model()

    def _try_load_embed_model(self):
        """Optionally load sentence-transformers for embedding-based classification."""
        self.embed_model = None
        try:
            from sentence_transformers import SentenceTransformer
            import numpy as np
            self.embed_model = SentenceTransformer("all-MiniLM-L6-v2")
            self._np = np
            logger.info("Sentence-transformers loaded for embedding classification")
            print("  ✓ Embedding classifier: all-MiniLM-L6-v2 (high accuracy)")
        except ImportError:
            print("  ℹ  sentence-transformers not installed — using rule+TF-IDF only")
            print("     (Install: pip install sentence-transformers for better accuracy)")

    def categorize(self, record: Dict) -> Dict:
        text            = record.get("text", "")
        source_category = record.get("category")

        # Stage 1: Rule-based
        rule_cat, rule_conf = self.rule_clf.classify(text, source_category)

        # Stage 2: TF-IDF
        tfidf_cat, tfidf_conf = self.tfidf_clf.classify(text)

        # Ensemble: weighted vote
        scores: Dict[str, float] = {}
        for cat, conf, weight in [(rule_cat, rule_conf, 0.5), (tfidf_cat, tfidf_conf, 0.4)]:
            scores[cat] = scores.get(cat, 0) + conf * weight

        # Source hint as weak signal
        if source_category in ("code", "math", "audio_transcript", "science"):
            scores[source_category] = scores.get(source_category, 0) + 0.3

        final_cat  = max(scores, key=scores.get)
        final_conf = min(scores[final_cat], 1.0)

        signals = ["rule", "tfidf"]

        # Stage 3: Embedding model for low-confidence cases
        if final_conf < self.threshold and self.embed_model and len(text) > 100:
            embed_cat, embed_conf = self._embed_classify(text)
            if embed_conf > final_conf:
                scores[embed_cat] = scores.get(embed_cat, 0) + embed_conf * 0.5
                final_cat  = max(scores, key=scores.get)
                final_conf = min(scores[final_cat], 1.0)
                signals.append("embed")

        return {
            **record,
            "category":       final_cat,
            "confidence":     round(final_conf, 3),
            "signals":        signals,
            "category_scores": {k: round(v, 3) for k, v in scores.items()},
        }

    def _embed_classify(self, text: str) -> Tuple[str, float]:
        CATEGORY_SENTENCES = {
            "code":        "This is programming code with functions, variables, and algorithms.",
            "math":        "This is a mathematical problem with equations, proofs, and calculations.",
            "science":     "This is scientific content about biology, chemistry, or physics.",
            "reasoning":   "This demonstrates logical reasoning and step-by-step thinking.",
            "creative":    "This is creative writing like fiction, poetry, or storytelling.",
            "instruction": "This is an instruction guide or tutorial with steps to follow.",
            "dialogue":    "This is a conversation or dialogue between people.",
            "factual":     "This is factual encyclopedic information and reference material.",
            "text":        "This is general informational text or an article.",
        }
        try:
            sample = text[:512]
            categories = list(CATEGORY_SENTENCES.keys())
            sentences  = [CATEGORY_SENTENCES[c] for c in categories]
            all_texts  = [sample] + sentences
            embeddings = self.embed_model.encode(all_texts, convert_to_numpy=True)
            doc_emb    = embeddings[0]
            cat_embs   = embeddings[1:]

            def cosine(a, b):
                return float(self._np.dot(a, b) / (self._np.linalg.norm(a) * self._np.linalg.norm(b) + 1e-9))

            sims = [cosine(doc_emb, ce) for ce in cat_embs]
            best_idx = int(self._np.argmax(sims))
            return categories[best_idx], float(sims[best_idx])
        except Exception as e:
            logger.warning(f"Embed classify error: {e}")
            return "text", 0.40


# ─────────────────────────────────────────────────────────────────────────────
# Main DataCategorizer
# ─────────────────────────────────────────────────────────────────────────────
class DataCategorizer:
    def __init__(self, settings):
        self.settings   = settings
        self.input_dir  = settings.categorizer_input_dir
        self.output_dir = settings.categorizer_output_dir
        os.makedirs(self.output_dir, exist_ok=True)

    def run(self):
        print(f"\n{'='*60}")
        print(f"  SAIL DATA CATEGORIZER")
        print(f"{'='*60}")
        print(f"  Input     : {self.input_dir}")
        print(f"  Output    : {self.output_dir}")
        print(f"  Threshold : {self.settings.categorizer_confidence_threshold}")
        print(f"{'='*60}\n")

        if self.settings.dry_run:
            print("  DRY RUN — skipping categorization\n")
            return

        categorizer = EnsembleCategorizer(self.settings.categorizer_confidence_threshold)

        total = low_conf = 0
        cat_counts: Dict[str, int] = {}
        out_files:  Dict[str, object] = {}

        for fname in sorted(os.listdir(self.input_dir)):
            if not fname.endswith(".jsonl"):
                continue
            path = os.path.join(self.input_dir, fname)
            print(f"  Processing: {fname}")

            with open(path, "r", encoding="utf-8") as f:
                for line in f:
                    line = line.strip()
                    if not line:
                        continue
                    try:
                        record = json.loads(line)
                    except json.JSONDecodeError:
                        continue

                    result = categorizer.categorize(record)
                    total += 1
                    cat   = result["category"]
                    conf  = result["confidence"]

                    if conf < self.settings.categorizer_confidence_threshold:
                        low_conf += 1

                    cat_counts[cat] = cat_counts.get(cat, 0) + 1

                    # Write to per-category output file
                    out_path = os.path.join(self.output_dir, f"{cat}.jsonl")
                    if cat not in out_files:
                        out_files[cat] = open(out_path, "a", encoding="utf-8")
                    out_files[cat].write(json.dumps(result, ensure_ascii=False) + "\n")

                    if total % 5000 == 0:
                        print(f"    → {total:,} records categorized")

        for f in out_files.values():
            f.close()

        print(f"\n  ✓ Categorization complete!")
        print(f"    Total records    : {total:,}")
        print(f"    Low confidence   : {low_conf:,} ({100*low_conf/max(total,1):.1f}%)")
        print(f"\n  Category distribution:")
        for cat, n in sorted(cat_counts.items(), key=lambda x: -x[1]):
            bar = "█" * min(int(n / max(cat_counts.values()) * 30), 30)
            print(f"    {cat:<20} {n:>8,}  {bar}")

        manifest = {
            "timestamp":       datetime.utcnow().isoformat(),
            "total":           total,
            "low_confidence":  low_conf,
            "categories":      cat_counts,
            "threshold":       self.settings.categorizer_confidence_threshold,
        }
        with open(os.path.join(self.output_dir, "category_manifest.json"), "w") as f:
            json.dump(manifest, f, indent=2)

        return cat_counts