"""Quick browser demo: type Korean text → see emotion + VAD prediction.
Usage:
CUDA_VISIBLE_DEVICES=4 python scripts/demo_server.py \
--ckpt checkpoints/klue_teacher_clean_ctx2/best.pt \
--context_window 2 --port 8899
"""
from __future__ import annotations
import argparse
import json
import sys
from functools import partial
from pathlib import Path
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
import torch
from flask import Flask, jsonify, request, send_from_directory
from models.microalbert.config import EMOTION_LABELS, MicroAlbertConfig
from models.microalbert.teacher import KlueTeacherForEmotionVAD
from models.microalbert.tokenizer import MicroTokenizer
app = Flask(__name__)
model = None
tok = None
cfg = None
ctx_window = 0
EMOTION_COLORS = {
"neutral": "#9E9E9E", "joy": "#FFD54F", "laughter": "#FFF176",
"excitement": "#FF8A65", "agreement": "#81C784", "gratitude": "#A5D6A7",
"sadness": "#64B5F6", "crying": "#42A5F5", "sulk": "#CE93D8",
"apology": "#B0BEC5", "struggle": "#FFAB91", "anger": "#EF5350",
"refusal": "#F48FB1", "surprise": "#FFE082", "fluster": "#FFCC80",
"shy": "#F8BBD0",
}
HTML = """
Emotion Demo
MicroALBERT Emotion Demo
"""
@app.route("/")
def index():
html = HTML.replace("EMOTION_COLORS_JSON", json.dumps(EMOTION_COLORS))
return html
@app.route("/predict", methods=["POST"])
def predict():
data = request.json
text = data.get("text", "").strip()
context = data.get("context", [])
if not text:
return jsonify({"error": "empty text"}), 400
# Build input with context
if ctx_window > 0 and context:
parts = []
for i, prev in enumerate(context[-ctx_window:]):
parts.append(f"[OTHER] {prev}")
parts.append(f"[SELF] {text}")
full_text = " [SEP] ".join(parts)
elif ctx_window > 0:
full_text = f"[SELF] {text}"
else:
full_text = text
ids = tok.encode(full_text, cfg.max_seq_len)
input_ids = torch.tensor([ids], dtype=torch.long).to(next(model.parameters()).device)
attn = torch.ones_like(input_ids)
with torch.no_grad():
out = model(input_ids, attn)
logits = out["emotion_logits"][0]
probs = torch.softmax(logits, dim=-1).cpu().numpy()
vad = out["vad"][0].cpu().numpy().tolist()
pred_id = int(probs.argmax())
top5_ids = probs.argsort()[::-1][:5]
return jsonify({
"emotion": EMOTION_LABELS[pred_id],
"confidence": float(probs[pred_id]),
"vad": [round(v, 3) for v in vad],
"top5": [
{"emotion": EMOTION_LABELS[int(i)], "prob": round(float(probs[i]), 4)}
for i in top5_ids
],
})
def main():
global model, tok, cfg, ctx_window
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt", type=Path, required=True)
ap.add_argument("--context_window", type=int, default=0)
ap.add_argument("--port", type=int, default=8899)
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = ap.parse_args()
ctx_window = args.context_window
ckpt = torch.load(args.ckpt, map_location=args.device, weights_only=False)
cfg = MicroAlbertConfig(**ckpt["config"]) if isinstance(ckpt.get("config"), dict) else MicroAlbertConfig()
tok = MicroTokenizer.build(
save_dir=args.ckpt.parent / "tokenizer",
train_jsonl=Path("data/emotion/seed_val.jsonl"),
max_len=cfg.max_seq_len,
add_speaker_tokens=(ctx_window > 0),
)
model = KlueTeacherForEmotionVAD(
model_name=ckpt.get("model_name", "klue/roberta-base"),
num_emotions=cfg.num_emotions,
vad_dim=cfg.vad_dim,
dropout=0.0,
vad_head_hidden=cfg.vad_head_hidden,
attention_dropout=0.0,
).to(args.device)
if len(tok) != model.backbone.config.vocab_size:
model.backbone.resize_token_embeddings(len(tok))
model.load_state_dict(ckpt["model_state_dict"])
model.eval()
print(f"\nModel loaded: {args.ckpt}")
print(f"Context window: {ctx_window}")
print(f"Device: {args.device}")
print(f"Open http://localhost:{args.port}\n")
app.run(host="0.0.0.0", port=args.port, debug=False)
if __name__ == "__main__":
main()