from datetime import datetime import json import logging import random import os import threading import queue import fcntl from typing import List, Dict import httpx import argparse import re from transformers import AutoTokenizer # Configure logging # logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') # Configuration API_BASE = 'http://localhost:6002/v1' MODEL_NAME = 'gpqa' API_KEY = 'asdf' TOKENIZER_MODEL = 'google/gemma-2-9b-it' # You can change this to match your API's tokenizer # Initialize tokenizer tokenizer = AutoTokenizer.from_pretrained(TOKENIZER_MODEL) # Words to apply logit bias to (expanded list with variations) LOGIT_BIAS_WORDS = [] # Optional: Array of categories to process CATEGORIES_TO_PROCESS = [ "Science", "Technology", "History", "Literature" ] def tokenize_and_create_logit_bias(words: List[str], bias_value: float = -100) -> Dict[int, float]: logit_bias = {} for word in words: token_ids = tokenizer.encode(word, add_special_tokens=False) for token_id in token_ids: logit_bias[token_id] = bias_value return logit_bias # Create logit_bias dictionary with a stronger negative bias LOGIT_BIAS = tokenize_and_create_logit_bias(LOGIT_BIAS_WORDS, bias_value=-1000) def make_openai_request(messages: List[Dict[str, str]]) -> Dict: url = f"{API_BASE}/chat/completions" headers = { "Content-Type": "application/json", "Authorization": f"Bearer {API_KEY}" } data = { "model": MODEL_NAME, "messages": messages, "max_tokens": 1500, "temperature": 0.7, "logit_bias": LOGIT_BIAS } with httpx.Client(timeout=60.0) as client: response = client.post(url, json=data, headers=headers) response.raise_for_status() return response.json() def process_category(category: str, depth: int) -> Dict: depth_adjusted_prompt = f"Generate a Graduate-Level Google-Proof Multiple Choice Question for the specific category: {category}. " if '->' in category: depth_adjusted_prompt += f"Focus on the last part of the category chain: '{category.split('->')[-1].strip()}'. " depth_adjusted_prompt += f"Include very specific subcategories in your response. Write an extremely long and detailed question that is elaborate and context-rich. Build up the background in the question and make it very complex. Your question should be very difficult to answer and delve deep into specific aspects or applications of the category." messages = [ {"role": "user", "content": depth_adjusted_prompt}, ] response = make_openai_request(messages) document = response['choices'][0]['message']['content'] messages.append({"role": "assistant", "content": document}) messages.append({"role": "user", "content": "Convert the above to JSON format with fields: question, answer, incorrect_answer_1, incorrect_answer_2, incorrect_answer_3, explanation, and subcategories (as an array). Ensure the subcategories are very specific and related to the last part of the category chain."}) response = make_openai_request(messages) json_data = response['choices'][0]['message']['content'] try: parsed_json = json.loads(json_data) parsed_json['category'] = category parsed_json['document'] = document parsed_json['depth'] = depth return parsed_json except json.JSONDecodeError: logging.error(f"Failed to parse JSON for category: {category}") return None def worker(task_queue: queue.Queue, output_file: str): while True: try: category, depth = task_queue.get(block=False) except queue.Empty: break try: processed_entry = process_category(category, depth) if processed_entry: with open(output_file, 'a') as outfile: fcntl.flock(outfile, fcntl.LOCK_EX) json.dump(processed_entry, outfile) outfile.write('\n') fcntl.flock(outfile, fcntl.LOCK_UN) logging.info(f"Processed category: {category} at depth {depth}") except Exception as e: logging.error(f"Error processing category {category}: {str(e)}") finally: task_queue.task_done() def process_categories(categories: List[str], output_file: str, depth: int, num_threads: int = 60): os.makedirs(os.path.dirname(output_file), exist_ok=True) task_queue = queue.Queue() for category in categories: task_queue.put((category, depth)) threads = [] for _ in range(num_threads): t = threading.Thread(target=worker, args=(task_queue, output_file)) t.daemon = True t.start() threads.append(t) try: task_queue.join() except KeyboardInterrupt: logging.info("Keyboard interrupt received. Shutting down...") finally: for _ in range(num_threads): task_queue.put(None) for t in threads: t.join() def safe_filename(filename: str) -> str: filename = filename.replace(' ', '_') filename = re.sub(r'[^\w\-_.]', '', filename) return filename.lower() if __name__ == "__main__": parser = argparse.ArgumentParser(description="Generate GPQA for given categories with deep dive") parser.add_argument("--category", help="Input category") parser.add_argument("--use-array", action="store_true", help="Use the predefined category array") parser.add_argument("--depth", type=int, default=4, help="Maximum depth of category exploration") args = parser.parse_args() max_depth = args.depth if args.use_array: if not CATEGORIES_TO_PROCESS: logging.error("No categories defined in the CATEGORIES_TO_PROCESS array.") exit(1) categories = CATEGORIES_TO_PROCESS safe_category = safe_filename(categories[0]) elif args.category: categories = [args.category] safe_category = safe_filename(args.category) else: logging.error("Please provide either --category or --use-array") exit(1) timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") output_file = os.path.join("./output/gpqa", f"{safe_category}_deep_dive_{timestamp}.jsonl") os.makedirs(os.path.dirname(output_file), exist_ok=True) all_categories = set(categories) processed_categories = set() for depth in range(max_depth): logging.info(f"Processing depth level: {depth + 1}") categories_to_process = list(all_categories - processed_categories) random.shuffle(categories_to_process) process_categories(categories_to_process, output_file, depth, num_threads=36) processed_categories.update(categories_to_process) if depth < max_depth - 1: new_subcategories = set() try: with open(output_file, 'r') as f: for line in f: entry = json.loads(line) if 'subcategories' in entry and entry['depth'] == depth: parent_category = entry['category'] for subcat in entry['subcategories']: new_subcategories.add(f"{parent_category} -> {subcat}") except FileNotFoundError: logging.error(f"Output file not found: {output_file}") continue except json.JSONDecodeError: logging.error(f"Error decoding JSON in file: {output_file}") continue all_categories.update(new_subcategories) logging.info(f"Total unique categories after depth {depth + 1}: {len(all_categories)}") logging.info(f"Total processed categories: {len(processed_categories)}") logging.info(f"Processing complete. Results saved to {output_file}") logging.info(f"Total unique categories generated: {len(all_categories)}") logging.info(f"Total categories processed: {len(processed_categories)}")