Spaces:
Running
Running
Create intent_classifier.py
Browse files- intent_classifier.py +17 -0
intent_classifier.py
ADDED
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from transformers import AutoModelForSequenceClassification, AutoTokenizer
|
2 |
+
import torch
|
3 |
+
|
4 |
+
class IntentClassifier:
|
5 |
+
def __init__(self):
|
6 |
+
# Define the intents directly
|
7 |
+
self.intents = {"SQL Query": "database_query", "ask a question": "product_description"}
|
8 |
+
|
9 |
+
def classify(self, queryType):
|
10 |
+
# Map the dropdown selection to the appropriate intent
|
11 |
+
if queryType == "sql_query":
|
12 |
+
return self.intents["sql_query"], 1.0 # Full confidence for SQL queries
|
13 |
+
elif queryType == "ask_question":
|
14 |
+
return self.intents["ask_question"], 1.0 # Full confidence for product description
|
15 |
+
|
16 |
+
# Handle unexpected input (optional)
|
17 |
+
return None, 0.0 # No intent matched
|