GPT-4o의 추론 학습: 파인튜닝 vs. 예산 강제

import json
import random
from datasets import load_dataset

# Load the Hugging Face dataset
dataset = load_dataset("simplescaling/s1K")["train"]

# Shuffle the dataset for randomness
random.seed(42)
dataset = list(dataset)  # Convert to list for shuffling
random.shuffle(dataset)

# # Split into 95% training and 5% validation
split_idx = int(0.95 * len(dataset))
train_data, val_data = dataset[:split_idx], dataset[split_idx:]

# Function to format each example in the agreed multi-turn OpenAI structure
def format_example(example):
    question = example["question"].strip()
    thinking_steps = "\n".join(example["thinking_trajectories"]).strip()
    final_answer = example["attempt"].strip()

    # Ensure final answer is correctly formatted
    final_answer = "Answer: " + final_answer if not final_answer.startswith("Answer:") else final_answer

    # Construct multi-turn format
    formatted_messages = [
        {"role": "system", "content": "You are a helpful AI assistant."},
        {"role": "user", "content": question + "\n Think for up to 8096 tokens."},
        {"role": "assistant", "content": "<imstart>think\n" + thinking_steps + "\n<imend>"},
        {"role": "user", "content": "Now provide the final answer. Respond in under 2048 tokens!!!"},
        {"role": "assistant", "content": "<imstart>answer\n" + final_answer + "\n<imend>"}
    ]

    return {"messages": formatted_messages}

# Apply formatting to training and validation sets
formatted_train_data = [format_example(example) for example in train_data]
formatted_val_data = [format_example(example) for example in val_data]

# Save as JSONL for OpenAI fine-tuning
train_output_file = "fullnew_s1K_openai_finetune_train.jsonl"
val_output_file = "fullnews1K_openai_finetune_val.jsonl"

# Save training data
with open(train_output_file, "w") as f:
    for entry in formatted_train_data:
        json.dump(entry, f)
        f.write("\n")

# Save validation data
with open(val_output_file, "w") as f:
    for entry in formatted_val_data:
        json.dump(entry, f)
        f.write("\n")

print(f"Training dataset saved as {train_output_file}")
print(f"Validation dataset saved as {val_output_file}")


{
  "messages": [
    {
      "role": "system",
      "content": "You are a helpful AI assistant."
    },
    {
      "role": "user",
      "content": "What is the sum of the first 100 positive integers? Think for up to 2048 tokens."
    },
    {
      "role": "assistant",
      "content": "<imstart>think\n{the reasoning trace}\n<imend>"
    },
    {
      "role": "user",
      "content": "Now provide the final answer. Respond in under 256 tokens."
    },
    {
      "role": "assistant",
      "content": "<imstart>answer\n{the answer}\n<imend>"
    }
  ]
}

import os
import asyncio
import json
from datasets import load_dataset
from openai import OpenAI
import weave; weave.init("aime_evaluation")

# Initialize OpenAI client
openai_client = OpenAI(api_key="your api key")
# Model constants
MODEL_BASE = "gpt-4o-2024-08-06"  # Stock model
MODEL_FINETUNED = "ft: my fine tuned model id"  # Fine-tuned model
JUDGE_MODEL = "gpt-4o-2024-08-06"

class BaseGPTModel(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        """Run inference using stock GPT-4o model"""
        try:
            response = openai_client.chat.completions.create(
                model=MODEL_BASE,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": f"{text}\nThink for up to 8096 tokens."}
                ],
                temperature=0.0
            )
            return response.choices[0].message.content
        except Exception as e:
            print(f"Failed to get base model response: {e}")
            return None

class FineTunedGPTModel(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        """Run inference using fine-tuned GPT-4o model"""
        try:
            # First, get thinking steps
            think_response = openai_client.chat.completions.create(
                model=MODEL_FINETUNED,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": f"{text}\nThink for up to 8096 tokens."}
                ],
                temperature=1.0
            )
            thinking_steps = think_response.choices[0].message.content
            
            # Then, get final answer
            final_response = openai_client.chat.completions.create(
                model=MODEL_FINETUNED,
                messages=[
                    {"role": "system", "content": "You are a helpful assistant."},
                    {"role": "user", "content": text},
                    {"role": "assistant", "content": f"<imstart>think\n{thinking_steps}\n<imend>"},
                    {"role": "user", "content": "Now provide the final answer. Respond in under 2048 tokens!!!"}
                ],
                temperature=1.0
            )
            final_answer = final_response.choices[0].message.content
            
            # Combine thinking and answer
            full_response = f"""Thinking Steps:
{thinking_steps}

Final Answer:
{final_answer}"""
            
            return full_response
        except Exception as e:
            print(f"Failed to get fine-tuned model response: {e}")
            return None



class BudgetForcingGPTModel(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        """Run inference using budget forcing approach with fine-tuned GPT-4o"""
        try:
            # Constant for number of thinking iterations
            NUM_IGNORE = 6
            
            # Initialize messages
            messages = [
                {"role": "system", "content": "You are a helpful assistant."},
                {"role": "user", "content": f"{text}\nThink for up to 8096 tokens."}
            ]
            
            # Thinking phase with budget forcing
            thinking_trace = ""
            
            for i in range(NUM_IGNORE):
                try:
                    response = openai_client.chat.completions.create(
                        model=MODEL_FINETUNED,
                        messages=messages,
                        max_tokens=8096,
                        temperature=1.0
                    )
                    
                    thinking_text = response.choices[0].message.content
                    thinking_trace += thinking_text + "\nContinuing...\n"
                    
                    messages.append({"role": "assistant", "content": thinking_text + "\nWait? Did I make a mistake? "})
                    messages.append({"role": "user", "content": "Continue your thinking where you left off, correcting any mistakes if there is any. Think for up to 8096 tokens."})
                    
                    await asyncio.sleep(1)
                    
                except Exception as e:
                    print(f"Error during reasoning step {i}: {e}")
                    break
            
            # Final answer phase
            messages.append({"role": "user", "content": "Now provide the final answer. Respond in under 2048 tokens!!!"})
            
            final_response = openai_client.chat.completions.create(
                model=MODEL_FINETUNED,
                messages=messages,
                max_tokens=2048,
                temperature=1.0
            )
            
            final_answer = final_response.choices[0].message.content
            
            # Combine thinking and answer
            full_response = f"""Thinking Steps:
{thinking_trace}

Final Answer:
{final_answer}"""
            
            return full_response
            
        except Exception as e:
            print(f"Failed to get budget forcing model response: {e}")
            return None


@weave.op
async def gpt4o_scorer(label: str, model_output: str) -> dict:
    """Score the model's output by comparing it with the ground truth."""
    try:
        # Extract the final answer section
        final_answer = (model_output.split("Final Answer:")[-1].strip() if "Final Answer:" in model_output else model_output)[-100:]
        
        query = f"""
        YOU ARE A LLM JUDGE DETERMINING IF THE FOLLOWING MODEL GENERATED ANSWER IS THE SAME AS THE CORRECT ANSWER 

        Model's Answer: {final_answer}
        Correct Answer: {label}
        
        Your task:
        1. State the model's predicted answer (answer only).
        2. State the ground truth (answer only).
        3. Determine if the model's final answer is correct (ignore formatting differences, etc.). RESPOND with the predicted and ground truth answer, followed with a JSON object containing the correctness encapsulated within the following delimiters:
           ```json
           {{ "correctness": true/false }}
           ```
        """
        
        response = openai_client.chat.completions.create(
            model=JUDGE_MODEL,
            messages=[{"role": "user", "content": query}]
        )
        response_text = response.choices[0].message.content
        
        json_start = response_text.index("```json") + 7
        json_end = response_text.index("```", json_start)
        correctness = json.loads(response_text[json_start:json_end].strip()).get("correctness", False)
        
        return {"correctness": correctness, "reasoning": response_text}
    except Exception as e:
        print(f"Scoring failed: {e}")
        return {"correctness": False, "reasoning": str(e)}

def load_ds():
    print("Loading AIME dataset...")
    dataset = load_dataset("Maxwell-Jia/AIME_2024")["train"]
    return [{"text": row["Problem"], "label": row["Answer"]} for row in dataset]

async def run_evaluations():
    print("Loading dataset...")
    dataset = load_ds()
    print("Preparing dataset for evaluation...")
    dataset_prepared = [{"text": row["text"], "label": row["label"]} for row in dataset]

    scorers = [gpt4o_scorer]
    

    # Evaluate fine-tuned GPT-4o
    print("\nEvaluating fine-tuned GPT-4o...")
    finetuned_evaluation = weave.Evaluation(
        dataset=dataset_prepared,
        scorers=scorers,
        name="finetuned-gpt4o Evaluation"
    )
    finetuned_results = await finetuned_evaluation.evaluate(FineTunedGPTModel())
    print(f"Results for fine-tuned GPT-4o: {finetuned_results}")

    # Evaluate budget forcing GPT-4o
    print("\nEvaluating budget forcing GPT-4o...")
    budget_evaluation = weave.Evaluation(
        dataset=dataset_prepared,
        scorers=scorers,
        name="budget-forcing-gpt4o Evaluation"
    )
    budget_results = await budget_evaluation.evaluate(BudgetForcingGPTModel())
    print(f"Results for budget forcing GPT-4o: {budget_results}")

    # # Evaluate base GPT-4o
    print("\nEvaluating base GPT-4o...")
    base_evaluation = weave.Evaluation(
        dataset=dataset_prepared,
        scorers=scorers,
        name="base-gpt4o Evaluation"
    )
    base_results = await base_evaluation.evaluate(BaseGPTModel())
    print(f"Results for base GPT-4o: {base_results}")


    # Return all results
    return {
        "base": base_results,
        "finetuned": finetuned_results,
        "budget_forcing": budget_results
    }


if __name__ == "__main__":
    asyncio.run(run_evaluations())

"Wait? Did I make a mistake?"

"Continue your thinking where you left off, correcting any mistakes if there is any. Think for up to 8096 tokens."

"Now provide the final answer. Respond in under 2048 tokens!!!"