detector-ai / train_finetune.py
Slovand's picture
Update train_finetune.py
e08cad4 verified
Raw History Blame Contribute Delete
6.1 kB
# train_finetune.py - Fine-Tuning IndoBERT untuk Deteksi Teks AI
import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score
from transformers import AutoTokenizer, AutoModelForSequenceClassification, Trainer, TrainingArguments
from datasets import Dataset
import torch
import os
print("=" * 60)
print("FINE-TUNING INDOBERT UNTUK DETEKSI TEKS AI")
print("Lingkungan Perguruan Tinggi")
print("=" * 60)
# ============================================================
# 1. LOAD DATASET HC3 (pakai library Hugging Face langsung)
# ============================================================
print("\n📖 Loading HC3 dataset from Hugging Face...")
from datasets import load_dataset
# Load dataset HC3 (Human ChatGPT Comparison Corpus)
dataset = load_dataset("Hello-SimpleAI/HC3", "all", split="train")
print(f"✅ Dataset loaded: {len(dataset)} samples")
# Extract text and labels
texts = []
labels = []
for item in dataset:
# Human answers (label 0)
for human_answer in item['human_answers']:
if human_answer and len(human_answer.split()) > 20:
texts.append(human_answer)
labels.append(0)
# ChatGPT answers (label 1)
for chatgpt_answer in item['chatgpt_answers']:
if chatgpt_answer and len(chatgpt_answer.split()) > 20:
texts.append(chatgpt_answer)
labels.append(1)
print(f" Human (0): {labels.count(0)}")
print(f" AI (1): {labels.count(1)}")
# Convert to DataFrame
df = pd.DataFrame({'text': texts, 'label': labels})
# Balance dataset (ambil 3000 masing-masing)
df_human = df[df['label'] == 0].sample(n=3000, random_state=42)
df_ai = df[df['label'] == 1].sample(n=3000, random_state=42)
df = pd.concat([df_human, df_ai]).reset_index(drop=True)
print(f" Balanced: {len(df)} samples (Human: {len(df[df['label']==0])}, AI: {len(df[df['label']==1])})")
# Split data
X_train, X_test, y_train, y_test = train_test_split(
df['text'].values, df['label'].values, test_size=0.2, random_state=42, stratify=df['label'].values
)
print(f" Train: {len(X_train)}, Test: {len(X_test)}")
# ============================================================
# 2. LOAD INDOBERT MODEL
# ============================================================
print("\n🔧 Loading IndoBERT model...")
model_name = "indobenchmark/indobert-base-p2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)
# ============================================================
# 3. TOKENIZE
# ============================================================
print("\n📝 Tokenizing data...")
def tokenize_function(examples):
return tokenizer(examples['text'], padding='max_length', truncation=True, max_length=512)
train_dataset = Dataset.from_dict({'text': X_train.tolist(), 'label': y_train.tolist()})
test_dataset = Dataset.from_dict({'text': X_test.tolist(), 'label': y_test.tolist()})
train_dataset = train_dataset.map(tokenize_function, batched=True)
test_dataset = test_dataset.map(tokenize_function, batched=True)
train_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])
test_dataset.set_format('torch', columns=['input_ids', 'attention_mask', 'label'])
# ============================================================
# 4. TRAINING CONFIGURATION
# ============================================================
print("\n⚙️ Setting up training...")
training_args = TrainingArguments(
output_dir='./finetuned-indobert',
num_train_epochs=3,
per_device_train_batch_size=8,
per_device_eval_batch_size=8,
learning_rate=2e-5,
warmup_steps=100,
weight_decay=0.01,
eval_strategy="epoch",
save_strategy="epoch",
load_best_model_at_end=True,
metric_for_best_model="accuracy",
save_total_limit=2,
fp16=False,
report_to="none"
)
def compute_metrics(eval_pred):
logits, labels = eval_pred
predictions = np.argmax(logits, axis=-1)
accuracy = accuracy_score(labels, predictions)
precision = precision_score(labels, predictions, average='binary', zero_division=0)
recall = recall_score(labels, predictions, average='binary', zero_division=0)
f1 = f1_score(labels, predictions, average='binary', zero_division=0)
return {
'accuracy': accuracy,
'precision': precision,
'recall': recall,
'f1': f1
}
# ============================================================
# 5. START TRAINING
# ============================================================
print("\n🤖 Starting fine-tuning (this will take 10-20 minutes)...")
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
eval_dataset=test_dataset,
compute_metrics=compute_metrics,
)
trainer.train()
# ============================================================
# 6. EVALUATE
# ============================================================
print("\n📈 Evaluating model...")
eval_results = trainer.evaluate()
print("\n" + "=" * 60)
print("HASIL FINE-TUNING")
print("=" * 60)
print(f" Accuracy: {eval_results['eval_accuracy']:.4f} ({eval_results['eval_accuracy']*100:.2f}%)")
print(f" Precision: {eval_results['eval_precision']:.4f} ({eval_results['eval_precision']*100:.2f}%)")
print(f" Recall: {eval_results['eval_recall']:.4f} ({eval_results['eval_recall']*100:.2f}%)")
print(f" F1-Score: {eval_results['eval_f1']:.4f} ({eval_results['eval_f1']*100:.2f}%)")
# ============================================================
# 7. SAVE MODEL
# ============================================================
print("\n💾 Saving fine-tuned model...")
os.makedirs("models/indobert-finetuned", exist_ok=True)
model.save_pretrained("models/indobert-finetuned")
tokenizer.save_pretrained("models/indobert-finetuned")
print("✅ Model saved to models/indobert-finetuned")
print("\n" + "=" * 60)
print("FINE-TUNING SELESAI!")
print("Model siap digunakan untuk deteksi teks AI")
print("=" * 60)