import torch
from transformers import AutoTokenizer
from huggingface_hub import hf_hub_download
import json

# Загрузка модели и токенизатора
model_path = hf_hub_download(repo_id="AlucardV/kinopotok-toxicity-multitask-model2", filename="pytorch_model.bin")
config_path = hf_hub_download(repo_id="AlucardV/kinopotok-toxicity-multitask-model2", filename="config.json")
tokenizer = AutoTokenizer.from_pretrained("AlucardV/kinopotok-toxicity-multitask-model2")

# Определение класса модели
class MultiTaskToxicityEncoder(torch.nn.Module):
    def __init__(self, model_name):
        super().__init__()
        from transformers import AutoModel
        self.encoder = AutoModel.from_pretrained(model_name)
        hidden_size = self.encoder.config.hidden_size
        self.profanity_head = torch.nn.Linear(hidden_size, 1)
        self.threat_head = torch.nn.Linear(hidden_size, 1)
        self.illegal_head = torch.nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
        cls = outputs.last_hidden_state[:, 0, :]
        return {
            "profanity": self.profanity_head(cls).squeeze(-1),
            "threat": self.threat_head(cls).squeeze(-1),
            "illegal": self.illegal_head(cls).squeeze(-1),
        }

# Загрузка весов
config = json.load(open(config_path))
model = MultiTaskToxicityEncoder(config["model_name"])
model.load_state_dict(torch.load(model_path, map_location="cpu"))
model.eval()

# Функция предсказания
def predict(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    with torch.no_grad():
        outputs = model(inputs["input_ids"], inputs["attention_mask"])
    probs = {k: torch.sigmoid(v).item() for k, v in outputs.items()}
    thresholds = {
        "profanity": 0.7500000000000002,
        "threat": 0.5000000000000001,
        "illegal": 0.1,
    }
    results = {
        k: {
            "toxic": probs[k] >= thresholds[k],
            "confidence": f"{probs[k]*100:.1f}%"
        } for k in probs
    }
    return results

# Пример
print(predict("Ты что, совсем охренел, мудак?"))