# run_rag_agent.py import os import json import time import re import requests from tqdm import tqdm import numpy as np import torch from typing import Optional, Tuple, List, Dict import argparse from transformers import AutoTokenizer from vllm import LLM, SamplingParams from bing_search import bing_web_search, extract_relevant_info, fetch_page_content from evaluate import run_evaluation from prompts import ( get_singleqa_rag_agent_instruction, get_multiqa_rag_agent_instruction, get_gpqa_rag_agent_instruction, get_math_rag_agent_instruction, get_code_rag_agent_instruction, get_task_instruction_openqa, get_task_instruction_math, get_task_instruction_multi_choice, get_task_instruction_code, ) # Define special symbols BEGIN_SEARCH_QUERY = "<|begin_search_query|>" END_SEARCH_QUERY = "<|end_search_query|>" BEGIN_SEARCH_RESULT = "<|begin_search_result|>" END_SEARCH_RESULT = "<|end_search_result|>" BEGIN_URL = "<|begin_url|>" END_URL = "<|end_url|>" BEGIN_FULL_PAGE = "<|begin_full_page|>" END_FULL_PAGE = "<|end_full_page|>" def parse_args(): parser = argparse.ArgumentParser(description="Run RAG Agent for various datasets and models.") # Dataset and split configuration parser.add_argument( '--dataset_name', type=str, required=True, choices=['gpqa', 'math500', 'aime', 'amc', 'livecode', 'nq', 'triviaqa', 'hotpotqa', '2wiki', 'musique', 'bamboogle', 'medmcqa', 'pubhealth'], help="Name of the dataset to use." ) parser.add_argument( '--split', type=str, required=True, choices=['test', 'diamond', 'main', 'extended'], help="Dataset split to use." ) parser.add_argument( '--subset_num', type=int, default=-1, help="Number of examples to process. Defaults to all if not specified." ) # RAG Agent configuration parser.add_argument( '--max_search_limit', type=int, default=5, help="Maximum number of searches per question." ) parser.add_argument( '--max_url_fetch', type=int, default=5, help="Maximum number of URL fetches per question." ) parser.add_argument( '--max_turn', type=int, default=10, help="Maximum number of turns." ) parser.add_argument( '--top_k', type=int, default=10, help="Maximum number of search documents to return." ) # Model configuration parser.add_argument( '--model_path', type=str, required=True, help="Path to the pre-trained model." ) parser.add_argument( '--use_jina', type=bool, default=True, help="Whether to use Jina API for document fetching." ) parser.add_argument( '--jina_api_key', type=str, default='None', help="Your Jina API Key to Fetch URL Content." ) # Sampling parameters parser.add_argument( '--temperature', type=float, default=0.7, help="Sampling temperature." ) parser.add_argument( '--top_p', type=float, default=0.8, help="Top-p sampling parameter." ) parser.add_argument( '--top_k_sampling', type=int, default=20, help="Top-k sampling parameter." ) parser.add_argument( '--repetition_penalty', type=float, default=None, help="Repetition penalty. If not set, defaults based on the model." ) parser.add_argument( '--max_tokens', type=int, default=32768, help="Maximum number of tokens to generate. If not set, defaults based on the model and dataset." ) # Bing API Configuration parser.add_argument( '--bing_subscription_key', type=str, required=True, help="Bing Search API subscription key." ) parser.add_argument( '--bing_endpoint', type=str, default="https://api.bing.microsoft.com/v7.0/search", help="Bing Search API endpoint." ) return parser.parse_args() def main(): args = parse_args() # Extract arguments dataset_name = args.dataset_name split = args.split subset_num = args.subset_num MAX_SEARCH_LIMIT = args.max_search_limit MAX_URL_FETCH = args.max_url_fetch MAX_TURN = args.max_turn top_k = args.top_k model_path = args.model_path temperature = args.temperature top_p = args.top_p top_k_sampling = args.top_k_sampling repetition_penalty = args.repetition_penalty max_tokens = args.max_tokens bing_subscription_key = args.bing_subscription_key bing_endpoint = args.bing_endpoint use_jina = args.use_jina jina_api_key = args.jina_api_key # Adjust parameters based on dataset if dataset_name in ['nq', 'triviaqa', 'hotpotqa', 'musique', 'bamboogle', '2wiki', 'medmcqa', 'pubhealth']: MAX_SEARCH_LIMIT = 5 MAX_URL_FETCH = 5 MAX_TURN = 10 top_k = 5 # Set default repetition_penalty if not provided if repetition_penalty is None: repetition_penalty = 1.05 if 'qwq' in model_path.lower() else 1.0 if args.jina_api_key == 'None': jina_api_key = None # Data paths based on dataset if dataset_name == 'livecode': data_path = f'./data/LiveCodeBench/{split}.json' elif dataset_name in ['math500', 'gpqa', 'aime', 'amc']: data_path = f'./data/{dataset_name.upper()}/{split}.json' else: data_path = f'./data/QA_Datasets/{dataset_name}.json' # ---------------------- Caching Mechanism ---------------------- # Define cache directories and file paths cache_dir = './cache' search_cache_path = os.path.join(cache_dir, 'search_cache.json') url_cache_path = os.path.join(cache_dir, 'url_cache.json') # Ensure cache directory exists os.makedirs(cache_dir, exist_ok=True) # Load existing caches or initialize empty dictionaries if os.path.exists(search_cache_path): with open(search_cache_path, 'r', encoding='utf-8') as f: search_cache = json.load(f) else: search_cache = {} if os.path.exists(url_cache_path): with open(url_cache_path, 'r', encoding='utf-8') as f: url_cache = json.load(f) else: url_cache = {} # Define function to save caches def save_caches(): with open(search_cache_path, 'w', encoding='utf-8') as f: json.dump(search_cache, f, ensure_ascii=False, indent=2) with open(url_cache_path, 'w', encoding='utf-8') as f: json.dump(url_cache, f, ensure_ascii=False, indent=2) # ---------------------- Model Loading ---------------------- tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token tokenizer.padding_side = 'left' # Define output directory based on model and dataset if 'qwq' in model_path.lower(): if dataset_name in ['math500', 'gpqa', 'aime', 'amc', 'livecode']: output_dir = f'./outputs/{dataset_name}.qwq.rag_agent' else: output_dir = f'./outputs/runs.qa/{dataset_name}.qwq.rag_agent' else: model_short_name = model_path.split('/')[-1].lower().replace('-instruct', '') output_dir = f'./outputs/runs.baselines/{dataset_name}.{model_short_name}.rag_agent' os.makedirs(output_dir, exist_ok=True) # Initialize the LLM llm = LLM( model=model_path, tensor_parallel_size=torch.cuda.device_count(), gpu_memory_utilization=0.95, ) # ---------------------- Data Loading ---------------------- with open(data_path, 'r', encoding='utf-8') as json_file: filtered_data = json.load(json_file) # ---------------------- Prepare Input ---------------------- input_list = [] for item in filtered_data: question = item['Question'] if dataset_name in ['nq', 'triviaqa', 'hotpotqa', 'musique', 'bamboogle', '2wiki']: if dataset_name in ['nq', 'triviaqa']: instruction = get_singleqa_rag_agent_instruction(MAX_SEARCH_LIMIT, MAX_URL_FETCH) elif dataset_name in ['hotpotqa', 'musique', 'bamboogle', '2wiki']: instruction = get_multiqa_rag_agent_instruction(MAX_SEARCH_LIMIT, MAX_URL_FETCH) if 'qwq' in model_path.lower(): user_prompt = get_task_instruction_openqa(question, model_name='qwq') else: user_prompt = get_task_instruction_openqa(question) elif dataset_name in ['math500', 'aime', 'amc']: instruction = get_math_rag_agent_instruction(MAX_SEARCH_LIMIT, MAX_URL_FETCH) if 'qwq' in model_path.lower(): user_prompt = get_task_instruction_math(question, model_name='qwq') else: user_prompt = get_task_instruction_math(question) elif dataset_name == 'gpqa': instruction = get_gpqa_rag_agent_instruction(MAX_SEARCH_LIMIT, MAX_URL_FETCH) if 'qwq' in model_path.lower(): user_prompt = get_task_instruction_multi_choice(question, model_name='qwq') elif 'llama' in model_path.lower(): user_prompt = get_task_instruction_multi_choice(question, model_name='llama') else: user_prompt = get_task_instruction_multi_choice(question) elif dataset_name == 'livecode': instruction = get_code_rag_agent_instruction(MAX_SEARCH_LIMIT, MAX_URL_FETCH) question_title = item.get('question_title', '') if 'qwq' in model_path.lower(): user_prompt = get_task_instruction_code(question, question_title=question_title, model_name='qwq') else: user_prompt = get_task_instruction_code(question) else: user_prompt = "" # Default to empty if dataset not matched prompt = [{"role": "user", "content": instruction + user_prompt}] prompt = tokenizer.apply_chat_template(prompt, tokenize=False, add_generation_prompt=True) input_list.append(prompt) if subset_num != -1: input_list = input_list[:subset_num] filtered_data = filtered_data[:subset_num] # Initialize active sequences with search and URL fetch counters active_sequences = [{ 'item': item, 'prompt': prompt, 'output': '', 'finished': False, 'history': [], 'pending_operations': [], # Queue of operations to execute 'executed_search_queries': set(), 'executed_url_fetches': set(), 'search_count': 0 # Search counter } for item, prompt in zip(filtered_data, input_list)] # ---------------------- Set Max Tokens ---------------------- if 'qwq' in model_path.lower(): if dataset_name in ['aime', 'amc', 'livecode']: max_tokens = 32768 else: max_tokens = 20480 else: max_tokens = 8192 # ---------------------- Generation Function ---------------------- def run_generation(sequences, max_tokens): """ Run LLM generation on provided sequences. """ prompts = [s['prompt'] for s in sequences] sampling_params = SamplingParams( max_tokens=max_tokens, temperature=temperature, top_p=top_p, top_k=top_k_sampling, repetition_penalty=repetition_penalty, stop=[END_SEARCH_QUERY, END_URL, tokenizer.eos_token], include_stop_str_in_output=True, ) output_list = llm.generate(prompts, sampling_params=sampling_params) return output_list # Function to extract text between two tags def extract_between(text: str, start_tag: str, end_tag: str) -> Optional[str]: pattern = re.escape(start_tag) + r"(.*?)" + re.escape(end_tag) matches = re.findall(pattern, text, flags=re.DOTALL) if matches: return matches[-1].strip() return None start_time = time.time() turn = 0 # Main loop until all sequences are completed while True: # Separate sequences with pending operations and those needing generation sequences_with_pending_ops = [seq for seq in active_sequences if not seq['finished'] and seq['pending_operations']] sequences_needing_generation = [seq for seq in active_sequences if not seq['finished'] and not seq['pending_operations']] # First, handle pending operations if sequences_with_pending_ops: print(f"{len(sequences_with_pending_ops)} sequences have pending operations. Executing...") for seq in sequences_with_pending_ops: # Execute the next pending operation operation = seq['pending_operations'].pop(0) # FIFO op_type = operation['type'] content = operation['content'] if op_type == 'search': query = content if query in search_cache: results = search_cache[query] print(f"Using cached search results for query: {query}") else: try: # Execute search and cache results results = bing_web_search(query, bing_subscription_key, bing_endpoint, market='en-US', language='en') search_cache[query] = results print(f"Executed and cached search for query: {query}") except Exception as e: print(f"Error during search query '{query}': {e}") search_cache[query] = {} results = {} relevant_info = extract_relevant_info(results)[:top_k] search_result_str = json.dumps(relevant_info, ensure_ascii=False, indent=2) # Append search results to the prompt append_text = f"\n{BEGIN_SEARCH_RESULT}\n{search_result_str}\n{END_SEARCH_RESULT}\n" seq['prompt'] += append_text seq['output'] += append_text # Update history seq['history'].append(append_text) # Increment search count seq['search_count'] += 1 elif op_type == 'fetch_url': urls = content # Calculate remaining URL fetches remaining_fetches = MAX_URL_FETCH - len(seq['executed_url_fetches']) if remaining_fetches <= 0: # Reached URL fetch limit, add limit message and mark sequence as finished limit_message = f"\n{BEGIN_FULL_PAGE}\nThe maximum number of URL fetches has been reached. You are not allowed to fetch more URLs.\n{END_FULL_PAGE}\n" seq['prompt'] += limit_message seq['output'] += limit_message seq['history'].append(limit_message) print("Reached URL fetch limit. Sequence marked as finished.") continue # Split and clean URLs urls_to_fetch = [u.strip() for u in urls.split(",")] # Filter already fetched URLs urls_to_fetch = [u for u in urls_to_fetch if u not in seq['executed_url_fetches']] # Limit the number of URLs to fetch urls_to_fetch = urls_to_fetch[:remaining_fetches] if not urls_to_fetch: print("All requested URLs have been fetched or exceeded the limit.") continue # Batch fetch page content, considering cache urls_to_fetch_filtered = [u for u in urls_to_fetch if u not in url_cache] cached_urls = [u for u in urls_to_fetch if u in url_cache] fetched_contents = [] # Use cached URL content for url in cached_urls: content = url_cache[url] print(f"Using cached URL content for URL: {url}") fetched_contents.append((url, content)) # Batch fetch uncached URLs if urls_to_fetch_filtered: try: # Batch pass uncached URLs contents = fetch_page_content(urls_to_fetch_filtered, use_jina=use_jina, jina_api_key=jina_api_key) for url, content in contents.items(): url_cache[url] = content print(f"Fetched and cached URL content for URL: {url}") fetched_contents.append((url, content)) except Exception as e: for url in urls_to_fetch_filtered: content = f"Error fetching URL: {e}" url_cache[url] = content fetched_contents.append((url, content)) print(f"Error fetching URL '{url}': {e}") # Update fetched URLs for url, _ in fetched_contents: seq['executed_url_fetches'].add(url) # Construct full page content string fetched_pages = dict(fetched_contents) full_page_str = json.dumps(fetched_pages, ensure_ascii=False, indent=2) # Append full page content to the prompt append_text = f"\n{BEGIN_FULL_PAGE}\n{full_page_str}\n{END_FULL_PAGE}\n" seq['prompt'] += append_text seq['output'] += append_text # Update history seq['history'].append(append_text) print(f"Fetched and cached {len(fetched_contents)} URLs.") # Check if URL fetch limit is reached if len(seq['executed_url_fetches']) >= MAX_URL_FETCH: limit_message = f"\n{BEGIN_FULL_PAGE}\nThe maximum number of URL fetches has been reached. You are not allowed to fetch more URLs.\n{END_FULL_PAGE}\n" seq['prompt'] += limit_message seq['output'] += limit_message seq['history'].append(limit_message) print("Reached URL fetch limit. Sequence marked as finished.") # Continue to the next iteration if there are pending operations if sequences_with_pending_ops: continue # Process operations first # Handle sequences needing generation if sequences_needing_generation: turn += 1 print(f"Turn {turn}: {len(sequences_needing_generation)} sequences need generation. Generating with LLM...") outputs = run_generation(sequences_needing_generation, max_tokens) print("Generation complete. Processing outputs...") # Process each generated output for seq, out in zip(sequences_needing_generation, outputs): text = out.outputs[0].text seq['history'].append(text) # Append generated text to prompt and output seq['prompt'] += text seq['output'] += text # Check if the generated content contains search queries or URL fetch requests search_query = extract_between(text, BEGIN_SEARCH_QUERY, END_SEARCH_QUERY) url_fetch = extract_between(text, BEGIN_URL, END_URL) if search_query: # Check if search limit is not exceeded if seq['search_count'] < MAX_SEARCH_LIMIT: # Check if this search query has not been executed if search_query not in seq['executed_search_queries']: # Add search operation to pending queue seq['pending_operations'].append({'type': 'search', 'content': search_query}) seq['executed_search_queries'].add(search_query) print(f"Added pending search operation for query: {search_query}") else: print(f"Search query already executed: {search_query}") else: # Add limit message if search limit is exceeded limit_message = f"\n{BEGIN_SEARCH_RESULT}\nThe maximum search limit is exceeded. You are not allowed to search.\n{END_SEARCH_RESULT}\n" seq['prompt'] += limit_message seq['output'] += limit_message seq['history'].append(limit_message) print(f"Search limit exceeded for query: {search_query}") if url_fetch: # Check if URL fetch limit is not exceeded if len(seq['executed_url_fetches']) < MAX_URL_FETCH: # Split and check if URLs have already been fetched urls = [u.strip() for u in url_fetch.split(",")] urls_to_fetch = [u for u in urls if u not in seq['executed_url_fetches']] if urls_to_fetch: # Add URL fetch operation to pending queue seq['pending_operations'].append({'type': 'fetch_url', 'content': ', '.join(urls_to_fetch)}) print(f"Added pending URL fetch operation for URLs: {urls_to_fetch}") else: print(f"All requested URLs have been fetched or exceeded the limit: {urls}") else: # Add limit message if URL fetch limit is exceeded limit_message = f"\n{BEGIN_FULL_PAGE}\nThe maximum number of URL fetches has been reached. You are not allowed to fetch more URLs.\n{END_FULL_PAGE}\n" seq['prompt'] += limit_message seq['output'] += limit_message seq['history'].append(limit_message) print("URL fetch limit exceeded.") # If no new operations are added, mark sequence as finished if not search_query and not url_fetch: seq['finished'] = True print("Sequence marked as finished.") # Check if all sequences are finished unfinished = [seq for seq in active_sequences if not seq['finished']] if not unfinished: break else: if turn >= MAX_TURN: print(f"Exceeded maximum number of turns ({MAX_TURN}). Stopping.") break # Optionally, implement a delay or other logic to prevent infinite loops pass total_time = time.time() - start_time # Collect all outputs output_list = [seq['output'] for seq in active_sequences] # ---------------------- Evaluation ---------------------- print("Evaluating generated answers...") run_evaluation( filtered_data=filtered_data, input_list=input_list, output_list=output_list, dataset_name=dataset_name, output_dir=output_dir, total_time=total_time, split=split, ) # ---------------------- Update Search and URL Cache ---------------------- print('Updating Search and URL Cache...') # Load existing caches or initialize empty dictionaries if os.path.exists(search_cache_path): with open(search_cache_path, 'r', encoding='utf-8') as f: search_cache_new = json.load(f) else: search_cache_new = {} if os.path.exists(url_cache_path): with open(url_cache_path, 'r', encoding='utf-8') as f: url_cache_new = json.load(f) else: url_cache_new = {} search_cache.update(search_cache_new) url_cache.update(url_cache_new) save_caches() print("Process completed.") if __name__ == "__main__": main()