import paddle
import numpy as np
from paddlenlp.transformers import ErnieTokenizer, ErnieForSequenceClassification
import paddle.nn.functional as F

# 类别名称列表
labels = [
    "天气查询", "通用对话", "页面切换", "日计划数量查询", "周计划数量查询",
    "日计划作业内容", "周计划作业内容", "施工人数", "作业考勤人数", "知识问答"
]

# 加载模型和tokenizer
model = ErnieForSequenceClassification.from_pretrained(R"E:\workingSpace\PycharmProjects\Intention_dev\ernie\output\checkpoint-4160")  # 使用文本分类模型
tokenizer = ErnieTokenizer.from_pretrained(R"E:\workingSpace\PycharmProjects\Intention_dev\ernie\output\checkpoint-4160")

# 创建输入示例
text = "宇宙中发现的第一个脉冲星是由谁发现的？"
inputs = tokenizer(text, max_length=256, truncation=True, padding='max_length', return_tensors="pd")

# 将输入数据转化为 Paddle tensor 格式
input_ids = paddle.to_tensor(inputs["input_ids"])

# 进行推理
model.eval()  # 确保模型在推理模式
logits = model(input_ids)  # 模型推理得到logits

# 使用 Softmax 转换 logits 为概率
probabilities = F.softmax(logits, axis=-1)  # 归一化 logits 得到概率分布

# 获取最大概率的标签（整个句子的意图）
max_prob_idx = np.argmax(probabilities.numpy(), axis=-1)  # 获取最大概率的标签
max_prob_value = np.max(probabilities.numpy(), axis=-1)  # 获取最大概率值

# 根据预测的标签索引映射到类别名称
predicted_label = labels[max_prob_idx[0]]  # 根据索引获取对应的标签
predicted_probability = max_prob_value[0]  # 获取最大概率值
print(predicted_label, predicted_probability)