run_multi_react.py

import argparse
import json
import os
import threading
from concurrent.futures import ThreadPoolExecutor, TimeoutError, as_completed
from datetime import datetime
from pathlib import Path

from tqdm import tqdm
from utils.prompt import SYSTEM_PROMPT_MULTI, QWEN_USER_PROMPT, FC_USER_PROMPT
from utils.react_agent import MultiTurnReactAgent
from utils.constant import FC_MODELS
# registor tools
from utils.tool_search import *
from utils.tool_visit import *


def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--model", type=str, default="")
    parser.add_argument("--output", type=str, default="")
    parser.add_argument("--dataset", type=str, default="wiki_100.jsonl")
    parser.add_argument("--temperature", type=float, default=0.6)
    parser.add_argument("--top_p", type=float, default=0.95)
    parser.add_argument("--max_workers", type=int, default=40)
    parser.add_argument("--sys_prompt", type=str, default=SYSTEM_PROMPT_MULTI)
    parser.add_argument("--usr_prompt", type=str, default=QWEN_USER_PROMPT)
    parser.add_argument("--roll_out_count", type=int, default=3)
    parser.add_argument("--timeout", type=int, default=600)
    args = parser.parse_args()

    # Parse model name (last part after / in the path)
    model_name = os.path.basename(args.model.rstrip('/'))

    if model_name in FC_MODELS:
        args.usr_prompt = FC_USER_PROMPT
    
    # Create output directory structure: args.output/model_name/dataset_name/
    dataset_dir = os.path.join(os.path.join(args.output, f"{model_name}"), Path(args.dataset).stem)
    
    # Create directory
    os.makedirs(dataset_dir, exist_ok=True)
    
    print(f"Model name: {model_name}")
    print(f"Dataset name: {args.dataset}")
    print(f"Output directory: {dataset_dir}")
    print(f"Rollout count: {args.roll_out_count}")

    data_filepath = f"eval_data/{args.dataset}"
    try:
        if data_filepath.endswith(".json"):
            with open(data_filepath, "r", encoding="utf-8") as f:
                items = json.load(f)
            if not isinstance(items, list):
                raise ValueError("Input JSON must be a list of objects.")
            if items and not isinstance(items[0], dict):
                raise ValueError("Input JSON list items must be objects.")
        elif data_filepath.endswith(".jsonl"):
            with open(data_filepath, "r", encoding="utf-8") as f:
                items = [json.loads(line) for line in f]
        else:
            raise ValueError("Unsupported file extension. Please use .json or .jsonl files.")
    except FileNotFoundError:
        print(f"Error: Input file not found at {data_filepath}")
        exit(1)
    except (json.JSONDecodeError, ValueError) as e:
        print(f"Error reading or parsing input file {data_filepath}: {e}")
        exit(1)

    # Create tasks for each rollout
    for rollout_idx in range(1, args.roll_out_count + 1):
        output_file = os.path.join(dataset_dir, f"iter{rollout_idx}.jsonl")
        
        print(f"\nStarting rollout {rollout_idx}/{args.roll_out_count}")
        print(f"Output file: {output_file}")
        
        # Check already processed queries
        processed_queries = set()
        if os.path.exists(output_file):
            try:
                with open(output_file, "r", encoding="utf-8") as f:
                    for line in f:
                        try:
                            data = json.loads(line)
                            # if "question" in data and "[Failed]" not in data['prediction']:
                            if "question" in data:
                                processed_queries.add(data["question"].strip())
                        except json.JSONDecodeError:
                            print(f"Warning: Skipping invalid line in output file: {line.strip()}")
            except Exception as e:
                print(f"Error processing line: {line.strip()}")
                print(repr(e))

        tasks_to_run = []
        for item in items:
            question = item.get("question", "").strip()
            if question == "":
                try:
                    user_msg = item["messages"][1]["content"] 
                    question = user_msg.split("User:")[1].strip() if "User:" in user_msg else user_msg
                    item["question"] = question
                except Exception as e:
                    print(f"Extract question from user message failed: {e}")
            if not question:
                print(f"Warning: Skipping item with empty question: {item}")
                continue

            if question not in processed_queries:
                tasks_to_run.append({"item": item.copy(), "rollout_id": rollout_idx})
            else:
                print(f"Skipping already processed question: {question}")

        print(f"Total questions in input: {len(items)}")
        print(f"Already successfully processed: {len(processed_queries)}")
        print(f"Total tasks to run for this rollout: {len(tasks_to_run)}")

        if not tasks_to_run:
            print(f"Rollout {rollout_idx} already completed, skipping")
            continue

        llm_cfg = {
            'model': args.model,
            'generate_cfg': {
                'max_input_tokens': 32000,
                'num_retries': 10, 
                'temperature': args.temperature, 
                'top_p': args.top_p
            },
        }

        system_message = args.sys_prompt + "\nCurrent date: " + datetime.now().strftime("%Y-%m-%d")

        test_agent = MultiTurnReactAgent(
            llm=llm_cfg,
            function_list=["search", "visit"],
            system_message=system_message
        )

        # 创建文件写入锁
        write_lock = threading.Lock()

        with ThreadPoolExecutor(max_workers=args.max_workers) as executor:
            # Submit tasks
            future_to_task = {
                executor.submit(
                    test_agent.run,
                    task,
                    args.usr_prompt
                ): task
                for task in tasks_to_run
            }

            for future in tqdm(as_completed(future_to_task), total=len(future_to_task), desc=f"Processing Rollout {rollout_idx}"):
                task_info = future_to_task[future]
                try:
                    result = future.result(timeout=args.timeout)
                    # 使用锁保护文件写入操作
                    with write_lock:
                        with open(output_file, "a", encoding="utf-8") as f:
                            f.write(json.dumps(result, ensure_ascii=False) + "\n")
                except TimeoutError:
                    print(f'Timeout (>10 mins): "{task_info["item"]["question"]}" '
                          f'(Rollout {task_info["rollout_id"]})')
                    future.cancel()
                    error_result = {
                        "rollout_id": task_info["rollout_id"],
                        **task_info["item"],
                        "messages": [],
                        "prediction": "[Failed]",
                        "termination": "Future resolution failed: Timeout (>10 mins)",
                    }
                    print("===============================")
                    print(error_result)
                    print("===============================")

                    # Also use lock to protect error writing
                    with write_lock:
                        with open(output_file, "a", encoding="utf-8") as f:
                            f.write(json.dumps(error_result, ensure_ascii=False) + "\n")

                except Exception as exc:
                    print(f'Task for question "{task_info["item"]["question"]}" (Rollout {task_info["rollout_id"]}) generated an exception: {exc}')
                    # Log error to the output file
                    error_result = {
                        "rollout_id": task_info["rollout_id"],
                        **task_info["item"],
                        "messages": [],
                        "prediction": "[Failed]",
                        "termination": f"Future resolution failed: {exc}",                       
                    }
                    print("===============================")
                    print(error_result)
                    print("===============================")
                    
                    # Also use lock to protect error writing
                    with write_lock:
                        with open(output_file, "a", encoding="utf-8") as f:
                            f.write(json.dumps(error_result, ensure_ascii=False) + "\n")
        
        print(f"Rollout {rollout_idx} completed")
    
    print(f"\nAll {args.roll_out_count} rollouts completed!")


if __name__ == "__main__":
    main()