# 导入包
from transformers import BertForSequenceClassification, AutoTokenizer

# 加载模型和分词器
model_path = "../rag_qa/models/bert_query_classifier"
model = BertForSequenceClassification.from_pretrained(model_path)  # 加载模型
tokenizer = AutoTokenizer.from_pretrained(model_path)

# 单个句子编码
text = "黑马的Java课程安排有哪些？"
inputs = tokenizer(text, return_tensors="pt") # 返回pytorch
print('inputs-->', inputs)

# 句子对编码
texts = ["黑马的Java课程安排有哪些？", "黑马的Python课程安排有哪些？"]
inputs = tokenizer(texts[0], texts[1], return_tensors="pt")
print("inputs-->", inputs)

texts = ["黑马的Java课程安排有哪些？", "黑马的Python有哪些？"]
inputs = tokenizer(texts, return_tensors="pt", padding=True)
print("inputs-->", inputs)

# 测试模型输出
text = "Java数据的基本类型" # 意图分类 专业知识
inputs = tokenizer(text, return_tensors="pt")
outputs = model(**inputs) # **表示将字典中的键值对作为参数传递给函数
print("outputs-->", outputs)

# 获取预测结果
result = outputs.logits.argmax(dim=-1)
print("result-->", result)