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許可證。

📋 模型信息