# -*- coding: utf-8 -*-
"""code_10-testing dev transfer model.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1YYajoZ_9v-0JQNl3QkTtdQlUDBVPxKDM
"""

!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)
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",
)

from transformers import BertConfig

# Assuming num_classes is defined earlier in the code. If not, define it:
num_classes = len(label_map)

# Transfer
model_path = '/content/drive/MyDrive/bert_model/model.pth'
tokenizer_path = '/content/drive/MyDrive/bert_tokenizer'
transfer_tokenizer = BertTokenizer(vocab_file='/content/drive/MyDrive/bert_tokenizer/vocab.txt')
config_original = BertConfig.from_pretrained('bert-base-multilingual-cased', num_labels=60)
transfer_model = BertForSequenceClassification(config=config_original)
model_state_dict = torch.load(model_path)
transfer_model.load_state_dict(model_state_dict)
transfer_model.num_labels = num_classes
transfer_model.classifier = torch.nn.Linear(transfer_model.config.hidden_size, num_classes)

# Create encodings for transfer learning
transfer_train_encodings = transfer_tokenizer(utt_train, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
transfer_test_encodings = transfer_tokenizer(utt_test, truncation=True, padding=True, max_length=max_length, return_tensors='pt')
transfer_train_dataset = IntentDataset(transfer_train_encodings, labels_train)
transfer_test_dataset = IntentDataset(transfer_test_encodings, labels_test)

# Optional: Freeze all layers except the last few
for param in transfer_model.parameters():
    param.requires_grad = False

# Unfreeze the classifier layers
for param in transfer_model.classifier.parameters():
    param.requires_grad = True

# Train the transfer model
transfer_trainer = Trainer(
    model=transfer_model,
    args=training_args,
    train_dataset=transfer_train_dataset,
    eval_dataset=transfer_test_dataset,
    compute_metrics=compute_metrics
)
transfer_trainer.train()
transfer_results = transfer_trainer.evaluate()

# Display results for the transfer model
display_results("Transfer BERT", transfer_results)