欢迎访问宙启技术站
智能推送

如何使用transformers库进行中文情感分类

发布时间:2024-01-17 23:42:56

要使用transformers库进行中文情感分类,需要以下几个步骤:

1. 安装依赖库:在使用transformers库之前,需要先安装依赖库。你可以使用pip命令来安装所需的库。在终端中运行以下命令:

pip install transformers

2. 加载预训练模型:transformers库提供了多种预训练的模型,如BERT、GPT、RoBERTa等。你可以根据自己的需求来选择合适的预训练模型。这里以BERT模型为例进行说明。首先需要导入必要的库:

from transformers import BertTokenizer, BertForSequenceClassification

model_name = 'bert-base-chinese'
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name)

3. 数据准备:为了使用模型进行训练和预测,需要将文本数据转化为模型所需的输入格式。在中文情感分类任务中,通常采用的输入格式是句子级别的分类,即将输入的句子转化为词嵌入向量的序列。可以使用tokenizer对输入文本进行分词和编码。

def prepare_input(text):
  encoded_input = tokenizer.encode_plus(
    text,
    text_pair=None,
    add_special_tokens=True,
    padding='max_length',
    truncation=True,
    max_length=128,
    return_tensors='pt',
  )
  return encoded_input

4. 训练模型:使用transformers库训练中文情感分类模型的具体过程与通用的文本分类任务类似。首先需要定义模型的优化器和损失函数。然后,使用数据集进行训练和验证。下面是一个简单的训练函数示例:

import torch
from torch.utils.data import DataLoader

def train(model, optimizer, loss_fn, train_dataset, val_dataset, batch_size=32, epochs=10):
  train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
  val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  model = model.to(device)

  for epoch in range(epochs):
    model.train()
    train_loss = 0.0
    for batch in train_dataloader:
      batch = {k: v.to(device) for k, v in batch.items()}
      optimizer.zero_grad()
      outputs = model(**batch)
      loss = outputs.loss
      train_loss += loss.item()
      loss.backward()
      optimizer.step()

    val_loss = 0.0
    model.eval()
    with torch.no_grad():
      for batch in val_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        val_loss += loss.item()
    
    print('Epoch {}/{} - Train loss: {:.4f} - Val loss: {:.4f}'.format(epoch+1, epochs, train_loss, val_loss))

5. 预测结果:完成模型的训练后,可以使用训练好的模型对新的文本数据进行情感分类。下面是一个简单的预测函数示例:

def predict(model, text):
  model.eval()
  input_data = prepare_input(text)
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  input_data = {k: v.to(device) for k, v in input_data.items()}

  with torch.no_grad():
    outputs = model(**input_data)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=1).item()
  
  return predicted_class

6. 完整示例代码如下:

from transformers import BertTokenizer, BertForSequenceClassification
import torch
from torch.utils.data import Dataset, DataLoader

class SentimentDataset(Dataset):
  def __init__(self, texts, labels):
    self.texts = texts
    self.labels = labels
    
  def __len__(self):
    return len(self.texts)
  
  def __getitem__(self, index):
    return {'text': self.texts[index], 'label': self.labels[index]}

def prepare_input(text):
  encoded_input = tokenizer.encode_plus(
    text,
    text_pair=None,
    add_special_tokens=True,
    padding='max_length',
    truncation=True,
    max_length=128,
    return_tensors='pt',
  )
  return encoded_input

def train(model, optimizer, loss_fn, train_dataset, val_dataset, batch_size=32, epochs=10):
  train_dataloader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
  val_dataloader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)

  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  model = model.to(device)

  for epoch in range(epochs):
    model.train()
    train_loss = 0.0
    for batch in train_dataloader:
      batch = {k: v.to(device) for k, v in batch.items()}
      optimizer.zero_grad()
      outputs = model(**batch)
      loss = outputs.loss
      train_loss += loss.item()
      loss.backward()
      optimizer.step()

    val_loss = 0.0
    model.eval()
    with torch.no_grad():
      for batch in val_dataloader:
        batch = {k: v.to(device) for k, v in batch.items()}
        outputs = model(**batch)
        loss = outputs.loss
        val_loss += loss.item()
    
    print('Epoch {}/{} - Train loss: {:.4f} - Val loss: {:.4f}'.format(epoch+1, epochs, train_loss, val_loss))

def predict(model, text):
  model.eval()
  input_data = prepare_input(text)
  device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
  input_data = {k: v.to(device) for k, v in input_data.items()}

  with torch.no_grad():
    outputs = model(**input_data)
    logits = outputs.logits
    predicted_class = torch.argmax(logits, dim=1).item()
  
  return predicted_class

if __name__ == '__main__':
  model_name = 'bert-base-chinese'
  tokenizer = BertTokenizer.from_pretrained(model_name)
  model = BertForSequenceClassification.from_pretrained(model_name)

  # 准备训练数据和验证数据
  train_texts = ['我喜欢这个电影', '这是一家很好的餐厅', '这个产品很实用', '这个景点真美']
  train_labels = [1, 1, 1, 1]
  val_texts = ['这部电影太差了', '这家餐厅太贵了', '这个产品质量很差', '这个景点太糟糕了']
  val_labels = [0, 0, 0, 0]
  
  train_dataset = SentimentDataset(train_texts, train_labels)
  val_dataset = SentimentDataset(val_texts, val_labels)

  # 定义模型的优化器和损失函数
  optimizer = torch.optim.AdamW(model.parameters(), lr=1e-5)
  loss_fn = torch.nn.CrossEntropyLoss()

  # 训练模型
  train(model, optimizer, loss_fn, train_dataset, val_dataset, batch_size=4, epochs=5)

  # 预测示例文本
  example_text = '这个电影很有意思'
  predicted_class = predict(model, example_text)
  print('Predicted class:', predicted_class)

以上是使用transformers库进行中文情感分类的示例代码。通过这个例子可以快速上手使用transformers库进行中文文本分类任务。