generate.py

import json
import os
import asyncio
import aiofiles
import time
import datetime
from pathlib import Path
from datasets import load_dataset
from openai import AsyncAzureOpenAI
import argparse

client = AsyncAzureOpenAI(
)

# concurrent control parameter
MAX_CONCURRENT = 50

# runtime control
MAX_RUNTIME_MINUTES = 600
program_start_time = time.time()
stop_event = asyncio.Event()

# statistics variable
processed_count = 0
error_count = 0

async def check_runtime():
    while True:
        current_runtime_minutes = (time.time() - program_start_time) / 60
        if current_runtime_minutes >= MAX_RUNTIME_MINUTES:
            print(f"Reach the maximum runtime {MAX_RUNTIME_MINUTES} minutes, stop all new requests")
            stop_event.set()
            break
        await asyncio.sleep(10)

async def generate_response(messages, model_name, semaphore):
    global processed_count, error_count
    
    if stop_event.is_set():
        return "ERROR"
    
    async with semaphore:
        if stop_event.is_set():
            return "ERROR"
        
        try:
            response = await client.chat.completions.create(
                model=model_name,
                messages=messages,
                max_tokens=2048,
                temperature=1.0
            )
            
            processed_count += 1
            return response.choices[0].message.content
                
        except Exception as e:
            error_count += 1
            error_type = type(e).__name__
            
            if "Content Exists Risk" in str(e):
                print("Content Exists Risk. Skipping.")
            else:
                print(f"Error ({error_type}): {e}")
            
            return "ERROR"

def create_extra_info_key(extra_info):
    if not extra_info:
        return None
    
    key_fields = []
    
    if 'problem_id' in extra_info:
        key_fields.append(extra_info['problem_id'])
    
    if 'submission_id' in extra_info:
        key_fields.append(str(extra_info['submission_id']))
    
    if 'task_type' in extra_info:
        key_fields.append(extra_info['task_type'])
    
    return '|'.join(key_fields) if key_fields else None


def load_existing_records(jsonl_file):
    existing_records = {}
    
    if os.path.exists(jsonl_file):
        try:
            with open(jsonl_file, 'r', encoding='utf-8') as f:
                for line_num, line in enumerate(f, 1):
                    if line.strip():
                        try:
                            record = json.loads(line.strip())
                            if 'extra_info' in record and record['extra_info']:
                                key = create_extra_info_key(record['extra_info'])
                                if key:
                                    responses_count = len(record.get('responses', []))
                                    existing_records[key] = responses_count
                        except json.JSONDecodeError as e:
                            print(f"Line {line_num} JSON parsing error: {e}")
                            continue
            print(f"Read {len(existing_records)} existing records")
            for key, count in existing_records.items():
                print(f"  {key}: {count} responses")
        except Exception as e:
            print(f"Error reading existing jsonl file: {e}")
    
    return existing_records


async def process_dataset(dataset_name, task_id, model_name, num_generations):
    global processed_count, error_count
    processed_count = 0
    error_count = 0
    
    print(f"Start processing dataset: {dataset_name}, task: task{task_id}, model: {model_name}")
    
    try:
        dataset = load_dataset(dataset_name, split="train")
        print(f"Dataset loaded successfully, {len(dataset)} records")
    except Exception as e:
        print(f"Failed to load dataset: {e}")
        return
    
    output_dir = f"results/{model_name}"
    os.makedirs(output_dir, exist_ok=True)
    output_file = os.path.join(output_dir, f"task{task_id}_gen.jsonl")
    
    existing_records = load_existing_records(output_file)
    
    semaphore = asyncio.Semaphore(MAX_CONCURRENT)
    
    runtime_check_task = asyncio.create_task(check_runtime())
    
    successful_records = 0
    
    all_tasks = []
    record_task_mapping = []
    
    for idx, record in enumerate(dataset):
        if stop_event.is_set():
            break
            
        extra_info = record.get('extra_info')
        existing_count = 0
        existing_responses = []
        
        if extra_info:
            record_key = create_extra_info_key(extra_info)
            if record_key and record_key in existing_records:
                existing_count = existing_records[record_key]
                if existing_count >= num_generations:
                    continue
                if os.path.exists(output_file):
                    with open(output_file, 'r', encoding='utf-8') as f:
                        for line in f:
                            if line.strip():
                                try:
                                    existing_record = json.loads(line.strip())
                                    if (existing_record.get('extra_info') and 
                                        create_extra_info_key(existing_record['extra_info']) == record_key):
                                        existing_responses = existing_record.get('responses', [])
                                        break
                                except json.JSONDecodeError:
                                    continue
        
        remaining_generations = num_generations - existing_count
        if remaining_generations <= 0:
            continue
            
        print(f"Record {idx}: {existing_count} responses, {remaining_generations} more to generate")
        
        prompt = record['prompt']
        
        record_tasks = []
        for gen_idx in range(remaining_generations):
            task = asyncio.create_task(
                generate_response(prompt, model_name, semaphore)
            )
            all_tasks.append(task)
            record_tasks.append(len(all_tasks) - 1)
        
        record_task_mapping.append({
            'record': record,
            'record_idx': idx,
            'task_indices': record_tasks,
            'existing_responses': existing_responses,
            'existing_count': existing_count
        })
    
    if len(all_tasks) == 0:
        print("No records to process")
        return

    print(f"Prepare to execute {len(all_tasks)} API calls")
    
    all_responses = await asyncio.gather(*all_tasks)
    
    all_existing_records = []
    processed_keys = set()
    
    if os.path.exists(output_file):
        with open(output_file, 'r', encoding='utf-8') as f:
            for line in f:
                if line.strip():
                    try:
                        existing_record = json.loads(line.strip())
                        all_existing_records.append(existing_record)
                    except json.JSONDecodeError:
                        continue
    
    for mapping in record_task_mapping:
        if stop_event.is_set():
            break
            
        record = mapping['record']
        idx = mapping['record_idx']
        task_indices = mapping['task_indices']
        existing_responses = mapping['existing_responses']
        existing_count = mapping['existing_count']
        
        try:
            new_responses = [all_responses[i] for i in task_indices]
            
            valid_new_responses = [resp for resp in new_responses if resp != "ERROR"]
            
            if not valid_new_responses and existing_count == 0:
                print(f"Record {idx} all generations failed, skip")
                continue
            
            all_responses_for_record = existing_responses + valid_new_responses
            
            output_record = dict(record)
            output_record['responses'] = all_responses_for_record
            
            extra_info = record.get('extra_info')
            if extra_info:
                record_key = create_extra_info_key(extra_info)
                if record_key:
                    processed_keys.add(record_key)
            
            if existing_count > 0:
                all_existing_records = [
                    r for r in all_existing_records 
                    if not (r.get('extra_info') and 
                           create_extra_info_key(r['extra_info']) == record_key)
                ]
            
            all_existing_records.append(output_record)
            
            successful_records += 1
            
            print(f"Record {idx}: {len(all_responses_for_record)} responses")
            
            if successful_records % 10 == 0:
                current_time = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
                elapsed = time.time() - program_start_time
                print(f"[{model_name}] [{current_time}] [elapsed: {elapsed:.2f}s] {successful_records} records processed")
                
        except Exception as e:
            print(f"Error processing record {idx}: {e}")
            continue
    
    async with aiofiles.open(output_file, 'w', encoding='utf-8') as f:
        for record in all_existing_records:
            await f.write(json.dumps(record, ensure_ascii=False) + '\n')

    
    runtime_check_task.cancel()
    try:
        await runtime_check_task
    except asyncio.CancelledError:
        pass
    
    total_runtime = (time.time() - program_start_time) / 60
    print(f"Task completed, total runtime: {total_runtime:.2f} minutes")
    print(f"End time: {datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"{successful_records} records processed, saved to {output_file}")
    print(f"Success: {processed_count}, Error: {error_count}")

async def main():
    parser = argparse.ArgumentParser(description="Generate test case responses")
    parser.add_argument("--model", type=str, required=True, help="Model name")
    parser.add_argument("--task", type=int, choices=[1, 2], help="Task ID (1 or 2)")
    parser.add_argument("--all", action="store_true", help="Process all tasks")
    
    args = parser.parse_args()
    
    datasets_config = {
        1: {
            "name": "Raywithyou/TestCase-Eval-Task1", # or Raywithyou/TestCase-Eval-Task1-DO (DO: direct-output prompt)
            "num_generations": 20
        },
        2: {
            "name": "Raywithyou/TestCase-Eval-Task2", # or Raywithyou/TestCase-Eval-Task2-DO (DO: direct-output prompt)
            "num_generations": 1
        }
    }
    
    if args.all:
        for task_id, config in datasets_config.items():
            print(f"\n{'='*60}")
            print(f"Start processing task {task_id}")
            print(f"{'='*60}")
            
            await process_dataset(
                config["name"], 
                task_id, 
                args.model, 
                config["num_generations"]
            )
    else:
        if args.task is None:
            print("Please specify --task or use --all parameter")
            return
            
        config = datasets_config[args.task]
        await process_dataset(
            config["name"], 
            args.task, 
            args.model, 
            config["num_generations"]
        )

if __name__ == "__main__":
    asyncio.run(main())