sentimart / src /utils /model_loader.py
myreport12's picture
Upload 20 files
5360f9f verified
Raw
History Blame Contribute Delete
3.75 kB
"""
Loader model IndoBERT untuk SentiMart.
Model hasil fine-tuning (dari notebook, cell 6: `trainer.save_model(...)`)
diharapkan berada di folder `model/indobert_sentiment_final/` pada root
project ini. Struktur folder tersebut biasanya berisi:
config.json, model.safetensors (atau pytorch_model.bin),
tokenizer_config.json, vocab.txt, special_tokens_map.json
Cara mengisi folder ini:
1. Di notebook (Kaggle/Colab), setelah training selesai, download folder
'./indobert_sentiment_final/' (klik kanan -> Download, atau zip dulu).
2. Extract ke: sentimart/model/indobert_sentiment_final/
Jika model belum ada, app tetap bisa dijalankan dalam MODE DEMO (prediksi
dummy berbasis kata kunci) supaya wireframe & alur UI tetap bisa dicek.
"""
import os
import numpy as np
import streamlit as st
MODEL_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), "model", "indobert_sentiment_final")
MAX_LENGTH = 128
LABEL_MAP = {0: "Negative", 1: "Positive"}
def model_is_available() -> bool:
return os.path.isdir(MODEL_DIR) and any(
f.startswith("config.json") for f in os.listdir(MODEL_DIR)
) if os.path.isdir(MODEL_DIR) else False
@st.cache_resource(show_spinner=False)
def load_model():
"""Load tokenizer + model sekali saja, disimpan di cache Streamlit."""
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR)
model.eval()
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
return tokenizer, model, device
def _demo_predict(text: str):
"""Fallback berbasis kata kunci sederhana, dipakai kalau model belum di-copy."""
positive_words = ["bagus", "cepat", "puas", "recommended", "ramah", "sesuai", "mantap", "keren"]
negative_words = ["jelek", "rusak", "lambat", "kecewa", "buruk", "tidak sesuai", "lama", "cacat"]
t = text.lower()
pos_hits = sum(w in t for w in positive_words)
neg_hits = sum(w in t for w in negative_words)
if pos_hits == neg_hits:
label, conf = ("Positive", 0.55) if len(t) % 2 == 0 else ("Negative", 0.55)
elif pos_hits > neg_hits:
label, conf = "Positive", min(0.6 + 0.1 * pos_hits, 0.97)
else:
label, conf = "Negative", min(0.6 + 0.1 * neg_hits, 0.97)
probs = {label: conf, ("Negative" if label == "Positive" else "Positive"): 1 - conf}
return label, conf, probs
def predict_sentiment(text: str):
"""Prediksi satu review. Return: (label, confidence, {'Positive': p, 'Negative': p})"""
from .preprocessing import light_normalize
if not model_is_available():
return _demo_predict(text)
import torch
tokenizer, model, device = load_model()
clean_text = light_normalize(text)
inputs = tokenizer(
clean_text, truncation=True, padding="max_length",
max_length=MAX_LENGTH, return_tensors="pt",
).to(device)
with torch.no_grad():
logits = model(**inputs).logits
probs = torch.softmax(logits, dim=-1).cpu().numpy()[0]
pred_idx = int(np.argmax(probs))
label = LABEL_MAP[pred_idx]
confidence = float(probs[pred_idx])
prob_dict = {"Negative": float(probs[0]), "Positive": float(probs[1])}
return label, confidence, prob_dict
def predict_batch(texts: list[str], progress_callback=None):
"""Prediksi banyak review sekaligus. Return list of (label, confidence, prob_dict)."""
results = []
total = len(texts)
for i, t in enumerate(texts):
results.append(predict_sentiment(t))
if progress_callback is not None:
progress_callback((i + 1) / total)
return results