Initial commit

dockerfile

@@ -0,0 +1,14 @@
+# Use the official Python 3.10.9 image
+FROM python:3.10.14
+
+# Copy the current directory contents into the container at .
+COPY . .
+
+# Set the working directory to /
+WORKDIR /
+
+# Install requirements.txt 
+RUN pip install --no-cache-dir --upgrade -r /requirements.txt
+
+# Start the FastAPI app on port 7860, the default port expected by Spaces
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

routers/distilbert.py

@@ -0,0 +1,14 @@
+from fastapi import APIRouter
+from src.distilbert import Distilbert
+ 
+
+MODEL_PATH = "SebaK13/DistilBERT-finetuned-customer-queries-balanced"
+predict = APIRouter(prefix="/predict", tags=["predict"])
+distilbert_service = Distilbert(model_path=MODEL_PATH)
+
+
+
+@predict.get("/predict")
+def predict(query: str):
+    output = distilbert_service.predict_query_type(query)
+    return {"label": output}

src/distilbert.py

@@ -0,0 +1,56 @@
+import json
+import pandas as pd
+from transformers import pipeline
+
+class Distilbert:
+    """
+    Class for distilBERT
+    """
+    def __init__(self, model_path:str):
+        self.pipe = pipeline("text-classification", model=model_path)
+    
+    
+    def predict_query_type(self, custom_query):
+        """
+        Predict the type of a given custom query using a pre-trained classifier.
+
+        Parameters:
+        custom_query (str): The input query for which the type needs to be predicted.
+
+        Returns:
+        str: The predicted label for the input query.
+        """
+        # Get predictions from the classifier
+        preds = self.pipe(custom_query, top_k=None)
+        
+        # Get the list of labels from the balanced dataset
+        labels_dict = self.label_dict()
+        labels = labels_dict.keys()
+        
+        # Create a DataFrame from the predictions
+        preds_df = pd.DataFrame(preds)
+        
+        # Process the labels to remove 'LABEL_' prefix and convert to integer
+        preds_df['label'] = preds_df['label'].str.replace('LABEL_', '')
+        preds_df['label'] = pd.to_numeric(preds_df['label'], errors='coerce').astype('Int64')
+        preds_df = preds_df.sort_values('label').reset_index(drop=True)
+        
+        # Find the index of the maximum score
+        max_score_index = preds_df['score'].idxmax()
+        
+        # Return the label corresponding to the maximum score
+        return labels[max_score_index]
+
+    def load_lables(self):
+        """
+        This function reads the 'labels.json' file, which contains a dictionary mapping integers to string labels. 
+        The function then returns this dictionary.
+
+        Returns:
+            dict: label dictionary with key as label string value and dictionary value as int
+        """
+        label_dict = {}
+        with open('labels.json') as json_file:
+            label_dict = json.load(json_file)
+        return label_dict
+                             

src/labels.json

@@ -0,0 +1,32 @@
+{
+    "Airport Services": 0,
+    "Baggage Policies": 1,
+    "Business Travel": 2,
+    "COVID-19 Policies": 3,
+    "Check-in Procedures": 4,
+    "Child and Infant Travel": 5,
+    "Complaints and Feedback": 6,
+    "Customer Account Issues": 7,
+    "Duty-Free Shopping": 8,
+    "Flight Bookings": 9,
+    "Flight Changes": 10,
+    "Flight Status": 11,
+    "Frequent Flyer Miles": 12,
+    "Group Bookings": 13,
+    "In-flight Services": 14,
+    "Lost and Found": 15,
+    "Loyalty Programs": 16,
+    "Mobile App Issues": 17,
+    "Partnerships and Alliances": 18,
+    "Payment Issues": 19,
+    "Pet Travel": 20,
+    "Promotions and Discounts": 21,
+    "Refunds and Compensation": 22,
+    "Seat Selection": 23,
+    "Special Assistance": 24,
+    "Travel Documentation": 25,
+    "Travel Insurance": 26,
+    "Travel Restrictions": 27,
+    "Travel Vouchers": 28,
+    "Weather-related Disruptions": 29
+}