hint_fine_tuning.py

import argparse
import os
import sys
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, random_split, TensorDataset
from src.dataset import TokenizerDataset
from src.bert import BERT
from src.pretrainer import BERTFineTuneTrainer1
from src.vocab import Vocab
import pandas as pd


# class CustomBERTModel(nn.Module):
#     def __init__(self, vocab_size, output_dim, pre_trained_model_path):
#         super(CustomBERTModel, self).__init__()
#         hidden_size = 768
#         self.bert = BERT(vocab_size=vocab_size, hidden=hidden_size, n_layers=12, attn_heads=12, dropout=0.1)
#         checkpoint = torch.load(pre_trained_model_path, map_location=torch.device('cpu'))
#         if isinstance(checkpoint, dict):
#             self.bert.load_state_dict(checkpoint)
#         elif isinstance(checkpoint, BERT):
#             self.bert = checkpoint
#         else:
#             raise TypeError(f"Expected state_dict or BERT instance, got {type(checkpoint)} instead.")
#         self.fc = nn.Linear(hidden_size, output_dim)

#     def forward(self, sequence, segment_info):
#         sequence = sequence.to(next(self.parameters()).device)
#         segment_info = segment_info.to(sequence.device)

#         if sequence.size(0) == 0 or sequence.size(1) == 0:
#             raise ValueError("Input sequence tensor has 0 elements. Check data preprocessing.")

#         x = self.bert(sequence, segment_info)
#         print(f"BERT output shape: {x.shape}")

#         if x.size(0) == 0 or x.size(1) == 0:
#             raise ValueError("BERT output tensor has 0 elements. Check input dimensions.")

#         cls_embeddings = x[:, 0]
#         logits = self.fc(cls_embeddings)
#         return logits

# class CustomBERTModel(nn.Module):
#     def __init__(self, vocab_size, output_dim, pre_trained_model_path):
#         super(CustomBERTModel, self).__init__()
#         hidden_size = 764  # Ensure this is 768
#         self.bert = BERT(vocab_size=vocab_size, hidden=hidden_size, n_layers=12, attn_heads=12, dropout=0.1)
        
#         # Load the pre-trained model's state_dict
#         checkpoint = torch.load(pre_trained_model_path, map_location=torch.device('cpu'))
#         if isinstance(checkpoint, dict):
#             self.bert.load_state_dict(checkpoint)
#         else:
#             raise TypeError(f"Expected state_dict, got {type(checkpoint)} instead.")
        
#         # Fully connected layer with input size 768
#         self.fc = nn.Linear(hidden_size, output_dim)

#     def forward(self, sequence, segment_info):
#         sequence = sequence.to(next(self.parameters()).device)
#         segment_info = segment_info.to(sequence.device)

#         x = self.bert(sequence, segment_info)
#         print(f"BERT output shape: {x.shape}")  # Should output (batch_size, seq_len, 768)

#         cls_embeddings = x[:, 0]  # Extract CLS token embeddings
#         print(f"CLS Embeddings shape: {cls_embeddings.shape}")  # Should output (batch_size, 768)

#         logits = self.fc(cls_embeddings)  # Should now pass a tensor of size (batch_size, 768) to `fc`
        
#         return logits


# for test
class CustomBERTModel(nn.Module):
    def __init__(self, vocab_size, output_dim, pre_trained_model_path):
        super(CustomBERTModel, self).__init__()
        self.hidden = 764  # Ensure this is defined correctly
        self.bert = BERT(vocab_size=vocab_size, hidden=self.hidden, n_layers=12, attn_heads=12, dropout=0.1)

        # Load the pre-trained model's state_dict
        checkpoint = torch.load(pre_trained_model_path, map_location=torch.device('cpu'))
        if isinstance(checkpoint, dict):
            self.bert.load_state_dict(checkpoint)
        else:
            raise TypeError(f"Expected state_dict, got {type(checkpoint)} instead.")

        self.fc = nn.Linear(self.hidden, output_dim)

    def forward(self, sequence, segment_info):
        x = self.bert(sequence, segment_info)
        cls_embeddings = x[:, 0]  # Extract CLS token embeddings
        logits = self.fc(cls_embeddings)  # Pass to fully connected layer
        return logits

def preprocess_labels(label_csv_path):
    try:
        labels_df = pd.read_csv(label_csv_path)
        labels = labels_df['last_hint_class'].values.astype(int)
        return torch.tensor(labels, dtype=torch.long)
    except Exception as e:
        print(f"Error reading dataset file: {e}")
        return None


def preprocess_data(data_path, vocab, max_length=128):
    try:
        with open(data_path, 'r') as f:
            sequences = f.readlines()
    except Exception as e:
        print(f"Error reading data file: {e}")
        return None, None

    if len(sequences) == 0:
        raise ValueError(f"No sequences found in data file {data_path}. Check the file content.")

    tokenized_sequences = []

    for sequence in sequences:
        sequence = sequence.strip()
        if sequence:
            encoded = vocab.to_seq(sequence, seq_len=max_length)
            encoded = encoded[:max_length] + [vocab.vocab.get('[PAD]', 0)] * (max_length - len(encoded))
            segment_label = [0] * max_length

            tokenized_sequences.append({
                'input_ids': torch.tensor(encoded),
                'segment_label': torch.tensor(segment_label)
            })

    if not tokenized_sequences:
        raise ValueError("Tokenization resulted in an empty list. Check the sequences and tokenization logic.")

    tokenized_sequences = [t for t in tokenized_sequences if len(t['input_ids']) == max_length]

    if not tokenized_sequences:
        raise ValueError("All tokenized sequences are of unexpected length. This suggests an issue with the tokenization logic.")

    input_ids = torch.cat([t['input_ids'].unsqueeze(0) for t in tokenized_sequences], dim=0)
    segment_labels = torch.cat([t['segment_label'].unsqueeze(0) for t in tokenized_sequences], dim=0)

    print(f"Input IDs shape: {input_ids.shape}")
    print(f"Segment labels shape: {segment_labels.shape}")

    return input_ids, segment_labels


def collate_fn(batch):
    inputs = []
    labels = []
    segment_labels = []

    for item in batch:
        if item is None:
            continue

        if isinstance(item, dict):
            inputs.append(item['input_ids'].unsqueeze(0))
            labels.append(item['label'].unsqueeze(0))
            segment_labels.append(item['segment_label'].unsqueeze(0))

    if len(inputs) == 0 or len(segment_labels) == 0:
        print("Empty batch encountered. Returning None to skip this batch.")
        return None

    try:
        inputs = torch.cat(inputs, dim=0)
        labels = torch.cat(labels, dim=0)
        segment_labels = torch.cat(segment_labels, dim=0)
    except Exception as e:
        print(f"Error concatenating tensors: {e}")
        return None

    return {
        'input': inputs,
        'label': labels,
        'segment_label': segment_labels
    }

def custom_collate_fn(batch):
    processed_batch = collate_fn(batch)
    
    if processed_batch is None or len(processed_batch['input']) == 0:
        # Return a valid batch with at least one element instead of an empty one
        return {
            'input': torch.zeros((1, 128), dtype=torch.long),
            'label': torch.zeros((1,), dtype=torch.long),
            'segment_label': torch.zeros((1, 128), dtype=torch.long)
        }
    
    return processed_batch


def train_without_progress_status(trainer, epoch, shuffle):
    for epoch_idx in range(epoch):
        print(f"EP_train:{epoch_idx}:")
        for batch in trainer.train_data:
            if batch is None:
                continue

            # Check if batch is a string (indicating an issue)
            if isinstance(batch, str):
                print(f"Error: Received a string instead of a dictionary in batch: {batch}")
                raise ValueError(f"Unexpected string in batch: {batch}")

            # Validate the batch structure before passing to iteration
            if isinstance(batch, dict):
                # Verify that all expected keys are present and that the values are tensors
                if all(key in batch for key in ['input_ids', 'segment_label', 'labels']):
                    if all(isinstance(batch[key], torch.Tensor) for key in batch):
                        try:
                            print(f"Batch Structure: {batch}")  # Debugging batch before iteration
                            trainer.iteration(epoch_idx, batch)
                        except Exception as e:
                            print(f"Error during batch processing: {e}")
                            sys.stdout.flush()
                            raise e  # Propagate the exception for better debugging
                    else:
                        print(f"Error: Expected all values in batch to be tensors, but got: {batch}")
                        raise ValueError("Batch contains non-tensor values.")
                else:
                    print(f"Error: Batch missing expected keys. Batch keys: {batch.keys()}")
                    raise ValueError("Batch does not contain expected keys.")
            else:
                print(f"Error: Expected batch to be a dictionary but got {type(batch)} instead.")
                raise ValueError(f"Invalid batch structure: {batch}")

# def main(opt):
#     # device = torch.device("cpu")
#     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

#     vocab = Vocab(opt.vocab_file)
#     vocab.load_vocab()

#     input_ids, segment_labels = preprocess_data(opt.data_path, vocab, max_length=128)
#     labels = preprocess_labels(opt.dataset)

#     if input_ids is None or segment_labels is None or labels is None:
#         print("Error in preprocessing data. Exiting.")
#         return

#     dataset = TensorDataset(input_ids, segment_labels, torch.tensor(labels, dtype=torch.long))
#     val_size = len(dataset) - int(0.8 * len(dataset))
#     val_dataset, train_dataset = random_split(dataset, [val_size, len(dataset) - val_size])

#     train_dataloader = DataLoader(
#         train_dataset,
#         batch_size=32,
#         shuffle=True,
#         collate_fn=custom_collate_fn
#     )
#     val_dataloader = DataLoader(
#         val_dataset,
#         batch_size=32,
#         shuffle=False,
#         collate_fn=custom_collate_fn
#     )

#     custom_model = CustomBERTModel(
#         vocab_size=len(vocab.vocab),
#         output_dim=2,
#         pre_trained_model_path=opt.pre_trained_model_path
#     ).to(device)

#     trainer = BERTFineTuneTrainer1(
#         bert=custom_model.bert,
#         vocab_size=len(vocab.vocab),
#         train_dataloader=train_dataloader,
#         test_dataloader=val_dataloader,
#         lr=5e-5,
#         num_labels=2,
#         with_cuda=torch.cuda.is_available(),
#         log_freq=10,
#         workspace_name=opt.output_dir,
#         log_folder_path=opt.log_folder_path 
#     )

#     trainer.train(epoch=20)

#     # os.makedirs(opt.output_dir, exist_ok=True)
#     # output_model_file = os.path.join(opt.output_dir, 'fine_tuned_model.pth')
#     # torch.save(custom_model.state_dict(), output_model_file)
#     # print(f'Model saved to {output_model_file}')
    
#     os.makedirs(opt.output_dir, exist_ok=True)
#     output_model_file = os.path.join(opt.output_dir, 'fine_tuned_model_2.pth')
#     torch.save(custom_model, output_model_file)
#     print(f'Model saved to {output_model_file}')


def main(opt):
    # Set device to GPU if available, otherwise use CPU
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    
    print(torch.cuda.is_available())  # Should return True if GPU is available
    print(torch.cuda.device_count())  

    # Load vocabulary
    vocab = Vocab(opt.vocab_file)
    vocab.load_vocab()

    # Preprocess data and labels
    input_ids, segment_labels = preprocess_data(opt.data_path, vocab, max_length=128)
    labels = preprocess_labels(opt.dataset)

    if input_ids is None or segment_labels is None or labels is None:
        print("Error in preprocessing data. Exiting.")
        return

    # Transfer tensors to the correct device (GPU/CPU)
    input_ids = input_ids.to(device)
    segment_labels = segment_labels.to(device)
    labels = torch.tensor(labels, dtype=torch.long).to(device)

    # Create TensorDataset and split into train and validation sets
    dataset = TensorDataset(input_ids, segment_labels, labels)
    val_size = len(dataset) - int(0.8 * len(dataset))
    val_dataset, train_dataset = random_split(dataset, [val_size, len(dataset) - val_size])

    # Create DataLoaders for training and validation
    train_dataloader = DataLoader(
        train_dataset,
        batch_size=32,
        shuffle=True,
        collate_fn=custom_collate_fn
    )
    val_dataloader = DataLoader(
        val_dataset,
        batch_size=32,
        shuffle=False,
        collate_fn=custom_collate_fn
    )

    # Initialize custom BERT model and move it to the device
    custom_model = CustomBERTModel(
        vocab_size=len(vocab.vocab),
        output_dim=2,
        pre_trained_model_path=opt.pre_trained_model_path
    ).to(device)

    # Initialize the fine-tuning trainer
    trainer = BERTFineTuneTrainer1(
        bert=custom_model.bert,
        vocab_size=len(vocab.vocab),
        train_dataloader=train_dataloader,
        test_dataloader=val_dataloader,
        lr=5e-5,
        num_labels=2,
        with_cuda=torch.cuda.is_available(),
        log_freq=10,
        workspace_name=opt.output_dir,
        log_folder_path=opt.log_folder_path 
    )

    # Train the model
    trainer.train(epoch=20)

    # Save the model to the specified output directory
    # os.makedirs(opt.output_dir, exist_ok=True)
    # output_model_file = os.path.join(opt.output_dir, 'fine_tuned_model_2.pth')
    # torch.save(custom_model.state_dict(), output_model_file)
    # print(f'Model saved to {output_model_file}')
    os.makedirs(opt.output_dir, exist_ok=True)
    output_model_file = os.path.join(opt.output_dir, 'fine_tuned_model_2.pth')
    torch.save(custom_model, output_model_file)
    print(f'Model saved to {output_model_file}')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Fine-tune BERT model.')
    parser.add_argument('--dataset', type=str, default='/home/jupyter/bert/dataset/hint_based/ratio_proportion_change_3/er/er_train.csv', help='Path to the dataset file.')
    parser.add_argument('--data_path', type=str, default='/home/jupyter/bert/ratio_proportion_change3_1920/_Aug23/gt/er.txt', help='Path to the input sequence file.')
    parser.add_argument('--output_dir', type=str, default='/home/jupyter/bert/ratio_proportion_change3_1920/_Aug23/output/hint_classification', help='Directory to save the fine-tuned model.')
    parser.add_argument('--pre_trained_model_path', type=str, default='/home/jupyter/bert/ratio_proportion_change3_1920/output/pretrain:1800ms:64hs:4l:8a:50s:64b:1000e:-5lr/bert_trained.seq_encoder.model.ep68', help='Path to the pre-trained BERT model.')
    parser.add_argument('--vocab_file', type=str, default='/home/jupyter/bert/ratio_proportion_change3_1920/_Aug23/pretraining/vocab.txt', help='Path to the vocabulary file.')
    parser.add_argument('--log_folder_path', type=str, default='/home/jupyter/bert/ratio_proportion_change3_1920/logs/oct_logs', help='Path to the folder for saving logs.')


    opt = parser.parse_args()
    main(opt)