""" scraper/space_a/inference.py Работа с датами и timestamp-ами (МСК), валидация диапазонов, построение URL отчёта. """ from __future__ import annotations 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, AutoModelForSequenceClassification from huggingface_hub import hf_hub_download # ========================= # Пороги тональности # ========================= # Повышенный порог для негатива снижает ложные срабатывания: # модель должна быть уверена на 75%+, чтобы пометить текст как негативный. # Позитив требует только 45% — он исторически предсказывается консервативнее. _SENTIMENT_NEGATIVE_THRESHOLD = 0.75 # было ~0.50 _SENTIMENT_POSITIVE_THRESHOLD = 0.45 # было ~0.50 # ========================= # Общая очистка текста # ========================= 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 = 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_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): 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 # ========================= # Тональность # ========================= # Индексы классов softmax: 0 = негатив, 1 = нейтрально, 2 = позитив # Используется для кастомной архитектуры RuBERTSentimentClassifier. # Для tbsa-модели (TBSAInferenceClassifier) маппинг фиксирован: # tone_to_idx в обучении: -1→0, 0→1, 1→2, поэтому col0=neg, col1=neu, col2=pos. _SENTIMENT_CLASS_NEGATIVE = 0 _SENTIMENT_CLASS_NEUTRAL = 1 _SENTIMENT_CLASS_POSITIVE = 2 # TBSA-таргет — слово, которое подставляется перед текстом при токенизации. # Должно совпадать с тем, что использовалось при обучении (_tbsa_train_engine.py). _TBSA_TARGET = "Авито" def index_to_tone(i: int) -> int: i = int(i) if i not in (0, 1, 2): raise ValueError(f"Некорректный индекс класса: {i}") return i - 1 def _apply_sentiment_thresholds( preds_tone: np.ndarray, probs: np.ndarray, negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, neg_col: int = _SENTIMENT_CLASS_NEGATIVE, pos_col: int = _SENTIMENT_CLASS_POSITIVE, ) -> np.ndarray: """Применяет раздельные пороги к вектору тональностей. Правила: - Предсказан негатив (-1): если prob_negative < negative_threshold → нейтраль. - Предсказан позитив (+1): если prob_positive < positive_threshold → нейтраль. - Предсказана нейтраль (0): порог не применяется — остаётся нейтраль. Args: preds_tone: массив int из {-1, 0, 1}, shape (N,) probs: softmax-вероятности, shape (N, 3); порядок колонок определяется neg_col / pos_col. negative_threshold: минимальная уверенность для сохранения негатива positive_threshold: минимальная уверенность для сохранения позитива neg_col: индекс колонки с вероятностью негатива pos_col: индекс колонки с вероятностью позитива Returns: Скорректированный массив тональностей, shape (N,). """ result = preds_tone.copy() for i, tone in enumerate(preds_tone): if tone == -1: if probs[i, neg_col] < negative_threshold: result[i] = 0 elif tone == 1: if probs[i, pos_col] < positive_threshold: result[i] = 0 return result class SentimentDatasetInfer(Dataset): """Датасет инференса тональности. Параметр tbsa_target: если задан (например "Авито"), токенизация выполняется как sentence-pair tokenizer(target, text) → [CLS] Авито [SEP] текст [SEP], что соответствует формату обучения TBSAClassifier. Также возвращает token_type_ids если токенайзер их поддерживает. Если tbsa_target=None — обычная одиночная токенизация. """ def __init__( self, texts_clean: List[str], tokenizer, max_length: int, tbsa_target: Optional[str] = None, ): self.texts = texts_clean self.tokenizer = tokenizer self.max_length = max_length self.tbsa_target = tbsa_target def __len__(self): return len(self.texts) def __getitem__(self, idx): text = self.texts[idx] if self.tbsa_target is not None: enc = self.tokenizer( self.tbsa_target, text, add_special_tokens=True, max_length=self.max_length, truncation=True, padding="max_length", return_tensors="pt", return_token_type_ids=True, ) else: enc = self.tokenizer( text, add_special_tokens=True, max_length=self.max_length, truncation=True, padding="max_length", return_tensors="pt", ) item = { "input_ids": enc["input_ids"].squeeze(0), "attention_mask": enc["attention_mask"].squeeze(0), } if "token_type_ids" in enc: item["token_type_ids"] = enc["token_type_ids"].squeeze(0) return item class RuBERTSentimentClassifier(nn.Module): """Кастомная голова для старых чекпоинтов (не-tbsa).""" 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) class TBSAInferenceClassifier(nn.Module): """Зеркало TBSAClassifier из _tbsa_train_engine.py для инференса. Архитектура: AutoModel (encoder) → Dropout → Linear(hidden, num_labels). Поддерживает token_type_ids, если модель их принимает (type_vocab_size > 1). Маппинг классов фиксирован обучением: col0=негатив, col1=нейтраль, col2=позитив. """ def __init__( self, model_name: str, num_labels: int = 3, dropout: float = 0.1, use_mean_pooling: bool = False, ): super().__init__() self.use_mean_pooling = use_mean_pooling self.encoder = AutoModel.from_pretrained(model_name) enc_cfg = self.encoder.config self.supports_token_type_ids = getattr(enc_cfg, "type_vocab_size", 0) > 1 print(f"[TBSAInferenceClassifier] supports_token_type_ids={self.supports_token_type_ids}") hidden = enc_cfg.hidden_size self.dropout = nn.Dropout(dropout) self.classifier = nn.Linear(hidden, num_labels) def forward(self, input_ids, attention_mask, token_type_ids=None): kwargs = dict(input_ids=input_ids, attention_mask=attention_mask) if self.supports_token_type_ids and token_type_ids is not None: kwargs["token_type_ids"] = token_type_ids out = self.encoder(**kwargs) if self.use_mean_pooling: m = attention_mask.unsqueeze(-1).float() pooled = (out.last_hidden_state * m).sum(1) / m.sum(1).clamp(min=1e-9) else: pooled = out.last_hidden_state[:, 0] return self.classifier(self.dropout(pooled)) # ── вспомогательная функция: маппинг id2label → neg/pos колонки ────────────── def _resolve_label_cols(id2label: Dict[int, str]) -> Tuple[int, int, int]: """Возвращает (neg_col, neu_col, pos_col) по словарю id2label из config. Поддерживает любой порядок лейблов в HF-конфиге: например {0: 'NEGATIVE', 1: 'NEUTRAL', 2: 'POSITIVE'} или {0: 'LABEL_0', 1: 'LABEL_1', 2: 'LABEL_2'} (fallback → 0,1,2). """ neg_col = neu_col = pos_col = None for idx, label in id2label.items(): l = label.lower() if any(x in l for x in ["neg", "-1", "негат"]): neg_col = int(idx) elif any(x in l for x in ["neu", "0", "нейтр"]): neu_col = int(idx) elif any(x in l for x in ["pos", "1", "позит"]): pos_col = int(idx) # fallback: если автодетект не сработал if neg_col is None: neg_col = _SENTIMENT_CLASS_NEGATIVE if neu_col is None: neu_col = _SENTIMENT_CLASS_NEUTRAL if pos_col is None: pos_col = _SENTIMENT_CLASS_POSITIVE return neg_col, neu_col, pos_col def _is_tbsa_checkpoint(repo_id: str) -> bool: """Проверяет наличие 'tbsa' в имени репозитория (регистронезависимо).""" return "tbsa" in repo_id.lower() def load_sentiment_model(sentiment_repo_id: str, device: torch.device): """Загружает модель тональности. Поддерживает два формата: 1. tbsa / TBSAInferenceClassifier — когда репозиторий содержит 'tbsa' в названии или в config.json есть поле "tbsa": true. config.json — кастомный обучающий конфиг (_tbsa_train_engine.py): содержит model_name, max_length, dropout, use_mean_pooling, tbsa=True. Веса загружаются из best_pytorch_model.bin. Токенизация: sentence-pair (target, text) → [CLS] Авито [SEP] text [SEP]. 2. Кастомная архитектура RuBERTSentimentClassifier (старый формат) — config.json содержит ключи 'model_name' и 'max_length' (tbsa отсутствует). Веса загружаются из pytorch_model.bin. Возвращает (model, tokenizer, max_length). """ # ── загружаем config.json ──────────────────────────────────────────────── config_path = hf_hub_download(sentiment_repo_id, "config.json") with open(config_path, "r", encoding="utf-8") as f: conf = json.load(f) # ── определяем формат ──────────────────────────────────────────────────── # tbsa=True: либо по имени репо, либо по явному флагу в конфиге is_tbsa = _is_tbsa_checkpoint(sentiment_repo_id) or bool(conf.get("tbsa", False)) tokenizer = AutoTokenizer.from_pretrained(sentiment_repo_id) if is_tbsa: # ── tbsa: кастомный обучающий конфиг (_tbsa_train_engine.py) ───────── # config.json НЕ является HF PretrainedConfig — нет model_type/architectures. # Загружаем через TBSAInferenceClassifier (зеркало TBSAClassifier). model_name = conf["model_name"] # "ai-forever/ruBERT-large" max_length = int(conf["max_length"]) # 256 dropout = float(conf.get("dropout", 0.1)) use_mean_pooling = bool(conf.get("use_mean_pooling", False)) print(f"[INFO] TBSA-режим: model_name={model_name}, max_length={max_length}") model_inner = TBSAInferenceClassifier( model_name=model_name, num_labels=3, dropout=dropout, use_mean_pooling=use_mean_pooling, ) weights_path = hf_hub_download(sentiment_repo_id, "best_pytorch_model.bin") state = torch.load(weights_path, map_location=device) model_inner.load_state_dict(state, strict=True) model_inner.to(device) model_inner.eval() # Маппинг фиксирован обучением (tone_to_idx: -1→0, 0→1, 1→2): # col 0 = негатив, col 1 = нейтраль, col 2 = позитив model = TbsaModelWrapper( model_inner, neg_col=0, neu_col=1, pos_col=2, tbsa_target=_TBSA_TARGET, ) print(f"[INFO] tbsa neg_col=0 neu_col=1 pos_col=2 target='{_TBSA_TARGET}'") else: # ── кастомная архитектура (старый формат) ──────────────────────────── model_name = conf["model_name"] max_length = int(conf["max_length"]) model_inner = RuBERTSentimentClassifier(model_name) weights_path = hf_hub_download(sentiment_repo_id, "pytorch_model.bin") state = torch.load(weights_path, map_location=device) model_inner.load_state_dict(state, strict=True) model_inner.to(device) model_inner.eval() model = model_inner return model, tokenizer, max_length class TbsaModelWrapper(nn.Module): """Обёртка над TBSAInferenceClassifier (или любой TBSA-моделью). Приводит интерфейс к общему виду: logits = model(input_ids, attention_mask) → shape (B, num_labels) Хранит: - маппинг колонок (neg_col, neu_col, pos_col) - tbsa_target: строка-таргет для sentence-pair токенизации ("Авито") или None для обычной токенизации. """ def __init__( self, hf_model: nn.Module, neg_col: int = _SENTIMENT_CLASS_NEGATIVE, neu_col: int = _SENTIMENT_CLASS_NEUTRAL, pos_col: int = _SENTIMENT_CLASS_POSITIVE, tbsa_target: Optional[str] = None, ): super().__init__() self.hf_model = hf_model self.neg_col = neg_col self.neu_col = neu_col self.pos_col = pos_col self.tbsa_target = tbsa_target def forward(self, input_ids, attention_mask, token_type_ids=None): if isinstance(self.hf_model, TBSAInferenceClassifier): return self.hf_model( input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, ) # Для HF AutoModelForSequenceClassification (устаревший путь) out = self.hf_model(input_ids=input_ids, attention_mask=attention_mask) return out.logits # ── единый конвертер «индекс предсказанного класса → тон {-1,0,1}» ─────────── def _preds_idx_to_tone( preds_idx: np.ndarray, model: nn.Module, ) -> np.ndarray: """Конвертирует argmax-индексы в тональность {-1, 0, 1}. Для TbsaModelWrapper использует neg_col/pos_col из обёртки. Для остальных моделей — index_to_tone (0→-1, 1→0, 2→+1). """ if isinstance(model, TbsaModelWrapper): result = np.zeros(len(preds_idx), dtype=int) # default нейтраль result[preds_idx == model.neg_col] = -1 result[preds_idx == model.pos_col] = 1 return result else: return np.array([index_to_tone(i) for i in preds_idx]) def infer_sentiment( texts_clean: List[str], model: nn.Module, tokenizer, max_length: int, device: torch.device, batch_size: int = 32, negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, ) -> Dict[str, np.ndarray]: """Пакетный инференс тональности с раздельными порогами по классам.""" _empty = { "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), } if not texts_clean: return _empty # Определяем tbsa_target для sentence-pair токенизации tbsa_target = model.tbsa_target if isinstance(model, TbsaModelWrapper) else None ds = SentimentDatasetInfer(texts_clean, tokenizer, max_length, tbsa_target=tbsa_target) 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) token_type_ids = batch.get("token_type_ids") if token_type_ids is not None: token_type_ids = token_type_ids.to(device) if isinstance(model, TbsaModelWrapper): logits = model(input_ids, attention_mask, token_type_ids=token_type_ids) else: logits = model(input_ids, attention_mask) all_logits.append(logits.cpu().numpy()) if not all_logits: return _empty 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_raw = _preds_idx_to_tone(preds_idx, model) # ── определяем neg/pos колонки для порогов ──────────────────────────────── if isinstance(model, TbsaModelWrapper): neg_col = model.neg_col pos_col = model.pos_col else: neg_col = _SENTIMENT_CLASS_NEGATIVE pos_col = _SENTIMENT_CLASS_POSITIVE preds_tone = _apply_sentiment_thresholds( preds_tone_raw, probs, negative_threshold=negative_threshold, positive_threshold=positive_threshold, neg_col=neg_col, pos_col=pos_col, ) 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, negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, # Устаревший параметр для обратной совместимости: # если передан — используется как negative_threshold. tone_threshold: Optional[float] = None, # Новый явный параметр для управления порогом позитива из вне. tone_threshold_neg: Optional[float] = None, ) -> Dict[str, Any]: """Инференс тональности для одного текста.""" _neg_thr = tone_threshold if tone_threshold is not None else negative_threshold _pos_thr = tone_threshold_neg if tone_threshold_neg is not None else positive_threshold res = infer_sentiment( [text_clean], model, tokenizer, max_length, device, batch_size=1, negative_threshold=_neg_thr, positive_threshold=_pos_thr, ) probs = res["probs"][0] tone_int = int(res["preds_tone"][0]) tone_str = tone_to_str(tone_int) # ── строим tone_probs в едином порядке (neg, neu, pos) ─────────────────── if isinstance(model, TbsaModelWrapper): neg_col = model.neg_col neu_col = model.neu_col pos_col = model.pos_col else: neg_col = _SENTIMENT_CLASS_NEGATIVE neu_col = _SENTIMENT_CLASS_NEUTRAL pos_col = _SENTIMENT_CLASS_POSITIVE tone_probs = { "негатив": float(probs[neg_col]), "нейтрально": float(probs[neu_col]), "позитив": float(probs[pos_col]), } 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, sentiment_negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, sentiment_positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, ) -> 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 # Обратная совместимость: старый tone_threshold переопределяет neg-порог. neg_thr = tone_threshold if tone_threshold is not None else sentiment_negative_threshold pos_thr = sentiment_positive_threshold 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, negative_threshold=neg_thr, positive_threshold=pos_thr, ) for local_i, global_i in enumerate(relevant_indices): tone_int = int(sent_res["preds_tone"][local_i]) 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, sentiment_negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, sentiment_positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, ): """ Считает 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] if len(relevant_indices) == 0: print("Нет релевантных записей.") return 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, ) pred_tags_all = [[] for _ in records] 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 # ── per-tag метрики ─────────────────────────────────────────────────────── metrics = {} for tag in all_tags: tp = fp = tn = fn = 0 for i in relevant_indices: pred = tag in pred_tags_all[i] true = tag in true_tags_all[i] if pred and true: tp += 1 elif pred and not true: fp += 1 elif not pred and true: fn += 1 else: tn += 1 precision = tp / (tp + fp) if (tp + fp) > 0 else 0.0 recall = tp / (tp + fn) if (tp + fn) > 0 else 0.0 f1 = ( 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0 ) metrics[tag] = { "TP": tp, "FP": fp, "TN": tn, "FN": fn, "Precision": precision, "Recall": recall, "F1": f1, } # ── micro-метрики ───────────────────────────────────────────────────────── total_tp = sum(m["TP"] for m in metrics.values()) total_fp = sum(m["FP"] for m in metrics.values()) total_fn = sum(m["FN"] for m in metrics.values()) micro_precision = total_tp / (total_tp + total_fp) if (total_tp + total_fp) > 0 else 0.0 micro_recall = total_tp / (total_tp + total_fn) if (total_tp + total_fn) > 0 else 0.0 micro_f1 = ( 2 * micro_precision * micro_recall / (micro_precision + micro_recall) if (micro_precision + micro_recall) > 0 else 0.0 ) print(f"\n{'='*60}") print(f" Micro Precision : {micro_precision:.4f}") print(f" Micro Recall : {micro_recall:.4f}") print(f" Micro F1 : {micro_f1:.4f}") print(f"{'='*60}\n") print(f"{'Тег':<40} {'P':>6} {'R':>6} {'F1':>6} {'TP':>5} {'FP':>5} {'FN':>5}") print("-" * 75) for tag, m in sorted(metrics.items(), key=lambda x: -x[1]["F1"]): print( f"{tag:<40} {m['Precision']:>6.3f} {m['Recall']:>6.3f} " f"{m['F1']:>6.3f} {m['TP']:>5} {m['FP']:>5} {m['FN']:>5}" ) return metrics