Upload 7 files

analyzer.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python
+"""
+Evaluation script for Next Word Prediction model.
+Loads the trained model and SentencePiece model,
+prepares the validation dataset, and computes:
+    - Perplexity (using average loss)
+    - Top-k Accuracy (e.g., top-3 accuracy)
+Usage:
+    python evaluate_next_word.py --data_path data.csv \
+         --sp_model_path spm.model --model_save_path best_model.pth \
+         [--batch_size 512] [--top_k 3]
+"""
+
+import os
+import sys
+import math
+import argparse
+import logging
+import pandas as pd
+
+import torch
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.utils.rnn import pad_sequence
+
+import sentencepiece as spm
+
+# ---------------------- Logging Configuration ----------------------
+logging.basicConfig(
+    stream=sys.stdout,
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+# ---------------------- Dataset Definition ----------------------
+class NextWordSPDataset(Dataset):
+    def __init__(self, sentences, sp):
+        self.sp = sp
+        self.samples = []
+        self.prepare_samples(sentences)
+    
+    def prepare_samples(self, sentences):
+        for sentence in sentences:
+            token_ids = self.sp.encode(sentence.strip(), out_type=int)
+            # For each sentence, create (input_sequence, target) pairs.
+            for i in range(1, len(token_ids)):
+                self.samples.append((
+                    torch.tensor(token_ids[:i], dtype=torch.long),
+                    torch.tensor(token_ids[i], dtype=torch.long)
+                ))
+    
+    def __len__(self):
+        return len(self.samples)
+    
+    def __getitem__(self, idx):
+        return self.samples[idx]
+
+def sp_collate_fn(batch):
+    inputs, targets = zip(*batch)
+    padded_inputs = pad_sequence(inputs, batch_first=True, padding_value=0)
+    targets = torch.stack(targets)
+    return padded_inputs, targets
+
+# ---------------------- Model Definition ----------------------
+class LSTMNextWordModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, dropout, fc_dropout=0.3):
+        super(LSTMNextWordModel, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
+        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers, 
+                            batch_first=True, dropout=dropout)
+        self.layer_norm = nn.LayerNorm(hidden_dim)
+        self.dropout = nn.Dropout(fc_dropout)
+        self.fc1 = nn.Linear(hidden_dim, hidden_dim // 2)
+        self.fc2 = nn.Linear(hidden_dim // 2, vocab_size)
+
+    def forward(self, x):
+        emb = self.embedding(x)
+        output, _ = self.lstm(emb)
+        last_output = output[:, -1, :]
+        norm_output = self.layer_norm(last_output)
+        norm_output = self.dropout(norm_output)
+        fc1_out = torch.relu(self.fc1(norm_output))
+        fc1_out = self.dropout(fc1_out)
+        logits = self.fc2(fc1_out)
+        return logits
+
+# ---------------------- Evaluation Functions ----------------------
+def evaluate_perplexity(model, dataloader, criterion, device):
+    model.eval()
+    total_loss = 0.0
+    total_samples = 0
+    with torch.no_grad():
+        for inputs, targets in dataloader:
+            inputs = inputs.to(device)
+            targets = targets.to(device)
+            logits = model(inputs)
+            loss = criterion(logits, targets)
+            total_loss += loss.item() * inputs.size(0)
+            total_samples += inputs.size(0)
+    avg_loss = total_loss / total_samples
+    perplexity = math.exp(avg_loss)
+    return perplexity
+
+def evaluate_topk_accuracy(model, dataloader, k, device):
+    model.eval()
+    correct = 0
+    total = 0
+    with torch.no_grad():
+        for inputs, targets in dataloader:
+            inputs = inputs.to(device)
+            targets = targets.to(device)
+            logits = model(inputs)
+            # Get top-k predictions for each sample
+            _, topk_indices = torch.topk(logits, k, dim=-1)
+            for i in range(len(targets)):
+                if targets[i] in topk_indices[i]:
+                    correct += 1
+            total += targets.size(0)
+    accuracy = correct / total if total > 0 else 0
+    return accuracy
+
+# ---------------------- Main Evaluation Routine ----------------------
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logging.info("Using device: %s", device)
+
+    # Load SentencePiece model
+    if not os.path.exists(args.sp_model_path):
+        logging.error("SentencePiece model not found at %s", args.sp_model_path)
+        sys.exit(1)
+    sp = spm.SentencePieceProcessor()
+    sp.load(args.sp_model_path)
+    logging.info("Loaded SentencePiece model from %s", args.sp_model_path)
+    
+    # Load data and prepare validation set
+    if not os.path.exists(args.data_path):
+        logging.error("Data CSV file not found at %s", args.data_path)
+        sys.exit(1)
+    df = pd.read_csv(args.data_path)
+    if 'data' not in df.columns:
+        logging.error("CSV file must contain a 'data' column.")
+        sys.exit(1)
+    sentences = df['data'].tolist()
+    # Use a portion for validation. Here, we assume last 10% is validation.
+    split_index = int(len(sentences) * 0.9)
+    valid_sentences = sentences[split_index:]
+    logging.info("Validation sentences: %d", len(valid_sentences))
+    
+    valid_dataset = NextWordSPDataset(valid_sentences, sp)
+    valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size,
+                              shuffle=False, collate_fn=sp_collate_fn)
+    
+    # Initialize model. You may need to adjust these parameters to match your training.
+    vocab_size = sp.get_piece_size()
+    embed_dim = args.embed_dim
+    hidden_dim = args.hidden_dim
+    num_layers = args.num_layers
+    dropout = args.dropout
+    model = LSTMNextWordModel(vocab_size, embed_dim, hidden_dim, num_layers, dropout)
+    model.to(device)
+    
+    # Load the trained model weights
+    if not os.path.exists(args.model_save_path):
+        logging.error("Model checkpoint not found at %s", args.model_save_path)
+        sys.exit(1)
+    model.load_state_dict(torch.load(args.model_save_path, map_location=device))
+    logging.info("Loaded model checkpoint from %s", args.model_save_path)
+    
+    # Define the loss criterion.
+    # Note: If you used label smoothing during training, you can reuse that here.
+    class LabelSmoothingLoss(nn.Module):
+        def __init__(self, smoothing=0.1):
+            super(LabelSmoothingLoss, self).__init__()
+            self.smoothing = smoothing
+
+        def forward(self, pred, target):
+            confidence = 1.0 - self.smoothing
+            vocab_size = pred.size(1)
+            one_hot = torch.zeros_like(pred).scatter(1, target.unsqueeze(1), 1)
+            smoothed_target = one_hot * confidence + self.smoothing / (vocab_size - 1)
+            log_prob = torch.log_softmax(pred, dim=-1)
+            loss = -(smoothed_target * log_prob).sum(dim=1).mean()
+            return loss
+
+    criterion = LabelSmoothingLoss(smoothing=args.label_smoothing)
+    
+    # Evaluate perplexity and top-k accuracy
+    val_perplexity = evaluate_perplexity(model, valid_loader, criterion, device)
+    topk_accuracy = evaluate_topk_accuracy(model, valid_loader, args.top_k, device)
+    
+    logging.info("Validation Perplexity: %.4f", val_perplexity)
+    logging.info("Top-%d Accuracy: %.4f", args.top_k, topk_accuracy)
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Evaluate Next Word Prediction Model")
+    parser.add_argument('--data_path', type=str, default='data.csv', help="Path to CSV file with a 'data' column")
+    parser.add_argument('--sp_model_path', type=str, default='spm.model', help="Path to the SentencePiece model file")
+    parser.add_argument('--model_save_path', type=str, default='best_model.pth', help="Path to the trained model checkpoint")
+    parser.add_argument('--batch_size', type=int, default=512, help="Batch size for evaluation")
+    parser.add_argument('--top_k', type=int, default=3, help="Top-k value for computing accuracy")
+    # Model hyperparameters (should match those used in training)
+    parser.add_argument('--embed_dim', type=int, default=256, help="Embedding dimension")
+    parser.add_argument('--hidden_dim', type=int, default=256, help="Hidden dimension")
+    parser.add_argument('--num_layers', type=int, default=2, help="Number of LSTM layers")
+    parser.add_argument('--dropout', type=float, default=0.3, help="Dropout rate")
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help="Label smoothing factor")
+    
+    args = parser.parse_args()
+    main(args)
+

app.py

@@ -0,0 +1,70 @@
+import streamlit as st
+import torch
+import sentencepiece as spm
+
+# ---------------------- Model & SentencePiece Loading ----------------------
+@st.cache_resource
+def load_model():
+    """Load the TorchScript model for inference."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model = torch.jit.load("best_model_scripted.pt", map_location=device)
+    model.to(device)
+    return model, device
+
+@st.cache_resource
+def load_sp_model():
+    """Load the SentencePiece model."""
+    sp = spm.SentencePieceProcessor()
+    sp.load("spm.model")
+    return sp
+
+# ---------------------- Prediction Function ----------------------
+def predict_next_words(model, sp, device, text, topk=3):
+    if not text.strip():
+        return []
+    token_ids = sp.encode(text.strip(), out_type=int)
+    if len(token_ids) == 0:
+        return []
+    input_seq = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = model(input_seq)
+        probabilities = torch.softmax(logits, dim=-1)
+        topk_result = torch.topk(probabilities, k=topk, dim=-1)
+        top_indices = topk_result.indices.squeeze(0).tolist()
+    predicted_pieces = [sp.id_to_piece(idx).lstrip("▁") for idx in top_indices]
+    return predicted_pieces
+
+# ---------------------- Streamlit App Layout ----------------------
+def main():
+    st.title("Real-Time Next Word Prediction")
+    st.write(
+        """
+        Start typing your sentence below. When you finish a word (i.e. type a space at the end),
+        the app will suggest three possible next words. Click on a suggestion to auto-complete your sentence.
+        """
+    )
+    
+    model, device = load_model()
+    sp = load_sp_model()
+
+    if "input_text" not in st.session_state:
+        st.session_state.input_text = ""
+    
+    user_input = st.text_input("Enter your sentence:", st.session_state.input_text, key="text_input")
+    st.session_state.input_text = user_input
+
+    if user_input.endswith(" "):
+        predictions = predict_next_words(model, sp, device, user_input, topk=3)
+        if predictions:
+            st.markdown("### Predictions:")
+            cols = st.columns(len(predictions))
+            for i, word in enumerate(predictions):
+                if cols[i].button(word):
+                    st.session_state.input_text = user_input + word + " "
+                    st.rerun()  # This triggers the refresh correctly
+    else:
+        st.write("Type a space at the end of your sentence to get next-word suggestions.")
+
+if __name__ == "__main__":
+    main()
+

best_model.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64a7b488dfce765aa9e59aa16eba1353409db2fecbe7de66c6059ce5f9667433
+size 19748260

best_model_scripted.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:80ac9a35fe8c8f1bc0f2cde2d9fced1064b97cfbd3cc424c20bb36f902a455d7
+size 19769323

next_word_prediction.py

@@ -0,0 +1,365 @@
+#!/usr/bin/env python
+"""
+Next Word Prediction using an LSTM model in PyTorch with advanced improvements.
+---------------------------------------------------------------------------------
+This script supports two modes:
+
+Training Mode (with --train):
+  - Loads data from CSV (must contain a 'data' column)
+  - Trains a SentencePiece model for subword tokenization (if not already available)
+  - Uses SentencePiece to tokenize text and create a Dataset of (input_sequence, target) pairs
+  - Builds and trains an LSTM-based model enhanced with:
+      * Extra fully connected layer (with ReLU and dropout)
+      * Layer Normalization after LSTM outputs
+      * Label Smoothing Loss for improved regularization
+      * Gradient clipping, Adam optimizer with weight decay, and ReduceLROnPlateau scheduling
+  - Saves training/validation loss graphs
+  - Converts and saves the model to TorchScript for production deployment
+
+Inference Mode (with --inference "Your sentence"):
+  - Loads the saved SentencePiece model and the TorchScript (or checkpoint) model
+  - Runs inference to predict the top 3 next words/subwords
+
+Usage:
+  Training mode:
+      python next_word_prediction.py --data_path data.csv --train
+  Inference mode:
+      python next_word_prediction.py --inference "How do you"
+"""
+
+import os
+import sys
+import argparse
+import logging
+import random
+import pickle
+from collections import Counter
+
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torch.nn.utils.rnn import pad_sequence
+
+# Import SentencePiece
+import sentencepiece as spm
+
+# ---------------------- Global Definitions ----------------------
+PAD_TOKEN = '<PAD>'  # For padding (id will be 0)
+UNK_TOKEN = '<UNK>'
+# We use SentencePiece so our tokens come from the trained model
+
+# Set up logging to stdout for Colab compatibility
+logging.basicConfig(
+    stream=sys.stdout,
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+
+# ---------------------- Label Smoothing Loss ----------------------
+class LabelSmoothingLoss(nn.Module):
+    def __init__(self, smoothing=0.1):
+        super(LabelSmoothingLoss, self).__init__()
+        self.smoothing = smoothing
+
+    def forward(self, pred, target):
+        confidence = 1.0 - self.smoothing
+        vocab_size = pred.size(1)
+        one_hot = torch.zeros_like(pred).scatter(1, target.unsqueeze(1), 1)
+        smoothed_target = one_hot * confidence + self.smoothing / (vocab_size - 1)
+        log_prob = torch.log_softmax(pred, dim=-1)
+        loss = -(smoothed_target * log_prob).sum(dim=1).mean()
+        return loss
+
+# ---------------------- SentencePiece Functions ----------------------
+def train_sentencepiece(corpus, model_prefix, vocab_size):
+    temp_file = "sp_temp.txt"
+    with open(temp_file, "w", encoding="utf-8") as f:
+        for sentence in corpus:
+            f.write(sentence.strip() + "\n")
+    spm.SentencePieceTrainer.train(
+        input=temp_file,
+        model_prefix=model_prefix,
+        vocab_size=vocab_size,
+        character_coverage=1.0,
+        model_type='unigram'
+    )
+    os.remove(temp_file)
+    logging.info("SentencePiece model trained and saved with prefix '%s'", model_prefix)
+
+def load_sentencepiece_model(model_path):
+    sp = spm.SentencePieceProcessor()
+    sp.load(model_path)
+    logging.info("Loaded SentencePiece model from %s", model_path)
+    return sp
+
+# ---------------------- Dataset using SentencePiece ----------------------
+class NextWordSPDataset(Dataset):
+    def __init__(self, sentences, sp):
+        logging.info("Initializing NextWordSPDataset with %d sentences", len(sentences))
+        self.sp = sp
+        self.samples = []
+        self.prepare_samples(sentences)
+        logging.info("Total samples generated: %d", len(self.samples))
+    
+    def prepare_samples(self, sentences):
+        for idx, sentence in enumerate(sentences):
+            token_ids = self.sp.encode(sentence.strip(), out_type=int)
+            for i in range(1, len(token_ids)):
+                self.samples.append((
+                    torch.tensor(token_ids[:i], dtype=torch.long),
+                    torch.tensor(token_ids[i], dtype=torch.long)
+                ))
+            if (idx + 1) % 1000 == 0:
+                logging.debug("Processed %d/%d sentences", idx + 1, len(sentences))
+    
+    def __len__(self):
+        return len(self.samples)
+    
+    def __getitem__(self, idx):
+        return self.samples[idx]
+
+def sp_collate_fn(batch):
+    inputs, targets = zip(*batch)
+    padded_inputs = pad_sequence(inputs, batch_first=True, padding_value=0)
+    targets = torch.stack(targets)
+    logging.debug("Batch collated: inputs shape %s, targets shape %s", padded_inputs.shape, targets.shape)
+    return padded_inputs, targets
+
+# ---------------------- Model Definition ----------------------
+class LSTMNextWordModel(nn.Module):
+    def __init__(self, vocab_size, embed_dim, hidden_dim, num_layers, dropout, fc_dropout=0.3):
+        super(LSTMNextWordModel, self).__init__()
+        self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
+        self.lstm = nn.LSTM(embed_dim, hidden_dim, num_layers=num_layers, 
+                            batch_first=True, dropout=dropout)
+        self.layer_norm = nn.LayerNorm(hidden_dim)
+        self.dropout = nn.Dropout(fc_dropout)
+        self.fc1 = nn.Linear(hidden_dim, hidden_dim // 2)
+        self.fc2 = nn.Linear(hidden_dim // 2, vocab_size)
+
+    def forward(self, x):
+        # Logging calls removed to allow TorchScript conversion.
+        emb = self.embedding(x)
+        output, _ = self.lstm(emb)
+        last_output = output[:, -1, :]
+        norm_output = self.layer_norm(last_output)
+        norm_output = self.dropout(norm_output)
+        fc1_out = torch.relu(self.fc1(norm_output))
+        fc1_out = self.dropout(fc1_out)
+        logits = self.fc2(fc1_out)
+        return logits
+
+# ---------------------- Training and Evaluation ----------------------
+def train_model(model, train_loader, valid_loader, optimizer, criterion, scheduler, device,
+                num_epochs, patience, model_save_path, clip_value=5):
+    best_val_loss = float('inf')
+    patience_counter = 0
+    train_losses = []
+    val_losses = []
+    logging.info("Starting training for %d epochs", num_epochs)
+
+    for epoch in range(num_epochs):
+        logging.info("Epoch %d started...", epoch + 1)
+        model.train()
+        total_loss = 0.0
+        for batch_idx, (inputs, targets) in enumerate(train_loader):
+            inputs = inputs.to(device)
+            targets = targets.to(device)
+            optimizer.zero_grad()
+            outputs = model(inputs)
+            loss = criterion(outputs, targets)
+            loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), clip_value)
+            optimizer.step()
+            total_loss += loss.item()
+            if (batch_idx + 1) % 50 == 0:
+                logging.debug("Epoch %d, Batch %d: Loss = %.4f", epoch + 1, batch_idx + 1, loss.item())
+        avg_train_loss = total_loss / len(train_loader)
+        train_losses.append(avg_train_loss)
+        logging.info("Epoch %d training completed. Avg Train Loss: %.4f", epoch + 1, avg_train_loss)
+        
+        model.eval()
+        total_val_loss = 0.0
+        with torch.no_grad():
+            for batch_idx, (inputs, targets) in enumerate(valid_loader):
+                inputs = inputs.to(device)
+                targets = targets.to(device)
+                outputs = model(inputs)
+                loss = criterion(outputs, targets)
+                total_val_loss += loss.item()
+                if (batch_idx + 1) % 50 == 0:
+                    logging.debug("Validation Epoch %d, Batch %d: Loss = %.4f", epoch + 1, batch_idx + 1, loss.item())
+        avg_val_loss = total_val_loss / len(valid_loader)
+        val_losses.append(avg_val_loss)
+        logging.info("Epoch %d validation completed. Avg Val Loss: %.4f", epoch + 1, avg_val_loss)
+        
+        scheduler.step(avg_val_loss)
+        
+        if avg_val_loss < best_val_loss:
+            best_val_loss = avg_val_loss
+            patience_counter = 0
+            torch.save(model.state_dict(), model_save_path)
+            logging.info("Checkpoint saved at epoch %d with Val Loss: %.4f", epoch + 1, avg_val_loss)
+        else:
+            patience_counter += 1
+            logging.info("No improvement in validation loss for %d consecutive epoch(s).", patience_counter)
+            if patience_counter >= patience:
+                logging.info("Early stopping triggered at epoch %d", epoch + 1)
+                break
+
+    plt.figure()
+    plt.plot(range(1, len(train_losses)+1), train_losses, label="Train Loss")
+    plt.plot(range(1, len(val_losses)+1), val_losses, label="Validation Loss")
+    plt.xlabel("Epoch")
+    plt.ylabel("Loss")
+    plt.legend()
+    plt.title("Training and Validation Loss")
+    plt.savefig("loss_graph.png")
+    logging.info("Loss graph saved as loss_graph.png")
+    
+    return train_losses, val_losses
+
+def predict_next_word(model, sentence, sp, device, topk=3):
+    """
+    Given a partial sentence, uses SentencePiece to tokenize and predicts the top k next words.
+    """
+    logging.info("Predicting top %d next words for input sentence: '%s'", topk, sentence)
+    model.eval()
+    token_ids = sp.encode(sentence.strip(), out_type=int)
+    logging.debug("Token IDs for prediction: %s", token_ids)
+    if len(token_ids) == 0:
+        logging.warning("No tokens found in input sentence.")
+        return []
+    input_seq = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0).to(device)
+    with torch.no_grad():
+        logits = model(input_seq)
+        probabilities = torch.softmax(logits, dim=-1)
+        topk_result = torch.topk(probabilities, k=topk, dim=-1)
+        top_indices = topk_result.indices.squeeze(0).tolist()
+    predicted_pieces = [sp.id_to_piece(idx) for idx in top_indices]
+    cleaned_predictions = [piece.lstrip("▁") for piece in predicted_pieces]
+    logging.info("Predicted top %d next words/subwords: %s", topk, cleaned_predictions)
+    return cleaned_predictions
+
+# ---------------------- Main Function ----------------------
+def main(args):
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    logging.info("Using device: %s", device)
+
+    # Inference-only mode
+    if args.inference is not None:
+        logging.info("Running in inference-only mode with input: '%s'", args.inference)
+        if not os.path.exists(args.sp_model_path):
+            logging.error("SentencePiece model not found at %s. Cannot run inference.", args.sp_model_path)
+            return
+        sp = load_sentencepiece_model(args.sp_model_path)
+        if os.path.exists(args.scripted_model_path):
+            logging.info("Loading TorchScript model from %s", args.scripted_model_path)
+            model = torch.jit.load(args.scripted_model_path, map_location=device)
+        elif os.path.exists(args.model_save_path):
+            logging.info("Loading model checkpoint from %s", args.model_save_path)
+            model = LSTMNextWordModel(vocab_size=sp.get_piece_size(),
+                                      embed_dim=args.embed_dim,
+                                      hidden_dim=args.hidden_dim,
+                                      num_layers=args.num_layers,
+                                      dropout=args.dropout,
+                                      fc_dropout=0.3)
+            model.load_state_dict(torch.load(args.model_save_path, map_location=device))
+            model.to(device)
+        else:
+            logging.error("No model checkpoint found. Exiting.")
+            return
+        predictions = predict_next_word(model, args.inference, sp, device, topk=1)
+        logging.info("Input: '%s' -> Predicted next words: %s", args.inference, predictions)
+        return
+
+    # Training mode
+    logging.info("Loading data from %s...", args.data_path)
+    df = pd.read_csv(args.data_path)
+    if 'data' not in df.columns:
+        logging.error("CSV file must contain a 'data' column. Exiting.")
+        return
+    sentences = df['data'].tolist()
+    logging.info("Total sentences loaded: %d", len(sentences))
+    
+    if not os.path.exists(args.sp_model_path):
+        logging.info("SentencePiece model not found at %s. Training new model...", args.sp_model_path)
+        train_sentencepiece(sentences, args.sp_model_prefix, args.vocab_size)
+    sp = load_sentencepiece_model(args.sp_model_path)
+    
+    train_sentences = sentences[:int(len(sentences) * args.train_split)]
+    valid_sentences = sentences[int(len(sentences) * args.train_split):]
+    train_dataset = NextWordSPDataset(train_sentences, sp)
+    valid_dataset = NextWordSPDataset(valid_sentences, sp)
+    train_loader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=sp_collate_fn)
+    valid_loader = DataLoader(valid_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=sp_collate_fn)
+    logging.info("DataLoaders created: %d training batches, %d validation batches",
+                 len(train_loader), len(valid_loader))
+    
+    vocab_size = sp.get_piece_size()
+    model = LSTMNextWordModel(vocab_size=vocab_size,
+                              embed_dim=args.embed_dim,
+                              hidden_dim=args.hidden_dim,
+                              num_layers=args.num_layers,
+                              dropout=args.dropout,
+                              fc_dropout=0.3)
+    model.to(device)
+    
+    criterion = LabelSmoothingLoss(smoothing=args.label_smoothing)
+    optimizer = optim.Adam(model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=1, verbose=True)
+    logging.info("Loss function, optimizer, and scheduler initialized.")
+    
+    if args.train:
+        logging.info("Training mode is ON.")
+        if os.path.exists(args.model_save_path):
+            logging.info("Existing checkpoint found at %s. Loading weights...", args.model_save_path)
+            model.load_state_dict(torch.load(args.model_save_path, map_location=device))
+        else:
+            logging.info("No checkpoint found. Training from scratch.")
+        train_losses, val_losses = train_model(model, train_loader, valid_loader, optimizer, criterion,
+                                                scheduler, device, args.num_epochs, args.patience,
+                                                args.model_save_path)
+        scripted_model = torch.jit.script(model)
+        scripted_model.save(args.scripted_model_path)
+        logging.info("Model converted to TorchScript and saved to %s", args.scripted_model_path)
+    else:
+        logging.info("Training flag not set. Skipping training and running inference demo.")
+        if not os.path.exists(args.model_save_path):
+            logging.error("No model checkpoint found. Exiting.")
+            return
+    
+
+# ---------------------- Entry Point ----------------------
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Next Word Prediction using LSTM in PyTorch with SentencePiece and advanced techniques")
+    parser.add_argument('--data_path', type=str, default='data.csv', help="Path to CSV file with a 'data' column (required for training)")
+    parser.add_argument('--vocab_size', type=int, default=10000, help="Vocabulary size for SentencePiece")
+    parser.add_argument('--train_split', type=float, default=0.9, help="Fraction of data to use for training")
+    parser.add_argument('--batch_size', type=int, default=512, help="Batch size for training")
+    parser.add_argument('--embed_dim', type=int, default=256, help="Dimension of word embeddings")
+    parser.add_argument('--hidden_dim', type=int, default=256, help="Hidden dimension for LSTM")
+    parser.add_argument('--num_layers', type=int, default=2, help="Number of LSTM layers")
+    parser.add_argument('--dropout', type=float, default=0.3, help="Dropout rate in LSTM")
+    parser.add_argument('--learning_rate', type=float, default=0.001, help="Learning rate for optimizer")
+    parser.add_argument('--weight_decay', type=float, default=1e-5, help="Weight decay (L2 regularization) for optimizer")
+    parser.add_argument('--num_epochs', type=int, default=25, help="Number of training epochs")
+    parser.add_argument('--patience', type=int, default=5, help="Early stopping patience")
+    parser.add_argument('--label_smoothing', type=float, default=0.1, help="Label smoothing factor")
+    parser.add_argument('--model_save_path', type=str, default='best_model.pth', help="Path to save the best model checkpoint")
+    parser.add_argument('--scripted_model_path', type=str, default='best_model_scripted.pt', help="Path to save the TorchScript model")
+    parser.add_argument('--sp_model_prefix', type=str, default='spm', help="Prefix for SentencePiece model files")
+    parser.add_argument('--sp_model_path', type=str, default='spm.model', help="Path to load/save the SentencePiece model")
+    parser.add_argument('--seed', type=int, default=42, help="Random seed for reproducibility")
+    parser.add_argument('--train', action='store_true', help="Flag to enable training mode. If not set, runs inference/demo using saved checkpoint.")
+    parser.add_argument('--inference', type=str, default=None, help="Input sentence for inference-only mode")
+    
+    args, unknown = parser.parse_known_args()
+    logging.info("Arguments parsed: %s", args)
+    main(args)

spm.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fe3060038cf9883da1a90d9a4770b57e82c537903000dcb7c07cee5acd7e68e8
+size 411288