chidamnat2002/intent_classifier

1import onnxruntime as ort2from transformers import AutoTokenizer3import numpy as np4import requests5import os6 7VERSION = "v0.1.1"8class IntentClassifier:9    def __init__(self):10        self.id2label = {0: 'information_intent',11                         1: 'yelp_intent',12                         2: 'navigation_intent',13                         3: 'travel_intent',14                         4: 'purchase_intent',15                         5: 'weather_intent',16                         6: 'translation_intent',17                         7: 'unknown'}18        self.label2id = {label: id for id, label in self.id2label.items()}19 20        self.tokenizer = AutoTokenizer.from_pretrained("Mozilla/mobilebert-uncased-finetuned-LoRA-intent-classifier")21        22        model_url = f"https://huggingface.co/Mozilla/mobilebert-uncased-finetuned-LoRA-intent-classifier/resolve/{VERSION}/onnx/model_quantized.onnx"23        model_dir_path = "models"24        model_path = f"{model_dir_path}/mobilebert-uncased-finetuned-LoRA-intent-classifier_model_quantized.onnx"25        if not os.path.exists(model_dir_path):26            os.makedirs(model_dir_path)27        if not os.path.exists(model_path):28            print("Downloading ONNX model...")29            response = requests.get(model_url)30            with open(model_path, "wb") as f:31                f.write(response.content)32            print("ONNX model downloaded.")33 34        # Load the ONNX model35        self.ort_session = ort.InferenceSession(model_path)36 37    def find_intent(self, sequence, verbose=False):38        inputs = self.tokenizer(sequence,39                                return_tensors="np",  # ONNX requires inputs in NumPy format40                                padding="max_length",  # Pad to max length41                                truncation=True,       # Truncate if the text is too long42                                max_length=64)43 44        # Convert inputs to NumPy arrays45        onnx_inputs = {k: v for k, v in inputs.items()}46 47        # Run the ONNX model48        logits = self.ort_session.run(None, onnx_inputs)[0]49 50        # Get the prediction51        prediction = np.argmax(logits, axis=1)[0]52        probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)53        rounded_probabilities = np.round(probabilities, decimals=3)54 55        pred_result = self.id2label[prediction]56        proba_result = dict(zip(self.label2id.keys(), rounded_probabilities[0].tolist()))57        58        if verbose:59            print(sequence + " -> " + pred_result)60            print(proba_result, "\n")61        62        return pred_result, proba_result63 64def main():65    text_list = [66        'floor repair cost',67        'pet store near me',68        'who is the us president',69        'italian food',70        'sandwiches for lunch',71        "cheese burger cost",72        "What is the weather today",73        "what is the capital of usa",74        "cruise trip to carribean",75    ]76    cls = IntentClassifier()77    for sequence in text_list:78        cls.find_intent(sequence, verbose=True)79 80if __name__ == '__main__':81    main()82