Upload data3/generate_problems_batch.py with huggingface_hub

data3/generate_problems_batch.py

@@ -0,0 +1,655 @@
+#!/usr/bin/env python3
+"""
+Generate programming problems from function_dataset_v2.csv using OpenAI Batch API.
+Batch API offers 50% cost savings compared to standard API.
+"""
+
+import csv
+import json
+import os
+import sys
+from openai import OpenAI
+from datetime import datetime
+from typing import Dict, Optional, List
+import time
+
+# Configuration
+MODEL_NAME = "gpt-4o-mini"
+MIN_RELEVANCE_SCORE = 60
+MAX_BUDGET_USD = 10.0
+
+# OpenAI Batch API pricing (50% off standard pricing)
+# Official pricing: https://openai.com/api/pricing/
+BATCH_PRICING = {
+    # GPT-5 series with Batch API discount
+    "gpt-5.2": {
+        "input": 0.875 / 1_000_000,    # $0.875 per 1M (50% off $1.75)
+        "output": 7.00 / 1_000_000,    # $7.00 per 1M (50% off $14.00)
+    },
+    "gpt-5.1": {
+        "input": 0.625 / 1_000_000,    # $0.625 per 1M (50% off $1.25)
+        "output": 5.00 / 1_000_000,    # $5.00 per 1M (50% off $10.00)
+    },
+    "gpt-5": {
+        "input": 0.625 / 1_000_000,    # $0.625 per 1M (50% off $1.25)
+        "output": 5.00 / 1_000_000,    # $5.00 per 1M (50% off $10.00)
+    },
+    "gpt-5-mini": {
+        "input": 0.125 / 1_000_000,    # $0.125 per 1M (50% off $0.25)
+        "output": 1.00 / 1_000_000,    # $1.00 per 1M (50% off $2.00)
+    },
+    "gpt-5-nano": {
+        "input": 0.025 / 1_000_000,    # $0.025 per 1M (50% off $0.05)
+        "output": 0.20 / 1_000_000,    # $0.20 per 1M (50% off $0.40)
+    },
+    # GPT-4o series with Batch API discount
+    "gpt-4o": {
+        "input": 1.25 / 1_000_000,     # $1.25 per 1M (50% off $2.50)
+        "output": 5.00 / 1_000_000,    # $5.00 per 1M (50% off $10.00)
+    },
+    "gpt-4o-2024-05-13": {
+        "input": 2.50 / 1_000_000,     # $2.50 per 1M (50% off $5.00)
+        "output": 7.50 / 1_000_000,    # $7.50 per 1M (50% off $15.00)
+    },
+    "gpt-4o-mini": {
+        "input": 0.075 / 1_000_000,    # $0.075 per 1M (50% off $0.15)
+        "output": 0.30 / 1_000_000,    # $0.30 per 1M (50% off $0.60)
+    },
+    # GPT-4 Turbo
+    "gpt-4-turbo": {
+        "input": 5.00 / 1_000_000,     # $5.00 per 1M (50% off $10.00)
+        "output": 15.00 / 1_000_000,   # $15.00 per 1M (50% off $30.00)
+    },
+    # GPT-3.5 Turbo
+    "gpt-3.5-turbo": {
+        "input": 0.25 / 1_000_000,     # $0.25 per 1M (50% off $0.50)
+        "output": 0.75 / 1_000_000,    # $0.75 per 1M (50% off $1.50)
+    },
+}
+
+PROMPT_TEMPLATE = """You are an expert in scientific computing and computational chemistry/biology/physics. Please create a high-quality programming problem inspired by the following code snippet from a real scientific computing project.
+
+The problem should focus on scientific computing concepts such as:
+- Numerical algorithms and simulations
+- Data analysis and visualization
+- Mathematical modeling
+- Scientific data processing
+- Computational methods in chemistry, biology, or physics
+
+Code snippet for inspiration:
+```python
+{code}
+```
+
+Present your output in two distinct sections:
+
+[Problem Description]
+Create a **completely self-contained** problem description that:
+- Does NOT directly reference the code snippet above
+- Provides all necessary context and background
+- Clearly states what needs to be implemented
+- Specifies input/output format and constraints
+- Is inspired by the scientific computing concepts in the code but creates a NEW, interesting problem
+- Assumes common programming knowledge but explains any domain-specific concepts
+
+[Solution]
+Provide a comprehensive, **correct** Python solution that:
+- Accurately solves the problem described
+- Includes clear comments explaining the approach
+- Uses appropriate scientific computing libraries (numpy, scipy, etc.) when relevant
+- Is complete and runnable
+- Follows best practices for scientific computing
+
+Remember: The problem should be INSPIRED by the code, not a direct copy. Create something educational and interesting for scientific computing practitioners."""
+
+
+class BatchAPIClient:
+    """Client for OpenAI Batch API with cost tracking."""
+    
+    def __init__(self, model_name: str = MODEL_NAME, api_key: Optional[str] = None):
+        """Initialize OpenAI Batch API client.
+        
+        Args:
+            model_name: Name of the OpenAI model to use
+            api_key: OpenAI API key (if None, will use OPENAI_API_KEY env variable)
+        """
+        self.model_name = model_name
+        self.client = OpenAI(api_key=api_key)
+        
+        # Get pricing for the model (Batch API is 50% off)
+        if model_name in BATCH_PRICING:
+            self.input_price = BATCH_PRICING[model_name]["input"]
+            self.output_price = BATCH_PRICING[model_name]["output"]
+        else:
+            print(f"Warning: No Batch pricing info for {model_name}, using gpt-4o-mini prices")
+            self.input_price = BATCH_PRICING["gpt-4o-mini"]["input"]
+            self.output_price = BATCH_PRICING["gpt-4o-mini"]["output"]
+        
+        print(f"📊 Batch API Pricing (50% off standard rates):")
+        print(f"   Input:  ${self.input_price * 1_000_000:.4f} per 1M tokens")
+        print(f"   Output: ${self.output_price * 1_000_000:.4f} per 1M tokens")
+        print()
+    
+    def create_batch_file(self, requests: List[Dict], output_path: str) -> str:
+        """Create a JSONL file for batch processing.
+        
+        Args:
+            requests: List of request dictionaries
+            output_path: Path to save the JSONL file
+            
+        Returns:
+            Path to the created file
+        """
+        with open(output_path, 'w', encoding='utf-8') as f:
+            for req in requests:
+                f.write(json.dumps(req, ensure_ascii=False) + '\n')
+        
+        print(f"✅ Created batch file: {output_path}")
+        print(f"   Total requests: {len(requests)}")
+        return output_path
+    
+    def upload_batch_file(self, file_path: str) -> str:
+        """Upload batch file to OpenAI.
+        
+        Args:
+            file_path: Path to the JSONL file
+            
+        Returns:
+            File ID
+        """
+        print(f"⬆️  Uploading batch file to OpenAI...")
+        with open(file_path, 'rb') as f:
+            batch_file = self.client.files.create(
+                file=f,
+                purpose='batch'
+            )
+        
+        print(f"✅ File uploaded: {batch_file.id}")
+        return batch_file.id
+    
+    def create_batch(self, file_id: str, description: Optional[str] = None) -> str:
+        """Create a batch job.
+        
+        Args:
+            file_id: ID of the uploaded file
+            description: Optional description for the batch
+            
+        Returns:
+            Batch ID
+        """
+        print(f"🚀 Creating batch job...")
+        batch = self.client.batches.create(
+            input_file_id=file_id,
+            endpoint="/v1/chat/completions",
+            completion_window="24h",
+            metadata={
+                "description": description or "Programming problems generation",
+                "created_at": datetime.now().isoformat()
+            }
+        )
+        
+        print(f"✅ Batch created: {batch.id}")
+        print(f"   Status: {batch.status}")
+        print(f"   Total requests: {batch.request_counts.total}")
+        return batch.id
+    
+    def check_batch_status(self, batch_id: str) -> Dict:
+        """Check the status of a batch job.
+        
+        Args:
+            batch_id: ID of the batch
+            
+        Returns:
+            Batch status information
+        """
+        batch = self.client.batches.retrieve(batch_id)
+        
+        status_info = {
+            'id': batch.id,
+            'status': batch.status,
+            'created_at': batch.created_at,
+            'completed_at': batch.completed_at,
+            'failed_at': batch.failed_at,
+            'expired_at': batch.expired_at,
+            'request_counts': {
+                'total': batch.request_counts.total,
+                'completed': batch.request_counts.completed,
+                'failed': batch.request_counts.failed,
+            },
+            'output_file_id': batch.output_file_id,
+            'error_file_id': batch.error_file_id,
+        }
+        
+        return status_info
+    
+    def download_results(self, file_id: str, output_path: str):
+        """Download batch results.
+        
+        Args:
+            file_id: ID of the output file
+            output_path: Path to save the results
+        """
+        print(f"⬇️  Downloading results...")
+        content = self.client.files.content(file_id)
+        
+        with open(output_path, 'wb') as f:
+            f.write(content.content)
+        
+        print(f"✅ Results saved to: {output_path}")
+    
+    def estimate_cost(self, num_requests: int, avg_input_tokens: int, avg_output_tokens: int) -> Dict:
+        """Estimate the cost of a batch job.
+        
+        Args:
+            num_requests: Number of requests
+            avg_input_tokens: Average input tokens per request
+            avg_output_tokens: Average output tokens per request
+            
+        Returns:
+            Cost estimation dictionary
+        """
+        total_input_tokens = num_requests * avg_input_tokens
+        total_output_tokens = num_requests * avg_output_tokens
+        
+        input_cost = total_input_tokens * self.input_price
+        output_cost = total_output_tokens * self.output_price
+        total_cost = input_cost + output_cost
+        
+        # Compare with standard API (2x the batch price)
+        standard_cost = total_cost * 2
+        savings = standard_cost - total_cost
+        
+        return {
+            'num_requests': num_requests,
+            'total_input_tokens': total_input_tokens,
+            'total_output_tokens': total_output_tokens,
+            'total_tokens': total_input_tokens + total_output_tokens,
+            'input_cost': input_cost,
+            'output_cost': output_cost,
+            'total_cost': total_cost,
+            'standard_api_cost': standard_cost,
+            'savings': savings,
+            'savings_percentage': 50.0
+        }
+
+
+def prepare_batch_requests(
+    input_file: str,
+    min_score: int = MIN_RELEVANCE_SCORE,
+    max_samples: Optional[int] = None,
+    start_from: int = 0,
+) -> List[Dict]:
+    """Prepare batch requests from function dataset.
+    
+    Args:
+        input_file: Path to function_dataset_v2.csv
+        min_score: Minimum relevance score to process
+        max_samples: Maximum number of samples to process
+        start_from: Skip first N rows
+        
+    Returns:
+        List of batch request dictionaries
+    """
+    print(f"📋 Preparing batch requests...")
+    print(f"   Input: {input_file}")
+    print(f"   Min Score: {min_score}")
+    if max_samples:
+        print(f"   Max Samples: {max_samples}")
+    print()
+    
+    requests = []
+    total_rows = 0
+    skipped_low_score = 0
+    skipped_no_code = 0
+    
+    with open(input_file, 'r', encoding='utf-8') as infile:
+        reader = csv.DictReader(infile)
+        
+        for row in reader:
+            total_rows += 1
+            
+            # Skip if resuming
+            if total_rows <= start_from:
+                continue
+            
+            # Check if we've reached max samples
+            if max_samples and len(requests) >= max_samples:
+                break
+            
+            # Filter by relevance score
+            try:
+                relevance_score = int(row.get('relevance_score', 0))
+            except (ValueError, TypeError):
+                relevance_score = 0
+            
+            if relevance_score < min_score:
+                skipped_low_score += 1
+                continue
+            
+            # Get function content
+            function_content = row.get('function_content', '').strip()
+            if not function_content or len(function_content) < 50:
+                skipped_no_code += 1
+                continue
+            
+            # Prepare metadata (OpenAI Batch API requires all metadata values to be strings)
+            metadata = {
+                'original_index': str(row.get('original_index', '')),
+                'function_name': str(row.get('function_name', '')),
+                'repo_name': str(row.get('repo_name', '')),
+                'path': str(row.get('path', '')),
+                'language': str(row.get('language', '')),
+                'relevance_score': str(relevance_score),  # Convert to string!
+                'function_start_line': str(row.get('function_start_line', '')),
+                'function_end_line': str(row.get('function_end_line', '')),
+            }
+            
+            # Generate prompt
+            prompt = PROMPT_TEMPLATE.format(code=function_content)
+            
+            # Create batch request in OpenAI Batch API format
+            request = {
+                "custom_id": f"request-{len(requests)}",
+                "method": "POST",
+                "url": "/v1/chat/completions",
+                "body": {
+                    "model": MODEL_NAME,
+                    "messages": [
+                        {
+                            "role": "system",
+                            "content": "You are an expert in scientific computing and programming education."
+                        },
+                        {
+                            "role": "user",
+                            "content": prompt
+                        }
+                    ],
+                    "temperature": 0.7,
+                    "metadata": metadata  # All values are now strings
+                }
+            }
+            
+            requests.append(request)
+    
+    print(f"✅ Prepared {len(requests)} requests")
+    print(f"   Total rows: {total_rows}")
+    print(f"   Skipped (low score): {skipped_low_score}")
+    print(f"   Skipped (no/short code): {skipped_no_code}")
+    print()
+    
+    return requests
+
+
+def process_batch_results(
+    results_file: str,
+    output_file: str,
+    model_name: str,
+    input_price: float,
+    output_price: float,
+    requests_file: Optional[str] = None
+):
+    """Process batch results and save to JSONL format.
+    
+    Args:
+        results_file: Path to batch results file
+        output_file: Path to output JSONL file
+        model_name: Model name used
+        input_price: Input token price
+        output_price: Output token price
+        requests_file: Optional path to original batch requests file (to restore prompts)
+    """
+    print(f"📊 Processing batch results...")
+    
+    # Load prompts from requests file if provided
+    prompts_map = {}
+    if requests_file and os.path.exists(requests_file):
+        print(f"   Loading prompts from: {requests_file}")
+        with open(requests_file, 'r', encoding='utf-8') as f:
+            for line in f:
+                req = json.loads(line)
+                custom_id = req['custom_id']
+                # Extract prompt from messages
+                for msg in req['body']['messages']:
+                    if msg['role'] == 'user':
+                        prompts_map[custom_id] = msg['content']
+                        break
+        print(f"   Loaded {len(prompts_map)} prompts")
+    
+    processed = 0
+    errors = 0
+    total_input_tokens = 0
+    total_output_tokens = 0
+    total_cost = 0.0
+    
+    with open(results_file, 'r', encoding='utf-8') as infile, \
+         open(output_file, 'w', encoding='utf-8') as outfile:
+        
+        for line in infile:
+            batch_result = json.loads(line)
+            
+            # Check if request was successful
+            if batch_result.get('error'):
+                errors += 1
+                print(f"❌ Error in {batch_result['custom_id']}: {batch_result['error']}")
+                continue
+            
+            response = batch_result['response']
+            custom_id = batch_result['custom_id']
+            
+            # Extract usage information
+            usage = response['body']['usage']
+            input_tokens = usage['prompt_tokens']
+            output_tokens = usage['completion_tokens']
+            
+            # Calculate cost
+            input_cost = input_tokens * input_price
+            output_cost = output_tokens * output_price
+            request_cost = input_cost + output_cost
+            
+            # Update totals
+            total_input_tokens += input_tokens
+            total_output_tokens += output_tokens
+            total_cost += request_cost
+            
+            # Get metadata from the original request
+            metadata = response['body'].get('metadata', {})
+            
+            # Extract the response text
+            response_text = response['body']['choices'][0]['message']['content']
+            
+            # Build result - include prompt if available
+            result = {
+                'metadata': metadata,
+                'response': response_text,
+                'usage': {
+                    'input_tokens': input_tokens,
+                    'output_tokens': output_tokens,
+                    'total_tokens': input_tokens + output_tokens,
+                    'input_cost': input_cost,
+                    'output_cost': output_cost,
+                    'request_cost': request_cost
+                },
+                'model': model_name,
+                'timestamp': datetime.now().isoformat(),
+                'custom_id': custom_id
+            }
+            
+            # Add prompt if we have it
+            if custom_id in prompts_map:
+                result['prompt'] = prompts_map[custom_id]
+            
+            outfile.write(json.dumps(result, ensure_ascii=False) + '\n')
+            processed += 1
+    
+    print(f"\n✅ Processed {processed} results")
+    print(f"   Errors: {errors}")
+    print()
+    
+    # Print usage summary
+    print("=" * 70)
+    print("BATCH API USAGE SUMMARY")
+    print("=" * 70)
+    print(f"Model:                 {model_name}")
+    print(f"Total Requests:        {processed}")
+    print(f"Total Input Tokens:    {total_input_tokens:,}")
+    print(f"Total Output Tokens:   {total_output_tokens:,}")
+    print(f"Total Tokens:          {total_input_tokens + total_output_tokens:,}")
+    print(f"\nBatch API Cost:        ${total_cost:.6f}")
+    print(f"Standard API Cost:     ${total_cost * 2:.6f}")
+    print(f"Savings (50%):         ${total_cost:.6f}")
+    print("=" * 70)
+
+
+def main():
+    import argparse
+    
+    parser = argparse.ArgumentParser(
+        description='Generate programming problems using OpenAI Batch API (50% cost savings)'
+    )
+    
+    subparsers = parser.add_subparsers(dest='command', help='Command to run')
+    
+    # Prepare command
+    prepare_parser = subparsers.add_parser('prepare', help='Prepare batch requests')
+    prepare_parser.add_argument('--input', default='function_dataset_v2.csv')
+    prepare_parser.add_argument('--output', default='batch_requests.jsonl')
+    prepare_parser.add_argument('--min-score', type=int, default=MIN_RELEVANCE_SCORE)
+    prepare_parser.add_argument('--max-samples', type=int, default=None)
+    prepare_parser.add_argument('--start-from', type=int, default=0)
+    prepare_parser.add_argument('--model', default=MODEL_NAME)
+    
+    # Submit command
+    submit_parser = subparsers.add_parser('submit', help='Submit batch job to OpenAI')
+    submit_parser.add_argument('--input', default='batch_requests.jsonl')
+    submit_parser.add_argument('--model', default=MODEL_NAME)
+    submit_parser.add_argument('--description', default='Programming problems generation')
+    
+    # Status command
+    status_parser = subparsers.add_parser('status', help='Check batch job status')
+    status_parser.add_argument('batch_id', help='Batch ID to check')
+    
+    # Download command
+    download_parser = subparsers.add_parser('download', help='Download batch results')
+    download_parser.add_argument('batch_id', help='Batch ID to download')
+    download_parser.add_argument('--output', default='batch_results.jsonl')
+    
+    # Process command
+    process_parser = subparsers.add_parser('process', help='Process downloaded results')
+    process_parser.add_argument('--input', default='batch_results.jsonl')
+    process_parser.add_argument('--output', default='programming_problems_batch.jsonl')
+    process_parser.add_argument('--model', default=MODEL_NAME)
+    process_parser.add_argument('--requests', default='batch_requests_full.jsonl',
+                                help='Original batch requests file (to restore prompts)')
+    
+    # Estimate command
+    estimate_parser = subparsers.add_parser('estimate', help='Estimate batch cost')
+    estimate_parser.add_argument('--num-requests', type=int, required=True)
+    estimate_parser.add_argument('--avg-input-tokens', type=int, default=1917)
+    estimate_parser.add_argument('--avg-output-tokens', type=int, default=2552)
+    estimate_parser.add_argument('--model', default=MODEL_NAME)
+    
+    args = parser.parse_args()
+    
+    if not args.command:
+        parser.print_help()
+        sys.exit(1)
+    
+    # Check API key
+    if not os.getenv('OPENAI_API_KEY'):
+        print("❌ Error: OPENAI_API_KEY environment variable not set.")
+        print("   Please set it with: export OPENAI_API_KEY='your-api-key'")
+        sys.exit(1)
+    
+    client = BatchAPIClient(model_name=args.model if hasattr(args, 'model') else MODEL_NAME)
+    
+    if args.command == 'prepare':
+        requests = prepare_batch_requests(
+            input_file=args.input,
+            min_score=args.min_score,
+            max_samples=args.max_samples,
+            start_from=args.start_from
+        )
+        
+        client.create_batch_file(requests, args.output)
+        
+        # Estimate cost
+        print("\n💰 Cost Estimation:")
+        estimate = client.estimate_cost(
+            num_requests=len(requests),
+            avg_input_tokens=1917,  # From your test
+            avg_output_tokens=2552   # From your test
+        )
+        print(f"   Estimated Batch API Cost:   ${estimate['total_cost']:.2f}")
+        print(f"   Standard API Cost:          ${estimate['standard_api_cost']:.2f}")
+        print(f"   Savings (50%):              ${estimate['savings']:.2f}")
+        print()
+        
+    elif args.command == 'submit':
+        file_id = client.upload_batch_file(args.input)
+        batch_id = client.create_batch(file_id, args.description)
+        
+        print(f"\n📝 Save this Batch ID: {batch_id}")
+        print(f"   Check status with: python3 {sys.argv[0]} status {batch_id}")
+        
+    elif args.command == 'status':
+        status = client.check_batch_status(args.batch_id)
+        
+        print("\n📊 Batch Status:")
+        print(f"   ID: {status['id']}")
+        print(f"   Status: {status['status']}")
+        print(f"   Total: {status['request_counts']['total']}")
+        print(f"   Completed: {status['request_counts']['completed']}")
+        print(f"   Failed: {status['request_counts']['failed']}")
+        
+        if status['status'] == 'completed':
+            print(f"\n✅ Batch completed!")
+            print(f"   Download with: python3 {sys.argv[0]} download {args.batch_id}")
+        elif status['status'] == 'failed':
+            print(f"\n❌ Batch failed!")
+        else:
+            print(f"\n⏳ Batch is still processing...")
+        
+    elif args.command == 'download':
+        status = client.check_batch_status(args.batch_id)
+        
+        if status['status'] != 'completed':
+            print(f"❌ Batch is not completed yet (status: {status['status']})")
+            sys.exit(1)
+        
+        client.download_results(status['output_file_id'], args.output)
+        print(f"\n✅ Downloaded to: {args.output}")
+        print(f"   Process with: python3 {sys.argv[0]} process --input {args.output}")
+        
+    elif args.command == 'process':
+        process_batch_results(
+            results_file=args.input,
+            output_file=args.output,
+            model_name=args.model,
+            input_price=client.input_price,
+            output_price=client.output_price,
+            requests_file=args.requests
+        )
+        print(f"\n✅ Final results saved to: {args.output}")
+        
+    elif args.command == 'estimate':
+        estimate = client.estimate_cost(
+            num_requests=args.num_requests,
+            avg_input_tokens=args.avg_input_tokens,
+            avg_output_tokens=args.avg_output_tokens
+        )
+        
+        print("\n💰 COST ESTIMATION")
+        print("=" * 70)
+        print(f"Number of Requests:    {estimate['num_requests']:,}")
+        print(f"Total Input Tokens:    {estimate['total_input_tokens']:,}")
+        print(f"Total Output Tokens:   {estimate['total_output_tokens']:,}")
+        print(f"Total Tokens:          {estimate['total_tokens']:,}")
+        print()
+        print(f"Batch API Cost:        ${estimate['total_cost']:.2f}")
+        print(f"Standard API Cost:     ${estimate['standard_api_cost']:.2f}")
+        print(f"💰 Savings (50%):      ${estimate['savings']:.2f}")
+        print("=" * 70)
+
+
+if __name__ == "__main__":
+    main()