import torch

from itertools import chain
import wandb

from datasets import load_dataset, load_from_disk
from transformers import DataCollatorForLanguageModeling
from transformers import GPT2LMHeadModel, GPT2Tokenizer, GPT2Config
from transformers import TrainingArguments, Trainer

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# ds = load_dataset("bookcorpus/bookcorpus", "default")
ds = load_from_disk("bookcorpus")
ds = ds["train"].select(range(10000))
ds = ds.train_test_split(test_size=0.1)

tokenizer = GPT2Tokenizer.from_pretrained("openai-community/gpt2")
tokenizer.pad_token = tokenizer.eos_token

def tokenize_fn(data):
    return tokenizer(text=data["text"])


tokenized_ds = ds.map(tokenize_fn, batched=True, remove_columns="text")
# tokenized_ds.save_to_disk("bookcorpus/tokenized_ds")
# tokenized_ds = load_from_disk("bookcorpus/tokenized_ds")

data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

config = GPT2Config()
model = GPT2LMHeadModel(config).to(device)

training_args = TrainingArguments(
    output_dir="checkpoints",                 # 模型检查点保存目录
    run_name="baseline",                      # WandB 运行名称
    eval_strategy="steps",                    # 评估策略
    eval_steps=50,                            # 每 50 步评估一次
    logging_strategy="steps",                 # 日志记录策略
    logging_steps=50,                         # 每 50 步记录一次日志
    num_train_epochs=3,                       # 训练 3 个epoch
    per_device_train_batch_size=32,           # 每个设备的训练批次大小
    per_device_eval_batch_size=32,            # 每个设备的评估批次大小
    learning_rate=2.5e-4,                     # 学习率
    weight_decay=1e-4,                        # 权重衰减
    report_to="wandb",                        # 报告到WandB
    lr_scheduler_type="cosine",               # 学习率调度器
    warmup_ratio=0.05,                        # 预热比例
)

trainer.train()

+ from accelerate import Accelerator
+ accelerator = Accelerator(mixed_precision="fp16")

- device = "cuda"
+ device = accelerator.device
  model.to(device)

model, optimizer, training_dataloader, scheduler = accelerator.prepare(
    model, optimizer, training_dataloader, scheduler
)

- inputs = inputs.to(device)
- targets = targets.to(device)
  outputs = model(inputs)
  loss = loss_function(outputs, targets)
- loss.backward()
+ accelerator.backward(loss)

import torch
import re
from tqdm.notebook import tqdm
from time import time


device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from datasets import load_dataset, load_from_disk

# ds = load_dataset("openai/MMMLU", "ZH_CN")["test"]
# ds.save_to_disk("MMMLU")
ds = load_from_disk("MMMLU")

filtered_ds = ds.filter(lambda example: example["Subject"] == "moral_scenarios").select(range(100))

from transformers import AutoTokenizer, AutoModelForCausalLM

# model_name = "Qwen/Qwen2.5-0.5B"
model_name = "Qwen"

tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype="auto", device_map="auto").to(device)

def create_prompt(example):
    return {"prompt": f"问题: {example['Question']}\nA) {example['A']}\nB) {example['B']}\nC) {example['C']}\nD) {example['D']}\n请你给出答案（只输出一个字母）: ", "answer": example["Answer"]}


prompts = filtered_ds.map(create_prompt)
total = len(prompts)
correct = 0

start_time = time()
for i in tqdm(range(0, total, batch_size)):
    text_batch = prompts[i : i + batch_size]["prompt"]
    model_inputs_batch = tokenizer(text_batch, return_tensors="pt", padding=True).to(model.device)

    generated_ids_batch = model.generate(**model_inputs_batch, max_new_tokens=512, pad_token_id=tokenizer.eos_token_id)
    generated_ids_batch = generated_ids_batch[:, model_inputs_batch.input_ids.shape[1] :]
    response_batch = tokenizer.batch_decode(generated_ids_batch, skip_special_tokens=True)

    responses = [response.strip()[-1] for response in response_batch]

    correct += sum(response == answer for response, answer in zip(responses, prompts[i : i + batch_size]["answer"]))

end_time = time()
execution_time = end_time - start_time
accuracy = correct / total
print(f"Execution time: {execution_time:.2f} seconds")
print(f"Accuracy: {accuracy:.2f}")

pipe = pipeline("text-generation", model_name, torch_dtype=torch.bfloat16, device_map="auto")
pipe.tokenizer.padding_side = "left"

results = pipe(prompts["prompt"], max_new_tokens=512, batch_size=batch_size)

from accelerate import infer_auto_device_map, dispatch_model


model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
)

model = dispatch_model(model, device_map=infer_auto_device_map(model))

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",
)

from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
from accelerate import infer_auto_device_map, dispatch_model

model_name = "Qwen"
quantization_config = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16)

model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
    quantization_config=quantization_config,
)

model = dispatch_model(model, device_map=infer_auto_device_map(model))

tokenizer = AutoTokenizer.from_pretrained(model_name, padding_side="left")

batch_size = 100