inference.py

import os
import json
import re
from datetime import datetime
from typing import List, Dict, Any, Tuple, Optional

import numpy as np
import pandas as pd

import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.nn.functional import sigmoid

from transformers import AutoTokenizer, AutoModel

from huggingface_hub import hf_hub_download

# =========================
# Общая очистка текста
# =========================

URL_PATTERN = re.compile(r"http\S+|www\.\S+")
MULTISPACE_PATTERN = re.compile(r"\s+")


def clean_text(text: str) -> str:
    if not isinstance(text, str):
        return ""
    text = text.replace("\n", " ").replace("\r", " ")
    text = URL_PATTERN.sub(" ", text)
    text = MULTISPACE_PATTERN.sub(" ", text)
    return text.strip()


# =========================
# Релевантность
# =========================

class RelevanceDatasetInfer(Dataset):
    def __init__(self, texts: List[str], tokenizer, max_length: int):
        self.texts = texts
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        enc = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        return {
            "input_ids": enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
        }


class RuBERTBinaryClassifier(nn.Module):
    def __init__(self, model_name: str, hidden_size: int = 512):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size
        self.dropout1 = nn.Dropout(0.3)
        self.dense = nn.Linear(h, hidden_size)
        self.relu = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(hidden_size, 1)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]
        x = self.dropout1(pooled)
        x = self.dense(x)
        x = self.relu(x)
        x = self.dropout2(x)
        logits = self.classifier(x).squeeze(-1)
        return logits


def load_relevance_model(relevance_repo_id: str, device: torch.device):
    """
    Ожидается в репо:
    - config.json с полями: model_name, max_length, threshold
    - pytorch_model.bin
    - файлы токенайзера (config, vocab и т.п.)
    """
    config_path = hf_hub_download(relevance_repo_id, "config.json")
    with open(config_path, "r", encoding="utf-8") as f:
        conf = json.load(f)

    model_name = conf["model_name"]
    max_length = int(conf["max_length"])
    threshold = float(conf["threshold"])

    tokenizer = AutoTokenizer.from_pretrained(relevance_repo_id)
    model = RuBERTBinaryClassifier(model_name)

    weights_path = hf_hub_download(relevance_repo_id, "pytorch_model.bin")
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()
    return model, tokenizer, max_length, threshold


def infer_relevance(
    texts: List[str],
    model: nn.Module,
    tokenizer,
    max_length: int,
    threshold: float,
    device: torch.device,
    batch_size: int = 32,
) -> Tuple[np.ndarray, np.ndarray]:
    if not texts:
        return np.array([], dtype=int), np.array([], dtype=float)

    ds = RelevanceDatasetInfer(texts, tokenizer, max_length)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)

    all_logits = []
    with torch.no_grad():
        for batch in loader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            logits = model(input_ids, attention_mask)
            all_logits.append(logits.cpu().numpy())

    if not all_logits:
        return np.array([], dtype=int), np.array([], dtype=float)

    all_logits = np.concatenate(all_logits, axis=0)
    probs = 1 / (1 + np.exp(-all_logits))
    preds = (probs >= threshold).astype(int)
    return preds, probs


# =========================
# Теги (multilabel)
# =========================

class RuBERTMultiLabelModel(nn.Module):
    def __init__(self, model_name: str, num_tags: int, hidden_size: int = 512):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size
        self.dropout1 = nn.Dropout(0.3)
        self.dense = nn.Linear(h, hidden_size)
        self.relu = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(hidden_size, num_tags)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]
        x = self.dropout1(pooled)
        x = self.dense(x)
        x = self.relu(x)
        x = self.dropout2(x)
        return self.classifier(x)


def load_tags_model(tags_repo_id: str, device: torch.device):
    """
    Ожидается структура репо:
    - training_config.json: {model_name, max_length, threshold, hidden_size, ...}
    - id2tag.json: {"0": "tag_name", ...}
    - pytorch_model.bin: state_dict RuBERTMultiLabelModel
    """
    tc_path = hf_hub_download(tags_repo_id, "training_config.json")
    with open(tc_path, "r", encoding="utf-8") as f:
        tc = json.load(f)

    model_name = tc.get("model_name", "DeepPavlov/rubert-base-cased")
    max_length = int(tc.get("max_length", 512))
    threshold = float(tc.get("threshold", 0.3))
    hidden_size = int(tc.get("hidden_size", 512))

    id2tag_path = hf_hub_download(tags_repo_id, "id2tag.json")
    with open(id2tag_path, "r", encoding="utf-8") as f:
        id2tag_raw = json.load(f)
    id2tag = {int(k): v for k, v in id2tag_raw.items()}
    num_tags = len(id2tag)
    all_tags = [id2tag[i] for i in sorted(id2tag.keys())]

    tokenizer = AutoTokenizer.from_pretrained(tags_repo_id)

    model = RuBERTMultiLabelModel(model_name, num_tags, hidden_size=hidden_size)
    weights_path = hf_hub_download(tags_repo_id, "pytorch_model.bin")
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()

    return model, tokenizer, max_length, all_tags, threshold


def predict_for_long_text_tags(
    model: nn.Module,
    text: str,
    tokenizer,
    device: torch.device,
    max_length: int,
) -> np.ndarray:
    enc_full = tokenizer(
        text,
        add_special_tokens=True,
        truncation=False,
        return_tensors="pt",
    )
    input_ids_full = enc_full["input_ids"][0]
    attention_full = enc_full["attention_mask"][0]
    seq_len = input_ids_full.shape[0]

    if seq_len <= max_length:
        input_ids = input_ids_full.unsqueeze(0).to(device)
        attention_mask = attention_full.unsqueeze(0).to(device)
        with torch.no_grad():
            logits = model(input_ids, attention_mask)
            probs = sigmoid(logits).squeeze(0).cpu().numpy()
        return probs

    logits_list = []
    for start in range(0, seq_len, max_length):
        end = start + max_length
        ids_chunk = input_ids_full[start:end]
        att_chunk = attention_full[start:end]
        if ids_chunk.shape[0] == 0:
            continue
        input_ids = ids_chunk.unsqueeze(0).to(device)
        attention_mask = att_chunk.unsqueeze(0).to(device)
        with torch.no_grad():
            logits = model(input_ids, attention_mask)
            logits_list.append(logits.squeeze(0).cpu().numpy())

    if not logits_list:
        enc = tokenizer(
            text,
            add_special_tokens=True,
            max_length=max_length,
            truncation=True,
            return_tensors="pt",
        )
        with torch.no_grad():
            logits = model(enc["input_ids"].to(device), enc["attention_mask"].to(device))
        return sigmoid(logits).squeeze(0).cpu().numpy()

    logits_avg = np.mean(np.stack(logits_list, axis=0), axis=0)
    probs = 1 / (1 + np.exp(-logits_avg))
    return probs


def infer_tags_for_texts(
    texts_clean: List[str],
    model: nn.Module,
    tokenizer,
    max_length: int,
    all_tags: List[str],
    threshold: float,
    device: torch.device,
    return_probs: bool = False,
):
    results_tags = []
    results_probs = []

    for text in texts_clean:
        probs = predict_for_long_text_tags(model, text, tokenizer, device, max_length)
        mask = probs >= threshold
        idxs = np.where(mask)[0]
        tags = [all_tags[i] for i in idxs]

        results_tags.append(tags)
        results_probs.append(probs)

    if return_probs:
        return results_tags, results_probs
    return results_tags


# =========================
# Тональность
# =========================

def index_to_tone(i: int) -> int:
    i = int(i)
    if i not in (0, 1, 2):
        raise ValueError(f"Некорректный индекс класса: {i}")
    return i - 1


class SentimentDatasetInfer(Dataset):
    def __init__(self, texts_clean: List[str], tokenizer, max_length: int):
        self.texts = texts_clean
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        text = str(self.texts[idx])
        enc = self.tokenizer(
            text,
            add_special_tokens=True,
            max_length=self.max_length,
            truncation=True,
            padding="max_length",
            return_tensors="pt",
        )
        return {
            "input_ids": enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
        }


class RuBERTSentimentClassifier(nn.Module):
    def __init__(self, model_name: str, num_labels: int = 3, hidden_size: int = 512):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size
        self.dropout1 = nn.Dropout(0.3)
        self.dense = nn.Linear(h, hidden_size)
        self.relu = nn.ReLU()
        self.dropout2 = nn.Dropout(0.2)
        self.classifier = nn.Linear(hidden_size, num_labels)

    def forward(self, input_ids, attention_mask):
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        pooled = out.last_hidden_state[:, 0]
        x = self.dropout1(pooled)
        x = self.dense(x)
        x = self.relu(x)
        x = self.dropout2(x)
        return self.classifier(x)


def load_sentiment_model(sentiment_repo_id: str, device: torch.device):
    """
    Ожидается в репо:
    - config.json с полями: model_name, max_length
    - pytorch_model.bin
    """
    config_path = hf_hub_download(sentiment_repo_id, "config.json")
    with open(config_path, "r", encoding="utf-8") as f:
        conf = json.load(f)

    model_name = conf["model_name"]
    max_length = int(conf["max_length"])

    tokenizer = AutoTokenizer.from_pretrained(sentiment_repo_id)
    model = RuBERTSentimentClassifier(model_name)
    weights_path = hf_hub_download(sentiment_repo_id, "pytorch_model.bin")
    state = torch.load(weights_path, map_location=device)
    model.load_state_dict(state, strict=True)
    model.to(device)
    model.eval()
    return model, tokenizer, max_length


def infer_sentiment(
    texts_clean: List[str],
    model: nn.Module,
    tokenizer,
    max_length: int,
    device: torch.device,
    batch_size: int = 32,
) -> Dict[str, np.ndarray]:
    if not texts_clean:
        return {
            "logits": np.zeros((0, 3), dtype=float),
            "probs": np.zeros((0, 3), dtype=float),
            "preds_idx": np.zeros((0,), dtype=int),
            "preds_tone": np.zeros((0,), dtype=int),
        }

    ds = SentimentDatasetInfer(texts_clean, tokenizer, max_length)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False)

    all_logits = []
    with torch.no_grad():
        for batch in loader:
            input_ids = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            logits = model(input_ids, attention_mask)
            all_logits.append(logits.cpu().numpy())

    all_logits = np.concatenate(all_logits, axis=0)
    probs = torch.softmax(torch.tensor(all_logits), dim=-1).numpy()
    preds_idx = probs.argmax(axis=1)
    preds_tone = np.array([index_to_tone(i) for i in preds_idx])

    return {
        "logits": all_logits,
        "probs": probs,
        "preds_idx": preds_idx,
        "preds_tone": preds_tone,
    }


def tone_to_str(tone_int: int) -> str:
    if tone_int == -1:
        return "негатив"
    elif tone_int == 0:
        return "нейтрально"
    elif tone_int == 1:
        return "позитив"
    return ""


def infer_sentiment_single(
    text_clean: str,
    model: nn.Module,
    tokenizer,
    max_length: int,
    device: torch.device,
    tone_threshold: Optional[float] = None,
) -> Dict[str, Any]:
    res = infer_sentiment(
        [text_clean],
        model,
        tokenizer,
        max_length,
        device,
        batch_size=1,
    )
    probs = res["probs"][0]
    preds_idx = int(res["preds_idx"][0])
    tone_int = int(res["preds_tone"][0])

    max_prob = float(probs[preds_idx])
    if tone_threshold is not None and max_prob < tone_threshold:
        tone_int = 0

    tone_str = tone_to_str(tone_int)
    tone_probs = {
        "негатив": float(probs[0]),
        "нейтрально": float(probs[1]),
        "позитив": float(probs[2]),
    }
    return {
        "tone_int": tone_int,
        "tone_str": tone_str,
        "tone_probs": tone_probs,
    }


# =========================
# Работа с JSON входом
# =========================

def load_input_json(json_path: str) -> List[Dict[str, Any]]:
    with open(json_path, "r", encoding="utf-8") as f:
        data = json.load(f)
    if isinstance(data, list):
        return data
    elif isinstance(data, dict) and "items" in data:
        return data["items"]
    else:
        raise ValueError("Неизвестный формат входного JSON")


def extract_full_text_from_record(record: Dict[str, Any]) -> str:
    title = record.get("title", "") or ""
    text = record.get("text", "") or ""
    full_text_raw = (title + ". " + text).strip(". ")
    return clean_text(full_text_raw)


def format_ba_datetime(ts: Any) -> str:
    if not ts:
        return ""
    try:
        dt = datetime.fromtimestamp(ts)
        return dt.strftime("%d.%m.%Y %H:%M")
    except Exception:
        return ""


# =========================
# Пайплайн для JSON -> CSV
# =========================

def run_pipeline_json_to_new_ba_like_csv(
    input_json: str,
    relevance_model_dir: str,
    tags_model_dir: str,
    sentiment_model_dir: str,
    output_csv_path: str,
    limit_messages: Optional[int] = None,
    relevance_threshold: Optional[float] = None,
    tags_threshold: Optional[float] = None,
    tone_threshold: Optional[float] = None,
) -> str:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    rel_model, rel_tokenizer, rel_max_len, rel_thr_cfg = load_relevance_model(relevance_model_dir, device)
    tags_model, tags_tokenizer, tags_max_len, all_tags, tags_thr_cfg = load_tags_model(tags_model_dir, device)
    sent_model, sent_tokenizer, sent_max_len = load_sentiment_model(sentiment_model_dir, device)

    rel_thr = relevance_threshold if relevance_threshold is not None else rel_thr_cfg
    tags_thr = tags_threshold if tags_threshold is not None else tags_thr_cfg

    records = load_input_json(input_json)
    if limit_messages is not None:
        records = records[:limit_messages]

    rows_base: List[Dict[str, Any]] = []
    texts_clean: List[str] = []
    hubtype_all: List[str] = []

    for i, rec in enumerate(records):
        title = rec.get("title", "") or ""
        text = rec.get("text", "") or ""
        full_text_raw = (title + ". " + text).strip(". ")
        full_text_clean = clean_text(full_text_raw)

        time_ts = rec.get("timeCreate") or rec.get("date") or 0
        if isinstance(time_ts, (int, float)) and time_ts > 0:
            date_str = datetime.fromtimestamp(time_ts).strftime("%d.%m.%Y %H:%M")
        else:
            date_str = ""

        base_row = {
            "Дата": date_str,
            "ID сообщения": rec.get("id", i),
            "Hash сообщения": rec.get("hash", ""),
            "Заголовок": title,
            "Текст": text,
            "Источник": rec.get("hub", ""),
            "Url": rec.get("url", rec.get("href", "")),
            "Тип источника": rec.get("sourceType", ""),
            "Тип сообщения": rec.get("type", ""),
            "Сюжет": rec.get("topic", ""),
            "Автор": rec.get("authorName", ""),
            "Url автора": rec.get("authorUrl", ""),
            "Тип автора": rec.get("authorType", ""),
            "Место публикации": rec.get("placeName", ""),
            "Url места публикации": rec.get("placeUrl", ""),
            "Пол": rec.get("gender", ""),
            "Возраст": rec.get("age", ""),
            "Аудитория": rec.get("audience", ""),
            "Комментариев": rec.get("commentsCount", ""),
            "Цитируемость СМИ": "",
            "Репостов": rec.get("repostsCount", ""),
            "Лайков": rec.get("likesCount", ""),
            "Вовлеченность": "",
            "Просмотров": rec.get("viewsCount", ""),
            "Оценка": "",
            "Дублей": "",
            "Аудитория СМИ": "",
            "Тональность": "",
            "Роль объекта": "",
            "Агрессия": "",
            "Страна": rec.get("country", ""),
            "Регион": rec.get("region", ""),
            "Город": rec.get("city", ""),
            "Язык": rec.get("lang", "Русский"),
            "WOM": "",
            "Обработано": "",
            "Место": rec.get("place", ""),
            "Адрес": rec.get("address", ""),
            "Product": rec.get("product", ""),
            "General": rec.get("general", ""),
        }

        rows_base.append(base_row)
        texts_clean.append(full_text_clean)
        hubtype_all.append(rec.get("hubtype", "") or "")

    rel_preds, _ = infer_relevance(
        texts_clean, rel_model, rel_tokenizer, rel_max_len, rel_thr, device
    )

    relevant_indices = np.where(rel_preds == 1)[0]
    tags_all = [[] for _ in records]
    tone_str_all = ["" for _ in records]

    if len(relevant_indices) > 0:
        texts_rel = [texts_clean[idx] for idx in relevant_indices]

        tags_rel = infer_tags_for_texts(
            texts_rel, tags_model, tags_tokenizer, tags_max_len, all_tags, tags_thr, device
        )
        for local_i, global_i in enumerate(relevant_indices):
            tags_for_item = tags_rel[local_i]
            if hubtype_all[global_i] == "Онлайн-СМИ" and "Новостные публикации" in all_tags:
                if "Новостные публикации" not in tags_for_item:
                    tags_for_item.append("Новостные публикации")
            tags_all[global_i] = tags_for_item

        sent_res = infer_sentiment(
            texts_rel, sent_model, sent_tokenizer, sent_max_len, device
        )
        for local_i, global_i in enumerate(relevant_indices):
            tone_int = int(sent_res["preds_tone"][local_i])
            probs_vec = sent_res["probs"][local_i]
            max_prob = float(probs_vec[sent_res["preds_idx"][local_i]])
            if tone_threshold is not None and max_prob < tone_threshold:
                tone_int = 0
            tone_str_all[global_i] = tone_to_str(tone_int)

    df = pd.DataFrame(rows_base)
    df["Релевантность"] = rel_preds.astype(int)
    df["Тональность"] = tone_str_all

    for tag in all_tags:
        df[tag] = 0

    for i, tags in enumerate(tags_all):
        for tag in tags:
            if tag in df.columns:
                df.at[i, tag] = 1

    df.to_csv(output_csv_path, sep=";", index=False, encoding="utf-8")
    return output_csv_path


# =========================
# Оценка модели (метрики по тегам)
# =========================

def evaluate_json_file_with_metrics(
    input_json: str,
    relevance_model_dir: str,
    tags_model_dir: str,
    sentiment_model_dir: str,
    limit_messages: Optional[int] = None,
    relevance_threshold: Optional[float] = None,
    tags_threshold: Optional[float] = None,
    tone_threshold: Optional[float] = None,
):
    """
    Считает micro-метрики и per-tag TP/FP/TN/FN/Precision/Recall/F1
    только по тегам (тональность здесь не участвует).
    Истинные теги берём из поля 'tags' записи (список строк).
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    rel_model, rel_tokenizer, rel_max_len, rel_thr_cfg = load_relevance_model(relevance_model_dir, device)
    tags_model, tags_tokenizer, tags_max_len, all_tags, tags_thr_cfg = load_tags_model(tags_model_dir, device)
    _sent_model, _sent_tokenizer, _sent_max_len = load_sentiment_model(sentiment_model_dir, device)

    rel_thr = relevance_threshold if relevance_threshold is not None else rel_thr_cfg
    tags_thr = tags_threshold if tags_threshold is not None else tags_thr_cfg

    records = load_input_json(input_json)
    if limit_messages is not None:
        records = records[:limit_messages]

    texts_clean: List[str] = []
    true_tags_all: List[List[str]] = []
    hubtype_all: List[str] = []

    for rec in records:
        full_text = extract_full_text_from_record(rec)
        texts_clean.append(full_text)
        raw_tags = rec.get("tags", [])
        if isinstance(raw_tags, list):
            true_tags = [str(t).strip() for t in raw_tags if str(t).strip()]
        elif isinstance(raw_tags, str):
            true_tags = [t.strip() for t in raw_tags.split(";") if t.strip()]
        else:
            true_tags = []
        true_tags_all.append(true_tags)
        hubtype_all.append(rec.get("hubtype", "") or "")

    rel_preds, _ = infer_relevance(
        texts_clean, rel_model, rel_tokenizer, rel_max_len, rel_thr, device
    )

    relevant_indices = np.where(rel_preds == 1)[0]
    pred_tags_all = [[] for _ in records]

    if len(relevant_indices) > 0:
        texts_rel = [texts_clean[idx] for idx in relevant_indices]
        tags_rel = infer_tags_for_texts(
            texts_rel, tags_model, tags_tokenizer, tags_max_len, all_tags, tags_thr, device
        )
        for local_i, global_i in enumerate(relevant_indices):
            tags_for_item = tags_rel[local_i]
            if hubtype_all[global_i] == "Онлайн-СМИ" and "Новостные публикации" in all_tags:
                if "Новостные публикации" not in tags_for_item:
                    tags_for_item.append("Новостные публикации")
            pred_tags_all[global_i] = tags_for_item

    tag_metrics: List[Dict[str, Any]] = []
    total_tp = total_fp = total_fn = total_tn = 0
    n_examples = len(records)

    for tag in all_tags:
        tp = fp = fn = tn = 0
        for true_tags, pred_tags in zip(true_tags_all, pred_tags_all):
            y_true = tag in true_tags
            y_pred = tag in pred_tags
            if y_true and y_pred:
                tp += 1
            elif not y_true and y_pred:
                fp += 1
            elif y_true and not y_pred:
                fn += 1
            else:
                tn += 1
        total_tp += tp
        total_fp += fp
        total_fn += fn
        total_tn += tn
        prec = tp / (tp + fp) if tp + fp > 0 else 0.0
        rec = tp / (tp + fn) if tp + fn > 0 else 0.0
        f1 = 2 * prec * rec / (prec + rec) if prec + rec > 0 else 0.0
        tag_metrics.append(
            {
                "tag": tag,
                "TP": tp,
                "FP": fp,
                "TN": tn,
                "FN": fn,
                "Precision": prec,
                "Recall": rec,
                "F1": f1,
            }
        )

    micro_p = total_tp / (total_tp + total_fp) if total_tp + total_fp > 0 else 0.0
    micro_r = total_tp / (total_tp + total_fn) if total_tp + total_fn > 0 else 0.0
    micro_f1 = 2 * micro_p * micro_r / (micro_p + micro_r) if micro_p + micro_r > 0 else 0.0

    micro_metrics = {
        "micro_precision": micro_p,
        "micro_recall": micro_r,
        "micro_f1": micro_f1,
        "n_examples": n_examples,
    }
    df_per_tag = pd.DataFrame(tag_metrics).sort_values("tag")
    return micro_metrics, df_per_tag


# =========================
# Утилита для одиночного текста
# =========================

class InferenceService:
    def __init__(
        self,
        relevance_model_dir: str,
        tags_model_dir: str,
        sentiment_model_dir: str,
        use_cuda: bool = False,
    ):
        self.device = torch.device(
            "cuda" if use_cuda and torch.cuda.is_available() else "cpu"
        )

        self.rel_model, self.rel_tok, self.rel_max_len, self.rel_thr = load_relevance_model(
            relevance_model_dir, self.device
        )
        self.tags_model, self.tags_tok, self.tags_max_len, self.all_tags, self.tags_thr = load_tags_model(
            tags_model_dir, self.device
        )
        self.sent_model, self.sent_tok, self.sent_max_len = load_sentiment_model(
            sentiment_model_dir, self.device
        )

    def analyze_text(
        self,
        text: str,
        relevance_threshold: Optional[float] = None,
        tags_threshold: Optional[float] = None,
        tone_threshold: Optional[float] = None,
        hubtype: Optional[str] = None,
    ) -> Dict[str, Any]:
        text_clean = clean_text(text)
        if not text_clean:
            return {
                "relevance": 0,
                "relevance_prob": 0.0,
                "tags": [],
                "tag_probs": [],
                "tone": "",
                "tone_probs": {},
            }

        rel_thr = relevance_threshold if relevance_threshold is not None else self.rel_thr
        tags_thr = tags_threshold if tags_threshold is not None else self.tags_thr

        rel_preds, rel_probs = infer_relevance(
            [text_clean],
            self.rel_model,
            self.rel_tok,
            self.rel_max_len,
            rel_thr,
            self.device,
            batch_size=1,
        )
        rel = int(rel_preds[0])
        rel_prob = float(rel_probs[0])

        if rel == 0:
            return {
                "relevance": 0,
                "relevance_prob": rel_prob,
                "tags": [],
                "tag_probs": [],
                "tone": "",
                "tone_probs": {},
            }

        tags_list, tag_probs_list = infer_tags_for_texts(
            [text_clean],
            self.tags_model,
            self.tags_tok,
            self.tags_max_len,
            self.all_tags,
            tags_thr,
            self.device,
            return_probs=True,
        )
        tags_list = tags_list[0]
        probs_all = tag_probs_list[0]
        tag_prob_pairs = list(zip(self.all_tags, probs_all))
        tag_prob_pairs.sort(key=lambda x: x[1], reverse=True)

        if hubtype == "Онлайн-СМИ" and "Новостные публикации" in self.all_tags:
            if "Новостные публикации" not in tags_list:
                tags_list.append("Новостные публикации")

        sent_single = infer_sentiment_single(
            text_clean,
            self.sent_model,
            self.sent_tok,
            self.sent_max_len,
            self.device,
            tone_threshold=tone_threshold,
        )
 
        return {
            "relevance": 1,
            "relevance_prob": rel_prob,
            "tags": tags_list,
            "tag_probs": tag_prob_pairs,
            "tone": sent_single["tone_str"],
            "tone_probs": sent_single["tone_probs"],
        }