from torch import nn
from torch.optim import Adam
import torch

model = nn.Linear(2, 1)  # 模型定义，内部是 2 * 1的矩阵
optimizer = Adam(model.parameters(), lr=1e-1)
loss_fn = nn.MSELoss()

## 输入是 10 * 2矩阵，可以理解为10个样本输入，每个样本是一对值组成，譬如:样本A(0.3,0,6) 样本B(0.2,0,8)
## 目标是 10 * 1矩阵，可以理解为10个样本对应的值， 这里简单把每个样本的两个值相加，并加入一个随机变量。
## 对于单个变量， Y = sum[x1,x2] + 随机值， 希望模型能够学会的是 Y = w1 * x1 + w2 * x2，这里的w1, w2就是模型的2*1矩阵的两个值，希望w1、w2最终学为1
input = torch.randn(10,2) * 10
#普朗克常数
bias = 6.6260693
target = torch.add(input.sum(dim=1, keepdim=True),bias)
print('训练样本输入',input,'\n训练样本输入-偏移量', bias, '\n训练样本目标值', target)

print("\n模型内部参数的初始随机值")
for name, param in model.named_parameters():
  print(name, param.data)

print("\n开始训练，会发现差异值越来越小，说明模型在收敛")
for epoch in range(100): # 练习100轮
    pred = model(input) # 输入 10*2 乘积模型的 2*1 得到预测值 10*1

    loss = loss_fn(pred, target) # 和目标值差异

    # 根据差异，修正模型参数。就是调w1, w2两个参数
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if epoch%10==9:
        print (f"Epoch: {epoch} | loss: {loss:.2f}")

# 拿一对随机值，测下模型， 是否产出值 = sum([x1,x2])
test = torch.randn(1,2) * 10
target = test.sum(dim=1, keepdim=True)
pred = model(test)
print('\n训练完毕，测试模型。\n可以看到预测值和目标值近似，说明模型训练成功\n测试输入:',test, '模型预测:',pred.detach(), '目标值:', target)

print("\n可以看到经过训练后的模型值w1、w2接近于1")
for name, param in model.named_parameters():
  print(name, param.data)

import torch
from torch import nn
from torch import optim
from torch.utils import data as Data
import numpy as np

d_model = 6 # Embedding的大小
max_len = 1024 # 输入序列的最长大小
d_ff = 12 # 前馈神经网络的隐藏层大小, 一般是d_model的四倍
d_k = d_v = 3 # 自注意力中K和V的维度, Q的维度直接用K的维度代替, 因为这二者必须始终相等
n_layers = 1 # Encoder和Decoder的层数
n_heads = 8 # 自注意力多头的头数
p_drop = 0.1 # propability of dropout

# 对encoder_input的PAD(0)做Mask，可以支持对训练样本打掩码
def get_attn_pad_mask(seq_q, seq_k):
  batch, len_q = seq_q.size()
  batch, len_k = seq_k.size()
  # we define index of PAD is 0, if tensor equals (zero) PAD tokens
  pad_attn_mask = seq_k.data.eq(0).unsqueeze(1) # [batch, 1, len_k]

  return pad_attn_mask.expand(batch, len_q, len_k) # [batch, len_q, len_k]

# 在deocder_input做上三角掩码
def get_attn_subsequent_mask(seq):
  attn_shape = [seq.size(0), seq.size(1), seq.size(1)] # [batch, target_len, target_len]
  subsequent_mask = np.triu(np.ones(attn_shape), k=1) # [batch, target_len, target_len]
  subsequent_mask = torch.from_numpy(subsequent_mask)

  return subsequent_mask # [batch, target_len, target_len]

# 做encoder_input添加位置编码
class PositionalEncoding(nn.Module):

  def __init__(self, d_model, dropout=.1, max_len=1024):
    super(PositionalEncoding, self).__init__()
    self.dropout = nn.Dropout(p=p_drop)

    positional_encoding = torch.zeros(max_len, d_model) # [max_len, d_model]
    position = torch.arange(0, max_len).float().unsqueeze(1) # [max_len, 1]

    div_term = torch.exp(torch.arange(0, d_model, 2).float() *
                         (-torch.log(torch.Tensor([10000])) / d_model)) # [max_len / 2]

    positional_encoding[:, 0::2] = torch.sin(position * div_term) # even
    positional_encoding[:, 1::2] = torch.cos(position * div_term) # odd

    # [max_len, d_model] -> [1, max_len, d_model] -> [max_len, 1, d_model]
    positional_encoding = positional_encoding.unsqueeze(0).transpose(0, 1)

    # register_buffer能够申请一个缓冲区中的常量, 并且它不会被加入到计算图中, 也就不会参与反向传播.
    self.register_buffer('pe', positional_encoding)

  def forward(self, x):
    # x: [seq_len, batch, d_model]
    # we can add positional encoding to x directly, and ignore other dimension
    x = x + self.pe[:x.size(0), ...]

    return self.dropout(x)

# Encoder和Deocder后的前馈层（含归一层）
class FeedForwardNetwork(nn.Module):
  def __init__(self):
    super(FeedForwardNetwork, self).__init__()
    self.ff1 = nn.Linear(d_model, d_ff)
    # 线性变化还原
    self.ff2 = nn.Linear(d_ff, d_model)
    self.relu = nn.ReLU()

    self.dropout = nn.Dropout(p=p_drop)
    self.layer_norm = nn.LayerNorm(d_model)

  def forward(self, x):
    x = self.ff1(x)
    x = self.relu(x)
    x = self.ff2(x)
    return self.layer_norm(x)

# 多头注意力（含归一层）
class MultiHeadAttention(nn.Module):

  def __init__(self, n_heads=8):
    super(MultiHeadAttention, self).__init__()
    # do not use more instance to implement multihead attention
    # it can be complete in one matrix
    self.n_heads = n_heads

    # we can't use bias because there is no bias term in formular
    # 多头放在同一个矩阵计算
    self.W_Q = nn.Linear(d_model, d_k * n_heads, bias=False)
    self.W_K = nn.Linear(d_model, d_k * n_heads, bias=False)
    self.W_V = nn.Linear(d_model, d_v * n_heads, bias=False)
    self.fc = nn.Linear(d_v * n_heads, d_model, bias=False)
    self.layer_norm = nn.LayerNorm(d_model)

  def forward(self, input_Q, input_K, input_V, attn_mask):
    '''
    To make sure multihead attention can be used both in encoder and decoder,
    we use Q, K, V respectively.
    input_Q: [batch, len_q, d_model]
    input_K: [batch, len_k, d_model]
    input_V: [batch, len_v, d_model]
    '''
    residual, batch = input_Q, input_Q.size(0)

    # [batch, len_q, d_model] -- matmul W_Q --> [batch, len_q, d_q * n_heads] -- view -->
    # [batch, len_q, n_heads, d_k,] -- transpose --> [batch, n_heads, len_q, d_k]

    Q = self.W_Q(input_Q).view(batch, -1, n_heads, d_k).transpose(1, 2) # [batch, n_heads, len_q, d_k]
    K = self.W_K(input_K).view(batch, -1, n_heads, d_k).transpose(1, 2) # [batch, n_heads, len_k, d_k]
    V = self.W_V(input_V).view(batch, -1, n_heads, d_v).transpose(1, 2) # [batch, n_heads, len_v, d_v]

    attn_mask = attn_mask.unsqueeze(1).repeat(1, n_heads, 1, 1) # [batch, n_heads, seq_len, seq_len]

    # prob: [batch, n_heads, len_q, d_v] attn: [batch, n_heads, len_q, len_k]
    prob, attn = ScaledDotProductAttention()(Q, K, V, attn_mask)

    prob = prob.transpose(1, 2).contiguous() # [batch, len_q, n_heads, d_v]
    prob = prob.view(batch, -1, n_heads * d_v).contiguous() # [batch, len_q, n_heads * d_v]

    output = self.fc(prob) # [batch, len_q, d_model]

    return self.layer_norm(residual + output), attn

# 点积，QKV
class ScaledDotProductAttention(nn.Module):
  def __init__(self):
    super(ScaledDotProductAttention, self).__init__()

  def forward(self, Q, K, V, attn_mask):
    '''
    Q: [batch, n_heads, len_q, d_k]
    K: [batch, n_heads, len_k, d_k]
    V: [batch, n_heads, len_v, d_v]
    attn_mask: [batch, n_heads, seq_len, seq_len]
    '''
    scores = torch.matmul(Q, K.transpose(-1, -2)) / np.sqrt(d_k) # [batch, n_heads, len_q, len_k]
    # -1e9是用很大的负数，使得其在Softmax中可以被忽略，实现mask效果
    scores.masked_fill_(attn_mask, -1e9)

    attn = nn.Softmax(dim=-1)(scores) # [batch, n_heads, len_q, len_k]
    prob = torch.matmul(attn, V) # [batch, n_heads, len_q, d_v]
    return prob, attn

# Encoder层，可以构建多层Encoder，由多头注意力和前馈层构成
class EncoderLayer(nn.Module):

  def __init__(self):
    super(EncoderLayer, self).__init__()
    self.encoder_self_attn = MultiHeadAttention()
    self.ffn = FeedForwardNetwork()

  def forward(self, encoder_input, encoder_pad_mask):
    '''
    encoder_input: [batch, source_len, d_model]
    encoder_pad_mask: [batch, n_heads, source_len, source_len]

    encoder_output: [batch, source_len, d_model]
    attn: [batch, n_heads, source_len, source_len]
    '''
    encoder_output, attn = self.encoder_self_attn(encoder_input, encoder_input, encoder_input, encoder_pad_mask)
    encoder_output = self.ffn(encoder_output) # [batch, source_len, d_model]

    return encoder_output, attn

# Encoder
class Encoder(nn.Module):

  def __init__(self):
    super(Encoder, self).__init__()
    self.source_embedding = nn.Embedding(source_vocab_size, d_model)
    self.positional_embedding = PositionalEncoding(d_model)
    self.layers = nn.ModuleList([EncoderLayer() for layer in range(n_layers)])

  def forward(self, encoder_input):
    # encoder_input: [batch, source_len]
    encoder_output = self.source_embedding(encoder_input) # [batch, source_len, d_model]
    encoder_output = self.positional_embedding(encoder_output.transpose(0, 1)).transpose(0, 1) # [batch, source_len, d_model]

    encoder_self_attn_mask = get_attn_pad_mask(encoder_input, encoder_input) # [batch, source_len, source_len]
    encoder_self_attns = list()
    for layer in self.layers:
      # encoder_output: [batch, source_len, d_model]
      # encoder_self_attn: [batch, n_heads, source_len, source_len]
      encoder_output, encoder_self_attn = layer(encoder_output, encoder_self_attn_mask)
      encoder_self_attns.append(encoder_self_attn)

    return encoder_output, encoder_self_attns
 # Decoder层，可以构建多层Decoder，由多头注意力和前馈层构成
 class DecoderLayer(nn.Module):

  def __init__(self):
    super(DecoderLayer, self).__init__()
    self.decoder_self_attn = MultiHeadAttention()
    self.encoder_decoder_attn = MultiHeadAttention()
    self.ffn = FeedForwardNetwork()

  def forward(self, decoder_input, encoder_output, decoder_self_mask, decoder_encoder_mask):
    '''
    decoder_input: [batch, target_len, d_mdoel]
    encoder_output: [batch, source_len, d_model]
    decoder_self_mask: [batch, target_len, target_len]
    decoder_encoder_mask: [batch, target_len, source_len]
    '''
    # masked mutlihead attention
    # Q, K, V all from decoder it self
    # decoder_output: [batch, target_len, d_model]
    # decoder_self_attn: [batch, n_heads, target_len, target_len]
    decoder_output, decoder_self_attn = self.decoder_self_attn(decoder_input, decoder_input, decoder_input, decoder_self_mask)

    # Q from decoder, K, V from encoder
    # decoder_output: [batch, target_len, d_model]
    # decoder_encoder_attn: [batch, n_heads, target_len, source_len]
    decoder_output, decoder_encoder_attn = self.encoder_decoder_attn(decoder_output, encoder_output, encoder_output, decoder_encoder_mask)
    decoder_output = self.ffn(decoder_output) # [batch, target_len, d_model]

    return decoder_output, decoder_self_attn, decoder_encoder_attn

# Decoder
class Decoder(nn.Module):

  def __init__(self):
    super(Decoder, self).__init__()
    self.target_embedding = nn.Embedding(target_vocab_size, d_model)
    self.positional_embedding = PositionalEncoding(d_model)
    self.layers = nn.ModuleList([DecoderLayer() for layer in range(n_layers)])

  def forward(self, decoder_input, encoder_input, encoder_output):
    '''
    decoder_input: [batch, target_len]
    encoder_input: [batch, source_len]
    encoder_output: [batch, source_len, d_model]
    '''
    decoder_output = self.target_embedding(decoder_input) # [batch, target_len, d_model]
    decoder_output = self.positional_embedding(decoder_output.transpose(0, 1)).transpose(0, 1) # [batch, target_len, d_model]
    decoder_self_attn_mask = get_attn_pad_mask(decoder_input, decoder_input) # [batch, target_len, target_len]
    decoder_subsequent_mask = get_attn_subsequent_mask(decoder_input) # [batch, target_len, target_len]

    decoder_encoder_attn_mask = get_attn_pad_mask(decoder_input, encoder_input) # [batch, target_len, source_len]

    decoder_self_mask = torch.gt(decoder_self_attn_mask + decoder_subsequent_mask, 0)
    decoder_self_attns, decoder_encoder_attns = [], []

    for layer in self.layers:
      # decoder_output: [batch, target_len, d_model]
      # decoder_self_attn: [batch, n_heads, target_len, target_len]
      # decoder_encoder_attn: [batch, n_heads, target_len, source_len]
      decoder_output, decoder_self_attn, decoder_encoder_attn = layer(decoder_output, encoder_output, decoder_self_mask, decoder_encoder_attn_mask)
      decoder_self_attns.append(decoder_self_attn)
      decoder_encoder_attns.append(decoder_encoder_attn)

    return decoder_output, decoder_self_attns, decoder_encoder_attns

class Transformer(nn.Module):

  def __init__(self):
    super(Transformer, self).__init__()

    self.encoder = Encoder()
    self.decoder = Decoder()
    self.fc = nn.Linear(d_model, target_vocab_size, bias=False)

  def forward(self, encoder_input, decoder_input):
    '''
    encoder_input: [batch, source_len]
    decoder_input: [batch, target_len]
    '''
    # encoder_output: [batch, source_len, d_model]
    # encoder_attns: [n_layers, batch, n_heads, source_len, source_len]
    encoder_output, encoder_attns = self.encoder(encoder_input)
    # decoder_output: [batch, target_len, d_model]
    # decoder_self_attns: [n_layers, batch, n_heads, target_len, target_len]
    # decoder_encoder_attns: [n_layers, batch, n_heads, target_len, source_len]
    decoder_output, decoder_self_attns, decoder_encoder_attns = self.decoder(decoder_input, encoder_input, encoder_output)
    decoder_logits = self.fc(decoder_output) # [batch, target_len, target_vocab_size]

    # decoder_logits: [batch * target_len, target_vocab_size]
    return decoder_logits.view(-1, decoder_logits.size(-1))

class Tokenizer:
  def __init__(self, sentences):
    super(Tokenizer, self).__init__()
    self.sentences = sentences

  def get_source_vocab(self):
    return self.source_vocab

  def get_target_vocab(self):
    return self.target_vocab

  def convert_token_to_ids(self):
    source_inputs = " ".join([sentences[i][0] for i in range(len(sentences))]).replace('E', '').split()
    source_inputs.insert(0, 'E')
    self.source_vocab = {k: v for v, k in enumerate(source_inputs)}
    target_inputs = " ".join([sentences[i][1] for i in range(len(sentences))]).replace('E', '').replace('S', '').split()
    target_inputs.insert(0, 'E')
    target_inputs.insert(1, 'S')
    self.target_vocab = {k: v for v, k in enumerate(target_inputs)}

    encoder_inputs, decoder_inputs, decoder_outputs = [], [], []
    for i in range(len(sentences)):
      encoder_input = [self.source_vocab[word] for word in sentences[i][0].split()]
      decoder_input = [self.target_vocab[word] for word in sentences[i][1].split()]
      decoder_output = [self.target_vocab[word] for word in sentences[i][2].split()]
      encoder_inputs.append(encoder_input)
      decoder_inputs.append(decoder_input)
      decoder_outputs.append(decoder_output)
    return torch.LongTensor(encoder_inputs), torch.LongTensor(decoder_inputs), torch.LongTensor(decoder_outputs)

  def convert_ids_to_source_sentences(self, ids, split_word=' '):
    return split_word.join([key for key in self.source_vocab][ids[i].item()] for i in range(len(ids)))

  def convert_ids_to_target_sentences(self, ids, split_word=' '):
    return split_word.join([key for key in self.target_vocab][ids[i].item()] for i in range(len(ids)))

# 训练
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
epochs = 400
lr = 1e-1

model = Transformer().to(device)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=lr)

sentences = [
        # Encoder的输入          Decoder的输入（teaching force）   #Decoder的输出（样本值）
        ['LLM with banzang E',  'S 半臧 和 大模型',               '半臧 和 大模型 E']
]
tokenizer = Tokenizer(sentences)
encoder_inputs, decoder_inputs, decoder_outputs = tokenizer.convert_token_to_ids()
dataset = Seq2SeqDataset(encoder_inputs, decoder_inputs, decoder_outputs)
data_loader = Data.DataLoader(dataset, 2, True)

source_vocab_size = len(tokenizer.get_source_vocab())
target_vocab_size = len(tokenizer.get_target_vocab())

for epoch in range(epochs):
  for encoder_input, decoder_input, decoder_output in data_loader:
    encoder_input = encoder_input.to(device)
    decoder_input = decoder_input.to(device)
    decoder_output = decoder_output.to(device)

    output = model(encoder_input, decoder_input)
    loss = criterion(output, decoder_output.view(-1))

    print('Epoch:', '%04d' % (epoch + 1), 'loss =', '{:.6f}'.format(loss), '\t预测:', tokenizer.convert_ids_to_target_sentences(output.max(dim=1, keepdim=False)[1].data), '\t(训练样本:', tokenizer.convert_ids_to_target_sentences(decoder_output.view(-1).data), ')')

    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

# 预测
output_len = len(decoder_outputs.squeeze(0))
for encoder_input, decoder_input, decoder_output in data_loader:
    encoder_input = encoder_input.to(device)
    decoder_input = torch.zeros(1, output_len).type_as(encoder_input.data)
    # 从"S"开始，source_vocab={'E':0, 'S':1, ...}
    next_symbol = 1
    print('输入:', tokenizer.convert_ids_to_source_sentences(encoder_input.data.squeeze(0)))

    for i in range(output_len):
        decoder_input[0][i] = next_symbol
        output = model(encoder_input, decoder_input)
        prob = output.max(dim=1, keepdim=False)[1]
        next_symbol = prob.data[i].item()

        print('输出:', tokenizer.convert_ids_to_target_sentences(prob.data[:i+1], ''))

        if next_symbol == 0:
          break

#第一层的Encoder的输出会做为第二层Encoder的输入
for layer in self.layers:
    encoder_output, encoder_self_attn = layer(encoder_output, encoder_self_attn_mask)
    encoder_self_attns.append(encoder_self_attn)

import torch
from transformers import AdamW, AutoTokenizer, AutoModelForSequenceClassification
from torch.utils.data import Dataset, DataLoader
from tqdm.auto import tqdm
checkpoint = "hfl/minirbt-h288"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=3)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)
# 1. 训练和评估数据准备
class TrainDataset(Dataset):
    def __init__(self, sentences):
        self.sentences = sentences
    def __len__(self):
        return len(self.sentences)
    def __getitem__(self, idx):
        return sentences['train'][idx][0], sentences['train'][idx][1], sentences['train'][idx][2]
class ValidationDataset(Dataset):
    def __init__(self, sentences):
        self.sentences = sentences
    def __len__(self):
        return len(self.sentences)
    def __getitem__(self, idx):
        return sentences['validation'][idx][0], sentences['validation'][idx][1], sentences['validation'][idx][2]
def data_collator(batch):
    sentence ,labels = [],[]
    for item in batch:
        sentence.append([item[0], item[1]])
        labels.append(item[2])
    inputs = tokenizer(sentence, padding=True, truncation=True, return_tensors="pt")
    inputs['labels'] = torch.tensor(labels)
    return inputs
# 2. 准备训练和评估数据
sentences = {
    "train": [   
        ["I have a green apple", "apple", 0],
        ["I have a black apple", "apple", 1],
        ["I have a red apple", "apple", 2],
        ["I have a red banner", "banner", 2]
    ],
    "validation": [
        ["I have a black apple", "apple", 1],
        ["I have a green banner", "banner", 0]
    ]
}
# 3. 评估模型
import evaluate
from transformers import TrainingArguments
from transformers import Trainer
import numpy as np
def compute_metrics(eval_preds):
    metric = evaluate.load("glue", "mrpc")
    logits, labels = eval_preds
    predictions = np.argmax(logits, axis=-1)
    return metric.compute(predictions=predictions, references=labels)
training_args = TrainingArguments("test-trainer", evaluation_strategy="epoch")
trainer = Trainer(
    model,
    training_args,
    train_dataset = TrainDataset(sentences),
    eval_dataset = ValidationDataset(sentences),
    data_collator = data_collator,
    tokenizer = tokenizer,
    compute_metrics=compute_metrics
)
trainer.train()
# 4. 进行训练
train_dataset = TrainDataset(sentences)
train_dataloader = DataLoader(train_dataset, shuffle=True, batch_size=8, collate_fn=data_collator)
num_epochs = 10
progress_bar = tqdm(range(num_epochs))
model.train()
# 观测指标
metric = evaluate.load("glue", "mrpc")
optimizer = AdamW(model.parameters(), lr=5e-5)
for epoch in range(num_epochs):
    for batch in train_dataloader:
        optimizer.zero_grad()
        output = model(**batch)
        loss, logits = output[:2]
        loss.backward()
        optimizer.step()
        predictions = torch.argmax(logits, dim=-1)
        metric.add_batch(predictions=predictions, references=batch["labels"])
        
        # 过程日志
        print(f"Epoch {epoch + 1}, Loss: {loss.item()}")
        progress_bar.update(1)
# 观测指标结果
metric.compute()
# 5. 预测
model.eval()
sentences = ["I have a black banner", "banner"]
logits = model(**tokenizer(sentences[0], sentences[1], return_tensors="pt")).logits
pred = torch.argmax(logits,dim=-1)
print(pred)

checkpoint = "hfl/minirbt-h288"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=3)

import torch
from transformers import AdamW, AutoTokenizer, AutoModelForSequenceClassification
from torch.utils.data import Dataset, DataLoader
from tqdm.auto import tqdm
from datasets import load_dataset
device = 'mps' if torch.backends.mps.is_available() else 'cpu'
class UniqueLabelsDataset(Dataset):
    def __init__(self, data, labels):
        self.data = data
        self.labels = labels
    def __len__(self):
        return len(self.data)
    def __getitem__(self, idx):
        return self.labels[idx], pre_process_input(self.data[idx].values())
# 训练和评估数据准备
def data_collator(batch):
    sentence, labels = [],[]
    for l, item in batch:
        sentence.append(item)
        labels.append(l.item())
    inputs = tokenizer(sentence, padding=True, truncation=True, return_tensors="pt")
    inputs['labels'] = torch.tensor(labels)
    return inputs
def pre_process_input(inputs):
    return f" ".join([str(value).replace(' ', '') for value in inputs])
# 准备训练和评估数据
'''
{
  '标题': Value(dtype='string', id=None), 
  '需求指派人': Value(dtype='string', id=None), 
  '***': Value(dtype='int64', id=None), 
  '***': Value(dtype='string', id=None), 
  '***': Value(dtype='string', id=None), 
  '***': Value(dtype='string', id=None), 
  '子需求技术pm': Value(dtype='int64', id=None), 
  'team_alias': Value(dtype='string', id=None)}
'''
train_dataset = load_dataset("csv", data_files = "./tt_l_demands.csv", delimiter=",")['train']
labels_vocab, unique_labels = torch.unique(torch.tensor(train_dataset['子需求技术pm']), return_inverse=True)
num_labels = len(labels_vocab)
train_dataset.remove_columns(["子需求技术pm"])
train_dataloader = DataLoader(UniqueLabelsDataset(train_dataset, unique_labels), shuffle=True, batch_size=64, collate_fn=data_collator)
# 初始化模型
checkpoint = "/Users/nanzhang/minirbt-h288"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForSequenceClassification.from_pretrained(checkpoint, num_labels=num_labels)
model.to(device)
from transformers import get_scheduler
# 进行训练
num_epochs = 100
# 显示训练进度，dynamic_ncols=True保持在一行
progress_bar = tqdm(range(num_epochs), dynamic_ncols=True)
#  lr从 1e-1 的预热，爬坡到 5e-5 
num_training_steps = num_epochs * len(train_dataloader)
lr_scheduler = get_scheduler(
    "linear",
    optimizer=optimizer,
    num_warmup_steps=1e-1,
    num_training_steps=num_training_steps,
)
model.train()
# 观测指标
optimizer = AdamW(model.parameters(), lr=5e-5)
for epoch in range(num_epochs):
    for batch in train_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        output = model(**batch)
        loss, logits = output[:2]
        loss.backward()
        optimizer.step()
        lr_scheduler.step()
        optimizer.zero_grad()
    # 显示进度
    #print(f"Epoch {epoch + 1}, Loss: {loss.item()}")
    progress_bar.set_postfix(loss=loss.item())
    progress_bar.update(1)
# 持久化训练好的参数
model.save_pretrained('./model')
tokenizer.save_pretrained('./model')
# 预测
sentences = ["***", '407080', '407080', '2722', '4066', '淘天项目', '50932674', '***']
input = tokenizer(pre_process_input(sentences), padding=True, truncation=True, return_tensors="pt")
input.to(device)
logits = model(**input).logits
pred = torch.argmax(logits,dim=-1)
print(labels_vocab[pred.item()])

{
  '标题': Value(dtype='string', id=None), 
  '需求指派人': Value(dtype='string', id=None), 
  '***': Value(dtype='int64', id=None), 
  '***': Value(dtype='string', id=None), 
  '***': Value(dtype='string', id=None), 
  '子需求pd': Value(dtype='string', id=None), 
  '子需求技术pm': Value(dtype='int64', id=None), 
  '***': Value(dtype='string', id=None)
}