Fixed some model loading errors in train function

ThoughtVectors.py

@@ -1,3 +1,4 @@
+# thought_vectors.py
 import random
 import math
 import torch
@@ -440,20 +441,8 @@ class ThoughtVectors:
         
         if model and os.path.exists(model):
             print("loading model...")
-            checkpoint = torch.load(model, map_location=self.device)
-            self.sp.Load(checkpoint['spm_model_path'])
-            self.vocab_size = self.sp.GetPieceSize()
-            self.encoder = ThoughtEncoder(
-                vocab_size=self.vocab_size, d_model=d_model, max_thoughts=max_thoughts,
-                nhead=encoder_nhead, num_layers=encoder_layers, dropout=dropout,
-                termination_threshold=termination_threshold, max_len=self.max_len
-            ).to(self.device)
-            self.decoder = ThoughtDecoder(
-                vocab_size=self.vocab_size, d_model=d_model, num_layers=decoder_layers,
-                nhead=decoder_nhead, dropout=dropout, max_len=self.max_len
-            ).to(self.device)
-            self.encoder.load_state_dict(checkpoint['encoder_state_dict'])
-            self.decoder.load_state_dict(checkpoint['decoder_state_dict'])
+            loaded_instance = self.load(model)
+            self.__dict__.update(loaded_instance.__dict__)  # actually update self
             print(f"Loaded pre-existing model from {model}")
         else:
             print("building models")
@@ -507,100 +496,122 @@ class ThoughtVectors:
         patience_counter = 0
         
         print("Beginning training")
-        for epoch in range(num_epochs):
-            total_train_loss = 0.0
-            num_train_batches = (len(train_lazy_data) + batch_size - 1) // batch_size
-            
-            for i in range(0, len(train_lazy_data), batch_size):
-                raw_batch = train_lazy_data[i:i + batch_size]
-                tokenized_batch = [
-                    [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
-                    for group in raw_batch
-                ]
-                group_batch = group_collate_fn(tokenized_batch)
-                
-                batch_loss = 0.0
-                with autocast('cuda'):
-                    for group_tensor in group_batch:
-                        if group_tensor.shape[1] <= self.max_len and group_tensor.numel() > 0:
-                            group_tensor = group_tensor.to(self.device)
-                            force_single_vector = random.random() < single_vector_prob
-                            thought_vectors = self.encoder(group_tensor, force_single_vector)
-                            output_logits = self.decoder(thought_vectors, group_tensor[:, :-1])
-                            
-                            # Masked loss
-                            mask = (group_tensor[:, 1:] != self.sp.pad_id()).float()
-                            loss = criterion(output_logits.reshape(-1, self.vocab_size),
-                                             group_tensor[:, 1:].reshape(-1))
-                            loss = (loss * mask.reshape(-1)).sum() / mask.sum().clamp(min=1.0)
-                            if not force_single_vector:
-                                loss += length_penalty * thought_vectors.shape[1]
-                            batch_loss += loss / accum_steps
-                
-                scaler.scale(batch_loss).backward()
-                total_train_loss += batch_loss.item() * accum_steps
-                if (i // batch_size + 1) % accum_steps == 0:
-                    scaler.step(optimizer)
-                    scaler.update()
-                    optimizer.zero_grad()
-                    torch.cuda.empty_cache()
-                
-                print(f"Batch {i // batch_size + 1}/{num_train_batches} - Loss: {batch_loss.item():.4f}")
-                # Debugging
-                print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
-                print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n\n")
-            
-            avg_train_loss = total_train_loss / num_train_batches if num_train_batches > 0 else float('inf')
-            
-            # Validation with test_lazy_data
-            if test_lazy_data is not None:
-                self.encoder.eval()
-                self.decoder.eval()
-                total_val_loss = 0.0
-                num_val_batches = (len(test_lazy_data) + batch_size - 1) // batch_size
-                print(f"Starting validation with {num_val_batches} batches")
+        try:
+            for epoch in range(num_epochs):
+                total_train_loss = 0.0
+                num_train_batches = (len(train_lazy_data) + batch_size - 1) // batch_size
                 
-                with torch.no_grad():
-                    for i in range(0, len(test_lazy_data), batch_size):
-                        raw_batch = test_lazy_data[i:i + batch_size]
-                        tokenized_batch = [
-                            [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
-                            for group in raw_batch
-                        ]
-                        group_batch = group_collate_fn(tokenized_batch)
-                        
-                        batch_val_loss = 0.0
+                # training
+                for i in range(0, len(train_lazy_data), batch_size):
+                    raw_batch = train_lazy_data[i:i + batch_size]
+                    tokenized_batch = [
+                        [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
+                        for group in raw_batch
+                    ]
+                    group_batch = group_collate_fn(tokenized_batch)
+                    
+                    batch_loss = 0.0
+                    with autocast('cuda'):
                         for group_tensor in group_batch:
                             if group_tensor.shape[1] <= self.max_len and group_tensor.numel() > 0:
                                 group_tensor = group_tensor.to(self.device)
-                                thought_vectors = self.encoder(group_tensor)
+                                force_single_vector = random.random() < single_vector_prob
+                                thought_vectors = self.encoder(group_tensor, force_single_vector)
                                 output_logits = self.decoder(thought_vectors, group_tensor[:, :-1])
-                                val_loss = criterion(output_logits.reshape(-1, self.vocab_size),
-                                                     group_tensor[:, 1:].reshape(-1))
-                                val_loss += length_penalty * thought_vectors.shape[1]
-                                batch_val_loss += val_loss.item()
-                        total_val_loss += batch_val_loss
-                        print(f"VAL batch loss: {batch_val_loss:.4f}")
-                        print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
-                        print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n\n")
-                avg_val_loss = total_val_loss / num_val_batches if num_val_batches > 0 else float('inf')
-            else:
-                avg_val_loss = avg_train_loss
-            
-            print(f"Epoch {epoch + 1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f}\n\n")
-            
-            if avg_val_loss < best_val_loss:
-                best_val_loss = avg_val_loss
-                patience_counter = 0
-                self.save(save_path, spm_model_prefix, clean=False)
-            else:
-                patience_counter += 1
-                if patience_counter >= patience:
-                    print(f"Early stopping triggered after {epoch + 1} epochs.")
+                                
+                                # Masked loss
+                                mask = (group_tensor[:, 1:] != self.sp.pad_id()).float()
+                                loss = criterion(output_logits.reshape(-1, self.vocab_size),
+                                                 group_tensor[:, 1:].reshape(-1))
+                                loss = (loss * mask.reshape(-1)).sum() / mask.sum().clamp(min=1.0)
+                                if not force_single_vector:
+                                    loss += length_penalty * thought_vectors.shape[1]
+                                batch_loss += loss / accum_steps
+                    
+                    scaler.scale(batch_loss).backward()
+                    total_train_loss += batch_loss.item() * accum_steps
+                    if (i // batch_size + 1) % accum_steps == 0:
+                        scaler.step(optimizer)
+                        scaler.update()
+                        optimizer.zero_grad()
+                        torch.cuda.empty_cache()
+                    
+                    print(f"Batch {i // batch_size + 1}/{num_train_batches} - Loss: {batch_loss.item():.4f}")
+                    # Debugging
+                    print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
+                    print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n\n")
+                
+                avg_train_loss = total_train_loss / num_train_batches if num_train_batches > 0 else float('inf')
+                
+                # Validation
+                if test_lazy_data is not None:
+                    self.encoder.eval()
+                    self.decoder.eval()
+                    total_val_loss = 0.0
+                    num_val_batches = (len(test_lazy_data) + batch_size - 1) // batch_size
+                    print(f"Starting validation with {num_val_batches} batches")
+                    
+                    with torch.no_grad():
+                        for i in range(0, len(test_lazy_data), batch_size):
+                            raw_batch = test_lazy_data[i:i + batch_size]
+                            tokenized_batch = [
+                                [[self.sp.bos_id()] + self.sp.EncodeAsIds(s) + [self.sp.eos_id()] for s in group]
+                                for group in raw_batch
+                            ]
+                            group_batch = group_collate_fn(tokenized_batch)
+                            
+                            batch_val_loss = 0.0
+                            for group_tensor in group_batch:
+                                if group_tensor.shape[1] <= self.max_len and group_tensor.numel() > 0:
+                                    group_tensor = group_tensor.to(self.device)
+                                    thought_vectors = self.encoder(group_tensor)
+                                    output_logits = self.decoder(thought_vectors, group_tensor[:, :-1])
+                                    val_loss = criterion(output_logits.reshape(-1, self.vocab_size),
+                                                         group_tensor[:, 1:].reshape(-1))
+                                    val_loss += length_penalty * thought_vectors.shape[1]
+                                    batch_val_loss += val_loss.item()
+                            total_val_loss += batch_val_loss
+                            print(f"VAL batch loss: {batch_val_loss:.4f}")
+                            print(f" Input: {self.sp.DecodeIds(group_tensor[0].tolist())}")
+                            print(f"Output: {self.sp.DecodeIds(output_logits.argmax(-1)[0].tolist())}\n\n")
+                    avg_val_loss = total_val_loss / num_val_batches if num_val_batches > 0 else float('inf')
+                else:
+                    avg_val_loss = avg_train_loss
+                
+                print(f"Epoch {epoch + 1}/{num_epochs} - Train Loss: {avg_train_loss:.4f} - Val Loss: {avg_val_loss:.4f}\n\n")
+                
+                if avg_val_loss < best_val_loss:
+                    best_val_loss = avg_val_loss
+                    patience_counter = 0
+                    self.save(save_path, spm_model_prefix, clean=False)
+                else:
+                    self.load(save_path)
+                    self.save(save_path, spm_model_prefix, clean=True)
+                    patience_counter += 1
+                    if patience_counter >= patience:
+                        print(f"Early stopping triggered after {epoch + 1} epochs.")
+                        break
+                
+                scheduler.step()
+        # Allow saving model on keyboard interrupt
+        except KeyboardInterrupt:
+            while True:
+                saving = input("Would you like to save the best (b), current (c), or no (n) model?")
+                saving = saving.strip().lower
+                if saving == "best" or saving == "b":
+                    self.load(save_path)
+                    self.save(save_path, spm_model_prefix, clean=True)
                     break
-            
-            scheduler.step()
-        
+                elif saving == "current" or saving == "c" or saving == "curr":
+                    self.save(save_path, spm_model_prefix, clean=True)
+                    break
+                elif saving == "n" or saving == "no" or saving == "":
+                    self.save(save_path, spm_model_prefix, clean=True)
+                    os.remove(save_path)
+                    break
+                else:
+                    print("invalid input, type 'b' 'c' or 'n'")
+                
         print(f"Best model saved to {save_path}")
     
     def encode(self, text: Union[str, List[str]], force_single_vector: bool = False) -> torch.Tensor:
@@ -802,10 +813,33 @@ class ThoughtVectors:
 
 
 if __name__ == "__main__":
+    def process_train(input_file):
+        data = []
+        with open(input_file, 'r', encoding='utf-8') as f_in:
+            reader = csv.reader(f_in)
+            for row in reader:
+                sentence1 = row[0].strip()
+                data.append([sentence1])
+
+        return data
+    
+    def process_test(input_file):
+        data = []
+        with open(input_file, 'r', encoding='utf-8') as f_in:
+            reader = csv.reader(f_in)
+            for row in reader:
+                sentence1 = row[0].strip()
+                data.append([sentence1])
+        return data
+        
+    data = process_train("SNLItrain.csv")
+    val = process_test("SNLIval.csv")
+
     tv = ThoughtVectors()
     tv.train(
-        group_data="train.csv",
-        test_data="val.csv",
+        group_data=data,
+        test_data=val,
+        model="thought_vectors_prototype-0.1.tar",
         num_epochs=1000,
         batch_size=128,
         accum_steps=1,
@@ -813,7 +847,7 @@ if __name__ == "__main__":
         weight_decay=2e-5,
         length_penalty=0.001,
         single_vector_prob=0.1,
-        save_path="thought_vectors_prototype.tar",
+        save_path="thought_vectors_prototype-0.2.tar",
         spm_model_prefix="spm",
         vocab_size=8192,
         d_model=512,
@@ -827,7 +861,7 @@ if __name__ == "__main__":
         termination_threshold=0.8,
         patience=5
     )
-    loaded_tv = ThoughtVectors.load("thought_vectors_prototype.tar")
+    loaded_tv = ThoughtVectors.load("thought_vectors_prototype-0.2.tar")
     thought_vectors = loaded_tv.encode("AI is smart")
     generated_text_greedy = loaded_tv.decode(thought_vectors, temperature=0.7, beam_width=0)
     generated_text_beam = loaded_tv.decode(thought_vectors, beam_width=5)