handler.py

"""HuggingFace Inference Endpoint handler for the DAM model."""

import io
from pathlib import Path

from pipeline import Pipeline


class EndpointHandler:
    def __init__(self, path=""):
        checkpoint = Path(path) / "dam3.1.ckpt"
        self.pipeline = Pipeline(checkpoint=checkpoint)

    def __call__(self, data):
        inputs = data.get("inputs")
        if isinstance(inputs, bytes):
            source = io.BytesIO(inputs)
        else:
            raise ValueError("Expected raw audio bytes in data['inputs']")

        quantized = self.pipeline.run_on_file(source, quantize=True)
        source.seek(0)
        raw = self.pipeline.run_on_file(source, quantize=False)

        return {
            "depression": quantized["depression"],
            "anxiety": quantized["anxiety"],
            "raw_scores": {k: v.mean().item() for k, v in raw.items()},
        }