audit/task_classifier.py

"""ARCHON SFT — Task Classifier

Mappe chaque sample (dataset source + contenu) vers un task_type
defini dans MTP_TASK_PROFILES. Approche hybride:
- Niveau 1 (rapide): lookup direct par dataset source
- Niveau 2 (raffine): heuristiques regex sur contenu (function-calling JSON, code, etc.)
- Niveau 3 (optionnel, lent): DistilBERT-66M classifier sur subset ambigu

Usage:
    classifier = TaskClassifier()
    task_type = classifier.classify(sample={"messages": [...]}, source_ds="oasst2")
"""

from __future__ import annotations

import re
from typing import Any


# --- Niveau 1: lookup direct dataset source -> task_type
DATASET_TASK_MAP: dict[str, str] = {
    # Orchestration / function-calling
    "glaiveai/glaive-function-calling-v2": "function_calling",
    "Salesforce/xlam-function-calling-60k": "function_calling",
    "NousResearch/hermes-function-calling-v1": "function_calling",
    "NousResearch/hermes-3-dataset": "tool_use_multi",
    "ToolBench/ToolBench": "tool_use_multi",
    "ai2-adapt-dev/MetaTool-Bench": "tool_use_multi",
    "openbmb/UltraInteract_sft": "tool_use_multi",
    "jondurbin/airoboros-3.2": "orchestration_seq",

    # Conversation multi-turn
    "OpenAssistant/oasst2": "general_conv",
    "OpenAssistant/oasst1": "general_conv",
    "HuggingFaceH4/ultrachat_200k": "general_conv",
    "LDJnr/Capybara": "general_conv",
    "HuggingFaceH4/no_robots": "general_conv",
    "databricks/databricks-dolly-15k": "general_conv",
    "teknium/OpenHermes-2.5": "tool_use_multi",
    "allenai/tulu-3-sft-mixture": "general_conv",

    # Reasoning CoT
    "microsoft/orca-math-word-problems-200k": "math_symbolic",
    "meta-math/MetaMathQA": "math_symbolic",
    "nvidia/OpenMathInstruct-2": "math_symbolic",
    "Open-Orca/SlimOrca": "reasoning_cot",
    "WizardLM/WizardLM_evol_instruct_70k": "reasoning_cot",

    # Code
    "bigcode/the-stack-dedup": "code",
    "codeparrot/github-code": "code",
    "codeparrot/apps": "code",
    "ise-uiuc/Magicoder-OSS-Instruct-75K": "code",
    "m-a-p/CodeFeedback-Filtered-Instruction": "code",
    "iamtarun/python_code_instructions_18k_alpaca": "code",
    "sahil2801/CodeAlpaca-20k": "code",
    "jtatman/python-code-dataset-500k": "code",

    # DPO / alignment
    "argilla/distilabel-intel-orca-dpo-pairs": "dpo_pair",
    "mlabonne/orpo-dpo-mix-40k": "dpo_pair",
    "HuggingFaceH4/ultrafeedback_binarized": "dpo_pair",
    "Intel/orca_dpo_pairs": "dpo_pair",

    # Networking (jescy525 own)
    "jescy525/archon-comm-v1-sources-light": "packets",
    "jescy525/archon-comm-v1-sources-packets": "packets",
    "jescy525/archon-etherlink-synth-v1": "orchestration_seq",

    # Trading shared (TOUS ASIs trade aussi, contributor light)
    "jescy525/paper-trades-feedback": "trading_ohlcv",
    "jescy525/binance-trading-272gb": "trading_ohlcv",
    "jescy525/synthetic-trading-reasoning": "trading_smc_ict",
    "jescy525/market-data-multi": "trading_ohlcv",
    "jescy525/crypto-multi-exchange": "trading_ohlcv",
    "jescy525/izanagi-sft-v1-sources": "trading_smc_ict",
    "gbharti/finance-alpaca": "finance_instruction",
    "AdaptLLM/finance-tasks": "finance_qa",
    "zeroshot/twitter-financial-news-sentiment": "trading_sentiment",
    "financial_phrasebank": "trading_sentiment",
    "nickmuchi/financial-classification": "trading_sentiment",
    "tastypear/unfiltered-financial-news-2010-2023": "trading_context",

    # NEXUS (coding/programmation)
    "nickrosh/Evol-Instruct-Code-80k-v1": "code",
    "smangrul/code-chat-assistant-v1": "code",
    "theblackcat102/evol-codealpaca-v1": "code",
    "bigcode/self-oss-instruct-sc2-exec-filter-50k": "code_verified",
    "Vezora/Tested-143k-Python-Alpaca": "code_verified",
    "jinaai/code_exercises": "code",
    "flytech/python-codes-25k": "code",
    "ajibawa-2023/Code-290k-ShareGPT": "code",
    "TokenBender/code_instructions_122k_alpaca_style": "code",
    "glaiveai/glaive-code-assistant-v3": "code",
    "LLM360/Code-Mix": "code",

    # AETHER (generalist + king reasoner + king trader)
    "open-thoughts/OpenThoughts-114k": "reasoning_cot",
    "AI-MO/NuminaMath-CoT": "math_symbolic",
    "AI-MO/NuminaMath-TIR": "math_symbolic",
    "TIGER-Lab/MATH-plus": "math_symbolic",
    "lighteval/MATH": "math_symbolic",
    "open-r1/OpenR1-Math-220k": "math_symbolic",
    "WizardLM/WizardLM_evol_instruct_V2_196k": "reasoning_cot",
    "garage-bAInd/Open-Platypus": "reasoning_cot",
    "stingning/ultrachat": "general_conv",
    "pankajmathur/orca_minis_uncensored": "reasoning_cot",
    "allenai/ai2_arc": "reasoning_cot",
    "Rowan/hellaswag": "reasoning_cot",
    "gsm8k": "math_symbolic",
    "camel-ai/math": "math_symbolic",
    "camel-ai/code": "code",
    "b-mc2/sql-create-context": "code",
    "allenai/dolmino-mix-1124": "general_conv",
    "CohereForAI/aya_dataset": "general_conv",
    "Helsinki-NLP/opus-100": "general_conv",

    # GENESIS (formal proofs / theorem proving)
    "hoskinson-center/proofnet": "theorem_proof",
    "EleutherAI/proof-pile-2": "theorem_proof",
    "deepmind/math_dataset": "math_symbolic",
    "MARIO-Math-Reasoning/MARIO_EVAL": "theorem_proof",
    "Goader/lean4-mathlib-theorem-statement": "theorem_proof",
    "wellecks/naturalproofs": "theorem_proof",
    "cai-lw/FOLIO": "theorem_proof",
    "tasksource/logical-deduction": "reasoning_cot",
    "hendrycks/competition_math": "math_symbolic",
    "zhengxuanzenwu/lean4-verified-proofs": "theorem_proof",

    # NOUS (S-expr + GP)
    "deepmind/code_contests": "sexpr_program",
    "codeparrot/codeparrot-clean": "sexpr_program",
    "bigcode/humanevalpack": "code_verified",
    "bigcode/bigcodebench": "code_verified",
    "nuprl/EditPackFT": "gp_mutation",
    "nuprl/MultiPL-E": "gp_crossover",
    "TIGER-Lab/TheoremQA": "theorem_proof",
    "nvidia/OpenMathInstruct-2": "math_symbolic",
    "greengerong/leetcode": "code_verified",
    "arcee-ai/EvolKit-20k": "gp_mutation",
    "jescy525/synthetic-code": "code",

    # CYPHER (cybersec)
    "mitre/cti": "threat_intel",
    "CyberNative/Cybersecurity_Data_Science_Dataset": "malware_analysis",
    "jphwang/cve_data": "cve_explanation",
    "Andyrasika/network-intrusion-detection": "network_attack",
    "dnth/network-intrusion-detection-dataset": "network_attack",
    "jescy525/cypher-corpus": "malware_analysis",
    "jescy525/cypher-malware-research-cc-by-nc": "malware_analysis",
    "jescy525/cypher-sft-v3-sources": "malware_analysis",
    "jescy525/cypher-trades-own": "trading_ohlcv",

    # SHIZUNE (CRM FR / secretaire)
    "bofenghuang/vigogne": "general_conv",
    "legmlai/openhermes-fr": "general_conv",
    "gretelai/synthetic-text-to-sql": "code",
    "argilla/distilabel-math-preference-dpo": "dpo_pair",
}


# --- Niveau 2: heuristiques regex sur contenu (raffinage)
RE_JSON_FUNCTION_CALL = re.compile(
    r'(?:"name"\s*:\s*"[\w_]+"\s*,\s*"arguments"|<tool_call>|<function_call>|"function"\s*:\s*\{)',
    re.IGNORECASE
)
RE_TOOL_RESULT = re.compile(r'(?:<tool_result>|"tool_response"|"function_response")', re.IGNORECASE)
RE_PROTOBUF = re.compile(r'(?:syntax\s*=\s*"proto[23]"|message\s+\w+\s*\{|rpc\s+\w+\s*\()', re.IGNORECASE)
RE_GRPC = re.compile(r'(?:grpc\.|GrpcChannel|grpcurl|protoc)', re.IGNORECASE)
RE_PACKET_FIELD = re.compile(r'(?:tcp\.flags|udp\.length|ip\.src|tcp\.seq|0x[0-9a-f]{8})', re.IGNORECASE)
RE_LOG_TEMPLATE = re.compile(r'(?:^\[\d{4}-\d{2}-\d{2}|^\d{2}:\d{2}:\d{2}\.\d+|level=\w+\s+msg=)', re.MULTILINE)
RE_CODE_BOILERPLATE = re.compile(
    r'(?:^import\s+\w+|^from\s+\w+\s+import|^package\s+\w+|^use\s+\w+::|^#include\s*<|^fn\s+\w+\s*\()',
    re.MULTILINE
)
RE_COT_STEPS = re.compile(r'(?:Step\s+\d+[:.]|Étape\s+\d+|Let me think|First[,.]|Second[,.]|Third)', re.IGNORECASE)
RE_MATH_SYMBOLIC = re.compile(r'(?:\\frac\{|\\int|\\sum|\\sqrt|\$.+\$|=\s*\\?[a-z]\w*\s*\(|let\s+\w+\s*=)', re.IGNORECASE)


def heuristic_refine(sample_text: str, base_task: str) -> str:
    """Affine le task_type via heuristiques contenu. Garde base_task si pas de signal fort."""
    if not sample_text:
        return base_task

    # Function calling has highest priority detection
    if RE_JSON_FUNCTION_CALL.search(sample_text) or RE_TOOL_RESULT.search(sample_text):
        return "function_calling"
    if RE_PROTOBUF.search(sample_text) or RE_GRPC.search(sample_text):
        return "protobuf_grpc"
    if RE_PACKET_FIELD.search(sample_text):
        return "packets"
    if RE_LOG_TEMPLATE.search(sample_text):
        return "logs_traces"
    if RE_MATH_SYMBOLIC.search(sample_text):
        return "math_symbolic"
    if RE_COT_STEPS.search(sample_text):
        if base_task == "general_conv":
            return "reasoning_cot"
    if RE_CODE_BOILERPLATE.search(sample_text):
        if base_task in ("general_conv", "default"):
            return "code"
    return base_task


def extract_text(sample: dict[str, Any]) -> str:
    """Concat tout le contenu d'un sample (messages chat ou champ text plat)."""
    if "messages" in sample:
        parts = []
        for m in sample["messages"]:
            parts.append(m.get("content", "") or "")
        return "\n".join(parts)
    return sample.get("text") or sample.get("content") or ""


class TaskClassifier:
    """Classifier deterministe (Niveau 1 + 2)."""

    def __init__(self, enable_heuristics: bool = True):
        self.enable_heuristics = enable_heuristics

    def classify(self, sample: dict[str, Any], source_ds: str) -> str:
        # Niveau 1
        base = DATASET_TASK_MAP.get(source_ds, "default")
        if not self.enable_heuristics:
            return base
        # Niveau 2 refinement
        text = extract_text(sample)
        return heuristic_refine(text, base)


if __name__ == "__main__":
    cls = TaskClassifier()
    tests = [
        ({"messages": [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hello"}]},
         "OpenAssistant/oasst2"),
        ({"messages": [{"role": "user", "content": "Call weather"},
                       {"role": "assistant", "content": '<tool_call>{"name":"get_weather","arguments":{"city":"Paris"}}</tool_call>'}]},
         "OpenAssistant/oasst2"),
        ({"text": "import torch\nfrom transformers import AutoModel\nmodel = AutoModel.from_pretrained('bert')"},
         "bigcode/the-stack-dedup"),
        ({"text": "Step 1: identify the variable.\nStep 2: solve for x.\nStep 3: verify."},
         "default"),
    ]
    for s, src in tests:
        print(f"  src={src:40s} -> {cls.classify(s, src)}")