npcpy 1.1.28__py3-none-any.whl → 1.2.32__py3-none-any.whl

npcpy/ft/model_ensembler.py

@@ -0,0 +1,357 @@
+import time
+import copy
+import random
+from dataclasses import dataclass, field
+from typing import List, Dict, Any, Optional
+from npcpy.llm_funcs import get_llm_response
+
+try:
+    from npcpy.ft.sft import predict_sft, load_sft_model
+except:
+    pass
+
+@dataclass
+class ModelGene:
+    """
+    Represents a specialized model with trigger patterns 
+    and confidence threshold
+    """
+    sft_path: Optional[str] = None
+    rl_path: Optional[str] = None
+    base_model: str = "Qwen/Qwen3-0.6B"
+    specialization: str = "general"
+    trigger_patterns: List[str] = field(default_factory=list)
+    confidence_threshold: float = 0.7
+
+
+def generate_trigger_patterns(specialization: str) -> List[str]:
+    """
+    Generate trigger patterns for a given specialization domain
+    """
+    patterns = {
+        'math': ['calculate', 'solve', 'equation', 'number'],
+        'code': ['function', 'class', 'bug', 'debug', 'code'],
+        'creative': ['story', 'poem', 'creative', 'imagine'],
+        'factual': ['what is', 'who is', 'when did', 'where is'],
+        'analysis': ['analyze', 'compare', 'evaluate', 'assess']
+    }
+    
+    return patterns.get(specialization, ['general'])
+
+
+def create_model_genome(
+    specializations: List[str],
+    base_model: str = "Qwen/Qwen3-0.6B"
+) -> List[ModelGene]:
+    """
+    Initialize a genome of specialized models
+    """
+    genome = []
+    
+    for spec in specializations:
+        gene = ModelGene(
+            base_model=base_model,
+            specialization=spec,
+            trigger_patterns=generate_trigger_patterns(spec),
+            confidence_threshold=random.uniform(0.6, 0.9)
+        )
+        genome.append(gene)
+    
+    return genome
+
+
+def mutate_model_genome(
+    genome: List[ModelGene],
+    mutation_type: str = 'random'
+) -> List[ModelGene]:
+    """
+    Apply genetic mutation to model genome
+    """
+    new_genome = copy.deepcopy(genome)
+    
+    mutations = [
+        'adjust_threshold',
+        'add_trigger',
+        'remove_gene',
+        'duplicate_gene'
+    ]
+    
+    if mutation_type == 'random':
+        mutation_type = random.choice(mutations)
+    
+    if mutation_type == 'adjust_threshold':
+        gene = random.choice(new_genome)
+        gene.confidence_threshold += random.uniform(-0.1, 0.1)
+        gene.confidence_threshold = max(
+            0.5,
+            min(0.95, gene.confidence_threshold)
+        )
+    
+    elif mutation_type == 'add_trigger':
+        gene = random.choice(new_genome)
+        new_trigger = f"pattern_{random.randint(1, 100)}"
+        if new_trigger not in gene.trigger_patterns:
+            gene.trigger_patterns.append(new_trigger)
+    
+    elif mutation_type == 'remove_gene' and len(new_genome) > 1:
+        new_genome.pop(random.randint(0, len(new_genome) - 1))
+    
+    elif mutation_type == 'duplicate_gene':
+        gene = random.choice(new_genome)
+        new_gene = copy.deepcopy(gene)
+        new_gene.specialization = f"{gene.specialization}_variant"
+        new_genome.append(new_gene)
+    
+    return new_genome
+
+
+def crossover_model_genomes(
+    genome1: List[ModelGene],
+    genome2: List[ModelGene]
+) -> List[ModelGene]:
+    """
+    Crossover two model genomes to create child genome
+    """
+    if not genome1 or not genome2:
+        return genome1 or genome2
+    
+    split = random.randint(1, min(len(genome1), len(genome2)) - 1)
+    
+    child = genome1[:split] + genome2[split:]
+    
+    return child
+
+
+def evaluate_model_genome(
+    genome: List[ModelGene],
+    test_cases: List[Dict[str, Any]],
+    router: 'ResponseRouter'
+) -> float:
+    """
+    Evaluate fitness of a model genome based on accuracy, 
+    speed and efficiency
+    """
+    correct = 0
+    total_time = 0
+    fast_responses = 0
+    
+    for test_case in test_cases:
+        result = router.route_query(
+            test_case['query'],
+            genome,
+            test_case.get('ground_truth')
+        )
+        
+        if result['correct']:
+            correct += 1
+        
+        total_time += result['response_time']
+        
+        if result['used_fast_path']:
+            fast_responses += 1
+    
+    accuracy = correct / len(test_cases)
+    speed_bonus = fast_responses / len(test_cases)
+    efficiency = 1.0 / (total_time / len(test_cases))
+    
+    fitness = (
+        accuracy * 0.6 +
+        speed_bonus * 0.2 +
+        efficiency * 0.2
+    )
+    
+    return fitness
+
+
+class ResponseRouter:
+    """
+    Routes queries through fast path, ensemble or full reasoning 
+    based on confidence thresholds
+    """
+    def __init__(
+        self,
+        fast_threshold: float = 0.8,
+        ensemble_threshold: float = 0.6
+    ):
+        self.fast_threshold = fast_threshold
+        self.ensemble_threshold = ensemble_threshold
+        self.response_cache = {}
+    
+    def route_query(
+        self,
+        query: str,
+        genome: List[ModelGene],
+        ground_truth: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Route query through system 1 fast path, 
+        ensemble or system 2 reasoning
+        """
+        start_time = time.time()
+        
+        fast_response = self._try_fast_path(query, genome)
+        
+        if fast_response and fast_response['confidence'] > (
+            self.fast_threshold
+        ):
+            response_time = time.time() - start_time
+            
+            return {
+                'response': fast_response['answer'],
+                'confidence': fast_response['confidence'],
+                'used_fast_path': True,
+                'response_time': response_time,
+                'correct': (
+                    ground_truth is None or 
+                    self._check_correctness(
+                        fast_response['answer'],
+                        ground_truth
+                    )
+                )
+            }
+        
+        ensemble_response = self._try_ensemble(query, genome)
+        
+        if ensemble_response['confidence'] > (
+            self.ensemble_threshold
+        ):
+            response_time = time.time() - start_time
+            
+            return {
+                'response': ensemble_response['answer'],
+                'confidence': ensemble_response['confidence'],
+                'used_fast_path': False,
+                'used_ensemble': True,
+                'response_time': response_time,
+                'correct': (
+                    ground_truth is None or
+                    self._check_correctness(
+                        ensemble_response['answer'],
+                        ground_truth
+                    )
+                )
+            }
+        
+        full_response = self._full_reasoning(query)
+        response_time = time.time() - start_time
+        
+        return {
+            'response': full_response,
+            'confidence': 0.5,
+            'used_fast_path': False,
+            'used_ensemble': False,
+            'response_time': response_time,
+            'correct': (
+                ground_truth is None or
+                self._check_correctness(
+                    full_response,
+                    ground_truth
+                )
+            )
+        }
+    
+    def _try_fast_path(
+        self,
+        query: str,
+        genome: List[ModelGene]
+    ) -> Optional[Dict[str, Any]]:
+        """
+        Try fast system 1 gut reaction using pattern matching
+        """
+        query_lower = query.lower()
+        
+        for gene in genome:
+            if any(
+                pattern in query_lower 
+                for pattern in gene.trigger_patterns
+            ):
+                if gene.sft_path:
+                    model, tokenizer = load_sft_model(gene.sft_path)
+                    
+                    response = predict_sft(
+                        model,
+                        tokenizer,
+                        query,
+                        temperature=0.1
+                    )
+                    
+                    return {
+                        'answer': response,
+                        'confidence': gene.confidence_threshold
+                    }
+        
+        return None
+    
+    def _try_ensemble(
+        self,
+        query: str,
+        genome: List[ModelGene]
+    ) -> Dict[str, Any]:
+        """
+        Try ensemble voting across specialized models
+        """
+        responses = []
+        
+        for gene in genome:
+            if gene.sft_path or gene.rl_path:
+                model_path = gene.rl_path or gene.sft_path
+                
+                model, tokenizer = load_sft_model(model_path)
+                
+                response = predict_sft(
+                    model,
+                    tokenizer,
+                    query,
+                    temperature=0.3
+                )
+                
+                responses.append({
+                    'answer': response,
+                    'weight': gene.confidence_threshold
+                })
+        
+        if not responses:
+            return {'answer': '', 'confidence': 0.0}
+        
+        best_response = max(responses, key=lambda x: x['weight'])
+        
+        avg_confidence = sum(
+            r['weight'] for r in responses
+        ) / len(responses)
+        
+        return {
+            'answer': best_response['answer'],
+            'confidence': avg_confidence
+        }
+    
+    def _full_reasoning(
+        self,
+        query: str,
+        model: str = "qwen3:1.7b",
+        provider: str = "ollama"
+    ) -> str:
+        """
+        Fall back to full system 2 reasoning
+        """
+        response = get_llm_response(
+            query,
+            model=model,
+            provider=provider
+        )
+        
+        return response.get('response', '')
+    
+    def _check_correctness(
+        self,
+        response: str,
+        ground_truth: str
+    ) -> bool:
+        """
+        Check if response matches ground truth
+        """
+        response_lower = response.lower().strip()
+        truth_lower = ground_truth.lower().strip()
+        
+        return response_lower == truth_lower or (
+            truth_lower in response_lower
+        )

npcpy/ft/rl.py

@@ -0,0 +1,360 @@
+from dataclasses import dataclass
+
+from datetime import datetime
+import glob
+import json
+import os
+import pandas as pd
+try:
+    from datasets import Dataset
+
+    from peft import LoraConfig, PeftModel
+    import torch
+    from transformers import (
+        AutoModelForCausalLM,
+        AutoTokenizer
+    )
+    from trl import DPOTrainer, DPOConfig
+except:
+    Dataset = None
+    PeftModel = None
+    DPOConfig = None
+    DPOTrainer = None
+    torch = None
+    AutoModelForCausalLM = None
+    AutoTokenizer = None
+
+    
+import random
+from typing import List, Dict, Any, Optional, Callable
+from npcpy.npc_compiler import NPC
+from npcpy.llm_funcs import get_llm_response
+
+
+@dataclass
+class RLConfig:
+    base_model_name: str = "Qwen/Qwen3-0.6B"
+    adapter_path: str = "./rl_adapter"
+    max_iterations: int = 8
+    min_reward_gap: float = 0.4
+    num_train_epochs: int = 20
+    per_device_train_batch_size: int = 1
+    gradient_accumulation_steps: int = 2
+    learning_rate: float = 1e-6
+    beta: float = 0.5
+    max_length: int = 512
+    max_prompt_length: int = 256
+
+
+class TaskExecutor:
+
+    def __init__(
+        self,
+        agent: NPC,
+        max_iterations: int = 8
+    ):
+        self.agent = agent
+        self.max_iterations = max_iterations
+    
+    def execute_task(
+        self,
+        task_prompt: str
+    ) -> Dict[str, Any]:
+
+        messages = [
+            {
+                "role": "system",
+                "content": self.agent.primary_directive
+            }
+        ]
+        
+        raw_responses = []
+        current_prompt = task_prompt
+        
+        for i in range(self.max_iterations):
+            response_obj = self.agent.get_llm_response(
+                current_prompt,
+                messages=messages,
+                auto_process_tool_calls=True
+            )
+            
+            raw_responses.append(response_obj)
+            messages = response_obj.get('messages', messages)
+            
+            last_content = messages[-1].get('content', '')
+            
+            if self._is_complete(last_content):
+                return {
+                    "raw_responses": raw_responses,
+                    "final_output": last_content,
+                    "total_iterations": i + 1,
+                    "completed": True
+                }
+            
+            current_prompt = (
+                "Continue or provide final answer."
+            )
+        
+        return {
+            "raw_responses": raw_responses,
+            "final_output": messages[-1].get('content', ''),
+            "total_iterations": self.max_iterations,
+            "completed": False
+        }
+    
+    def _is_complete(self, content: str) -> bool:
+
+        completion_markers = [
+            "final answer:",
+            "conclusion:",
+            "result:",
+            "therefore",
+            "in summary"
+        ]
+        content_lower = content.lower()
+        return any(
+            marker in content_lower 
+            for marker in completion_markers
+        )
+
+
+def collect_traces(
+    tasks: List[Dict[str, Any]],
+    agents: List[NPC],
+    reward_fn: Callable[[Dict], float],
+    config: Optional[RLConfig] = None
+) -> List[Dict[str, Any]]:
+
+    if config is None:
+        config = RLConfig()
+    
+    traces = []
+    
+    for task in tasks:
+        task_prompt = task.get('prompt', task.get('input', ''))
+        
+        for agent in agents:
+            executor = TaskExecutor(
+                agent,
+                max_iterations=config.max_iterations
+            )
+            
+            result = executor.execute_task(task_prompt)
+            
+            trace = {
+                "agent_name": agent.name,
+                "task_prompt": task_prompt,
+                "final_output": result['final_output'],
+                "total_iterations": result['total_iterations'],
+                "completed": result['completed'],
+                "task_metadata": task
+            }
+            
+            trace['reward'] = reward_fn(trace)
+            
+            traces.append(trace)
+            
+            print(
+                f"Agent {agent.name}: "
+                f"Reward={trace['reward']:.2f}"
+            )
+    
+    return traces
+
+
+def create_preference_pairs(
+    traces: List[Dict[str, Any]],
+    min_reward_gap: float = 0.4
+) -> Dataset:
+
+    df = pd.DataFrame(traces)
+    df = df[df['reward'] > -1.0].copy()
+    
+    if len(df) < 2:
+        return None
+    
+    df = df.sort_values('reward', ascending=False)
+    
+    top_quantile = df['reward'].quantile(
+        0.8,
+        interpolation='higher'
+    )
+    low_quantile = df['reward'].quantile(
+        0.2,
+        interpolation='lower'
+    )
+    
+    high_traces = df[df['reward'] >= top_quantile]
+    low_traces = df[df['reward'] <= low_quantile]
+    
+    pairs = []
+    
+    for _, high_trace in high_traces.iterrows():
+        for _, low_trace in low_traces.iterrows():
+            reward_gap = (
+                high_trace['reward'] - low_trace['reward']
+            )
+            
+            if reward_gap >= min_reward_gap:
+                pairs.append({
+                    "prompt": str(high_trace['task_prompt']),
+                    "chosen": str(high_trace['final_output']),
+                    "rejected": str(low_trace['final_output'])
+                })
+    
+    if len(pairs) < 5:
+        print(
+            f"Warning: Only {len(pairs)} pairs found. "
+            "May overfit."
+        )
+    
+    return Dataset.from_list(pairs[:100])
+
+
+def train_with_dpo(
+    traces: List[Dict[str, Any]],
+    config: Optional[RLConfig] = None
+) -> str:
+
+    if config is None:
+        config = RLConfig()
+    
+    preference_dataset = create_preference_pairs(
+        traces,
+        min_reward_gap=config.min_reward_gap
+    )
+    
+    if preference_dataset is None or len(preference_dataset) == 0:
+        print("No valid preference pairs. Cannot train.")
+        return None
+    
+    model = AutoModelForCausalLM.from_pretrained(
+        config.base_model_name,
+        torch_dtype=torch.float32,
+        device_map="auto",
+        low_cpu_mem_usage=True
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        config.base_model_name,
+        trust_remote_code=True
+    )
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    peft_config = LoraConfig(
+        r=8,
+        lora_alpha=16,
+        lora_dropout=0.1,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=[
+            "q_proj",
+            "k_proj",
+            "v_proj",
+            "o_proj"
+        ]
+    )
+    
+    training_args = DPOConfig(
+        output_dir="./dpo_results",
+        per_device_train_batch_size=(
+            config.per_device_train_batch_size
+        ),
+        gradient_accumulation_steps=(
+            config.gradient_accumulation_steps
+        ),
+        learning_rate=config.learning_rate,
+        num_train_epochs=config.num_train_epochs,
+        weight_decay=0.1,
+        beta=config.beta,
+        logging_steps=2,
+        save_steps=10,
+        remove_unused_columns=False,
+        max_length=config.max_length,
+        max_prompt_length=config.max_prompt_length,
+        dataloader_num_workers=0,
+        fp16=False,
+        bf16=False,
+        optim="adamw_torch",
+        warmup_steps=2,
+        save_strategy="steps",
+        save_total_limit=3
+    )
+    
+    trainer = DPOTrainer(
+        model,
+        args=training_args,
+        train_dataset=preference_dataset,
+        peft_config=peft_config
+    )
+    
+    print("Starting DPO training...")
+    trainer.train()
+    
+    trainer.save_model(config.adapter_path)
+    print(f"Adapter saved to {config.adapter_path}")
+    
+    return config.adapter_path
+
+
+def run_rl_training(
+    tasks: List[Dict[str, Any]],
+    agents: List[NPC],
+    reward_fn: Callable[[Dict], float],
+    config: Optional[RLConfig] = None,
+    save_traces: bool = True
+) -> str:
+
+    if config is None:
+        config = RLConfig()
+    
+    print(f"Collecting traces from {len(tasks)} tasks...")
+    traces = collect_traces(
+        tasks,
+        agents,
+        reward_fn,
+        config
+    )
+    
+    if save_traces:
+        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
+        traces_file = f"rl_traces_{timestamp}.csv"
+        df = pd.DataFrame(traces)
+        df.to_csv(traces_file, index=False)
+        print(f"Traces saved to {traces_file}")
+    
+    print("Training with DPO...")
+    adapter_path = train_with_dpo(traces, config)
+    
+    return adapter_path
+
+
+def load_rl_model(
+    base_model_id: str,
+    adapter_path: str
+):
+
+    print(f"Loading base model: {base_model_id}")
+    model = AutoModelForCausalLM.from_pretrained(
+        base_model_id,
+        torch_dtype=torch.float32,
+        device_map="auto",
+        attn_implementation='eager'
+    )
+    
+    tokenizer = AutoTokenizer.from_pretrained(
+        base_model_id,
+        trust_remote_code=True
+    )
+    
+    if tokenizer.pad_token is None:
+        tokenizer.pad_token = tokenizer.eos_token
+    
+    if adapter_path and os.path.exists(adapter_path):
+        print(f"Loading adapter: {adapter_path}")
+        model = PeftModel.from_pretrained(model, adapter_path)
+        model = model.merge_and_unload()
+    
+    return model, tokenizer