Update app.py

app.py

@@ -1,91 +1,100 @@
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
+import torch
 import joblib
+import json
 import pandas as pd
-from typing import Optional
+import numpy as np
+import re
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, BertConfig
+from huggingface_hub import hf_hub_download
+import os
+from safetensors.torch import load_file
 import uvicorn
 
-# Load the model and label encoder
-model, label_encoder = joblib.load('queue_prediction_model.joblib')
+# Initialize FastAPI app
+app = FastAPI(title="Ticket Field Predictor API")
 
-app = FastAPI(title="Queue Prediction API")
+# Get Hugging Face token (set via Space Secrets as HF_TOKEN)
+HUGGINGFACE_TOKEN = os.getenv("HF_TOKEN")
 
+# Model Repository on Hugging Face
+REPO_ID = "AtaishNehra/ticket_field_predictor_model"
+
+# Define features
+FEATURES = [
+    "Company", "Contract", "Priority", "Issue Type",
+    "Service Level Agreement", "Sub-Issue Type", "Queue"
+]
+
+# Load label encoders
+encoder_path = hf_hub_download(repo_id=REPO_ID, filename="label_encoders.joblib", token=HUGGINGFACE_TOKEN)
+label_encoders = joblib.load(encoder_path)
+
+# Load tokenizers
+tokenizers = {feature: AutoTokenizer.from_pretrained("bert-base-uncased") for feature in FEATURES}
+
+# Function to load models
+def load_model(feature):
+    config_filename = f"{feature}/config.json"
+    model_filename = f"{feature}/model.safetensors"
+
+    # Download and load the configuration
+    config_path = hf_hub_download(repo_id=REPO_ID, filename=config_filename, token=HUGGINGFACE_TOKEN)
+    with open(config_path, "r") as f:
+        config = json.load(f)
+
+    # Download and load the model state dict
+    model_path = hf_hub_download(repo_id=REPO_ID, filename=model_filename, token=HUGGINGFACE_TOKEN)
+    state_dict = load_file(model_path, device="cpu")
+
+    # Ensure configuration has a model_type
+    config.setdefault("model_type", "bert")
+    model_config = BertConfig.from_dict(config)
+    model = AutoModelForSequenceClassification.from_config(model_config)
+
+    # Load state dict without strict checking
+    model.load_state_dict(state_dict, strict=False)
+    model.eval()
+    return model
+
+# Load models for each feature
+models = {feature: load_model(feature) for feature in FEATURES}
+
+# Input model for API request
 class TicketInput(BaseModel):
     title: str
     description: str
 
-class TicketPrediction(BaseModel):
-    predicted_queue: str
-    confidence: float
-    top_3_predictions: list
-
-# Copy your helper functions
-def clean_text(text):
-    if pd.isna(text):
-        return ""
-    text = str(text)
-    text = re.sub(r'\S+@\S+', 'EMAIL', text)
-    text = re.sub(r'\b(?:\d{1,3}\.){3}\d{1,3}\b', 'IP_ADDRESS', text)
-    text = re.sub(r'\d{1,2}/\d{1,2}/\d{4}', 'DATE', text)
-    text = re.sub(r'\b\d+\b', 'NUM', text)
-    text = text.lower()
-    return text
-
-def extract_features(df):
-    # Copy your extract_features function here
-    # (The same function from your original code)
-    pass
-
-@app.post("/predict", response_model=TicketPrediction)
-async def predict(ticket: TicketInput):
+# Prediction function
+@torch.no_grad()
+def predict(title: str, description: str):
+    inputs = f"{title} {description}"
+    predictions = {}
+
+    for feature in FEATURES:
+        tokenizer = tokenizers[feature]
+        model = models[feature]
+        label_encoder = label_encoders[feature]
+
+        tokens = tokenizer(inputs, return_tensors="pt", padding=True, truncation=True)
+        output = model(**tokens).logits
+        pred_idx = torch.argmax(output, dim=1).item()
+
+        # Decode prediction
+        predictions[feature] = label_encoder.inverse_transform([pred_idx])[0]
+
+    return predictions
+
+@app.post("/predict")
+async def api_predict(ticket: TicketInput):
+    """API endpoint for external HTTP requests"""
     try:
-        # Create a DataFrame with the input data
-        sample = pd.DataFrame({
-            'Description': [ticket.description],
-            'Title': [ticket.title],
-            'Priority': ['Normal']
-        })
-
-        # Process features
-        sample_processed = extract_features(sample)
-        
-        feature_columns = [
-            'cleaned_description', 'cleaned_title', 'has_ups', 'has_battery',
-            'has_problem', 'has_ip', 'has_serial', 'is_security_related',
-            'is_network_related', 'is_programming_related', 'is_support_related',
-            'description_length', 'title_length', 'word_count', 'is_high_priority'
-        ]
-
-        X = sample_processed[feature_columns]
-
-        # Make prediction
-        prediction = model.predict(X)
-        probabilities = model.predict_proba(X)
-
-        predicted_queue = label_encoder.inverse_transform(prediction)[0]
-        confidence = float(np.max(probabilities[0]))
-
-        # Get top 3 predictions
-        top_3_idx = np.argsort(probabilities[0])[-3:][::-1]
-        top_3_predictions = [
-            {
-                "queue": queue,
-                "probability": float(prob)
-            }
-            for queue, prob in zip(
-                label_encoder.inverse_transform(top_3_idx),
-                probabilities[0][top_3_idx]
-            )
-        ]
-
-        return {
-            "predicted_queue": predicted_queue,
-            "confidence": confidence,
-            "top_3_predictions": top_3_predictions
-        }
+        predictions = predict(ticket.title, ticket.description)
+        return {"predictions": predictions}
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    uvicorn.run(app, host="0.0.0.0", port=7860)