Don't think I need this

general_training_helper.py

@@ -1,172 +0,0 @@
-from torch.utils.data import DataLoader
-from level_dataset import LevelDataset
-import random
-from plotter import Plotter
-from datetime import datetime
-import os
-import threading
-import json
-import torch.nn.functional as F
-import torch
-
-
-
-
-def create_dataloaders(json_path, val_json, tokenizer, data_mode, augment, num_tiles, 
-                       negative_prompt_training, block_embeddings, batch_size):
-    # Initialize dataset
-    train_dataset = LevelDataset(
-        json_path=json_path,
-        tokenizer=tokenizer,
-        shuffle=True,
-        mode=data_mode,
-        augment=augment,
-        num_tiles=num_tiles,
-        negative_captions=negative_prompt_training,
-        block_embeddings=block_embeddings
-    )
-    val_dataset = None
-    if val_json is not None:
-        val_dataset = LevelDataset(
-            json_path=val_json,
-            tokenizer=tokenizer,
-            shuffle=False,
-            mode=data_mode,
-            augment=False,
-            num_tiles=num_tiles,
-            negative_captions=negative_prompt_training,
-            block_embeddings=block_embeddings
-        )
-
-    # Create dataloader
-    train_dataloader = DataLoader(
-        train_dataset,
-        batch_size=batch_size,
-        shuffle=True,
-        num_workers=4,
-        drop_last=True,
-        persistent_workers=True
-    )
-    
-    val_dataloader = None
-    if val_dataset is not None:
-        val_dataloader = DataLoader(
-            val_dataset,
-            batch_size=batch_size,
-            shuffle=False,
-            num_workers=4,
-            drop_last=False,
-            persistent_workers=True
-        )
-    
-    return train_dataloader, val_dataloader
-
-
-def get_random_training_samples(train_dataloader, negative_prompt_training, output_dir = None):
-    train_dataset = train_dataloader.dataset
-    # Sample four random captions from the dataset
-    sample_indices = [random.randint(0, len(train_dataset) - 1) for _ in range(4)]
-
-    sample_captions = [train_dataset[i][1] for i in sample_indices]
-    print("Sample captions:")
-    for caption in sample_captions:
-        print(caption)
-
-    sample_negative_captions = ""
-    if negative_prompt_training:
-        sample_negative_captions = [train_dataset[i][2] for i in sample_indices]
-        print("Sample negative captions:")
-        for caption in sample_negative_captions:
-            print(f"  NEG: {caption}")
-
-    #Write captions to a file
-    if output_dir is not None:
-        os.makedirs(output_dir, exist_ok=True)
-        out_path = os.path.join(output_dir, "sample_captions.txt")
-        with open(out_path, "w", encoding="utf-8") as f:
-            f.write("Sample captions:\n")
-            for caption in sample_captions:
-                f.write(str(caption) + "\n")
-            if negative_prompt_training:
-                f.write("\nSample negative captions:\n")
-                for caption in sample_negative_captions:
-                    f.write(str(caption) + "\n")
-        print(f"Sample captions written to {out_path}")
-
-
-    return sample_captions, sample_negative_captions
-
-
-def start_plotter(log_file, output_dir, left_key, right_key, left_label, right_label, png_name):
-    formatted_date = datetime.now().strftime(r'%Y%m%d-%H%M%S')
-
-    plotter = Plotter(log_file, update_interval=5.0, left_key=left_key, right_key=right_key,
-                            left_label=left_label, right_label=right_label, output_png=f'{png_name}_{formatted_date}.png')
-    plot_thread = threading.Thread(target=plotter.start_plotting)
-    plot_thread.daemon = True
-    plot_thread.start()
-    print(f"{png_name} plotting enabled. Progress will be saved to {os.path.join(output_dir, f'{png_name}_{formatted_date}.png')}")
-    return plotter, plot_thread
-
-
-def kill_plotter(plotter, plot_thread):
-    if plot_thread and plot_thread.is_alive():
-        plotter.stop_plotting()
-        plot_thread.join(timeout=5.0)
-        if plot_thread.is_alive():
-            print("Warning: Plot thread did not terminate properly")
-
-
-def load_config_from_json(config_path):
-    """Load hyperparameters from a JSON config file."""
-    try:
-        with open(config_path, 'r') as f:
-            config = json.load(f)
-            print(f"Configuration loaded from {config_path}")
-            
-            # Print the loaded config for verification
-            print("Loaded hyperparameters:")
-            for key, value in config.items():
-                print(f"  {key}: {value}")
-                
-            return config
-    except (json.JSONDecodeError, FileNotFoundError) as e:
-        print(f"Error loading config file: {e}")
-        raise e
-
-
-def update_args_from_config(args, config):
-    """Update argparse namespace with values from config."""
-    # Convert config dict to argparse namespace
-    for key, value in config.items():
-        if hasattr(args, key):
-            setattr(args, key, value)
-    return args
-
-
-def get_scene_from_embeddings(image, block_embeddings):
-    """Code copied over from level_dataset, should give limited support for block embeddings"""
-    # Reshape sample to [batch_size * height * width, embedding_dim]
-    batch_size, embedding_dim, height, width = image.shape
-    
-    flat_samples = image.permute(0, 2, 3, 1).reshape(-1, embedding_dim)
-    
-    # Normalize vectors for cosine similarity
-    flat_samples = F.normalize(flat_samples, p=2, dim=1).cpu()
-    block_embeddings = F.normalize(block_embeddings, p=2, dim=1)
-
-    # Calculate cosine similarity between each position and all tile embeddings
-    similarities = torch.matmul(flat_samples, block_embeddings.t())
-    
-    # Get indices of most similar tiles
-    indices = torch.softmax(similarities, dim=1)
-    
-    
-    # Reshape back to [batch_size, height, width]
-    indices = indices.reshape(batch_size, height, width, 13)
-    indices = indices.permute(0, 3, 1, 2)
-
-    image=indices.detach().cpu()
-    return image
-
-