Spaces:
Running
Running
File size: 2,350 Bytes
a653bb2 | 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 | """Lyrical emotion classification model."""
import torch
from transformers import AutoModelForSequenceClassification, AutoConfig
from pathlib import Path
from .config import config
class LyricEmotionClassifier:
"""RoBERTa-based emotion classifier for song lyrics."""
def __init__(self, model_name: str = None, num_labels: int = None):
self.model_name = model_name or config.model_name
self.num_labels = num_labels or config.num_labels
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.model = None
def load_pretrained(self) -> "LyricEmotionClassifier":
"""Load pretrained model for fine-tuning."""
model_config = AutoConfig.from_pretrained(
self.model_name,
num_labels=self.num_labels,
id2label={i: label for i, label in enumerate(config.label_names)},
label2id={label: i for i, label in enumerate(config.label_names)},
)
self.model = AutoModelForSequenceClassification.from_pretrained(
self.model_name,
config=model_config
)
self.model.to(self.device)
return self
def load_finetuned(self, path: Path) -> "LyricEmotionClassifier":
"""Load a fine-tuned model from disk."""
self.model = AutoModelForSequenceClassification.from_pretrained(path)
self.model.to(self.device)
return self
def save(self, path: Path):
"""Save model to disk."""
path.mkdir(parents=True, exist_ok=True)
self.model.save_pretrained(path)
def predict(self, inputs: dict) -> dict:
"""Run inference on tokenized inputs."""
self.model.eval()
# Move inputs to device
inputs = {k: v.to(self.device) for k, v in inputs.items()}
with torch.no_grad():
outputs = self.model(**inputs)
probs = torch.softmax(outputs.logits, dim=-1)
predictions = torch.argmax(probs, dim=-1)
return {
"predictions": predictions.cpu().numpy(),
"probabilities": probs.cpu().numpy(),
"labels": [config.label_names[p] for p in predictions.cpu().numpy()]
}
def get_model() -> LyricEmotionClassifier:
"""Factory function to get classifier instance."""
return LyricEmotionClassifier().load_pretrained()
|