NLP Text Classification

import torch
from transformers import (
    BertTokenizer, BertForSequenceClassification,
    AdamW, get_linear_schedule_with_warmup
)
from torch.utils.data import Dataset, DataLoader
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
import numpy as np

# =============================================
# 1. CUSTOM DATASET
# =============================================
class TextClassificationDataset(Dataset):
    def __init__(self, texts, labels, tokenizer, max_len=128):
        self.texts = texts
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_len = max_len
    
    def __len__(self):
        return len(self.texts)
    
    def __getitem__(self, idx):
        encoding = self.tokenizer(
            self.texts[idx],
            max_length=self.max_len,
            padding='max_length',
            truncation=True,
            return_tensors='pt'
        )
        return {
            'input_ids': encoding['input_ids'].squeeze(0),
            'attention_mask': encoding['attention_mask'].squeeze(0),
            'labels': torch.tensor(self.labels[idx], dtype=torch.long)
        }

# =============================================
# 2. PREPARE DATA
# =============================================
texts = [
    "Produk ini sangat bagus dan berkualitas tinggi",
    "Saya sangat puas dengan pembelian ini",
    "Barang cepat sampai dan sesuai deskripsi",
    "Kualitas sangat baik, recommended",
    "Pelayanan ramah dan pengiriman cepat",
    "Sangat senang dengan hasilnya",
    "Produk sangat mengecewakan dan rusak",
    "Barang tidak sesuai foto, sangat buruk",
    "Pengiriman sangat lambat",
    "Kualitas jelek, buang uang saja",
    "Pelayanan buruk dan tidak ramah",
    "Barang palsu dan tidak berfungsi",
] * 30  # Repeat for enough data

labels = [1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0] * 30

X_train, X_val, y_train, y_val = train_test_split(
    texts, labels, test_size=0.2, random_state=42, stratify=labels
)

# =============================================
# 3. LOAD BERT
# =============================================
MODEL_NAME = 'indobenchmark/indobert-base-p1'  # BERT bahasa Indonesia!
# Alternatif: 'bert-base-multilingual-cased' (multilingual)

tokenizer = BertTokenizer.from_pretrained(MODEL_NAME)
model = BertForSequenceClassification.from_pretrained(
    MODEL_NAME, num_labels=2
)

print(f"Model: {MODEL_NAME}")
print(f"Parameters: {sum(p.numel() for p in model.parameters()):,}")

# Create DataLoaders
train_dataset = TextClassificationDataset(X_train, y_train, tokenizer, max_len=128)
val_dataset = TextClassificationDataset(X_val, y_val, tokenizer, max_len=128)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=16, shuffle=False)

# =============================================
# 4. FINE-TUNING
# =============================================
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)

# Optimizer with weight decay
optimizer = AdamW(model.parameters(), lr=2e-5, weight_decay=0.01)

# Learning rate scheduler
num_epochs = 5
total_steps = len(train_loader) * num_epochs
scheduler = get_linear_schedule_with_warmup(
    optimizer, 
    num_warmup_steps=int(0.1 * total_steps),  # 10% warmup
    num_training_steps=total_steps
)

# Training loop
print("\nFine-tuning BERT...")
for epoch in range(num_epochs):
    # TRAIN
    model.train()
    total_loss = 0
    train_preds, train_true = [], []
    
    for batch in train_loader:
        input_ids = batch['input_ids'].to(device)
        attention_mask = batch['attention_mask'].to(device)
        labels_batch = batch['labels'].to(device)
        
        optimizer.zero_grad()
        outputs = model(input_ids=input_ids, attention_mask=attention_mask,
                       labels=labels_batch)
        loss = outputs.loss
        logits = outputs.logits
        
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
        optimizer.step()
        scheduler.step()
        
        total_loss += loss.item()
        preds = torch.argmax(logits, dim=1).cpu().numpy()
        train_preds.extend(preds)
        train_true.extend(labels_batch.cpu().numpy())
    
    train_acc = accuracy_score(train_true, train_preds)
    avg_loss = total_loss / len(train_loader)
    
    # VALIDATION
    model.eval()
    val_preds, val_true = [], []
    
    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 = batch['labels'].to(device)
            
            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            preds = torch.argmax(outputs.logits, dim=1).cpu().numpy()
            
            val_preds.extend(preds)
            val_true.extend(labels_batch.cpu().numpy())
    
    val_acc = accuracy_score(val_true, val_preds)
    
    print(f"Epoch [{epoch+1}/{num_epochs}] "
          f"Loss: {avg_loss:.4f} | "
          f"Train Acc: {train_acc:.4f} | "
          f"Val Acc: {val_acc:.4f}")

# =============================================
# 5. INFERENCE
# =============================================
model.eval()
new_texts = [
    "Produk luar biasa bagus, sangat puas!",
    "Barang rusak dan pelayanan sangat buruk",
    "Cukup bagus untuk harga segitu"
]

print("\n=== Predictions ===")
for text in new_texts:
    encoding = tokenizer(text, return_tensors='pt', max_length=128,
                        padding='max_length', truncation=True)
    input_ids = encoding['input_ids'].to(device)
    attention_mask = encoding['attention_mask'].to(device)
    
    with torch.no_grad():
        outputs = model(input_ids=input_ids, attention_mask=attention_mask)
        probs = torch.softmax(outputs.logits, dim=1)
        pred = torch.argmax(probs, dim=1).item()
    
    label = "Positif" if pred == 1 else "Negatif"
    conf = probs[0][pred].item()
    print(f"  [{label} ({conf:.2%})] {text}")