Advertisement

结合pytorch+transformers进行文本分类

阅读量:

文本分类属于自然语言处理领域的基础性任务,其应用频率极高,本文旨在提供一个即插即用的文本分类工具。

模型model.py文件如下:

复制代码
    # File model.py
    # -*- coding: utf-8 -*-
    import torch
    import os
    from torch import nn
    from transformers import BertForSequenceClassification, BertConfig
    
    class BertModel(nn.Module):
    def __init__(self, num_labels):
        super(BertModel, self).__init__()
        self.bert = BertForSequenceClassification.from_pretrained("hfl/chinese-roberta-wwm-ext", num_labels=num_labels)
        self.device = torch.device("cuda")
        for param in self.bert.parameters():
            param.requires_gra

全部评论 (0)

还没有任何评论哟~