Spaces:
Runtime error
Runtime error
File size: 2,782 Bytes
5f946b0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
import torch
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
import os
import json
# Get project root directory
PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
def load_model():
print("📦 Loading pre-trained intent classification model...")
model_name = "distilbert-base-uncased"
tokenizer = DistilBertTokenizerFast.from_pretrained(model_name)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Define intent labels
intent_labels = [
"summary", "comparison", "trend", "anomaly", "forecast"
]
num_labels = len(intent_labels)
# Create label mapping
label_mapping = {label: idx for idx, label in enumerate(intent_labels)}
# Load model with our number of labels
model = DistilBertForSequenceClassification.from_pretrained(
model_name, num_labels=num_labels
)
model = model.to(device)
model.eval()
return model, tokenizer, device, label_mapping
def classify_intent(question, model, tokenizer, device, label_mapping):
# Tokenize input
inputs = tokenizer(
question,
return_tensors="pt",
truncation=True,
padding=True,
max_length=128
).to(device)
# Get prediction
with torch.no_grad():
outputs = model(**inputs)
predicted_class_id = outputs.logits.argmax().item()
# Convert back to label
id2label = {v: k for k, v in label_mapping.items()}
intent = id2label[predicted_class_id]
return intent
if __name__ == "__main__":
# Load the model
model, tokenizer, device, label_mapping = load_model()
# Save the model and label mapping
output_dir = os.path.join(PROJECT_ROOT, "model_intent_classifier")
print(f"💾 Saving model to {output_dir}")
os.makedirs(output_dir, exist_ok=True)
model.save_pretrained(output_dir)
tokenizer.save_pretrained(output_dir)
# Save label mapping
with open(os.path.join(output_dir, "label_mapping.json"), "w") as f:
json.dump(label_mapping, f)
print(f"✅ Model successfully saved to {output_dir}")
# Example usage
test_questions = [
"What is the total sales amount for each product category?",
"Compare sales between March and April",
"Show me the sales trend over the last 6 months",
"Which products have unusual sales patterns?",
"What will be the sales forecast for next month?"
]
print("\nTesting intent classification:")
for question in test_questions:
intent = classify_intent(question, model, tokenizer, device, label_mapping)
print(f"Question: {question}")
print(f"Predicted intent: {intent}\n")
|