indoml / code2_testing_libraries.py
Code Files
225a1cb about 3 years ago
26.4 kB
# -*- coding: utf-8 -*-
"""code2_testing-libraries.ipynb
Automatically generated by Colaboratory.
Original file is located at
    https://colab.research.google.com/drive/1_yrDJgQUAOpyGgXSf1Bp1L9erNs45IqV
"""

!pip install transformers torch
!pip install accelerate -U

from google.colab import files
import zipfile
import os

uploaded = files.upload()

for fn in uploaded.keys():
    print('User uploaded file "{name}" with length {length} bytes'.format(
        name=fn, length=len(uploaded[fn])))

    if fn.endswith('.zip'):
        with zipfile.ZipFile(fn, 'r') as zip_ref:
            zip_ref.extractall()

        print(f"Files from '{fn}' unzipped in the current directory!")
    else:
        print(f"'{fn}' is not a zip file.")

os.listdir()

import json
import numpy as np
import torch
from transformers import (BertTokenizer, BertForSequenceClassification,
                          RobertaTokenizer, RobertaForSequenceClassification,
                          XLNetTokenizer, XLNetForSequenceClassification,
                          DebertaTokenizer, DebertaForSequenceClassification,
                          BartTokenizer, BartForSequenceClassification,
                          TrainingArguments, Trainer,
                          XLMRobertaXLConfig, XLMRobertaXLModel, AutoTokenizer,
                          XLMRobertaTokenizer, XLMRobertaForSequenceClassification) # Added for XLM-RoBERTa
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset

from sklearn.metrics import classification_report

def compute_metrics(p):
    pred_labels = np.argmax(p.predictions, axis=1)
    true_labels = p.label_ids
    metrics = classification_report(true_labels, pred_labels, output_dict=True)
    accuracy = metrics['accuracy']
    weighted_avg = metrics['weighted avg']
    precision = weighted_avg['precision']
    f1 = weighted_avg['f1-score']
    return {
        'accuracy': accuracy,
        'f1': f1,
        'precision': precision
    }


# Load surprise data
with open('surprise.data', 'r') as data_file:
    data = [json.loads(line) for line in data_file]
with open('surprise.solution', 'r') as solution_file:
    solutions = [json.loads(line) for line in solution_file]

intent_map = {item['indoml_id']: item['intent'] for item in solutions}
indoml_ids = [item['indoml_id'] for item in data]
intents = [intent_map[indoml_id] for indoml_id in indoml_ids]
utt = [item['utt'] for item in data]

label_map = {intent: i for i, intent in enumerate(set(intents))}
num_labels = len(label_map)  # This is the added line
train_labels = [label_map[label] for label in intents]

utt_train, utt_test, labels_train, labels_test = train_test_split(utt, train_labels, test_size=0.3, random_state=42)

max_length = 64

class IntentDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

# Common training arguments
training_args = TrainingArguments(
    output_dir='./results',
    save_strategy="epoch",
    num_train_epochs=12,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
)


# XLM-RoBERTa
xlmroberta_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large")
xlmroberta_model = XLMRobertaForSequenceClassification.from_pretrained("xlm-roberta-large", num_labels=num_labels)
xlmroberta_train_encodings = xlmroberta_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
xlmroberta_test_encodings = xlmroberta_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
xlmroberta_train_dataset = IntentDataset(xlmroberta_train_encodings, labels_train)
xlmroberta_test_dataset = IntentDataset(xlmroberta_test_encodings, labels_test)

xlmroberta_trainer = Trainer(
    model=xlmroberta_model,
    args=training_args,
    train_dataset=xlmroberta_train_dataset,
    eval_dataset=xlmroberta_test_dataset,
    compute_metrics=compute_metrics
)
xlmroberta_trainer.train()
xlmroberta_results = xlmroberta_trainer.evaluate()

def display_results(model_name, results):
    print(f"{model_name} Results:")
    print(f"  Accuracy: {results['eval_accuracy']:.4f}")
    print(f"  F1-Score: {results['eval_f1']:.4f}")
    print(f"  Precision: {results['eval_precision']:.4f}")
    print(f"  Training Steps: {results['training_steps']}")

display_results("XLM-RoBERTa", xlmroberta_results)

import json
import numpy as np
import torch
from transformers import (BertTokenizer, BertForSequenceClassification,
                          RobertaTokenizer, RobertaForSequenceClassification,
                          XLNetTokenizer, XLNetForSequenceClassification,
                          DebertaTokenizer, DebertaForSequenceClassification,
                          BartTokenizer, BartForSequenceClassification,
                          TrainingArguments, Trainer)
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset

from sklearn.metrics import classification_report

def compute_metrics(p):
    pred_labels = np.argmax(p.predictions, axis=1)
    true_labels = p.label_ids
    metrics = classification_report(true_labels, pred_labels, output_dict=True)
    accuracy = metrics['accuracy']
    weighted_avg = metrics['weighted avg']
    precision = weighted_avg['precision']
    f1 = weighted_avg['f1-score']
    return {
        'accuracy': accuracy,
        'f1': f1,
        'precision': precision
    }


# Load surprise data
with open('surprise.data', 'r') as data_file:
    data = [json.loads(line) for line in data_file]
with open('surprise.solution', 'r') as solution_file:
    solutions = [json.loads(line) for line in solution_file]

intent_map = {item['indoml_id']: item['intent'] for item in solutions}
indoml_ids = [item['indoml_id'] for item in data]
intents = [intent_map[indoml_id] for indoml_id in indoml_ids]
utt = [item['utt'] for item in data]

label_map = {intent: i for i, intent in enumerate(set(intents))}
train_labels = [label_map[label] for label in intents]

utt_train, utt_test, labels_train, labels_test = train_test_split(utt, train_labels, test_size=0.3, random_state=42)

max_length = 64

class IntentDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

# Common training arguments
training_args = TrainingArguments(
    output_dir='./results',
    save_strategy="epoch",
    num_train_epochs=12,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
)

# BERT
bert_tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
bert_model = BertForSequenceClassification.from_pretrained("bert-large-uncased", num_labels=len(label_map))
bert_train_encodings = bert_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bert_test_encodings = bert_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bert_train_dataset = IntentDataset(bert_train_encodings, labels_train)
bert_test_dataset = IntentDataset(bert_test_encodings, labels_test)

bert_trainer = Trainer(
    model=bert_model,
    args=training_args,
    train_dataset=bert_train_dataset,
    eval_dataset=bert_test_dataset,
    compute_metrics=compute_metrics
)
bert_trainer.train()
bert_results = bert_trainer.evaluate()

# RoBERTa
roberta_tokenizer = RobertaTokenizer.from_pretrained("roberta-large")
roberta_model = RobertaForSequenceClassification.from_pretrained("roberta-large", num_labels=len(label_map))
roberta_train_encodings = roberta_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
roberta_test_encodings = roberta_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
roberta_train_dataset = IntentDataset(roberta_train_encodings, labels_train)
roberta_test_dataset = IntentDataset(roberta_test_encodings, labels_test)

roberta_trainer = Trainer(
    model=roberta_model,
    args=training_args,
    train_dataset=roberta_train_dataset,
    eval_dataset=roberta_test_dataset,
    compute_metrics=compute_metrics
)
roberta_trainer.train()
roberta_results = roberta_trainer.evaluate()

# DeBERTa
deb_tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-large")
deb_model = DebertaForSequenceClassification.from_pretrained("microsoft/deberta-large", num_labels=len(label_map))
deb_train_encodings = deb_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
deb_test_encodings = deb_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
deb_train_dataset = IntentDataset(deb_train_encodings, labels_train)
deb_test_dataset = IntentDataset(deb_test_encodings, labels_test)

deb_trainer = Trainer(
    model=deb_model,
    args=training_args,
    train_dataset=deb_train_dataset,
    eval_dataset=deb_test_dataset,
    compute_metrics=compute_metrics
)
deb_trainer.train()
deb_results = deb_trainer.evaluate()

# BART
bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
bart_model = BartForSequenceClassification.from_pretrained("facebook/bart-large", num_labels=len(label_map))
bart_train_encodings = bart_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bart_test_encodings = bart_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bart_train_dataset = IntentDataset(bart_train_encodings, labels_train)
bart_test_dataset = IntentDataset(bart_test_encodings, labels_test)

bart_trainer = Trainer(
    model=bart_model,
    args=training_args,
    train_dataset=bart_train_dataset,
    eval_dataset=bart_test_dataset,
    compute_metrics=compute_metrics
)
bart_trainer.train()
bart_results = bart_trainer.evaluate()

def display_results(model_name, results):
    print(f"{model_name} Results:")
    print(f"  Accuracy: {results['eval_accuracy']:.4f}")
    print(f"  F1-Score: {results['eval_f1']:.4f}")
    print(f"  Precision: {results['eval_precision']:.4f}")
    print(f"  Training Steps: {results['training_steps']}")

display_results("BERT", bert_results)
display_results("RoBERTa", roberta_results)
display_results("DeBERTa", deb_results)
display_results("BART", bart_results)

import json
import numpy as np
import torch
from transformers import (BertTokenizer, BertForSequenceClassification,
                          RobertaTokenizer, RobertaForSequenceClassification,
                          XLNetTokenizer, XLNetForSequenceClassification,
                          DebertaTokenizer, DebertaForSequenceClassification,
                          BartTokenizer, BartForSequenceClassification,
                          TrainingArguments, Trainer)
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset

from sklearn.metrics import classification_report

def compute_metrics(p):
    # Checking if predictions are stored in a tuple
    predictions = p.predictions[0] if isinstance(p.predictions, tuple) else p.predictions
    # Ensuring predictions are in numpy format
    if isinstance(predictions, torch.Tensor):
        predictions = predictions.detach().numpy()
    # Converting predictions to labels
    pred_labels = np.argmax(predictions, axis=1)

    # Getting labels and ensuring they're in numpy format
    true_labels = p.label_ids
    if isinstance(true_labels, torch.Tensor):
        true_labels = true_labels.detach().numpy()

    # Printing shapes for debugging
    print(predictions.shape)
    print(true_labels.shape)

    metrics = classification_report(true_labels, pred_labels, output_dict=True)
    accuracy = metrics['accuracy']
    weighted_avg = metrics['weighted avg']
    precision = weighted_avg['precision']
    f1 = weighted_avg['f1-score']

    return {
        'accuracy': accuracy,
        'f1': f1,
        'precision': precision
    }






# Load surprise data
with open('surprise.data', 'r') as data_file:
    data = [json.loads(line) for line in data_file]
with open('surprise.solution', 'r') as solution_file:
    solutions = [json.loads(line) for line in solution_file]

intent_map = {item['indoml_id']: item['intent'] for item in solutions}
indoml_ids = [item['indoml_id'] for item in data]
intents = [intent_map[indoml_id] for indoml_id in indoml_ids]
utt = [item['utt'] for item in data]

label_map = {intent: i for i, intent in enumerate(set(intents))}
train_labels = [label_map[label] for label in intents]

utt_train, utt_test, labels_train, labels_test = train_test_split(utt, train_labels, test_size=0.3, random_state=42)

max_length = 64

class IntentDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: val[idx].clone().detach() for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

# Common training arguments
training_args = TrainingArguments(
    output_dir='./results',
    save_strategy="epoch",
    num_train_epochs=12,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
)
bart_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
bart_model = BartForSequenceClassification.from_pretrained("facebook/bart-large", num_labels=len(label_map))
bart_train_encodings = bart_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bart_test_encodings = bart_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bart_train_dataset = IntentDataset(bart_train_encodings, labels_train)
bart_test_dataset = IntentDataset(bart_test_encodings, labels_test)

bart_trainer = Trainer(
    model=bart_model,
    args=training_args,
    train_dataset=bart_train_dataset,
    eval_dataset=bart_test_dataset,
    compute_metrics=compute_metrics
)

bart_trainer.train()
bart_results = bart_trainer.evaluate()

def display_results(model_name, results):
    print(f"{model_name} Results:")
    print(f"  Accuracy: {results['eval_accuracy']:.4f}")
    print(f"  F1-Score: {results['eval_f1']:.4f}")
    print(f"  Precision: {results['eval_precision']:.4f}")
    print(f"  Training Steps: {results['training_steps']}")

display_results("BERT", bert_results)
display_results("RoBERTa", roberta_results)
display_results("DeBERTa", deb_results)
display_results("BART", bart_results)

import json
import numpy as np
import torch
from transformers import (BertTokenizer, BertForSequenceClassification,
                          RobertaTokenizer, RobertaForSequenceClassification,
                          XLNetTokenizer, XLNetForSequenceClassification,
                          DebertaTokenizer, DebertaForSequenceClassification,
                          BartTokenizer, BartForSequenceClassification,
                          TrainingArguments, Trainer)
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset

from sklearn.metrics import classification_report

def compute_metrics(p):
    pred_labels = np.argmax(p.predictions, axis=1)
    true_labels = p.label_ids
    metrics = classification_report(true_labels, pred_labels, output_dict=True)
    accuracy = metrics['accuracy']
    weighted_avg = metrics['weighted avg']
    precision = weighted_avg['precision']
    f1 = weighted_avg['f1-score']
    return {
        'accuracy': accuracy,
        'f1': f1,
        'precision': precision
    }


# Load surprise data
with open('surprise.data', 'r') as data_file:
    data = [json.loads(line) for line in data_file]
with open('surprise.solution', 'r') as solution_file:
    solutions = [json.loads(line) for line in solution_file]

intent_map = {item['indoml_id']: item['intent'] for item in solutions}
indoml_ids = [item['indoml_id'] for item in data]
intents = [intent_map[indoml_id] for indoml_id in indoml_ids]
utt = [item['utt'] for item in data]

label_map = {intent: i for i, intent in enumerate(set(intents))}
train_labels = [label_map[label] for label in intents]

utt_train, utt_test, labels_train, labels_test = train_test_split(utt, train_labels, test_size=0.3, random_state=42)

max_length = 64

class IntentDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

# Common training arguments
training_args = TrainingArguments(
    output_dir='./results',
    save_strategy="epoch",
    num_train_epochs=20,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
)

# BERT
bert_tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
bert_model = BertForSequenceClassification.from_pretrained("bert-large-uncased", num_labels=len(label_map))
bert_train_encodings = bert_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bert_test_encodings = bert_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
bert_train_dataset = IntentDataset(bert_train_encodings, labels_train)
bert_test_dataset = IntentDataset(bert_test_encodings, labels_test)

bert_trainer = Trainer(
    model=bert_model,
    args=training_args,
    train_dataset=bert_train_dataset,
    eval_dataset=bert_test_dataset,
    compute_metrics=compute_metrics
)
bert_trainer.train()
bert_results = bert_trainer.evaluate()

def display_results(model_name, results):
    print(f"{model_name} Results:")
    print(f"  Accuracy: {results['eval_accuracy']:.4f}")
    print(f"  F1-Score: {results['eval_f1']:.4f}")
    print(f"  Precision: {results['eval_precision']:.4f}")

display_results("BERT", bert_results)

import json
import numpy as np
import torch
from transformers import (DebertaTokenizer, DebertaForSequenceClassification,
                          TrainingArguments, Trainer)
from sklearn.model_selection import train_test_split
from torch.utils.data import Dataset
from sklearn.metrics import classification_report

def compute_metrics(p):
    pred_labels = np.argmax(p.predictions, axis=1)
    true_labels = p.label_ids
    metrics = classification_report(true_labels, pred_labels, output_dict=True)
    accuracy = metrics['accuracy']
    weighted_avg = metrics['weighted avg']
    precision = weighted_avg['precision']
    f1 = weighted_avg['f1-score']
    return {
        'accuracy': accuracy,
        'f1': f1,
        'precision': precision
    }


# Load surprise data
with open('surprise.data', 'r') as data_file:
    data = [json.loads(line) for line in data_file]
with open('surprise.solution', 'r') as solution_file:
    solutions = [json.loads(line) for line in solution_file]

intent_map = {item['indoml_id']: item['intent'] for item in solutions}
indoml_ids = [item['indoml_id'] for item in data]
intents = [intent_map[indoml_id] for indoml_id in indoml_ids]
utt = [item['utt'] for item in data]

label_map = {intent: i for i, intent in enumerate(set(intents))}
train_labels = [label_map[label] for label in intents]

utt_train, utt_test, labels_train, labels_test = train_test_split(utt, train_labels, test_size=0.3, random_state=42)

max_length = 64

class IntentDataset(Dataset):
    def __init__(self, encodings, labels):
        self.encodings = encodings
        self.labels = labels

    def __getitem__(self, idx):
        item = {key: torch.tensor(val[idx]) for key, val in self.encodings.items()}
        item['labels'] = torch.tensor(self.labels[idx])
        return item

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

# Common training arguments
training_args = TrainingArguments(
    output_dir='./results',
    save_strategy="epoch",
    num_train_epochs=12,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=64,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
)

# DeBERTa
deb_tokenizer = DebertaTokenizer.from_pretrained("microsoft/deberta-xlarge")
deb_model = DebertaForSequenceClassification.from_pretrained("microsoft/deberta-xlarge", num_labels=len(label_map))
deb_train_encodings = deb_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
deb_test_encodings = deb_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
deb_train_dataset = IntentDataset(deb_train_encodings, labels_train)
deb_test_dataset = IntentDataset(deb_test_encodings, labels_test)

deb_trainer = Trainer(
    model=deb_model,
    args=training_args,
    train_dataset=deb_train_dataset,
    eval_dataset=deb_test_dataset,
    compute_metrics=compute_metrics
)
deb_trainer.train()
deb_results = deb_trainer.evaluate()

def display_results(model_name, results):
    print(f"{model_name} Results:")
    print(f"  Accuracy: {results['eval_accuracy']:.4f}")
    print(f"  F1-Score: {results['eval_f1']:.4f}")
    print(f"  Precision: {results['eval_precision']:.4f}")

display_results("DeBERTa", deb_results)

import itertools
import warnings
warnings.filterwarnings("ignore")

def ensemble_predictions(models, trainers, test_dataset):
    total_predictions = None
    for trainer in trainers:
        predictions, _, _ = trainer.predict(test_dataset)

        # If predictions are in tuple, extract the necessary part
        if isinstance(predictions, tuple):
            predictions = predictions[0]

        # Convert torch tensors to numpy arrays
        if isinstance(predictions, torch.Tensor):
            predictions = predictions.detach().numpy()

        # Sum the predictions
        if total_predictions is None:
            total_predictions = predictions
        else:
            total_predictions += predictions

    # Compute the average prediction by dividing by number of models
    ensemble_predictions = total_predictions / len(models)
    return ensemble_predictions

model_names = ["BERT", "RoBERTa", "DeBERTa"]

all_combinations = []
for r in [1, 2, 3]:
    for subset in itertools.combinations(model_names, r):
        all_combinations.append(subset)

for combo in all_combinations:
    trainers = []
    if "BERT" in combo:
        trainers.append(bert_trainer)
    if "RoBERTa" in combo:
        trainers.append(roberta_trainer)
    if "DeBERTa" in combo:
        trainers.append(deb_trainer)

    ensemble_pred = ensemble_predictions(models=combo, trainers=trainers, test_dataset=bert_test_dataset)  # Assuming all the datasets have same ordering

    # Convert averaged logits to labels
    pred_labels = np.argmax(ensemble_pred, axis=1)
    metrics = classification_report(labels_test, pred_labels, output_dict=True, zero_division=0)
    accuracy = metrics['accuracy']
    weighted_avg = metrics['weighted avg']
    precision = weighted_avg['precision']
    f1 = weighted_avg['f1-score']

    print(f"Results for {', '.join(combo)} ensemble:")
    print(f"  Accuracy: {accuracy:.4f}")
    print(f"  Precision: {precision:.4f}")
    print(f"  F1-Score: {f1:.4f}")
    print("="*40)

def get_predictions(trainer, dataset):
    outputs = trainer.predict(dataset)

    # Checking if predictions are stored in a tuple
    predictions = outputs.predictions[0] if isinstance(outputs.predictions, tuple) else outputs.predictions
    # Ensuring predictions are in numpy format
    if isinstance(predictions, torch.Tensor):
        predictions = predictions.detach().numpy()

    # Converting predictions to labels
    pred_labels = np.argmax(predictions, axis=1)

    # Printing shapes for debugging
    print(predictions.shape)
    print(pred_labels.shape)

    return pred_labels

bert_preds = get_predictions(bert_trainer, bert_test_dataset)
roberta_preds = get_predictions(roberta_trainer, roberta_test_dataset)
deb_preds = get_predictions(deb_trainer, deb_test_dataset)
bart_preds = get_predictions(bart_trainer, bart_test_dataset)


from itertools import combinations
from scipy.stats import mode
from sklearn.metrics import accuracy_score, f1_score, precision_score

# Combine model predictions and evaluate
models = ["BERT", "RoBERTa", "DeBERTa", "BART"]
preds_list = [bert_preds, roberta_preds, deb_preds, bart_preds]

for r in range(1, 5):  # For 1 to 4 models
    for subset in combinations(range(4), r):  # Get all combinations of models of size r
        subset_names = [models[i] for i in subset]
        subset_preds = [preds_list[i] for i in subset]

        # Majority voting
        ensemble_preds, _ = mode(subset_preds, axis=0)
        ensemble_preds = ensemble_preds.squeeze()

        # Calculate metrics
        accuracy = accuracy_score(labels_test, ensemble_preds)
        f1 = f1_score(labels_test, ensemble_preds, average='weighted')
        precision = precision_score(labels_test, ensemble_preds, average='weighted')

        # Print results
        print(f"Ensemble of {', '.join(subset_names)}:")
        print(f"  Accuracy: {accuracy:.4f}")
        print(f"  F1-Score: {f1:.4f}")
        print(f"  Precision: {precision:.4f}")
        print("-"*50)