from model import TransformerConfig, TransformerLanguageModel
from tokenizer import load_tokenizer
import torch
from torch.utils.data import DataLoader
from dataset_loader import MultiSourceDatasetV2
import random
from tqdm import tqdm

# 模型参数
config = TransformerConfig(
    50304, # vocab_size
    1024, # block_size
    768, # n_embed
    12,  # n_heads
    12,  # n_layers
    0.0, # dropout
    True # bias
)


# 训练参数
batch_size = 8
max_iters = 150000
gradient_accumulation_steps = 5
eval_interval = 100
save_interval = 500
learning_rate = 1e-4
device = 'cuda:0' # if torch.cuda.is_available() else 'cpu'

# 建立模型
model = TransformerLanguageModel(config)
model = model.to(device)

ckpt_id = 43000
checkpoint = f"checkpoints/new/{ckpt_id}.pt"
model.load_state_dict(torch.load(checkpoint))

# 加载分词器
tokenizer = load_tokenizer("tokenizer.model")

# 数据加载
recipe_files = [
    [f"data/enwiki/enwiki-{page}.jsonl" for page in range(6400)],
    [f"data/fineweb/fineweb-{page}.jsonl" for page in range(14850)],
    [f"data/zhwiki/zhwiki-{page}.jsonl" for page in range(1350)],
    [f"data/zhihu/zhihu-{page}.jsonl" for page in range(975)],
    [f"data/allnovels-split/ans-{page}.jsonl" for page in range(1330)]
]
probs = [
    0.2,
    0.3,
    0.2,
    0.1,
    0.2
]

# 建立数据加载器
ds = MultiSourceDatasetV2(recipe_files, probs)
loader = DataLoader(ds, batch_size)

# 建立数据处理函数
# 1. 使用tokenizer转化为整数id
# 2. 添加eos token
# 3. 对超出长度限制+1的数据进行随机截取
# 4. 计算最大长度
# 4. 对不足最大长度的数据用pad token(此处等于eos token)补足
# 5. 合成一个int64格式的tensor ids, 形状为(B,T+1), 使用ids[:,:-1]和ids[:,1:]作为x,y
def get_input_ids(text_batch, eos_token_id=50303, block_size=config.block_size):
    texts = text_batch["text"]
    ids = [tokenizer.encode(text) + [eos_token_id] for text in texts]
    for i in range(len(ids)):
        if len(ids[i]) > block_size+1:
            start = random.randint(0,len(ids[i])-100)
            ids[i] = ids[i][start:start+block_size+1]
    max_len = max([len(item) for item in ids])
    ids = [item + [eos_token_id] * (max_len - len(item)) for item in ids]
    ids = torch.tensor(ids, dtype=torch.int64)
    return ids[:,:-1],ids[:,1:]

# 建立优化器
optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate)
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 1000, 2, 5e-7)

# 文本生成测试函数
@torch.no_grad()
def gen_text(text):
    model.eval()
    ids = torch.tensor(tokenizer.encode(text)).to(device).view(1,-1)
    output_ids = model.generate(ids)[0,:]
    model.train()
    return tokenizer.decode(output_ids.tolist())[0]

# # 进行训练
# # 初始化进度条
# pbar = tqdm(total=max_iters+1)
# # 初始化每步的loss
# all_loss = 0.0
# # 初始化梯度累加
# grad_steps = 0

# # 使用数据加载器获取一个新的数据batch
# for iter_num, (x,y) in enumerate(loader):
#     # x,y = get_input_ids(batch)

#     # 每隔eval_interval轮检查模型生成效果，每隔save_interval保存一次
#     steps = iter_num // gradient_accumulation_steps
#     if iter_num % gradient_accumulation_steps == 0 and (steps % save_interval == 0 or steps == max_iters):
#         print(gen_text("I love you, "))
#         torch.save(model.state_dict(),f'checkpoints/mixed/mixed-{steps}.pt')
#         print(f"Step {steps} saved.")

#     # 调用模型计算logits和loss
#     _, loss = model(x.to(device), targets = y.to(device), device=device)
#     loss = loss / gradient_accumulation_steps

#     # 反向传播计算梯度
#     loss.backward()
#     grad_steps += 1
#     all_loss += loss.item()

#     # 到达梯度累加步数以后更新参数
#     if grad_steps >= gradient_accumulation_steps:

#         # 更新参数
#         optimizer.step()

#         # 梯度归零
#         optimizer.zero_grad(set_to_none=True)

#         # 重置梯度累加步数
#         grad_steps = 0

#         # 更新进度条
#         pbar.update()

#         # 每轮输出一次loss
#         print(f"\nLoss: {all_loss}")

#         # 重置loss
#         all_loss = 0.0

#     # 达到步数以后结束训练
#     if iter_num == max_iters * gradient_accumulation_steps:
#         break

# 进行训练
ds_iter = iter(loader)
for iter in tqdm(range(max_iters+1)):
    if iter < ckpt_id:
        continue
    all_loss = 0.0
    # 梯度归零
    optimizer.zero_grad(set_to_none=True)
    for _ in range(gradient_accumulation_steps):
        # 使用数据加载器获取一个新的数据batch
        x, y = next(ds_iter)
        # 调用模型计算logits和loss
        logits,loss = model(x.to(device), y.to(device), device=device)
        loss = loss / gradient_accumulation_steps
        all_loss += loss.item()
        # 反向传播计算梯度
        loss.backward()
    # 更新参数
    optimizer.step()
    scheduler.step()
    # 每隔save_interval保存一次
    if iter % save_interval == 0 or iter == max_iters:
        torch.save(model.state_dict(),f'checkpoints/new/{iter}.pt')
        print(f"Step {iter} saved.")
    # 每隔eval_iter步评估一次
    if iter % eval_interval == 0 or iter == max_iters:
        print(f"Step: {iter}, Loss: {all_loss}")
        print(gen_text("我喜欢你，"))