app.py

import streamlit as st
import time
import os
import gc
import torch
from src.data_processing import load_huggingface_faq_data, load_faq_data, preprocess_faq, augment_faqs
from src.embedding import FAQEmbedder
from src.llm_response import ResponseGenerator
from src.utils import time_function, format_memory_stats, evaluate_response, evaluate_retrieval, baseline_keyword_search

# Suppress CUDA warning and Torch path errors
os.environ["CUDA_VISIBLE_DEVICES"] = ""
os.environ["TORCH_NO_PATH_CHECK"] = "1"

st.set_page_config(page_title="E-Commerce FAQ Chatbot", layout="wide", initial_sidebar_state="expanded")

@time_function
def initialize_components(use_huggingface: bool = True, model_name: str = "microsoft/phi-2", enable_augmentation: bool = True):
    """
    Initialize RAG system components
    """
    try:
        if use_huggingface:
            faqs = load_huggingface_faq_data("NebulaByte/E-Commerce_FAQs")
        else:
            faqs = load_faq_data("data/faq_data.csv")
        
        processed_faqs = augment_faqs(preprocess_faq(faqs), enable_augmentation=enable_augmentation)
        embedder = FAQEmbedder()
        
        if os.path.exists("embeddings"):
            embedder.load("embeddings")
        else:
            embedder.create_embeddings(processed_faqs)
            embedder.save("embeddings")
        
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        
        response_generator = ResponseGenerator(model_name=model_name)
        response_generator.generate_response("Warmup query", [{"question": "Test", "answer": "Test"}])
        
        return embedder, response_generator, len(processed_faqs)
    except Exception as e:
        st.error(f"Initialization failed: {e}")
        raise

def main():
    st.title("E-Commerce Customer Support FAQ Chatbot")
    st.subheader("Ask about orders, shipping, returns, or other e-commerce queries")
    
    st.sidebar.title("Configuration")
    use_huggingface = st.sidebar.checkbox("Use Hugging Face Dataset", value=True)
    enable_augmentation = st.sidebar.checkbox("Enable FAQ Augmentation", value=True, help="Generate paraphrased questions to expand dataset")
    target_lang = st.sidebar.selectbox("Language", ["en", "es", "fr"], index=0)
    
    model_options = {
        "Phi-2 (Recommended for 16GB RAM)": "microsoft/phi-2",
        "TinyLlama-1.1B (Fastest)": "TinyLlama/TinyLlama-1.1B-Chat-v1.0",
        "Mistral-7B (For 15GB+ GPU)": "mistralai/Mistral-7B-Instruct-v0.1"
    }
    selected_model = st.sidebar.selectbox("Select LLM Model", list(model_options.keys()), index=0)
    model_name = model_options[selected_model]
    
    if st.sidebar.checkbox("Show Memory Usage", value=True):
        st.sidebar.subheader("Memory Usage")
        for key, value in format_memory_stats().items():
            st.sidebar.text(f"{key}: {value}")
    
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "query_cache" not in st.session_state:
        st.session_state.query_cache = {}
    if "feedback" not in st.session_state:
        st.session_state.feedback = []
    
    if "system_initialized" not in st.session_state or st.sidebar.button("Reload System"):
        with st.spinner("Initializing system..."):
            try:
                st.session_state.embedder, st.session_state.response_generator, num_faqs = initialize_components(
                    use_huggingface=use_huggingface,
                    model_name=model_name,
                    enable_augmentation=enable_augmentation
                )
                st.session_state.system_initialized = True
                st.sidebar.success(f"System initialized with {num_faqs} FAQs!")
            except Exception as e:
                st.error(f"System initialization failed: {e}")
                return
    
    col1, col2 = st.columns([2, 1])
    
    with col1:
        st.subheader("Conversation")
        chat_container = st.container(height=400)
        with chat_container:
            for i, message in enumerate(st.session_state.chat_history):
                if message["role"] == "user":
                    st.markdown(f"**You**: {message['content']}")
                else:
                    st.markdown(f"**Bot**: {message['content']}")
                if i < len(st.session_state.chat_history) - 1:
                    st.markdown("---")
        
        with st.form(key="chat_form"):
            user_query = st.text_input("Type your question:", key="user_input", placeholder="e.g., How do I track my order?")
            submit_button = st.form_submit_button("Ask")
        
        if len(st.session_state.chat_history) > 0:
            with st.form(key=f"feedback_form_{len(st.session_state.chat_history)}"):
                rating = st.slider("Rate this response (1-5)", 1, 5, key=f"rating_{len(st.session_state.chat_history)}")
                comments = st.text_area("Comments", key=f"comments_{len(st.session_state.chat_history)}")
                if st.form_submit_button("Submit Feedback"):
                    st.session_state.feedback.append({
                        "rating": rating,
                        "comments": comments,
                        "response": st.session_state.chat_history[-1]["content"]
                    })
                    with open("feedback.json", "w") as f:
                        json.dump(st.session_state.feedback, f)
                    st.success("Feedback submitted!")
    
    with col2:
        if st.session_state.get("system_initialized", False):
            st.subheader("Retrieved Information")
            info_container = st.container(height=500)
            with info_container:
                if "current_faqs" in st.session_state:
                    for i, faq in enumerate(st.session_state.current_faqs):
                        st.markdown(f"**Relevant FAQ #{i+1}**")
                        st.markdown(f"**Q**: {faq['question']}")
                        st.markdown(f"**A**: {faq['answer'][:150]}..." if len(faq['answer']) > 150 else f"**A**: {faq['answer']}")
                        st.markdown(f"*Similarity Score*: {faq['similarity']:.2f}")
                        if 'category' in faq and faq['category']:
                            st.markdown(f"*Category*: {faq['category']}")
                        st.markdown("---")
                else:
                    st.markdown("Ask a question to see relevant FAQs.")
    
    if "retrieval_time" in st.session_state and "generation_time" in st.session_state:
        st.sidebar.subheader("Performance Metrics")
        st.sidebar.markdown(f"Retrieval time: {st.session_state.retrieval_time:.2f} seconds")
        st.sidebar.markdown(f"Response generation: {st.session_state.generation_time:.2f} seconds")
        st.sidebar.markdown(f"Total time: {st.session_state.retrieval_time + st.session_state.generation_time:.2f} seconds")
    
    if submit_button and user_query:
        from src.data_processing import translate_faq
        from googletrans import Translator
        translator = Translator()
        if target_lang != "en":
            user_query_translated = translator.translate(user_query, dest="en").text
        else:
            user_query_translated = user_query
        
        if user_query_translated in st.session_state.query_cache:
            response, relevant_faqs = st.session_state.query_cache[user_query_translated]
        else:
            gc.collect()
            if torch.cuda.is_available():
                torch.cuda.empty_cache()
            
            start_time = time.time()
            relevant_faqs = st.session_state.embedder.retrieve_relevant_faqs(user_query_translated)
            retrieval_time = time.time() - start_time
            
            if target_lang != "en":
                relevant_faqs = [translate_faq(faq, target_lang) for faq in relevant_faqs]
            
            start_time = time.time()
            response = st.session_state.response_generator.generate_response(user_query_translated, relevant_faqs)
            generation_time = time.time() - start_time
            
            if target_lang != "en":
                response = translator.translate(response, dest=target_lang).text
            
            st.session_state.query_cache[user_query_translated] = (response, relevant_faqs)
            st.session_state.retrieval_time = retrieval_time
            st.session_state.generation_time = generation_time
            st.session_state.current_faqs = relevant_faqs
        
        st.session_state.chat_history.append({"role": "user", "content": user_query})
        st.session_state.chat_history.append({"role": "assistant", "content": response})
    
    if st.button("Clear Chat History"):
        st.session_state.chat_history = []
        st.session_state.query_cache = {}
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    
    if st.session_state.get("system_initialized", False):
        st.sidebar.subheader("Baseline Comparison")
        baseline_faqs = baseline_keyword_search(user_query_translated if 'user_query_translated' in locals() else "", st.session_state.embedder.faqs)
        st.sidebar.write(f"RAG FAQs: {[faq['question'][:50] for faq in st.session_state.get('current_faqs', [])]}")
        st.sidebar.write(f"Keyword FAQs: {[faq['question'][:50] for faq in baseline_faqs]}")
    
    st.subheader("Sample Questions")
    sample_questions = [
        "How do I track my order?",
        "What should I do if my delivery is delayed?",
        "How do I return a product?",
        "Can I cancel my order after placing it?",
        "How quickly will my order be delivered?"
    ]
    cols = st.columns(2)
    for i, question in enumerate(sample_questions):
        col_idx = i % 2
        if cols[col_idx].button(question, key=f"sample_{i}"):
            st.session_state.user_input = question
            st.session_state.chat_history.append({"role": "user", "content": question})
            
            from src.data_processing import translate_faq
            from googletrans import Translator
            translator = Translator()
            if target_lang != "en":
                question_translated = translator.translate(question, dest="en").text
            else:
                question_translated = question
            
            if question_translated in st.session_state.query_cache:
                response, relevant_faqs = st.session_state.query_cache[question_translated]
            else:
                gc.collect()
                if torch.cuda.is_available():
                    torch.cuda.empty_cache()
                
                start_time = time.time()
                relevant_faqs = st.session_state.embedder.retrieve_relevant_faqs(question_translated)
                retrieval_time = time.time() - start_time
                
                if target_lang != "en":
                    relevant_faqs = [translate_faq(faq, target_lang) for faq in relevant_faqs]
                
                start_time = time.time()
                response = st.session_state.response_generator.generate_response(question_translated, relevant_faqs)
                generation_time = time.time() - start_time
                
                if target_lang != "en":
                    response = translator.translate(response, dest=target_lang).text
                
                st.session_state.query_cache[question_translated] = (response, relevant_faqs)
                st.session_state.retrieval_time = retrieval_time
                st.session_state.generation_time = generation_time
                st.session_state.current_faqs = relevant_faqs
            
            st.session_state.chat_history.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()