app.py

import os
os.environ["TF_KERAS"] = "1"  # ✅ Forces Transformers to use TensorFlow-Keras

import streamlit as st
import tensorflow as tf
from transformers import TFDistilBertForSequenceClassification, DistilBertTokenizer
from huggingface_hub import snapshot_download


# ✅ Student Information
My_info = "Student ID: 6319250G, Name: Aung Hlaing Tun"

# ✅ Define Hugging Face Model Repo
MODEL_REPO_ID = "ZAM-ITI-110/Dist_Bert"

# ✅ Load Model & Tokenizer from Hugging Face
@st.cache_resource
def load_model(repo_id):
    """Download and load the TensorFlow model and tokenizer."""

    # 📌 Define cache directory
    cache_dir = "hf_models"

    # 📌 Ensure directory exists
    os.makedirs(cache_dir, exist_ok=True)

    # 📌 Download model from Hugging Face (if not cached)
    try:
        st.write("🔄 Downloading Model from Hugging Face...")
        download_dir = snapshot_download(repo_id, cache_dir=cache_dir, local_files_only=False)
        st.write(f"✅ Model downloaded at: {download_dir}")

        # 📌 Debugging: Print contents of model directory
        st.write("📂 Model Directory Contents:", os.listdir(download_dir))

    except Exception as e:
        st.error(f"❌ Error downloading model: {str(e)}")
        return None, None
        
    # 📌 Load Model and Tokenizer (Remove `from_tf=True`)
    try:
        model = TFDistilBertForSequenceClassification.from_pretrained(download_dir)  # ✅ FIX: Removed `from_tf=True`
        tokenizer = DistilBertTokenizer.from_pretrained(download_dir)
    except Exception as e:
        st.error(f"❌ Error loading model: {str(e)}")
        model, tokenizer = None, None
        
    return model, tokenizer

# ✅ Load Model
st.write("🔄 Loading Model...")
model, tokenizer = load_model(MODEL_REPO_ID)

if model is None or tokenizer is None:
    st.error("❌ Model loading failed. Please check logs above.")
    st.stop()

st.success("✅ Model Loaded Successfully!")

# ✅ Prediction Function
def predict_team_and_email(text):
    """Predict the team and corresponding email for a given ticket description."""
    inputs = tokenizer(text, return_tensors="tf", truncation=True, padding=True)
    
    # Perform inference
    preds = model(inputs["input_ids"])
    predicted_label = tf.argmax(preds.logits, axis=1).numpy()[0]

    # Mapping Labels to Team Names and Emails
    label_mapping = {
        0: "Code Review Team",
        1: "Functional Team",
        2: "Infrastructure Team",
        3: "Performance Team",
        4: "Security Team"
    }
    email_mapping = {
        0: "codereview.team@company.com",
        1: "functional.team@company.com",
        2: "infra.team@company.com",
        3: "performance.team@company.com",
        4: "security.team@company.com"
    }
    
    return label_mapping.get(predicted_label, "Unknown"), email_mapping.get(predicted_label, "Unknown")

# ✅ Title and Student Info
st.title("📩 Development of an AI Ticket Classifier Model Using DistilBERT")
st.markdown(f"*{My_info}*")

st.markdown("""
### 🔍 About this App  
- This system predicts the appropriate **team assignment** and **contact email** based on the ticket description.  
- Simply enter up to **6 ticket descriptions**, and the AI will classify them accordingly.
""")

# ✅ Layout for Ticket Inputs and Predictions
col1, col2, col3 = st.columns(3)

# ✅ Dictionary to store ticket descriptions
ticket_inputs = {}

for i in range(6):  # Up to 6 tickets
    with [col1, col2, col3][i % 3]:  # Distribute across 3 columns
        ticket_inputs[f"ticket_{i+1}"] = st.text_area(f"🎟️ Ticket {i+1}", placeholder=f"Enter ticket description {i+1}...")

# ✅ Prediction Button
if st.button("🔮 Predict All Tickets"):
    st.subheader("📌 Prediction Results:")
    for i, (key, text) in enumerate(ticket_inputs.items()):
        if text.strip():
            team, email = predict_team_and_email(text)
            st.success(f"🎟️ **Ticket {i+1}:** Predicted Team → {team}, Contact: {email}")
        else:
            st.warning(f"🚨 Ticket {i+1} is empty. Please enter a description.")

st.write("📌 Enter your ticket descriptions and click **Predict All Tickets** to classify them.")