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
        })
def split_response(text):
    #获取Prediction的内容
    prediction_pattern = r"Prediction:\s*(False|True|NEI)\s*$"
    prediction_match = re.search(prediction_pattern, text, re.MULTILINE)
    if prediction_match:
        prediction = prediction_match.group(1).strip()
    else:
        prediction = 'None'
        print("没有找到匹配的内容")
    #获取Explanation的内容
    explanation_pattern = r"Explanation:\s*(.*)"
    explanation_match = re.search(explanation_pattern, text, re.MULTILINE)
    if explanation_match:
        explanation = explanation_match.group(1).strip()
    else:
        explanation = None  # 如果没有匹配项，设置为 None
    return prediction, explanation

if results:
    df = pd.DataFrame(results)
    
    for index, row in df.iterrows():
        text = row['response']
        prediction, explanation= split_response(text)
        df.at[index, 'Prediction'] = prediction
        df.at[index, 'Explanation'] = explanation

    df['ID'] = df['ID'].str.replace('FMD_test_', '', regex=False)
    df = df.rename(columns={'ID': 'id','Prediction': 'pred','Explanation': 'explanation'})
    df = df.drop('response',axis=1)
    mapping = {
        'False': 0,
        'True': 1,
        'NEI': 2
    }
    df['pred'] = df['pred'].replace(mapping)
    df.to_csv("/gemini/code/FMD/inference/result_final_model_4/result.csv",index = False)