app.py

import gradio as gr
import logging
import time
from datetime import datetime
from typing import List, Optional, Tuple
import random
import nltk
# nltk.download('punkt') # Ensure punkt is downloaded if needed
from nltk.tokenize import sent_tokenize
import io
# from joblib import dump, load # Not used currently, commented out

# Import Hugging Face libraries
from transformers import AutoTokenizer, AutoModelForCausalLM
from sentence_transformers import SentenceTransformer
from datasets import load_dataset # Added for dataset loading

# Import ML/Data libraries
from sklearn.metrics.pairwise import cosine_similarity
import numpy as np

# Standard libraries
from concurrent.futures import ThreadPoolExecutor # Still useful for embedding generation

# Configure logging
logging.basicConfig(
    level=logging.INFO, format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__) # Use __name__ for logger

# Download NLTK data (optional, might not be strictly needed depending on chunking)
# try:
#     nltk.download('punkt', quiet=True)
# except Exception as e:
#     logger.warning(f"Failed to download NLTK data: {e}")

# --- Configuration ---
class Config:
    MODEL_NAME = "microsoft/DialoGPT-medium"
    EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
    MAX_TOKENS_RESPONSE = 150  # Max tokens for the generated response part
    MAX_TOKENS_INPUT = 800    # Max tokens allowed for context + query (adjust based on model limits)
    SIMILARITY_THRESHOLD = 0.3 # Adjusted threshold, tune as needed
    CHUNK_SIZE = 300          # Smaller chunk size might be better for dataset entries
    MAX_WORKERS = 5           # For parallel embedding generation
    DATASET_NAME = "acecalisto3/sspnc" # Hugging Face Dataset ID
    DATASET_SPLIT = "train"     # Which split of the dataset to use
    TEXT_COLUMNS = ["Subject", "Body"] # Columns containing text to index
    SOURCE_INFO_COLUMNS = ["Subject", "Date"] # Columns to use for source attribution

# --- Data Structures ---
class ResourceItem:
    def __init__(self, source_id: str, content: str, resource_type: str):
        self.source_id = source_id # Changed 'url' to 'source_id' for clarity
        self.content = content
        self.type = resource_type
        self.embedding = None # Overall embedding (optional now, as we use chunk embeddings)
        self.chunks = []
        self.chunk_embeddings = []

    def __str__(self):
        return f"ResourceItem(type={self.type}, source_id={self.source_id}, content_length={len(self.content)})"

    def create_chunks(self, chunk_size=Config.CHUNK_SIZE):
        """Split content into overlapping chunks using sentence tokenization for better boundaries"""
        if not self.content:
            logger.warning(f"Content is empty for source_id: {self.source_id}. Skipping chunk creation.")
            return

        try:
            sentences = sent_tokenize(self.content)
        except LookupError:
            logger.warning("NLTK 'punkt' tokenizer not found. Falling back to simple whitespace splitting. Consider running nltk.download('punkt')")
            # Fallback to word splitting if sentence tokenization fails
            words = self.content.split()
            overlap = chunk_size // 4
            for i in range(0, len(words), chunk_size - overlap):
                chunk = ' '.join(words[i : i + chunk_size])
                if chunk:
                    self.chunks.append(chunk)
            return
        except Exception as e:
             logger.error(f"Error during sentence tokenization for {self.source_id}: {e}. Skipping chunk creation.")
             return


        current_chunk = ""
        overlap_sentences = max(1, chunk_size // 100) # Overlap a few sentences
        last_sentences = []

        for sentence in sentences:
            # If adding the next sentence exceeds chunk size (considering words approx)
            if len((current_chunk + " " + sentence).split()) > chunk_size:
                if current_chunk: # Add the completed chunk
                    self.chunks.append(current_chunk.strip())
                # Start new chunk with overlap
                current_chunk = " ".join(last_sentences) + " " + sentence
            else:
                current_chunk += " " + sentence

            # Keep track of last sentences for overlap
            last_sentences.append(sentence)
            if len(last_sentences) > overlap_sentences:
                last_sentences.pop(0)

        # Add the last remaining chunk
        if current_chunk.strip():
            self.chunks.append(current_chunk.strip())

        if not self.chunks:
             logger.warning(f"No chunks created for source_id: {self.source_id}. Content might be too short or tokenization failed.")


# --- Chatbot Core Logic ---
class SchoolChatbot:
    def __init__(self):
        logger.info("Initializing SchoolChatbot...")
        self.setup_models()
        self.resources: List[ResourceItem] = []
        self.load_and_index_dataset() # Changed from crawl_and_index_resources

    def setup_models(self):
        try:
            logger.info("Setting up models...")
            # Consider adding device mapping if GPU is available: device_map="auto"
            self.tokenizer = AutoTokenizer.from_pretrained(Config.MODEL_NAME)
            self.model = AutoModelForCausalLM.from_pretrained(Config.MODEL_NAME)
            self.embedding_model = SentenceTransformer(Config.EMBEDDING_MODEL)
            # Ensure tokenizer has a padding token
            if self.tokenizer.pad_token is None:
                self.tokenizer.pad_token = self.tokenizer.eos_token
                self.model.config.pad_token_id = self.model.config.eos_token_id
            logger.info("Models setup completed successfully.")
        except Exception as e:
            logger.error(f"Failed to setup models: {e}")
            raise RuntimeError("Failed to initialize required models") from e

    def load_and_index_dataset(self):
        logger.info(f"Loading dataset: {Config.DATASET_NAME}, split: {Config.DATASET_SPLIT}")
        try:
            # Load the dataset
            dataset = load_dataset(Config.DATASET_NAME, split=Config.DATASET_SPLIT)
            logger.info(f"Dataset loaded successfully. Number of rows: {len(dataset)}")

            # Process dataset rows in parallel (for embedding generation)
            with ThreadPoolExecutor(max_workers=Config.MAX_WORKERS) as executor:
                futures = []
                for i, row in enumerate(dataset):
                    # Combine text from specified columns
                    text_content = " ".join([str(row[col]) for col in Config.TEXT_COLUMNS if row.get(col)])
                    text_content = text_content.strip() # Remove leading/trailing whitespace

                    # Create a source identifier
                    source_parts = [f"{col}: {row[col]}" for col in Config.SOURCE_INFO_COLUMNS if row.get(col)]
                    source_id = f"Dataset Entry {i} ({'; '.join(source_parts)})" # More informative ID

                    if not text_content:
                        logger.warning(f"Row {i} has no content in specified columns. Skipping.")
                        continue

                    # Submit the processing task
                    futures.append(executor.submit(self.process_and_store_resource, source_id, text_content, 'dataset_entry'))

                # Wait for all futures to complete and collect results
                for future in futures:
                    try:
                        result_item = future.result()
                        if result_item:
                            self.resources.append(result_item)
                    except Exception as e:
                        logger.error(f"Error processing dataset entry in thread: {e}")

            logger.info(f"Dataset processing completed. Indexed {len(self.resources)} resources.")

        except Exception as e:
            logger.error(f"Failed to load or process dataset {Config.DATASET_NAME}: {e}")
            # Decide if the app should continue without data or raise an error
            # raise RuntimeError("Failed to load data") from e # Option: halt if data fails

    def process_and_store_resource(self, source_id: str, text_data: str, resource_type: str) -> Optional[ResourceItem]:
        """Creates ResourceItem, chunks, and generates embeddings for a single data entry."""
        try:
            # Create resource item and split into chunks
            item = ResourceItem(source_id, text_data, resource_type)
            item.create_chunks()

            if not item.chunks:
                logger.warning(f"No chunks generated for {source_id}. Skipping storage.")
                return None

            # Generate embeddings for chunks (can be slow, hence the thread pool)
            chunk_embeddings_list = self.embedding_model.encode(item.chunks, show_progress_bar=False) # Batch encode
            item.chunk_embeddings = chunk_embeddings_list

            # Calculate average embedding (optional, might not be needed if only using chunk search)
            # if item.chunk_embeddings:
            #     item.embedding = np.mean(item.chunk_embeddings, axis=0)

            logger.debug(f"Processed resource: {source_id} (type={resource_type}), {len(item.chunks)} chunks.")
            return item # Return the processed item

        except Exception as e:
            logger.error(f"Error processing/storing resource {source_id}: {e}")
            return None # Return None on error

    # store_resource is now process_and_store_resource and called within the thread pool

    def find_best_matching_chunks(self, query: str, n_chunks: int = 3) -> List[Tuple[str, float, str]]:
        """Finds the most relevant text chunks based on semantic similarity."""
        if not self.resources:
            logger.warning("No resources loaded or indexed. Cannot find matches.")
            return []

        try:
            query_embedding = self.embedding_model.encode(query)
            all_chunks_with_scores = []

            for resource in self.resources:
                if not resource.chunks or not resource.chunk_embeddings:
                    continue # Skip resources with no chunks/embeddings

                # Calculate similarity between query and all chunks of the current resource
                similarities = cosine_similarity([query_embedding], resource.chunk_embeddings)[0]

                for chunk, score in zip(resource.chunks, similarities):
                    if score > Config.SIMILARITY_THRESHOLD:
                        all_chunks_with_scores.append((chunk, float(score), resource.source_id)) # Use source_id

            # Sort by similarity score (descending) and return top n
            all_chunks_with_scores.sort(key=lambda x: x[1], reverse=True)
            return all_chunks_with_scores[:n_chunks]

        except Exception as e:
            logger.error(f"Error finding matching chunks: {e}")
            return []

    def generate_response(self, user_input: str) -> str:
        """Generates a response based on user input and retrieved context."""
        try:
            # 1. Find relevant context chunks
            best_chunks = self.find_best_matching_chunks(user_input)

            if not best_chunks:
                logger.info(f"No relevant chunks found for query: '{user_input}'")
                return "I couldn't find specific information related to your question in the provided documents. Could you please rephrase or ask about a different topic?"

            # 2. Prepare context and source attribution
            context = "\n".join([chunk[0] for chunk in best_chunks])
            # Use source_id from the chunk tuple (index 2)
            source_ids = sorted(list(set(chunk[2] for chunk in best_chunks)))
            sources_text = "\n\nSources:\n" + "\n".join([f"- {sid}" for sid in source_ids])

            # 3. Prepare input for the language model
            # Ensure the input doesn't exceed model limits
            prompt_template = f"Based on the following information:\n{context}\n\nAnswer the question: {user_input}\nAnswer:"
            # prompt_template = f"Context: {context}\nUser: {user_input}\nAssistant:" # Alternative simpler prompt

            # 4. Tokenize and truncate if necessary
            input_ids = self.tokenizer.encode(prompt_template, return_tensors='pt', max_length=Config.MAX_TOKENS_INPUT, truncation=True)

            # 5. Generate response using the language model
            logger.info("Generating response with LLM...")
            output_sequences = self.model.generate(
                input_ids=input_ids,
                max_new_tokens=Config.MAX_TOKENS_RESPONSE, # Control length of *new* tokens
                pad_token_id=self.tokenizer.pad_token_id,
                eos_token_id=self.tokenizer.eos_token_id,
                temperature=0.7,
                top_p=0.9,
                do_sample=True,
                num_return_sequences=1 # Generate one response
            )

            # Decode the generated part of the response
            # response_text = self.tokenizer.decode(output_sequences[0], skip_special_tokens=True)
            # Decode only the newly generated tokens, excluding the prompt
            response_text = self.tokenizer.decode(output_sequences[0][input_ids.shape[-1]:], skip_special_tokens=True)


            # Basic post-processing (optional)
            response_text = response_text.strip()
            # Remove potential repetition of the question if the model includes it
            if user_input.lower() in response_text.lower()[:len(user_input)+10]:
                 response_text = response_text.split(user_input, 1)[-1].strip("? ")


            logger.info(f"Generated response (before sources): {response_text}")

            # 6. Combine response and sources
            full_response = response_text + sources_text
            return full_response

        except Exception as e:
            logger.exception(f"Error generating response: {e}") # Use logger.exception to include stack trace
            return "I apologize, but I encountered an error while processing your question. Please check the logs or try again later."

# --- Gradio Interface ---
def create_gradio_interface(chatbot: SchoolChatbot):
    """Creates and returns the Gradio web interface."""
    def respond(user_input: str) -> str:
        if not user_input:
            return "Please enter a question."
        # Add basic input sanitization if needed
        return chatbot.generate_response(user_input)

    interface = gr.Interface(
        fn=respond,
        inputs=gr.Textbox(
            label="Ask a Question",
            placeholder="Type your question about the school information...",
            lines=3, # Increased lines slightly
        ),
        outputs=gr.Textbox(
            label="Answer",
            placeholder="Response will appear here...",
            lines=10, # Increased lines for longer answers + sources
        ),
        title="School Information Chatbot (Dataset Powered)",
        description="Ask about information contained in the school dataset. The chatbot uses AI to find relevant details and generate answers.",
        examples=[ # Update examples based on dataset content
            ["What are the main subjects covered in the documents?"],
            ["Are there any mentions of specific events or dates?"],
            ["Summarize the key points about [topic from dataset]."]
        ],
        theme=gr.themes.Soft(),
        allow_flagging="never", # Changed from flagging_mode
        # Optional: Add feedback capabilities
        # feedback=["thumbs", "textbox"],
    )
    return interface

# --- Main Execution ---
if __name__ == "__main__":
    # Install necessary libraries if running for the first time
    # pip install gradio transformers sentence-transformers torch datasets scikit-learn nltk numpy beautifulsoup4 requests PyPDF2 icalendar fake-useragent joblib # Ensure all are installed
    print("Starting application...")
    try:
        # 1. Initialize the chatbot (loads models and data)
        school_chatbot = SchoolChatbot()

        # 2. Create the Gradio interface
        app_interface = create_gradio_interface(school_chatbot)

        # 3. Launch the interface
        print("Launching Gradio Interface...")
        app_interface.launch(
            server_name="0.0.0.0", # Accessible on the local network
            server_port=7860,
            share=False,           # Set to True to get a public link (use with caution)
            debug=False            # Set to True for more detailed Gradio logs (can be verbose)
        )
        print("Interface launched. Access it at http://localhost:7860 (or the relevant IP)")

    except ImportError as ie:
        logger.error(f"ImportError: {ie}. Make sure all required libraries are installed.")
        print(f"ImportError: {ie}. Please install the missing library (e.g., pip install {ie.name}).")
    except Exception as e:
        logger.critical(f"Failed to start the application: {e}", exc_info=True) # Log critical error with stack trace
        print(f"Critical error during startup: {e}")