Scicom-intl/Malaysian-Turn-Detector-Qwen3-0.6B

import torch
import math
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelForCausalLM

model_id = "Scicom-intl/Malaysian-Turn-Detector-Qwen3-0.6B"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16).cuda().eval()

IM_END_ID = tokenizer.convert_tokens_to_ids("<|im_end|>")

def get_turn_end_prob(text):
    """Returns probability that the speaker's turn is complete."""
    # Strip trailing <|im_end|> so the model predicts whether to emit it
    if text.endswith("<|im_end|>"):
        text = text[:-len("<|im_end|>")]
    inputs = tokenizer(text, return_tensors="pt").to("cuda")
    with torch.no_grad():
        logits = model(**inputs).logits
    prob = F.softmax(logits[0, -1], dim=-1)[IM_END_ID].item()
    return prob

# Complete turn - should be high probability
text = "<|im_start|>user\nHello, saya nak tanya pasal bil saya.<|im_end|>\n<|im_start|>assistant\nBoleh, sila berikan nombor akaun anda."
prob = get_turn_end_prob(text)
print(f"P(turn complete) = {prob:.4f}")  # ~0.74

# Incomplete turn - should be low probability
text = "<|im_start|>user\nHello, saya nak tanya pasal bil saya.<|im_end|>\n<|im_start|>assistant\nBoleh, sila berikan nombor"
prob = get_turn_end_prob(text)
print(f"P(turn complete) = {prob:.4f}")  # ~0.00