m

best_of_n.py

@@ -0,0 +1,368 @@
+import aiohttp
+import asyncio
+import json
+import os
+import re
+from tqdm.asyncio import tqdm_asyncio
+from tqdm import tqdm
+from typing import List, Dict, Any, Tuple
+import gc
+from datetime import datetime
+from datasets import load_dataset
+
+# 1 sample
+# Overall Statistics:
+# Total correct answers: 58/100
+# Accuracy: 0.58
+# Average reward score: 0.0929
+
+# best of 2
+# Overall Statistics:
+# Total correct answers: 63/100
+# Accuracy: 0.63
+# Average reward score: 0.1148
+
+# best of 4
+# Overall Statistics:
+# Total correct answers: 64/100
+# Accuracy: 0.64
+# Average reward score: 0.1257
+
+# best of 8
+# Overall Statistics:
+# Total correct answers: 63/100
+# Accuracy: 0.63
+# Average reward score: 0.1307
+
+# best of 16
+# Overall Statistics:
+# Total correct answers: 70/100
+# Accuracy: 0.70
+# Average reward score: 0.1345
+
+# best of 32
+# Overall Statistics:
+# Total correct answers: 67/100
+# Accuracy: 0.67
+# Average reward score: 0.1380
+
+# best of 64
+# Overall Statistics:
+# Total correct answers: 67/100
+# Accuracy: 0.67
+# Average reward score: 0.1461
+
+
+
+# Configuration
+FIREWORKS_API_KEY = os.getenv('FIREWORKS_API_KEY')
+if not FIREWORKS_API_KEY:
+    raise ValueError("FIREWORKS_API_KEY environment variable must be set")
+
+FIREWORKS_API_ENDPOINT = "https://api.fireworks.ai/inference/v1/chat/completions"
+REWARD_MODEL_ENDPOINT = "https://rawsh--reward-api-model-score.modal.run"
+
+# Separate rate limits for different APIs
+LLM_MAX_CONCURRENT = 100
+REWARD_MAX_CONCURRENT = 32
+BATCH_SIZE = 10
+
+BEST_OF_N = 2
+MODEL_NAME = "accounts/fireworks/models/llama-v3p1-8b-instruct"
+
+# Timeout configurations
+REWARD_MODEL_TIMEOUT = 20
+LLM_TIMEOUT = 10
+REWARD_MODEL_MAX_RETRIES = 3
+
+class APIError(Exception):
+    """Custom exception for API-related errors"""
+    pass
+
+async def with_retry(func, max_retries=3, base_delay=1):
+    """Generic retry wrapper with exponential backoff"""
+    for i in range(max_retries):
+        try:
+            return await func()
+        except Exception as e:
+            if i == max_retries - 1:
+                raise
+            delay = base_delay * (2 ** i)
+            print(f"Attempt {i+1} failed: {str(e)}. Retrying in {delay}s...")
+            await asyncio.sleep(delay)
+
+def extract_answer(completion: str) -> Tuple[str, str]:
+    """Extract the final answer from the completion."""
+    match = re.search(r"Answer:\s*([A-Z])", completion, re.IGNORECASE)
+    if match:
+        return completion.strip(), match.group(1).upper()
+    # Fallback: look for the last letter A-Z in the completion
+    letters = re.findall(r'[A-Z]', completion, re.IGNORECASE)
+    return completion.strip(), letters[-1].upper() if letters else ""
+
+async def get_reward_score(
+    reward_sem: asyncio.Semaphore,
+    session: aiohttp.ClientSession,
+    messages: List[Dict[str, str]]
+) -> float:
+    """Get reward model score for a completion."""
+    async def _get_score():
+        async with reward_sem:
+            try:
+                async with session.post(
+                    REWARD_MODEL_ENDPOINT,
+                    json={"messages": messages},
+                    headers={"Content-Type": "application/json"},
+                    timeout=aiohttp.ClientTimeout(total=REWARD_MODEL_TIMEOUT)
+                ) as response:
+                    if response.status != 200:
+                        text = await response.text()
+                        print(f"Error {response.status}: {text}")
+                        raise APIError(f"Reward API returned status {response.status}")
+                    result = await response.json()
+                    return float(result.get('score', 0))
+            except asyncio.TimeoutError:
+                print("Reward model request timed out")
+                raise
+            except Exception as e:
+                print(f"Exception in get_reward_score: {str(e)}")
+                raise
+
+    try:
+        return await with_retry(_get_score, max_retries=REWARD_MODEL_MAX_RETRIES)
+    except Exception:
+        print("All reward score attempts failed")
+        return 0.0
+
+async def verify_answer(
+    llm_sem: asyncio.Semaphore,
+    session: aiohttp.ClientSession,
+    student_answer: str,
+    correct_idx: int
+) -> float:
+    """Verify if the student's answer is correct."""
+    # Convert index to letter (0 -> A, 1 -> B, etc.)
+    correct_letter = chr(65 + correct_idx)  # 65 is ASCII for 'A'
+    return 1.0 if student_answer.upper() == correct_letter else 0.0
+
+async def get_completions(
+    llm_sem: asyncio.Semaphore,
+    reward_sem: asyncio.Semaphore,
+    session: aiohttp.ClientSession,
+    question: str,
+    choices: List[str],
+    n: int
+) -> List[Tuple[str, str, float]]:
+    """Generate n completions and get their reward scores."""
+    # Format question with options
+    formatted_question = f"{question}\n\nOptions:\n"
+    for idx, choice in enumerate(choices):
+        formatted_question += f"{chr(65 + idx)}) {choice}\n"
+
+    print(f"\n[Generating {n} completions for question]")
+    print(f"Q: {formatted_question}")
+
+    USER_PROMPT = ("You are an expert at truthful reasoning and you always pick the most accurate answer. "
+                  "Think step by step and output your reasoning followed by your final answer using the following format:\n"
+                  "Answer: X where X is one of the available letter options.\n\n")
+
+    async def _get_completions():
+        async with llm_sem:
+            messages = [
+                {"role": "user", "content": (
+                    USER_PROMPT +
+                    f"{formatted_question}"
+                )}
+            ]
+
+            async with session.post(
+                FIREWORKS_API_ENDPOINT,
+                headers={
+                    "Accept": "application/json",
+                    "Content-Type": "application/json",
+                    "Authorization": f"Bearer {FIREWORKS_API_KEY}"
+                },
+                json={
+                    "model": MODEL_NAME,
+                    "messages": messages,
+                    "n": n,
+                    "temperature": 0.7,
+                    "max_tokens": 4096,
+                    "top_p": 1,
+                    "top_k": 40,
+                    "presence_penalty": 0,
+                    "frequency_penalty": 0
+                },
+                timeout=aiohttp.ClientTimeout(total=LLM_TIMEOUT)
+            ) as response:
+                if response.status != 200:
+                    text = await response.text()
+                    raise APIError(f"OpenAI API returned status {response.status}: {text}")
+
+                return await response.json()
+
+    try:
+        result = await with_retry(_get_completions)
+        completions = []
+
+        # Get reward scores for each completion
+        for choice in result["choices"]:
+            full_completion, extracted_answer = extract_answer(
+                choice["message"]["content"].strip()
+            )
+
+            # Get reward score
+            reward_score = await get_reward_score(
+                reward_sem,
+                session,
+                [
+                    {"role": "user", "content": USER_PROMPT + formatted_question},
+                    {"role": "assistant", "content": full_completion}
+                ]
+            )
+
+            completions.append((full_completion, extracted_answer, reward_score))
+
+        # Log results
+        print("\n[Completion Results]")
+        for i, (_, answer, score) in enumerate(completions, 1):
+            print(f"  {i}. {answer:<40} [reward: {score:.4f}]")
+
+        return completions
+
+    except Exception as e:
+        print(f"Exception in get_completions: {str(e)}")
+        return [("", "", 0.0)] * n
+
+async def evaluate_question(
+    llm_sem: asyncio.Semaphore,
+    reward_sem: asyncio.Semaphore,
+    session: aiohttp.ClientSession,
+    example
+) -> Dict[str, Any]:
+    """Evaluate a single question with best-of-n completions."""
+    question = example['question']
+    mc1_targets = example['mc1_targets']
+    choices = mc1_targets['choices']
+    correct_idx = mc1_targets['labels'].index(1)  # Find index where label is 1
+
+    # Get n completions with reasoning, extracted answers, and reward scores
+    completion_data = await get_completions(llm_sem, reward_sem, session, question, choices, BEST_OF_N)
+    completions, extracted_answers, reward_scores = zip(*completion_data)
+
+    # Use reward scores to pick the best completion
+    best_idx = reward_scores.index(max(reward_scores))
+    best_completion = completions[best_idx]
+    best_extracted = extracted_answers[best_idx]
+
+    # Verify correctness of the best answer
+    correctness_score = await verify_answer(
+        llm_sem, session, best_extracted, correct_idx
+    )
+
+    return {
+        'question': question,
+        'choices': choices,
+        'correct_answer': chr(65 + correct_idx),  # Convert index to letter
+        'completions': completions,
+        'extracted_answers': extracted_answers,
+        'reward_scores': reward_scores,
+        'best_reward_score': reward_scores[best_idx],
+        'best_completion': best_completion,
+        'best_extracted_answer': best_extracted,
+        'is_correct': bool(correctness_score)
+    }
+
+async def process_batch(
+    llm_sem: asyncio.Semaphore,
+    reward_sem: asyncio.Semaphore,
+    session: aiohttp.ClientSession,
+    batch_data: List[Dict]
+) -> List[Dict[str, Any]]:
+    """Process a batch of questions."""
+    batch_requests = [
+        evaluate_question(llm_sem, reward_sem, session, example) 
+        for example in batch_data
+    ]
+    return await tqdm_asyncio.gather(*batch_requests)
+
+async def evaluate_all(session: aiohttp.ClientSession, dataset) -> List[Dict[str, Any]]:
+    """Evaluate all questions in the dataset using batching."""
+    llm_sem = asyncio.Semaphore(LLM_MAX_CONCURRENT)
+    reward_sem = asyncio.Semaphore(REWARD_MAX_CONCURRENT)
+
+    # Convert dataset to list of dictionaries for easier processing
+    dataset_dicts = [
+        {
+            'question': item['question'],
+            'mc1_targets': item['mc1_targets']
+        }
+        for item in dataset
+    ]
+
+    import random
+    random.seed(42)
+    random.shuffle(dataset_dicts)
+    dataset_dicts = dataset_dicts[:100]
+
+    results = []
+    print(f"\nEvaluating {len(dataset_dicts)} questions with {BEST_OF_N} completions each...")
+
+    # Process in batches
+    for i in range(0, len(dataset_dicts), BATCH_SIZE):
+        batch_data = dataset_dicts[i:i + BATCH_SIZE]
+        print(f"\nProcessing batch {i//BATCH_SIZE + 1}/{(len(dataset_dicts) + BATCH_SIZE - 1)//BATCH_SIZE}")
+
+        batch_results = await process_batch(llm_sem, reward_sem, session, batch_data)
+        results.extend(batch_results)
+
+        # Periodic cleanup
+        gc.collect()
+        await asyncio.sleep(1)  # Small delay between batches
+
+    return results
+
+async def main():
+    try:
+        # Load TruthfulQA dataset
+        dataset = load_dataset("truthful_qa", "multiple_choice")
+        validation_set = dataset["validation"]
+        print(f"Loaded {len(validation_set)} questions from TruthfulQA validation set")
+
+        # Configure session with connection pooling
+        connector = aiohttp.TCPConnector(
+            limit=max(LLM_MAX_CONCURRENT, REWARD_MAX_CONCURRENT),
+            force_close=True
+        )
+        timeout = aiohttp.ClientTimeout(total=60)
+
+        # Create timestamp for output file
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+
+        async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session:
+            results = await evaluate_all(session, validation_set)
+
+            if results:
+                print("\nOverall Statistics:")
+                correct_count = sum(1 for r in results if r['is_correct'])
+                total_count = len(results)
+
+                print(f"Total correct answers: {correct_count}/{total_count}")
+                print(f"Accuracy: {correct_count/total_count:.2f}")
+                print(f"Average reward score: {sum(r['best_reward_score'] for r in results)/total_count:.4f}")
+
+                # Save results
+                output_file = f'truthfulqa_mc_results.json'
+                with open(output_file, 'w') as f:
+                    json.dump(results, f, indent=2)
+                print(f"\nDetailed results saved to {output_file}")
+
+    except Exception as e:
+        print(f"Error in main: {str(e)}")
+        raise
+    finally:
+        if 'connector' in locals() and hasattr(connector, 'close'):
+            await connector.close()
+
+if __name__ == "__main__":
+    asyncio.run(main())