|
--- |
|
license: mit |
|
datasets: |
|
- cnmoro/QuestionClassification-v2 |
|
language: |
|
- en |
|
- pt |
|
tags: |
|
- classification |
|
- questioning |
|
- directed |
|
- generic |
|
pipeline_tag: text-classification |
|
base_model: |
|
- ibm-granite/granite-embedding-30m-english |
|
library_name: transformers |
|
--- |
|
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 |
|
|
|
```python |
|
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 |
|
``` |