cnmoro / granite-question-classifier

A finetuned version of ibm-granite/granite-embedding-30m-english.

The goal is to classify questions into "Directed" or "Generic".

If a question is not directed, we would change the actions we perform on a RAG pipeline (if it is generic, semantic search wouldn't be useful directly; e.g. asking for a summary).

(Class 0 is Generic; Class 1 is Directed)

The accuracy achieved during training was 94%.

This model is designed to be an upgrade of the previous model: https://huggingface.co/cnmoro/bert-tiny-question-classifier

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