Update inference.py

inference.py

@@ -1,13 +1,11 @@
-import os
-from typing import Any, Dict, List, Optional
+# requirements
+!pip install torch==2.8.0 torchvision==0.17.2
+!pip install transformers==4.48.3 tokenizers sentencepiece accelerate
 
-import torch
-from transformers import AutoModel, AutoTokenizer
 
-try:
-    from huggingface_hub import hf_hub_download
-except Exception:
-    hf_hub_download = None
+import torch
+from typing import List, Optional
+from transformers import AutoTokenizer, AutoModel
 
 
 ERROR_NAMES_RU = {
@@ -20,335 +18,156 @@ ERROR_NAMES_RU = {
 }
 
 
-def _resolve_calibration_path(model_path: str) -> Optional[str]:
-    local_path = os.path.join(model_path, "calibration_data.pth")
-    if os.path.exists(local_path):
-        return local_path
-
-    if hf_hub_download is None or os.path.isdir(model_path):
-        return None
-
-    try:
-        return hf_hub_download(
-            repo_id=model_path,
-            filename="calibration_data.pth",
-        )
-    except Exception:
-        return None
-
-
-class RQAInferenceHF:
-    def __init__(
-        self,
-        model_path: str,
-        device: Optional[torch.device] = None,
-        max_length: int = 512,
-        issue_uncertain_margin: float = 0.05,
-        hidden_uncertain_margin: float = 0.05,
-        error_uncertain_margin: float = 0.05,
-    ):
-        self.model_path = model_path
-        self.device = device or torch.device(
-            "cuda" if torch.cuda.is_available() else "cpu"
-        )
+class RQAJudge:
+    def __init__(self, model_name="skatzR/RQA-R2", device=None, max_length: int = 512):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
         self.max_length = int(max_length)
-        self.issue_uncertain_margin = float(issue_uncertain_margin)
-        self.hidden_uncertain_margin = float(hidden_uncertain_margin)
-        self.error_uncertain_margin = float(error_uncertain_margin)
 
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            model_name,
+            trust_remote_code=True
+        )
         self.model = AutoModel.from_pretrained(
-            model_path,
-            trust_remote_code=True,
-        ).to(self.device).eval()
-        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+            model_name,
+            trust_remote_code=True
+        ).to(self.device)
+
+        self.model.eval()
 
         cfg = self.model.config
-        self.schema_version = str(getattr(cfg, "schema_version", "unknown"))
-        self.error_types = list(getattr(cfg, "error_types", []))
-        self.t_issue = float(getattr(cfg, "temperature_has_issue", 1.0))
-        self.t_hidden = float(getattr(cfg, "temperature_is_hidden", 1.0))
-        self.t_errors = list(
-            getattr(cfg, "temperature_errors", [1.0] * len(self.error_types))
-        )
-        self.th_issue = float(getattr(cfg, "threshold_has_issue", 0.5))
-        self.th_hidden = float(getattr(cfg, "threshold_is_hidden", 0.5))
-        self.th_error = float(getattr(cfg, "threshold_error", 0.5))
-        self.th_errors = list(
-            getattr(cfg, "threshold_errors", [self.th_error] * len(self.error_types))
-        )
+        self.error_types = list(cfg.error_types)
 
-        calibration_path = _resolve_calibration_path(model_path)
-        if calibration_path:
-            calibration = torch.load(calibration_path, map_location="cpu")
-            calibration_error_types = calibration.get("error_types", None)
-            if calibration_error_types is not None:
-                if list(calibration_error_types) != self.error_types:
-                    raise ValueError(
-                        "Calibration artifact error_types mismatch with model.config.error_types."
-                    )
-
-            self.schema_version = str(
-                calibration.get("schema_version", self.schema_version)
-            )
-            self.t_issue = float(
-                calibration.get("temperature_has_issue", self.t_issue)
-            )
-            self.t_hidden = float(
-                calibration.get("temperature_is_hidden", self.t_hidden)
-            )
-            self.t_errors = list(
-                calibration.get("temperature_errors", self.t_errors)
-            )
-            self.th_issue = float(
-                calibration.get("threshold_has_issue", self.th_issue)
-            )
-            self.th_hidden = float(
-                calibration.get("threshold_is_hidden", self.th_hidden)
-            )
-            self.th_error = float(
-                calibration.get("threshold_error", self.th_error)
-            )
-            self.th_errors = list(
-                calibration.get("threshold_errors", self.th_errors)
-            )
+        self.temp_issue = float(cfg.temperature_has_issue)
+        self.temp_hidden = float(cfg.temperature_is_hidden)
+        self.temp_errors = list(cfg.temperature_errors)
 
-    def _apply_temperature(
-        self,
-        issue_logits: torch.Tensor,
-        hidden_logits: torch.Tensor,
-        errors_logits: torch.Tensor,
-    ):
-        calibrated_issue = issue_logits / float(self.t_issue)
-        calibrated_hidden = hidden_logits / float(self.t_hidden)
-        calibrated_errors = errors_logits.clone()
-        for idx in range(calibrated_errors.size(1)):
-            temperature = float(self.t_errors[idx]) if idx < len(self.t_errors) else 1.0
-            calibrated_errors[:, idx] = calibrated_errors[:, idx] / temperature
-        return calibrated_issue, calibrated_hidden, calibrated_errors
+        self.threshold_issue = float(cfg.threshold_has_issue)
+        self.threshold_hidden = float(cfg.threshold_is_hidden)
+        self.threshold_error = float(cfg.threshold_error)
+        self.threshold_errors = list(cfg.threshold_errors)
 
     @torch.no_grad()
-    def predict(
+    def infer(
         self,
         text: str,
-        return_probs: bool = False,
-        threshold_issue: Optional[float] = None,
-        threshold_hidden: Optional[float] = None,
-        threshold_error: Optional[float] = None,
-        threshold_errors: Optional[List[float]] = None,
-    ) -> Dict[str, Any]:
-        issue_threshold = self.th_issue if threshold_issue is None else float(threshold_issue)
-        hidden_threshold = self.th_hidden if threshold_hidden is None else float(threshold_hidden)
-        error_threshold = self.th_error if threshold_error is None else float(threshold_error)
-        error_thresholds = self.th_errors if threshold_errors is None else list(threshold_errors)
-
-        encoded = self.tokenizer(
+        issue_threshold: Optional[float] = None,
+        hidden_threshold: Optional[float] = None,
+        error_threshold: Optional[float] = None,
+        error_thresholds: Optional[List[float]] = None,
+        issue_uncertain_margin: float = 0.05,
+        hidden_uncertain_margin: float = 0.05,
+        error_uncertain_margin: float = 0.05,
+    ):
+        issue_threshold = self.threshold_issue if issue_threshold is None else float(issue_threshold)
+        hidden_threshold = self.threshold_hidden if hidden_threshold is None else float(hidden_threshold)
+        error_threshold = self.threshold_error if error_threshold is None else float(error_threshold)
+        error_thresholds = self.threshold_errors if error_thresholds is None else list(error_thresholds)
+
+        inputs = self.tokenizer(
             text,
             truncation=True,
             max_length=self.max_length,
             padding="max_length",
-            return_tensors="pt",
-        )
-        input_ids = encoded["input_ids"].to(self.device)
-        attention_mask = encoded["attention_mask"].to(self.device)
-
-        outputs = self.model(input_ids=input_ids, attention_mask=attention_mask)
-        issue_logits, hidden_logits, errors_logits = self._apply_temperature(
-            outputs["has_issue_logits"],
-            outputs["is_hidden_logits"],
-            outputs["errors_logits"],
-        )
+            return_tensors="pt"
+        ).to(self.device)
+
+        outputs = self.model(**inputs)
+
+        issue_logit = outputs["has_issue_logits"] / self.temp_issue
+        hidden_logit = outputs["is_hidden_logits"] / self.temp_hidden
 
-        issue_probability = float(torch.sigmoid(issue_logits).item())
-        has_issue = issue_probability >= issue_threshold
+        error_logits = outputs["errors_logits"][0].clone()
+        for i in range(len(self.error_types)):
+            error_logits[i] = error_logits[i] / self.temp_errors[i]
 
-        result: Dict[str, Any] = {
-            "schema_version": self.schema_version,
+        issue_prob = torch.sigmoid(issue_logit).item()
+        has_issue = issue_prob >= issue_threshold
+
+        result = {
             "text": text,
             "class": None,
             "status": "ok",
             "review_required": False,
-            "has_logical_issue": bool(has_issue),
-            "has_issue_probability": issue_probability,
-            "threshold_has_issue": issue_threshold,
-            "temperature_has_issue": float(self.t_issue),
-            "is_hidden_problem": False,
+            "has_issue": has_issue,
+            "issue_probability": issue_prob,
+            "hidden_problem": False,
             "hidden_probability": None,
-            "threshold_is_hidden": hidden_threshold,
-            "temperature_is_hidden": float(self.t_hidden),
             "errors": [],
             "num_errors": 0,
+            "threshold_issue": issue_threshold,
+            "threshold_hidden": hidden_threshold,
             "threshold_error": error_threshold,
             "threshold_errors": error_thresholds,
-            "calibrated": (
-                abs(self.t_issue - 1.0) > 1e-6
-                or abs(self.t_hidden - 1.0) > 1e-6
-                or any(abs(float(t) - 1.0) > 1e-6 for t in self.t_errors)
-            ),
+            "schema_version": getattr(self.model.config, "schema_version", "unknown"),
         }
 
-        if abs(issue_probability - issue_threshold) <= self.issue_uncertain_margin:
+        if abs(issue_prob - issue_threshold) <= issue_uncertain_margin:
             result["status"] = "uncertain"
             result["review_required"] = True
 
         if not has_issue:
             result["class"] = "logical"
-            if return_probs:
-                result["raw"] = {"p_issue": issue_probability}
             return result
 
-        hidden_probability = float(torch.sigmoid(hidden_logits).item())
-        is_hidden = hidden_probability >= hidden_threshold
-        result["hidden_probability"] = hidden_probability
-        result["is_hidden_problem"] = bool(is_hidden)
+        hidden_prob = torch.sigmoid(hidden_logit).item()
+        is_hidden = hidden_prob >= hidden_threshold
+
+        result["hidden_problem"] = is_hidden
+        result["hidden_probability"] = hidden_prob
 
-        if abs(hidden_probability - hidden_threshold) <= self.hidden_uncertain_margin:
+        if abs(hidden_prob - hidden_threshold) <= hidden_uncertain_margin:
             result["status"] = "uncertain"
             result["review_required"] = True
 
         if is_hidden:
             result["class"] = "hidden"
-            if return_probs:
-                result["raw"] = {
-                    "p_issue": issue_probability,
-                    "p_hidden": hidden_probability,
-                }
             return result
 
-        error_probabilities = torch.sigmoid(errors_logits).cpu().numpy()[0]
-        detected_errors = []
-        for idx, error_type in enumerate(self.error_types):
-            probability = float(error_probabilities[idx])
-            threshold_i = float(
-                error_thresholds[idx] if idx < len(error_thresholds) else error_threshold
-            )
-            if abs(probability - threshold_i) <= self.error_uncertain_margin:
+        error_probs = torch.sigmoid(error_logits).tolist()
+        detected = []
+        for i, err_name in enumerate(self.error_types):
+            prob = float(error_probs[i])
+            threshold_i = float(error_thresholds[i] if i < len(error_thresholds) else error_threshold)
+
+            if abs(prob - threshold_i) <= error_uncertain_margin:
                 result["status"] = "uncertain"
                 result["review_required"] = True
-            if probability >= threshold_i:
-                detected_errors.append(
-                    {
-                        "type": error_type,
-                        "probability": probability,
-                        "threshold": threshold_i,
-                        "temperature": float(self.t_errors[idx]) if idx < len(self.t_errors) else 1.0,
-                    }
-                )
-
-        detected_errors.sort(key=lambda item: item["probability"], reverse=True)
-        result["class"] = "explicit"
-        result["errors"] = detected_errors
-        result["num_errors"] = len(detected_errors)
-
-        if return_probs:
-            result["error_probabilities"] = {
-                error_type: float(probability)
-                for error_type, probability in zip(self.error_types, error_probabilities)
-            }
-            result["raw"] = {
-                "p_issue": issue_probability,
-                "p_hidden": hidden_probability,
-            }
 
-        return result
+            if prob >= threshold_i:
+                detected.append((err_name, prob))
 
-    def pretty_print(self, prediction: Dict[str, Any], use_russian_names: bool = True) -> None:
-        print("-" * 70)
-        print(
-            f"Class: {prediction['class']} | status={prediction['status']} "
-            f"| review_required={prediction['review_required']}"
-        )
-        print(
-            f"Issue: {prediction['has_logical_issue']} "
-            f"({prediction['has_issue_probability'] * 100:.2f}%)  "
-            f"th={prediction['threshold_has_issue']:.3f}"
-        )
-        if prediction["hidden_probability"] is not None:
-            print(
-                f"Hidden: {prediction['is_hidden_problem']} "
-                f"({prediction['hidden_probability'] * 100:.2f}%)  "
-                f"th={prediction['threshold_is_hidden']:.3f}"
-            )
+        detected.sort(key=lambda x: x[1], reverse=True)
 
-        if prediction["errors"]:
-            printable_errors = []
-            for item in prediction["errors"]:
-                label = (
-                    ERROR_NAMES_RU.get(item["type"], item["type"])
-                    if use_russian_names
-                    else item["type"]
-                )
-                printable_errors.append((label, round(item["probability"], 3)))
-            print(f"Top errors: {printable_errors}")
+        result["class"] = "explicit"
+        result["errors"] = detected
+        result["num_errors"] = len(detected)
+        return result
 
+    def pretty_print(self, r):
+        print("\n" + "=" * 72)
+        print("📄 Текст:")
+        print(r["text"])
 
-class RQAJudge:
-    def __init__(
-        self,
-        model_name: str = "skatzR/RQA-R2",
-        device: Optional[torch.device] = None,
-        max_length: int = 512,
-    ):
-        self.runner = RQAInferenceHF(
-            model_path=model_name,
-            device=device,
-            max_length=max_length,
+        print(
+            f"\n🔎 Обнаружена проблема: {'ДА' if r['has_issue'] else 'НЕТ'} "
+            f"({r['issue_probability'] * 100:.2f}%)"
         )
+        print(f"🧠 Класс: {r['class']}")
 
-    def infer(
-        self,
-        text: str,
-        issue_threshold: Optional[float] = None,
-        hidden_threshold: Optional[float] = None,
-        error_threshold: Optional[float] = None,
-        error_thresholds: Optional[List[float]] = None,
-    ) -> Dict[str, Any]:
-        prediction = self.runner.predict(
-            text=text,
-            return_probs=True,
-            threshold_issue=issue_threshold,
-            threshold_hidden=hidden_threshold,
-            threshold_error=error_threshold,
-            threshold_errors=error_thresholds,
-        )
-        return {
-            "text": text,
-            "class": prediction["class"],
-            "status": prediction["status"],
-            "review_required": prediction["review_required"],
-            "has_issue": prediction["has_logical_issue"],
-            "issue_probability": prediction["has_issue_probability"],
-            "hidden_problem": prediction["is_hidden_problem"],
-            "hidden_probability": prediction["hidden_probability"],
-            "errors": [
-                (item["type"], item["probability"])
-                for item in prediction["errors"]
-            ],
-            "num_errors": prediction["num_errors"],
-            "threshold_has_issue": prediction["threshold_has_issue"],
-            "threshold_is_hidden": prediction["threshold_is_hidden"],
-            "threshold_error": prediction["threshold_error"],
-        }
+        if r["status"] == "uncertain":
+            print("⚠️ Статус: uncertain")
 
-    def pretty_print(self, result: Dict[str, Any], use_russian_names: bool = True) -> None:
-        converted = {
-            "class": result["class"],
-            "status": result["status"],
-            "review_required": result["review_required"],
-            "has_logical_issue": result["has_issue"],
-            "has_issue_probability": result["issue_probability"],
-            "threshold_has_issue": result["threshold_has_issue"],
-            "is_hidden_problem": result["hidden_problem"],
-            "hidden_probability": result["hidden_probability"],
-            "threshold_is_hidden": result["threshold_is_hidden"],
-            "errors": [
-                {
-                    "type": error_type,
-                    "probability": probability,
-                }
-                for error_type, probability in result["errors"]
-            ],
-        }
-        self.runner.pretty_print(converted, use_russian_names=use_russian_names)
+        if r["hidden_probability"] is not None:
+            print(
+                f"🟡 Hidden: {'ДА' if r['hidden_problem'] else 'НЕТ'} "
+                f"({r['hidden_probability'] * 100:.2f}%)"
+            )
 
+        if r["errors"]:
+            print("\n❌ Явные логические ошибки:")
+            for name, prob in r["errors"]:
+                print(f"  • {ERROR_NAMES_RU.get(name, name)} — {prob * 100:.2f}%")
+        else:
+            print("\n✅ Явных логических ошибок не обнаружено")
 
-__all__ = ["RQAInferenceHF", "RQAJudge", "ERROR_NAMES_RU"]
+        print("=" * 72)