Hugging Face's logo

Snapkitty
/
bert-agent

Model card Files
xet
 Community
bert-agent
File size: 5,423 Bytes
30f011f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
//! server.rs — Axum HTTP server
//!
//! POST /verify
//! Body: { "premise": "...", "hypothesis": "...", "chunk_id": "..." }
//! Response: { "score": 0.94, "verdict": "Entailment", "hash": "abc123..." }
//!
//! The handler tokenises the (premise, hypothesis) pair using the same
//! cross-encoder format as training: [CLS] premise [SEP] hypothesis [SEP].
//! It then sends a VerifyRequest through the MPSC channel to the inference
//! daemon and awaits the oneshot response.

use std::sync::Arc;

use axum::{
    extract::State,
    http::StatusCode,
    response::IntoResponse,
    routing::post,
    Json, Router,
};
use serde::{Deserialize, Serialize};
use tokio::sync::{mpsc, oneshot};

use crate::types::{DaemonConfig, VerifyRequest, VerifyResponse, Verdict};

// ── Request / Response DTOs ─────────────────────────────────────────────────

#[derive(Debug, Deserialize)]
pub struct VerifyBody {
    pub premise:    String,    // retrieved source chunk
    pub hypothesis: String,    // LLM generated claim
    pub chunk_id:   String,
}

#[derive(Debug, Serialize)]
pub struct VerifyReply {
    pub score:   f32,
    pub verdict: String,
    pub hash:    String,       // BLAKE3 hex for the audit ledger
}

// ── Shared state ─────────────────────────────────────────────────────────────

pub struct AppState {
    pub tx:       mpsc::Sender<VerifyRequest>,
    pub cfg:      Arc<DaemonConfig>,
    pub tokenizer: Arc<tokenizers::Tokenizer>,
}

// ── Tokenisation ─────────────────────────────────────────────────────────────

/// Encode (premise, hypothesis) as a cross-encoder input:
///   [CLS] premise_tokens [SEP] hypothesis_tokens [SEP]
fn encode_pair(
    tokenizer: &tokenizers::Tokenizer,
    premise: &str,
    hypothesis: &str,
    max_length: usize,
) -> (Vec<i64>, Vec<i64>, Vec<i64>) {
    use tokenizers::EncodeInput;

    let encoding = tokenizer
        .encode(
            EncodeInput::Dual(
                tokenizers::InputSequence::Raw(premise.into()),
                tokenizers::InputSequence::Raw(hypothesis.into()),
            ),
            true,
        )
        .expect("tokenisation failed");

    let ids:   Vec<i64> = encoding.get_ids().iter().map(|&x| x as i64).collect();
    let mask:  Vec<i64> = encoding.get_attention_mask().iter().map(|&x| x as i64).collect();
    let types: Vec<i64> = encoding.get_type_ids().iter().map(|&x| x as i64).collect();

    // Truncate to max_length
    let trunc = |v: Vec<i64>| v.into_iter().take(max_length).collect::<Vec<_>>();
    (trunc(ids), trunc(mask), trunc(types))
}

// ── Handler ──────────────────────────────────────────────────────────────────

async fn verify_handler(
    State(state): State<Arc<AppState>>,
    Json(body): Json<VerifyBody>,
) -> impl IntoResponse {
    let (input_ids, attention_mask, token_type_ids) = encode_pair(
        &state.tokenizer,
        &body.premise,
        &body.hypothesis,
        512,
    );

    let (resp_tx, resp_rx) = oneshot::channel::<VerifyResponse>();

    let request = VerifyRequest {
        input_ids,
        attention_mask,
        token_type_ids,
        chunk_id:   body.chunk_id,
        claim_text: body.hypothesis,
        responder:  resp_tx,
    };

    if state.tx.send(request).await.is_err() {
        return (
            StatusCode::SERVICE_UNAVAILABLE,
            Json(serde_json::json!({"error": "inference daemon unavailable"})),
        );
    }

    match resp_rx.await {
        Ok(resp) => {
            let hash_hex = resp.attestation_hash
                .iter()
                .map(|b| format!("{:02x}", b))
                .collect::<String>();
            (
                StatusCode::OK,
                Json(serde_json::json!({
                    "score":   resp.entailment_score,
                    "verdict": format!("{:?}", resp.label),
                    "hash":    hash_hex,
                })),
            )
        }
        Err(_) => (
            StatusCode::INTERNAL_SERVER_ERROR,
            Json(serde_json::json!({"error": "inference worker dropped"})),
        ),
    }
}

// ── Health check ─────────────────────────────────────────────────────────────

async fn health_handler() -> impl IntoResponse {
    Json(serde_json::json!({"status": "ok"}))
}

// ── Router builder ───────────────────────────────────────────────────────────

pub fn build_router(state: Arc<AppState>) -> Router {
    Router::new()
        .route("/verify", post(verify_handler))
        .route("/health", axum::routing::get(health_handler))
        .with_state(state)
}