# -*- coding: utf-8 -*-
"""code_final2_rob+deb2.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1mVKtHWMKEFD32tvVXNtdQdmAO-EQDWG3
"""

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

!pip install sentencepiece

import json
import numpy as np
import torch
from transformers import (
    DebertaTokenizer, DebertaForSequenceClassification,
    RobertaTokenizer, RobertaForSequenceClassification,
    DebertaV2Tokenizer, DebertaV2ForSequenceClassification,
    BartTokenizer, BartForSequenceClassification,
    TrainingArguments, Trainer
)
from torch.utils.data import Dataset
from sklearn.metrics import classification_report

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]

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)

training_args = TrainingArguments(
    output_dir='./results',
    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=1.5e-5,
    save_strategy="no"
)

# 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, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
roberta_train_dataset = IntentDataset(roberta_train_encodings, train_labels)
roberta_trainer = Trainer(model=roberta_model, args=training_args, train_dataset=roberta_train_dataset)
roberta_trainer.train()

# Common TrainingArguments
training_args = TrainingArguments(
    output_dir='./results',
    num_train_epochs=10,
    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=1.5e-5,
    save_strategy="no"
)

# DeBERTa v2 XL
debv2_tokenizer = DebertaV2Tokenizer.from_pretrained("microsoft/deberta-v2-xlarge")
debv2_model = DebertaV2ForSequenceClassification.from_pretrained("microsoft/deberta-v2-xlarge", num_labels=len(label_map))
debv2_train_encodings = debv2_tokenizer(utt, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
debv2_train_dataset = IntentDataset(debv2_train_encodings, train_labels)
debv2_trainer = Trainer(model=debv2_model, args=training_args, train_dataset=debv2_train_dataset)
debv2_trainer.train()

# Test predictions
with open('massive_test.data', 'r') as test_file:
    test_data = [json.loads(line) for line in test_file]
test_utts = [item['utt'] for item in test_data]
test_ids = [item['indoml_id'] for item in test_data]

roberta_test_encodings = roberta_tokenizer(test_utts, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
debv2_test_encodings = debv2_tokenizer(test_utts, truncation=True, padding=True, max_length=max_length, return_tensors='pt')

roberta_preds = roberta_trainer.predict(IntentDataset(roberta_test_encodings, [0]*len(test_utts))).predictions
debv2_preds = debv2_trainer.predict(IntentDataset(debv2_test_encodings, [0]*len(test_utts))).predictions

ensemble_preds = (roberta_preds + debv2_preds) / 2
predicted_labels = np.argmax(ensemble_preds, axis=1)

predicted_intents = [list(label_map.keys())[list(label_map.values()).index(label)] for label in predicted_labels]

# Create submission file
submission_data = [{"indoml_id": int(test_id), "intent": intent} for test_id, intent in zip(test_ids, predicted_intents)]
submission_filename = 'massive_test_submission.predict'

with open(submission_filename, 'w') as submission_file:
    for item in submission_data:
        submission_file.write(json.dumps(item) + '\n')

# Zip the submission file
with zipfile.ZipFile("submission.zip", "w") as zipf:
    zipf.write(submission_filename)