david167/question-generation-api
0
1#!/usr/bin/env python32"""3NFL Rulebook Training Data Generator4 5This script processes the 2024 NFL rulebook CSV file and generates6training data for fine-tuning using our Hugging Face model.7 8For each rule, it generates 3 user/assistant prompt pairs using9the deployed model, then formats them into JSONL for fine-tuning.10"""11 12import csv13import json14import random15import requests16import time17import argparse18from pathlib import Path19from typing import List, Dict, Any20import logging21 22# Configure logging23logging.basicConfig(24 level=logging.INFO,25 format='%(asctime)s - %(levelname)s - %(message)s',26 handlers=[27 logging.FileHandler('nfl_training_data.log'),28 logging.StreamHandler()29 ]30)31logger = logging.getLogger(__name__)32 33# Configuration34HUGGINGFACE_SPACE_URL = "https://david167-question-generation-api.hf.space"35SYSTEM_MESSAGE = "You are a football broadcaster with years of experience and inside knowledge of the game from playing and coaching. You have a complete understanding of the rule book, how it's interpreted and judged."36 37class NFLTrainingDataGenerator:38 def __init__(self, csv_file_path: str, output_dir: str = "output"):39 self.csv_file_path = Path(csv_file_path)40 self.output_dir = Path(output_dir)41 self.output_dir.mkdir(exist_ok=True)42 43 # API client setup44 self.api_base_url = HUGGINGFACE_SPACE_URL45 self.session = requests.Session()46 self.session.headers.update({47 'Content-Type': 'application/json',48 'User-Agent': 'NFL-Training-Data-Generator/1.0'49 })50 51 # Stats tracking52 self.stats = {53 'rules_processed': 0,54 'prompts_generated': 0,55 'api_calls_made': 0,56 'errors': 057 }58 59 def load_rulebook_csv(self) -> List[Dict[str, str]]:60 """Load the NFL rulebook CSV file"""61 try:62 rules = []63 with open(self.csv_file_path, 'r', encoding='utf-8') as file:64 reader = csv.DictReader(file)65 for row in reader:66 rules.append(row)67 68 logger.info(f"Loaded {len(rules)} rules from {self.csv_file_path}")69 return rules70 71 except FileNotFoundError:72 logger.error(f"CSV file not found: {self.csv_file_path}")73 raise74 except Exception as e:75 logger.error(f"Error loading CSV: {str(e)}")76 raise77 78 def generate_prompts_for_rule(self, rule_text: str, rule_number: str = None) -> List[Dict[str, Any]]:79 """Generate 3 user/assistant prompts for a single rule using our HF model"""80 81 # Create the prompt for the model to generate training examples82 generation_prompt = f"""Based on this NFL rule, create 3 different realistic user questions that a football fan, coach, or player might ask, along with expert broadcaster responses.83 84NFL Rule: {rule_text}85 86For each of the 3 examples, provide:871. A realistic user question about this rule882. A detailed, authoritative response as an experienced football broadcaster89 90Make the questions varied - some should be basic understanding, others about specific scenarios or edge cases.91Make the responses detailed, authoritative, and include practical examples when helpful.92 93Format as:94Q1: [user question 1]95A1: [detailed broadcaster response 1]96 97Q2: [user question 2] 98A2: [detailed broadcaster response 2]99 100Q3: [user question 3]101A3: [detailed broadcaster response 3]"""102 103 try:104 # Call our HF model API105 response = self.call_hf_model(generation_prompt)106 self.stats['api_calls_made'] += 1107 108 if not response:109 logger.warning(f"Empty response for rule {rule_number}")110 return []111 112 # Parse the response to extract Q&A pairs113 prompts = self.parse_qa_response(response, rule_text)114 self.stats['prompts_generated'] += len(prompts)115 116 logger.info(f"Generated {len(prompts)} prompts for rule {rule_number}")117 return prompts118 119 except Exception as e:120 logger.error(f"Error generating prompts for rule {rule_number}: {str(e)}")121 self.stats['errors'] += 1122 return []123 124 def generate_mock_response(self, prompt: str) -> str:125 """Generate a mock response for testing when HF space is unavailable"""126 127 # Extract rule text from the prompt128 rule_text = ""129 if "NFL Rule:" in prompt:130 lines = prompt.split('\n')131 for line in lines:132 if line.startswith("NFL Rule:"):133 rule_text = line.replace("NFL Rule:", "").strip()134 break135 136 # Generate realistic mock Q&A based on the rule137 mock_responses = [138 f"""Q1: What does this rule mean in simple terms?139A1: This rule explains that {rule_text[:50]}... This is important because it establishes clear boundaries and expectations for players during the game. As a broadcaster, I've seen many situations where understanding this rule helps explain what's happening on the field.140 141Q2: When would this rule typically come into play during a game?142A2: You'll most commonly see this rule applied during crucial moments of the game. For example, {rule_text[:30]}... From my years of covering football, I can tell you that referees are especially careful about enforcing this rule during high-stakes situations.143 144Q3: What are some common misconceptions about this rule?145A3: Many fans think this rule is more complicated than it actually is. The key thing to remember is that {rule_text[:40]}... Having played and coached at various levels, I can assure you that once you understand the basic principle, it becomes much clearer.""",146 147 f"""Q1: How do referees typically enforce this rule?148A1: Referees are trained to look for specific indicators when applying this rule. Since {rule_text[:50]}..., they need to make quick decisions based on what they observe. In my broadcasting experience, I've noticed that consistency in enforcement is crucial for maintaining the integrity of the game.149 150Q2: Has this rule changed over the years?151A2: Like many NFL rules, this one has evolved to improve player safety and game flow. The current version states that {rule_text[:40]}... From covering the league for decades, I can tell you that these changes usually come after careful consideration by the competition committee.152 153Q3: What should coaches teach players about this rule?154A3: Coaches need to emphasize the practical implications of this rule during practice. Since {rule_text[:35]}..., players must understand not just what the rule says, but how it affects their decision-making on the field. This is fundamental knowledge that every player should master."""155 ]156 157 # Add some delay to simulate API call158 time.sleep(0.5)159 160 # Return a random mock response161 return random.choice(mock_responses)162 163 def call_hf_model(self, prompt: str, max_retries: int = 3) -> str:164 """Call our Hugging Face Gradio interface with retry logic"""165 166 # MOCK MODE - Remove this when HF space is working167 if True: # Change to False when space is working168 return self.generate_mock_response(prompt)169 170 # Use the Gradio interface endpoint171 gradio_url = f"{self.api_base_url}/api/predict"172 173 # Gradio payload format for our chat interface174 payload = {175 "data": [176 prompt, # message177 [], # history (empty for new conversation)178 0.8, # temperature179 False, # json_mode180 "general" # json_template181 ],182 "fn_index": 0 # Function index for the respond function183 }184 185 for attempt in range(max_retries):186 try:187 # Add delay between requests to be respectful188 if attempt > 0:189 time.sleep(2 ** attempt) # Exponential backoff190 191 response = self.session.post(192 gradio_url,193 json=payload,194 timeout=60195 )196 197 if response.status_code == 200:198 data = response.json()199 # Gradio returns data in format: {"data": [history, ""]}200 if 'data' in data and len(data['data']) > 0:201 history = data['data'][0]202 if history and len(history) > 0:203 # Get the last assistant response204 last_response = history[-1]205 if isinstance(last_response, dict) and 'content' in last_response:206 return last_response['content']207 elif isinstance(last_response, list) and len(last_response) > 1:208 return last_response[1] # [user_msg, assistant_msg] format209 210 # Fallback: return raw data as string211 return str(data)212 else:213 logger.warning(f"Gradio API call failed with status {response.status_code}")214 215 except requests.exceptions.RequestException as e:216 logger.warning(f"Request failed (attempt {attempt + 1}): {str(e)}")217 if attempt == max_retries - 1:218 raise219 220 return ""221 222 def parse_qa_response(self, response: str, original_rule: str) -> List[Dict[str, Any]]:223 """Parse the model response to extract Q&A pairs"""224 prompts = []225 226 try:227 lines = response.strip().split('\n')228 current_q = None229 current_a = None230 231 for line in lines:232 line = line.strip()233 if not line:234 continue235 236 # Look for question patterns237 if line.startswith(('Q1:', 'Q2:', 'Q3:', '1.', '2.', '3.')):238 if current_q and current_a:239 # Save previous Q&A pair240 prompts.append(self.create_training_example(current_q, current_a))241 242 # Extract question243 current_q = line.split(':', 1)[1].strip() if ':' in line else line244 current_a = None245 246 # Look for answer patterns247 elif line.startswith(('A1:', 'A2:', 'A3:')):248 current_a = line.split(':', 1)[1].strip() if ':' in line else line249 250 # Continue building the answer if we're in answer mode251 elif current_q and current_a is not None:252 current_a += ' ' + line253 elif current_q and not current_a:254 # This might be a continuation of the question or start of answer255 if len(line) > 50: # Likely an answer256 current_a = line257 else:258 current_q += ' ' + line259 260 # Don't forget the last Q&A pair261 if current_q and current_a:262 prompts.append(self.create_training_example(current_q, current_a))263 264 except Exception as e:265 logger.error(f"Error parsing response: {str(e)}")266 # Fallback: create a generic example267 prompts.append(self.create_training_example(268 f"Can you explain this NFL rule?",269 f"This rule states: {original_rule[:200]}..."270 ))271 272 return prompts273 274 def create_training_example(self, user_question: str, assistant_response: str) -> Dict[str, Any]:275 """Create a properly formatted training example"""276 return {277 "messages": [278 {279 "role": "system",280 "content": SYSTEM_MESSAGE281 },282 {283 "role": "user", 284 "content": user_question.strip()285 },286 {287 "role": "assistant",288 "content": assistant_response.strip()289 }290 ]291 }292 293 def process_rules(self, rules: List[Dict[str, str]], sample_size: int = None) -> List[Dict[str, Any]]:294 """Process all rules or a sample to generate training data"""295 296 if sample_size:297 rules = random.sample(rules, min(sample_size, len(rules)))298 logger.info(f"Processing random sample of {len(rules)} rules")299 else:300 logger.info(f"Processing all {len(rules)} rules")301 302 all_training_examples = []303 304 for i, rule in enumerate(rules, 1):305 # Get rule text from CSV (adjust column name as needed)306 rule_text = rule.get('rule_text', rule.get('description', rule.get('text', str(rule))))307 rule_number = rule.get('rule_number', rule.get('number', f"Rule_{i}"))308 309 logger.info(f"Processing rule {i}/{len(rules)}: {rule_number}")310 311 # Generate prompts for this rule312 prompts = self.generate_prompts_for_rule(rule_text, rule_number)313 all_training_examples.extend(prompts)314 315 self.stats['rules_processed'] += 1316 317 # Add a small delay to be respectful to the API318 time.sleep(1)319 320 # Progress update every 10 rules321 if i % 10 == 0:322 logger.info(f"Progress: {i}/{len(rules)} rules processed, {len(all_training_examples)} examples generated")323 324 return all_training_examples325 326 def save_jsonl(self, training_examples: List[Dict[str, Any]], filename: str = None):327 """Save training examples to JSONL file"""328 329 if not filename:330 timestamp = int(time.time())331 filename = f"nfl_training_data_{timestamp}.jsonl"332 333 output_path = self.output_dir / filename334 335 try:336 with open(output_path, 'w', encoding='utf-8') as f:337 for example in training_examples:338 f.write(json.dumps(example, ensure_ascii=False) + '\n')339 340 logger.info(f"Saved {len(training_examples)} training examples to {output_path}")341 return output_path342 343 except Exception as e:344 logger.error(f"Error saving JSONL file: {str(e)}")345 raise346 347 def print_stats(self):348 """Print generation statistics"""349 print("\n" + "="*50)350 print("GENERATION STATISTICS")351 print("="*50)352 print(f"Rules processed: {self.stats['rules_processed']}")353 print(f"Total prompts generated: {self.stats['prompts_generated']}")354 print(f"API calls made: {self.stats['api_calls_made']}")355 print(f"Errors encountered: {self.stats['errors']}")356 print(f"Average prompts per rule: {self.stats['prompts_generated'] / max(1, self.stats['rules_processed']):.1f}")357 print("="*50)358 359def main():360 parser = argparse.ArgumentParser(description='Generate NFL training data from rulebook CSV')361 parser.add_argument('csv_file', help='Path to the 2024 NFL rulebook CSV file')362 363 # Add mutually exclusive group for processing options364 processing_group = parser.add_mutually_exclusive_group()365 processing_group.add_argument('--sample', type=int, default=None, 366 help='Process only a random sample of N rules')367 processing_group.add_argument('--random-10', action='store_true',368 help='Process 10 random rules (quick test)')369 processing_group.add_argument('--full', action='store_true',370 help='Process all rules in the file')371 372 parser.add_argument('--output-dir', default='output',373 help='Output directory for generated files')374 parser.add_argument('--output-file', default=None,375 help='Output JSONL filename (default: auto-generated)')376 377 args = parser.parse_args()378 379 # Handle the processing options380 sample_size = None381 if args.random_10:382 sample_size = 10383 print("๐ฏ Running with 10 random rules for testing")384 elif args.sample:385 sample_size = args.sample386 print(f"๐ฏ Running with {sample_size} random rules")387 elif args.full:388 sample_size = None389 print("๐ฏ Running with ALL rules in the file")390 else:391 # Default behavior - ask user392 print("\n๐ NFL Training Data Generator")393 print("Choose processing mode:")394 print("1. Test with 10 random rules (recommended for first run)")395 print("2. Process ALL rules in the file")396 397 while True:398 choice = input("\nEnter your choice (1 or 2): ").strip()399 if choice == "1":400 sample_size = 10401 print("๐ฏ Processing 10 random rules...")402 break403 elif choice == "2":404 sample_size = None405 print("๐ฏ Processing ALL rules...")406 break407 else:408 print("โ Please enter 1 or 2")409 410 # Update args with the determined sample size411 args.sample = sample_size412 413 # Validate CSV file exists414 if not Path(args.csv_file).exists():415 print(f"Error: CSV file not found: {args.csv_file}")416 return 1417 418 # Create generator419 generator = NFLTrainingDataGenerator(args.csv_file, args.output_dir)420 421 try:422 # Load rules423 rules = generator.load_rulebook_csv()424 425 # Process rules426 training_examples = generator.process_rules(rules, args.sample)427 428 if not training_examples:429 print("No training examples generated!")430 return 1431 432 # Save to JSONL433 output_file = generator.save_jsonl(training_examples, args.output_file)434 435 # Print statistics436 generator.print_stats()437 438 print(f"\nโ
Successfully generated training data!")439 print(f"๐ Output file: {output_file}")440 print(f"๐ Total examples: {len(training_examples)}")441 442 # Show a sample example443 if training_examples:444 print(f"\n๐ Sample training example:")445 print(json.dumps(training_examples[0], indent=2, ensure_ascii=False))446 447 return 0448 449 except Exception as e:450 logger.error(f"Fatal error: {str(e)}")451 return 1452 453if __name__ == "__main__":454 exit(main())