记录三模型话轮检测实验、44点评估及延迟对比

experiments/smart-turn-v3.2/inference.py

@@ -0,0 +1,75 @@
+from pathlib import Path
+
+import numpy as np
+import onnxruntime as ort
+from transformers import WhisperFeatureExtractor
+
+from audio_utils import truncate_audio_to_last_n_seconds
+
+ONNX_MODEL_PATH = str(Path(__file__).with_name("smart-turn-v3.2-cpu.onnx"))
+
+def build_session(onnx_path):
+    so = ort.SessionOptions()
+    so.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
+    so.inter_op_num_threads = 1
+    so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+    return ort.InferenceSession(onnx_path, sess_options=so, providers=["CPUExecutionProvider"])
+
+feature_extractor = WhisperFeatureExtractor(chunk_length=8)
+session = build_session(ONNX_MODEL_PATH)
+
+
+def predict_endpoint(audio_array):
+    """
+    Predict whether an audio segment is complete (turn ended) or incomplete.
+
+    Args:
+        audio_array: Numpy array containing audio samples at 16kHz
+
+    Returns:
+        Dictionary containing prediction results:
+        - prediction: 1 for complete, 0 for incomplete
+        - probability: Probability of completion (sigmoid output)
+    """
+
+    # Truncate to 8 seconds (keeping the end) or pad to 8 seconds
+    audio_array = truncate_audio_to_last_n_seconds(audio_array, n_seconds=8)
+
+    # Process audio using Whisper's feature extractor
+    inputs = feature_extractor(
+        audio_array,
+        sampling_rate=16000,
+        return_tensors="np",
+        padding="max_length",
+        max_length=8 * 16000,
+        truncation=True,
+        do_normalize=True,
+    )
+
+    # Extract features and ensure correct shape for ONNX
+    input_features = inputs.input_features.squeeze(0).astype(np.float32)
+    input_features = np.expand_dims(input_features, axis=0)  # Add batch dimension
+
+    # Run ONNX inference
+    outputs = session.run(None, {"input_features": input_features})
+
+    # Extract probability (ONNX model returns sigmoid probabilities)
+    probability = outputs[0][0].item()
+
+    # Make prediction (1 for Complete, 0 for Incomplete)
+    prediction = 1 if probability > 0.5 else 0
+
+    return {
+        "prediction": prediction,
+        "probability": probability,
+    }
+
+
+# Example usage
+if __name__ == "__main__":
+    # Create a dummy audio array for testing (1 second of random audio)
+    dummy_audio = np.random.randn(16000).astype(np.float32)
+
+    result = predict_endpoint(dummy_audio)
+    print(f"Prediction: {result['prediction']}")
+    print(f"Probability: {result['probability']:.4f}")