wehub-resource-sync
27 lines
991 B
Python
27 lines
991 B
Python
import torch
|
|
from deepeval.models.base_model import DeepEvalBaseModel
|
|
from detoxify import Detoxify
|
|
|
|
|
|
class DetoxifyModel(DeepEvalBaseModel):
|
|
def __init__(self, model_name: str | None = None, *args, **kwargs):
|
|
if model_name is not None:
|
|
assert model_name in [
|
|
"original",
|
|
"unbiased",
|
|
"multilingual",
|
|
], "Invalid model. Available variants: original, unbiased, multilingual"
|
|
model_name = "original" if model_name is None else model_name
|
|
super().__init__(model_name, *args, **kwargs)
|
|
|
|
def load_model(self):
|
|
device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
return Detoxify(self.model_name, device=device)
|
|
|
|
def _call(self, text: str):
|
|
toxicity_score_dict = self.model.predict(text)
|
|
mean_toxicity_score = sum(list(toxicity_score_dict.values())) / len(
|
|
toxicity_score_dict
|
|
)
|
|
return mean_toxicity_score, toxicity_score_dict
|