feat(*): add all for like bert

config.py

@@ -0,0 +1,35 @@
+# coding: UTF-8
+
+import os
+import torch
+
+class Config(object):
+    def __init__(self, data_dir):
+        assert os.path.exists(data_dir)
+        self.train_file = os.path.join(data_dir, "train.txt")
+        self.dev_file = os.path.join(data_dir, "dev.txt")
+        self.label_file = os.path.join(data_dir, "label.txt")
+        assert os.path.isfile(self.train_file)
+        assert os.path.isfile(self.dev_file)
+        assert os.path.isfile(self.label_file)
+
+        self.saved_model_dir = os.path.join(data_dir, "model")
+        self.saved_model = os.path.join(self.saved_model_dir, "bert_model.pth")
+        if not os.path.exists(self.saved_model_dir):
+            os.mkdir(self.saved_model_dir)
+
+        self.label_list = [label.strip() for label in open(self.label_file, "r", encoding="UTF-8").readlines()]
+        self.num_labels = len(self.label_list)
+
+        self.num_epochs = 3
+        self.log_batch = 100
+        self.batch_size = 128
+        self.max_seq_len = 32
+        self.require_improvement = 1000
+
+        self.warmup_steps = 0
+        self.weight_decay = 0.01
+        self.max_grad_norm = 1.0
+        self.learning_rate = 5e-5
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+

data/input.txt

@@ -0,0 +1,3 @@
+调查显示：29.5%的人不满意当年所选高考专业
+广汽今日整体上市 最大短板在过度依赖丰田本田
+梦游之王再现湖人大滑坡 金州小快枪刷分气懵科比

data/label.txt

@@ -0,0 +1,10 @@
+finance
+realty
+stocks
+education
+science
+society
+politics
+sports
+game
+entertainment

main.py

@@ -0,0 +1,106 @@
+# coding: UTF-8
+
+import os
+import time
+import torch
+import argparse
+import numpy as np
+from tqdm import tqdm
+from train import train
+from config import Config
+from preprocess import DataProcessor, get_time_dif
+from transformers import BertConfig, BertTokenizer, BertForSequenceClassification
+
+parser = argparse.ArgumentParser(description="Bert Chinese Text Classification")
+parser.add_argument("--mode", type=str, required=True, help="train/demo/predict")
+parser.add_argument("--data_dir", type=str, default="./data", help="training data and saved model path")
+parser.add_argument("--pretrained_bert_dir", type=str, default="./pretrained_bert", help="pretrained bert model path")
+parser.add_argument("--seed", type=int, default=1, help="random seed for initialization")
+parser.add_argument("--input_file", type=str, default="./data/input.txt", help="input file to be predicted")
+args = parser.parse_args()
+
+def set_seed(seed):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+
+def main():
+    set_seed(args.seed)
+    config = Config(args.data_dir)
+
+    tokenizer = BertTokenizer.from_pretrained(args.pretrained_bert_dir)
+    bert_config = BertConfig.from_pretrained(args.pretrained_bert_dir, num_labels=config.num_labels)
+    model = BertForSequenceClassification.from_pretrained(
+        os.path.join(args.pretrained_bert_dir, "pytorch_model.bin"),
+        config=bert_config
+    )
+    model.to(config.device)
+
+    if args.mode == "train":
+        print("loading data...")
+        start_time = time.time()
+        train_iterator = DataProcessor(config.train_file, config.device, tokenizer, config.batch_size, config.max_seq_len, args.seed)
+        dev_iterator = DataProcessor(config.dev_file, config.device, tokenizer, config.batch_size, config.max_seq_len, args.seed)
+        time_dif = get_time_dif(start_time)
+        print("time usage:", time_dif)
+
+        # train
+        train(model, config, train_iterator, dev_iterator)
+    
+    elif args.mode == "demo":
+        model.load_state_dict(torch.load(config.saved_model))
+        model.eval()
+        while True:
+            sentence = input("请输入文本:\n")
+            inputs = tokenizer(
+                sentence, 
+                max_length=config.max_seq_len,
+                truncation="longest_first",
+                return_tensors="pt")
+            inputs = inputs.to(config.device)
+            with torch.no_grad():
+                outputs = model(**inputs)
+                logits = outputs[0]
+                label = torch.max(logits.data, 1)[1].tolist()
+                print("分类结果:" + config.label_list[label[0]])
+            flag = str(input("continue? (y/n):"))
+            if flag == "Y" or flag == "y":
+                continue
+            else:
+                break
+    else:
+        model.load_state_dict(torch.load(config.saved_model))
+        model.eval()
+
+        text = []
+        with open(args.input_file, mode="r", encoding="UTF-8") as f:
+            for line in tqdm(f):
+                sentence = line.strip()
+                if not sentence:    continue
+                text.append(sentence)
+
+        num_samples = len(text)
+        num_batches = (num_samples - 1) // config.batch_size + 1
+        for i in range(num_batches):
+            start = i * config.batch_size
+            end = min(num_samples, (i + 1) * config.batch_size)
+            inputs = tokenizer.batch_encode_plus(
+                text[start: end],
+                padding=True,
+                max_length=config.max_seq_len,
+                truncation="longest_first",
+                return_tensors="pt")
+            inputs = inputs.to(config.device)
+
+            outputs = model(**inputs)
+            logits = outputs[0]
+
+            preds = torch.max(logits.data, 1)[1].tolist()
+            labels = [config.label_list[_] for _ in preds]
+            for j in range(start, end):
+                print("%s\t%s" % (text[j], labels[j - start]))
+                
+
+if __name__ == "__main__":
+    main()

preprocess.py

@@ -0,0 +1,85 @@
+# coding: UTF-8
+
+import time
+import torch
+import random
+from tqdm import tqdm
+from datetime import timedelta
+
+def get_time_dif(start_time):
+    end_time = time.time()
+    time_dif = end_time - start_time
+    return timedelta(seconds=int(round(time_dif)))
+
+class DataProcessor(object):
+    def __init__(self, path, device, tokenizer, batch_size, max_seq_len, seed):
+        self.seed = seed
+        self.device = device
+        self.tokenizer = tokenizer
+        self.batch_size = batch_size
+        self.max_seq_len = max_seq_len
+
+        self.data = self.load(path)
+
+        self.index = 0
+        self.residue = False
+        self.num_samples = len(self.data[0])
+        self.num_batches = self.num_samples // self.batch_size
+        if self.num_samples % self.batch_size != 0:
+            self.residue = True
+
+    def load(self, path):
+        contents = []
+        labels = []
+        with open(path, mode="r", encoding="UTF-8") as f:
+            for line in tqdm(f):
+                line = line.strip()
+                if not line:    continue
+                if line.find('\t') == -1:   continue
+                content, label = line.split("\t")
+                contents.append(content)
+                labels.append(int(label))
+        #random shuffle
+        index = list(range(len(labels)))
+        random.seed(self.seed)
+        random.shuffle(index)
+        contents = [contents[_] for _ in index]
+        labels = [labels[_] for _ in index]
+        return (contents, labels)
+
+    def __next__(self):
+        if self.residue and self.index == self.num_batches:
+            batch_x = self.data[0][self.index * self.batch_size: self.num_samples]
+            batch_y = self.data[1][self.index * self.batch_size: self.num_samples]
+            batch = self._to_tensor(batch_x, batch_y)
+            self.index += 1
+            return batch
+        elif self.index >= self.num_batches:
+            self.index = 0
+            raise StopIteration
+        else:
+            batch_x = self.data[0][self.index * self.batch_size: (self.index + 1) * self.batch_size]
+            batch_y = self.data[1][self.index * self.batch_size: (self.index + 1) * self.batch_size]
+            batch = self._to_tensor(batch_x, batch_y)
+            self.index += 1
+            return batch
+
+    def _to_tensor(self, batch_x, batch_y):
+        inputs = self.tokenizer.batch_encode_plus(
+            batch_x,
+            padding="max_length",
+            max_length=self.max_seq_len,
+            truncation="longest_first",
+            return_tensors="pt")
+        inputs = inputs.to(self.device)
+        labels = torch.LongTensor(batch_y).to(self.device)
+        return (inputs, labels)
+
+    def __iter__(self):
+        return self
+    
+    def __len__(self):
+        if self.residue:
+            return self.num_batches + 1
+        else:
+            return self.num_batches

pretrained_bert/README.md

@@ -0,0 +1,3 @@
+在 pretrained_bert 文件夹中放入 huggingface 的 bert-base-chinese 模型权重 pytorch_model.bin、配置文件 config.json 和词典 vocab.txt
+
+huggingface bert-base-chinese 下载地址：https://huggingface.co/bert-base-chinese/tree/main

train.py

@@ -0,0 +1,122 @@
+# coding: UTF-8
+
+from typing import Iterator
+from transformers import AdamW, get_linear_schedule_with_warmup
+from preprocess import get_time_dif
+from sklearn import metrics
+import time
+import torch
+import numpy as np
+
+def eval(model, config, iterator, flag=False):
+    model.eval()
+
+    total_loss = 0
+    all_preds = np.array([], dtype=int)
+    all_labels = np.array([], dtype=int)
+    with torch.no_grad():
+        for batch, labels in iterator:
+            outputs = model(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                token_type_ids=batch["token_type_ids"],
+                labels=labels)
+
+            loss = outputs[0]
+            logits = outputs[1]
+
+            total_loss += loss
+            true = labels.data.cpu().numpy()
+            pred = torch.max(logits.data, 1)[1].cpu().numpy()
+            all_labels = np.append(all_labels, true)
+            all_preds = np.append(all_preds, pred)
+    
+    acc = metrics.accuracy_score(all_labels, all_preds)
+    if flag:
+        report = metrics.classification_report(all_labels, all_preds, target_names=config.label_list, digits=4)
+        confusion = metrics.confusion_matrix(all_labels, all_preds)
+        return acc, total_loss / len(iterator), report, confusion
+    return acc, total_loss / len(iterator)
+
+
+def test(model, config, iterator):
+    model.load_state_dict(torch.load(config.saved_model))
+    start_time = time.time()
+    acc, loss, report, confusion = eval(model, config, iterator, flag=True)
+    msg = "Test Loss: {0:>5.2},  Test Acc: {1:>6.2%}"
+    print(msg.format(loss, acc))
+    print("Precision, Recall and F1-Score...")
+    print(report)
+    print("Confusion Matrix...")
+    print(confusion)
+    time_dif = get_time_dif(start_time)
+    print("Time usage:", time_dif)
+
+
+def train(model, config, train_iterator, dev_iterator):
+    model.train()
+    start_time = time.time()
+
+    no_decay = ["bias", "LayerNorm.bias", "LayerNorm.weight"]
+    param_optimizer = model.named_parameters()
+    optimizer_grouped_parameters = [
+        {"params": [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': config.weight_decay},
+        {'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
+    ]
+
+    t_total = len(train_iterator) * config.num_epochs
+    optimizer = AdamW(optimizer_grouped_parameters, lr=config.learning_rate)
+    scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=config.warmup_steps, num_training_steps=t_total)
+
+    total_batch = 0
+    last_improve = 0
+    break_flag = False
+    best_dev_loss = float('inf')
+    for epoch in range(config.num_epochs):
+        print("Epoch [{}/{}]".format(epoch + 1, config.num_epochs))
+        for _, (batch, labels) in enumerate(train_iterator):
+
+            outputs = model(
+                input_ids=batch["input_ids"],
+                attention_mask=batch["attention_mask"],
+                token_type_ids=batch["token_type_ids"],
+                labels=labels)
+            
+            loss = outputs[0]
+            logits = outputs[1]
+
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), config.max_grad_norm)
+
+            optimizer.step()
+            scheduler.step() 
+            optimizer.zero_grad()
+
+            if total_batch % config.log_batch == 0:
+                true = labels.data.cpu()
+                pred = torch.max(logits.data, 1)[1].cpu()
+                acc = metrics.accuracy_score(true, pred)
+                dev_acc, dev_loss = eval(model, config, dev_iterator)
+                if dev_loss < best_dev_loss:
+                    best_dev_loss = dev_loss
+                    torch.save(model.state_dict(), config.saved_model)
+                    improve = "*"
+                    last_improve = total_batch
+                else:
+                    improve = ""
+
+                time_dif = get_time_dif(start_time)
+                msg = 'Iter: {0:>6}, Batch Train Loss: {1:>5.2}, Batch Train Acc: {2:>6.2%}, Val Loss: {3:>5.2}, Val Acc: {4:>6.2%}, Time: {5} {6}'
+                print(msg.format(total_batch, loss.item(), acc, dev_loss, dev_acc, time_dif, improve))
+                model.train()
+
+            total_batch += 1
+            if total_batch - last_improve > config.require_improvement:
+                print("No improvement for a long time, auto-stopping...")
+                break_flag = True
+                break
+        if break_flag:
+            break
+    
+    test(model, config, dev_iterator)
+