# -*- coding: utf-8 -*-
"""code7_gpt2.ipynb

Automatically generated by Colaboratory.

Original file is located at
    https://colab.research.google.com/drive/1LlzuHqU-HhK-u9FJmKMJKmwSUvksslQ8
"""

!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 transformers
import json
from transformers import GPT2Tokenizer, GPT2LMHeadModel

# 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}
intents = [intent_map[item['indoml_id']] for item in data]
utt = [item['utt'] for item in data]

label_map = {intent: i for i, intent in enumerate(set(intents))}

# Load GPT-2 model and tokenizer
tokenizer = GPT2Tokenizer.from_pretrained("gpt2-medium")
model = GPT2LMHeadModel.from_pretrained("gpt2-medium").eval()

def get_intent(utterance):
    input_text = f"Detect the intent of the following sentence:\n'{utterance}'\nThe possible intents are: {', '.join(label_map.keys())}."
    input_ids = tokenizer.encode(input_text, return_tensors='pt')
    output = model.generate(input_ids, max_length=200, num_return_sequences=1, temperature=0.8)
    output_text = tokenizer.decode(output[0], skip_special_tokens=True)
    # Extract the intent from the output text
    intent = output_text.split('\n')[-1].strip()
    return intent

# Get intents for the test data
with open('massive_test.data', 'r') as test_file:
    test_data = [json.loads(line) for line in test_file]

test_utterances = [item['utt'] for item in test_data]
predicted_intents = [get_intent(utterance) for utterance in test_utterances]

# Create submission file
submission_data = [{"indoml_id": int(item["indoml_id"]), "intent": intent} for item, intent in zip(test_data, 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')

# If you want to zip the file as in the sample
import zipfile

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