Spaces:

parchiev
/

SIRIN

Running

App Files Files Community

parchiev commited on Jul 14

Commit

6c98539

verified ·

1 Parent(s): 3fd0904

Convert static replay to Streamlit app (hosted profile)

Browse files

This view is limited to 50 files because it contains too many changes. See raw diff

Files changed (50) hide show

.gitattributes +5 -0
.streamlit/config.toml +20 -0
README.md +14 -7
index.html +0 -144
pyproject.toml +21 -0
requirements.txt +15 -0
sirin/__init__.py +0 -0
sirin/classification/__init__.py +1 -0
sirin/classification/linear.py +335 -0
sirin/configs/detector/probing.yaml +8 -0
sirin/configs/feature_processor/hiddens.yaml +7 -0
sirin/configs/model_adapter/default.yaml +8 -0
sirin/configs/model_adapter/hf_large_model.yaml +17 -0
sirin/configs/model_adapter/hf_with_batching.yaml +18 -0
sirin/configs/model_adapter/openai_with_parallel.yaml +16 -0
sirin/configs/model_adapter/openrouter.yaml +15 -0
sirin/configs/model_adapter/vllm_with_batching.yaml +18 -0
sirin/configs/pipeline/probing.yaml +23 -0
sirin/configs/token_locator_config/default.yaml +1 -0
sirin/configs/train.yaml +24 -0
sirin/configs/train_args/default.yaml +10 -0
sirin/definitions/__init__.py +3 -0
sirin/definitions/constants.py +20 -0
sirin/definitions/enums.py +204 -0
sirin/definitions/types.py +35 -0
sirin/detection/__init__.py +0 -0
sirin/detection/_lm_polygraph_compat.py +168 -0
sirin/detection/approximators/__init__.py +18 -0
sirin/detection/approximators/base.py +13 -0
sirin/detection/approximators/sep.py +224 -0
sirin/detection/base.py +497 -0
sirin/detection/judging/__init__.py +40 -0
sirin/detection/judging/judges/__init__.py +4 -0
sirin/detection/judging/judges/base.py +292 -0
sirin/detection/judging/judges/claim/__init__.py +3 -0
sirin/detection/judging/judges/claim/decoder.py +201 -0
sirin/detection/judging/judges/claim/encoder.py +169 -0
sirin/detection/judging/judges/claim/openai.py +134 -0
sirin/detection/judging/judges/sequence/__init__.py +3 -0
sirin/detection/judging/judges/sequence/decoder.py +139 -0
sirin/detection/judging/judges/sequence/encoder.py +108 -0
sirin/detection/judging/judges/sequence/openai.py +99 -0
sirin/detection/judging/judges/token/__init__.py +3 -0
sirin/detection/judging/judges/token/decoder.py +186 -0
sirin/detection/judging/judges/token/encoder.py +131 -0
sirin/detection/judging/judges/token/openai.py +139 -0
sirin/detection/judging/judges/utils/__init__.py +23 -0
sirin/detection/judging/judges/utils/decoder.py +67 -0
sirin/detection/judging/judges/utils/metrics.py +106 -0
sirin/detection/judging/judges/utils/prompts.py +127 -0

.gitattributes CHANGED Viewed

@@ -48,3 +48,8 @@ static/videos/teaser.mp4 filter=lfs diff=lfs merge=lfs -text
 static/videos/toby.mp4 filter=lfs diff=lfs merge=lfs -text
 assets/logo.png filter=lfs diff=lfs merge=lfs -text
 assets/silk_bg.jpg filter=lfs diff=lfs merge=lfs -text

 static/videos/toby.mp4 filter=lfs diff=lfs merge=lfs -text
 assets/logo.png filter=lfs diff=lfs merge=lfs -text
 assets/silk_bg.jpg filter=lfs diff=lfs merge=lfs -text
+sirin/ui/assets/logo.png filter=lfs diff=lfs merge=lfs -text
+sirin/ui/assets/silk_bg.jpg filter=lfs diff=lfs merge=lfs -text
+sirin/ui/static/silk_bg.jpg filter=lfs diff=lfs merge=lfs -text
+sirin/ui/workspace/frontend/build/NotoColorEmoji-nature-Rpfd13Si.woff2 filter=lfs diff=lfs merge=lfs -text
+sirin/ui/workspace/frontend/build/NotoColorEmoji-objects-EKxheXEn.woff2 filter=lfs diff=lfs merge=lfs -text

.streamlit/config.toml ADDED Viewed

	@@ -0,0 +1,20 @@

+# Native-widget palette + fonts mirror sirin/ui/tokens.json (palette.light + fonts). The @font-face
+# rules that back these families are injected by sirin/ui/styles.py from app/static/fonts/*.
+[theme]
+base = "light"
+primaryColor = "#E562A8"
+backgroundColor = "#FAF3F8"
+secondaryBackgroundColor = "#FFFFFF"
+textColor = "#231A21"
+linkColor = "#0E9384"
+baseRadius = "12px"
+font = "Manrope"
+codeFont = "IBM Plex Mono"
+[browser]
+gatherUsageStats = false
+[server]
+# Serve sirin/ui/static/* at /app/static/* so the silk backdrop loads from a cached static file
+# instead of a ~259KB base64 data URI re-shipped in the injected CSS on every rerun.
+enableStaticServing = true

README.md CHANGED Viewed

@@ -1,21 +1,28 @@
 ---
-title: SIRIN UI Demo
 emoji: 🔎
 colorFrom: purple
 colorTo: pink
-sdk: static
 tags:
   - hallucination-detection
   - answerability
   - longmemeval
   - demo
-short_description: Interactive replay of verified SIRIN inspection artifacts
 ---
-# SIRIN UI Demo
-An interactive, offline replay of verified SIRIN artifacts for contextual-inconsistency inspection.
-Choose a recorded LongMemEval case to inspect decisive evidence, the recorded response, its uncertainty or sequence score visualization, and run provenance. This public demo does not load model weights, call external APIs, or accept user data.
-For the complete project and live UI workflow, see the [SIRIN repository](https://github.com/sb-ai-lab/SIRIN).

 ---
+title: SIRIN
 emoji: 🔎
 colorFrom: purple
 colorTo: pink
+sdk: streamlit
+sdk_version: "1.59.1"
+python_version: "3.11"
+app_file: sirin/ui/streamlit_app.py
+pinned: false
 tags:
   - hallucination-detection
   - answerability
   - longmemeval
   - demo
+short_description: Hallucination inspection — replay + live API judges
 ---
+# SIRIN
+Interactive workspace for contextual-inconsistency inspection: replay curated
+LongMemEval / PsiloQA cases with recorded evidence and scores, or run live
+LLM-judge detectors against API providers.
+External API judge calls are opt-in per session and use the Space's own key;
+no user data is stored. Local model inference is disabled in this hosted build.
+Full project: https://github.com/sb-ai-lab/SIRIN

index.html DELETED Viewed

@@ -1,144 +0,0 @@
-<!doctype html>
-<html lang="en">
-  <head>
-    <meta charset="utf-8">
-    <meta name="viewport" content="width=device-width, initial-scale=1">
-    <meta name="description" content="Offline replay of verified SIRIN inspection artifacts.">
-    <title>SIRIN UI Demo</title>
-    <style>
-      :root {
-        --bg: #10091c;
-        --panel: rgba(26, 16, 44, .82);
-        --line: rgba(255, 255, 255, .14);
-        --text: #f6f1fc;
-        --muted: #c9bedb;
-        --mint: #a9f0d6;
-        --pink: #f07bbb;
-        --violet: #af8dff;
-        --danger: #ff637f;
-        color-scheme: dark;
-      }
-      * { box-sizing: border-box; }
-      body {
-        min-height: 100vh;
-        margin: 0;
-        color: var(--text);
-        background: linear-gradient(rgba(16, 9, 28, .88), rgba(16, 9, 28, .93)),
-          url("assets/silk_bg.jpg") center / cover fixed, var(--bg);
-        font: 16px/1.55 Inter, ui-sans-serif, system-ui, -apple-system, BlinkMacSystemFont,
-          "Segoe UI", sans-serif;
-      }
-      .shell { width: min(1160px, calc(100% - 32px)); margin: 0 auto; padding: 40px 0 64px; }
-      .hero { display: flex; align-items: center; gap: 16px; margin-bottom: 8px; }
-      .logo { width: 58px; height: 58px; object-fit: contain; }
-      h1, h2, h3, p { margin: 0; }
-      h1 { font-size: clamp(2rem, 5vw, 3.2rem); letter-spacing: -.05em; }
-      h2 { font-size: 1.15rem; margin-bottom: 12px; }
-      .lede, .muted { color: var(--muted); }
-      .notice { margin: 26px 0 18px; padding: 14px 16px; border-left: 3px solid var(--mint); color: var(--muted); background: rgba(169, 240, 214, .08); }
-      .tabs { display: flex; flex-wrap: wrap; gap: 10px; margin: 18px 0 24px; }
-      button {
-        border: 1px solid var(--line); border-radius: 999px; padding: 9px 15px; cursor: pointer;
-        color: var(--text); background: rgba(255, 255, 255, .06); font: inherit; transition: .18s ease;
-      }
-      button:hover, button[aria-selected="true"] { border-color: var(--pink); background: rgba(240, 123, 187, .18); box-shadow: 0 0 0 3px rgba(240, 123, 187, .1); }
-      .grid { display: grid; grid-template-columns: minmax(0, 1.55fr) minmax(280px, .9fr); gap: 18px; }
-      .panel { border: 1px solid var(--line); border-radius: 18px; padding: clamp(18px, 3vw, 28px); background: var(--panel); backdrop-filter: blur(14px); box-shadow: 0 16px 48px rgba(0, 0, 0, .2); }
-      .case-head { display: flex; align-items: start; justify-content: space-between; gap: 16px; }
-      .question { margin-top: 4px; color: var(--muted); font-size: 1.04rem; }
-      .pill { flex: none; border-radius: 999px; padding: 5px 10px; color: #190b20; background: var(--mint); font-size: .8rem; font-weight: 750; }
-      .response { margin-top: 22px; padding: 18px; border-radius: 14px; border: 1px solid rgba(240, 123, 187, .32); background: rgba(240, 123, 187, .08); }
-      .eyebrow { display: block; margin-bottom: 7px; color: var(--pink); font-size: .73rem; font-weight: 800; letter-spacing: .09em; text-transform: uppercase; }
-      .answer { font-size: 1.24rem; font-weight: 750; line-height: 1.9; }
-      .heat { border-bottom: 3px solid var(--pink); border-radius: 4px; padding: .05em .15em; }
-      .gold { margin-top: 11px; color: var(--muted); font-size: .9rem; }
-      .gold strong { color: var(--mint); }
-      .section { margin-top: 22px; }
-      .evidence { display: grid; gap: 10px; }
-      .evidence article { padding: 13px 14px; border-left: 3px solid var(--violet); border-radius: 0 10px 10px 0; background: rgba(175, 141, 255, .09); }
-      .evidence h3 { color: var(--mint); font-size: .85rem; }
-      .evidence p { margin-top: 4px; color: var(--muted); font-size: .94rem; }
-      .metric { display: grid; grid-template-columns: 1fr 1fr; gap: 10px; margin-top: 18px; }
-      .metric div { padding: 12px; border-radius: 10px; background: rgba(255, 255, 255, .055); }
-      .metric span { display: block; color: var(--muted); font-size: .76rem; text-transform: uppercase; letter-spacing: .06em; }
-      .metric strong { display: block; margin-top: 2px; font-size: .9rem; }
-      details { margin-top: 18px; border-top: 1px solid var(--line); padding-top: 14px; }
-      summary { cursor: pointer; color: var(--mint); font-weight: 700; }
-      pre { overflow: auto; margin: 12px 0 0; padding: 12px; border-radius: 10px; background: rgba(0, 0, 0, .28); color: var(--muted); font: 12px/1.5 ui-monospace, SFMono-Regular, Menlo, monospace; white-space: pre-wrap; word-break: break-word; }
-      footer { margin-top: 26px; color: var(--muted); font-size: .85rem; }
-      a { color: var(--mint); }
-      .error { color: #ffd5dd; border-left-color: var(--danger); }
-      @media (max-width: 760px) { .shell { width: min(100% - 24px, 1160px); padding-top: 24px; } .grid { grid-template-columns: 1fr; } .case-head { display: block; } .pill { display: inline-block; margin-top: 10px; } }
-    </style>
-  </head>
-  <body>
-    <main class="shell">
-      <header class="hero">
-        <img class="logo" src="assets/logo.png" alt="SIRIN logo">
-        <div>
-          <h1>SIRIN</h1>
-          <p class="lede">Semantic Inconsistency Recognition &amp; Inspection Nexus</p>
-        </div>
-      </header>
-      <p class="notice">Offline replay of verified LongMemEval artifacts. This public Space does not load model weights, run detectors, make API calls, or retain user data.</p>
-      <nav id="tabs" class="tabs" aria-label="Recorded cases"></nav>
-      <section id="app" aria-live="polite"><p class="muted">Loading verified artifacts…</p></section>
-      <footer>For the full live workflow, visit <a href="https://github.com/sb-ai-lab/SIRIN">the SIRIN repository</a>.</footer>
-    </main>
-    <script>
-      const tabs = document.querySelector('#tabs');
-      const app = document.querySelector('#app');
-      const escapeHtml = value => String(value ?? '').replace(/[&<>'"]/g, character => ({'&':'&amp;','<':'&lt;','>':'&gt;',"'":'&#039;','"':'&quot;'}[character]));
-      const scoreColor = score => `hsl(${Math.round(148 - Math.max(0, Math.min(1, score)) * 148)} 86% 66%)`;
-      function renderedAnswer(caseData) {
-        const trace = caseData.token_trace;
-        if (trace?.token_pieces?.length) {
-          return trace.token_pieces.map((piece, index) => {
-            const score = trace.normalized_scores[index] ?? 0;
-            return `<span class="heat" title="relative token uncertainty: ${score.toFixed(3)}" style="border-bottom-color:${scoreColor(score)};background:${scoreColor(score)}22">${escapeHtml(piece)}</span>`;
-          }).join('');
-        }
-        const score = Number(caseData.score ?? 0);
-        return `<span class="heat" title="recorded sequence score: ${score.toFixed(3)}" style="border-bottom-color:${scoreColor(score)};background:${scoreColor(score)}22">${escapeHtml(caseData.prediction)}</span>`;
-      }
-      function render(caseData) {
-        const provenance = caseData.provenance || {};
-        const metric = caseData.token_trace
-          ? 'Token uncertainty · localization only'
-          : `Recorded sequence score · ${Number(caseData.score).toFixed(3)}`;
-        app.innerHTML = `
-          <div class="grid">
-            <article class="panel">
-              <div class="case-head"><div><h2>${escapeHtml(caseData.title)}</h2><p class="question">${escapeHtml(caseData.question)}</p></div><span class="pill">${metric}</span></div>
-              <section class="response"><span class="eyebrow">Recorded response</span><div class="answer">${renderedAnswer(caseData)}</div><p class="gold">Dataset reference: <strong>${escapeHtml(caseData.gold)}</strong></p></section>
-              <section class="section"><span class="eyebrow">Decisive evidence</span><div class="evidence">${(caseData.evidence_excerpts || []).map(evidence => `<article><h3>${escapeHtml(evidence.cue || evidence.label)} · rank ${escapeHtml(evidence.rank)}${evidence.date ? ` · ${escapeHtml(evidence.date)}` : ''}</h3><p>${escapeHtml(evidence.content)}</p></article>`).join('')}</div></section>
-            </article>
-            <aside class="panel">
-              <span class="eyebrow">Recorded scorer</span><h2>${escapeHtml(caseData.detector)}</h2>
-              <div class="metric"><div><span>Dataset</span><strong>${escapeHtml(provenance.dataset || '—')}-${escapeHtml(provenance.dataset_variant || '—')}</strong></div><div><span>Model</span><strong>${escapeHtml(provenance.model || '—')}</strong></div><div><span>Protocol</span><strong>${escapeHtml(provenance.score_protocol || '—')}</strong></div><div><span>Generation temp.</span><strong>${escapeHtml(provenance.generation_temperature ?? '—')}</strong></div></div>
-              <details><summary>Artifact provenance</summary><pre>${escapeHtml(JSON.stringify({ sample_id: caseData.sample_id, prompt_sha256: caseData.prompt_sha256, messages_sha256: caseData.messages_sha256, detector: caseData.detector, score: caseData.score, provenance }, null, 2))}</pre></details>
-            </aside>
-          </div>`;
-        [...tabs.children].forEach(tab => tab.setAttribute('aria-selected', String(tab.dataset.id === caseData.sample_id)));
-      }
-      fetch('assets/demo_cases.json')
-        .then(response => response.ok ? response.json() : Promise.reject(new Error(`HTTP ${response.status}`)))
-        .then(payload => {
-          const cases = payload.cases || [];
-          if (!cases.length) throw new Error('No recorded cases were found.');
-          cases.forEach((caseData, index) => {
-            const tab = document.createElement('button');
-            tab.type = 'button'; tab.dataset.id = caseData.sample_id; tab.textContent = caseData.title;
-            tab.setAttribute('aria-selected', String(index === 0)); tab.addEventListener('click', () => render(caseData)); tabs.append(tab);
-          });
-          render(cases[0]);
-        })
-        .catch(error => { app.innerHTML = `<p class="notice error">Unable to load the verified demo artifacts: ${escapeHtml(error.message)}</p>`; });
-    </script>
-  </body>
-</html>

pyproject.toml ADDED Viewed

	@@ -0,0 +1,21 @@

+[build-system]
+requires = ["setuptools>=45", "wheel"]
+build-backend = "setuptools.build_meta"
+[project]
+name = "sirin"
+version = "0.1.0"
+description = "SIRIN UI - hosted Space build (deps pinned in requirements.txt)"
+requires-python = ">=3.11,<3.14"
+dependencies = []
+[[tool.streamlit.component.components]]
+name = "sirin_workspace"
+asset_dir = "ui/workspace/frontend/build"
+[tool.setuptools.packages.find]
+include = ["sirin*"]
+[tool.setuptools.package-data]
+"sirin.ui" = ["assets/*", "static/*", "static/fonts/*", "tokens.json"]
+"sirin.ui.workspace" = ["pyproject.toml", "frontend/package.json", "frontend/build/*"]

requirements.txt ADDED Viewed

	@@ -0,0 +1,15 @@

+--extra-index-url https://download.pytorch.org/whl/cpu
+torch==2.9.0+cpu
+streamlit==1.59.1
+transformers==4.57.6
+datasets==3.6.0
+peft==0.17.1
+numpy==1.26.4
+scipy==1.17.1
+scikit-learn==1.6.1
+hydra-core==1.3.2
+joblib==1.5.3
+loguru==0.7.3
+pydantic==2.13.0
+httpx==0.28.1
+-e .

sirin/__init__.py ADDED Viewed

File without changes

sirin/classification/__init__.py ADDED Viewed

	@@ -0,0 +1 @@


1	+ from .linear import LinearClassifier

sirin/classification/linear.py ADDED Viewed

	@@ -0,0 +1,335 @@

+from typing import List, Optional, Tuple
+import torch
+import torch.nn as nn
+class Identity(nn.Module):
+    def __init__(self):
+        super(Identity, self).__init__()
+    def forward(self, x):
+        return x
+class LinearClassifier(nn.Module):
+    def __init__(
+        self,
+        embedding_dim: List[int],
+        num_features: List[int],
+        attention_pooling: List[bool],
+        ensemble: List[int],
+        num_classes: int = 1,
+        use_contrastive: bool = True,
+        projection_dim: int = 128,
+        contrastive_layers: Optional[List[int]] = None,
+        projection_hidden_dim: Optional[int] = None,
+        projection_num_layers: int = 2,
+        use_projection_dropout: bool = False,
+        projection_dropout: float = 0.1,
+    ):
+        super().__init__()
+        self.use_contrastive = use_contrastive
+        self.projection_dim = projection_dim
+        self.projection_hidden_dim = projection_hidden_dim
+        self.projection_num_layers = projection_num_layers
+        self.use_projection_dropout = use_projection_dropout
+        self.projection_dropout = projection_dropout
+        self.contrastive_layers = contrastive_layers or range(len(embedding_dim))
+        self.q = nn.ModuleList([])
+        self.attention_pooling = attention_pooling
+        self.layer_classifiers = nn.ModuleList([Identity() for _ in embedding_dim])
+        for i in ensemble:
+            self.layer_classifiers[i] = nn.Linear(
+                embedding_dim[i] * num_features[i], num_classes
+            )
+        non_ensemble_size = sum([
+            embedding_dim[i] * num_features[i]
+            for i in range(len(num_features))
+            if i not in ensemble
+        ])
+        ensemble_size = len(ensemble)
+        self.final_vector_size = non_ensemble_size + ensemble_size
+        self.classifier = nn.Linear(self.final_vector_size, num_classes)
+        if self.use_contrastive:
+            self.projection_head = self._build_projection_head()
+        for e_d, n_f in zip(embedding_dim, num_features):
+            queries = nn.ParameterList()
+            for i in range(n_f):
+                query = nn.Parameter(torch.empty(1, 1, e_d))
+                nn.init.xavier_uniform_(query)
+                queries.append(query)
+            self.q.append(queries)
+    def _build_projection_head(self) -> nn.Module:
+        """Builds multi-layer projection head for contrastive learning."""
+        projection_layers = []
+        projection_input_dim = self.final_vector_size
+        hidden_dim = self.projection_hidden_dim or (projection_input_dim // 2)
+        # First layer
+        projection_layers.append(nn.Linear(projection_input_dim, hidden_dim))
+        projection_layers.append(nn.BatchNorm1d(hidden_dim))
+        projection_layers.append(nn.ReLU(inplace=True))
+        if self.use_projection_dropout:
+            projection_layers.append(nn.Dropout(self.projection_dropout))
+        # Middle layers (if num_layers > 2)
+        for _ in range(self.projection_num_layers - 2):
+            projection_layers.append(nn.Linear(hidden_dim, hidden_dim))
+            projection_layers.append(nn.BatchNorm1d(hidden_dim))
+            projection_layers.append(nn.ReLU(inplace=True))
+            if self.use_projection_dropout:
+                projection_layers.append(nn.Dropout(self.projection_dropout))
+        # Final projection layer
+        projection_layers.append(nn.Linear(hidden_dim, self.projection_dim))
+        return nn.Sequential(*projection_layers)
+    def _extract_contrastive_features(
+        self,
+        hiddens_list: List[torch.Tensor],
+        attention_masks_list: List[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        contrastive_features = []
+        for layer_idx in self.contrastive_layers:
+            if layer_idx < len(hiddens_list):
+                hiddens = hiddens_list[layer_idx]
+                attention_masks = attention_masks_list[layer_idx] if attention_masks_list else None
+                layer_features = self._pool_layer_features(
+                    hiddens, attention_masks, layer_idx
+                )
+                contrastive_features.append(layer_features)
+        if not contrastive_features:
+            return None
+        if len(contrastive_features) > 1:
+            contrastive_features = torch.cat(contrastive_features, dim=-1)
+        else:
+            contrastive_features = contrastive_features[0]
+        return contrastive_features
+    def _pool_layer_features(
+        self,
+        hiddens: torch.Tensor,
+        attention_masks: Optional[torch.Tensor],
+        layer_idx: int,
+    ) -> torch.Tensor:
+        features = []
+        for i, query in enumerate(self.q[layer_idx]):
+            hidden = hiddens[:, i]
+            if self.attention_pooling[layer_idx]:
+                attention_scores = torch.matmul(hidden, query.transpose(-1, -2)).squeeze(-1)
+                if attention_masks is not None:
+                    attention_mask = attention_masks[:, i]
+                    attention_scores = attention_scores.masked_fill(
+                        attention_mask == 0, float('-inf')
+                    )
+                attention_weights = torch.nn.functional.softmax(attention_scores, dim=-1)
+                pooled = torch.sum(hidden * attention_weights.unsqueeze(-1), dim=1)
+            else:
+                if attention_masks is not None:
+                    attention_mask = attention_masks[:, i]
+                    mask_expanded = attention_mask.unsqueeze(-1)
+                    hidden_masked = hidden * mask_expanded
+                    sum_pooled = torch.sum(hidden_masked, dim=1)
+                    lengths = torch.sum(mask_expanded, dim=1)
+                    pooled = sum_pooled / lengths.clamp(min=1e-9)
+                else:
+                    pooled = torch.mean(hidden, dim=1)
+            features.append(pooled)
+        layer_features = torch.cat(features, dim=-1)
+        return layer_features
+    def forward(
+        self,
+        hiddens_list: List[torch.Tensor],
+        attention_masks_list: List[torch.Tensor] = None,
+        return_projection: bool = False,
+        **kwargs,
+    ) -> torch.Tensor:
+        final = []
+        for index, (hiddens, attention_masks) in enumerate(zip(hiddens_list, attention_masks_list)):
+            context = []
+            for i, query in enumerate(self.q[index]):
+                hidden = hiddens[:, i]
+                attention_mask = (
+                    attention_masks[:, i]
+                    if attention_masks is not None
+                    else torch.ones(hidden.shape[:2])
+                )
+                if self.attention_pooling[index]:
+                    attention_scores = torch.matmul(
+                        hidden, query.transpose(-1, -2)
+                    ).squeeze(-1)
+                    if attention_masks is not None:
+                        attention_mask = attention_masks[:, i]
+                        attention_scores = attention_scores.masked_fill(
+                            attention_mask == 0, float('-inf')
+                        )
+                    attention_weights = torch.nn.functional.softmax(
+                        attention_scores, dim=-1
+                    )
+                    hidden = torch.sum(
+                        hidden * attention_weights.unsqueeze(-1), dim=1
+                    )
+                if hidden.ndim == 4:
+                    valid_counts = attention_mask.sum(dim=1)
+                    last_valid_idx = valid_counts - 1
+                    hidden = hidden[:, last_valid_idx, :]
+                hidden = self.layer_classifiers[index](hidden)
+                context.append(hidden)
+            context_vector = torch.concat(context, dim=1)
+            final.append(context_vector)
+        final_vector = torch.cat(final, dim=-1)
+        logits = self.classifier(final_vector.float())
+        if self.use_contrastive and return_projection:
+            # Get projection for contrastive loss
+            projection = self.get_contrastive_projection(
+                hiddens_list, attention_masks_list
+            )
+            return projection, logits
+        return logits
+    def get_contrastive_projection(
+        self,
+        hiddens_list: List[torch.Tensor],
+        attention_masks_list: List[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if not self.use_contrastive:
+            raise ValueError('Model was not initialized with use_contrastive=True')
+        contrastive_features = self._extract_contrastive_features(
+            hiddens_list, attention_masks_list
+        )
+        if contrastive_features is None:
+            final_vector = self._get_final_vector(hiddens_list, attention_masks_list)
+            projection = self.projection_head(final_vector)
+        else:
+            projection = self.projection_head(contrastive_features)
+        return projection
+    def forward_with_projection(
+        self,
+        hiddens_list: List[torch.Tensor],
+        attention_masks_list: List[torch.Tensor] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        return self.forward(
+            hiddens_list,
+            attention_masks_list,
+            return_projection=True
+        )
+    def _get_final_vector(
+        self,
+        hiddens_list: List[torch.Tensor],
+        attention_masks_list: List[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        final = []
+        for index, (hiddens, attention_masks) in enumerate(zip(hiddens_list, attention_masks_list)):
+            context = []
+            for i, query in enumerate(self.q[index]):
+                hidden = hiddens[:, i]
+                attention_mask = (
+                    attention_masks[:, i]
+                    if attention_masks is not None
+                    else torch.ones(hidden.shape[:2])
+                )
+                if self.attention_pooling[index]:
+                    attention_scores = torch.matmul(
+                        hidden, query.transpose(-1, -2)
+                    ).squeeze(-1)
+                    if attention_masks is not None:
+                        attention_mask = attention_masks[:, i]
+                        attention_scores = attention_scores.masked_fill(
+                            attention_mask == 0, float('-inf')
+                        )
+                    attention_weights = torch.nn.functional.softmax(
+                        attention_scores, dim=-1
+                    )
+                    hidden = torch.sum(
+                        hidden * attention_weights.unsqueeze(-1), dim=1
+                    )
+                if hidden.ndim == 4:
+                    valid_counts = attention_mask.sum(dim=1)
+                    last_valid_idx = valid_counts - 1
+                    hidden = hidden[:, last_valid_idx, :]
+                context.append(hidden)
+            context_vector = torch.concat(context, dim=1)
+            final.append(context_vector)
+        final_vector = torch.cat(final, dim=-1)
+        return final_vector
+    def get_layer_features(
+        self,
+        hiddens_list: List[torch.Tensor],
+        attention_masks_list: List[torch.Tensor] = None,
+        layer_indices: Optional[List[int]] = None,
+    ) -> List[torch.Tensor]:
+        """
+        Extracts features at specified layers for multi-level contrastive learning.
+        Args:
+            hiddens_list: List of hidden states from different layers
+            attention_masks_list: List of attention masks
+            layer_indices: Indices of layers to extract features from (None = all layers)
+        Returns:
+            List of feature tensors, one per specified layer
+        """
+        if layer_indices is None:
+            layer_indices = list(range(len(hiddens_list)))
+        layer_features = []
+        for layer_idx in layer_indices:
+            if layer_idx >= len(hiddens_list):
+                continue
+            hiddens = hiddens_list[layer_idx]
+            attention_masks = attention_masks_list[layer_idx] if attention_masks_list else None
+            # Pool features for this layer
+            features = self._pool_layer_features(hiddens, attention_masks, layer_idx)
+            layer_features.append(features)
+        return layer_features

sirin/configs/detector/probing.yaml ADDED Viewed

	@@ -0,0 +1,8 @@

+_target_: sirin.detection.probing.SequenceTabPFNProbingDetector
+config:
+  num_cpus: 1
+  embedding_dim: null
+  num_features: null
+  dropout_rate: 0.1
+  ensemble: null
+# feature_processor will be passed during instantiation

sirin/configs/feature_processor/hiddens.yaml ADDED Viewed

	@@ -0,0 +1,7 @@

+_target_: sirin.detection.processors.hidden.HiddensProcessor
+config:
+  feature_cache_dir: "${oc.env:HOME}/.cache/feature_cache"
+  token_locator_config: ${token_locator_config}
+  layers: null
+  separate: false
+cache_dir: ${cache_dir}

sirin/configs/model_adapter/default.yaml ADDED Viewed

	@@ -0,0 +1,8 @@

+_target_: sirin.inference.adapters.HfModelAdapter
+config:
+  model_path: "Qwen/Qwen2.5-3B-Instruct"
+  device: "cuda"
+  # Optional: batch_size for automatic batch splitting (prevents OOM)
+  # batch_size: 8  # Uncomment to enable batch processing
+  # See: sirin/configs/model_adapter/hf_with_batching.yaml for example

sirin/configs/model_adapter/hf_large_model.yaml ADDED Viewed

	@@ -0,0 +1,17 @@

+# HuggingFace adapter for large models (7B+)
+# Small batch size + multi-GPU support
+_target_: sirin.inference.adapters.HfModelAdapter
+config:
+  model_path: "meta-llama/Llama-2-7b-chat-hf"
+  device: "cuda"
+  batch_size: 4  # Small batch size for large models
+  # Memory optimization
+  max_length: 1024
+  model_dtype: "bf16"
+  low_cpu_mem_usage: true
+  # Multi-GPU distribution
+  device_map: "auto"  # Automatically distribute across GPUs
+  # max_memory: {0: "20GiB", 1: "20GiB"}  # Uncomment to set per-GPU limits

sirin/configs/model_adapter/hf_with_batching.yaml ADDED Viewed

	@@ -0,0 +1,18 @@

+# HuggingFace adapter with automatic batch processing
+# Use this config for memory-efficient inference with large batches
+_target_: sirin.inference.adapters.HfModelAdapter
+config:
+  model_path: "Qwen/Qwen2.5-3B-Instruct"
+  device: "cuda"
+  batch_size: 8  # Process 8 samples at a time to prevent OOM
+  # Optional: Additional HF-specific parameters
+  max_length: 2048
+  model_dtype: "bf16"
+  padding: "max_length"
+  truncation: true
+  # Optional: Multi-GPU support
+  # device_map: "auto"  # Uncomment for multi-GPU inference
+  # max_memory: {0: "20GiB", 1: "20GiB"}  # Per-GPU memory limits

sirin/configs/model_adapter/openai_with_parallel.yaml ADDED Viewed

	@@ -0,0 +1,16 @@

+# OpenAI adapter with parallel async processing
+# Processes multiple requests concurrently for 10x speedup
+_target_: sirin.inference.adapters.OpenAIModelAdapter
+config:
+  model_path: "gpt-4"
+  base_url: "https://api.openai.com/v1"
+  # batch_size: null  # Not used for OpenAI (uses async instead)
+  # API parameters
+  max_retries: 3
+  timeout: 60
+  # Note: Use use_async=True and max_concurrent in sample() call
+  # Example:
+  #   adapter.sample(inputs, use_async=True, max_concurrent=10)

sirin/configs/model_adapter/openrouter.yaml ADDED Viewed

	@@ -0,0 +1,15 @@

+# OpenRouter adapter (OpenAI-compatible API)
+# Use this for accessing various models through OpenRouter
+_target_: sirin.inference.adapters.OpenAIModelAdapter
+config:
+  model_path: "meta-llama/llama-2-70b-chat"  # Or any OpenRouter model
+  base_url: "https://openrouter.ai/api/v1"
+  # API parameters
+  max_retries: 3
+  timeout: 60
+  proxy_url: null  # Optional: HTTP proxy
+  # Note: Parallel processing via use_async=True in sample() call
+  # adapter.sample(inputs, use_async=True, max_concurrent=5)

sirin/configs/model_adapter/vllm_with_batching.yaml ADDED Viewed

	@@ -0,0 +1,18 @@

+# vLLM adapter with batch processing
+# vLLM provides high-throughput inference with efficient batching
+_target_: sirin.inference.adapters.VllmModelAdapter
+config:
+  model_path: "meta-llama/Llama-2-7b-chat-hf"
+  device: "cuda"
+  batch_size: 16  # vLLM can handle larger batches efficiently
+  # vLLM-specific parameters
+  max_length: 2048
+  enforce_eager: true
+  gpu_memory_utilization: 0.85  # Use 85% of GPU memory
+  # Optional: Additional model kwargs
+  # model_kwargs:
+  #   tensor_parallel_size: 2  # For multi-GPU
+  #   dtype: "bfloat16"

sirin/configs/pipeline/probing.yaml ADDED Viewed

	@@ -0,0 +1,23 @@

+_target_: sirin.detection.probing.ProbingPipeline
+config:
+  save_dir: "./outputs"
+  sampling:
+    temperature: 1.0
+    top_p: 1.0
+    top_k: 50
+    max_length: 512
+    do_sample: true
+  model_manager:
+    max_active_models: 3
+    memory_threshold: 0.8
+    max_retries: 3
+    retry_delay: 1.0
+  save_intermediate: true
+  experiment_name: "example_probe"
+  f_beta: 1.0
+  train_args: ${train_args}
+task_type: ${task_type}
+target_approximator: null
+experiment_logger: null
+# detector, train_dataset, eval_dataset, and generator_adapter will be passed during instantiation

sirin/configs/token_locator_config/default.yaml ADDED Viewed

	@@ -0,0 +1 @@


1	+ locate_answer_start: true

sirin/configs/train.yaml ADDED Viewed

	@@ -0,0 +1,24 @@

+# Top-level config which composes smaller configs via Hydra defaults
+defaults:
+  - model_adapter: default
+  - token_locator_config: default
+  - feature_processor: hiddens
+  - detector: probing
+  - pipeline: probing
+  - train_args: default
+# Runtime params
+cuda_visible_devices: "0"
+cache_dir: "./cache"
+use_cache: true
+task_type: null
+log_level: DEBUG
+random_seed: 42
+change_nested_random_seeds: true
+train_dataset_path: ???
+eval_dataset_path: ???
+hf_token : null
+openai_api_key: null

sirin/configs/train_args/default.yaml ADDED Viewed

	@@ -0,0 +1,10 @@

+max_epochs: 5
+validation_interval: 1
+val_size: 0.1
+learning_rate: 0.001
+train_batch_size: 8
+val_batch_size: 8
+threshold: 0.5
+loss_function: "bce"
+alpha_scheduler: "exp"
+metrics: null

sirin/definitions/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@

+from .constants import *
+from .enums import *
+from .types import *

sirin/definitions/constants.py ADDED Viewed

	@@ -0,0 +1,20 @@

+HF_TOKEN_ENV = "HF_TOKEN"
+OPENAI_API_KEY_ENV = "OPENAI_API_KEY"
+INPUT_COL = "input"
+INPUT_PROC_COL = "processed_input"
+REFERENCE_COL = "reference"
+TARGET_COL = "target"
+ANSWER_INDICES = "answer_indices"
+OFFSETS_COL = "offsets"
+GROUP_ID_COL = "group_id"
+GENERATED_DSET_KEY = "generated_dataset"
+CLEANED_DSET_KEY = "cleaned_dataset"
+DSET_KEY = "dataset"
+METRICS_KEY = "metrics"
+TORCH_MODULE_KEY = "torch_module"
+HIDDENS_FILE_KEY = "hidden_file"
+ARTIFACTS_YAML = "artifacts.yaml"
+SAVE_FEAT_TEMPLATE = "{feature_type}_{sample_hash}_layer_{layer_idx}.pkl"

sirin/definitions/enums.py ADDED Viewed

	@@ -0,0 +1,204 @@

+from enum import Enum
+from typing import Callable
+from sklearn.metrics import (
+    accuracy_score,
+    auc,
+    f1_score,
+    fbeta_score,
+    precision_recall_curve,
+    precision_score,
+    recall_score,
+    roc_auc_score,
+    average_precision_score
+)
+class LogLevel(Enum):
+    DEBUG = 'debug'
+    INFO = 'info'
+    WARNING = 'warning'
+    ERROR = 'error'
+    CRITICAL = 'critical'
+class ModelType(Enum):
+    CAUSAL = 'causal'
+    BASE = 'base'
+    TOKEN_CLASSIFICATION = 'token_classification'
+    SEQUENCE_CLASSIFICATION = 'sequence_classification'
+class ModelAdapterType(Enum):
+    HUGGINGFACE = 'hf'
+    VLLM = 'vllm'
+    OPENAI = 'openai'
+class DetectionTaskType(Enum):
+    HALLUCINATION_DETECTION = 'hallucination'
+    QUERY_ANSWERABILITY = 'answerability'
+class FeatureType(Enum):
+    HIDDEN = 'hidden'
+    ATTENTION = 'attention'
+    LOOKBACK = 'lookback'
+    LOGIT = 'logit'
+    SUBLAYER = 'sublayer'
+    TOKEN_UNCERTAINTY = 'token_uncertainty'
+    SEQUENCE_UNCERTAINTY = 'sequence_uncertainty'
+class SplitStrategy(Enum):
+    SENTENCE = 'sentence'
+    PARAGRAPH = 'paragraph'
+    CHARACTER = 'character'
+    LANGCHAIN = 'langchain'
+    ATOMIC = 'atomic_facts'
+class AggregationMethod(Enum):
+    MEAN = 'mean'
+    MAX = 'max'
+    MIN = 'min'
+    VOTE = 'vote'
+class SideType(Enum):
+    LEFT = 'left'
+    RIGHT = 'right'
+    INNER = 'inner'
+    OUTER = 'outer'
+class AggregationType(Enum):
+    MICRO = 'micro'
+    MACRO = 'macro'
+    CONCAT = 'concat'
+class DetectionLevel(Enum):
+    TOKEN = 'token'
+    SEQUENCE = 'sequence'
+    CLAIM = 'claim'
+class DataCollatorType(Enum):
+    """Data collator types for model training.
+    - TOKEN: For token-level classification with -100 padding mask
+    - PADDING: For sequence-level classification with scalar labels
+    """
+    TOKEN = 'token'
+    PADDING = 'padding'
+class Phase(Enum):
+    TRAIN = 'train'
+    TEST = 'test'
+    VAL = 'val'
+class CompressionMethod(Enum):
+    """Supported compression methods."""
+    PCA = 'pca'
+    UMAP = 'umap'
+    NONE = 'none'
+class ScalingMethod(Enum):
+    """Supported scaling methods."""
+    STANDARD = 'standard'  # Z-score normalization: (x - mean) / std
+    MINMAX = 'minmax'  # Min-Max scaling: (x - min) / (max - min)
+    ROBUST = 'robust'  # Robust scaling using median and IQR
+    NONE = 'none'  # No scaling
+class LmMetric(Enum):
+    """Standard NLP evaluation metrics constants."""
+    ROUGE_1 = 'rouge_1'
+    ROUGE_2 = 'rouge_2'
+    ROUGE_L = 'rouge_l'
+    BLEU = 'bleu'
+    TER = 'ter'
+    BERT_SCORE = 'bert_score'
+    BERT_SCORE_PRECISION = 'bert_score_precision'
+    BERT_SCORE_RECALL = 'bert_score_recall'
+    BERT_SCORE_F1 = 'bert_score_f1'
+    METEOR = 'meteor'
+    @classmethod
+    def get_rouge_metrics(cls) -> list['LmMetric']:
+        """Get all ROUGE-based metrics."""
+        return [cls.ROUGE_1, cls.ROUGE_2, cls.ROUGE_L]
+    @classmethod
+    def get_bert_metrics(cls) -> list['LmMetric']:
+        """Get all BERT-based metrics."""
+        return [cls.BERT_SCORE_PRECISION, cls.BERT_SCORE_RECALL, cls.BERT_SCORE_F1]
+class ClassificationMetric(Enum):
+    ROC_AUC = 'roc_auc'
+    F1 = 'f1'
+    ACCURACY = 'accuracy'
+    PR_AUC = 'pr_auc'
+    PRECISION = 'precision'
+    RECALL = 'recall'
+    FBETA = 'fbeta'
+    AP = 'average_precision'
+    @staticmethod
+    def get_metric_function(metric: str) -> Callable:
+        def pr_auc(y, y_pred_score):
+            precision, recall, _ = precision_recall_curve(y, y_pred_score)
+            return auc(recall, precision)
+        if metric == ClassificationMetric.ROC_AUC.value:
+            return roc_auc_score
+        elif metric == ClassificationMetric.F1.value:
+            return f1_score
+        elif metric == ClassificationMetric.ACCURACY.value:
+            return accuracy_score
+        elif metric == ClassificationMetric.PR_AUC.value:
+            return pr_auc
+        elif metric == ClassificationMetric.PRECISION.value:
+            return precision_score
+        elif metric == ClassificationMetric.RECALL.value:
+            return recall_score
+        elif metric == ClassificationMetric.FBETA.value:
+            return fbeta_score
+        elif metric == ClassificationMetric.AP.value:
+            return average_precision_score
+        else:
+            raise NotImplementedError(f'Unsupported metric: {metric}')
+    def __str__(self):
+        return self.value
+    def startswith(self, prefix):
+        return self.value.startswith(prefix)
+BASIC_METRICS = [
+    ClassificationMetric.F1,
+    ClassificationMetric.ROC_AUC,
+    ClassificationMetric.ACCURACY,
+    ClassificationMetric.PR_AUC,
+    ClassificationMetric.PRECISION,
+    ClassificationMetric.RECALL,
+    ClassificationMetric.AP
+]
+class TokenLocation(Enum):
+    ANS_END = 'answer_end'
+    ANS_START = 'answer_start'
+    ANS_MID = 'answer_middle'
+    EOS = 'eos'
+    N_TO_ANS_START = 'n_to_answer_start'
+    N_TO_ANS_END = 'n_to_answer_end'
+    SUBSTRING = 'substring'

sirin/definitions/types.py ADDED Viewed

	@@ -0,0 +1,35 @@

+from enum import Enum
+from typing import Union
+import torch
+from omegaconf import DictConfig, ListConfig
+CfgDictType = Union[dict, DictConfig]
+CfgListType = Union[list, ListConfig]
+class TorchDtype(Enum):
+    bf16 = torch.bfloat16
+    f64 = torch.float64
+    f32 = torch.float32
+    f16 = torch.float16
+    c32 = torch.complex32
+    c64 = torch.complex64
+    c128 = torch.complex128
+    i8 = torch.int8
+    i16 = torch.int16
+    i32 = torch.int32
+    i64 = torch.int64
+    b8 = torch.bool
+    u8 = torch.uint8
+    @staticmethod
+    def from_str(dtype_name: str) -> torch.dtype:
+        return TorchDtype[dtype_name].value
+    @staticmethod
+    def from_dtype(dtype: torch.dtype) -> str:
+        for dtype_member in TorchDtype:
+            if dtype_member.value == dtype:
+                return dtype_member.name
+        raise ValueError(f"Unknown dtype: {dtype}.")

sirin/detection/__init__.py ADDED Viewed

File without changes

sirin/detection/_lm_polygraph_compat.py ADDED Viewed

	@@ -0,0 +1,168 @@

+import sys
+import types
+from transformers import AutoProcessor, GenerationConfig
+import transformers
+transformers = sys.modules['transformers']
+if not hasattr(transformers, 'HybridCache') and hasattr(transformers, 'DynamicCache'):
+    transformers.HybridCache = transformers.DynamicCache
+from transformers.generation.utils import (
+    GenerateBeamDecoderOnlyOutput,
+    GenerateDecoderOnlyOutput,
+)
+gen_utils = sys.modules['transformers.generation.utils']
+# ensemble_beam.py: from transformers.generation.beam_search import BeamScorer
+# Module was removed in transformers 5.x
+if 'transformers.generation.beam_search' not in sys.modules:
+    try:
+        from transformers.utils.dummy_pt_objects import BeamScorer as _BeamScorer
+    except ImportError:
+        _BeamScorer = None
+    if _BeamScorer is not None:
+        beam_search_mod = types.ModuleType('transformers.generation.beam_search')
+        beam_search_mod.BeamScorer = _BeamScorer
+        sys.modules['transformers.generation.beam_search'] = beam_search_mod
+# ensemble_beam.py: BeamSearchOutput, BeamSearchDecoderOnlyOutput
+if not hasattr(gen_utils, 'BeamSearchOutput'):
+    gen_utils.BeamSearchOutput = GenerateBeamDecoderOnlyOutput
+if not hasattr(gen_utils, 'BeamSearchDecoderOnlyOutput'):
+    gen_utils.BeamSearchDecoderOnlyOutput = GenerateBeamDecoderOnlyOutput
+# ensemble_greedy.py: GreedySearchOutput, GreedySearchDecoderOnlyOutput
+if not hasattr(gen_utils, 'GreedySearchOutput'):
+    gen_utils.GreedySearchOutput = GenerateDecoderOnlyOutput
+if not hasattr(gen_utils, 'GreedySearchDecoderOnlyOutput'):
+    gen_utils.GreedySearchDecoderOnlyOutput = GenerateDecoderOnlyOutput
+# ensemble_sample.py: SampleOutput, SampleDecoderOnlyOutput
+if not hasattr(gen_utils, 'SampleOutput'):
+    gen_utils.SampleOutput = GenerateDecoderOnlyOutput
+if not hasattr(gen_utils, 'SampleDecoderOnlyOutput'):
+    gen_utils.SampleDecoderOnlyOutput = GenerateDecoderOnlyOutput
+# visual_whitebox_model.py: AutoModelForVision2Seq
+if not hasattr(transformers, 'AutoModelForVision2Seq'):
+    transformers.AutoModelForVision2Seq = transformers.AutoModelForImageTextToText
+# greedy_probs.py: out.scores are raw logits, not log-probabilities.
+# GreedyProbsCalculator treats them as log-probs, so Perplexity/MaxSequenceProb
+# get garbage values (and inf from -inf logits in MoE models).
+# VisualWhiteboxModel has a _ScoresProcessor that applies log_softmax, but
+# WhiteboxModelBasic does not.
+import torch as _torch
+from lm_polygraph.model_adapters.whitebox_model_basic import WhiteboxModelBasic
+_original_wb_generate = WhiteboxModelBasic.generate
+def _patched_wb_generate(self, *args, **kwargs):
+    out = _original_wb_generate(self, *args, **kwargs)
+    if hasattr(out, 'scores') and out.scores is not None:
+        out.scores = tuple(s.log_softmax(dim=-1) for s in out.scores)
+    return out
+WhiteboxModelBasic.generate = _patched_wb_generate
+import openai as _openai
+from lm_polygraph.utils.model import BlackboxModel
+_original_blackbox_init = BlackboxModel.__init__
+def _patched_blackbox_init(self, openai_api_key=None, model_path=None,
+                           hf_api_token=None, generation_parameters=None,
+                           supports_logprobs=False, base_url=None):
+    if generation_parameters is None:
+        from lm_polygraph.utils.generation_parameters import GenerationParameters
+        generation_parameters = GenerationParameters()
+    _original_blackbox_init(
+        self,
+        openai_api_key=openai_api_key,
+        model_path=model_path,
+        hf_api_token=hf_api_token,
+        generation_parameters=generation_parameters,
+        supports_logprobs=supports_logprobs,
+    )
+    self.base_url = base_url
+    if openai_api_key is not None and base_url is not None:
+        self.openai_api = _openai.OpenAI(api_key=openai_api_key, base_url=base_url)
+BlackboxModel.__init__ = _patched_blackbox_init
+@staticmethod
+def _patched_from_openai(openai_api_key=None, model_path=None,
+                         supports_logprobs=False, **kwargs):
+    from lm_polygraph.utils.generation_parameters import GenerationParameters
+    generation_parameters = kwargs.pop('generation_parameters', GenerationParameters())
+    base_url = kwargs.pop('base_url', None)
+    return BlackboxModel(
+        openai_api_key=openai_api_key,
+        model_path=model_path,
+        supports_logprobs=supports_logprobs,
+        generation_parameters=generation_parameters,
+        base_url=base_url,
+    )
+BlackboxModel.from_openai = _patched_from_openai
+def load_stat_calculator(cfg, env):
+    from lm_polygraph.stat_calculators.greedy_probs_blackbox import BlackboxGreedyTextsCalculator
+    top_logprobs = cfg.get('top_logprobs', 5)
+    return BlackboxGreedyTextsCalculator(top_logprobs=top_logprobs)
+from lm_polygraph.defaults import register_default_stat_calculators as _reg_module
+_original_register = _reg_module.register_default_stat_calculators
+def _patched_register(model_type, language='en', hf_cache=None,
+                      blackbox_supports_logprobs=False, top_logprobs=5,
+                      output_attentions=True, output_hidden_states=True,
+                      deberta_batch_size=10):
+    result = _original_register(
+        model_type=model_type,
+        language=language,
+        hf_cache=hf_cache,
+        blackbox_supports_logprobs=blackbox_supports_logprobs,
+        output_attentions=output_attentions,
+        output_hidden_states=output_hidden_states,
+        deberta_batch_size=deberta_batch_size,
+    )
+    for container in result:
+        if container.name == 'BlackboxGreedyTextsCalculator':
+            from omegaconf import OmegaConf
+            cfg = OmegaConf.to_container(container.cfg, resolve=True)
+            cfg['top_logprobs'] = top_logprobs
+            container.cfg = OmegaConf.create(cfg)
+            container.builder = 'sirin.detection._lm_polygraph_compat'
+    return result
+_reg_module.register_default_stat_calculators = _patched_register
+# greedy_alternatives_nli.py: deberta.deberta_tokenizer.batch_encode_plus(...)
+# DebertaTokenizer removed batch_encode_plus in transformers 5.x.
+# Patch: add it back as a wrapper around __call__ (same semantics).
+from transformers import DebertaTokenizer as _DebertaTokenizer
+if not hasattr(_DebertaTokenizer, 'batch_encode_plus'):
+    def _batch_encode_plus(self, *args, **kwargs):
+        return self(*args, **kwargs)
+    _DebertaTokenizer.batch_encode_plus = _batch_encode_plus

sirin/detection/approximators/__init__.py ADDED Viewed

	@@ -0,0 +1,18 @@

+from .base import TargetApproximatorBase
+__all__ = ["TargetApproximatorBase", "SEPTargetApproximator"]
+def __getattr__(name):
+    if name != "SEPTargetApproximator":
+        raise AttributeError(name)
+    try:
+        from .sep import SEPTargetApproximator
+    except ModuleNotFoundError as exc:
+        if exc.name == "lm_polygraph":
+            raise ModuleNotFoundError(
+                "SEPTargetApproximator requires the optional lm_polygraph dependency"
+            ) from exc
+        raise
+    return SEPTargetApproximator

sirin/detection/approximators/base.py ADDED Viewed

	@@ -0,0 +1,13 @@

+from abc import ABC, abstractmethod
+from typing import Any, Dict, List
+from sirin.inference.adapters import ModelAdapterBase
+class TargetApproximatorBase(ABC):
+    def __init__(self, extractor: ModelAdapterBase):
+        self.extractor = extractor
+    @abstractmethod
+    def __call__(self, samples: List[List[Dict]], **kwargs) -> Any:
+        pass

sirin/detection/approximators/sep.py ADDED Viewed

	@@ -0,0 +1,224 @@

+import inspect
+from typing import Dict, List, Optional
+import numpy as np
+import torch
+from loguru import logger as lg
+import sirin.detection._lm_polygraph_compat  # noqa: F401 — must be before lm-polygraph imports
+from sirin.detection.utils.math import binarize_entropy
+from sirin.inference.adapters import ModelAdapterBase
+from sirin.detection.approximators import TargetApproximatorBase
+from sirin.metrics.classification import calculate_classification_metrics
+from sirin.definitions import ClassificationMetric
+def _log_roc_auc(labels: List[int], scores) -> None:
+    try:
+        lg.info(
+            f"ROC_AUC_SCORE {calculate_classification_metrics(labels, scores, metrics=[ClassificationMetric.ROC_AUC])} #"
+        )
+    except ValueError as e:
+        lg.debug(f"ROC-AUC logging skipped: {e}")
+class SEPTargetApproximator(TargetApproximatorBase):
+    def __init__(
+        self,
+        extractor: ModelAdapterBase,
+        num_return_sequences: int = 5,
+        sampling_kwargs: Optional[Dict] = None,
+    ):
+        super().__init__(extractor=extractor)
+        from lm_polygraph.utils.deberta import Deberta
+        self.nli_model = Deberta(batch_size=8)
+        self.stats = {}
+        self.num_return_sequences = num_return_sequences
+        self.sampling_kwargs = sampling_kwargs or {}
+    def _ensure_logprob_support(self) -> None:
+        params = inspect.signature(self.extractor.sample).parameters
+        supports_kwargs = any(
+            param.kind == inspect.Parameter.VAR_KEYWORD
+            for param in params.values()
+        )
+        if 'return_logprobs' not in params and not supports_kwargs:
+            raise NotImplementedError(
+                f"SEP requires an adapter with return_logprobs support; got {type(self.extractor).__name__}"
+            )
+    def _sample_with_logprobs(self, samples: List[List[Dict]]):
+        self._ensure_logprob_support()
+        return self.extractor.sample(
+            samples,
+            return_logprobs=True,
+            num_return_sequences=self.num_return_sequences,
+            max_tokens=100,
+            temperature=1.0,
+            **getattr(self, 'sampling_kwargs', {}),
+        )
+    @staticmethod
+    def _is_token_logprob(token_logprob) -> bool:
+        if isinstance(token_logprob, (int, float, np.floating)):
+            return True
+        if not isinstance(token_logprob, (tuple, list)) or len(token_logprob) == 0:
+            return False
+        first = token_logprob[0]
+        return isinstance(first, (str, int, float, np.floating))
+    @staticmethod
+    def _is_sequence_logprobs(value) -> bool:
+        return isinstance(value, list) and (
+            len(value) == 0 or SEPTargetApproximator._is_token_logprob(value[0])
+        )
+    @staticmethod
+    def _primary_logprob(token_logprob) -> float:
+        if isinstance(token_logprob, (int, float, np.floating)):
+            return float(token_logprob)
+        if len(token_logprob) >= 2 and isinstance(token_logprob[0], str):
+            return float(token_logprob[1])
+        return float(token_logprob[0])
+    def _normalize_sample_outputs(self, responses, logits, sample_count: int):
+        if responses and isinstance(responses[0], str):
+            if (
+                self.num_return_sequences > 1
+                and len(responses) == sample_count * self.num_return_sequences
+            ):
+                sample_texts = [
+                    responses[i : i + self.num_return_sequences]
+                    for i in range(0, len(responses), self.num_return_sequences)
+                ]
+            else:
+                sample_texts = [[response] for response in responses]
+        else:
+            sample_texts = responses
+        if (
+            logits
+            and self.num_return_sequences > 1
+            and len(logits) == sample_count * self.num_return_sequences
+            and all(self._is_sequence_logprobs(logits_for_sample) for logits_for_sample in logits)
+        ):
+            sample_logits = [
+                logits[i : i + self.num_return_sequences]
+                for i in range(0, len(logits), self.num_return_sequences)
+            ]
+        else:
+            sample_logits = []
+            for logits_for_sample in logits:
+                if self._is_sequence_logprobs(logits_for_sample):
+                    sample_logits.append([logits_for_sample])
+                else:
+                    sample_logits.append(logits_for_sample)
+        return sample_texts, self._reshape_logprobs(sample_logits)
+    @staticmethod
+    def _reshape_logprobs(logits) -> List[List[float]]:
+        """Convert adapter logprob output to per-sequence token logprob sums."""
+        result = []
+        for gen in logits:
+            seq_scores = []
+            for sample in gen:
+                if not sample:
+                    seq_scores.append(0.0)
+                    continue
+                seq_scores.append(
+                    float(
+                        np.array(
+                            [
+                                SEPTargetApproximator._primary_logprob(token)
+                                for token in sample
+                            ]
+                        ).sum()
+                    )
+                )
+            result.append(seq_scores)
+        return result
+    @staticmethod
+    def _call_stat_calculator(calculator, *args):
+        signature = inspect.signature(calculator.__call__)
+        if any(
+            param.kind == inspect.Parameter.VAR_POSITIONAL
+            for param in signature.parameters.values()
+        ):
+            return calculator(*args)
+        positional_count = len(
+            [
+                param
+                for param in signature.parameters.values()
+                if param.kind
+                in (
+                    inspect.Parameter.POSITIONAL_ONLY,
+                    inspect.Parameter.POSITIONAL_OR_KEYWORD,
+                )
+            ]
+        )
+        return calculator(*args[:positional_count])
+    def _compute_semantic_entropy(
+        self, samples: List[List[Dict]], labels: Optional[List[int]] = None
+    ) -> List[int]:
+        from lm_polygraph.estimators import SemanticEntropy
+        from lm_polygraph.stat_calculators import (
+            SemanticClassesCalculator,
+            SemanticMatrixCalculator,
+        )
+        responses, logits = self._sample_with_logprobs(samples)
+        sample_texts, logprob_sums = self._normalize_sample_outputs(
+            responses,
+            logits,
+            sample_count=len(samples),
+        )
+        new_stats = {
+            'sample_texts': sample_texts,
+            'sample_log_probs': logprob_sums,
+        }
+        self.stats.update(new_stats)
+        semantic_matrix_calculator = SemanticMatrixCalculator(nli_model=self.nli_model)
+        semantic_matrix = self._call_stat_calculator(
+            semantic_matrix_calculator,
+            self.stats,
+            samples,
+            self.extractor,
+        )
+        self.stats.update(semantic_matrix)
+        semantic_classes_calculator = SemanticClassesCalculator()
+        semantic_classes = self._call_stat_calculator(
+            semantic_classes_calculator,
+            self.stats,
+            samples,
+            self.extractor,
+        )
+        self.stats.update(semantic_classes)
+        semantic_entropy_calculator = SemanticEntropy()
+        semantic_entropy = semantic_entropy_calculator(self.stats)
+        new_labels, split = binarize_entropy(torch.tensor(semantic_entropy))
+        self.split = split
+        if labels is not None:
+            _log_roc_auc(labels, semantic_entropy)
+        return new_labels.tolist()
+    def fit(
+        self, samples: List[List[Dict]], labels: Optional[List[int]] = None
+    ) -> List[int]:
+        return self._compute_semantic_entropy(samples, labels)
+    def __call__(
+        self, samples: List[List[Dict]], labels: Optional[List[int]] = None
+    ) -> List[int]:
+        return self._compute_semantic_entropy(samples, labels)

sirin/detection/base.py ADDED Viewed

	@@ -0,0 +1,497 @@

+import random
+from abc import ABC, abstractmethod
+from typing import Any, Dict, List, Optional, Tuple
+from pathlib import Path
+import datasets
+import numpy as np
+import torch
+import torch.multiprocessing as mp
+import transformers
+from loguru import logger as lg
+from sirin.definitions import (
+    DSET_KEY,
+    INPUT_COL,
+    INPUT_PROC_COL,
+    REFERENCE_COL,
+    TARGET_COL,
+    GROUP_ID_COL,
+    DetectionTaskType,
+    LmMetric,
+)
+from sirin.detection.approximators import TargetApproximatorBase
+from sirin.detection.utils.torch import InputsDataset
+from sirin.utils.hf import get_dataset_identifier
+from sirin.inference.adapters import ModelAdapterBase
+from sirin.inference.model_manager import ModelManager
+from sirin.inference.cleaners import (
+    BaseCleaner,
+    CodeCleaner,
+    CustomCleaner,
+    HTMLCleaner,
+    RegexCleaner,
+)
+from sirin.detection.splitters import SplitManager
+from sirin.detection.utils.logits import logits_to_probs_preds
+from sirin.loggers import LoggerBase
+from sirin.metrics import calculate_lm_metrics
+from sirin.models.detection import (
+    BertScoreConfig,
+    DetectionResult,
+    DetectorBaseConfig,
+    PipelineBaseConfig,
+    TrainingArgsConfig,
+)
+from sirin.models.inference import CleanerConfig
+from sirin.utils.savers import DatasetSaver
+class DetectorBase(ABC):
+    """Abstract base class for different detection methods."""
+    model: Any
+    def __init__(
+        self,
+        config: DetectorBaseConfig,
+    ):
+        self.threshold = 0.5
+        self.config = config
+        self.device = self.config.device
+        self.num_cpus = max(1, min(self.config.num_cpus, mp.cpu_count() - 2))
+        self.set_all_seeds(self.config.seed)
+        self._context_splitter = (
+            SplitManager(
+                config=config.context_split_config,
+                split_response=getattr(
+                    config.context_split_config, 'split_response', False
+                ),
+            )
+            if config.context_split_config is not None
+            else None
+        )
+    @classmethod
+    def set_all_seeds(cls, seed: int):
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)  # for multi-GPU
+        torch.backends.cudnn.deterministic = True
+        torch.backends.cudnn.benchmark = False
+        transformers.set_seed(seed)
+    @abstractmethod
+    def setup_model(self, model: Optional[ModelAdapterBase] = None):
+        pass
+    @abstractmethod
+    def detect(self, sample: Any, **kwargs) -> Dict[str, Any]:
+        pass
+    def _validate_and_create_save_dir(self, filepath: Optional[str]) -> Path:
+        if filepath is None:
+            raise ValueError("Filepath must be provided for saving.")
+        save_dir = Path(filepath)
+        # Ensure we're saving to a directory, not a file
+        if save_dir.suffix:
+            raise ValueError(
+                "Please provide a directory path (folder name) for saving, "
+                "not a file path with extension."
+            )
+        save_dir.mkdir(parents=True, exist_ok=True)
+        return save_dir
+    def _split_context_samples(
+        self, samples: List[Any]
+    ) -> Tuple[List[Any], Optional[List[int]]]:
+        if not self._context_splitter:
+            return samples, None
+        split_samples = self._context_splitter.split_inputs(samples)
+        group_ids = []
+        flattened_samples = []
+        for index, splits in enumerate(split_samples):
+            group_ids.extend([index] * len(splits))
+            flattened_samples.extend(splits)
+        return flattened_samples, group_ids
+    def _aggregate_context_predictions(
+        self,
+        group_ids: Optional[List[int]],
+        preds: Any,
+        probs: Any,
+        binary: bool = True,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        if not self._context_splitter or group_ids is None:
+            return preds, probs
+        return self._context_splitter.aggregate_predictions(
+            group_ids, preds, probs, binary=binary
+        )
+    def _logits_to_probs_preds(
+        self, logits: torch.Tensor, threshold: Optional[float] = None
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+        threshold = self.threshold if threshold is None else threshold
+        return logits_to_probs_preds(logits, threshold)
+    @abstractmethod
+    def _save_config(self, save_dir: Path) -> None:
+        pass
+    @abstractmethod
+    def _load_config(self, load_dir: Path) -> None:
+        pass
+    def save(self, filepath: Optional[str] = None):
+        # Use model_save_path from config if filepath not provided
+        if filepath is None:
+            filepath = getattr(self.config, 'model_save_path', None)
+        save_dir = self._validate_and_create_save_dir(filepath)
+        # Save config (required, implemented by subclass)
+        self._save_config(save_dir)
+        # Save additional components if implemented
+        if hasattr(self, '_save_additional_components'):
+            self._save_additional_components(save_dir)
+        lg.info(f"Successfully saved to {save_dir}")
+    def load(self, filepath: str):
+        load_dir = Path(filepath)
+        if not load_dir.exists():
+            raise FileNotFoundError(f"Load directory not found: {load_dir}")
+        if not load_dir.is_dir():
+            raise ValueError(f"Expected directory path, got file: {load_dir}")
+        # Load config (required, implemented by subclass)
+        self._load_config(load_dir)
+        # Load additional components if implemented
+        if hasattr(self, '_load_additional_components'):
+            self._load_additional_components(load_dir)
+        lg.info(f"Successfully loaded from {load_dir}")
+    @abstractmethod
+    def train(
+        self,
+        cfg: TrainingArgsConfig,
+        train_data: InputsDataset,
+        val_data: Optional[InputsDataset],
+        logger: Any = None,
+    ) -> DetectionResult:
+        pass
+class PipelineBase(ABC):
+    """Abstract base class for different trainers of detectors."""
+    def __init__(
+        self,
+        config: PipelineBaseConfig,
+        detector: DetectorBase,
+        train_dataset: datasets.Dataset,
+        target_col: str = TARGET_COL,
+        eval_dataset: Optional[datasets.Dataset] = None,
+        task_type: str = DetectionTaskType.HALLUCINATION_DETECTION,
+        generator_adapter: Optional[ModelAdapterBase] = None,
+        target_approximator: Optional[TargetApproximatorBase] = None,
+        experiment_logger: Optional[LoggerBase] = None,
+    ):
+        self.config = config
+        self.detector = detector
+        self.train_dataset = train_dataset
+        self.eval_dataset = eval_dataset
+        self.task_type = task_type
+        self.target_col = target_col
+        self._generator_adapter = generator_adapter
+        self.target_approximator = target_approximator
+        self.experiment_logger = experiment_logger
+    @abstractmethod
+    def train(self) -> DetectionResult:
+        pass
+    @abstractmethod
+    def eval(self) -> DetectionResult:
+        pass
+    def _load_dataset(
+        self,
+        dataset: datasets.Dataset | datasets.DatasetDict,
+        split: Optional[str] = None,
+    ) -> Tuple[InputsDataset, DatasetSaver]:
+        """Load and process dataset through the pipeline."""
+        split_name = split or ''
+        dataset_name = get_dataset_identifier(dataset)
+        if split and hasattr(dataset, 'keys') and split in dataset:
+            dataset = dataset[split]
+        if self._generator_adapter and self._generator_adapter.name:
+            model_name = self._generator_adapter.name.split('/')[-1].replace('.', '')
+        else:
+            model_name = ''
+        assert self.target_col in dataset.column_names, (
+            f"Target column '{self.target_col}' not found in dataset."
+        )
+        assert INPUT_COL in dataset.column_names, (
+            f"Input column '{INPUT_COL}' not found in dataset."
+        )
+        dataset_saver = (
+            DatasetSaver(
+                dataset_name,
+                save_dir=self.config.save_dir,
+                file_type_alias=DSET_KEY,
+            )
+            if self.config.save_intermediate
+            else None
+        )
+        stages = [
+            ('generation', self._generate_answers, {}),
+            ('cleaning', self._clean_generation_outputs, {}),
+            ('metrics', self._compute_lm_metrics, {}),
+        ]
+        if split == 'train' and self.config.split_context_train:
+            stages.append(('splitting', self._split_context, {}))
+        for stage_name, processor, kwargs in stages:
+            lg.info(f"Starting stage: {stage_name}")
+            dataset = processor(dataset, **kwargs)
+            if dataset_saver is not None:
+                dataset_saver.save(
+                    dataset,
+                    label='/'.join([model_name, stage_name, split_name]).strip('/'),
+                    description=f"Dataset {dataset_name} (split={split_name}, model={model_name}) after {stage_name} stage.",
+                )
+        dataset = InputsDataset(
+            dataset[INPUT_COL],
+            dataset[self.target_col],
+            dataset[GROUP_ID_COL] if GROUP_ID_COL in dataset.column_names else None,
+        )
+        return dataset, dataset_saver
+    def _split_context(self, dataset: datasets.Dataset) -> datasets.Dataset:
+        if self.detector._context_splitter:
+            return self.detector._context_splitter.split_dataset(dataset)
+        lg.info("No context splitting specified.")
+        return dataset
+    def _load_lm_metrics(self) -> Dict[LmMetric, Any]:
+        """Extract metrics to calculate from config with their configurations."""
+        metrics = self.config.lm_metrics
+        metrics_config = {}
+        if metrics.rouge:
+            for metric in LmMetric.get_rouge_metrics():
+                metrics_config[metric] = None
+        if metrics.bleu:
+            metrics_config[LmMetric.BLEU] = None
+        if metrics.meteor:
+            metrics_config[LmMetric.METEOR] = None
+        if metrics.ter:
+            metrics_config[LmMetric.TER] = None
+        if metrics.bert_score:
+            bert_config = BertScoreConfig(
+                model_path=metrics.bert_score.model_path,
+                batch_size=metrics.bert_score.batch_size,
+                nthreads=metrics.bert_score.nthreads,
+                device=metrics.bert_score.device
+                if hasattr(metrics.bert_score, 'device')
+                else None,
+                lang=metrics.bert_score.lang
+                if hasattr(metrics.bert_score, 'lang')
+                else 'en',
+            )
+            for metric in LmMetric.get_bert_metrics():
+                metrics_config[metric] = bert_config
+        return metrics_config
+    def _generate_answers(
+        self,
+        dataset: datasets.Dataset,
+    ) -> datasets.Dataset:
+        if self.task_type == DetectionTaskType.QUERY_ANSWERABILITY:
+            return dataset
+        generation_idxs = []
+        generation_inputs = []
+        for idx, example in enumerate(dataset):
+            input_val = example[INPUT_COL]
+            if input_val and input_val[-1]['role'] != 'assistant':
+                generation_idxs.append(idx)
+                generation_inputs.append(input_val)
+        if generation_idxs:
+            if self._generator_adapter is None:
+                raise ValueError(
+                    "Dataset contains samples without assistant answers, "
+                    "but no generator_adapter was provided to the pipeline."
+                )
+            generator_model = ModelManager.load_model(self._generator_adapter)
+            generated_answers = []
+            batch_size = self.detector.config.batch_size or 1
+            sampling = self.config.sampling
+            for i in range(0, len(generation_inputs), batch_size):
+                batch_inputs = generation_inputs[i : i + batch_size]
+                batch_answers = generator_model.sample(
+                    batch_inputs,
+                    max_tokens=sampling.max_length,
+                    temperature=sampling.temperature,
+                    top_p=sampling.top_p,
+                    top_k=sampling.top_k,
+                    **sampling.kwargs,
+                )
+                generated_answers.extend(batch_answers)
+            generated_dict = dict(zip(generation_idxs, generated_answers))
+            def add_generation(example, idx):
+                if idx in generated_dict:
+                    example[INPUT_COL].append(
+                        {
+                            'role': 'assistant',
+                            'content': generated_dict[idx],
+                        }
+                    )
+                return example
+            dataset = dataset.map(
+                add_generation,
+                with_indices=True,
+                load_from_cache_file=False,
+                desc="Updating with generated answers",
+                batched=False,
+                num_proc=self.detector.num_cpus
+                if self.detector.config.use_multiprocessing
+                else None,
+            )
+        else:
+            lg.info("There are no missing LLM answers. No need to generate answers.")
+        return dataset
+    def _compute_lm_metrics(
+        self,
+        dataset: datasets.Dataset,
+    ) -> datasets.Dataset:
+        if self.config.lm_metrics is not None:
+            metrics_to_calculate = self._load_lm_metrics()
+            existing_columns = set(dataset.column_names)
+            missing_metrics = {}
+            for metric, config in metrics_to_calculate.items():
+                if metric.value not in existing_columns:
+                    missing_metrics[metric] = config
+            if not missing_metrics:
+                lg.info("All requested metrics already present in dataset.")
+                return dataset
+            lg.info(
+                f"Calculating missing metrics: {[m.value for m in missing_metrics.keys()]}"
+            )
+            # Prefer cleaned inputs if available.
+            source_col = (
+                INPUT_PROC_COL if INPUT_PROC_COL in dataset.column_names else INPUT_COL
+            )
+            generations = []
+            for item in dataset[source_col]:
+                if isinstance(item, list):
+                    if not item:
+                        generations.append('')
+                        continue
+                    last_item = item[-1]
+                    if isinstance(last_item, dict):
+                        generations.append(last_item.get('content', ''))
+                    else:
+                        generations.append(str(last_item))
+                elif isinstance(item, dict):
+                    generations.append(item.get('content', ''))
+                else:
+                    generations.append(item)
+            scores_dict = calculate_lm_metrics(
+                missing_metrics,
+                generations,
+                dataset[REFERENCE_COL],
+            )
+            for metric_name, metric_values in scores_dict.items():
+                dataset = dataset.add_column(metric_name, metric_values)
+        else:
+            lg.info("Metrics calculation is disabled.")
+        return dataset
+    def _clean_generation_outputs(
+        self,
+        dataset: datasets.Dataset,
+    ) -> datasets.Dataset:
+        if self.config.cleaner_configs:
+            def clean_example(example: List[Dict[str, str]], cleaner):
+                example[-1]['content'] = cleaner.clean(example[-1]['content'])
+                return example
+            def _create_cleaner(config: CleanerConfig) -> BaseCleaner:
+                """Factory method to create the appropriate cleaner based on config"""
+                cleaner_map = {
+                    'regex': RegexCleaner,
+                    'html': HTMLCleaner,
+                    'code': CodeCleaner,
+                    'custom': CustomCleaner,
+                }
+                cleaner_class = cleaner_map.get(config.type)
+                if not cleaner_class:
+                    raise ValueError(f"Unknown cleaner type: {config.type}")
+                return cleaner_class(config.params or {})
+            for cleaner_config in self.config.cleaner_configs:
+                cleaner = _create_cleaner(config=cleaner_config)
+                dataset = dataset.map(
+                    lambda x: {INPUT_PROC_COL: clean_example(x[INPUT_COL], cleaner)},
+                    batched=False,
+                    num_proc=self.detector.num_cpus
+                    if self.detector.config.use_multiprocessing
+                    else None,
+                    desc=f"Using cleaner: {cleaner_config.name}",
+                )
+        else:
+            lg.info("No cleaning specified.")
+            dataset = dataset.add_column(INPUT_PROC_COL, dataset[INPUT_COL])
+        return dataset
+    def __del__(self):
+        if self.experiment_logger:
+            try:
+                self.experiment_logger.finish()
+            except Exception as e:
+                lg.debug(f"Experiment logger finish failed during pipeline cleanup: {e}")

sirin/detection/judging/__init__.py ADDED Viewed

	@@ -0,0 +1,40 @@

+def _ensure_transformers_hybrid_cache(transformers_module=None):
+    previous_module = None
+    if transformers_module is None:
+        import sys
+        previous_module = sys.modules.get('transformers')
+        from transformers import AutoProcessor  # noqa: F401
+        transformers_module = sys.modules['transformers']
+    if hasattr(transformers_module, 'HybridCache'):
+        return
+    dynamic_cache = getattr(transformers_module, 'DynamicCache', None)
+    if dynamic_cache is None:
+        try:
+            from transformers import DynamicCache as dynamic_cache
+        except ImportError:
+            return
+    transformers_module.HybridCache = dynamic_cache
+    if previous_module is not None and previous_module is not transformers_module:
+        previous_module.HybridCache = dynamic_cache
+_ensure_transformers_hybrid_cache()
+from .judges import (
+    SequenceEncoderJudge,
+    SequenceDecoderJudge,
+    SequenceOpenAIJudge,
+    TokenDecoderJudge,
+    TokenEncoderJudge,
+    TokenOpenAIJudge,
+    ClaimDecoderJudge,
+    ClaimEncoderJudge,
+    ClaimOpenAIJudge,
+    HfJudgeBase,
+    JudgeAnnotationError,
+    OpenAIJudgeBase,
+)
+from .pipeline import JudgePipeline

sirin/detection/judging/judges/__init__.py ADDED Viewed

	@@ -0,0 +1,4 @@

+from .token import TokenDecoderJudge, TokenEncoderJudge, TokenOpenAIJudge
+from .sequence import SequenceDecoderJudge, SequenceEncoderJudge, SequenceOpenAIJudge
+from .claim import ClaimDecoderJudge, ClaimEncoderJudge, ClaimOpenAIJudge
+from .base import HfJudgeBase, JudgeAnnotationError, OpenAIJudgeBase

sirin/detection/judging/judges/base.py ADDED Viewed

	@@ -0,0 +1,292 @@

+from abc import abstractmethod
+from typing import Any, Dict, List, Optional
+from pathlib import Path
+import joblib
+from datasets import Dataset
+from loguru import logger as lg
+from peft import get_peft_model
+from transformers import (
+    DataCollatorForTokenClassification,
+    DataCollatorWithPadding,
+    TrainingArguments,
+)
+from sirin.utils.config_manager import validate_hydra_config
+from sirin.models.detection import DetectionResult, JudgeBaseConfig
+from sirin.utils.config_serialization import (
+    deserialize_judge_config,
+    serialize_judge_config,
+)
+from sirin.inference.adapters import ModelAdapterBase, HfModelAdapter, OpenAIModelAdapter
+from sirin.inference.model_manager import ModelManager
+from sirin.detection.base import DetectorBase
+from sirin.definitions import (
+    INPUT_COL,
+    TARGET_COL,
+    ClassificationMetric,
+    BASIC_METRICS,
+    DetectionLevel,
+    DataCollatorType,
+)
+class JudgeAnnotationError(Exception):
+    """Raised when a judge cannot produce a usable annotation for a sample.
+    E.g. every sampled generation failed the reference-echo check, so there is no honest
+    per-character consensus to emit. Callers must surface this rather than let an empty
+    annotation render as "all clear".
+    """
+class JudgeBase(DetectorBase):
+    """Base class for all judge implementations with shared configuration logic."""
+    detection_level: DetectionLevel
+    @validate_hydra_config
+    def __init__(self, config: JudgeBaseConfig, model_adapter: ModelAdapterBase):
+        super().__init__(config)
+        self.model_adapter = ModelManager.load_model(model_adapter)
+        if self.config.model_load_path:
+            self.load(self.config.model_load_path)
+        self.setup_model()
+    @abstractmethod
+    def setup_model(self):
+        pass
+    @abstractmethod
+    def _compute_metrics(self, eval_pred: Any) -> Dict[str, Any]:
+        pass
+    @abstractmethod
+    def _preprocess(self, samples: Dataset) -> Dict[str, Any]:
+        pass
+    def _save_config(self, save_dir: Path):
+        """
+        Save judge configuration to directory.
+        Args:
+            save_dir: Directory where to save the config
+        """
+        model_path = str(save_dir / 'model') if hasattr(self, '_save_model') else None
+        config_dict = serialize_judge_config(
+            config=self.config,
+            threshold=self.threshold,
+            model_adapter=self.model_adapter.config,
+            model_path=model_path,
+        )
+        joblib.dump(config_dict, save_dir / 'config.joblib')
+        lg.info(f'Saved config to {save_dir / "config.joblib"}')
+    def _load_config(self, load_dir: Path):
+        config_path = load_dir / 'config.joblib'
+        if not config_path.exists():
+            raise FileNotFoundError(f'Config file not found at: {config_path}')
+        loaded_objects = joblib.load(config_path)
+        config, threshold, adapter_config, model_path = deserialize_judge_config(loaded_objects)
+        # Update config attributes
+        for key, value in config.__dict__.items():
+            if key not in ['model_load_path', 'model_save_path']:
+                setattr(self.config, key, value)
+        self.threshold = threshold
+        lg.info(f'Loaded config from {config_path}')
+        # Store adapter config for subclass use
+        self._loaded_adapter_config = adapter_config
+        self._loaded_model_path = model_path
+class HfJudgeBase(JudgeBase):
+    """Base class for HuggingFace model-based judges with training support."""
+    # Override in subclasses to specify data collator type
+    data_collator_type: DataCollatorType = DataCollatorType.PADDING  # Default for sequence encoder judges
+    def __init__(self, config: JudgeBaseConfig, model_adapter: ModelAdapterBase):
+        super().__init__(config, model_adapter)
+    def setup_model(self):
+        if self.config.peft_config is not None:
+            self.model_adapter.model = get_peft_model(self.model_adapter.model, self.config.peft_config)
+            lg.info(self.model_adapter.model.print_trainable_parameters())
+        else:
+            lg.info("Training without peft")
+        self.class_token_ids = [
+            self.model_adapter.tokenizer.convert_tokens_to_ids(str(i))
+            for i in range(self.config.num_classification_heads)
+        ]
+    def _save_model(self, save_dir: Path):
+        model_path = save_dir / 'model'
+        model_path.mkdir(exist_ok=True)
+        # Save model
+        if hasattr(self.model_adapter, 'model'):
+            self.model_adapter.model.save_pretrained(str(model_path))
+            lg.info(f'Saved model to {model_path}')
+        # Save tokenizer
+        if hasattr(self.model_adapter, 'tokenizer'):
+            self.model_adapter.tokenizer.save_pretrained(str(model_path))
+            lg.info(f'Saved tokenizer to {model_path}')
+    def _load_model(self, load_dir: Path):
+        model_path = load_dir / 'model'
+        if not model_path.exists():
+            raise FileNotFoundError(f'Model directory not found at: {model_path}')
+        # Load model using adapter config from _load_config
+        adapter_config = getattr(self, '_loaded_adapter_config', None)
+        if adapter_config is None:
+            raise ValueError('Config must be loaded before model. Call _load_config() first.')
+        # Update model path to point to saved model
+        adapter_config.model_path = str(model_path)
+        model_manager = ModelManager()
+        self.model_adapter = model_manager.load(
+            HfModelAdapter(config=adapter_config, model=str(model_path))
+        )
+        lg.info(f'Loaded model from {model_path}')
+    def _save_additional_components(self, save_dir: Path):
+        self._save_model(save_dir)
+    def _load_additional_components(self, load_dir: Path):
+        self._load_model(load_dir)
+    def _check_truncation_warning(self, samples: List[Dict[str, Any]]):
+        if not self.model_adapter.tokenizer:
+            return
+        preprocessed_samples = self.model_adapter._preprocess_input(samples)
+        if self.model_adapter.config.truncation:
+            for i, sample in enumerate(preprocessed_samples):
+                tokens = self.model_adapter.tokenizer.encode(
+                    sample,
+                    add_special_tokens=True,
+                    truncation=False,
+                )
+                max_length = self.model_adapter.tokenizer.model_max_length
+                if len(tokens) > max_length:
+                    lg.warning(
+                        f"Sample {i} exceeds model's max length ({len(tokens)} > {max_length}). "
+                        f"Truncation will be applied. Consider reducing input length."
+                    )
+    def train(
+        self,
+        training_args: TrainingArguments,
+        train_data: Dataset,
+        val_data: Optional[Dataset],
+        metrics: Optional[List[ClassificationMetric]] = None,
+        **kwargs,
+    ) -> DetectionResult:
+        self.metrics_to_compute = metrics or BASIC_METRICS
+        train_data_tokenized = train_data.batch_process(
+            lambda x: self._preprocess(x),
+            batch_size=self.config.batch_size,
+            remove_columns=[INPUT_COL, TARGET_COL],
+        )
+        train_data_tokenized = Dataset.from_dict(train_data_tokenized)
+        val_data_tokenized = None
+        if val_data is not None:
+            val_data_tokenized = val_data.batch_process(
+                lambda x: self._preprocess(x),
+                batch_size=self.config.batch_size,
+                remove_columns=[INPUT_COL, TARGET_COL],
+            )
+            val_data_tokenized = Dataset.from_dict(val_data_tokenized)
+        # Select data collator based on judge type (specified by subclass)
+        if self.data_collator_type == DataCollatorType.TOKEN:
+            data_collator = DataCollatorForTokenClassification(
+                tokenizer=self.model_adapter.tokenizer,
+                padding=self.model_adapter.config.padding,
+                label_pad_token_id=-100,
+            )
+        elif self.data_collator_type == DataCollatorType.PADDING:
+            data_collator = DataCollatorWithPadding(
+                tokenizer=self.model_adapter.tokenizer,
+                padding=self.model_adapter.config.padding,
+            )
+        else:
+            raise ValueError(f'Invalid data_collator_type: {self.data_collator_type}')
+        trainer = self.trainer(
+            model=self.model_adapter.model,
+            data_collator=data_collator,
+            args=training_args,
+            train_dataset=train_data_tokenized,
+            eval_dataset=val_data_tokenized,
+            processing_class=self.model_adapter.tokenizer,
+            compute_metrics=self._compute_metrics if val_data_tokenized else None,
+            **kwargs,
+        )
+        train_result = trainer.train()
+        self.model_adapter.model = trainer.model
+        if training_args.output_dir:
+            trainer.save_model()
+            self.model_adapter.tokenizer.save_pretrained(training_args.output_dir)
+        return DetectionResult(
+            metrics=train_result.metrics,
+            probs=None,
+            threshold=self.threshold,
+        )
+class OpenAIJudgeBase(JudgeBase):
+    """Base class for OpenAI API-based judges (no training support)."""
+    def __init__(self, config: JudgeBaseConfig, model_adapter: ModelAdapterBase):
+        super().__init__(config, model_adapter)
+    def setup_model(self):
+        pass
+    def _load_config(self, load_dir: Path):
+        """Override to reload OpenAI model adapter after config load."""
+        super()._load_config(load_dir)
+        # Reload model adapter with saved config
+        model_manager = ModelManager()
+        self.model_adapter = model_manager.load(
+            OpenAIModelAdapter(
+                config=self._loaded_adapter_config,
+                model=self._loaded_adapter_config.model_path or self._loaded_model_path
+            )
+        )
+        lg.info('Reloaded OpenAI model adapter')
+    def train(
+        self,
+        training_args: TrainingArguments,
+        train_data: Dataset,
+        val_data: Optional[Dataset],
+        metrics: Optional[List[ClassificationMetric]] = None,
+        **kwargs,
+    ) -> DetectionResult:
+        lg.info("There is no training for API judges")
+        return DetectionResult()
+    def _preprocess(self, samples: Dataset) -> Dataset:
+        return samples
+    def _compute_metrics(self, eval_pred: Any) -> Dict[str, Any]:
+        lg.info("There is no metrics compution during training for API judges")
+        return {}

sirin/detection/judging/judges/claim/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@

+from .decoder import ClaimDecoderJudge
+from .encoder import ClaimEncoderJudge
+from .openai import ClaimOpenAIJudge

sirin/detection/judging/judges/claim/decoder.py ADDED Viewed

	@@ -0,0 +1,201 @@

+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+import joblib
+import numpy as np
+import torch
+from loguru import logger as lg
+from transformers import EvalPrediction
+from sirin.definitions import INPUT_COL, TARGET_COL, DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import HfJudgeBase
+from sirin.detection.judging.judges.utils import (
+    format_dialogue_for_training,
+    prepare_decoder_inputs_with_labels,
+    build_prompt_messages,
+)
+from sirin.detection.judging.training import SequenceDecoderJudgeTrainer
+from sirin.detection.splitters import SplitManager
+from sirin.detection.utils.basic import calibrate_threshold
+from sirin.inference.adapters import HfModelAdapter, ModelAdapterBase
+from sirin.metrics.classification import calculate_classification_metrics
+from sirin.models.detection import HfJudgeConfig, SplitConfig
+from sirin.utils.hf import get_assistant_prefix
+class ClaimDecoderJudge(HfJudgeBase):
+    """
+    Decoder judge implementation for claim-level hallucination detection.
+    Decomposes each response into atomic claims via SplitManager, classifies
+    each claim independently using next-token prediction from a generative
+    decoder, then re-aggregates per-claim predictions back to response level.
+    """
+    detection_level = DetectionLevel.CLAIM
+    data_collator_type = DataCollatorType.TOKEN
+    trainer = SequenceDecoderJudgeTrainer
+    def __init__(
+        self,
+        config: HfJudgeConfig,
+        model_adapter: HfModelAdapter,
+        response_splitter_config: SplitConfig,
+        split_model: Optional[ModelAdapterBase] = None,
+    ):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.response_splitter_config = response_splitter_config
+        self.response_splitter = SplitManager(config=response_splitter_config)
+        self.split_model = split_model or model_adapter
+        self.assistant_prefix = get_assistant_prefix(self.model_adapter._model_name)
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Detects hallucinations at claim level using a generative decoder model.
+        Each response is split into atomic claims, then the decoder predicts
+        a class token (0, 1, etc.) for each claim independently. Per-claim
+        predictions are re-aggregated to response level.
+        """
+        original_samples = list(samples)
+        samples_fact_splits = self.response_splitter.split_inputs(
+            samples, self.split_model
+        )
+        response_group_ids = []
+        for index, splits in enumerate(samples_fact_splits):
+            response_group_ids.extend([index] * len(splits))
+        samples = [split for splits in samples_fact_splits for split in splits]
+        samples, group_ids = self._split_context_samples(samples)
+        self._check_truncation_warning(samples)
+        formatted_input = build_prompt_messages(self.config, samples)
+        preprocessed_input = self.model_adapter._preprocess_input(formatted_input)
+        preprocessed_input = [
+            f"{sample}{self.assistant_prefix}" for sample in preprocessed_input
+        ]
+        inputs = self.model_adapter.tokenizer(
+            preprocessed_input,
+            return_tensors=self.config.return_tensors,
+            truncation=self.model_adapter.config.truncation,
+            padding=self.model_adapter.config.padding,
+        ).to(self.config.device)
+        with torch.no_grad():
+            outputs = self.model_adapter.model.generate(
+                **inputs,
+                max_new_tokens=self.config.max_new_tokens,
+                do_sample=False,
+                temperature=self.config.temperature,
+                top_p=self.config.top_p,
+                pad_token_id=self.model_adapter.tokenizer.eos_token_id,
+                eos_token_id=self.model_adapter.tokenizer.eos_token_id,
+                return_dict_in_generate=True,
+                output_scores=True,
+            )
+        first_token_logits = outputs.scores[0]
+        selected_logits = first_token_logits[:, self.class_token_ids]
+        probs, preds = self._logits_to_probs_preds(selected_logits)
+        probs = probs.tolist()
+        preds = preds.tolist()
+        preds, probs = self._aggregate_context_predictions(
+            group_ids, preds, probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        overall_preds, overall_probs = self.response_splitter.aggregate_predictions(
+            response_group_ids,
+            preds,
+            probs,
+            binary=(self.config.num_classification_heads <= 2),
+        )
+        results = []
+        for sample_idx, sample in enumerate(original_samples):
+            sample_mask = [
+                i for i, gid in enumerate(response_group_ids) if gid == sample_idx
+            ]
+            splitted_samples = samples_fact_splits[sample_idx]
+            sample_probs = [probs[i] for i in sample_mask]
+            sample_preds = [preds[i] for i in sample_mask]
+            facts = [
+                {"fact": split[1]["content"], "pred": pred, "prob": prob}
+                for split, pred, prob in zip(splitted_samples, sample_probs, sample_preds)
+            ]
+            results.append(
+                {
+                    "sample": sample,
+                    "overall_pred": overall_preds[sample_idx],
+                    "overall_prob": overall_probs[sample_idx],
+                    "facts": facts,
+                }
+            )
+        self.claim_results = results
+        return overall_probs, overall_preds, labels
+    def _compute_metrics(self, eval_pred: EvalPrediction) -> Dict[str, Any]:
+        """Compute metrics for claim-level classification."""
+        probs = eval_pred.predictions
+        labels = eval_pred.label_ids
+        if self.config.num_classification_heads == 2:
+            self.threshold = calibrate_threshold(
+                probs,
+                labels,
+                self.config.threshold_method,
+                self.config.threshold_percentile,
+                self.config.fixed_threshold,
+            )
+            preds = (probs > self.threshold).astype(int)
+        else:
+            self.threshold = None
+            preds = np.argmax(probs, axis=1)
+        metrics = calculate_classification_metrics(
+            labels,
+            probs,
+            preds,
+            metrics=self.metrics_to_compute,
+        )
+        lg.info(f"Evaluation metrics: {metrics}")
+        return metrics
+    def _preprocess(self, samples: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess samples for claim-level classification with decoder models."""
+        inputs = samples[INPUT_COL]
+        targets = samples[TARGET_COL]
+        self._check_truncation_warning(inputs)
+        formatted_messages = format_dialogue_for_training(
+            inputs, targets, self.config, is_token_level=False
+        )
+        return prepare_decoder_inputs_with_labels(
+            formatted_messages, self.model_adapter, self.config
+        )
+    def _save_config(self, save_dir: Path):
+        super()._save_config(save_dir)
+        joblib.dump(self.response_splitter_config, save_dir / "splitter_config.joblib")
+        lg.info(f"Saved splitter config to {save_dir / 'splitter_config.joblib'}")
+    def _load_config(self, load_dir: Path):
+        super()._load_config(load_dir)
+        splitter_config_path = load_dir / "splitter_config.joblib"
+        if splitter_config_path.exists():
+            self.response_splitter_config = joblib.load(splitter_config_path)
+            self.response_splitter = SplitManager(config=self.response_splitter_config)
+            lg.info(f"Loaded splitter config from {splitter_config_path}")

sirin/detection/judging/judges/claim/encoder.py ADDED Viewed

	@@ -0,0 +1,169 @@

+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+import joblib
+import numpy as np
+from loguru import logger as lg
+from transformers import EvalPrediction, Trainer
+from sirin.definitions import INPUT_COL, TARGET_COL, DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import HfJudgeBase
+from sirin.detection.judging.judges.utils import (
+    process_logits_to_probs,
+    calibrate_and_compute_metrics,
+)
+from sirin.detection.splitters import SplitManager
+from sirin.inference.adapters import HfModelAdapter, ModelAdapterBase
+from sirin.models.detection import HfJudgeConfig, SplitConfig
+class ClaimEncoderJudge(HfJudgeBase):
+    """
+    Encoder judge implementation for claim-level hallucination detection.
+    Decomposes each response into atomic claims via SplitManager, classifies
+    each claim independently using an encoder classification head, then
+    re-aggregates per-claim predictions back to response level.
+    """
+    detection_level = DetectionLevel.CLAIM
+    data_collator_type = DataCollatorType.PADDING
+    trainer = Trainer
+    def __init__(
+        self,
+        config: HfJudgeConfig,
+        model_adapter: HfModelAdapter,
+        response_splitter_config: SplitConfig,
+        split_model: Optional[ModelAdapterBase] = None,
+    ):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.response_splitter_config = response_splitter_config
+        self.response_splitter = SplitManager(config=response_splitter_config)
+        self.split_model = split_model or model_adapter
+        self.device = self.config.device
+        if self.model_adapter.tokenizer.pad_token is None:
+            self.model_adapter.tokenizer.pad_token = (
+                self.model_adapter.tokenizer.eos_token
+            )
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Detects hallucinations at claim level using encoder embeddings.
+        Each response is split into atomic claims, classified independently,
+        then predictions are re-aggregated to response level.
+        """
+        original_samples = list(samples)
+        samples_fact_splits = self.response_splitter.split_inputs(
+            samples, self.split_model
+        )
+        response_group_ids = []
+        for index, splits in enumerate(samples_fact_splits):
+            response_group_ids.extend([index] * len(splits))
+        samples = [split for splits in samples_fact_splits for split in splits]
+        samples, group_ids = self._split_context_samples(samples)
+        self._check_truncation_warning(samples)
+        model_states = self.model_adapter.generate_hiddens(
+            inputs=samples,
+            return_hiddens=False,
+            return_logits=True,
+        )
+        logits = model_states.logits
+        probs, preds = self._logits_to_probs_preds(logits)
+        probs = probs.tolist()
+        preds = preds.tolist()
+        preds, probs = self._aggregate_context_predictions(
+            group_ids, preds, probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        overall_preds, overall_probs = self.response_splitter.aggregate_predictions(
+            response_group_ids,
+            preds,
+            probs,
+            binary=(self.config.num_classification_heads <= 2),
+        )
+        results = []
+        for sample_idx, sample in enumerate(original_samples):
+            sample_mask = [
+                i for i, gid in enumerate(response_group_ids) if gid == sample_idx
+            ]
+            splitted_samples = samples_fact_splits[sample_idx]
+            sample_probs = [probs[i] for i in sample_mask]
+            sample_preds = [preds[i] for i in sample_mask]
+            facts = [
+                {"fact": split[1]["content"], "pred": pred, "prob": prob}
+                for split, pred, prob in zip(splitted_samples, sample_probs, sample_preds)
+            ]
+            results.append(
+                {
+                    "sample": sample,
+                    "overall_pred": overall_preds[sample_idx],
+                    "overall_prob": overall_probs[sample_idx],
+                    "facts": facts,
+                }
+            )
+        self.claim_results = results
+        return overall_probs, overall_preds, labels
+    def _compute_metrics(self, eval_pred: EvalPrediction) -> Dict[str, Any]:
+        """Compute metrics for claim-level classification."""
+        logits, labels = eval_pred
+        probs, _ = process_logits_to_probs(logits, labels, filter_padding=False)
+        is_binary = logits.shape[-1] <= 2
+        metrics, threshold = calibrate_and_compute_metrics(
+            probs, labels, self.config, self.metrics_to_compute, is_binary
+        )
+        self.threshold = threshold
+        return metrics
+    def _preprocess(self, samples: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess samples for claim-level classification."""
+        inputs = samples[INPUT_COL]
+        targets = samples[TARGET_COL]
+        self._check_truncation_warning(inputs)
+        processed_inputs = self.model_adapter._preprocess_input(inputs)
+        tokenized = self.model_adapter.tokenizer(
+            processed_inputs,
+            padding=False,
+            truncation=self.model_adapter.config.truncation,
+            add_special_tokens=False,
+        )
+        tokenized["labels"] = targets
+        return tokenized
+    def _save_config(self, save_dir: Path):
+        super()._save_config(save_dir)
+        joblib.dump(self.response_splitter_config, save_dir / "splitter_config.joblib")
+        lg.info(f"Saved splitter config to {save_dir / 'splitter_config.joblib'}")
+    def _load_config(self, load_dir: Path):
+        super()._load_config(load_dir)
+        splitter_config_path = load_dir / "splitter_config.joblib"
+        if splitter_config_path.exists():
+            self.response_splitter_config = joblib.load(splitter_config_path)
+            self.response_splitter = SplitManager(config=self.response_splitter_config)
+            lg.info(f"Loaded splitter config from {splitter_config_path}")

sirin/detection/judging/judges/claim/openai.py ADDED Viewed

	@@ -0,0 +1,134 @@

+from pathlib import Path
+from typing import Dict, List, Optional, Tuple, Union
+import joblib
+import numpy as np
+from loguru import logger as lg
+from sirin.definitions import DetectionLevel
+from sirin.detection.judging.judges.base import OpenAIJudgeBase
+from sirin.detection.judging.judges.utils import build_prompt_messages
+from sirin.detection.splitters import SplitManager
+from sirin.inference.adapters import ModelAdapterBase
+from sirin.models.detection import OpenAIJudgeConfig, SplitConfig
+class ClaimOpenAIJudge(OpenAIJudgeBase):
+    """
+    OpenAI API judge implementation for claim-level hallucination detection.
+    Decomposes each response into atomic claims via SplitManager, classifies
+    each claim independently via the OpenAI API, then re-aggregates per-claim
+    predictions back to response level. No training support.
+    """
+    detection_level = DetectionLevel.CLAIM
+    def __init__(
+        self,
+        config: OpenAIJudgeConfig,
+        model_adapter: ModelAdapterBase,
+        response_splitter_config: SplitConfig,
+        split_model: Optional[ModelAdapterBase] = None,
+    ):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.response_splitter_config = response_splitter_config
+        self.response_splitter = SplitManager(config=response_splitter_config)
+        self.split_model = split_model or model_adapter
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Detects hallucinations at claim level using the OpenAI API.
+        Each response is split into atomic claims, classified independently
+        via API logprobs, then predictions are re-aggregated to response level.
+        """
+        original_samples = list(samples)
+        samples_fact_splits = self.response_splitter.split_inputs(
+            samples, self.split_model
+        )
+        response_group_ids = []
+        for index, splits in enumerate(samples_fact_splits):
+            response_group_ids.extend([index] * len(splits))
+        samples = [split for splits in samples_fact_splits for split in splits]
+        samples, group_ids = self._split_context_samples(samples)
+        formatted_input = build_prompt_messages(self.config, samples)
+        results, logprobs_results = self.model_adapter.sample(
+            inputs=formatted_input,
+            max_tokens=1,
+            temperature=self.config.temperature,
+            top_p=self.config.top_p,
+            return_logprobs=True,
+            top_logprobs=2,
+            **kwargs,
+        )
+        probs = []
+        preds = []
+        for idx, logprob_result in enumerate(logprobs_results):
+            # models that omit logprobs get a nan score (not a crash); the verdict still holds.
+            if logprob_result and logprob_result[0]:
+                probs.append(-1 * logprob_result[0][0])
+            else:
+                probs.append(float('nan'))
+            text = str(results[idx]).strip()
+            preds.append(int(text) if text.isdigit() else 0)
+        preds, probs = self._aggregate_context_predictions(
+            group_ids, preds, probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        overall_preds, overall_probs = self.response_splitter.aggregate_predictions(
+            response_group_ids,
+            preds,
+            probs,
+            binary=(self.config.num_classification_heads <= 2),
+        )
+        results = []
+        for sample_idx, sample in enumerate(original_samples):
+            sample_mask = [
+                i for i, gid in enumerate(response_group_ids) if gid == sample_idx
+            ]
+            splitted_samples = samples_fact_splits[sample_idx]
+            sample_probs = [probs[i] for i in sample_mask]
+            sample_preds = [preds[i] for i in sample_mask]
+            facts = [
+                {"fact": split[1]["content"], "pred": pred, "prob": prob}
+                for split, pred, prob in zip(splitted_samples, sample_probs, sample_preds)
+            ]
+            results.append(
+                {
+                    "sample": sample,
+                    "overall_pred": overall_preds[sample_idx],
+                    "overall_prob": overall_probs[sample_idx],
+                    "facts": facts,
+                }
+            )
+        self.claim_results = results
+        return overall_probs, overall_preds, labels
+    def _save_config(self, save_dir: Path):
+        super()._save_config(save_dir)
+        joblib.dump(self.response_splitter_config, save_dir / "splitter_config.joblib")
+        lg.info(f"Saved splitter config to {save_dir / 'splitter_config.joblib'}")
+    def _load_config(self, load_dir: Path):
+        super()._load_config(load_dir)
+        splitter_config_path = load_dir / "splitter_config.joblib"
+        if splitter_config_path.exists():
+            self.response_splitter_config = joblib.load(splitter_config_path)
+            self.response_splitter = SplitManager(config=self.response_splitter_config)
+            lg.info(f"Loaded splitter config from {splitter_config_path}")

sirin/detection/judging/judges/sequence/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@

+from .decoder import SequenceDecoderJudge
+from .encoder import SequenceEncoderJudge
+from .openai import SequenceOpenAIJudge

sirin/detection/judging/judges/sequence/decoder.py ADDED Viewed

	@@ -0,0 +1,139 @@

+from typing import Any, Dict, List, Optional, Tuple, Union
+import numpy as np
+import torch
+from loguru import logger as lg
+from transformers import EvalPrediction
+from sirin.definitions import INPUT_COL, TARGET_COL, DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import HfJudgeBase
+from sirin.detection.judging.judges.utils import (
+    format_dialogue_for_training,
+    prepare_decoder_inputs_with_labels,
+    build_prompt_messages,
+)
+from sirin.detection.judging.training import SequenceDecoderJudgeTrainer
+from sirin.detection.utils.basic import calibrate_threshold
+from sirin.inference.adapters import HfModelAdapter
+from sirin.metrics.classification import calculate_classification_metrics
+from sirin.models.detection import HfJudgeConfig
+from sirin.utils.hf import get_assistant_prefix
+class SequenceDecoderJudge(HfJudgeBase):
+    """
+    Decoder judge implementation for sequence-level hallucination detection.
+    Uses generative decoder models (e.g., Llama, GPT) with next-token prediction.
+    Predicts a class token (0, 1, etc.) for the entire sequence.
+    """
+    detection_level = DetectionLevel.SEQUENCE
+    data_collator_type = DataCollatorType.TOKEN
+    trainer = SequenceDecoderJudgeTrainer
+    def __init__(self, config: HfJudgeConfig, model_adapter: HfModelAdapter):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.assistant_prefix = get_assistant_prefix(self.model_adapter._model_name)
+        self.last_generations: list[str] | None = None  # generated text is not locally decoded here.
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Architecture: Uses generative decoder model (e.g., Llama, GPT).
+        Differs from encoder judges by using next-token prediction instead of classification head.
+        Predicts a class token (0, 1, etc.) as the first generated token.
+        """
+        samples, group_ids = self._split_context_samples(samples)
+        self._check_truncation_warning(samples)
+        formatted_input = build_prompt_messages(self.config, samples)
+        preprocessed_input = self.model_adapter._preprocess_input(formatted_input)
+        preprocessed_input = [
+            f"{sample}{self.assistant_prefix}" for sample in preprocessed_input
+        ]
+        inputs = self.model_adapter.tokenizer(
+            preprocessed_input,
+            return_tensors=self.config.return_tensors,
+            truncation=self.model_adapter.config.truncation,
+            padding=self.model_adapter.config.padding,
+        ).to(self.config.device)
+        # Decoder approach: Generate next token and extract logits for class tokens
+        with torch.no_grad():
+            outputs = self.model_adapter.model.generate(
+                **inputs,
+                max_new_tokens=self.config.max_new_tokens,
+                do_sample=False,
+                temperature=self.config.temperature,
+                top_p=self.config.top_p,
+                pad_token_id=self.model_adapter.tokenizer.eos_token_id,
+                eos_token_id=self.model_adapter.tokenizer.eos_token_id,
+                return_dict_in_generate=True,
+                output_scores=True,
+            )
+        # Extract logits for class tokens (e.g., "0", "1") from first generated position
+        first_token_logits = outputs.scores[0]
+        selected_logits = first_token_logits[:, self.class_token_ids]
+        probs, preds = self._logits_to_probs_preds(selected_logits)
+        probs = probs.tolist()
+        preds = preds.tolist()
+        preds, probs = self._aggregate_context_predictions(
+            group_ids, preds, probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        return probs, preds, labels
+    def _compute_metrics(self, eval_pred: EvalPrediction) -> Dict[str, Any]:
+        """Compute metrics for sequence-level classification."""
+        probs = eval_pred.predictions
+        labels = eval_pred.label_ids
+        if self.config.num_classification_heads == 2:
+            self.threshold = calibrate_threshold(
+                probs,
+                labels,
+                self.config.threshold_method,
+                self.config.threshold_percentile,
+                self.config.fixed_threshold,
+            )
+            preds = (probs > self.threshold).astype(int)
+        else:
+            self.threshold = None
+            preds = np.argmax(probs, axis=1)
+        metrics = calculate_classification_metrics(
+            labels,
+            probs,
+            preds,
+            metrics=self.metrics_to_compute,
+        )
+        lg.info(f"Evaluation metrics: {metrics}")
+        return metrics
+    def _preprocess(self, samples: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess samples for sequence-level classification with decoder models."""
+        inputs = samples[INPUT_COL]
+        targets = samples[TARGET_COL]
+        self._check_truncation_warning(inputs)
+        # Format dialogue samples with system/user/assistant roles
+        formatted_messages = format_dialogue_for_training(
+            inputs, targets, self.config, is_token_level=False
+        )
+        # Prepare tokenized inputs with proper label masking
+        return prepare_decoder_inputs_with_labels(
+            formatted_messages, self.model_adapter, self.config
+        )

sirin/detection/judging/judges/sequence/encoder.py ADDED Viewed

	@@ -0,0 +1,108 @@

+from typing import Any, Dict, List, Optional, Tuple, Union
+import numpy as np
+from loguru import logger as lg
+from transformers import EvalPrediction, Trainer
+from sirin.definitions import INPUT_COL, TARGET_COL, DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import HfJudgeBase
+from sirin.detection.judging.judges.utils import (
+    process_logits_to_probs,
+    calibrate_and_compute_metrics,
+)
+from sirin.models.detection import HfJudgeConfig
+from sirin.inference.adapters import HfModelAdapter
+class SequenceEncoderJudge(HfJudgeBase):
+    """
+    Encoder judge implementation for sequence-level hallucination detection.
+    Uses AutoModelForSequenceClassification with encoder embeddings + classification head.
+    """
+    detection_level = DetectionLevel.SEQUENCE
+    data_collator_type = DataCollatorType.PADDING
+    trainer = Trainer
+    def __init__(self, config: HfJudgeConfig, model_adapter: HfModelAdapter):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.device = self.config.device
+        if self.model_adapter.tokenizer.pad_token is None:
+            self.model_adapter.tokenizer.pad_token = (
+                self.model_adapter.tokenizer.eos_token
+            )
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Detects hallucinations at sequence level using encoder embeddings.
+        For binary: Returns 0 (no hallucination) or 1 (hallucination)
+        For multiclass: Returns predicted class index
+        Architecture: Uses encoder embeddings with classification head (AutoModelForSequenceClassification).
+        This differs from decoder judges which use generative models with next-token prediction.
+        """
+        samples, group_ids = self._split_context_samples(samples)
+        self._check_truncation_warning(samples)
+        # Encoder approach: Get logits directly from classification head
+        model_states = self.model_adapter.generate_hiddens(
+            inputs=samples,
+            return_hiddens=False,
+            return_logits=True,
+        )
+        logits = model_states.logits
+        probs, preds = self._logits_to_probs_preds(logits)
+        probs = probs.tolist()
+        preds = preds.tolist()
+        preds, probs = self._aggregate_context_predictions(
+            group_ids, preds, probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        return probs, preds, labels
+    def _compute_metrics(self, eval_pred: EvalPrediction) -> Dict[str, Any]:
+        """Compute metrics for sequence-level classification."""
+        logits, labels = eval_pred
+        # Convert logits to probabilities (handles binary/multiclass)
+        probs, _ = process_logits_to_probs(logits, labels, filter_padding=False)
+        # Calibrate threshold and compute metrics
+        is_binary = logits.shape[-1] <= 2
+        metrics, threshold = calibrate_and_compute_metrics(
+            probs, labels, self.config, self.metrics_to_compute, is_binary
+        )
+        self.threshold = threshold
+        return metrics
+    def _preprocess(self, samples: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess samples for sequence-level classification."""
+        inputs = samples[INPUT_COL]
+        targets = samples[TARGET_COL]
+        self._check_truncation_warning(inputs)
+        processed_inputs = self.model_adapter._preprocess_input(inputs)
+        tokenized = self.model_adapter.tokenizer(
+            processed_inputs,
+            padding=False,
+            truncation=self.model_adapter.config.truncation,
+            add_special_tokens=False,
+        )
+        tokenized["labels"] = targets
+        return tokenized

sirin/detection/judging/judges/sequence/openai.py ADDED Viewed

	@@ -0,0 +1,99 @@

+import numpy as np
+from typing import Dict, List, Tuple, Union, Optional
+from sirin.models.detection import OpenAIJudgeConfig
+from sirin.detection.judging.judges.base import JudgeAnnotationError, OpenAIJudgeBase
+from sirin.detection.judging.judges.utils import build_prompt_messages
+from sirin.definitions import DetectionLevel, DataCollatorType
+from sirin.inference.adapters import ModelAdapterBase
+class SequenceOpenAIJudge(OpenAIJudgeBase):
+    """
+    OpenAI API judge implementation for sequence-level hallucination detection.
+    Architecture: Uses OpenAI API (e.g., GPT-4) instead of local HuggingFace models.
+    No training support - inference only via API calls.
+    """
+    detection_level = DetectionLevel.SEQUENCE
+    def __init__(self, config: OpenAIJudgeConfig, model_adapter: ModelAdapterBase):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.last_generations: list[str] | None = None
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Uses OpenAI API for sequence-level classification.
+        Returns logprobs from API instead of running local model inference.
+        """
+        samples, group_ids = self._split_context_samples(samples)
+        formatted_input = build_prompt_messages(self.config, samples)
+        def _sample(return_logprobs: bool):
+            return self.model_adapter.sample(
+                inputs=formatted_input,
+                max_tokens=self.config.verdict_max_tokens,
+                temperature=self.config.temperature,
+                top_p=self.config.top_p,
+                return_logprobs=return_logprobs,
+                top_logprobs=2,
+                **kwargs
+            )
+        try:
+            results, logprobs_results = _sample(return_logprobs=True)
+        except Exception as exc:
+            # Some providers reject logprobs outright (400): retry once without them and
+            # keep the honest no-score verdict (nan). Anything else propagates unchanged.
+            try:
+                import openai
+                logprobs_rejected = (
+                    isinstance(exc, openai.BadRequestError)
+                    and 'logprob' in str(exc).lower()
+                )
+            except ImportError:
+                logprobs_rejected = False
+            if not logprobs_rejected:
+                raise
+            results = _sample(return_logprobs=False)
+            logprobs_results = [None] * len(results)
+        self.last_generations = list(results)
+        probs = []
+        preds = []
+        num_classes = max(self.config.num_classification_heads, 2)
+        for idx, logprob_result in enumerate(logprobs_results):
+            text = str(results[idx]).strip()
+            # Reasoning models return prose in content; take the first valid class digit.
+            digit = next(
+                (int(c) for c in text if c.isdigit() and int(c) < num_classes), None
+            )
+            if digit is None:
+                # Never fabricate a verdict when no valid class digit appears anywhere.
+                raise JudgeAnnotationError(
+                    f'The judge model did not answer with a bare class digit and no class '
+                    f'digit appears anywhere in its answer (got {text!r}). Choose a judge '
+                    'model that states the verdict digit, or raise verdict_max_tokens '
+                    f'(currently {self.config.verdict_max_tokens}) so a reasoning model '
+                    'can finish thinking before the digit.'
+                )
+            preds.append(digit)
+            # The adapter's logprobs carry floats only (no token text), so the first-token
+            # logprob is provably the digit's only when the whole answer IS the digit.
+            # Some models also omit logprobs entirely. Either way: nan, verdict still holds.
+            if logprob_result and logprob_result[0] and text == str(digit):
+                probs.append(-1 * logprob_result[0][0])
+            else:
+                probs.append(float('nan'))
+        preds, probs = self._aggregate_context_predictions(
+            group_ids, preds, probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        return probs, preds, labels

sirin/detection/judging/judges/token/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@

+from .decoder import TokenDecoderJudge
+from .encoder import TokenEncoderJudge
+from .openai import TokenOpenAIJudge

sirin/detection/judging/judges/token/decoder.py ADDED Viewed

	@@ -0,0 +1,186 @@

+from typing import Any, Dict, List, Optional, Tuple, Union
+import numpy as np
+import torch
+from transformers import Trainer
+from sirin.definitions import INPUT_COL, TARGET_COL, DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import HfJudgeBase
+from sirin.detection.judging.judges.utils import (
+    format_dialogue_for_training,
+    prepare_decoder_inputs_with_labels,
+    build_prompt_messages,
+    calibrate_and_compute_metrics,
+    calculate_character_probabilities,
+    create_char_binary_vector,
+    find_span_segments,
+)
+from sirin.detection.utils.token import convert_spans_to_labels
+from sirin.inference.adapters import ModelAdapterBase
+from sirin.models.detection import HfJudgeConfig
+from sirin.utils.hf import get_assistant_prefix
+class TokenDecoderJudge(HfJudgeBase):
+    """
+    Decoder judge implementation for token-level hallucination detection.
+    Uses generative decoder models with multiple beam samples to identify
+    hallucinated character spans based on generation consistency.
+    """
+    detection_level = DetectionLevel.TOKEN
+    data_collator_type = DataCollatorType.TOKEN
+    trainer = Trainer
+    def __init__(self, config: HfJudgeConfig, model_adapter: ModelAdapterBase):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.model_adapter.tokenizer.add_tokens(['[SPAN]', '[/SPAN]'], special_tokens=True)
+        self.model_adapter.model.resize_token_embeddings(len(self.model_adapter.tokenizer))
+        self.assistant_prefix = get_assistant_prefix(self.model_adapter._model_name)
+        self.class_token_ids = None
+        self.last_generations: list[str] | None = None
+        self.last_spans: list[list[tuple[int, int]]] | None = None
+    def detect(
+        self, samples: Union[List[str], List[List[Dict]]], labels: Optional[np.ndarray] = None, **kwargs
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Architecture: Uses generative decoder model with multiple beam samples.
+        Differs from sequence decoder by identifying character-level hallucination spans.
+        Uses generation consistency across beams: characters wrapped in [SPAN] tags
+        by multiple generations are marked as hallucinated.
+        """
+        samples, group_ids = self._split_context_samples(samples)
+        self._check_truncation_warning(samples)
+        formatted_input = build_prompt_messages(self.config, samples)
+        preprocessed_input = self.model_adapter._preprocess_input(formatted_input)
+        preprocessed_input = [
+            f'{sample}{self.assistant_prefix}' for sample in preprocessed_input
+        ]
+        # Generate multiple sequences per sample for consistency checking
+        generated_texts = self.model_adapter.sample(
+            inputs=preprocessed_input,
+            max_tokens=self.config.max_new_tokens,
+            temperature=self.config.temperature,
+            top_p=self.config.top_p,
+            num_return_sequences=self.config.num_beams,
+            do_sample=True,
+            pad_token_id=self.model_adapter.tokenizer.eos_token_id,
+            eos_token_id=self.model_adapter.tokenizer.eos_token_id,
+            **kwargs
+        )
+        self.last_generations = list(generated_texts)
+        self.last_spans = [find_span_segments(text) for text in generated_texts]
+        # Group generated sequences by input sample
+        all_generated_sequences = []
+        num_beams = self.config.num_beams
+        for i in range(len(samples)):
+            start_idx = i * num_beams
+            sample_sequences = generated_texts[start_idx:start_idx + num_beams]
+            all_generated_sequences.append(sample_sequences)
+        # Calculate character-level probabilities from span tag consistency
+        all_char_probs = []
+        all_char_preds = []
+        for generated_texts in all_generated_sequences:
+            sample_char_probs = calculate_character_probabilities(generated_texts)
+            sample_char_probs = torch.tensor(sample_char_probs)
+            all_char_probs.append(sample_char_probs)
+            all_char_preds.append((sample_char_probs > self.threshold).long())
+        char_probs = [prob.tolist() for prob in all_char_probs]
+        char_preds = [pred.tolist() for pred in all_char_preds]
+        char_preds, char_probs = self._aggregate_context_predictions(
+            group_ids, char_preds, char_probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        # Convert span annotations to character-level labels using shared utility
+        char_labels = convert_spans_to_labels(labels, char_probs) if labels is not None else None
+        return char_probs, char_preds, char_labels
+    def _compute_metrics(self, eval_pred: Any) -> Dict[str, Any]:
+        """
+        Compute metrics for token-level decoder judge.
+        Converts token-level logits into predicted span-tagged text (teacher-forced),
+        then derives character-level binary vectors for metric computation.
+        """
+        logits = getattr(eval_pred, "predictions", eval_pred[0])
+        labels = getattr(eval_pred, "label_ids", eval_pred[1])
+        if isinstance(logits, (tuple, list)):
+            logits = logits[0]
+        pred_ids = np.argmax(logits, axis=-1)
+        probs_list = []
+        labels_list = []
+        tokenizer = self.model_adapter.tokenizer
+        for pred_ids_sample, labels_sample in zip(pred_ids, labels):
+            active_mask = labels_sample != -100
+            if not np.any(active_mask):
+                continue
+            label_ids = labels_sample[active_mask]
+            pred_ids_active = pred_ids_sample[active_mask]
+            label_text = tokenizer.decode(label_ids, skip_special_tokens=False)
+            pred_text = tokenizer.decode(pred_ids_active, skip_special_tokens=False)
+            label_vector = create_char_binary_vector(label_text)
+            pred_vector = create_char_binary_vector(pred_text)
+            if not label_vector or not pred_vector:
+                continue
+            min_len = min(len(label_vector), len(pred_vector))
+            if min_len == 0:
+                continue
+            labels_list.append(np.array(label_vector[:min_len]))
+            probs_list.append(np.array(pred_vector[:min_len], dtype=float))
+        if not labels_list:
+            return {}
+        labels_flat = np.concatenate(labels_list)
+        probs_flat = np.concatenate(probs_list)
+        metrics, threshold = calibrate_and_compute_metrics(
+            probs_flat,
+            labels_flat,
+            self.config,
+            self.metrics_to_compute,
+            is_binary=True,
+        )
+        self.threshold = threshold
+        return metrics
+    def _preprocess(self, samples: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess samples for token-level classification with decoder models."""
+        inputs = samples[INPUT_COL]
+        targets = samples[TARGET_COL]
+        self._check_truncation_warning(inputs)
+        # Format dialogue samples with system/user/assistant roles (with span wrapping)
+        formatted_messages = format_dialogue_for_training(
+            inputs, targets, self.config, is_token_level=True
+        )
+        # Prepare tokenized inputs with proper label masking
+        return prepare_decoder_inputs_with_labels(
+            formatted_messages, self.model_adapter, self.config
+        )

sirin/detection/judging/judges/token/encoder.py ADDED Viewed

	@@ -0,0 +1,131 @@

+from typing import Any, Dict, List, Optional, Tuple
+import numpy as np
+import torch
+from loguru import logger as lg
+from transformers import EvalPrediction, Trainer
+from sirin.definitions import INPUT_COL, TARGET_COL, DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import HfJudgeBase
+from sirin.detection.judging.judges.utils import (
+    process_logits_to_probs,
+    calibrate_and_compute_metrics,
+    rearrange_token_predictions_with_indices,
+    repeat_labels_with_offsets,
+)
+from sirin.detection.utils.token import get_token_labels, get_answer_offsets
+from sirin.models.detection import HfJudgeConfig
+from sirin.inference.adapters import HfModelAdapter
+class TokenEncoderJudge(HfJudgeBase):
+    """
+    Encoder judge implementation for token-level hallucination detection.
+    Uses AutoModelForTokenClassification with encoder embeddings + token classification head.
+    Handles per-token predictions with offset mapping for character-level spans.
+    """
+    detection_level = DetectionLevel.TOKEN
+    data_collator_type = DataCollatorType.TOKEN
+    trainer = Trainer
+    def __init__(self, config: HfJudgeConfig, model_adapter: HfModelAdapter):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.device = self.config.device
+        if self.model_adapter.tokenizer.pad_token is None:
+            self.model_adapter.tokenizer.pad_token = self.model_adapter.tokenizer.eos_token
+    def detect(
+        self,
+        samples: List[str],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Detects hallucinations at token level using encoder embeddings.
+        For binary: Returns 0 (no hallucination) or 1 (hallucination)
+        Architecture: Uses encoder embeddings with token classification head (AutoModelForTokenClassification).
+        Differs from sequence-level which returns one prediction per sample.
+        This approach requires offset mapping to align token predictions with character spans.
+        """
+        samples, group_ids = self._split_context_samples(samples)
+        self._check_truncation_warning(samples)
+        # Token-level approach: Get character offsets for each token in answer
+        offsets, _, answer_indices = get_answer_offsets(samples, self.model_adapter)
+        if labels is not None:
+            labels = get_token_labels(offsets, labels)
+        # Get per-token logits from token classification head
+        model_states = self.model_adapter.generate_hiddens(
+            inputs=samples,
+            return_hiddens=False,
+            return_logits=True,
+            **kwargs,
+        )
+        logits = model_states.logits
+        probs, preds = self._logits_to_probs_preds(logits)
+        probs = probs.cpu()
+        preds = preds.cpu()
+        offsets = [torch.tensor(offset).long() for offset in offsets]
+        # Rearrange token predictions to character-level using shared utility
+        char_probs, char_preds = rearrange_token_predictions_with_indices(
+            probs, preds, offsets, answer_indices
+        )
+        char_preds, char_probs = self._aggregate_context_predictions(
+            group_ids, char_preds, char_probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        if labels is not None:
+            repeats = [offset[:, 1] - offset[:, 0] for offset in offsets]
+            char_labels = repeat_labels_with_offsets(labels, repeats)
+        else:
+            char_labels = None
+        return char_probs, char_preds, char_labels
+    def _compute_metrics(self, eval_pred: EvalPrediction) -> Dict[str, Any]:
+        """
+        Compute metrics for token-level classification.
+        Filters out padding tokens (-100) before computing metrics.
+        """
+        logits, labels = eval_pred
+        # Convert logits to probabilities, filtering padding tokens
+        probs, active_labels = process_logits_to_probs(logits, labels, filter_padding=True)
+        # Token-level is always binary classification
+        metrics, threshold = calibrate_and_compute_metrics(
+            probs, active_labels, self.config, self.metrics_to_compute, is_binary=True
+        )
+        self.threshold = threshold
+        return metrics
+    def _preprocess(self, samples: Dict[str, Any]) -> Dict[str, Any]:
+        """Preprocess samples for token-level classification."""
+        inputs = samples[INPUT_COL]
+        targets = samples[TARGET_COL]
+        self._check_truncation_warning(inputs)
+        answer_offsets, tokenized, _ = get_answer_offsets(inputs, self.model_adapter)
+        labels_raw = get_token_labels(answer_offsets=answer_offsets, target_spans=targets)
+        labels = [[-100] * len(sample) for sample in tokenized['input_ids']]
+        for i in range(len(labels)):
+            labels[i][-len(labels_raw[i]):] = labels_raw[i]
+        tokenized['labels'] = labels
+        return tokenized

sirin/detection/judging/judges/token/openai.py ADDED Viewed

	@@ -0,0 +1,139 @@

+from typing import Dict, List, Optional, Tuple, Union
+import numpy as np
+import torch
+from sirin.definitions import DetectionLevel, DataCollatorType
+from sirin.detection.judging.judges.base import JudgeAnnotationError, OpenAIJudgeBase
+from sirin.detection.judging.judges.utils import (
+    build_prompt_messages,
+    calculate_character_probabilities,
+    extract_answer_from_generation,
+    find_span_segments,
+)
+from sirin.detection.utils.token import convert_spans_to_labels
+from sirin.inference.adapters import OpenAIModelAdapter
+from sirin.models.detection import OpenAIJudgeConfig
+class TokenOpenAIJudge(OpenAIJudgeBase):
+    """
+    OpenAI API judge implementation for token-level hallucination detection.
+    Architecture: Uses OpenAI API with multiple generations to identify character-level spans.
+    Similar to TokenDecoderJudge but uses API instead of local model.
+    No training support - inference only via API calls.
+    """
+    detection_level = DetectionLevel.TOKEN
+    def __init__(self, config: OpenAIJudgeConfig, model_adapter: OpenAIModelAdapter):
+        super().__init__(config=config, model_adapter=model_adapter)
+        self.class_token_ids = None
+        self.last_generations: list[str] | None = None
+        self.last_spans: list[list[tuple[int, int]]] | None = None
+        self.last_consensus: list[dict] | None = None
+    @staticmethod
+    def _reference_answer(sample: Union[str, List[Dict]]) -> str:
+        """The assistant answer the char scores must align to (matches format_dialogue_samples)."""
+        if isinstance(sample, list) and len(sample) >= 2 and isinstance(sample[1], dict):
+            return sample[1]['content']
+        return sample
+    @staticmethod
+    def _echo_of(generation: Optional[str]) -> str:
+        """A generation's answer with reasoning wrapper and span tags removed, for echo checking.
+        A ``None``/empty generation (a provider may return empty content, e.g. when a reasoning
+        budget is exhausted) is just an invalid vote: it returns '' and fails the echo check.
+        """
+        if not generation:
+            return ''
+        text = extract_answer_from_generation(generation)  # strips <think>… and outer whitespace
+        return text.replace('[SPAN]', '').replace('[/SPAN]', '').strip()
+    def detect(
+        self,
+        samples: Union[List[str], List[List[Dict]]],
+        labels: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> Tuple[np.ndarray, np.ndarray, Optional[np.ndarray]]:
+        """
+        Uses OpenAI API for token-level classification with multiple generations.
+        Each sample is annotated by ``num_beams`` independent generations. A generation only
+        votes if it echoes the reference answer verbatim (after stripping any reasoning wrapper
+        and the span tags); paraphrases/misaligned echoes are dropped. Each character's score is
+        the fraction of VALID generations that flagged it, over the reference's characters. If no
+        generation validates for a sample we raise instead of emitting a misleading all-clear.
+        """
+        samples, group_ids = self._split_context_samples(samples)
+        references = [self._reference_answer(sample) for sample in samples]
+        formatted_input = build_prompt_messages(self.config, samples)
+        n = self.config.num_beams
+        generated = self.model_adapter.sample(
+            inputs=formatted_input,
+            max_tokens=self.config.max_new_tokens,  # 100-token default truncates a real answer.
+            temperature=self.config.temperature,
+            top_p=self.config.top_p,
+            n=n,
+        )
+        # Adapter contract: n==1 -> list[str] (one per sample); n>1 -> list[list[str]] (n per sample).
+        per_sample_gens = [[g] for g in generated] if n == 1 else generated
+        # Adapter side channel (n>1): per-sample finish_reasons aligned 1:1 with the generations;
+        # 'length' means the completion budget truncated that generation mid-reasoning.
+        per_sample_reasons = getattr(self.model_adapter, 'last_finish_reasons', None) or []
+        all_char_probs = []
+        all_char_preds = []
+        self.last_consensus = []
+        for i, (reference, gens) in enumerate(zip(references, per_sample_gens)):
+            reasons = list(per_sample_reasons[i]) if i < len(per_sample_reasons) else []
+            reasons += [None] * (len(gens) - len(reasons))
+            valid = []
+            invalid = {'truncated': 0, 'empty': 0, 'not_verbatim': 0}
+            for gen, reason in zip(gens, reasons):
+                if self._echo_of(gen) == str(reference).strip():
+                    valid.append(gen)
+                elif reason == 'length':
+                    invalid['truncated'] += 1
+                elif not gen:
+                    invalid['empty'] += 1
+                else:
+                    invalid['not_verbatim'] += 1
+            entry = {
+                'requested': n, 'valid': len(valid), 'temperature': self.config.temperature,
+            }
+            counts = {kind: count for kind, count in invalid.items() if count}
+            if counts:
+                entry['invalid'] = counts
+            self.last_consensus.append(entry)
+            if not valid:
+                raise JudgeAnnotationError(
+                    f'No judge generation echoed the answer verbatim '
+                    f'({len(gens)} generation(s) all dropped); cannot annotate spans.'
+                )
+            sample_char_probs = torch.tensor(
+                calculate_character_probabilities(valid, reference=reference)
+            )
+            all_char_probs.append(sample_char_probs)
+            all_char_preds.append((sample_char_probs > self.threshold).long())
+        # Expose sample 0's generations/spans for the UI (reads gens[0]/spans[0]).
+        self.last_generations = list(per_sample_gens[0]) if per_sample_gens else []
+        # `or ''`: a provider may return None content (e.g. exhausted reasoning budget) for one
+        # generation while the sample still validates on the others.
+        self.last_spans = [find_span_segments(g or '') for g in self.last_generations]
+        char_probs = [prob.tolist() for prob in all_char_probs]
+        char_preds = [pred.tolist() for pred in all_char_preds]
+        char_preds, char_probs = self._aggregate_context_predictions(
+            group_ids, char_preds, char_probs, binary=(self.config.num_classification_heads <= 2)
+        )
+        # Convert span annotations to character-level labels using shared utility
+        char_labels = convert_spans_to_labels(labels, char_probs) if labels is not None else None
+        return char_probs, char_preds, char_labels

sirin/detection/judging/judges/utils/__init__.py ADDED Viewed

	@@ -0,0 +1,23 @@

+from sirin.detection.judging.judges.utils.decoder import (
+    prepare_decoder_inputs_with_labels,
+)
+from sirin.detection.judging.judges.utils.metrics import (
+    calibrate_and_compute_metrics,
+    process_logits_to_probs,
+)
+from sirin.detection.judging.judges.utils.prompts import (
+    build_prompt_messages,
+    format_dialogue_for_training,
+    format_dialogue_samples,
+)
+from sirin.detection.judging.judges.utils.token_level import (
+    calculate_character_probabilities,
+    create_char_binary_vector,
+    extract_answer_from_generation,
+    find_span_segments,
+    merge_overlapping_spans,
+    rearrange_token_predictions_with_indices,
+    repeat_labels_with_offsets,
+    wrap_spans,
+)
+from sirin.detection.utils.logits import logits_to_probs_preds

sirin/detection/judging/judges/utils/decoder.py ADDED Viewed

	@@ -0,0 +1,67 @@

+from typing import Any, Dict, List
+from sirin.utils.hf import get_assistant_prefix
+def prepare_decoder_inputs_with_labels(
+    formatted_messages: List[List[Dict[str, str]]],
+    model_adapter: Any,
+    config: Any,
+) -> Dict[str, Any]:
+    """
+    Prepare inputs for decoder models with proper label masking.
+    Creates tokenized inputs where only the assistant response tokens are used for loss.
+    Context tokens (system + user prompts) are masked with -100.
+    Args:
+        formatted_messages: Formatted message lists (system, user, assistant)
+        model_adapter: Model adapter with tokenizer and preprocessing
+        config: Judge configuration with tokenization settings
+    Returns:
+        Dict with 'input_ids', 'attention_mask', 'labels'
+    """
+    # Get assistant prefix (e.g., "assistant\n" for the model's format)
+    assistant_prefix = get_assistant_prefix(model_adapter._model_name)
+    # Preprocess messages without assistant response to find context length
+    preprocessed_input_without_answer = [
+        f'{sample}{assistant_prefix}'
+        for sample in model_adapter._preprocess_input(
+            [sample[:2] for sample in formatted_messages]  # Only system + user
+        )
+    ]
+    # Preprocess full messages including assistant response
+    preprocessed_input = model_adapter._preprocess_input(formatted_messages)
+    # Tokenize without target (to find where assistant response starts)
+    encoding_without_target = model_adapter.tokenizer(
+        preprocessed_input_without_answer,
+        truncation=model_adapter.config.truncation,
+        padding=False,
+        add_special_tokens=False,
+    )
+    # Tokenize with target
+    encoding = model_adapter.tokenizer(
+        preprocessed_input,
+        truncation=model_adapter.config.truncation,
+        padding=False,
+        return_tensors=config.return_tensors,
+        add_special_tokens=False,
+    )
+    # Create labels: copy input_ids but mask everything before assistant response
+    labels = encoding['input_ids'].clone()
+    for i, sample in enumerate(encoding_without_target['input_ids']):
+        labels[i, :len(sample)] = -100  # Mask context tokens
+    return {
+        'input_ids': encoding['input_ids'],
+        'attention_mask': encoding['attention_mask'],
+        'labels': labels,
+    }

sirin/detection/judging/judges/utils/metrics.py ADDED Viewed

	@@ -0,0 +1,106 @@

+from typing import Any, Dict, List, Optional
+import numpy as np
+import torch
+from loguru import logger as lg
+from sirin.definitions import ClassificationMetric
+from sirin.detection.utils.basic import calibrate_threshold
+from sirin.metrics.classification import calculate_classification_metrics
+def process_logits_to_probs(
+    logits: np.ndarray,
+    labels: Optional[np.ndarray] = None,
+    filter_padding: bool = False,
+) -> tuple[np.ndarray, np.ndarray]:
+    """
+    Convert logits to probabilities, handling binary/multiclass cases.
+    Args:
+        logits: Raw logits from model (shape: [batch, num_classes] or [batch, seq, num_classes])
+        labels: Optional labels for filtering padding tokens (token-level only)
+        filter_padding: Whether to filter out -100 padding labels (for token-level)
+    Returns:
+        tuple of (probs, active_labels) where:
+        - probs: probability values
+        - active_labels: filtered labels (only if filter_padding=True, else original labels)
+    """
+    # Handle token-level: flatten and filter padding
+    if filter_padding and labels is not None:
+        logits_flat = logits.reshape(-1, logits.shape[-1])
+        labels_flat = labels.reshape(-1)
+        active_mask = labels_flat != -100
+        logits = logits_flat[active_mask]
+        labels = labels_flat[active_mask]
+    logits_tensor = torch.from_numpy(logits) if isinstance(logits, np.ndarray) else logits
+    num_classes = logits_tensor.shape[-1]
+    if num_classes == 1:
+        # Binary classification with single output
+        probs = torch.sigmoid(logits_tensor).squeeze(-1)
+    elif num_classes == 2:
+        # Binary classification with two outputs
+        probs = torch.softmax(logits_tensor, dim=-1)[:, 1]
+    else:
+        # Multiclass classification
+        probs_full = torch.softmax(logits_tensor, dim=-1)
+        probs = torch.max(probs_full, dim=-1).values
+    probs_np = probs.numpy() if isinstance(probs, torch.Tensor) else probs
+    return probs_np, labels
+def calibrate_and_compute_metrics(
+    probs: np.ndarray,
+    labels: np.ndarray,
+    config: Any,
+    metrics_to_compute: List[ClassificationMetric],
+    is_binary: bool = True,
+) -> tuple[Dict[str, Any], Optional[float]]:
+    """
+    Calibrate threshold and compute classification metrics.
+    Args:
+        probs: Probability predictions
+        labels: Ground truth labels
+        config: Judge configuration with threshold settings
+        metrics_to_compute: List of metrics to calculate
+        is_binary: Whether this is binary classification (affects threshold usage)
+    Returns:
+        tuple of (metrics_dict, threshold) where:
+        - metrics_dict: Computed metric values
+        - threshold: Calibrated threshold (None for multiclass)
+    """
+    if is_binary:
+        threshold = calibrate_threshold(
+            probs,
+            labels,
+            config.threshold_method,
+            config.threshold_percentile,
+            config.fixed_threshold,
+        )
+        preds = (probs > threshold).astype(int)
+    else:
+        threshold = None
+        if len(probs.shape) > 1:
+            # Multiclass: use argmax for predictions
+            preds = np.argmax(probs, axis=-1)
+        else:
+            preds = probs.astype(int)
+    metrics = calculate_classification_metrics(
+        labels,
+        probs,
+        preds,
+        metrics=metrics_to_compute,
+    )
+    lg.info(f"Evaluation metrics: {metrics}")
+    return metrics, threshold

sirin/detection/judging/judges/utils/prompts.py ADDED Viewed

	@@ -0,0 +1,127 @@

+from typing import Any, Dict, List
+# Span-annotation prompts for the token-level API judge. The judge parses [SPAN]...[/SPAN]
+# tags out of the model's echoed answer, so the model MUST reproduce the answer verbatim and
+# only insert tags — anything else (paraphrase, corrections, verdict digits) breaks the
+# character alignment and is dropped by the reference-echo check.
+SPAN_TAG_SYSTEM_PROMPT = (
+    "You are a hallucination span annotator. You receive a dialogue: a user prompt "
+    "(which may include a context passage) and an assistant answer. Reproduce the "
+    "assistant's answer EXACTLY, character for character, and insert the markers [SPAN] "
+    "and [/SPAN] around every hallucinated part — any span not supported by, or "
+    "contradicted by, the context.\n"
+    "Rules:\n"
+    "- Copy the answer verbatim. Do NOT correct, paraphrase, reorder, translate, or add "
+    "anything; the ONLY characters you may add are the [SPAN] and [/SPAN] markers.\n"
+    "- Wrap only the hallucinated words; you may use the markers multiple times.\n"
+    "- If nothing is hallucinated, output the answer unchanged with no markers.\n"
+    "- Output ONLY the annotated answer: no explanations, no verdict digits, no quotes, "
+    "no code fences."
+)
+SPAN_TAG_USER_PROMPT = (
+    "Dialogue:\n{sample}\n\n"
+    "Return the assistant's answer verbatim, wrapping each hallucinated span in "
+    "[SPAN]...[/SPAN]. If no content is hallucinated, return the answer unchanged."
+)
+def format_dialogue_samples(config: Any, samples: List[Any]) -> List[Any]:
+    """
+    Format dialogue samples using the configured dialogue format.
+    Args:
+        config: Judge configuration with dialogue_format
+        samples: List of samples (either strings or list of message dicts)
+    Returns:
+        List of formatted dialogue strings
+    """
+    if not samples:
+        return []
+    first = samples[0]
+    if isinstance(first, list) and len(first) >= 2:
+        return [
+            config.dialogue_format.format(
+                **{'question': sample[0]['content'], 'answer': sample[1]['content']}
+            )
+            for sample in samples
+        ]
+    return samples
+def build_prompt_messages(config: Any, samples: List[Any]) -> List[List[Dict[str, str]]]:
+    """
+    Build full prompt messages with system and user roles.
+    Used by API-based judges (OpenAI) to create message lists for API calls.
+    Args:
+        config: Judge configuration with system_prompt, user_prompt, dialogue_format
+        samples: List of samples to format
+    Returns:
+        List of message lists, each containing system and user messages
+    """
+    formatted_samples = format_dialogue_samples(config, samples)
+    return [
+        [
+            {'role': 'system', 'content': config.system_prompt},
+            {'role': 'user', 'content': config.user_prompt.format(**{'sample': sample})},
+        ]
+        for sample in formatted_samples
+    ]
+def format_dialogue_for_training(
+    inputs: List[List[Dict[str, str]]],
+    targets: List[Any],
+    config: Any,
+    is_token_level: bool = False,
+) -> List[List[Dict[str, str]]]:
+    """
+    Format dialogue samples with system/user/assistant roles for training.
+    Used by decoder judges (both sequence and token-level) for preprocessing.
+    Args:
+        inputs: List of dialogue samples (each is list of messages)
+        targets: List of target labels (int for sequence, List[Tuple] for token)
+        config: Judge configuration with dialogue_format, system_prompt, user_prompt
+        is_token_level: Whether this is token-level (uses span wrapping) or sequence-level
+    Returns:
+        List of formatted message lists ready for tokenization
+    """
+    from sirin.detection.judging.judges.utils.token_level import wrap_spans
+    # Format dialogues using the config's dialogue format
+    formatted_input = [
+        config.dialogue_format.format(
+            **{'question': sample[0]['content'], 'answer': sample[1]['content']}
+        )
+        for sample in inputs
+    ]
+    # Build full message lists with system/user/assistant roles
+    formatted_messages = []
+    for i, sample_text in enumerate(formatted_input):
+        # For token-level, wrap spans in the assistant content
+        if is_token_level:
+            dialogue_sample = inputs[i]
+            spans = targets[i]
+            assistant_content = wrap_spans(dialogue_sample[1]['content'], spans)
+        else:
+            # For sequence-level, use the label directly
+            assistant_content = str(targets[i])
+        formatted_messages.append([
+            {'role': 'system', 'content': config.system_prompt},
+            {'role': 'user', 'content': config.user_prompt.format(**{'sample': sample_text})},
+            {'role': 'assistant', 'content': assistant_content},
+        ])
+    return formatted_messages