run_generation.py

"""
Batch Generation Script - Symbolic Optimization for Fake News Generation

Independently runs symbolic optimization process to generate optimized fake news.
Saves after processing each item to improve fault tolerance.

Usage:
    python run_generation.py --input_file data/train.json --output_file results/generated.json
"""

import json
import os
import argparse
from typing import List, Dict
from agents import Generator
from optimization import SymbolicOptimizer
from config import DEFAULT_LLM_MODEL


def load_data(input_file: str, resume_index: int = 0) -> List[Dict]:
    """Load input data"""
    with open(input_file, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    if resume_index > 0:
        print(f"Resuming from index {resume_index}")
        data = data[resume_index:]
    
    return data


def load_existing_results(output_file: str) -> List[Dict]:
    """Load existing results"""
    if os.path.exists(output_file) and os.path.getsize(output_file) > 0:
        try:
            with open(output_file, 'r', encoding='utf-8') as f:
                return json.load(f)
        except json.JSONDecodeError:
            print(f"Warning: Cannot parse {output_file}, will create new file")
            return []
    return []


def save_result(output_file: str, results: List[Dict]):
    """Save results to file"""
    os.makedirs(os.path.dirname(output_file), exist_ok=True)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(results, f, ensure_ascii=False, indent=2)


def run_generation(
    input_file: str,
    output_file: str,
    llm_model: str,
    generator_model: str,
    max_iterations: int,
    resume_index: int = 0,
    max_samples: int = None
):
    """
    Run batch generation workflow
    
    Args:
        input_file: Input file path
        output_file: Output file path
        llm_model: Model for optimizer
        generator_model: Model for generator
        max_iterations: Maximum optimization iterations per sample
        resume_index: Resume index
        max_samples: Maximum number of samples to process
    """
    print("="*60)
    print("SALF Batch Generation Script")
    print("="*60)
    
    # Load data
    print(f"\nLoading data: {input_file}")
    input_data = load_data(input_file, resume_index)
    
    if max_samples:
        input_data = input_data[:max_samples]
        print(f"Limiting samples to: {max_samples}")
    
    print(f"Samples to process: {len(input_data)}")
    
    # Load existing results
    results = load_existing_results(output_file)
    print(f"Existing results: {len(results)}")
    
    # Process each item
    for idx, item in enumerate(input_data):
        actual_idx = idx + resume_index
        
        print(f"\n{'='*60}")
        print(f"Processing sample {actual_idx + 1}/{resume_index + len(input_data)}")
        print(f"{'='*60}")
        
        news_content = item.get('content', '')
        label = item.get('label', 1)
        
        # If real news, skip
        if label == 0:
            print(f"[Skip] Sample {actual_idx} is real news")
            result = {
                'content': news_content,
                'label': label,
                'optimized_content': news_content,
                'optimization_iterations': 0
            }
            result.update({k: v for k, v in item.items() if k not in result})
            results.append(result)
            save_result(output_file, results)
            continue
        
        try:
            # Initialize components
            generator = Generator(llm_model_name=generator_model)
            optimizer = SymbolicOptimizer(
                llm_model_name=llm_model,
                generator_model_name=generator_model
            )
            
            # Run multiple optimization iterations
            current_news = news_content
            optimization_logs = []
            
            for t in range(max_iterations):
                print(f"\n[Iteration {t+1}/{max_iterations}]")
                
                # Create a simple debate record (for symbolic optimization)
                # Note: Here we assume the news is fake and needs optimization
                simple_debate_record = {
                    'topic': current_news,
                    'positive_opening': 'This news appears to be real based on its structure.',
                    'negative_opening': 'This news has characteristics of fake news.',
                    'positive_questioning_one': 'Can you provide evidence?',
                    'negative_answering_one': 'The writing style and lack of sources suggest it is fabricated.',
                    'negative_questioning_two': 'What specific elements make it suspicious?',
                    'positive_answering_two': 'Some details seem plausible.',
                    'positive_closing': 'The news has some credible elements.',
                    'negative_closing': 'Overall, this appears to be fake news.',
                    'judgement': 'The news is likely fake due to lack of verifiable sources.'
                }
                
                # Execute symbolic optimization
                current_prompt = generator.get_prompt()
                optimized_prompt, optimized_news, opt_log = optimizer.optimize(
                    current_news,
                    current_prompt,
                    simple_debate_record
                )
                
                # Update generator and current news
                generator.update_prompt(optimized_prompt)
                current_news = optimized_news
                
                optimization_logs.append({
                    'iteration': t + 1,
                    'loss_score': opt_log['loss_score'],
                    'optimized_content': optimized_news
                })
                
                print(f"[Iteration {t+1}] Loss Score: {opt_log['loss_score']}")
            
            # Save result
            result = {
                'content': news_content,
                'label': label,
                'optimized_content': current_news,
                'optimization_iterations': max_iterations,
                'optimization_logs': optimization_logs
            }
            result.update({k: v for k, v in item.items() if k not in result})
            
            results.append(result)
            
            # Save after each item
            save_result(output_file, results)
            print(f"\n[Saved] Results saved to {output_file}")
            
        except Exception as e:
            print(f"\n[Error] Error processing sample {actual_idx}: {str(e)}")
            import traceback
            traceback.print_exc()
            
            # Save progress even on error
            result = {
                'content': news_content,
                'label': label,
                'optimized_content': news_content,
                'optimization_iterations': 0,
                'error': str(e)
            }
            result.update({k: v for k, v in item.items() if k not in result})
            results.append(result)
            save_result(output_file, results)
            
            print(f"[Saved] Error record saved, continuing to next item")
            continue
    
    print(f"\n{'='*60}")
    print(f"Batch generation complete!")
    print(f"Total processed: {len(results)} items")
    print(f"Results saved to: {output_file}")
    print(f"{'='*60}")


def main():
    parser = argparse.ArgumentParser(description='SALF Batch Generation Script')
    
    parser.add_argument('--input_file', type=str, required=True,
                       help='Input data file (JSON format)')
    parser.add_argument('--output_file', type=str, required=True,
                       help='Output result file')
    parser.add_argument('--llm_model', type=str, default=DEFAULT_LLM_MODEL,
                       help='LLM model for optimizer')
    parser.add_argument('--generator_model', type=str, default=DEFAULT_LLM_MODEL,
                       help='LLM model for generator')
    parser.add_argument('--max_iterations', type=int, default=3,
                       help='Maximum optimization iterations per sample')
    parser.add_argument('--resume_index', type=int, default=0,
                       help='Resume from specified index')
    parser.add_argument('--max_samples', type=int, default=None,
                       help='Maximum number of samples to process (for testing)')
    
    args = parser.parse_args()
    
    run_generation(
        input_file=args.input_file,
        output_file=args.output_file,
        llm_model=args.llm_model,
        generator_model=args.generator_model,
        max_iterations=args.max_iterations,
        resume_index=args.resume_index,
        max_samples=args.max_samples
    )


if __name__ == "__main__":
    main()