import torch

from transformers import BertTokenizer, BertForSequenceClassification

from torch.utils.data import Dataset, DataLoader

import pandas as pd

import numpy as np

 

class NewsDataset(Dataset):

def __init__(self, texts, labels, tokenizer, max_length=512):

self.texts = texts; self.labels = labels; self.tokenizer = tokenizer; self.max_length = max_length

def __len__(self): return len(self.texts)

def __getitem__(self, idx):

text, label = str(self.texts[idx]), self.labels[idx]

encoding = self.tokenizer(text, add_special_tokens=True, max_length=self.max_length, padding='max_length', truncation=True, return_tensors='pt')

return {'input_ids': encoding['input_ids'].flatten(), 'attention_mask': encoding['attention_mask'].flatten(), 'labels': torch.tensor(label, dtype=torch.long)}

 

def train_model(model, train_loader, val_loader, device, epochs=3):

optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)

for epoch in range(epochs):

model.train(); total_loss = 0

for batch in train_loader:

input_ids = batch['input_ids'].to(device); attention_mask = batch['attention_mask'].to(device); labels = batch['labels'].to(device)

outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)

loss = outputs.loss; total_loss += loss.item()

loss.backward(); optimizer.step(); optimizer.zero_grad()

model.eval(); correct = 0; total = 0

with torch.no_grad():

for batch in val_loader:

input_ids = batch['input_ids'].to(device); attention_mask = batch['attention_mask'].to(device); labels = batch['labels'].to(device)

outputs = model(input_ids=input_ids, attention_mask=attention_mask, labels=labels)

_, predicted = torch.max(outputs.logits, 1); total += labels.size(0); correct += (predicted == labels).sum().item()

print(f'Epoch {epoch+1}: Loss={total_loss/len(train_loader):.4f}, Acc={100*correct/total:.2f}%')

 

def main():

df = pd.read_csv('news_dataset.csv')

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

model = BertForSequenceClassification.from_pretrained('bert-base-chinese', num_labels=len(set(df['category'])))

dataset = NewsDataset(df['text'].values, df['category'].values, tokenizer)

train_size = int(0.8 * len(dataset))

train_dataset, val_dataset = torch.utils.data.random_split(dataset, [train_size, len(dataset)-train_size])

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)

val_loader = DataLoader(val_dataset, batch_size=16)

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

model.to(device); train_model(model, train_loader, val_loader, device)

torch.save(model.state_dict(), 'news_classifier.pth')