from transformers import AutoTokenizer, AutoModelForCausalLM

# 加载模型
model = AutoModelForCausalLM.from_pretrained(
    "KE-Team/KE-SemanticVAD").to('cuda')
tokenizer = AutoTokenizer.from_pretrained("KE-Team/KE-SemanticVAD")


# System Prompt
AGNET_SPKING_SYS='# Role\n你是人机实时交互的**用户行为分析**模块，你将收到包含部分历史信息的 Human 和 Agent 最新实时对话记录 (Dialog)\n\n# 任务\n当前【Agent正在发言】，在此过程中，你需要基于对话分析 Human 的意图属于 <打断> 还是 <附和>\n\n# 输出\n不要有多余的分析，仅严格输出以下二者之一: <打断> 或 <附和>\n\n# 判断标准\n## <打断> 的情况\nHuman 行为: 试图抢夺话题主导权\n特征包括:\n- 提供新概念/词汇/判断（如命名、定性、对比）\n- 提出问题或异议\n- 引入与当前话题无关的新话题\n\n## <附和> 的情况\nHuman 行为: 赞同 Agent, 期望 Agent 继续说\n特征包括:\n- 使用零内容反馈（嗯/啊/对）\n- 机械重复 Agent 中的原词/同义词\n- 表达简单的确认或同意（如“是的”、“没错”）\n'
HUMAN_SPKING_SYS = '# Role\n你是人机实时交互的**用户行为分析**模块，你将收到包含部分历史信息的 Human 和 Agent 最新实时对话记录 (Dialog)\n\n# 任务\n当前【Human正在发言】，你需要基于对话判断 Human 是否已经完成发言\n\n# 输出\n严格输出以下二者之一: <完成> 或 <未完>\n\n# 判断标准\n## <完成> 的情况\nHuman 发言语义完整，说话很可能已经结束\n- 发言包含完整命题（如明确提问/请求/结论）\n- 出现结束性标记词（"好了"/"你觉得呢？"）\n\n## <未完> 的情况\nHuman 发言语义不完整，仍然可能继续说话\n- 语句末尾含连接词（"而且"/"不过"/"然后"）\n- 用户发言中夹杂思考词（"呃..."/"嗯..."）\n'
SYS_MAP = dict(
    agent = AGNET_SPKING_SYS,
    human = HUMAN_SPKING_SYS
    )

# Dialog Format
def dia_format(x):
    cur_spk = x['speaker']
    system = SYS_MAP[cur_spk]
    if cur_spk == 'agent':
        u1 = x['history']['query']
        a1 = x['history']['answer']
        u2 = x['query']
        dialog = f"# Dialog\nHuman[历史]:{u1}\nAgent:[实时]:{a1}\nHuman[实时]:{u2}\n"
    elif cur_spk == 'human':
        if len(x['history']) <= 1:
            u2 = x['query']
            dialog = f"# Dialog\nHuman[实时]:{u2}\n"
        else:
            u1 = x['history']['query']
            a1 = x['history']['answer']
            u2 = x['query']
            dialog = f"# Dialog\nHuman[历史]:{u1}\nAgent:[历史]:{a1}\nHuman[实时]:{u2}\n"
    else:
        raise ValueError('current speaker should in agent or human')
    return [{'role': 'system', 'content':system}, {'role': 'user', 'content': dialog}]


# 数据样例
example = {
      "speaker": "agent",
      "query": "那具体是怎么实现的？比如说，如",
      "history": {
          "query": "怎么把人工智能技术用在虚拟现实开发上呢？",
          "answer": "将人工智能技术应用到虚拟现实开发中，可以通过智能算法来提升用户体验，比如使用机器学习来创建更真实的虚拟角色"
          }
      }
messages = dia_format(example)
text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=64
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]
response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(f"检测结果: {response}")  # -> <打断>

Citation

Please cite our Hugging-Face when using our code, data or model.

Bibtext

@misc{KE-SemanticVAD,
      author = {KE-TEAM},
      title = {KE-SemanticVAD},
      year = {2025},
      publisher = {Hugging Face},
      url = {https://huggingface.co/KE-Team/KE-SemanticVAD}
}