Upload app.py

app.py

@@ -1,98 +1,156 @@
-import streamlit as st
-from llama_cpp import Llama
-from dotenv import load_dotenv
-import os
-
-
-load_dotenv()
-print(f"REPO_ID: {os.getenv('REPO_ID')}")
-
-llm = Llama.from_pretrained(
-    repo_id=os.getenv("REPO_ID"),
-    filename="unsloth.Q8_0.gguf"
-)
-
-def load_html(file_path):
-    with open(file_path, "r", encoding="utf-8") as file:
-        return file.read()
-
-st.markdown(load_html("chatbot_ui.html"), unsafe_allow_html=True)
-st.markdown('<div class="title">HR Chatbot</div>', unsafe_allow_html=True)
-st.write("\n\n\n")
-
-
-intents = {
-    
-    "Employee Benefits & Policies": "Assist with salary, leave, and policies.",
-    "Employee Support & Self-Service": "Provide support for HR-related queries.",
-    "Recruitment & Onboarding": "Help with job applications and hiring processes."
-}
-
-
-intent_keywords = {
-    
-    "Employee Benefits & Policies": ["leave", "insurance", "policy", "benefit", "salary", "payroll", "bonus"],
-    "Employee Support & Self-Service": ["complaint", "technical issue", "self-service", "support", "HR help"],
-    "Recruitment & Onboarding": ["hiring", "job", "apply", "interview", "onboarding", "new employee"]
-}
-
-
-def determine_intent(question):
-    question_lower = question.lower()
-    for intent, keywords in intent_keywords.items():
-        if any(keyword in question_lower for keyword in keywords):
-            return intent
-    return None  
-
-
-if "selected_intent" not in st.session_state:
-    st.session_state.selected_intent = list(intents.keys())[0]
-
-cols = st.columns(len(intents))
-for i, (intent, desc) in enumerate(intents.items()):
-    btn_label = f"✅ **{intent}**" if st.session_state.selected_intent == intent else intent
-    if cols[i].button(btn_label, key=intent, help=desc):
-        st.session_state.selected_intent = intent  
-        st.rerun()
-
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-
-if st.session_state.messages:
-    st.markdown('<div class="chat-container">', unsafe_allow_html=True)
-    for msg in st.session_state.messages:
-        role_class = "user-msg" if msg["role"] == "user" else "bot-msg"
-        with st.chat_message(msg["role"]):
-            st.markdown(f'<div class="chat-box {role_class}">{msg["content"]}</div>', unsafe_allow_html=True)
-
-
-if prompt := st.chat_input("Ask me anything related to HR..."):
-   
-    st.session_state.messages.append({"role": "user", "content": prompt})
-
-    with st.chat_message("user"):
-        st.markdown(f'<div class="chat-box user-msg">👤 {prompt}</div>', unsafe_allow_html=True)
-
-   
-    correct_intent = determine_intent(prompt)
-    
-    if correct_intent and correct_intent != st.session_state.selected_intent:
-      
-        reply = f"🔄 Redirecting... Your query belongs to '{correct_intent}'. Please select the correct category."
-    else:
-        with st.spinner("🤔 Thinking..."):
-            response = llm.create_chat_completion(
-                messages=[{"role": "user", "content": prompt}],  
-                max_tokens=280 
-            )
-            reply = response["choices"][0]["message"]["content"]
-
-    
-    st.session_state.messages.append({"role": "assistant", "content": reply})
-
-    with st.chat_message("assistant"):
-        st.markdown(f'<div class="chat-box bot-msg">🤖 {reply}</div>', unsafe_allow_html=True)
-
-if st.session_state.messages:
-    st.markdown('</div>', unsafe_allow_html=True)
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForSequenceClassification
+import streamlit as st
+from llama_cpp import Llama
+from dotenv import load_dotenv
+import os
+import torch
+
+load_dotenv()
+
+try:
+    REPO_ID = os.getenv('REPO_ID')
+    INTENT_MODEL = os.getenv('INTENT_MODEL')
+    CHITCHAT_MODEL = os.getenv('CHITCHAT_MODEL')
+
+    if not all([REPO_ID, INTENT_MODEL, CHITCHAT_MODEL]):
+        raise EnvironmentError("One or more required environment variables are missing.")
+
+except Exception as e:
+    st.error(f"Environment setup failed: {e}")
+    raise SystemExit(e)
+
+try:
+    llm_hr = Llama.from_pretrained(
+        repo_id=REPO_ID,
+        filename="unsloth.F16.gguf"
+    )
+    intent_tokenizer = AutoTokenizer.from_pretrained(INTENT_MODEL)
+    intent_model = AutoModelForSequenceClassification.from_pretrained(INTENT_MODEL)
+    tokenizer = AutoTokenizer.from_pretrained(CHITCHAT_MODEL)
+    model = AutoModelForSeq2SeqLM.from_pretrained(CHITCHAT_MODEL)
+    pipe_chitchat = pipeline('text2text-generation', model=model, tokenizer=tokenizer)
+
+except Exception as e:
+    st.error(f"Model loading failed: {e}")
+    raise SystemExit(e)
+
+
+intent_labels = ["Employee Benefits & Policies", "Employee Support & Self-Service", "Recruitment & Onboarding"]
+total_labels = ["😂 Let's Chat & Chill"] + intent_labels
+
+
+st.set_page_config(page_title="AI HR Chatbot", layout="wide")
+st.markdown("""
+<style>
+body { background-color: #f6f9fc; }
+.sidebar .sidebar-content { background-color: #ffffff; }
+.chat-container { max-height: 500px; overflow-y: auto; padding: 1rem; display: flex; flex-direction: column; }
+.chat-box { display: inline-block; border-radius: 10px; padding: 8px 12px; margin: 5px 0; max-width: 70%; word-wrap: break-word; }
+.user-msg { background-color: #708090; text-align: right; align-self: flex-end; color: white; }
+.bot-msg { background-color: #d3d3d3; text-align: left; align-self: flex-start; color: black; }
+div.title { font-size: 2rem; font-weight: bold; margin: 20px 0; color: #0077cc; }
+</style>
+""", unsafe_allow_html=True)
+
+if "intent" not in st.session_state:
+    st.session_state.intent = total_labels[0]
+if "messages" not in st.session_state:
+    st.session_state.messages = []
+
+
+def select_intent(current_intent):
+    display_labels = [f"✅ {intent}" if intent == current_intent else intent for intent in total_labels]
+    selected = st.sidebar.radio(
+        "💼 Select Your HR Intent",
+        display_labels,
+        index=display_labels.index(f"✅ {current_intent}") if f"✅ {current_intent}" in display_labels else 0
+    )
+    return selected.replace("✅ ", "")
+
+
+def determine_intent(text):
+    try:
+        inputs = intent_tokenizer(text, return_tensors="pt", truncation=True, padding=True)
+        with torch.no_grad():
+            logits = intent_model(**inputs).logits
+        predicted_label = torch.argmax(logits, dim=1).item()
+        return predicted_label
+    except Exception as e:
+        st.error(f"Intent determination failed: {e}")
+        return 0  
+
+
+def render_chat():
+    if st.session_state.messages:
+        st.markdown('<div class="chat-container">', unsafe_allow_html=True)
+        for msg in st.session_state.messages:
+            role_class = "user-msg" if msg["role"] == "user" else "bot-msg"
+            with st.chat_message(msg["role"], avatar=msg.get("avatar")):
+                st.markdown(f'<div class="chat-box {role_class}">{msg["content"]}</div>', unsafe_allow_html=True)
+        st.markdown('</div>', unsafe_allow_html=True)
+
+
+def generate_hr_response(user_prompt, intent):
+    hr_prompt = """You are an HR Assistant at our company.
+    Your role is to assist employees by providing accurate and concise responses regarding company policies and HR-related questions.
+
+    ### Instruction:
+    {}
+
+    ### Input:
+    {}
+
+    ### Response:
+    {}"""
+    instruction = f"Answer the HR-related query categorized as '{intent}'."
+    formatted_prompt = hr_prompt.format(instruction, user_prompt, "") + tokenizer.eos_token
+
+    try:
+        response = llm_hr.create_chat_completion(
+            messages=[{"role": "user", "content": formatted_prompt}],
+            max_tokens=100,
+        )
+        return response["choices"][0]["message"]["content"]
+    except Exception as e:
+        st.error(f"Failed to generate HR response: {e}")
+        return "⚠️ Sorry, I couldn't process your HR request at the moment."
+
+
+st.markdown('<div class="title">AI HR Chatbot </div>', unsafe_allow_html=True)
+st.markdown("<h6 style='color:rgb(131, 123, 160);'>Your personal assistant for HR queries and support.</h6>", unsafe_allow_html=True)
+
+st.session_state.intent = select_intent(st.session_state.intent)
+render_chat()
+
+
+try:
+    if prompt := st.chat_input("Ask me anything related to HR or just chat casually..."):
+        st.session_state.messages.append({"role": "user", "content": prompt, "avatar": "man.png"})
+        with st.chat_message("user", avatar="man.png"):
+            st.markdown(f'<div class="chat-box user-msg">{prompt}</div>', unsafe_allow_html=True)
+
+        if st.session_state.intent == "😂 Let's Chat & Chill":
+            with st.spinner("😄 Chitchatting..."):
+                try:
+                    result = pipe_chitchat(prompt)
+                    reply = result[0]['generated_text'].strip() 
+                except Exception as e:
+                    st.error(f"Chitchat generation failed: {e}")
+                    reply = "Oops, I couldn't come up with a witty reply! 😅"
+        else:
+            detected_intent = determine_intent(prompt)
+            predicted_label = intent_labels[detected_intent]
+
+            if predicted_label != st.session_state.intent:
+                
+                st.warning(f"🔁 Automatically switched to **{predicted_label}** based on your query.")
+
+            with st.spinner("Generating response with specialized HR model..."):
+                reply = generate_hr_response(prompt, predicted_label)
+
+        st.session_state.messages.append({"role": "bot", "content": reply, "avatar": "robot.png"})
+        with st.chat_message("bot", avatar="robot.png"):
+            st.markdown(f'<div class="chat-box bot-msg">{reply}</div>', unsafe_allow_html=True)
+
+except Exception as e:
+    st.error(f"Something went wrong while processing your input: {e}")