예산 강제 s1-32B: 기다리기만 하면 충분할까?

import os
from litellm import completion
from datasets import load_dataset
import asyncio
import weave; weave.init("dataset_comparison")
import time 
# Set OpenAI API key
os.environ["OPENAI_API_KEY"] = ""


def load_datasets():
    print("Loading datasets...")
    aime_dataset = load_dataset("Maxwell-Jia/AIME_2024")["train"]
    s1k_dataset = load_dataset("simplescaling/s1K")["train"]
    return aime_dataset, s1k_dataset

@weave.op
async def compare_questions(q1: str, q2: str) -> str:
    prompt = f"""Compare these two math problems and respond with ONLY the word 'same' if they are the same question (ignoring formatting differences or slight numerical differences/wording differences) or 'different' if they are different questions.

Question 1: {q1}
Question 2: {q2}"""

    response = completion(
        model="openai/gpt-4o-mini",
        messages=[{"content": prompt, "role": "user"}],
        max_tokens=10,
        temperature=0.1
    )
    return 'same' in response.choices[0].message.content.strip().lower()

async def run_comparison():
    aime_dataset, s1k_dataset = load_datasets()
    
    print("Starting comparisons...")
    for aime_q in aime_dataset:
        for s1k_q in s1k_dataset:
            result = await compare_questions(aime_q["Problem"], s1k_q["question"])

            if result:
                print(f"\nFound matching question!")
                print(f"AIME: {aime_q['Problem'][:100]}...")
                print(f"S1K: {s1k_q['question'][:100]}...")
                print("---")

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


import os
import asyncio
import json
from datasets import load_dataset
from vllm import LLM, SamplingParams
from transformers import AutoTokenizer
from openai import OpenAI
import time


import weave; weave.init("aime_evaluation")

# Initialize OpenAI client
openai_client = OpenAI(api_key="your api key")

# Model constants
MAX_TOKENS_THINKING = 132000
NUM_IGNORE = 6
MODEL_NAME = "simplescaling/s1-32B"
JUDGE_MODEL = "gpt-4o-2024-08-06"


# Initialize model and tokenizer globally
model = LLM(
    "Qwen/Qwen2.5-32B-Instruct",
    tensor_parallel_size=2,
)
tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-32B-Instruct")

class QwenBaseModel(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        """Run inference using base Qwen model without Wait iterations"""
        try:
            # Prepare the prompt
            prompt = f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
            
            # Single thinking phase
            stop_token_ids = tokenizer("<|im_start|><|im_end|>")["input_ids"]
            sampling_params = SamplingParams(
                max_tokens=MAX_TOKENS_THINKING,
                min_tokens=0,
                stop_token_ids=stop_token_ids,
                skip_special_tokens=False,
                temperature=0.0,
            )
            
            prompt += "<|im_start|>"
            output = model.generate(prompt, sampling_params=sampling_params)
            
            # Final answer phase
            prompt += output[0].outputs[0].text
            stop_token_ids = tokenizer("<|im_end|>")["input_ids"]
            sampling_params = SamplingParams(
                max_tokens=32768,
                min_tokens=0,
                stop_token_ids=stop_token_ids,
                skip_special_tokens=False,
                temperature=0.0,
            )
            
            final_output = model.generate(prompt, sampling_params=sampling_params)
            return final_output[0].outputs[0].text
            
        except Exception as e:
            print(f"Failed to get Qwen response: {e}")
            return None


class S1ModelWithWait(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        """Run inference using S1 model with Wait iterations"""
        try:
            # Prepare the prompt
            prompt = f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
            
            # First thinking phase
            stop_token_ids = tokenizer("<|im_start|><|im_end|>")["input_ids"]
            sampling_params = SamplingParams(
                max_tokens=MAX_TOKENS_THINKING,
                min_tokens=0,
                stop_token_ids=stop_token_ids,
                skip_special_tokens=False,
                temperature=0.0,
            )
            
            prompt += "<|im_start|>think"
            output = model.generate(prompt, sampling_params=sampling_params)
            
            # Handle multiple thinking iterations
            ignore_str = "Wait"
            max_tokens_thinking_tmp = MAX_TOKENS_THINKING
            for i in range(NUM_IGNORE):
                tokens_used = len(output[0].outputs[0].token_ids)
                max_tokens_thinking_tmp -= tokens_used
                
                if max_tokens_thinking_tmp < 1000:  # Safe buffer
                    break
                    
                prompt += output[0].outputs[0].text + ignore_str
                sampling_params = SamplingParams(
                    max_tokens=max_tokens_thinking_tmp,
                    min_tokens=1,
                    stop_token_ids=stop_token_ids,
                    skip_special_tokens=False,
                    temperature=0.0,
                )
                output = model.generate(prompt, sampling_params=sampling_params)
                
                if i > 0 and output[0].outputs[0].text in prompt:
                    break
            
            # Final answer phase
            prompt += output[0].outputs[0].text
            stop_token_ids = tokenizer("<|im_end|>")["input_ids"]
            sampling_params = SamplingParams(
                max_tokens=32768,
                min_tokens=0,
                stop_token_ids=stop_token_ids,
                skip_special_tokens=False,
                temperature=0.0,
            )
            
            final_output = model.generate(prompt, sampling_params=sampling_params)
            return final_output[0].outputs[0].text
            
        except Exception as e:
            print(f"Failed to get S1 response: {e}")
            return None

class S1ModelNoWait(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        """Run inference using S1 model without any Wait iterations"""
        try:
            # Prepare the prompt
            prompt = f"<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n{text}<|im_end|>\n<|im_start|>assistant\n"
            
            # Single thinking phase
            stop_token_ids = tokenizer("<|im_start|><|im_end|>")["input_ids"]
            sampling_params = SamplingParams(
                max_tokens=MAX_TOKENS_THINKING,
                min_tokens=0,
                stop_token_ids=stop_token_ids,
                skip_special_tokens=False,
                temperature=0.0,
            )
            
            prompt += "<|im_start|>think"
            output = model.generate(prompt, sampling_params=sampling_params)
            
            # Final answer phase
            prompt += output[0].outputs[0].text
            stop_token_ids = tokenizer("<|im_end|>")["input_ids"]
            sampling_params = SamplingParams(
                max_tokens=32768,
                min_tokens=0,
                stop_token_ids=stop_token_ids,
                skip_special_tokens=False,
                temperature=0.0,
            )
            
            final_output = model.generate(prompt, sampling_params=sampling_params)
            return final_output[0].outputs[0].text
            
        except Exception as e:
            print(f"Failed to get S1 response: {e}")
            return None

def run_o1_inference(prompt: str) -> str:
    """Run inference using OpenAI o1-preview"""
    try:
        response = openai_client.chat.completions.create(
            model="o1-preview",
            messages=[
                {"role": "user", "content": f"Solve the following problem. put your final answer within \\boxed{{}}:\n{prompt}"}
            ]
        )
        return response.choices[0].message.content
    except Exception as e:
        print(f"Failed to get o1-preview response: {e}")
        return None

class O1PreviewModel(weave.Model):
    @weave.op
    async def predict(self, text: str) -> str:
        print("running o1-preview inference")
        return run_o1_inference(text)

@weave.op
async def gpt4o_scorer(label: str, model_output: str) -> dict:
    """Score the model's output by comparing it with the ground truth."""
    query = f"""
    YOU ARE A LLM JUDGE DETERMINING IF THE FOLLOWING MODEL GENERATED ANSWER IS THE SAME AS THE CORRECT ANSWER 
    I WILL GIVE YOU THE LAST 100 CHARS OF THE MODEL'S REASONING PATH, WHICH WILL CONTAIN THE FINAL ANSWER -> 

    Model's Answer (last 100 chars): {str(model_output)[-100:]}
    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 }}
       ```
    """
    
    try:
        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():
    global model, tokenizer
    print("Loading dataset...")
    dataset = load_ds()
    print("Preparing dataset for evaluation...")
    dataset_prepared = [{"text": row["text"], "label": row["label"]} for row in dataset]

    # First evaluate base Qwen
    print("Starting Qwen evaluation...")
    scorers = [gpt4o_scorer]
    evaluation = weave.Evaluation(
        dataset=dataset_prepared,
        scorers=scorers,
        name="qwen-base Evaluation"
    )
    
    results = await evaluation.evaluate(QwenBaseModel())
    print(f"Results for base Qwen: {results}")

    # Switch to S1 model
    print("\nSwitching to S1-32B model...")
    model = LLM(
        "simplescaling/s1-32B",
        tensor_parallel_size=4,
    )
    tokenizer = AutoTokenizer.from_pretrained("simplescaling/s1-32B")

    # Now evaluate S1 variations and o1
    test_models = {
        "s1-32b-wait": S1ModelWithWait(),
        "s1-32b-nowait": S1ModelNoWait(),
        "o1-preview": O1PreviewModel()
    }

    for model_name, model_instance in test_models.items():
        print(f"\nEvaluating {model_name}...")
        evaluation = weave.Evaluation(
            dataset=dataset_prepared,
            scorers=scorers,
            name=f"{model_name} Evaluation"
        )
        results = await evaluation.evaluate(model_instance)
        print(f"Results for {model_name}: {results}")

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