import pandas as pd
from datasets import Dataset
from openmind import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer
import torch
from peft import LoraConfig, TaskType, get_peft_model
from transformers import DataCollatorForSeq2Seq
from torch.utils.data import DataLoader
df = pd.read_json('code_alpaca_2k.json')
ds = Dataset.from_pandas(df)
print(len(ds))

model_path = '/tmp/code/NousResearch/Hermes-2-Pro-Mistral-7B'

tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False, trust_remote_code=True)

def process_func(example):
    MAX_LENGTH = 2048    
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer(f"<|im_start|>system\n你是一个人工智能助手，请你根据要求回答问题。<|im_end|>\n<|im_start|>user\n{example['instruction'] + example['input']}<|im_end|>\n<|im_start|>assistant\n", add_special_tokens=False)  # 
    response = tokenizer(f"{example['output']}", add_special_tokens=False)
    input_ids = instruction["input_ids"] + response["input_ids"] + [tokenizer.pad_token_id]
    attention_mask = instruction["attention_mask"] + response["attention_mask"] + [1]  # 因为eos token咱们也是要关注的所以 补充为1
    labels = [-100] * len(instruction["input_ids"]) + response["input_ids"] + [tokenizer.pad_token_id]  
    if len(input_ids) > MAX_LENGTH:  # 做一个截断
        input_ids = input_ids[:MAX_LENGTH]
        attention_mask = attention_mask[:MAX_LENGTH]
        labels = labels[:MAX_LENGTH]
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }
tokenized_id = ds.map(process_func, remove_columns=ds.column_names)
## 加载训练模型
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="auto",torch_dtype=torch.bfloat16)

config = LoraConfig(
    task_type=TaskType.CAUSAL_LM, 
    target_modules=["v_proj","k_proj","gate_proj","q_proj"],
    inference_mode=False, # 训练模式
    r=8, # Lora 秩
    lora_alpha=16, # Lora alaph，具体作用参见 Lora 原理
    lora_dropout=0.1# Dropout 比例
)

model = get_peft_model(model, config)

model.enable_input_require_grads()
lora_path = "./output/deepseek-coder-6.7b-base"

args = TrainingArguments(
    output_dir=lora_path,
    overwrite_output_dir=True,
    per_device_train_batch_size=2,
    gradient_accumulation_steps=1,
    logging_steps=1,
    logging_strategy='steps',
    logging_dir = lora_path,
    num_train_epochs=5,
    logging_nan_inf_filter=True,
    save_steps=500,
    learning_rate=1e-4,
    save_strategy='steps',
    dataloader_num_workers=2,
    gradient_checkpointing=True
)

data_loader = DataLoader(tokenized_id, batch_size=args.per_device_train_batch_size, num_workers=0)

trainer = Trainer(
    model=model,
    args=args,
    train_dataset=tokenized_id,
    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True)
)

trainer.train()

import torch
from openmind import AutoTokenizer, AutoModelForCausalLM,is_torch_npu_available
import argparse

def mean_pooling(model_output, attention_mask):
    token_embeddings = model_output[0] # model_output的第一个元素包含所有token嵌入
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
    return torch.sum(token_embeddings * input_mask_expanded, 1) / torch.clamp(input_mask_expanded.sum(1), min=1e-9)

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_name_or_path",
        type=str,
        help="Path to model",
        default="Rose/llm-jp-3-1.8b-instruct",
    )
    args = parser.parse_args()
    return args
def main():
    args = parse_args()
    model_path = args.model_name_or_path

    if is_torch_npu_available():
        device = "npu:0"
    else:
        device = "cpu"
    device='npu'
    tokenizer = AutoTokenizer.from_pretrained('../', trust_remote_code=True)
    model = AutoModelForCausalLM.from_pretrained(
        "../",
        torch_dtype=torch.float16
    ).to(device)

    prompts = [
        """<|im_start|>system
    You are a sentient, superintelligent artificial general intelligence, here to teach and assist me.<|im_end|>
    <|im_start|>user
    Write a short story about Goku discovering kirby has teamed up with Majin Buu to destroy the world.<|im_end|>
    <|im_start|>assistant""",
        ]

    for chat in prompts:
        print(chat)
        
        # 使用 tokenizer 编码输入文本，并指定返回 PyTorch 张量
        inputs = tokenizer(chat, return_tensors="pt").to(device)
        
        # 将整个 inputs 字典传递给 model.generate() 方法
        generated_ids = model.generate(**inputs, max_new_tokens=750, temperature=0.8, repetition_penalty=1.1, do_sample=True, eos_token_id=tokenizer.eos_token_id)
        
        # 获取 input_ids 的形状以正确切片生成的 tokens
        input_length = inputs['input_ids'].shape[1]
        
        # 解码生成的文本并去除特殊标记
        response = tokenizer.decode(generated_ids[0][input_length:], skip_special_tokens=True, clean_up_tokenization_space=True)
        
        print(f"Response: {response}")


if __name__ == "__main__":
    main()