import os import json import re import unicodedata 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 # ============================================================================= # Нормализация казахского текста # Казахский латинский алфавит → кириллица, чтобы токенайзер ruBERT / E5 # воспринимал казахские слова в обоих алфавитах единообразно. # # Fix 1 (Issue #63): удалены дубли, исправлено ó → ө (не о!), убран ö→ö. # Порядок важен: сначала диграфы, затем одиночные символы. # ============================================================================= _KK_LAT_MAP = [ # Диграфы — строго первыми ("gh", "ғ"), ("kh", "х"), ("sh", "ш"), ("ch", "ч"), ("ng", "ң"), # Казахская латиница (стандарт MES-2021) — без дублей ("á", "ə"), ("ä", "ə"), ("é", "е"), ("í", "і"), ("ó", "ө"), # ИСПРАВЛЕНО: ó → ө (казахская «ö»), а не о (русская) ("ú", "ұ"), ("ü", "ү"), ("ñ", "ң"), ("ı", "і"), ("ń", "ң"), ("ǵ", "ғ"), ] def normalize_kz_text(text: str) -> str: """Транслитерирует казахскую латиницу → кириллицу (case-insensitive). Поддерживает: - Русский (без изменений) - Казахский кириллический (без изменений) - Казахский латинский (MES-2017, MES-2021, неформальная транслитерация) Fix 2 (Issue #63): добавлена Unicode NFC-нормализация перед replace, чтобы составные символы NFD (á = а + combining accent) корректно распознавались при поступлении текстов из Telegram, 2GIS, Instagram. """ # NFC: приводим составные символы к канонической форме text = unicodedata.normalize("NFC", text) lower = text.lower() for lat, cyr in _KK_LAT_MAP: lower = lower.replace(lat, cyr) return lower # ============================================================================= # Пороги тональности # ============================================================================= # Повышенный порог для негатива снижает ложные срабатывания. # Позитив требует меньшей уверенности — предсказывается консервативнее. _SENTIMENT_NEGATIVE_THRESHOLD = 0.75 _SENTIMENT_POSITIVE_THRESHOLD = 0.45 # Индексы классов softmax: 0 = негатив, 1 = нейтрально, 2 = позитив _SENTIMENT_CLASS_NEGATIVE = 0 _SENTIMENT_CLASS_NEUTRAL = 1 _SENTIMENT_CLASS_POSITIVE = 2 # ============================================================================= # Общая очистка текста # ============================================================================= URL_PATTERN = re.compile(r"http\S+|www\.\S+") MULTISPACE_PATTERN = re.compile(r"\s+") def clean_text(text: str) -> str: """Очищает и нормализует текст: поддерживает RU, KZ-cyr, KZ-lat.""" if not isinstance(text, str): return "" text = text.replace("\n", " ").replace("\r", " ") text = URL_PATTERN.sub(" ", text) text = MULTISPACE_PATTERN.sub(" ", text) text = text.strip() # Нормализуем казахскую латиницу → кириллицу (после удаления URL) text = normalize_kz_text(text) return text class RelevanceDatasetInfer(Dataset): def __init__(self, texts, tokenizer, max_length): self.texts = texts self.tokenizer = tokenizer self.max_length = max_length def __len__(self): return len(self.texts) def __getitem__(self, idx): enc = self.tokenizer(str(self.texts[idx]), 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, hidden_size=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): pooled = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0] return self.classifier(self.dropout2(self.relu(self.dense(self.dropout1(pooled))))).squeeze(-1) def load_relevance_model(relevance_repo_id, 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) tokenizer = AutoTokenizer.from_pretrained(relevance_repo_id) model = RuBERTBinaryClassifier(conf["model_name"]) state = torch.load( hf_hub_download(relevance_repo_id, "pytorch_model.bin"), map_location=device, weights_only=True, ) model.load_state_dict(state, strict=True) model.to(device).eval() return model, tokenizer, int(conf["max_length"]), float(conf["threshold"]) def infer_relevance(texts, model, tokenizer, max_length, threshold, device, batch_size=32): 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: all_logits.append(model(batch["input_ids"].to(device), batch["attention_mask"].to(device)).cpu().numpy()) all_logits = np.concatenate(all_logits) probs = 1 / (1 + np.exp(-all_logits)) return (probs >= threshold).astype(int), probs class RuBERTMultiLabelModel(nn.Module): def __init__(self, model_name, num_tags, hidden_size=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): pooled = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0] return self.classifier(self.dropout2(self.relu(self.dense(self.dropout1(pooled))))) def load_tags_model(tags_repo_id, device): """Загружает модель тегов. Returns ------- model, tokenizer, max_length, all_tags, threshold, head_tokens, tail_tokens, strategy """ with open(hf_hub_download(tags_repo_id, "training_config.json"), "r", encoding="utf-8") as f: tc = json.load(f) with open(hf_hub_download(tags_repo_id, "id2tag.json"), "r", encoding="utf-8") as f: id2tag = {int(k): v for k, v in json.load(f).items()} all_tags = [id2tag[i] for i in sorted(id2tag)] head_tokens = int(tc.get("head_tokens", 128)) tail_tokens = int(tc.get("tail_tokens", 382)) strategy = tc.get("strategy", "head_tail") tokenizer = AutoTokenizer.from_pretrained(tags_repo_id) model = RuBERTMultiLabelModel(tc.get("model_name", "DeepPavlov/rubert-base-cased"), len(id2tag), tc.get("hidden_size", 512)) state = torch.load( hf_hub_download(tags_repo_id, "pytorch_model.bin"), map_location=device, weights_only=True, ) model.load_state_dict(state, strict=True) model.to(device).eval() return model, tokenizer, int(tc.get("max_length", 512)), all_tags, float(tc.get("threshold", 0.3)), head_tokens, tail_tokens, strategy # ============================================================================= # Инференс: предсказание вероятностей для одного текста # Стратегия head_tail синхронизирована с GSMKZ-tags/inference.py # ============================================================================= def _encode_chunk(tokenizer, text, max_length, truncation=True, device=None): enc = tokenizer( text, add_special_tokens=True, max_length=max_length, truncation=truncation, padding="max_length", return_tensors="pt", ) if device is not None: return enc["input_ids"].to(device), enc["attention_mask"].to(device) return enc["input_ids"], enc["attention_mask"] def predict_for_long_text_tags( model: nn.Module, text: str, tokenizer, device: torch.device, max_length: int, head_tokens: int = 128, tail_tokens: int = 382, strategy: str = "head_tail", ) -> np.ndarray: """Предсказывает вероятности тегов для (потенциально длинного) текста. Стратегии: - "head_tail": берёт head_tokens токенов с начала + tail_tokens с конца, соединяет в один чанк <= max_length. Лучше всего для отзывов. - "sliding": скользящее окно, усредняет логиты по чанкам. """ enc_full = tokenizer( text, add_special_tokens=False, truncation=False, return_tensors="pt", ) ids_full = enc_full["input_ids"][0] att_full = enc_full["attention_mask"][0] seq_len = ids_full.shape[0] if seq_len <= max_length - 2: ids_t, att_t = _encode_chunk(tokenizer, text, max_length, device=device) with torch.no_grad(): logits = model(ids_t, att_t) return sigmoid(logits).squeeze(0).cpu().numpy() if strategy == "head_tail": avail = max_length - 2 h = min(head_tokens, avail) t = min(tail_tokens, avail - h) head_ids = ids_full[:h] tail_ids = ids_full[-t:] if t > 0 else torch.tensor([], dtype=torch.long) cls_id = tokenizer.cls_token_id sep_id = tokenizer.sep_token_id parts = [torch.tensor([cls_id]), head_ids] if t > 0: parts.append(tail_ids) parts.append(torch.tensor([sep_id])) chunk_ids = torch.cat(parts).unsqueeze(0).to(device) chunk_att = torch.ones_like(chunk_ids) pad_len = max_length - chunk_ids.shape[1] if pad_len > 0: pad_id = tokenizer.pad_token_id or 0 chunk_ids = torch.cat([ chunk_ids, torch.full((1, pad_len), pad_id, dtype=torch.long, device=device) ], dim=1) chunk_att = torch.cat([ chunk_att, torch.zeros((1, pad_len), dtype=torch.long, device=device) ], dim=1) with torch.no_grad(): logits = model(chunk_ids, chunk_att) return sigmoid(logits).squeeze(0).cpu().numpy() logits_list = [] for start in range(0, seq_len, max_length): chunk_ids = ids_full[start: start + max_length] chunk_att = att_full[start: start + max_length] if chunk_ids.shape[0] == 0: continue with torch.no_grad(): logits = model( chunk_ids.unsqueeze(0).to(device), chunk_att.unsqueeze(0).to(device), ) logits_list.append(logits.squeeze(0).cpu().numpy()) if not logits_list: ids_t, att_t = _encode_chunk(tokenizer, text, max_length, device=device) with torch.no_grad(): logits = model(ids_t, att_t) return sigmoid(logits).squeeze(0).cpu().numpy() avg_logits = np.mean(np.stack(logits_list, axis=0), axis=0) return 1.0 / (1.0 + np.exp(-avg_logits)) def infer_tags_for_texts( texts_clean, model, tokenizer, max_length, all_tags, threshold, device, return_probs=False, head_tokens: int = 128, tail_tokens: int = 382, strategy: str = "head_tail", ): """Fix 3 (Issue #63): guard на пустой текст — пропускаем вызов модели.""" results_tags, results_probs = [], [] for text in texts_clean: if not text or not text.strip(): results_tags.append([]) results_probs.append(np.zeros(len(all_tags), dtype=float)) continue probs = predict_for_long_text_tags( model, text, tokenizer, device, max_length, head_tokens=head_tokens, tail_tokens=tail_tokens, strategy=strategy, ) tags = [all_tags[i] for i in np.where(probs >= threshold)[0]] results_tags.append(tags) results_probs.append(probs) return (results_tags, results_probs) if return_probs else results_tags def index_to_tone(i): 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): порог не применяется. """ 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): def __init__(self, texts_clean, tokenizer, max_length): self.texts = texts_clean self.tokenizer = tokenizer self.max_length = max_length def __len__(self): return len(self.texts) def __getitem__(self, idx): enc = self.tokenizer(str(self.texts[idx]), 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, num_labels=3, hidden_size=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): pooled = self.bert(input_ids=input_ids, attention_mask=attention_mask).last_hidden_state[:, 0] return self.classifier(self.dropout2(self.relu(self.dense(self.dropout1(pooled))))) def load_sentiment_model(sentiment_repo_id, device): with open(hf_hub_download(sentiment_repo_id, "config.json"), "r", encoding="utf-8") as f: conf = json.load(f) tokenizer = AutoTokenizer.from_pretrained(sentiment_repo_id) model = RuBERTSentimentClassifier(conf["model_name"]) state = torch.load( hf_hub_download(sentiment_repo_id, "pytorch_model.bin"), map_location=device, weights_only=True, ) model.load_state_dict(state, strict=True) model.to(device).eval() return model, tokenizer, int(conf["max_length"]) def infer_sentiment( texts_clean, model, tokenizer, max_length, device, batch_size=32, negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, ): """Пакетный инференс тональности с раздельными порогами по классам.""" _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 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: all_logits.append(model(batch["input_ids"].to(device), batch["attention_mask"].to(device)).cpu().numpy()) all_logits = np.concatenate(all_logits) probs = torch.softmax(torch.tensor(all_logits), dim=-1).numpy() preds_idx = probs.argmax(axis=1) preds_tone_raw = np.array([index_to_tone(i) for i in preds_idx]) preds_tone = _apply_sentiment_thresholds( preds_tone_raw, probs, negative_threshold=negative_threshold, positive_threshold=positive_threshold, neg_col=_SENTIMENT_CLASS_NEGATIVE, pos_col=_SENTIMENT_CLASS_POSITIVE, ) return {"logits": all_logits, "probs": probs, "preds_idx": preds_idx, "preds_tone": preds_tone} def tone_to_str(tone_int): return {-1: "негатив", 0: "нейтрально", 1: "позитив"}.get(int(tone_int), "") def infer_sentiment_single( text_clean: str, model, tokenizer, max_length: int, device, negative_threshold: float = _SENTIMENT_NEGATIVE_THRESHOLD, positive_threshold: float = _SENTIMENT_POSITIVE_THRESHOLD, tone_threshold: Optional[float] = None, ) -> Dict[str, Any]: """Инференс тональности для одного текста с раздельными порогами.""" _neg_thr = tone_threshold if tone_threshold is not None else negative_threshold _pos_thr = 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]) return { "tone_int": tone_int, "tone_str": tone_to_str(tone_int), "tone_probs": { "негатив": float(probs[_SENTIMENT_CLASS_NEGATIVE]), "нейтрально": float(probs[_SENTIMENT_CLASS_NEUTRAL]), "позитив": float(probs[_SENTIMENT_CLASS_POSITIVE]), }, } def load_input_json(json_path): 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"] raise ValueError("Неизвестный формат входного JSON") def extract_full_text_from_record(record): title = record.get("title", "") or "" text = record.get("text", "") or "" return clean_text((title + ". " + text).strip(". ")) def format_ba_datetime(ts): if not ts: return "" try: return datetime.fromtimestamp(ts).strftime("%d.%m.%Y %H:%M") except Exception: return ""