inference.py

import os
import json
import re
import logging
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

logger = logging.getLogger(__name__)

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

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()


# =========================
# Канонизация тегов
# =========================

# Маппинг: написание из id2tag модели → каноническое написание политики BA.
# Пополняй при появлении новых missingpolicy в логах.
TAG_CANONICAL_MAP: Dict[str, str] = {
    "AppStoreAu ru":                        "AppStore/Au ru",
    "App Store/Au ru":                      "AppStore/Au ru",
    "СП -Шаблонный ответ":                  "СП Шаблонный ответ",
    "СП -Не решают проблемы":               "СП Не решают проблемы",
    "СП -Не получается дозвониться":        "СП Не получается дозвонится",
    "СП- Рекомендация":                     "СП Рекомендация",
    "СП – вопросы":                         "СП - вопросы",
}


def normalize_tag(tag: str) -> str:
    """Возвращает каноническое написание тега по TAG_CANONICAL_MAP."""
    return TAG_CANONICAL_MAP.get(tag.strip(), tag.strip())


def normalize_tags(tags: List[str]) -> List[str]:
    """
    Приводит список тегов к каноническим написаниям.
    Дедуплицирует: если два разных написания сводятся к одному каноническому,
    оставляет только первое вхождение.
    """
    seen: set = set()
    result: List[str] = []
    for tag in tags:
        canonical = normalize_tag(tag)
        if canonical not in seen:
            seen.add(canonical)
            result.append(canonical)
    return result


def check_tags_vs_policy(all_tags: List[str], policy: Dict) -> None:
    """
    Логирует все теги модели, отсутствующие в политике BA.
    Вызывать один раз при старте (после load_tags_model).
    Помогает превентивно выявлять расхождения до попадания в продакшн.
    """
    policy_tags = set(policy.keys())
    missing = []
    for tag in all_tags:
        canonical = normalize_tag(tag)
        if canonical not in policy_tags:
            missing.append((tag, canonical))
    if missing:
        for original, canonical in missing:
            logger.warning(
                "TAG '%s' (canonical: '%s') отсутствует в tag_fusion_policy "
                "→ добавь в TAG_CANONICAL_MAP или в политику",
                original, canonical,
            )
    else:
        logger.info("✅ Все теги модели найдены в политике BA")


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

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 = 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
    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)
# =========================

# ── НОВАЯ архитектура (ruRoBERTa-large, CLS+mean pooling) ──
class RuRoBERTaMultiLabelModel(nn.Module):
    """
    Соответствует Avito_train_tags.py:
    ruRoBERTa-large (hidden=1024) → (CLS + mean_pool)/2 → Dropout → Linear(1024, num_tags)

    Этап 1 (Issue #56/#57): добавлен для загрузки новых чекпоинтов.
    load_tags_model автоматически выбирает этот класс когда
    'arch'=='ruroberta' или 'roberta' в model_name.
    """
    def __init__(self, model_name: str, num_tags: int):
        super().__init__()
        self.bert = AutoModel.from_pretrained(model_name)
        h = self.bert.config.hidden_size  # 1024 для large
        self.dropout = nn.Dropout(0.1)
        self.classifier = nn.Linear(h, num_tags)

    def _mean_pool(self, last_hidden: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        mask = attention_mask.unsqueeze(-1).float()
        return (last_hidden * mask).sum(1) / mask.sum(1).clamp(min=1e-9)

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor:
        out = self.bert(input_ids=input_ids, attention_mask=attention_mask)
        cls_vec  = out.last_hidden_state[:, 0]
        mean_vec = self._mean_pool(out.last_hidden_state, attention_mask)
        pooled   = (cls_vec + mean_vec) / 2
        return self.classifier(self.dropout(pooled))


# ── СТАРАЯ архитектура (ruBERT-base, двухслойная голова) — оставляем для совместимости ──
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):
    """
    Загружает модель тегов с автоматическим выбором архитектуры.

    Автоматически выбирает архитектуру по 'arch' или 'model_name'
    в training_config.json:
      - 'ruroberta' / 'roberta' в model_name → RuRoBERTaMultiLabelModel (новая)
      - иначе → RuBERTMultiLabelModel (старая, для обратной совместимости)

    Этап 1 (Issue #56/#57): обновлена для поддержки новых чекпоинтов ruRoBERTa.
    """
    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))
    arch        = tc.get("arch", "")  # явный ключ архитектуры
    head_tokens = int(tc.get("head_tokens", 256))
    tail_tokens = int(tc.get("tail_tokens", 253))
    strategy    = tc.get("long_text_strategy", "head_tail")

    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)

    # Выбор архитектуры
    is_roberta = (arch == "ruroberta") or ("roberta" in model_name.lower())
    if is_roberta:
        logger.info("Загружаем RuRoBERTaMultiLabelModel (arch=ruroberta, hidden=1024)")
        model = RuRoBERTaMultiLabelModel(model_name, num_tags)
    else:
        logger.info("Загружаем RuBERTMultiLabelModel (arch=rubert, hidden=%d)", hidden_size)
        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, head_tokens, tail_tokens, strategy


def predict_for_long_text_tags(
    model: nn.Module,
    text: str,
    tokenizer,
    device: torch.device,
    max_length: int,
    head_tokens: int = 256,
    tail_tokens: int = 253,
    strategy: str = "head_tail",
) -> np.ndarray:
    """
    Предсказание тегов для (потенциально длинного) текста.

    Этап 1 (Issue #56/#57): заменён sliding window на head+tail стратегию,
    соответствующую методологии обучения Avito_train_tags.py.

    strategy='head_tail': берём первые head_tokens + последние tail_tokens токенов.
    strategy='truncate': простое усечение (fallback).
    """
    enc_full = tokenizer(text, add_special_tokens=False, truncation=False, return_tensors="pt")
    ids      = enc_full["input_ids"][0]
    seq_len  = ids.shape[0]

    cls_id = tokenizer.cls_token_id or tokenizer.bos_token_id
    sep_id = tokenizer.sep_token_id or tokenizer.eos_token_id
    pad_id = tokenizer.pad_token_id

    if seq_len <= max_length - 2:
        seq  = [cls_id] + ids.tolist() + [sep_id]
        pad_len = max_length - len(seq)
        seq  = seq + [pad_id] * pad_len
        mask = [1] * (max_length - pad_len) + [0] * pad_len
    else:
        if strategy == "head_tail":
            head = ids[:head_tokens].tolist()
            tail = ids[-tail_tokens:].tolist()
            seq  = [cls_id] + head + [sep_id] + tail + [sep_id]
            if len(seq) < max_length:
                seq = seq + [pad_id] * (max_length - len(seq))
            seq  = seq[:max_length]
            mask = [1] * max_length
        else:
            # truncate fallback
            seq  = [cls_id] + ids[:max_length - 2].tolist() + [sep_id]
            mask = [1] * max_length

    input_ids      = torch.tensor([seq],  dtype=torch.long).to(device)
    attention_mask = torch.tensor([mask], dtype=torch.long).to(device)

    with torch.no_grad():
        logits = model(input_ids, attention_mask)
        probs  = torch.sigmoid(logits).squeeze(0).cpu().numpy()
    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,
    head_tokens: int = 256,
    tail_tokens: int = 253,
    strategy: str = "head_tail",
):
    results_tags = []
    results_probs = []

    for text in texts_clean:
        probs = predict_for_long_text_tags(
            model, text, tokenizer, device, max_length,
            head_tokens=head_tokens, tail_tokens=tail_tokens, strategy=strategy,
        )
        mask = probs >= threshold
        idxs = np.where(mask)[0]
        tags = [all_tags[i] for i in idxs]
        # Канонизация: приводим все теги к написанию политики BA
        tags = normalize_tags(tags)

        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 = 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
    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]:
    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 ""