Save vocab size with model_state_dict

inference_lstm.py

@@ -51,6 +51,11 @@ if __name__ == "__main__":
                            output_dim=args.num_classes)
     
     model_state = torch.load(args.model_path)
+
+    if 'vocab_size' in model_state:
+        vocab_size = model_state['vocab_size']
+        tokenizer.vocab_size = vocab_size
+
     if 'model_state_dict' in model_state:
         model.load_state_dict(model_state['model_state_dict'])
     else:

knowledge_distillation.py

@@ -164,7 +164,10 @@ class DistillationTrainer:
             if val_f1 > self.best_val_f1:
                 self.best_val_f1 = val_f1
                 self.best_model_state = self.student_model.state_dict().copy()
-                torch.save(self.student_model.state_dict(), save_path)
+                torch.save({
+                    'model_state_dict': self.student_model.state_dict(),
+                    'vocab_size': self.student_model.vocab_size,
+                }, save_path)
                 logger.info(f"New best model saved with validation F1: {val_f1:.4f}")
             
             logger.info(f"Epoch {epoch+1}/{epochs}: "