Update app.py

app.py

@@ -1,12 +1,14 @@
 import os
 import warnings
+import logging
+import sys
 warnings.filterwarnings("ignore", category=UserWarning)
 
 import numpy as np
 import pandas as pd
 import torch
 from sentence_transformers import SentenceTransformer
-from typing import List, Callable
+from typing import List, Callable, Dict, Optional, Any
 import glob
 from tqdm import tqdm
 import pickle
@@ -19,9 +21,22 @@ import re
 import time
 import requests
 
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout)
+    ]
+)
+
 # Force CPU device
 torch.device('cpu')
 
+# Create necessary directories
+for directory in ['models', 'ESPN_data', 'embeddings_cache']:
+    os.makedirs(directory, exist_ok=True)
+
 # Logging configuration
 LOGGING_CONFIG = {
     'enabled': True,
@@ -49,36 +64,7 @@ def download_file_with_progress(url: str, filename: str):
         for data in response.iter_content(chunk_size=1024):
             size = file.write(data)
             progress_bar.update(size)
-            
-@st.cache_data
-def load_from_drive(file_id: str):
-    """Load pickle file directly from Google Drive"""
-    try:
-        # Direct download URL for Google Drive
-        url = f"https://drive.google.com/uc?id={file_id}&export=download"
-        
-        # First request to get the confirmation token
-        session = requests.Session()
-        response = session.get(url, stream=True)
-        
-        # Check if we need to confirm download
-        for key, value in response.cookies.items():
-            if key.startswith('download_warning'):
-                # Add confirmation parameter to the URL
-                url = f"{url}&confirm={value}"
-                response = session.get(url, stream=True)
-                break
-        
-        # Load the content and convert to pickle
-        content = response.content
-        print(f"Successfully downloaded {len(content)} bytes")
-        return pickle.loads(content)
-        
-    except Exception as e:
-        print(f"Detailed error: {str(e)}")  # This will help debug
-        st.error(f"Error loading file from Drive: {str(e)}")
-        return None
-    
+
 def log_function(func: Callable) -> Callable:
     """Decorator to log function inputs and outputs"""
     @functools.wraps(func)
@@ -135,48 +121,70 @@ def check_environment():
         st.stop()
         return False
 
-
-
 class SentenceTransformerRetriever:
-    @st.cache_resource
-    def __init__(_self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", cache_dir: str = "embeddings_cache"):
-        # Force CPU device and suppress warnings
+    def __init__(self, model_name: str = "sentence-transformers/all-MiniLM-L6-v2", cache_dir: str = "embeddings_cache"):
+        self.device = torch.device("cpu")
+        self.model_name = model_name
+        self.cache_dir = cache_dir
+        self.cache_file = "embeddings.pkl"
+        self.doc_embeddings = None
+        os.makedirs(cache_dir, exist_ok=True)
+        # Initialize model using cached method
+        self.model = self._load_model()
+    
+    @st.cache_resource(show_spinner=False)
+    def _load_model(self):
+        """Load and cache the sentence transformer model"""
         with warnings.catch_warnings():
             warnings.simplefilter("ignore")
-            _self.device = torch.device("cpu")
-            _self.model = SentenceTransformer(model_name, device="cpu")
-            _self.doc_embeddings = None
-            _self.cache_dir = cache_dir
-            _self.cache_file = "embeddings.pkl"
-            os.makedirs(cache_dir, exist_ok=True)
+            model = SentenceTransformer(self.model_name, device="cpu")
+            # Verify model is loaded correctly
+            test_embedding = model.encode("test", convert_to_tensor=True)
+            if not isinstance(test_embedding, torch.Tensor):
+                raise ValueError("Model initialization failed")
+            return model
     
     def get_cache_path(self, data_folder: str = None) -> str:
         return os.path.join(self.cache_dir, self.cache_file)
-
     
     @log_function
     def save_cache(self, data_folder: str, cache_data: dict):
-        cache_path = self.get_cache_path()
-        if os.path.exists(cache_path):
-            os.remove(cache_path)
-        with open(cache_path, 'wb') as f:
-            pickle.dump(cache_data, f)
-            print(f"Cache saved at: {cache_path}")
+        try:
+            cache_path = self.get_cache_path()
+            if os.path.exists(cache_path):
+                os.remove(cache_path)
+            with open(cache_path, 'wb') as f:
+                pickle.dump(cache_data, f)
+                logging.info(f"Cache saved at: {cache_path}")
+        except Exception as e:
+            logging.error(f"Error saving cache: {str(e)}")
+            raise
     
     @log_function
     @st.cache_data
-    def load_cache(_self, data_folder: str = None) -> dict:
-        cache_path = _self.get_cache_path()
-        if os.path.exists(cache_path):
-            with open(cache_path, 'rb') as f:
-                print(f"Loading cache from: {cache_path}")
-                return pickle.load(f)
-        return None
-    
+    def load_cache(self, data_folder: str = None) -> Optional[Dict]:
+        try:
+            cache_path = self.get_cache_path()
+            if os.path.exists(cache_path):
+                with open(cache_path, 'rb') as f:
+                    logging.info(f"Loading cache from: {cache_path}")
+                    cache_data = pickle.load(f)
+                    if isinstance(cache_data, dict) and 'embeddings' in cache_data and 'documents' in cache_data:
+                        return cache_data
+                    logging.warning("Invalid cache format")
+            return None
+        except Exception as e:
+            logging.error(f"Error loading cache: {str(e)}")
+            return None
+
     @log_function
     def encode(self, texts: List[str], batch_size: int = 32) -> torch.Tensor:
-        embeddings = self.model.encode(texts, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)
-        return F.normalize(embeddings, p=2, dim=1)
+        try:
+            embeddings = self.model.encode(texts, batch_size=batch_size, convert_to_tensor=True, show_progress_bar=True)
+            return F.normalize(embeddings, p=2, dim=1)
+        except Exception as e:
+            logging.error(f"Error encoding texts: {str(e)}")
+            raise
 
     @log_function
     def store_embeddings(self, embeddings: torch.Tensor):
@@ -187,23 +195,16 @@ class SentenceTransformerRetriever:
         if self.doc_embeddings is None:
             raise ValueError("No document embeddings stored!")
         
-        # Compute similarities
         similarities = F.cosine_similarity(query_embedding, self.doc_embeddings)
-        
-        # Get top k scores and indices
         k = min(k, len(documents))
         scores, indices = torch.topk(similarities, k=k)
         
-        # Log similarity statistics
-        print(f"\nSimilarity Stats:")
-        print(f"Max similarity: {similarities.max().item():.4f}")
-        print(f"Mean similarity: {similarities.mean().item():.4f}")
-        print(f"Selected similarities: {scores.tolist()}")
+        logging.info(f"\nSimilarity Stats:")
+        logging.info(f"Max similarity: {similarities.max().item():.4f}")
+        logging.info(f"Mean similarity: {similarities.mean().item():.4f}")
+        logging.info(f"Selected similarities: {scores.tolist()}")
         
         return indices.cpu(), scores.cpu()
-    
-
-
 
 class RAGPipeline:
     def __init__(self, data_folder: str, k: int = 5):
@@ -212,29 +213,25 @@ class RAGPipeline:
         self.retriever = SentenceTransformerRetriever()
         self.documents = []
         self.device = torch.device("cpu")
-        self.model_path = "mistral-7b-v0.1.Q4_K_M.gguf"
-        # Initialize model in init
+        self.model_path = os.path.join("models", "mistral-7b-v0.1.Q4_K_M.gguf")
         self.llm = None
-        self.initialize_model()
-
+        self._initialize_model()
 
-    st.cache_resource
-    def initialize_model(_self):
-        """Initialize the model with proper error handling and verification
-        
-        Note: Using _self instead of self for Streamlit caching compatibility
-        """
+    @st.cache_resource(show_spinner=False)
+    def _initialize_model(self):
+        """Initialize the model with proper error handling and verification"""
         try:
-            if not os.path.exists(_self.model_path):
+            os.makedirs(os.path.dirname(self.model_path), exist_ok=True)
+            
+            if not os.path.exists(self.model_path):
                 direct_url = "https://huggingface.co/TheBloke/Mistral-7B-v0.1-GGUF/resolve/main/mistral-7b-v0.1.Q4_K_M.gguf"
-                download_file_with_progress(direct_url, _self.model_path)
+                download_file_with_progress(direct_url, self.model_path)
 
-            # Verify file exists and has content
-            if not os.path.exists(_self.model_path):
-                raise FileNotFoundError(f"Model file {_self.model_path} not found after download attempts")
+            if not os.path.exists(self.model_path):
+                raise FileNotFoundError(f"Model file {self.model_path} not found after download attempts")
             
-            if os.path.getsize(_self.model_path) < 1000000:  # Less than 1MB
-                os.remove(_self.model_path)
+            if os.path.getsize(self.model_path) < 1000000:  # Less than 1MB
+                os.remove(self.model_path)
                 raise ValueError("Downloaded model file is too small, likely corrupted")
             
             llm_config = {
@@ -245,43 +242,73 @@ class RAGPipeline:
                 "verbose": False
             }
         
-            _self.llm = Llama(model_path=_self.model_path, **llm_config)
+            self.llm = Llama(model_path=self.model_path, **llm_config)
             st.success("Model loaded successfully!")
 
         except Exception as e:
+            logging.error(f"Error initializing model: {str(e)}")
             st.error(f"Error initializing model: {str(e)}")
             raise
-        
+
+    def check_model_health(self):
+        """Verify that the model is loaded and functioning"""
+        try:
+            if self.llm is None:
+                return False
+                
+            # Simple test prompt
+            test_response = self.llm(
+                "Test prompt",
+                max_tokens=10,
+                temperature=0.4,
+                echo=False
+            )
+            
+            return isinstance(test_response, dict) and 'choices' in test_response
+        except Exception:
+            return False
+
     @log_function
     @st.cache_data
-    def load_and_process_csvs(_self):
-        cache_data = _self.retriever.load_cache(_self.data_folder)
-        if cache_data is not None:
-            _self.documents = cache_data['documents']
-            _self.retriever.store_embeddings(cache_data['embeddings'])
-            return
-
-        csv_files = glob.glob(os.path.join(_self.data_folder, "*.csv"))
-        all_documents = []
-        
-        for csv_file in tqdm(csv_files, desc="Reading CSV files"):
-            try:
-                df = pd.read_csv(csv_file)
-                texts = df.apply(lambda x: " ".join(x.astype(str)), axis=1).tolist()
-                all_documents.extend(texts)
-            except Exception as e:
-                print(f"Error processing file {csv_file}: {e}")
-                continue
-        
-        _self.documents = all_documents
-        embeddings = _self.retriever.encode(all_documents)
-        _self.retriever.store_embeddings(embeddings)
-        
-        cache_data = {
-            'embeddings': embeddings,
-            'documents': _self.documents
-        }
-        _self.retriever.save_cache(_self.data_folder, cache_data)
+    def load_and_process_csvs(self):
+        try:
+            cache_data = self.retriever.load_cache(self.data_folder)
+            if cache_data is not None:
+                self.documents = cache_data['documents']
+                self.retriever.store_embeddings(cache_data['embeddings'])
+                return
+
+            csv_files = glob.glob(os.path.join(self.data_folder, "*.csv"))
+            if not csv_files:
+                raise FileNotFoundError(f"No CSV files found in {self.data_folder}")
+                
+            all_documents = []
+            
+            for csv_file in tqdm(csv_files, desc="Reading CSV files"):
+                try:
+                    df = pd.read_csv(csv_file)
+                    texts = df.apply(lambda x: " ".join(x.astype(str)), axis=1).tolist()
+                    all_documents.extend(texts)
+                except Exception as e:
+                    logging.error(f"Error processing file {csv_file}: {e}")
+                    continue
+            
+            if not all_documents:
+                raise ValueError("No documents were successfully loaded")
+                
+            self.documents = all_documents
+            embeddings = self.retriever.encode(all_documents)
+            self.retriever.store_embeddings(embeddings)
+            
+            cache_data = {
+                'embeddings': embeddings,
+                'documents': self.documents
+            }
+            self.retriever.save_cache(self.data_folder, cache_data)
+            
+        except Exception as e:
+            logging.error(f"Error in load_and_process_csvs: {str(e)}")
+            raise
 
     def preprocess_query(self, query: str) -> str:
         """Clean and prepare the query"""
@@ -299,6 +326,12 @@ class RAGPipeline:
     @log_function
     def process_query(self, query: str, placeholder) -> str:
         try:
+            # Check if models are properly initialized
+            if self.llm is None:
+                raise RuntimeError("LLM model not initialized")
+            if self.retriever.model is None:
+                raise RuntimeError("Sentence transformer model not initialized")
+                
             # Preprocess query
             query = self.preprocess_query(query)
             
@@ -311,9 +344,9 @@ class RAGPipeline:
             indices, scores = self.retriever.search(query_embedding, self.k, self.documents)
             
             # Print search results for debugging
-            print("\nSearch Results:")
+            logging.info("\nSearch Results:")
             for idx, score in zip(indices.tolist(), scores.tolist()):
-                print(f"Score: {score:.4f} | Document: {self.documents[idx][:100]}...")
+                logging.info(f"Score: {score:.4f} | Document: {self.documents[idx][:100]}...")
             
             relevant_docs = [self.documents[idx] for idx in indices.tolist()]
             
@@ -339,7 +372,6 @@ class RAGPipeline:
             
             # Generate response
             response_placeholder = placeholder.empty()
-            generated_text = ""
             
             try:
                 response = self.llm(
@@ -368,163 +400,183 @@ class RAGPipeline:
                     return message
                     
             except Exception as e:
-                print(f"Generation error: {str(e)}")
+                logging.error(f"Generation error: {str(e)}")
                 message = "Had some trouble generating the response. Please try again."
                 response_placeholder.warning(message)
                 return message
                 
         except Exception as e:
-            print(f"Process error: {str(e)}")
+            logging.error(f"Process error: {str(e)}")
             message = "Something went wrong. Please try again with a different question."
             placeholder.warning(message)
             return message
-        
 
-
-@st.cache_resource
+@st.cache_resource(show_spinner=False)
 def initialize_rag_pipeline():
     """Initialize the RAG pipeline once"""
-    data_folder = "ESPN_data"  # Update this path as needed
-    rag = RAGPipeline(data_folder)
-    rag.load_and_process_csvs()
-    return rag
+    try:
+        data_folder = "ESPN_data"
+        if not os.path.exists(data_folder):
+            os.makedirs(data_folder, exist_ok=True)
+            
+        rag = RAGPipeline(data_folder)
+        rag.load_and_process_csvs()
+        return rag
+    except Exception as e:
+        logging.error(f"Pipeline initialization error: {str(e)}")
+        st.error("Failed to initialize the system. Please check your data folder and try again.")
+        raise
 
 def main():
-    # Environment check
-    if not check_environment():
-        return
-
-    # Page config
-    st.set_page_config(
-        page_title="The Sport Chatbot",
-        page_icon="🏆",
-        layout="wide"  # Changed back to wide for more space
-    )
-
-    # Improved CSS styling
-    st.markdown("""
-        <style>
-        /* Container styling */
-        .block-container {
-            padding-top: 2rem;
-            padding-bottom: 2rem;
-        }
-        
-        /* Text input styling */
-        .stTextInput > div > div > input {
-            width: 100%;
-        }
-        
-        /* Button styling */
-        .stButton > button {
-            width: 200px;
-            margin: 0 auto;
-            display: block;
-            background-color: #FF4B4B;
-            color: white;
-            border-radius: 5px;
-            padding: 0.5rem 1rem;
-        }
-        
-        /* Title styling */
-        .main-title {
-            text-align: center;
-            padding: 1rem 0;
-            font-size: 3rem;
-            color: #1F1F1F;
-        }
-        
-        .sub-title {
-            text-align: center;
-            padding: 0.5rem 0;
-            font-size: 1.5rem;
-            color: #4F4F4F;
-        }
-        
-        /* Description styling */
-        .description {
-            text-align: center;
-            color: #666666;
-            padding: 0.5rem 0;
-            font-size: 1.1rem;
-            line-height: 1.6;
-            margin-bottom: 1rem;
-        }
-
-        /* Answer container styling */
-        .stMarkdown {
-            max-width: 100%;
-        }
-
-        /* Streamlit default overrides */
-        .st-emotion-cache-16idsys p {
-            font-size: 1.1rem;
-            line-height: 1.6;
-        }
-        
-        /* Container for main content */
-        .main-content {
-            max-width: 1200px;
-            margin: 0 auto;
-            padding: 0 1rem;
-        }
-        </style>
-    """, unsafe_allow_html=True)
-
-    # Header section with improved styling
-    st.markdown("<h1 class='main-title'>🏆 The Sport Chatbot</h1>", unsafe_allow_html=True)
-    st.markdown("<h3 class='sub-title'>Using ESPN API</h3>", unsafe_allow_html=True)
-    st.markdown("""
-        <p class='description'>
-            Hey there! 👋 I can help you with information on Ice Hockey, Baseball, American Football, Soccer, and Basketball. 
-            With access to the ESPN API, I'm up to date with the latest details for these sports up until October 2024.
-        </p>
-        <p class='description'>
-            Got any general questions? Feel free to ask—I'll do my best to provide answers based on the information I've been trained on!
-        </p>
-    """, unsafe_allow_html=True)
-
-    # Add some spacing
-    st.markdown("<br>", unsafe_allow_html=True)
-
-
-    # Initialize the pipeline
     try:
-        with st.spinner("Loading resources..."):
-            rag = initialize_rag_pipeline()
-    except Exception as e:
-        print(f"Initialization error: {str(e)}")
-        st.error("Unable to initialize the system. Please check if all required files are present.")
-        st.stop()
+        # Environment check
+        if not check_environment():
+            return
 
-    # Create columns for layout with golden ratio
-    col1, col2, col3 = st.columns([1, 6, 1])
-    
-    with col2:
-        # Query input with label styling
-        query = st.text_input("What would you like to know about sports?")
+        # Page config
+        st.set_page_config(
+            page_title="The Sport Chatbot",
+            page_icon="🏆",
+            layout="wide"
+        )
+
+        # Improved CSS styling
+        st.markdown("""
+            <style>
+            /* Container styling */
+            .block-container {
+                padding-top: 2rem;
+                padding-bottom: 2rem;
+            }
+            
+            /* Text input styling */
+            .stTextInput > div > div > input {
+                width: 100%;
+            }
+            
+            /* Button styling */
+            .stButton > button {
+                width: 200px;
+                margin: 0 auto;
+                display: block;
+                background-color: #FF4B4B;
+                color: white;
+                border-radius: 5px;
+                padding: 0.5rem 1rem;
+            }
+            
+            /* Title styling */
+            .main-title {
+                text-align: center;
+                padding: 1rem 0;
+                font-size: 3rem;
+                color: #1F1F1F;
+            }
+            
+            .sub-title {
+                text-align: center;
+                padding: 0.5rem 0;
+                font-size: 1.5rem;
+                color: #4F4F4F;
+            }
+            
+            /* Description styling */
+            .description {
+                text-align: center;
+                color: #666666;
+                padding: 0.5rem 0;
+                font-size: 1.1rem;
+                line-height: 1.6;
+                margin-bottom: 1rem;
+            }
+
+            /* Answer container styling */
+            .stMarkdown {
+                max-width: 100%;
+            }
+
+            /* Streamlit default overrides */
+            .st-emotion-cache-16idsys p {
+                font-size: 1.1rem;
+                line-height: 1.6;
+            }
+            
+            /* Container for main content */
+            .main-content {
+                max-width: 1200px;
+                margin: 0 auto;
+                padding: 0 1rem;
+            }
+            </style>
+        """, unsafe_allow_html=True)
+
+        # Header section
+        st.markdown("<h1 class='main-title'>🏆 The Sport Chatbot</h1>", unsafe_allow_html=True)
+        st.markdown("<h3 class='sub-title'>Using ESPN API</h3>", unsafe_allow_html=True)
+        st.markdown("""
+            <p class='description'>
+                Hey there! 👋 I can help you with information on Ice Hockey, Baseball, American Football, Soccer, and Basketball. 
+                With access to the ESPN API, I'm up to date with the latest details for these sports up until October 2024.
+            </p>
+            <p class='description'>
+                Got any general questions? Feel free to ask—I'll do my best to provide answers based on the information I've been trained on!
+            </p>
+        """, unsafe_allow_html=True)
+
+        # Add spacing
+        st.markdown("<br>", unsafe_allow_html=True)
+
+        # Initialize the pipeline
+        try:
+            with st.spinner("Loading resources..."):
+                rag = initialize_rag_pipeline()
+                
+                # Add a model health check
+                if not rag.check_model_health():
+                    st.error("Model initialization failed. Please try restarting the application.")
+                    return
+                    
+        except Exception as e:
+            logging.error(f"Initialization error: {str(e)}")
+            st.error("Unable to initialize the system. Please check if all required files are present.")
+            return
+
+        # Create columns for layout with golden ratio
+        col1, col2, col3 = st.columns([1, 6, 1])
         
-        # Centered button
-        if st.button("Get Answer"):
-            if query:
-                response_placeholder = st.empty()
-                try:
-                    response = rag.process_query(query, response_placeholder)
-                    print(f"Generated response: {response}")
-                except Exception as e:
-                    print(f"Query processing error: {str(e)}")
-                    response_placeholder.warning("Unable to process your question. Please try again.")
-            else:
-                st.warning("Please enter a question!")
-
-    # Footer with improved styling
-    st.markdown("<br><br>", unsafe_allow_html=True)
-    st.markdown("---")
-    st.markdown("""
-        <p style='text-align: center; color: #666666; padding: 1rem 0;'>
-            Powered by ESPN Data & Mistral AI 🚀
-        </p>
-    """, unsafe_allow_html=True)
+        with col2:
+            # Query input with label styling
+            query = st.text_input("What would you like to know about sports?")
+            
+            # Centered button
+            if st.button("Get Answer"):
+                if query:
+                    response_placeholder = st.empty()
+                    try:
+                        response = rag.process_query(query, response_placeholder)
+                        logging.info(f"Generated response: {response}")
+                    except Exception as e:
+                        logging.error(f"Query processing error: {str(e)}")
+                        response_placeholder.warning("Unable to process your question. Please try again.")
+                else:
+                    st.warning("Please enter a question!")
+
+        # Footer
+        st.markdown("<br><br>", unsafe_allow_html=True)
+        st.markdown("---")
+        st.markdown("""
+            <p style='text-align: center; color: #666666; padding: 1rem 0;'>
+                Powered by ESPN Data & Mistral AI 🚀
+            </p>
+        """, unsafe_allow_html=True)
+
+    except Exception as e:
+        logging.error(f"Application error: {str(e)}")
+        st.error("An unexpected error occurred. Please check the logs and try again.")
 
 if __name__ == "__main__":
-    main()
+    try:
+        main()
+    except Exception as e:
+        logging.error(f"Application error: {str(e)}")
+        st.error("An unexpected error occurred. Please check the logs and try again.")