from transformers import (
    AutoConfig,
    AutoTokenizer,
    BitsAndBytesConfig,
    AutoProcessor, 
    LlamaForCausalLM,
    MllamaForConditionalGeneration,
    AutoModelForCausalLM
)
import torch
from peft import PeftModel
from datasets import load_from_disk
import pandas as pd
from tqdm import tqdm
from torch.utils.data import DataLoader


mode_path = '/gemini/pretrain/meta-llamaLlama-3.2-11B-Vision-Instruct'
lora_path = '/gemini/code/FMD/model/final_model_4/checkpoint-2440' # lora 输出对应 checkpoint 路径

# 加载tokenizer
tokenizer = AutoTokenizer.from_pretrained(mode_path, trust_remote_code=True)

# 加载模型
model = MllamaForConditionalGeneration.from_pretrained(mode_path, device_map="auto",torch_dtype=torch.bfloat16, trust_remote_code=True).eval()

# 加载lora权重
model = PeftModel.from_pretrained(model, model_id=lora_path)
test_dataset = load_from_disk("/gemini/code/FMD/final_dataset/Test")
results = []
with torch.no_grad():
    for data in tqdm(test_dataset):
        model_input = tokenizer(
            data['instruction_1'],                  # 输入文本
            add_special_tokens=False,     # 不添加特殊标记
            truncation=True,              # 启用截断
            max_length=3000              # 设置最大长度
        )
        model_input = tokenizer.decode(model_input["input_ids"], skip_special_tokens=False)
        
        model_inputs = tokenizer(f"<|begin_of_text|><|start_header_id|>system<|end_header_id|>\n\nYou are an expert in financial misinformation detection.<|eot_id|><|start_header_id|>user<|end_header_id|>\n\n{model_input}\nimage information: {data['image_info']}<|eot_id|><|start_header_id|>assistant<|end_header_id|>\n\n", truncation=True, max_length=3600, add_special_tokens=False,return_tensors="pt").to('cuda')
        # 生成模型输出
        generated_ids = model.generate(**model_inputs, max_new_tokens=1024)
        
        # 去除输入部分的 token，以保留生成的预测结果
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]
        
        # 解码生成的预测结果
        responses = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
        print(responses)
        # 将每个结果按顺序存储到列表中
        results.append({
            "ID": data['ID'],
            "response": responses
        })
if results:
    df = pd.DataFrame(results)
    output_csv = "/gemini/code/FMD/inference/result_final_model_4/response.csv"
    df.to_csv(output_csv, index=False, encoding='utf-8')
    print(f"Results saved to {output_csv}")