main.py

"""
SALF (Symbolic Adversarial Learning Framework) Main Program

Implements the complete training loop of Algorithm 1 from the paper.

Corresponds to paper: Algorithm 1 - SALF Training Loop
Source: Refactored from original experimental code
"""

import json
import os
import argparse
import time
from typing import List, Dict
from agents import Generator, DebateDetector
from optimization import SymbolicOptimizer
from config import (
    DEFAULT_LLM_MODEL,
    DEFAULT_MAX_ITERATIONS,
    DEFAULT_MAX_RETRIES,
    DATA_DIR,
    RESULTS_DIR
)


def parse_list(string: str) -> List[str]:
    """Parse comma-separated string into list"""
    return [s.strip() for s in string.split(',')]


def run_salf_iteration(
    news_content: str,
    generator: Generator,
    detector_model: str,
    optimizer: SymbolicOptimizer,
    iteration: int
) -> Dict:
    """
    Run one SALF iteration
    
    Corresponds to paper: One loop of Algorithm 1
    
    Args:
        news_content: Current news content
        generator: Generator instance
        detector_model: Model name for detector
        optimizer: Symbolic optimizer instance
        iteration: Current iteration number
    
    Returns:
        Dictionary containing iteration results
    """
    print(f"\n{'='*60}")
    print(f"SALF Iteration {iteration}")
    print(f"{'='*60}")
    
    # 1. Generator generates fake news (or uses current content)
    # In first iteration, use input news_content
    # In subsequent iterations, use optimized content
    current_news = news_content
    
    # 2. Detector performs debate and makes judgement
    print(f"\n[Step 1] Running debate detection...")
    detector = DebateDetector(news_topic=current_news, llm_model_names=[detector_model])
    debate_record, judgement = detector.run()
    
    print(f"[Step 1] Judgement: {judgement}")
    
    # 3. Update Detector (if judgement fails - i.e., judged as real news)
    # Note: In current implementation, Detector prompts are fixed, Detector update not implemented here
    # If needed, can add prompt update mechanism similar to Generator
    
    # 4. Update Generator (symbolic optimization)
    print(f"\n[Step 2] Running symbolic optimization...")
    current_prompt = generator.get_prompt()
    optimized_prompt, optimized_news, optimization_log = optimizer.optimize(
        current_news,
        current_prompt,
        debate_record
    )
    
    # Update generator prompt
    generator.update_prompt(optimized_prompt)
    
    print(f"[Step 2] Optimization complete")
    print(f"[Step 2] Loss score: {optimization_log['loss_score']}")
    
    # Return iteration results
    return {
        'iteration': iteration,
        'original_news': news_content,
        'optimized_news': optimized_news,
        'debate_record': debate_record,
        'judgement': judgement,
        'optimization_log': optimization_log
    }


def run_salf_training(
    input_data: List[Dict],
    llm_model: str,
    generator_model: str,
    max_iterations: int,
    output_file: str,
    resume_index: int = 0
):
    """
    Run complete SALF training workflow
    
    Corresponds to paper: Algorithm 1
    
    Args:
        input_data: Input data list
        llm_model: LLM model name
        generator_model: Generator model name
        max_iterations: Maximum iterations
        output_file: Output file path
        resume_index: Resume index
    """
    # Ensure output directory exists
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    
    # Initialize results list
    results = []
    
    # If resuming, load existing results
    if resume_index > 0 and os.path.exists(output_file):
        with open(output_file, 'r', encoding='utf-8') as f:
            results = json.load(f)
    
    # Process each data item
    for idx, item in enumerate(input_data):
        actual_idx = idx + resume_index
        
        print(f"\n\n{'#'*60}")
        print(f"Processing item {actual_idx}")
        print(f"{'#'*60}")
        
        # Get news content and label
        news_content = item.get('content', '')
        label = item.get('label', 1)  # 0=real news, 1=fake news
        
        # If real news, skip optimization
        if label == 0:
            print(f"[Info] Item {actual_idx} is real news, skipping optimization")
            result = {
                'content': news_content,
                'label': label,
                'optimized_content': news_content,
                'iterations': [],
                'final_judgement': 'positive'  # Real news
            }
            # Preserve other fields from original data
            result.update({k: v for k, v in item.items() if k not in result})
            results.append(result)
            continue
        
        # Initialize components
        generator = Generator(llm_model_name=generator_model)
        optimizer = SymbolicOptimizer(
            llm_model_name=llm_model,
            generator_model_name=generator_model
        )
        
        # Run multiple iterations
        iterations_log = []
        current_news = news_content
        
        for t in range(max_iterations):
            try:
                iteration_result = run_salf_iteration(
                    news_content=current_news,
                    generator=generator,
                    detector_model=llm_model,
                    optimizer=optimizer,
                    iteration=t + 1
                )
                
                iterations_log.append(iteration_result)
                
                # Update current news to optimized version
                current_news = iteration_result['optimized_news']
                
            except Exception as e:
                print(f"[Error] Iteration {t+1} failed: {str(e)}")
                break
        
        # Save results
        result = {
            'content': news_content,
            'label': label,
            'optimized_content': current_news,
            'iterations': iterations_log,
            'final_judgement': iterations_log[-1]['judgement'] if iterations_log else ''
        }
        
        # Preserve other fields from original data
        result.update({k: v for k, v in item.items() if k not in result})
        
        results.append(result)
        
        # Save results in real-time
        with open(output_file, 'w', encoding='utf-8') as f:
            json.dump(results, f, ensure_ascii=False, indent=2)
        
        print(f"\n[Info] Item {actual_idx} completed and saved")
    
    print(f"\n\n{'='*60}")
    print(f"SALF Training Complete!")
    print(f"Results saved to: {output_file}")
    print(f"{'='*60}")


def main():
    """Main function"""
    parser = argparse.ArgumentParser(description='SALF (Symbolic Adversarial Learning Framework) Training')
    
    # Data parameters
    parser.add_argument('--input_file', type=str, default=str(DATA_DIR / 'train.json'),
                       help='Input data file (JSON format)')
    parser.add_argument('--output_file', type=str, default=str(RESULTS_DIR / 'salf_output.json'),
                       help='Output file for results')
    parser.add_argument('--resume_index', type=int, default=0,
                       help='Resume from this index (for interrupted training)')
    
    # Model parameters
    parser.add_argument('--llm_model', type=str, default=DEFAULT_LLM_MODEL,
                       help='LLM model name for detector and optimizer')
    parser.add_argument('--generator_model', type=str, default=DEFAULT_LLM_MODEL,
                       help='LLM model name for generator')
    
    # Training parameters
    parser.add_argument('--max_iterations', type=int, default=DEFAULT_MAX_ITERATIONS,
                       help='Maximum number of SALF iterations per sample')
    parser.add_argument('--max_samples', type=int, default=None,
                       help='Maximum number of samples to process (for testing)')
    
    args = parser.parse_args()
    
    # Load input data
    print(f"Loading data from: {args.input_file}")
    try:
        with open(args.input_file, 'r', encoding='utf-8') as f:
            input_data = json.load(f)
        
        # If resume_index specified, start from that position
        if args.resume_index > 0:
            print(f"Resuming from index {args.resume_index}")
            input_data = input_data[args.resume_index:]
        
        # If max_samples specified, only process first N samples
        if args.max_samples:
            print(f"Processing only {args.max_samples} samples")
            input_data = input_data[:args.max_samples]
        
        print(f"Loaded {len(input_data)} samples")
        
    except Exception as e:
        print(f"Error loading input file: {str(e)}")
        return
    
    # Run SALF training
    run_salf_training(
        input_data=input_data,
        llm_model=args.llm_model,
        generator_model=args.generator_model,
        max_iterations=args.max_iterations,
        output_file=args.output_file,
        resume_index=args.resume_index
    )


if __name__ == "__main__":
    main()