Whisper Hallucination Detection via Confidence Scores

A TensorFlow model for detecting hallucinated segments in transcriptions using confidence scores. It supports custom loss masking and layer expansion.

Whisper doesn't provide confidence scores out of the box. The trick here is to utilize word timestamp probabilities either generated by Faster Whisper or Whisper Timestamped as confidence scores.

Custom Objects

ExpandDimsLayer
masked_binary_crossentropy

Usage

from huggingface_hub import snapshot_download
import tensorflow as tf
import sys, os, json

# 1. Download entire repo to local cache
local_repo_path = snapshot_download(repo_id="avcton/whisper-hallucination-detection")

# 2. Add the folder to sys.path for imports
sys.path.append(local_repo_path)

# 3. Import custom objects
from custom_objects import ExpandDimsLayer, masked_binary_crossentropy
from predict_and_decode_utils import predict_and_decode

# 4. Load config
with open(os.path.join(local_repo_path, "config.json")) as f:
    cfg = json.load(f)

# 5. Load model
model = tf.keras.models.load_model(
    os.path.join(local_repo_path, "model.keras"),
    custom_objects={
        "ExpandDimsLayer": ExpandDimsLayer,
        "masked_binary_crossentropy": masked_binary_crossentropy
    }
)

# 6. Call predictions
results = predict_and_decode(
    model,
    transcriptions=...,
    max_segment_len=cfg["max_segment_len"],
    max_score_len=cfg["max_score_len"],
    binary_threshold=cfg["binary_threshold"],
    label_mapping=cfg["label_mapping"]
)

Documentation

Description for predict_and_decode function:

  > predict_and_decode(model, transcriptions, max_segment_len, max_score_len, binary_threshold, label_mapping)

  Makes predictions on a list of transcriptions using the trained model and decodes the results.

  Args:
    model: The trained Keras model.
    transcriptions: A list of transcriptions, where each transcription is a list of score lists.
                Example: [[[t1_seg1_score1, t1_seg1_score2], [t1_seg2_score1, t1_seg2_score2]], [[t2_seg1_score1, t2_seg1_score2], [t1_seg2_score1]]]
    max_segment_len: The maximum number of segments in a transcription (used for padding).
    max_score_len: The maximum number of scores in a segment (used for padding).
    binary_threshold: The best threshold to convert probabilities to binary predictions.
    label_mapping: The dictionary used to map labels ('H', 'NH') to integers (1, 0).

  Returns:
    A list of lists, where each inner list contains the predicted class labels
    ('H' or 'NH') for the segments in the corresponding input transcription. Padded transcriptions will not have a prediction.
    Example: [['H', 'NH'], ['NH', 'H']]

Limitations

Maximum segment length in a transcription can only be 83.
Maximum scores/words in a segment can only be 112.

Config

Binary Threshold = 0.4187 (Found using Youden's J statistic)
Label Mapping = {"H":1, "NH":0}