Views
No views yet
1"""
2用预训练的Bert模型微调数据集
3"""
4import torch
5from datasets import load_dataset
6from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
7id_to_label = {0: "positive", 1: "negative", 2: "neutral"}
8label_to_id = {"positive": 0, "negative": 1, "neutral": 2}
9DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
11def predict(text, model, tokenizer, CLASS_NAME):
12 inputs = tokenizer(text, return_tensors="pt")
13 with torch.no_grad():
14 outputs = model(**inputs)
15 logits = outputs.logits
16 predicted_class = torch.argmax(logits).item()
17 print(f"Input Text: {text}")
18 print(f"Predicted class: {int(predicted_class)} {CLASS_NAME[int(predicted_class)]}")
19 return int(predicted_class)
20# 加载数据集
21dataset = load_dataset("csv", data_files={'train': '/mnt/workspace/data2.csv'})
22# 加载预训练的BERT tokenizer
23pretrained_path=r"/mnt/workspace/.cache/modelscope/hub/models/tiansz/bert-base-chinese"
24tokenizer = AutoTokenizer.from_pretrained(pretrained_path)
25model = AutoModelForSequenceClassification.from_pretrained(pretrained_path, num_labels=3)
26# 将标签转换为整数索引
27def map_labels(examples):
28 examples["label"] = label_to_id[examples["label"]]
29 return examples
30
31# 定义tokenize函数
32def tokenize(batch):
33 return tokenizer(batch['text'], padding="max_length", truncation=True, max_length=8, return_tensors="pt")
34# 对数据集进行tokenization
35dataset = dataset.map(map_labels)
36tokenized_datasets=dataset.map(tokenize,batched=True)
37tokenized_datasets = tokenized_datasets.rename_column("label", "labels")
38#处理dataset
39split_dataset = tokenized_datasets['train'].train_test_split(test_size=0.2)
40#print(split_dataset)
41
42# 设置训练参数
43training_args = TrainingArguments(
44 output_dir='./results',
45 eval_strategy='epoch',
46 save_strategy='epoch',
47 learning_rate=2e-5,
48 per_device_train_batch_size=4,
49 per_device_eval_batch_size=4,
50 logging_first_step=100,
51 # 总的训练轮数
52 num_train_epochs=100,
53 weight_decay=0.01,
54 report_to="tensorboard",
55 logging_dir='./logs',
56 save_total_limit=1, # 最大保存数
57 dataloader_num_workers=4, # 增加数据加载器的工作线程数
58 load_best_model_at_end=True, # 训练完成后加载最优模型
59)
60
61
62# 定义Trainer
63trainer = Trainer(
64 model=model,
65 args=training_args,
66 train_dataset=split_dataset['train'],
67 eval_dataset=split_dataset['test'],
68 tokenizer=tokenizer,
69)
70print(DEVICE)
71model.to(DEVICE)
72# 训练模型
73trainer.train()
74# 保存模型
75model.save_pretrained('./classification_model')
76tokenizer.save_pretrained('./classification_model')
77# 测试模型
78test_reviews = [
79 "奥尔西是一名历史老师"
80]
81model.to('cpu')
82text_list = []
83for review in test_reviews:
84 label = predict(review, model, tokenizer, id_to_label)
85 text_list.append(f"{label}-{id_to_label[label]}")
86if text_list:
87 print({"predict": text_list})
88