Upload app.py

app.py

@@ -0,0 +1,581 @@
+import os
+import warnings
+warnings.filterwarnings("ignore", category=UserWarning)
+
+import streamlit as st
+import torch
+import torch.nn.functional as F
+import re
+import requests
+#from dotenv import load_dotenv
+from embedding_processor import SentenceTransformerRetriever, process_data
+import pickle
+
+import os
+import warnings
+import json  # Add this import
+
+# Add at the top with other imports
+from llama_cpp import Llama
+import requests
+from tqdm import tqdm
+
+
+import logging
+import sys
+
+
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[logging.StreamHandler(sys.stdout)]
+)
+
+# Create necessary directories at startup
+for directory in ['models', 'ESPN_data', 'embeddings_cache']:
+    os.makedirs(directory, exist_ok=True)
+
+
+
+# Load environment variables
+#load_dotenv()
+
+# Add the new function here, right after imports and before API configuration
+
+@st.cache_data
+def load_from_drive(file_id: str):
+    """Load pickle file directly from Google Drive"""
+    try:
+        url = f"https://drive.google.com/uc?id={file_id}&export=download"
+        session = requests.Session()
+        response = session.get(url, stream=True)
+        
+        for key, value in response.cookies.items():
+            if key.startswith('download_warning'):
+                url = f"{url}&confirm={value}"
+                response = session.get(url, stream=True)
+                break
+        
+        content = response.content
+        print(f"Successfully downloaded {len(content)} bytes")
+        return pickle.loads(content)
+        
+    except Exception as e:
+        print(f"Detailed error: {str(e)}")
+        st.error(f"Error loading file from Drive: {str(e)}")
+        return None
+
+# Hugging Face API configuration
+
+# API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-v0.1"
+# headers = {"Authorization": f"Bearer HF_TOKEN"}
+#model_name = 'mistralai/Mistral-7B-v0.1'
+
+
+class RAGPipeline:
+
+    def __init__(self, data_folder: str, k: int = 5):
+        try:
+            self.data_folder = data_folder
+            self.k = k
+            self.retriever = SentenceTransformerRetriever()
+            self.documents = []
+            self.device = torch.device("cpu")
+            self.model_path = "mistral-7b-v0.1.Q4_K_M.gguf"
+            self.llm = None
+            self.initialize_model()  # Using the class method
+            
+        except Exception as e:
+            logging.error(f"Error in RAGPipeline initialization: {str(e)}")
+            raise
+
+    @st.cache_resource
+    def initialize_model(_self):  # Changed 'self' to '_self' for Streamlit caching
+        """Initialize the model with proper error handling and verification"""
+        try:
+            if not os.path.exists(_self.model_path):
+                st.info("Downloading model... This may take a while.")
+                direct_url = "https://huggingface.co/TheBloke/Mistral-7B-v0.1-GGUF/resolve/main/mistral-7b-v0.1.Q4_K_M.gguf"
+                _self.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 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 = {
+                "n_ctx": 2048,
+                "n_threads": 4,
+                "n_batch": 512,
+                "n_gpu_layers": 0,
+                "verbose": False
+            }
+        
+            _self.llm = Llama(model_path=_self.model_path, **llm_config)
+            st.success("Model loaded successfully!")
+            
+        except Exception as e:
+            st.error(f"Error initializing model: {str(e)}")
+            raise
+
+    def download_file_with_progress(self, url: str, filename: str):
+        """Download a file with progress bar using requests"""
+        response = requests.get(url, stream=True)
+        total_size = int(response.headers.get('content-length', 0))
+        
+        with open(filename, 'wb') as file, tqdm(
+            desc=filename,
+            total=total_size,
+            unit='iB',
+            unit_scale=True,
+            unit_divisor=1024,
+        ) as progress_bar:
+            for data in response.iter_content(chunk_size=1024):
+                size = file.write(data)
+                progress_bar.update(size)
+
+    # Alternative API call with streaming
+    def query_model(self, prompt: str) -> str:
+        """Query the local Llama model instead of API"""
+        try:
+            if self.llm is None:
+                raise RuntimeError("Model not initialized")
+            
+            # Generate response using Llama model
+            response = self.llm(
+                prompt,
+                max_tokens=512,
+                temperature=0.4,
+                top_p=0.95,
+                echo=False,
+                stop=["Question:", "\n\n"]
+            )
+        
+            # Check and extract response
+            if response and 'choices' in response and len(response['choices']) > 0:
+                text = response['choices'][0].get('text', '').strip()
+                return text
+            else:
+                raise ValueError("No valid response generated")
+            
+        except Exception as e:
+            logging.error(f"Error in query_model: {str(e)}")
+            raise
+    def preprocess_query(self, query: str) -> str:
+        """Clean and prepare the query"""
+        query = query.lower().strip()
+        query = re.sub(r'\s+', ' ', query)
+        return query
+        
+    def process_query(self, query: str, placeholder) -> str:
+        try:
+            # Preprocess query
+            query = self.preprocess_query(query)
+        
+            # Show retrieval status
+            status = placeholder.empty()
+            status.write("🔍 Finding relevant information...")
+        
+            # Get embeddings and search
+            query_embedding = self.retriever.encode([query])
+            similarities = F.cosine_similarity(query_embedding, self.retriever.doc_embeddings)
+            scores, indices = torch.topk(similarities, k=min(self.k, len(self.documents)))
+        
+            relevant_docs = [self.documents[idx] for idx in indices.tolist()]
+        
+            # Update status
+            status.write("💭 Generating response...")
+        
+            # Prepare context and prompt
+            context = "\n".join(relevant_docs[:3])  # Use top 3 most relevant docs
+            prompt = f"""Context information is below:
+            {context}
+        
+            Given the context above, please answer the following question:
+            {query}
+
+            Guidelines:
+            - If you cannot answer based on the context, say so politely
+            - Keep the response concise and focused
+            - Only include sports-related information
+            - No dates or timestamps in the response
+            - Use clear, natural language
+        
+            Answer:"""
+        
+            # Generate response
+            response_placeholder = placeholder.empty()
+        
+            try:
+                response_text = self.query_model(prompt)
+                if response_text:
+                    final_response = self.postprocess_response(response_text)
+                    response_placeholder.markdown(final_response)
+                    return final_response
+                else:
+                    message = "No relevant answer found. Please try rephrasing your question."
+                    response_placeholder.warning(message)
+                    return message
+                
+            except Exception as 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:
+            logging.error(f"Process error: {str(e)}")
+            message = "Something went wrong. Please try again with a different question."
+            placeholder.warning(message)
+            return message
+
+    def postprocess_response(self, response: str) -> str:
+        """Clean up the generated response"""
+        response = response.strip()
+        response = re.sub(r'\s+', ' ', response)
+        response = re.sub(r'\d{4}-\d{2}-\d{2}\s\d{2}:\d{2}:\d{2}(?:\+\d{2}:?\d{2})?', '', response)
+        return response
+
+
+#     def process_query(self, query: str, placeholder) -> str:
+#         try:
+#             # Preprocess query
+#             query = self.preprocess_query(query)
+            
+#             # Show retrieval status
+#             status = placeholder.empty()
+#             status.write("🔍 Finding relevant information...")
+            
+#             # Get embeddings and search using tensor operations
+#             query_embedding = self.retriever.encode([query])
+#             similarities = F.cosine_similarity(query_embedding, self.retriever.doc_embeddings)
+#             scores, indices = torch.topk(similarities, k=min(self.k, len(self.documents)))
+            
+#             # Print search results for debugging
+#             print("\nSearch Results:")
+#             for idx, score in zip(indices.tolist(), scores.tolist()):
+#                 print(f"Score: {score:.4f} | Document: {self.documents[idx][:100]}...")
+            
+#             relevant_docs = [self.documents[idx] for idx in indices.tolist()]
+            
+#             # Update status
+#             status.write("💭 Generating response...")
+            
+#             # Prepare context and prompt
+#             context = "\n".join(relevant_docs[:3])  # Only use top 3 most relevant docs
+#             prompt = f"""Answer this question using the given context. Be specific and detailed.
+
+# Context: {context}
+
+# Question: {query}
+
+# Answer (provide a complete, detailed response):"""
+            
+#             # Generate response
+#             response_placeholder = placeholder.empty()
+            
+#             try:
+#                 response = requests.post(
+#                     model_name,
+#                     #headers=headers,
+#                     json={
+#                         "inputs": prompt,
+#                         "parameters": {
+#                             "max_new_tokens": 1024,
+#                             "temperature": 0.5,
+#                             "top_p": 0.9,
+#                             "top_k": 50,
+#                             "repetition_penalty": 1.03,
+#                             "do_sample": True
+#                         }
+#                     },
+#                     timeout=30
+#                 ).json()
+                
+#                 if response and isinstance(response, list) and len(response) > 0:
+#                     generated_text = response[0].get('generated_text', '').strip()
+#                     if generated_text:
+#                         # Find and extract only the answer part
+#                         if "Answer:" in generated_text:
+#                             answer_part = generated_text.split("Answer:")[-1].strip()
+#                         elif "Answer (provide a complete, detailed response):" in generated_text:
+#                             answer_part = generated_text.split("Answer (provide a complete, detailed response):")[-1].strip()
+#                         else:
+#                             answer_part = generated_text.strip()
+                            
+#                         # Clean up the answer
+#                         answer_part = answer_part.replace("Context:", "").replace("Question:", "")
+                        
+#                         final_response = self.postprocess_response(answer_part)
+#                         response_placeholder.markdown(final_response)
+#                         return final_response
+                    
+#                 message = "No relevant answer found. Please try rephrasing your question."
+#                 response_placeholder.warning(message)
+#                 return message
+                    
+#             except Exception as e:
+#                 print(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)}")
+#             message = "Something went wrong. Please try again with a different question."
+#             placeholder.warning(message)
+#             return message
+def check_environment():
+    """Check if the environment is properly set up"""
+    # if not headers['Authorization']:
+    #     st.error("HUGGINGFACE_API_KEY environment variable not set!")
+    #     st.stop()
+    #     return False
+    
+    try:
+        import torch
+        import sentence_transformers
+        return True
+    except ImportError as e:
+        st.error(f"Missing required package: {str(e)}")
+        st.stop()
+        return False
+
+# @st.cache_resource
+# def initialize_rag_pipeline():
+#     """Initialize the RAG pipeline once"""
+#     data_folder = "ESPN_data"
+#     return RAGPipeline(data_folder)
+def check_space_requirements():
+    """Check if we're running on HF Space and have necessary resources"""
+    try:
+        # Check if we're on HF Space
+        is_space = os.environ.get('SPACE_ID') is not None
+        
+        if is_space:
+            # Check disk space
+            disk_space = os.statvfs('/')
+            free_space_gb = (disk_space.f_frsize * disk_space.f_bavail) / (1024**3)
+            
+            if free_space_gb < 10:  # Need at least 10GB free
+                st.warning(f"Low disk space: {free_space_gb:.1f}GB free")
+            
+            # Check if model exists
+            model_path = "mistral-7b-v0.1.Q4_K_M.gguf"
+            if not os.path.exists(model_path):
+                st.info("Model will be downloaded on first run")
+                
+            # Check if embeddings exist
+            if not os.path.exists('embeddings_cache/embeddings.pkl'):
+                st.info("Embeddings will be loaded from Drive")
+                
+        return True
+        
+    except Exception as e:
+        logging.error(f"Space requirements check failed: {str(e)}")
+        return False
+
+@st.cache_resource(show_spinner=False)
+def initialize_rag_pipeline():
+    """Initialize the RAG pipeline once"""
+    try:
+        # First check/create necessary directories
+        for directory in ['models', 'ESPN_data', 'embeddings_cache']:
+            os.makedirs(directory, exist_ok=True)
+            
+        # Load embeddings from Drive
+        drive_file_id = "1MuV63AE9o6zR9aBvdSDQOUextp71r2NN"
+        with st.spinner("Loading embeddings from Google Drive..."):
+            cache_data = load_from_drive(drive_file_id)
+            if cache_data is None:
+                st.error("Failed to load embeddings from Google Drive")
+                st.stop()
+        
+        # Initialize pipeline
+        data_folder = "ESPN_data"
+        rag = RAGPipeline(data_folder)  # This will automatically initialize the model through __init__
+        
+        # Store embeddings
+        rag.documents = cache_data['documents']
+        rag.retriever.store_embeddings(cache_data['embeddings'])
+        
+        st.success("System initialized successfully!")
+        return rag
+        
+    except Exception as e:
+        logging.error(f"Pipeline initialization error: {str(e)}")
+        st.error(f"Failed to initialize the system: {str(e)}")
+        raise
+        
+    except Exception as e:
+        logging.error(f"Pipeline initialization error: {str(e)}")
+        st.error(f"Failed to initialize the system: {str(e)}")
+        raise
+
+def main():
+    try:
+        # Environment check
+        if not check_environment() or not check_space_requirements():
+            return
+
+        # Session state for initialization status
+        if 'initialized' not in st.session_state:
+            st.session_state.initialized = False
+
+        # 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 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
+        if not st.session_state.initialized:
+            try:
+                with st.spinner("Loading resources..."):
+                    # Create necessary directories
+                    for directory in ['models', 'ESPN_data', 'embeddings_cache']:
+                        os.makedirs(directory, exist_ok=True)
+                    
+                    # Initialize RAG pipeline
+                    st.session_state.rag = initialize_rag_pipeline()
+                    st.session_state.initialized = True
+                    
+                st.success("System initialized successfully!")
+            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.")
+                st.stop()
+
+        # 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?")
+            
+            # Centered button
+            if st.button("Get Answer"):
+                if query:
+                    response_placeholder = st.empty()
+                    try:
+                        # Get response from RAG pipeline
+                        response = st.session_state.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 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 🚀<br>
+                <small>Running on Hugging Face Spaces</small>
+            </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__":
+    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.")