语言模型的数学能力实验

具体技术来自xval：https://github.com/PolymathicAI/xVal/blob/main/xval_demo.ipynb

语言模型的各项能力各项能力的提升是趋势，月之暗面和claude的超长文本能力，大海捞针。
多模态能力，cot能力，数学能力，代码能力，总结能力， 等等。

其中让我比较关心的是数学能力，举个例子，一个人在路上步行速度为4km/h，然后经给123分钟后，到哪里了？
对于这里的4， 123, 在人类视角里分别代表一个数字, 只有涉及到计算的时候，才单独抽取出来进行计算。
但是当前普遍语言模型的做法，不太符合真正高效的处理。

一个人在路上步行速度为4km/h，然后经给123分钟后，到哪里了？
1 普遍的模型：
-> 【“一个”，“人”，“在”，“路上”，“步行”，“速度”，“为”，“4”，“km/h”，“然后”，“经给“，”123“，“分钟“，”后“，”到哪里”，“了”】

2 更好的模型：
-> 【“一个”，“人”，“在”，“路上”，“步行”，“速度”，“为”，“【num】”，“km/h”，“然后”，“经给“，”【num】“，“分钟“，”后“，”到哪里”，“了”】
-> 【-，-，-，-，-，-，-，4，-，-，-，123，-，-，-，-】
映射到
-> h-embed: 【】
也就是在tokenizer的时候，需要把数字不再认定是离散的单个字符，而是通过一个类似bos，eos的特殊token，命名为num来取代。
此举的目的主要是将离散的数字token变成连续的，更加符合现实情况，可以提高数学能力，我认为数学能力，和对数字的敏感性提高后，必然带来推理能力的提高。

这是xval的基本思想，在tokenize的时候对数字特殊处理，然后使用一个更改过的numformer，是transformer的变体，以此来提高对数字的掌控力，可惜这个xval并没有把事情完整的完成。
基本测试完numformer就完事了。

所以需要进一步进行实验，取phi小模型，分别两组，一个是没有进行numformer改造的，一个是引入numformer改造的。
然后开始训练引入numformer改造的，预料的选择和配比和原来的完全一致。
最后进行测试。

对于xval的了解

文件结构是
xval_demo.ipynb： 主要是使用numformer，和算法的实现，然后进行tokenizer训练，和使用numformer进行seq2seq的训练和预测。
tokenizer.json：tokenizer本体
tokenize-dataset.py：训练tokenizer的数据准备
env.yaml：conda python环境

还有xval包
init.py：初始化文件
analyze.py： ？
make_tokenizer.py: ？
numformer.py: 应该是numformer的实现，transformer的变体。
preprocess.py: 预处理文件

对于xval的了解应该可以进一步学会如何将transformer改造成numformer，如何使用该token训练方式、
也就是xval_demo.ipynb的内容。

• 1 首先导入库

• 2 查看数据内容和结构
  可以看到这些数据就是一叠一叠字符串，这里字符串是json格式的。

  ds['text'][0][:n_characters]

  总共12w5k个

• 3 tokenize（最重要的一环之一）
  因为这里处理的数据比较特殊，是一个有一个json格式的，所以只需要提取key值作为tokens。

  # Building a tokenizer with the extracted keys as individual tokens.
  make_tokenizer.make_tokenizer(
      save_file="./tokenizer.json", efficient_json=True, sample_keys=sample_keys
  )
  ````
  
  这个函数是xval实现的，sample_keys就是数据集合里面所有出现的key值。
  
  ```python
  tokenizer = PreTrainedTokenizerFast(
      tokenizer_file="./tokenizer.json",
      bos_token="[END]", # beginning of sentence
      eos_token="[END]", # end of sentence
      mask_token="[MASK]", # mask token
      pad_token="[PAD]", # pad token
  )
  ````
  就这样一个tokenizer就搞定了。
  
  - 4 接下来使用该tokenizer试试
    上面的tokenizer.json已经出现了【NUM】的token了，所以实际该tokenizer已经可以用了。
  ```python
  tokenized_x = preprocess.tokenize_fnc(ds['text'][0][:100], tokenizer)
  print('input_ids:', tokenized_x['input_ids'])
  print('numbers:', tokenized_x['numbers'])
  ````
  对于一个句子：{'description':{'planet0':{'m':+1.31e+0,'a':+1.94e+0,'e':+1.90e+0},'plane
  会形成两个长度一致数组：
  input_ids:【4, 18, 4, 26, 4, 21,      3,    8,  20,      3,       8, 16,     3,     5, 8】
  numbers: 【1. 1.  1.  1.  1.  1.  1.3125 1.   1.    1.944，1.   1.  1.897  1. 1.】
  其中token_id为3的其实是数字，对应到下面numbers数组上，会有小数存起来。
  
  ```python
  print("\nStarting tokenization...")
  tokenize_lambda = lambda x: preprocess.tokenize_fnc(x, tokenizer)
  tokenized_ds = ds.map(
      tokenize_lambda,
      batched=False,
      num_proc=30,
      remove_columns=["text"],
      load_from_cache_file=False,
  )
  tokenized_ds
  
  输出：
  Dataset({
      features: ['input_ids', 'numbers', 'len'],
      num_rows: 125000
  })

  对125000个json进行处理后得到这么一个tokenize_ds

• 5 对照实验
  后面xval为了对比，从而tokenize方式进行了对比
  P10(一个数字5个token) (词表大小：28)
  P1000(一个数字3个token) (词表大小：918)
  P1999(一个数字2个token) (词表大小：1816)
  FP15(一个数字1个token) (词表大小：28800)
  XVAL(一个数字1个token) (词表大小：1)

• 6 训练

  import torch
  from torch import optim
  from datasets import DatasetDict
  from torch.utils.data import DataLoader
  import torch.nn.functional as F
  from tqdm import tqdm
  import matplotlib.pyplot as plt
  import pandas as pd
  from xval import numformer
  
  ### Define model
  # The vocab_size is the number of different tokens in the tokenizer.
  # context length is the maximum sequence size. 
  model = numformer.Numformer(vocab_size=27, nhead=3, num_layers=3, d_model=384,  dim_feedforward=1536, context_length=955).cuda()
  lr = 1e-4
  weight_decay = 0.01
  optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
  
  ### Load the tokenizer 
  tokenizer_path = "./tokenizer.json"
  tokenizer = PreTrainedTokenizerFast(
      tokenizer_file=tokenizer_path,
      bos_token="[END]",
      eos_token="[END]",
      mask_token="[MASK]",
      pad_token="[PAD]",
  )
  pad_token_id = tokenizer.pad_token_id
  num_token_id = tokenizer.convert_tokens_to_ids("[NUM]")
  mask_token_id = tokenizer.mask_token_id
  mlm_probability = 0.3
  epochs = 10
  
  ### Load tokenized datasets 
  dataset_path = "./data/tokenized_ds_xval"
  tokenized_ds = DatasetDict.load_from_disk(dataset_path)
  
  # Define the masked xVal collator which takes samples of unequal length and masks out both the token_ids and the numbers.
  collator = numformer.define_masked_num_collator(pad_token_id, mask_token_id, mlm_probability)
  
  train_loader = DataLoader(
      tokenized_ds["train"],
      batch_size=32,
      shuffle=True,
      collate_fn=collator,
  )
  
  
  ### Run training loop
  
  loss_hist = []
  loss_mlm_hist = []
  loss_num_hist = []
  
  max_n_batches = 100 # without capping the number of batches, training takes many hours
  
  try: 
      for e in tqdm(range(epochs)):
          n_batches = 0
          for batch in train_loader:
              if n_batches > max_n_batches: 
                  break
              logit_preds, num_preds = model(batch["x"].cuda(), batch["x_num"].cuda())
              with torch.autocast(device_type="cuda"):
                  loss_mlm = F.cross_entropy(
                      logit_preds.view(-1, logit_preds.size(-1)),
                      batch["y"].cuda().view(-1),
                      ignore_index=-100,
                      reduction="mean",
                  )
                  num_mask = batch['y']==num_token_id
                  loss_num = F.mse_loss(
                      num_preds[num_mask],
                      batch["y_num"][num_mask].view(-1,1).cuda(),
                      reduction="mean",
                  )
              loss = loss_mlm + loss_num
              optimizer.zero_grad()
              loss.backward()
              optimizer.step()
              loss_hist.append(loss.item())
              loss_mlm_hist.append(loss_mlm.item())
              loss_num_hist.append(loss_num.item())
              n_batches += 1 
  
              # calculate the running average of the losses
              try:
                  loss_avg = 0.99*loss_avg + 0.01*loss.item()
                  loss_mlm_avg = 0.99*loss_mlm_avg + 0.01*loss_mlm.item()
                  loss_num_avg = 0.99*loss_num_avg + 0.01*loss_num.item()
              except:
                  loss_avg = loss.item()
                  loss_mlm_avg = loss_mlm.item()
                  loss_num_avg = loss_num.item()
          
          ### Save checkpoint at each epoch
          checkpoint = {
                  "model": model.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "loss": loss_avg,
                  "loss_hist": loss_hist,
                  "loss_mlm_hist": loss_mlm_hist,
                  "loss_num_hist": loss_num_hist,
              }
          torch.save(checkpoint, "./ckpt.pt")
          print(f"Epoch #{e}: loss_mlm = {loss_mlm_avg:.3f}; loss_num = {loss_num_avg:.3f}; loss_total = {loss_avg:.3f}")
  except KeyboardInterrupt:
      print('Interrupted')

  这份训练代码不理解的是
  loss_hist, loss_mlm_hist, loss_num_hist,
  tqdm的用法，
  logit_preds, num_preds, torch.autocast(),
  cross_entropy, mse_loss, num_mask
  loss = loss_mlm+loss_num （这个是唯一理解的）
  后续还有一系列不理解的变量和模块，但是check_point是可以理解的，每个epoch都可以。

• 评估模型
  基本所有的预测函数，都给封装到xval的

  masked_sample = analyze.mask_numbers(sample, tokenizer, n_list=i)    
  out = analyze.predict(model, masked_sample, device)    
  analyze.predict_numbers(model, sample, tokenizer, n_list=number_ids[:3], device='cuda', all_at_once=False)    

  这些预测函数基本都是为了预测单个被mask的number。