머신러닝으로 기사 분류 자동화하기: 종합 가이드

import os
import pandas as pd
import numpy as np
import random
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.preprocessing import LabelEncoder
from sklearn.model_selection import train_test_split
import re
import string
import nltk
from nltk.corpus import stopwords
from sklearn.metrics import accuracy_score, confusion_matrix, classification_report
import wandb

config = {
    "seed": 42,
    "lr": 0.001,
    "epochs": 5
}

wandb.init(project='article_classifier', config=config)

random.seed(config["seed"])
os.environ["PYTHONHASHSEED"] = str(config["seed"])
np.random.seed(config["seed"])
torch.manual_seed(config["seed"])
torch.cuda.manual_seed(config["seed"])
torch.backends.cudnn.deterministic = True

nltk.download('stopwords')

# Load dataset from CSV
data = pd.read_csv('Articles.csv', encoding='ISO-8859-1')

# Preprocessing
def clean_text(text):
    text = text.lower()  # Convert text to lowercase
    text = re.sub(r'\d+', '', text)  # Remove digits
    text = text.translate(str.maketrans('', '', string.punctuation))  # Remove punctuation
    text = ' '.join([word for word in text.split() if word not in stopwords.words('english')])  # Remove stop words
    text = text.strip()  # Remove leading and trailing whitespaces
    return text

data['Cleaned_Article'] = data['Article'].apply(clean_text)

X = data['Cleaned_Article']
y = data['NewsType']

# Label Encoding
label_encoder = LabelEncoder()
y = label_encoder.fit_transform(y)

class_labels = label_encoder.classes_

# Vectorization
vectorizer = CountVectorizer()
X = vectorizer.fit_transform(X)

# Convert sparse matrix to dense tensor
X = torch.tensor(X.toarray(), dtype=torch.float32)
y = torch.tensor(y, dtype=torch.long)

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=42)
print(f"Training samples: {len(X_train)} - Testing samples: {len(X_test)}")

class NewsClassifier(nn.Module):
    def __init__(self, input_dim, output_dim):
        super(NewsClassifier, self).__init__()
        self.fc = nn.Sequential(
            nn.Linear(input_dim, input_dim//8),
            nn.ReLU(),
            nn.Linear(input_dim//8, output_dim)
        )

    def forward(self, x):
        x = self.fc(x)
        return x

input_dim = X_train.shape[1]
output_dim = len(y_train.unique())
model = NewsClassifier(input_dim, output_dim)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Logging articles to the table
pred_table = wandb.Table(columns=["Name", "Sample 1", "Sample 2", "Sample 3", "Sample 4", "Sample 5", "Sample 6", "Sample 7", "Sample 8", "Sample 9", "Sample 10"])

# Split
_, test_x, _, test_y = train_test_split(data['Cleaned_Article'].values, data['NewsType'].values, test_size=0.5, random_state=42)
pred_table.add_data("Articles", test_x[0], test_x[1], test_x[2], test_x[3], test_x[4], test_x[5], test_x[6], test_x[7], test_x[8], test_x[9])

# Initial results
initial_samples = X_test[:10]  # Get 10 initial samples

initial_predictions = torch.argmax(torch.softmax(model(initial_samples), dim=1), dim=1)
initial_predictions_labels = [class_labels[i] for i in initial_predictions.tolist()]
print("Initial Predictions: \t", initial_predictions_labels)

pred_table.add_data("Initial Predictions", initial_predictions_labels[0], 
initial_predictions_labels[1], initial_predictions_labels[2], initial_predictions_labels[3], 
initial_predictions_labels[4], initial_predictions_labels[5], initial_predictions_labels[6],
initial_predictions_labels[7], initial_predictions_labels[8], initial_predictions_labels[9])

# Training loop
for epoch in range(config["epochs"]):
    model.train()
    optimizer.zero_grad()
    outputs = model(X_train)
    loss = criterion(outputs, y_train)
    print(f"Epoch: {epoch+1:2d} - Loss: {loss:1.4f}")
    wandb.log({"Epoch":epoch+1, "Loss":loss})
    loss.backward()
    optimizer.step()

    # Logging final results
    final_predictions = torch.argmax(torch.softmax(model(initial_samples), dim=1), dim=1)
    final_predictions_labels = [class_labels[i] for i in final_predictions.tolist()]
    pred_table.add_data(f"Epoch {epoch+1}", final_predictions_labels[0], final_predictions_labels[1], final_predictions_labels[2], final_predictions_labels[3], final_predictions_labels[4], final_predictions_labels[5], final_predictions_labels[6], final_predictions_labels[7], final_predictions_labels[8], final_predictions_labels[9])

print("Final Predictions: \t", final_predictions_labels)

model.eval()
with torch.no_grad():
    test_outputs = model(X_test)
    predicted = torch.argmax(torch.softmax(test_outputs, dim=1), dim=1)

    accuracy = accuracy_score(y_test, predicted)
    confusion_mat = confusion_matrix(y_test, predicted)
    classification_rep = classification_report(y_test, predicted)

# Logging evaluation metrics to W&B
wandb.log({"Accuracy": accuracy})
wandb.log({"Confusion Matrix": confusion_mat.tolist()})
wandb.log({"Classification Report": classification_rep})

print("Evaluation ->")
print("Accuracy: ", accuracy)
print("Confusion Matrix:\n", confusion_mat)
print("Classification Report:\n", classification_rep)

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

wandb.finish()