Update app.py

app.py

@@ -1,11 +1,20 @@
 import gradio as gr
 import os
+import re
 import torch
+import joblib
+import pandas as pd
+import nltk
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoModelForSeq2SeqLM
+from nltk.corpus import stopwords
+from nltk.stem import PorterStemmer
 
-print("Booting up PredictiX Inference API (Direct Model Load)...")
+print("Booting up PredictiX Inference API...")
 
 hf_token = os.environ.get("HF_TOKEN")
+nltk.download("stopwords", quiet=True)
+_stop_words = set(stopwords.words("english"))
+_stemmer = PorterStemmer()
 
 # 1. Load Ticket Categorization Model
 try:
@@ -13,79 +22,157 @@ try:
     cat_id = cat_path if os.path.exists(cat_path) else "Dinusha-Ekanayake/predictix-ticket_categorization_model"
     cat_tokenizer = AutoTokenizer.from_pretrained(cat_id, token=hf_token)
     cat_model = AutoModelForSequenceClassification.from_pretrained(cat_id, token=hf_token)
+    cat_model.eval()
+    print("Categorization model loaded.")
 except Exception as e:
     cat_model = None
     print(f"Failed to load categorizer: {e}")
 
-# 2. Load Ticket Summarization Model
+# 2. Load Ticket Priority Model
+try:
+    pri_path = "./ticket_priority_classification_model"
+    model_file = os.path.join(pri_path, "xgboost_vehicle_priority_model_3class.pkl")
+    encoder_file = os.path.join(pri_path, "priority_label_encoder.pkl")
+    pri_model = joblib.load(model_file)
+    pri_encoder = joblib.load(encoder_file)
+    print("Priority model loaded.")
+except Exception as e:
+    pri_model = None
+    pri_encoder = None
+    print(f"Failed to load priority model: {e}")
+
+# 3. Load Ticket Summarization Model
 try:
     sum_path = "./predictix-ticket_summarization_model"
     ts_id = sum_path if os.path.exists(sum_path) else "Dinusha-Ekanayake/predictix-ticket_summarization_model"
     ts_tokenizer = AutoTokenizer.from_pretrained(ts_id, token=hf_token)
     ts_model = AutoModelForSeq2SeqLM.from_pretrained(ts_id, token=hf_token)
+    print("Ticket summarization model loaded.")
 except Exception as e:
     ts_model = None
     print(f"Failed to load ticket summarizer: {e}")
 
-# 3. Load Asset Summarization Model
+# 4. Load Asset Summarization Model
 try:
     as_id = "Dinusha-Ekanayake/predictix-asset_summarization_model"
     as_tokenizer = AutoTokenizer.from_pretrained(as_id, token=hf_token)
     as_model = AutoModelForSeq2SeqLM.from_pretrained(as_id, token=hf_token)
+    print("Asset summarization model loaded.")
 except Exception as e:
     as_model = None
     print(f"Failed to load asset summarizer: {e}")
 
+
+# --- Priority helpers ---
+def _clean_text(text):
+    text = str(text).lower()
+    text = text.replace("pls", "please").replace("asap", "as soon as possible")
+    text = re.sub(r"[^a-zA-Z\s]", " ", text)
+    text = re.sub(r"\s+", " ", text).strip()
+    important = {"ac", "abs", "rpm"}
+    out = []
+    for tok in text.split():
+        if tok in important:
+            out.append(tok)
+        elif tok not in _stop_words and len(tok) > 2:
+            out.append(_stemmer.stem(tok))
+    return " ".join(out)
+
+_HIGH_KEYWORDS = {
+    "fire", "smoke", "explosion", "fuel leak", "brake failure", "no brakes",
+    "engine seized", "total failure", "accident", "crash", "rollover",
+    "unsafe to drive", "cannot drive", "vehicle stopped", "complete breakdown",
+    "coolant leak", "overheating", "electrical fire", "cannot start",
+}
+_LOW_KEYWORDS = {
+    "scratch", "dent", "cosmetic", "minor", "small crack", "paint", "sticker",
+    "mirror", "wiper", "next service", "no rush", "low urgency",
+    "seat cover", "floor mat", "trim", "logo", "decal", "cleaning",
+}
+
+
 # --- API Functions ---
 def categorize(text):
-    if not cat_model: return {"error": "Categorization model not loaded."}
+    if not cat_model:
+        return {"error": "Categorization model not loaded."}
     inputs = cat_tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
     with torch.no_grad():
-        logits = cat_model(**inputs).logits
-    # Get highest score
-    probs = torch.nn.functional.softmax(logits, dim=-1)[0]
-    best_idx = torch.argmax(probs).item()
-    label = cat_model.config.id2label[best_idx]
-    score = probs[best_idx].item()
-    return [{"label": label, "score": score}]
+        probs = torch.softmax(cat_model(**inputs).logits, dim=-1)[0]
+    scores = [
+        {"label": cat_model.config.id2label[i], "score": round(float(s), 4)}
+        for i, s in enumerate(probs.tolist())
+    ]
+    scores.sort(key=lambda x: x["score"], reverse=True)
+    return scores
+
+
+def prioritize(text):
+    if not pri_model:
+        return {"error": "Priority model not loaded."}
+    lower = text.lower()
+    if any(kw in lower for kw in _HIGH_KEYWORDS):
+        return {"priority": "High"}
+    if any(kw in lower for kw in _LOW_KEYWORDS):
+        return {"priority": "Low"}
+    row = pd.DataFrame([{
+        "combined_text": _clean_text(text),
+        "vehicle_type": "Truck",
+        "issue_category": "Engine",
+        "sensor_alert": "Check engine light",
+        "operating_environment": "Urban",
+        "weather_condition": "Normal",
+        "vehicle_age": 5,
+        "mileage_km": 100000,
+        "downtime_hours": 0.0,
+        "maintenance_overdue_days": 0,
+        "previous_failures": 0,
+    }])
+    pred = pri_model.predict(row)[0]
+    label = pri_encoder.inverse_transform([pred])[0]
+    return {"priority": label}
+
 
 def summarize_ticket(text):
-    if not ts_model: return {"error": "Ticket Summarization model not loaded."}
+    if not ts_model:
+        return {"error": "Ticket Summarization model not loaded."}
     inputs = ts_tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
     with torch.no_grad():
         outputs = ts_model.generate(**inputs, min_length=15, max_length=150, num_beams=4, early_stopping=True)
-    summary = ts_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"summary": summary}
+    return {"summary": ts_tokenizer.decode(outputs[0], skip_special_tokens=True)}
+
 
 def summarize_asset(text):
-    if not as_model: return {"error": "Asset Summarization model not loaded."}
+    if not as_model:
+        return {"error": "Asset Summarization model not loaded."}
     inputs = as_tokenizer(text, return_tensors="pt", truncation=True, max_length=1024)
     with torch.no_grad():
         outputs = as_model.generate(**inputs, min_length=20, max_length=150, num_beams=4, early_stopping=True)
-    summary = as_tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return {"summary": summary}
+    return {"summary": as_tokenizer.decode(outputs[0], skip_special_tokens=True)}
+
 
-# --- Server API Interface ---
+# --- Gradio UI ---
 with gr.Blocks(title="PredictiX API") as demo:
-    gr.Markdown("# PredictiX Internal Inference Server 🚀")
-    
+    gr.Markdown("# PredictiX Internal Inference Server")
+
     with gr.Tab("Ticket Categorization"):
         cat_in = gr.Textbox(label="Ticket Title & Description")
         cat_out = gr.JSON(label="Categorization Result")
-        cat_btn = gr.Button("Categorize")
-        cat_btn.click(categorize, inputs=cat_in, outputs=cat_out, api_name="categorize")
-        
+        gr.Button("Categorize").click(categorize, inputs=cat_in, outputs=cat_out, api_name="categorize")
+
+    with gr.Tab("Ticket Priority"):
+        pri_in = gr.Textbox(label="Ticket Title & Description")
+        pri_out = gr.JSON(label="Priority Result")
+        gr.Button("Prioritize").click(prioritize, inputs=pri_in, outputs=pri_out, api_name="prioritize")
+
     with gr.Tab("Ticket Summarization"):
         ts_in = gr.Textbox(label="Ticket Details")
         ts_out = gr.JSON(label="Summary")
-        ts_btn = gr.Button("Summarize Ticket")
-        ts_btn.click(summarize_ticket, inputs=ts_in, outputs=ts_out, api_name="summarize_ticket")
-        
+        gr.Button("Summarize Ticket").click(summarize_ticket, inputs=ts_in, outputs=ts_out, api_name="summarize_ticket")
+
     with gr.Tab("Asset Summarization"):
         as_in = gr.Textbox(label="Asset Details")
         as_out = gr.JSON(label="Summary")
-        as_btn = gr.Button("Summarize Asset")
-        as_btn.click(summarize_asset, inputs=as_in, outputs=as_out, api_name="summarize_asset")
+        gr.Button("Summarize Asset").click(summarize_asset, inputs=as_in, outputs=as_out, api_name="summarize_asset")
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()