""" 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]+[^>]*>.*", # 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