indoml / code3_testing_massive_optimised_libraries.py
Code Files
225a1cb about 3 years ago
27.5 kB
# -*- coding: utf-8 -*-
"""code3_testing-massive-optimised-libraries.ipynb
Automatically generated by Colaboratory.
Original file is located at
    https://colab.research.google.com/drive/1MqyIO7-h90kGPj47FxEoX3BRNT42elXZ
"""

!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()

# Install required libraries
!pip install sentencepiece

import json
import numpy as np
from transformers import (TrainingArguments, Trainer, AutoTokenizer, AutoModelForSequenceClassification, XLMRobertaConfig)
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 dataset
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)
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

# Custom dataset class
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)

# 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"
)

# Tokenization, Model Configuration, and Dataset Preparation
config = XLMRobertaConfig.from_pretrained("philschmid/habana-xlm-r-large-amazon-massive")
config.num_labels = num_labels
xlmroberta_tokenizer = AutoTokenizer.from_pretrained("philschmid/habana-xlm-r-large-amazon-massive")
xlmroberta_model = AutoModelForSequenceClassification.from_config(config)

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)

# Trainer Setup and Model Training
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()

# Displaying results
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("XLM-Roberta", xlmroberta_results)

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

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    metrics = classification_report(labels, predictions, 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
    }

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

tokenizer = AutoTokenizer.from_pretrained("projecte-aina/roberta-large-ca-v2-massive")

def tokenize_function(utt, labels):
    encodings = tokenizer(utt, truncation=True, padding='max_length', max_length=max_length, return_tensors="pt")
    dataset = [{"input_ids": val[0], "attention_mask": val[1], "labels": torch.tensor(label)} for val, label in zip(zip(encodings["input_ids"], encodings["attention_mask"]), labels)]
    return dataset

train_data = tokenize_function(utt_train, labels_train)
test_data = tokenize_function(utt_test, labels_test)

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

    def __getitem__(self, idx):
        return self.data[idx]

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

train_dataset = IntentDataset(train_data)
test_dataset = IntentDataset(test_data)

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",
    report_to="none"
)

tokenizer = AutoTokenizer.from_pretrained("projecte-aina/roberta-large-ca-v2-massive")
model = AutoModelForSequenceClassification.from_pretrained("projecte-aina/roberta-large-ca-v2-massive", num_labels=60)
import torch.nn as nn
num_unique_labels = len(label_map)
model.config.num_labels = num_unique_labels
model.classifier.out_proj = nn.Linear(model.config.hidden_size, num_unique_labels)


trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=tokenizer,
    compute_metrics=compute_metrics
)

trainer.train()

!pip install sentencepiece

import json
import numpy as np
import torch
from transformers import (TrainingArguments, Trainer, AutoTokenizer,
                          AutoModelForSequenceClassification, XLMRobertaConfig) # You don't need all the other imports for this example
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",
)



# Adjust the number of labels in the classifier layer
config = XLMRobertaConfig.from_pretrained("philschmid/habana-xlm-r-large-amazon-massive")
config.num_labels = num_labels

# XLM-RoBERTa
xlmroberta_tokenizer = AutoTokenizer.from_pretrained("philschmid/habana-xlm-r-large-amazon-massive")
xlmroberta_model = AutoModelForSequenceClassification.from_config(config)


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=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("cartesinus/bert-base-uncased-amazon-massive-intent")
bert_model = BertForSequenceClassification.from_pretrained(
    "cartesinus/bert-base-uncased-amazon-massive-intent",
    num_labels=len(label_map),
    ignore_mismatched_sizes=True
)
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)

!pip install sentencepiece

import json
import numpy as np
import torch
from transformers import (BertTokenizer, BertForSequenceClassification,
                          RobertaTokenizer, RobertaForSequenceClassification,
                          XLNetTokenizer, XLNetForSequenceClassification,
                          DebertaTokenizer, DebertaForSequenceClassification,
                          BartTokenizer, BartForSequenceClassification,
                          TrainingArguments, Trainer, pipeline, AutoModelForTokenClassification,
                          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

torch.cuda.empty_cache()

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=8,
    per_device_eval_batch_size=32,
    warmup_steps=500,
    weight_decay=0.01,
    logging_dir='./logs',
    logging_steps=50,
    learning_rate=2e-5,
    evaluation_strategy="epoch",
    gradient_accumulation_steps=2,
    fp16=True
)

import torch
from transformers import AutoTokenizer, XLMRobertaXLForSequenceClassification

# XLM-RoBERTa
tokenizer = AutoTokenizer.from_pretrained("Soonhwan-Kwon/xlm-roberta-xlarge")
xlmroberta_model = XLMRobertaXLForSequenceClassification.from_pretrained("Soonhwan-Kwon/xlm-roberta-xlarge", num_labels=num_labels)
xlmroberta_train_encodings = tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
xlmroberta_test_encodings = 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 (TrainingArguments, Trainer, AutoTokenizer,
                          AutoModelForSequenceClassification, XLMRobertaConfig) # You don't need all the other imports for this example
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",
)



# Adjust the number of labels in the classifier layer
config = XLMRobertaConfig.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
config.num_labels = num_labels

# XLM-RoBERTa
xlmroberta_tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
xlmroberta_model = AutoModelForSequenceClassification.from_config(config)


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, pipeline, AutoModelForTokenClassification,
                          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-finetuned-conll03-english")
xlmroberta_model = AutoModelForTokenClassification.from_pretrained("xlm-roberta-large-finetuned-conll03-english", 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)