AKSHAY BIRADAR

Add HF Inference API handler + model test frontend

60e735c 2 months ago

4.51 kB

"""
Hugging Face Inference API handler for MultimodalFusionModel.
Handles POST requests with:
  - "text": patient speech transcript (required)
  - "audio": base64-encoded WAV audio, 16kHz mono float32 (optional)
Returns:
  { "repetition": 0.xx, "confusion": 0.xx, "distress": 0.xx, "topic_switch": 0.xx }
"""
import os, io, base64, struct, json
import numpy as np
import torch

CLASSES = ["repetition", "confusion", "distress", "topic_switch"]

_model = None
_device = None


def load_model():
    """Load MultimodalFusionModel from the repo's best_model.pt."""
    global _model, _device
    if _model is not None:
        return _model

    # Import model class from repo (model.py is in the same directory)
    import sys
    sys.path.insert(0, os.path.dirname(__file__))
    from model import MultimodalFusionModel

    _device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    model = MultimodalFusionModel(
        audio_emb_dim=768,
        text_emb_dim=768,
        fusion_hidden=256,
        lstm_layers=2,
        num_heads=4,
        num_classes=4,
        freeze_encoders=True,
    )

    ckpt_path = os.path.join(os.path.dirname(__file__), "best_model.pt")
    state = torch.load(ckpt_path, map_location=_device, weights_only=False)

    # Handle different checkpoint formats
    if isinstance(state, dict) and "model_state_dict" in state:
        state = state["model_state_dict"]
    elif isinstance(state, dict) and "state_dict" in state:
        state = state["state_dict"]

    # Strip "module." prefix (DataParallel / DistributedDataParallel)
    cleaned = {}
    for k, v in state.items():
        cleaned[k.replace("module.", "", 1) if k.startswith("module.") else k] = v

    # Load only matching keys (skip encoder weights if not in checkpoint)
    model_state = model.state_dict()
    matched = {k: v for k, v in cleaned.items() if k in model_state and v.shape == model_state[k].shape}
    model_state.update(matched)
    model.load_state_dict(model_state)

    model.eval().to(_device)
    _model = model
    return _model


def _decode_audio_b64(b64_str):
    """Decode base64 WAV (16kHz mono) to float32 numpy array."""
    raw = base64.b64decode(b64_str)
    # Parse WAV header: skip 44 bytes for standard PCM WAV
    if raw[:4] == b"RIFF":
        # Find "data" chunk
        pos = 12
        while pos < len(raw) - 8:
            chunk_id = raw[pos:pos+4]
            chunk_size = struct.unpack("<I", raw[pos+4:pos+8])[0]
            if chunk_id == b"data":
                audio_data = raw[pos+8:pos+8+chunk_size]
                break
            pos += 8 + chunk_size
        else:
            audio_data = raw[44:]  # fallback
    else:
        audio_data = raw  # raw PCM

    # Try to parse as float32 PCM (model expects float waveforms)
    samples = np.frombuffer(audio_data, dtype=np.float32)
    if samples.shape[0] == 0:
        # Try int16 → float32
        samples = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
    return samples


def predict(data):
    """
    Run inference.
    Parameters:
        data: dict with keys:
            - "text" (str, required): transcribed speech
            - "audio" (str, optional): base64-encoded 16kHz mono float32 WAV
    Returns:
        dict: { class_name: probability float }
    """
    model = load_model()
    text = data.get("text", "")
    audio_b64 = data.get("audio", None)

    # Build waveform
    if audio_b64:
        waveform = _decode_audio_b64(audio_b64)
        waveform = torch.tensor(waveform, dtype=torch.float32).unsqueeze(0).to(_device)
    else:
        # No audio — use silence (model still works via text encoder)
        waveform = torch.zeros(1, 16000, dtype=torch.float32).to(_device)

    # Tokenize text
    enc = model.tokenizer(
        text, truncation=True, padding=True,
        max_length=128, return_tensors="pt"
    )
    input_ids = enc["input_ids"].to(_device)
    attention_mask = enc["attention_mask"].to(_device)

    # Compute embeddings → forward → sigmoid
    with torch.no_grad():
        audio_emb, text_emb = model.get_embeddings(waveform, input_ids, attention_mask)
        # Add turn dimension: (1, 1, emb_dim)
        audio_emb = audio_emb.unsqueeze(1)
        text_emb = text_emb.unsqueeze(1)
        logits = model.forward(audio_embeds=audio_emb, text_embeds=text_emb)
        probs = torch.sigmoid(logits).squeeze(0).cpu().numpy()

    return {CLASSES[i]: round(float(probs[i]), 4) for i in range(len(CLASSES))}

	"""
	Hugging Face Inference API handler for MultimodalFusionModel.

	Handles POST requests with:
	- "text": patient speech transcript (required)
	- "audio": base64-encoded WAV audio, 16kHz mono float32 (optional)

	Returns:
	{ "repetition": 0.xx, "confusion": 0.xx, "distress": 0.xx, "topic_switch": 0.xx }
	"""
	import os, io, base64, struct, json
	import numpy as np
	import torch

	CLASSES = ["repetition", "confusion", "distress", "topic_switch"]

	_model = None
	_device = None


	def load_model():
	"""Load MultimodalFusionModel from the repo's best_model.pt."""
	global _model, _device
	if _model is not None:
	return _model

	# Import model class from repo (model.py is in the same directory)
	import sys
	sys.path.insert(0, os.path.dirname(__file__))
	from model import MultimodalFusionModel

	_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

	model = MultimodalFusionModel(
	audio_emb_dim=768,
	text_emb_dim=768,
	fusion_hidden=256,
	lstm_layers=2,
	num_heads=4,
	num_classes=4,
	freeze_encoders=True,
	)

	ckpt_path = os.path.join(os.path.dirname(__file__), "best_model.pt")
	state = torch.load(ckpt_path, map_location=_device, weights_only=False)

	# Handle different checkpoint formats
	if isinstance(state, dict) and "model_state_dict" in state:
	state = state["model_state_dict"]
	elif isinstance(state, dict) and "state_dict" in state:
	state = state["state_dict"]

	# Strip "module." prefix (DataParallel / DistributedDataParallel)
	cleaned = {}
	for k, v in state.items():
	cleaned[k.replace("module.", "", 1) if k.startswith("module.") else k] = v

	# Load only matching keys (skip encoder weights if not in checkpoint)
	model_state = model.state_dict()
	matched = {k: v for k, v in cleaned.items() if k in model_state and v.shape == model_state[k].shape}
	model_state.update(matched)
	model.load_state_dict(model_state)

	model.eval().to(_device)
	_model = model
	return _model


	def _decode_audio_b64(b64_str):
	"""Decode base64 WAV (16kHz mono) to float32 numpy array."""
	raw = base64.b64decode(b64_str)
	# Parse WAV header: skip 44 bytes for standard PCM WAV
	if raw[:4] == b"RIFF":
	# Find "data" chunk
	pos = 12
	while pos < len(raw) - 8:
	chunk_id = raw[pos:pos+4]
	chunk_size = struct.unpack("<I", raw[pos+4:pos+8])[0]
	if chunk_id == b"data":
	audio_data = raw[pos+8:pos+8+chunk_size]
	break
	pos += 8 + chunk_size
	else:
	audio_data = raw[44:] # fallback
	else:
	audio_data = raw # raw PCM

	# Try to parse as float32 PCM (model expects float waveforms)
	samples = np.frombuffer(audio_data, dtype=np.float32)
	if samples.shape[0] == 0:
	# Try int16 → float32
	samples = np.frombuffer(audio_data, dtype=np.int16).astype(np.float32) / 32768.0
	return samples


	def predict(data):
	"""
	Run inference.

	Parameters:
	data: dict with keys:
	- "text" (str, required): transcribed speech
	- "audio" (str, optional): base64-encoded 16kHz mono float32 WAV
	Returns:
	dict: { class_name: probability float }
	"""
	model = load_model()
	text = data.get("text", "")
	audio_b64 = data.get("audio", None)

	# Build waveform
	if audio_b64:
	waveform = _decode_audio_b64(audio_b64)
	waveform = torch.tensor(waveform, dtype=torch.float32).unsqueeze(0).to(_device)
	else:
	# No audio — use silence (model still works via text encoder)
	waveform = torch.zeros(1, 16000, dtype=torch.float32).to(_device)

	# Tokenize text
	enc = model.tokenizer(
	text, truncation=True, padding=True,
	max_length=128, return_tensors="pt"
	)
	input_ids = enc["input_ids"].to(_device)
	attention_mask = enc["attention_mask"].to(_device)

	# Compute embeddings → forward → sigmoid
	with torch.no_grad():
	audio_emb, text_emb = model.get_embeddings(waveform, input_ids, attention_mask)
	# Add turn dimension: (1, 1, emb_dim)
	audio_emb = audio_emb.unsqueeze(1)
	text_emb = text_emb.unsqueeze(1)
	logits = model.forward(audio_embeds=audio_emb, text_embeds=text_emb)
	probs = torch.sigmoid(logits).squeeze(0).cpu().numpy()

	return {CLASSES[i]: round(float(probs[i]), 4) for i in range(len(CLASSES))}