# %%
from datasets import Dataset, load_dataset
import pandas as pd
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer
from peft import LoraConfig, TaskType, get_peft_model, AutoPeftModelForCausalLM

# %%
df = pd.read_csv('data/riddles_data.csv')
df =df.sample(frac = 1)
#df = df[:1000]

# %%
df.describe()

# %%
ds = Dataset.from_pandas(df)

# %%
ds[:3]

# %%
llm_model_name="Qwen/Qwen1.5-0.5B-Chat"
model = AutoModelForCausalLM.from_pretrained(llm_model_name)
tokenizer = AutoTokenizer.from_pretrained(llm_model_name,trust_remote_code=True, pad_token='<|endoftext|>')

tokenizer


# %%
def process_func(example):
    MAX_LENGTH = 512 
    input_ids, attention_mask, labels = [], [], []
    instruction = tokenizer(f"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n<|im_start|>user\n猜谜语：\n谜面：{example['riddle']}\n\n谜底是什么？<|im_end|>\n<|im_start|>assistant\n", add_special_tokens=False)  # add_special_tokens 不在开头加 special_tokens
    response = tokenizer(f"谜底是：{example['label']}", 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]
    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]
        print (f"{tokenizer.decode(input_ids)} Too Long")
    return {
        "input_ids": input_ids,
        "attention_mask": attention_mask,
        "labels": labels
    }

# %%
tokenized_id = ds.map(process_func, remove_columns=ds.column_names)
tokenized_id

# %%
tokenizer.decode(tokenized_id[0]['input_ids'])

# %%
tokenizer.decode(list(filter(lambda x: x != -100, tokenized_id[1]["labels"])))

# %%
config = LoraConfig(
    task_type=TaskType.CAUSAL_LM, 
    target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
    #target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
    inference_mode=False,
    r=32,
    lora_alpha=32,
    lora_dropout=0.05
)


# %%
model = get_peft_model(model, config)
config

# %%
model.print_trainable_parameters()

# %%
args = TrainingArguments(
    output_dir="./Qwen1.5_0.5B_Chat_sft_full",
    logging_steps=10,
    num_train_epochs=2,
    save_steps=10,
    learning_rate=1e-4,
    save_on_each_node=True,
    fp16=False
)

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

trainer.train(resume_from_checkpoint=True)

# %%
trainer.save_model("./qwen_sft_full")

# %%
llm_model_name="Qwen/Qwen1.5-0.5B-Chat"

#model = AutoModelForCausalLM.from_pretrained(llm_model_name)
# # Load PEFT model on CPU
model = AutoPeftModelForCausalLM.from_pretrained(
    "Qwen1.5_0.5B_Chat_sft_full_ckpt_200_ok/checkpoint-210",
    #low_cpu_mem_usage=True,
)
# # Merge LoRA and base model and save
#merged_model = model.merge_and_unload()
#merged_model.save_pretrained("./qwen_sft",safe_serialization=False, max_shard_size="2GB")

tokenizer = AutoTokenizer.from_pretrained(llm_model_name,trust_remote_code=True, pad_token='<|endoftext|>')


# %%
prompt = "谜面：一生受用（猜一字）\n谜底是什么？请解释。"
messages = [
    {"role": "system", "content": "You are a helpful assistant."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)

print(text)
model_inputs = tokenizer([text], return_tensors="pt").to("cpu")

generated_ids = model.generate(
        model_inputs.input_ids,
        max_new_tokens=128,
        do_sample=False,
        top_p=0.0
)

generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

# %%
response 

# %%