File size: 3,370 Bytes
97f37a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
import os
import torch
import pandas as pd
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoConfig
class PredictSentiment:
def __init__(self,
model_name='alwanrahmana/sentiment-classification-deberta-base',
label_map=None,
device=None):
# Local path buat nyimpan model
local_dir = f"./app/inference/pretrained_models/{model_name.replace('/', '_')}"
# Check apakah model sudah ada di local
if os.path.exists(local_dir) and os.listdir(local_dir):
print(f"✅ Local model found at {local_dir}. Loading from local...")
config = AutoConfig.from_pretrained(local_dir)
self.tokenizer = AutoTokenizer.from_pretrained(local_dir)
self.model = AutoModelForSequenceClassification.from_pretrained(local_dir, config=config)
else:
print(f"⬇️ Downloading model from Hugging Face Hub: {model_name} ...")
config = AutoConfig.from_pretrained(model_name)
self.tokenizer = AutoTokenizer.from_pretrained(model_name)
self.model = AutoModelForSequenceClassification.from_pretrained(model_name, config=config)
os.makedirs(local_dir, exist_ok=True)
self.tokenizer.save_pretrained(local_dir)
self.model.save_pretrained(local_dir)
print(f"✅ Model saved to {local_dir}")
self.model.eval()
# Device config
self.DEVICE = device if device else ("cuda" if torch.cuda.is_available() else "cpu")
self.model.to(self.DEVICE)
# Label mapping
self.label_map = label_map if label_map else {0: "Negative", 1: "Neutral", 2: "Positive"}
def _predict_sentiment(self, aspect, text, max_len=256):
encoding = self.tokenizer(
aspect,
text,
truncation=True,
padding='max_length',
max_length=max_len,
return_tensors='pt'
)
input_ids = encoding['input_ids'].to(self.DEVICE)
attention_mask = encoding['attention_mask'].to(self.DEVICE)
token_type_ids = encoding.get('token_type_ids', None)
if token_type_ids is not None:
token_type_ids = token_type_ids.to(self.DEVICE)
with torch.no_grad():
outputs = self.model(
input_ids=input_ids,
attention_mask=attention_mask,
token_type_ids=token_type_ids
)
logits = outputs.logits
pred = torch.argmax(logits, dim=1).item()
return self.label_map.get(pred, "Unknown")
def predict(self, df):
"""
Input: DataFrame hasil PredictAspect
Output: DataFrame dengan kolom text, aspect, predicted_sentiment
"""
results = []
for idx, row in df.iterrows():
text = row['text']
aspects = row['aspects']
for aspect in aspects:
sentiment = self._predict_sentiment(aspect, text)
results.append({
'text': text,
'aspect': aspect,
'predicted_sentiment': sentiment
})
result_df = pd.DataFrame(results)
return result_df
|