IndoBERT-ABSA / model_utils.py
Brian045's picture
Create model_utils.py
e34817b verified
import torch
import torch.nn as nn
from transformers import AutoTokenizer, BertModel, BertPreTrainedModel, BertConfig
from pathlib import Path
import logging
# Konfigurasi logging untuk menampilkan pesan error yang jelas
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# Definisi pemetaan (bisa juga diletakkan di sini agar terpusat)
ASPECT_MAP = {
0: 'Fasilitas', 1: 'Harga', 2: 'Hotel', 3: 'Kamar', 4: 'Kolam Renang',
5: 'Layanan', 6: 'Lokasi', 7: 'Makanan', 8: 'Sarapan', 9: 'Staf'
}
SENTIMENT_MAP = {0: 'Negative', 1: 'Positive'}
# Definisi kelas model, sama seperti sebelumnya
class IndoBERTABSA(BertPreTrainedModel):
def __init__(self, config):
super().__init__(config)
self.num_aspect_labels = len(ASPECT_MAP)
self.num_sentiment_labels = len(SENTIMENT_MAP)
self.bert = BertModel(config)
self.dropout = nn.Dropout(config.hidden_dropout_prob)
self.aspect_classifier = nn.Linear(config.hidden_size, self.num_aspect_labels)
self.sentiment_classifier = nn.Linear(config.hidden_size, self.num_sentiment_labels)
self.init_weights()
def forward(self, input_ids=None, attention_mask=None, token_type_ids=None):
outputs = self.bert(
input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids,
)
pooled_output = outputs[1]
pooled_output = self.dropout(pooled_output)
aspect_logits = self.aspect_classifier(pooled_output)
sentiment_logits = self.sentiment_classifier(pooled_output)
return aspect_logits, sentiment_logits
def load_model_and_tokenizer():
"""
Fungsi terpusat untuk memuat tokenizer dan model.
Mengembalikan (model, tokenizer) jika berhasil, atau (None, None) jika gagal.
"""
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f"Menggunakan perangkat: {device}")
try:
tokenizer = AutoTokenizer.from_pretrained("indobenchmark/indobert-base-p2")
logging.info("Tokenizer berhasil dimuat.")
except Exception as e:
logging.error(f"GAGAL memuat tokenizer: {e}")
return None, None
try:
config = BertConfig.from_pretrained("indobenchmark/indobert-base-p2")
logging.info("Konfigurasi BERT berhasil dimuat.")
except Exception as e:
logging.error(f"GAGAL memuat konfigurasi BERT: {e}")
return None, None
model = IndoBERTABSA(config)
model_save_path = Path('absa_indobert_model_75.pth')
if not model_save_path.exists():
logging.error(f"File model tidak ditemukan di '{model_save_path}'")
return None, None
try:
model.load_state_dict(torch.load(model_save_path, map_location=device))
model.to(device)
model.eval()
logging.info(f"Model berhasil dimuat dari {model_save_path} dan disetel ke mode evaluasi.")
return model, tokenizer
except Exception as e:
logging.error(f"GAGAL memuat bobot model dari {model_save_path}: {e}")
return None, None