#!/usr/bin/env python3
"""
VeRL 모델 생성 테스트
데이터를 딕셔너리 리스트로 변환하여 정상적인 응답이 나오는지 확인
"""
import os
import sys
import pandas as pd
import numpy as np
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

print("=" * 80)
print("VeRL 모델 생성 테스트")
print("=" * 80)

# GPU 설정
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Device: {device}")

# 1. 모델과 토크나이저 로드
print("\n모델 로딩 중...")
model_name = "Qwen/Qwen2.5-7B"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)
print(f"모델 로드 완료: {model_name}")

# 2. 데이터 로드
data_path = "/home/ubuntu/RLVR/TestTime-RLVR-v2/tmp/batch_results/ttrlvr_azr_unified_20250822_151912/mbpp/Mbpp_2/round_1/azr_training_data/induction.parquet"
df = pd.read_parquet(data_path)
print(f"\n데이터 로드: {len(df)} 샘플")

# 3. 테스트할 샘플들
test_samples = []

# 샘플 1: 실제 데이터에서 첫 번째 샘플 (처음 500자만 사용)
prompt_str = df.iloc[0]['prompt']
if isinstance(prompt_str, str):
    # 너무 길면 잘라서 사용
    truncated_prompt = prompt_str[:1000] + "\n\nAssistant:"
    prompt_dict = [{"role": "user", "content": truncated_prompt}]
else:
    prompt_dict = prompt_str
test_samples.append(("실제 데이터 샘플 (truncated)", prompt_dict))

# 샘플 2: 간단한 코딩 문제
simple_prompt = [{
    "role": "user",
    "content": "Write a Python function to calculate the factorial of a number."
}]
test_samples.append(("간단한 코딩 문제", simple_prompt))

# 샘플 3: AZR 스타일 프롬프트
azr_style_prompt = [{
    "role": "user",
    "content": """Write a function that takes a list of numbers and returns the sum of all even numbers.

def sum_even_numbers(numbers):"""
}]
test_samples.append(("AZR 스타일", azr_style_prompt))

# 4. 각 샘플에 대해 생성 테스트
print("\n" + "=" * 80)
print("생성 테스트 시작")
print("=" * 80)

for i, (name, prompt_dict) in enumerate(test_samples, 1):
    print(f"\n[테스트 {i}] {name}")
    print("-" * 60)
    
    # Chat template 적용
    prompt_with_template = tokenizer.apply_chat_template(
        prompt_dict,
        add_generation_prompt=True,
        tokenize=False
    )
    
    print(f"템플릿 적용 후 시작: {repr(prompt_with_template[:100])}...")
    
    # 토큰화
    inputs = tokenizer(
        prompt_with_template,
        return_tensors="pt",
        truncation=True,
        max_length=1024
    ).to(device)
    
    print(f"입력 토큰 수: {inputs['input_ids'].shape[1]}")
    
    # 생성
    with torch.no_grad():
        outputs = model.generate(
            **inputs,
            max_new_tokens=150,
            temperature=0.7,
            do_sample=True,
            top_p=0.95,
            pad_token_id=tokenizer.pad_token_id,
            eos_token_id=tokenizer.eos_token_id
        )
    
    # 디코딩
    generated = outputs[0][inputs['input_ids'].shape[1]:]
    response = tokenizer.decode(generated, skip_special_tokens=True)
    
    print(f"\n생성된 응답:")
    print(">" * 40)
    print(response[:500])  # 처음 500자만 출력
    print("<" * 40)
    
    # 응답 품질 체크
    if any(keyword in response.lower() for keyword in ['def ', 'return', 'function', 'python', '```']):
        print("✅ 코드 관련 응답 생성됨")
    elif any(keyword in response.lower() for keyword in ['stravinsky', 'department', 'openstring', '中文']):
        print("❌ 이상한 응답 생성됨")
    else:
        print("❓ 응답 타입 불명확")

print("\n" + "=" * 80)
print("테스트 완료")
print("=" * 80)

# 5. 결론
print("\n결론:")
print("-" * 60)
print("위 테스트에서 정상적인 코드가 생성되면:")
print("  → 딕셔너리 리스트 형식이 올바르게 작동함")
print("  → complete_pipeline.py 수정이 제대로 적용되지 않은 것이 문제")
print("\n이상한 응답이 계속 생성되면:")
print("  → 다른 근본적인 문제가 있을 수 있음")