php.py

import datetime
import pandas as pd
import numpy as np
import torch
from torch import nn
from transformers import DistilBertTokenizer, DistilBertModel, DistilBertConfig
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset, DataLoader
from tqdm import tqdm

# Initialization and Data Preparation
embedding_dim = 128
tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
# Load the configuration for DistilBERT
config = DistilBertConfig(dropout=0.2, attention_dropout=0.2)
dbert_model = DistilBertModel.from_pretrained('distilbert-base-uncased', config=config)

# Freeze all the parameters in the DistilBERT model.
for param in dbert_model.parameters():
 param.requires_grad = False

# Load your dataset. Assuming it's a CSV file with columns 'conversation', 'model_a', 'model_b', 'winner'
datapath = 'pair_wise.csv'
df = pd.read_csv(datapath)

unique_model_names = list(set(df['model_a'].unique()).union(set(df['model_b'].unique())))
model2id = {model: idx for idx, model in enumerate(unique_model_names)}

class PairWiseDataSet(Dataset):
 def __init__(self, dataframe, tokenizer, model2id, max_length=512):
 self.dataframe = dataframe
 self.tokenizer = tokenizer
 self.model2id = model2id
 self.max_length = max_length

 def __len__(self):
 return len(self.dataframe)

 def __getitem__(self, idx):
 row = self.dataframe.iloc[idx]
 text = row['conversation']
 model_a_id = self.model2id[row['model_a']]
 model_b_id = self.model2id[row['model_b']]
 winner_id = self.model2id[row['winner']]
 
 encoding = self.tokenizer(
 text,
 padding='max_length',
 truncation=True,
 max_length=self.max_length,
 return_tensors='pt'
 )
 input_ids = encoding['input_ids'].squeeze(0)
 attention_mask = encoding['attention_mask'].squeeze(0)
 
 return input_ids, attention_mask, torch.tensor([model_a_id, model_b_id]), torch.tensor(winner_id)

# Split the dataset
train_df, test_df = train_test_split(df, test_size=0.1, random_state=42)
train_dataset = PairWiseDataSet(train_df, tokenizer, model2id)
test_dataset = PairWiseDataSet(test_df, tokenizer, model2id)

train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=32, shuffle=False)

class PairwiseClassifier(nn.Module):
 def __init__(self, dbert_model, num_models, embedding_dim):
 super(PairwiseClassifier, self).__init__()
 self.dbert_model = dbert_model
 self.model_embeddings = nn.Embedding(num_models, embedding_dim)
 self.classifier = nn.Sequential(
 nn.Linear(dbert_model.config.hidden_size + 2 * embedding_dim, 768),
 nn.ReLU(),
 nn.Dropout(0.1),
 nn.Linear(768, num_models)
 )

 def forward(self, input_ids, attention_mask, model_pair_ids):
 outputs = self.dbert_model(input_ids=input_ids, attention_mask=attention_mask)
 pooled_output = outputs.last_hidden_state[:, 0]
 model_embeddings = self.model_embeddings(model_pair_ids)
 model_embeddings = model_embeddings.view(model_embeddings.size(0), -1)
 combined = torch.cat((pooled_output, model_embeddings), 1)
 logits = self.classifier(combined)
 return logits

# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
# Adjust for Metal GPU
device = "mps" if torch.backends.mps.is_available() else "cpu"
print(f"Using {device} device")
model = PairwiseClassifier(dbert_model, len(model2id), embedding_dim).to(device)

# Measure time for training
start_time = datetime.datetime.now()

# Training loop
epochs = 5 # Number of epochs
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001) # Define the optimizer
criterion = nn.CrossEntropyLoss() # Define the loss function
for epoch in range(epochs):
 model.train()
 train_loss = 0.0
 correct_predictions = 0
 total_predictions = 0

 train_loop = tqdm(train_loader, position=0, leave=True, desc=f'Epoch {epoch+1}/{epochs} [Training]')
 for inputs, attention_mask, model_pairs, labels in train_loop:
 inputs, attention_mask, model_pairs, labels = inputs.to(device), attention_mask.to(device), model_pairs.to(device), labels.to(device)
 
 optimizer.zero_grad() # Zero the gradients
 outputs = model(inputs, attention_mask, model_pairs) # Forward pass
 loss = criterion(outputs, labels) # Compute the loss
 loss.backward() # Backward pass
 optimizer.step() # Update weights
 
 train_loss += loss.item()
 _, preds = torch.max(outputs, dim=1)
 correct_predictions += torch.sum(preds == labels).item()
 total_predictions += labels.size(0)

 train_loop.set_postfix(loss=(train_loss / (train_loop.n + 1)))

 train_accuracy = correct_predictions / total_predictions
 
 # Validation loop
 model.eval() # Set the model to evaluation mode
 val_loss = 0.0
 correct_predictions = 0
 total_predictions = 0

 val_loop = tqdm(test_loader, position=0, leave=True, desc=f'Epoch {epoch+1}/{epochs} [Validation]')
 with torch.no_grad():
 for inputs, attention_mask, model_pairs, labels in val_loop:
 inputs, attention_mask, model_pairs, labels = inputs.to(device), attention_mask.to(device), model_pairs.to(device), labels.to(device)

 outputs = model(inputs, attention_mask, model_pairs)
 loss = criterion(outputs, labels)
 val_loss += loss.item()
 _, preds = torch.max(outputs, dim=1)
 correct_predictions += torch.sum(preds == labels).item()
 total_predictions += labels.size(0)

 val_loop.set_postfix(loss=(val_loss / (val_loop.n + 1)))

 val_accuracy = correct_predictions / total_predictions

 # Print formatted loss and accuracy
 print(f'\nEpoch {epoch+1} \t Training Loss: {train_loss / len(train_loader):.3f} \t Validation Loss: {val_loss / len(test_loader):.3f}')
 print(f'\t Training Accuracy: {train_accuracy:.3%} \t Validation Accuracy: {val_accuracy:.3%}')


# Save the model
torch.save(model.state_dict(), 'pairwise_classifier_model.pth')
print('Model saved to pairwise_classifier_model.pth')

# Measure training time
end_time = datetime.datetime.now()
print(f'Training completed in: {end_time - start_time}')