generate.py

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)}")