from pathlib import Path
import torch
from transformers import AutoConfig, AutoModelForCausalLM, AutoProcessor


model_path = Path(__file__).parent.resolve()
print(f"Loading model from: {model_path}")

# 加载模型配置
config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
print(f"Model config: {config.model_type}")
print(f"Architecture: {config.architectures}")

# 加载模型（使用 bfloat16 精度和自动设备映射）
print("\nLoading model...")
model = AutoModelForCausalLM.from_pretrained(
    model_path,
    dtype=torch.bfloat16,
    device_map="auto",
    attn_implementation="flash_attention_2",
    trust_remote_code=True
)
print(f"✓ Model loaded successfully!")
print(f"Model type: {type(model).__name__}")
print(f"Model device: {model.device}")

# 加载处理器（tokenizer + image processor）
print("\nLoading processor...")
processor = AutoProcessor.from_pretrained(model_path, trust_remote_code=True)


# ============================================================================
# 测试 1: 纯文本对话
# ============================================================================
print("\n" + "=" * 80)
print("测试 1: 纯文本对话")
print("=" * 80)

text_messages = [
    {
        "role": "user",
        "content": [
            {"type": "text", "text": "你好，请介绍一下自己，包括你的能力和用途。"}
        ]
    }
]

print("\n准备纯文本输入...")
text_inputs = processor.apply_chat_template(
    text_messages,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
    return_tensors="pt"
)
text_inputs = text_inputs.to(model.device)

print(f"Input shape: {text_inputs['input_ids'].shape}")
print(f"Has pixel values: {'pixel_values' in text_inputs}")

print("\n生成纯文本回复...")
with torch.inference_mode():
    text_generated_ids = model.generate(
        **text_inputs,
        max_new_tokens=256,
        do_sample=False,
        temperature=1.0,
    )

text_generated_ids_trimmed = [
    out_ids[len(in_ids):] for in_ids, out_ids in zip(text_inputs.input_ids, text_generated_ids)
]

text_output = processor.batch_decode(
    text_generated_ids_trimmed,
    skip_special_tokens=True,
    clean_up_tokenization_spaces=False
)

print("\n" + "-" * 80)
print("纯文本输出：")
print("-" * 80)
print(text_output[0])
print("-" * 80)
print("\n✅ 纯文本测试完成！")

# ============================================================================
# 测试 2: 图文混合输入
# ============================================================================
print("\n" + "=" * 80)
print("测试 2: 图文混合输入（多模态）")
print("=" * 80)

# 构建对话消息（图文混合输入）
multimodal_messages = [
    {
        "role": "user",
        "content": [
            {"type": "image", "image": "./panda.jpg"},
            # {"type": "image", "image": "./milk.jpeg"},
            {"type": "text", "text": "请描述这张图"},
        ],
    }
]

print("\n准备图文混合输入...")
# 应用对话模板并进行 tokenization
multimodal_inputs = processor.apply_chat_template(
    multimodal_messages,
    tokenize=True,
    add_generation_prompt=True,
    return_dict=True,
    return_tensors="pt"
)
multimodal_inputs = multimodal_inputs.to(model.device)

print(f"Input shape: {multimodal_inputs['input_ids'].shape}")
print(f"Pixel values shape: {multimodal_inputs['pixel_values'].shape if 'pixel_values' in multimodal_inputs else 'N/A'}")

# 生成输出
print("\n生成图像描述...")
with torch.inference_mode():
    multimodal_generated_ids = model.generate(
        **multimodal_inputs,
        max_new_tokens=512,
        do_sample=False,
        temperature=1.0,
    )

# 提取生成的 token（去除输入部分）
multimodal_generated_ids_trimmed = [
    out_ids[len(in_ids):] for in_ids, out_ids in zip(multimodal_inputs.input_ids, multimodal_generated_ids)
]

# 解码为文本
multimodal_output = processor.batch_decode(
    multimodal_generated_ids_trimmed,
    skip_special_tokens=True,
    clean_up_tokenization_spaces=False
)

print("\n" + "-" * 80)
print("图像描述输出：")
print("-" * 80)
print(multimodal_output[0])
print("-" * 80)
print("\n✅ 多模态测试完成！")