Update README.md

README.md

@@ -1,3 +1,60 @@
----
-license: mit
----
+---
+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%.
+
+```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
+```