结合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)
还没有任何评论哟~
