Test

dataset_lstm.py

@@ -76,6 +76,29 @@ class LSTMTokenizer:
             'input_ids': torch.tensor(ids, dtype=torch.long),
             'attention_mask': torch.tensor(attention_mask, dtype=torch.long)
         }
+    
+    def from_json(self, json_file):
+        """Load tokenizer from JSON file container word2idx dict"""
+        import json
+        with open(json_file, 'r') as f:
+            data = json.load(f)
+            self.word2idx = data
+            self.idx2word = {v: k for k, v in self.word2idx.items()}
+            # Add pad and unk tokens
+            self.word2idx['<pad>'] = 0
+            self.word2idx['<unk>'] = 1
+            self.idx2word[0] = '<pad>'
+            self.idx2word[1] = '<unk>'
+            # Update vocab size
+            self.vocab_size = len(self.word2idx)
+            logger.info(f"Loaded tokenizer with {self.vocab_size} tokens")
+
+    def save(self, json_file):
+        """Save tokenizer word2idx to JSON file"""
+        import json
+        with open(json_file, 'w') as f:
+            json.dump(self.word2idx, f, ensure_ascii=False, indent=4)
+            logger.info(f"Tokenizer saved to {json_file}")
 
 class LSTMDataset(Dataset):
     """Dataset for LSTM model"""
@@ -109,7 +132,7 @@ class LSTMDataset(Dataset):
 
 def prepare_lstm_data(data_path, text_col='text', label_col='label', 
                      max_vocab_size=30000, max_seq_length=512,
-                     val_split=0.1, test_split=0.1, batch_size=32, seed=42, return_datasets=False):
+                     val_split=0.1, test_split=0.1, batch_size=32, seed=42, return_datasets=False, return_tokenizer=False):
     """
     Load data and prepare for LSTM model
     """
@@ -162,8 +185,10 @@ def prepare_lstm_data(data_path, text_col='text', label_col='label',
     val_dataset = LSTMDataset(val_texts, val_labels, tokenizer)
     test_dataset = LSTMDataset(test_texts, test_labels, tokenizer)
 
-    if return_datasets:
+    if return_datasets and not return_tokenizer:
         return train_dataset, val_dataset, test_dataset, tokenizer.vocab_size
+    elif return_tokenizer and not return_datasets:
+        return train_dataset, val_dataset, test_dataset, tokenizer
     
     # Create data loaders
     if len(train_dataset.texts) == 0:
@@ -184,4 +209,7 @@ def prepare_lstm_data(data_path, text_col='text', label_col='label',
     else:
         test_loader = DataLoader(test_dataset, batch_size=batch_size)
     
-    return train_loader, val_loader, test_loader, tokenizer.vocab_size
+    if not return_tokenizer:
+        return train_loader, val_loader, test_loader, tokenizer.vocab_size
+    else:
+        return train_loader, val_loader, test_loader, tokenizer

distill_bert_to_lstm.py

@@ -3,6 +3,7 @@ import os
 import logging
 import torch
 import random
+import json
 import numpy as np
 from model import DocBERT
 from models.lstm_model import DocumentBiLSTM
@@ -127,15 +128,16 @@ def main():
     
     # Create LSTM data loaders
     logger.info("Creating LSTM data loaders...")
-    lstm_train_loader, lstm_val_loader, lstm_test_loader, vocab_size = prepare_lstm_data(
+    lstm_train_loader, lstm_val_loader, lstm_test_loader, tokenizer = prepare_lstm_data(
         args.data_path,
         text_col=args.text_column,
         label_col=args.label_column,
         max_vocab_size=30000,
         max_seq_length=args.max_seq_length,
         batch_size=args.batch_size,
-        seed=args.seed
+        seed=args.seed, return_tokenizer=True
     )
+    vocab_size = tokenizer.vocab_size
     
     logger.info(f"LSTM Vocabulary size: {vocab_size}")
     
@@ -187,6 +189,10 @@ def main():
     save_path = os.path.join(args.output_dir, "distilled_lstm_model.pth")
     trainer.train(epochs=args.epochs, save_path=save_path)
     
+    # Save the tokenizer
+    tokenizer_path = os.path.join(args.output_dir, "tokenizer.json")
+    tokenizer.save(tokenizer_path)
+
     logger.info("Knowledge distillation completed!")
 
 if __name__ == "__main__":

example_uses.md

@@ -17,5 +17,5 @@ python ./distill_bert_to_lstm.py --bert_model bert-base-uncased --bert_model_pat
 
 - Inference with distilled LSTM model (test_data.csv is test dataset with 4 classes like ag_news)
 ```
-python ./inference_lstm.py --model_path "./docbert_lstm/distilled_lstm_model.pth" --num_classes 4 --class_names "World" "Sports" "Business" "Science" --text_column "Description" --label_column "Class Index" --data_path "./test_data.csv" --inference_batch_limit 10
+python ./inference_lstm.py --model_path "./docbert_lstm/distilled_lstm_model.pth" --num_classes 4 --class_names "World" "Sports" "Business" "Science" --text_column "Description" --label_column "Class Index" --data_path "./test_data.csv" --inference_batch_limit 10 --tokenizer_path "./docbert_lstm/tokenizer.json"
 ```

inference_lstm.py

@@ -1,7 +1,7 @@
 from dataset_lstm import prepare_lstm_data, LSTMTokenizer, LSTMDataset
 from models.lstm_model import DocumentBiLSTM
 from sklearn import metrics
-import torch
+import torch, logging
 import numpy as np
 import argparse
 
@@ -9,6 +9,7 @@ if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Document Classification with LSTM")
     parser.add_argument("--data_path", type=str, required=True, help="Path to the dataset")
     parser.add_argument("--model_path", type=str, required=True, help="Path to the trained model")
+    parser.add_argument("--tokenizer_path", type=str, required=True, help="Path to the tokenizer")
     parser.add_argument("--max_seq_length", type=int, default=512, help="Maximum sequence length for LSTM")
     parser.add_argument("--batch_size", type=int, default=32, help="Batch size for training and evaluation")
     parser.add_argument("--num_classes", type=int, required=True, help="Number of classes for classification")
@@ -30,13 +31,26 @@ if __name__ == "__main__":
     # Set device
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Prepare data
+    model_state = torch.load(args.model_path)
+    
     tokenizer = LSTMTokenizer(max_seq_length=args.max_seq_length)
-    test_dataset, _, _, vocab_size = prepare_lstm_data(args.data_path,
+    #tokenizer.from_json(args.tokenizer_path)
+
+    _, _, _, tokenizer = prepare_lstm_data("./test_data.csv",
+                                    text_col=args.text_column,
+                                    label_col=args.label_column,
+                                    batch_size=args.batch_size,
+                                    max_seq_length=args.max_seq_length, val_split=0.0, test_split=0.0, return_datasets=False, return_tokenizer=True)
+    
+    tokenizer.save(args.tokenizer_path)
+
+    # Prepare data
+    
+    _, _, test_dataset, vocab_size = prepare_lstm_data(args.data_path,
                                     text_col=args.text_column,
                                     label_col=args.label_column,
                                     batch_size=args.batch_size,
-                                    max_seq_length=args.max_seq_length, val_split=0.0, test_split=0.0, return_datasets=True)
+                                    max_seq_length=args.max_seq_length, val_split=0.0, test_split=1.0, return_datasets=True)
 
     test_loader, _, _, vocab_size = prepare_lstm_data(args.data_path,
                                   text_col=args.text_column,
@@ -50,8 +64,6 @@ if __name__ == "__main__":
                            hidden_dim=args.hidden_dim,
                            n_layers=args.num_layers,
                            output_dim=args.num_classes)
-    
-    model_state = torch.load(args.model_path)
 
     if 'vocab_size' in model_state:
         vocab_size = model_state['vocab_size']
@@ -77,10 +89,6 @@ if __name__ == "__main__":
 
             outputs = model(input_ids)
 
-            if batch_count == 0 or batch_count == 1:
-                print(f"Labels: {labels}")
-                print(f"Outputs: {outputs}")
-
             predictions = torch.argmax(outputs, dim=1)
             all_predictions = np.append(all_predictions, predictions.cpu().numpy())