dam-AA / handler.py
aabouzeid
Add HuggingFace Inference Endpoint handler
0a56a24
Raw
History Blame
880 Bytes
"""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()},
}