main.py

import os
import json 
import random 

import nltk
import numpy as np 

import torch 
import torch.nn as nn
import torch.nn.functional as f
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

class ChatBotModel(nn.Module):
    """
    A simple feed-forward neural network for chatbot intent classification.
    
    This model takes a fixed-size input vector (e.g. sentence embedding, 
    bag-of-words, or TF-IDF features) and predicts probabilities over 
    possible intents/classes (e.g. greeting, goodbye, ask_price, etc.).
    """
    def __init__(self, input_size, output_size):
        """
        Initialize the chatbot model.
        
        Args:
            input_size (int): Dimension of the input feature vector 
            output_size (int): Number of possible output classes/intents
        """
        super(ChatBotModel, self).__init__()

        self.fc1 = nn.Linear(input_size, 128)
        self.fc2 = nn.Linear(128, 64)
        self.fc3 = nn.Linear(64, output_size) # produces raw scores for each class
        self.relu = nn.ReLU() # introduces non-linearity
        self.dropout = nn.Dropout(0.5) # randomly zeros out 50% of features during training to prevent overfitting

    def forward(self, x):
        x = self.relu(self.fc1(x))
        x = self.dropout(x)
        x = self.relu(self.fc2(x))
        x = self.fc3(x)

        return x 
    
class ChatbotAssistant:
    
    def __init__(self, intents_path, function_mappings = None):
        self.model = None 
        self.intents_path = intents_path

        self.documents = [] # List of all processed sentences (patterns)
        self.vocabulary = [] # Unique words collected from all patterns (after tokenization)
        self.intents = [] # List of intent tags for each document
        self.intents_responses = {} # Dictionary: intent_tag → list of possible responses

        self.function_mappings = function_mappings if function_mappings is not None else {}

        self.X = None # Feature matrix -> Shape: (num_documents, vocab_size)
        self.y = None # Label vector -> Shape: (num_documents,)

        self.parse_intents()
        self.prepare_data()

    @staticmethod
    def tokenize_and_lemmatize(text):
        """
        Tokenize input text and apply lemmatization

        Examples:
            "Running dogs are better!" → ['run', 'dog', 'be', 'good']
        
        Returns:
            list[str]: List of lemmatized, lowercase tokens
        """
        lemmatizer = nltk.WordNetLemmatizer()

        words = nltk.word_tokenize(text)
        # Lowercase + lemmatize (e.g. "running" → "run", "better" → "good")
        words = [lemmatizer.lemmatize(word.lower()) for word in words]

        return words

    def bag_of_words(self, words):
        """
        Convert a list of words into a bag-of-words vector (binary presence vector).
        
        Example:
            vocabulary = ['hello', 'bye', 'thanks']
            words = ['hello', 'world']
            → [1, 0, 0]
        
        Returns:
            list[int]: Binary vector of length = len(self.vocabulary)
        """
        return [1 if word in words else 0 for word in self.vocabulary]
    
    def parse_intents(self):
        """
        Load the intents JSON file and process it into internal data structures.
        - self.vocabulary:   sorted list of ALL unique lemmatized words from patterns
        - self.intents:      list of unique intent tags (e.g. ['greeting', 'goodbye', ...])
        - self.documents:    list of tuples → [(tokenized_pattern_words, intent_tag), ...]
        - self.intents_responses: dict → {tag: [possible response strings]}
        """

        if os.path.exists(self.intents_path):
            with open(self.intents_path, 'r') as f:
                intents_data = json.load(f)

            for intent in intents_data['intents']:
                if intent['tag'] not in self.intents:
                    self.intents.append(intent['tag'])
                    self.intents_responses[intent['tag']] = intent['responses']

                for pattern in intent['patterns']:
                    pattern_words = self.tokenize_and_lemmatize(pattern)
                    self.vocabulary.extend(pattern_words)
                    self.documents.append((pattern_words, intent['tag']))

                self.vocabulary = sorted(set(self.vocabulary))

    def prepare_data(self):
        """
        Convert all collected text examples (documents) into numeric training data
        Output:
        - self.X : NumPy array of shape (n_samples, vocab_size) → bag-of-words features
        - self.y : NumPy array of shape (n_samples,)         → integer class labels
        """
        if not self.documents:
            raise ValueError("No documents available. Ensure parse_intents() has been called and intents file is valid.")
        bags = []
        indices = []

        for document in self.documents:
            words = document[0]
            intent_tag = document[1]

            bag = self.bag_of_words(words)
            intent_index = self.intents.index(intent_tag)

            bags.append(bag)
            indices.append(intent_index)

        self.X = np.array(bags)
        self.y = np.array(indices)
    
    def train_model(self, batch_size, learning_rate, epochs):
        """
        Train the neural network using the prepared data (self.X and self.y).
        
        Args:
            batch_size (int):     How many examples to process before updating weights
            learning_rate (float): Step size for weight updates (too big → unstable, too small → slow)
            epochs (int):         How many full passes through the entire dataset
        """
        if self.X is None or self.y is None:
            raise ValueError("Data must be prepared before training. Call prepare_data() first.")
        
        X_tensor = torch.tensor(self.X, dtype=torch.float32)
        y_tensor = torch.tensor(self.y, dtype=torch.long)

        dataset = TensorDataset(X_tensor, y_tensor)
        loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

        # input_size  = number of words in vocabulary
        # output_size = number of different intents
        self.model = ChatBotModel(self.X.shape[1], len(self.intents))

        criterion = nn.CrossEntropyLoss()

        optimizer  = optim.Adam(self.model.parameters(), lr=learning_rate)

        for epoch in range(epochs):
            running_loss = 0.0
            num_batches = 0

            for batch_X, batch_y in loader:
                optimizer.zero_grad() # Clear old gradients
                outputs = self.model(batch_X)
                loss = criterion(outputs, batch_y)
                loss.backward()
                optimizer.step()
                running_loss += loss
                num_batches += 1
            
            avg_loss = running_loss / num_batches
            print(f"Epoch {epoch+1}: Loss: {running_loss / avg_loss:.4f}")
        
        print("Training finished.")
    
    def save_model(self, model_path, dimension_path):
        """
        Save the trained model weights and the required dimensions to disk.
        
        Two files are created:
        1. model_path (.pth)        → contains the actual learned weights/parameters
        2. dimension_path (.json)   → contains input_size and output_size
        
        Args:
            model_path (str):      File path where to save the PyTorch model weights
                                (usually ends with .pth or .pt)
            dimension_path (str):  File path where to save the dimensions JSON
        """
        if self.model is None: 
            raise ValueError("Model must be trained or loaded...")
        torch.save(self.model.state_dict(), model_path)
        
        if self.X is None :
            raise ValueError("Data must be prepared before training. Call prepare_data() first.")
        
        dimensions = {
            'input_size': self.X.shape[1],     # = len(vocabulary) = number of input features
            'output_size': len(self.intents)   # = number of possible intent classes
        }

        with open(dimension_path, 'w') as f:
            json.dump(dimensions, f, indent=4)
        
        print(f"Model weights saved to:     {model_path}")
        print(f"Dimensions saved to:        {dimension_path}")
        print(f"input_size:  {dimensions['input_size']}")
        print(f"output_size: {dimensions['output_size']}")
    
    def load_model(self, model_path, dimensions_path):
        """
        Load a previously saved model from disk so inference can be performed without retraining.
        
        Two files are required:
        1. dimensions_path (.json)  → contains 'input_size' and 'output_size'
        2. model_path (.pth)        → contains the actual trained weights (state_dict)
    
        Args:
            model_path (str):       Path to the saved PyTorch state_dict file (.pth)
            dimensions_path (str):  Path to the JSON file with input/output sizes
        """
        with open(dimensions_path, 'r') as f:
            dimensions = json.load(f)
        
        input_size  = dimensions['input_size']
        output_size = dimensions['output_size']

        self.model = ChatBotModel(input_size, output_size)
        state_dict = torch.load(model_path, weights_only=True)
        self.model.load_state_dict(state_dict)

        print(f"Model successfully loaded from: {model_path}")
        print(f"input_size:  {input_size}")
        print(f"output_size: {output_size}")
    
    def process_message(self, input_message):
        """
        Main entry point for processing a new user message and generating a reply.
        
        Full pipeline:
        1. Preprocess the input text (tokenize + lemmatize)
        2. Convert to bag-of-words vector using the training vocabulary
        3. Run inference through the neural network
        4. Select the most likely intent
        5. Execute any associated function (if registered in function_mappings)
        6. Return a random response from the possible replies for that intent
        
        Args:
            input_message (str): Raw text typed by the user
            
        Returns:
            str or None: A randomly chosen response string, or None if no responses exist
        """
        if self.model is None: 
            raise ValueError("Model must be trained or loaded...")
        
        words = self.tokenize_and_lemmatize(input_message)
        bag = self.bag_of_words(words)

        bag_tensor = torch.tensor([bag], dtype=torch.float32)

        self.model.eval()
        with torch.no_grad(): # no gradients needed during inference
            predictions = self.model(bag_tensor) # e.g. predictions/logits = tensor([[-1.23, -4.87,  0.12, -2.91,  6.84]])
        
        probabilities = torch.softmax(predictions, dim=1)[0] 

        predicted_class_index = int(torch.argmax(predictions, dim=1).item())
        predicted_intent = self.intents[predicted_class_index]  

        confidence = probabilities[predicted_class_index].item()
        
        if confidence < 0.50:
            return f"I'm only {confidence:.1%} sure... Did you mean something else?"

        if predicted_intent in self.function_mappings:
            func = self.function_mappings[predicted_intent]
            if callable(func):
                func()
        
        if self.intents_responses[predicted_intent]:
            return random.choice(self.intents_responses[predicted_intent])
        else:
            return None

def get_stocks():
    stocks = ['APPL', 'META', 'NVDA', 'GS', 'MSFT']

    print(random.sample(stocks, 3))

if __name__ == '__main__':
    # assistant = ChatbotAssistant('intents.json', function_mappings = {'stocks': get_stocks})
    # assistant.parse_intents()
    # assistant.prepare_data()
    # assistant.train_model(batch_size=8, learning_rate=0.001, epochs=100)
    #
    # assistant.save_model('chatbot_model.pth', 'dimensions.json')

    assistant = ChatbotAssistant('intents.json', function_mappings = {'stocks': get_stocks})
    assistant.parse_intents()
    assistant.load_model('chatbot_model.pth', 'dimensions.json')

    while True:
        message = input('Enter your message:')

        if message == '/quit':
            break

        print(assistant.process_message(message))