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 embedding_processor import SentenceTransformerRetriever, process_data import pickle import logging import sys from llama_cpp import Llama from tqdm import tqdm # Set page config first st.set_page_config( page_title="The Sport Chatbot", page_icon="🏆", layout="wide" ) # Configure logging logging.basicConfig( level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', handlers=[logging.StreamHandler(sys.stdout)] ) def download_file_with_progress(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) @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 @st.cache_resource(show_spinner=False) def load_llama_model(): """Load Llama model with caching""" try: model_path = "mistral-7b-v0.1.Q4_K_M.gguf" if not os.path.exists(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" download_file_with_progress(direct_url, model_path) llm_config = { "model_path": model_path, "n_ctx": 2048, "n_threads": 4, "n_batch": 512, "n_gpu_layers": 0, "verbose": False } model = Llama(**llm_config) st.success("Model loaded successfully!") return model except Exception as e: st.error(f"Error loading model: {str(e)}") raise def check_environment(): """Check if the environment is properly set up""" try: import torch import sentence_transformers return True except ImportError as e: st.error(f"Missing required package: {str(e)}") st.stop() return False class RAGPipeline: def __init__(self, data_folder: str, k: int = 5): self.data_folder = data_folder self.k = k self.retriever = SentenceTransformerRetriever() self.documents = [] self.device = torch.device("cpu") self.llm = load_llama_model() 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 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 query_model(self, prompt: str) -> str: """Query the local Llama model""" try: if self.llm is None: raise RuntimeError("Model not initialized") response = self.llm( prompt, max_tokens=512, temperature=0.4, top_p=0.95, echo=False, stop=["Question:", "\n\n"] ) 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 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]) 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 @st.cache_resource(show_spinner=False) def initialize_rag_pipeline(): """Initialize the RAG pipeline once""" try: # Create necessary directories os.makedirs("ESPN_data", 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) # Store embeddings rag.documents = cache_data['documents'] rag.retriever.store_embeddings(cache_data['embeddings']) 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 def main(): try: # Environment check if not check_environment(): return # Improved CSS styling st.markdown(""" """, unsafe_allow_html=True) # Header section st.markdown("

🏆 The Sport Chatbot

", unsafe_allow_html=True) st.markdown("

Using ESPN API

", unsafe_allow_html=True) st.markdown("""

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.

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!

""", unsafe_allow_html=True) # Initialize the pipeline if 'rag' not in st.session_state: with st.spinner("Loading resources..."): st.session_state.rag = initialize_rag_pipeline() # Create columns for layout col1, col2, col3 = st.columns([1, 6, 1]) with col2: # Query input query = st.text_input("What would you like to know about sports?") if st.button("Get Answer"): if query: response_placeholder = st.empty() try: 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 st.markdown("

", unsafe_allow_html=True) st.markdown("---") st.markdown("""

Powered by ESPN Data & Mistral AI 🚀

""", 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()