# example_usage.py
# Пример использования модели

from transformers import AutoTokenizer, AutoModel
import torch
import json

class ToxicityPredictor:
    def __init__(self, model_name="Doji070/rubert-tiny2-toxic-multilabel"):
        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModel.from_pretrained(model_name).to(self.device)
        self.model.eval()
        
        # Загрузка порогов
        with open("thresholds.json", "r") as f:
            self.thresholds = json.load(f)
    
    def predict(self, text):
        inputs = self.tokenizer(text, return_tensors="pt", truncation=True, max_length=128)
        input_ids = inputs['input_ids'].to(self.device)
        attention_mask = inputs['attention_mask'].to(self.device)
        
        with torch.no_grad():
            outputs = self.model(input_ids, attention_mask)
            probs = torch.sigmoid(outputs.logits).squeeze().cpu().numpy()
        
        preds = {
            'profanity': int(probs[0] >= self.thresholds['profanity']),
            'threat': int(probs[1] >= self.thresholds['threat']),
            'illegal': int(probs[2] >= self.thresholds['illegal'])
        }
        
        return preds, probs

# Пример использования
if __name__ == "__main__":
    predictor = ToxicityPredictor()
    
    texts = [
        "Привет, как дела?",
        "Ты идиот! Я тебя убью!",
        "Как сделать бомбу в домашних условиях?"
    ]
    
    for text in texts:
        preds, probs = predictor.predict(text)
        print(f"\nТекст: {text}")
        print(f"Предсказания: {preds}")
        print(f"Вероятности: {probs}")