Granite Question Classifier

🚀 花岗岩问题分类器

这是一个微调版的模型，旨在将问题分类为“定向”或“通用”，可优化RAG管道的操作，训练准确率达94%。

🚀 快速开始

这是ibm-granite/granite-embedding-30m-english的微调版本。其目标是将问题分类为“定向”或“通用”。如果一个问题不是定向的，我们会改变在RAG管道上执行的操作（例如，如果是通用问题，语义搜索可能不会直接起作用，如请求摘要）。（类别0是通用；类别1是定向）。训练期间达到的准确率为94%。

该模型旨在作为之前模型的升级版：点击查看

💻 使用示例

基础用法

import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_id = "cnmoro/granite-question-classifier"
model = AutoModelForSequenceClassification.from_pretrained(model_id)
tokenizer = AutoTokenizer.from_pretrained(model_id)
model.eval()

def predict_question_category(question):  
    inputs = tokenizer.encode_plus(
        question,
        add_special_tokens=True,
        max_length=512,
        return_tensors="pt",
        truncation=True
    )
    
    input_ids = inputs["input_ids"]
    attention_mask = inputs["attention_mask"]
    
    with torch.no_grad():
        outputs = model(input_ids, attention_mask=attention_mask)
        logits = outputs.logits.squeeze(-1)
        print(logits)
        prediction = (logits > 0).float().item()
    
    # Map prediction to category
    return "directed" if prediction == 1.0 else "generic"

predict_question_category("Qual o resumo do texto?") # generic
predict_question_category("Qual foi a crítica que o autor recebeu do jornal, em relação a sua opinião?") # directed

📄 许可证

本项目采用MIT许可证。

📋 模型信息