# %% [markdown]
# ## Introduction
# 
# Congratulations to everyone who persevered through this challenging competition!
# 
# I've been eagerly looking forward to reading the top solutions, and now that the competition has finally concluded, I am thrilled.
# 
# It has been a great learning journey for all of us. I am excited to share that I plan to publish a training notebook that demonstrates how to achieve a score of `Private Leaderboard 0.963` using a single model.
# 
# [Training Notebook PrivateLB 0.963](https://www.kaggle.com/code/haqishen/training-single-model-to-get-privatelb-0-963/)
# 
# And this notebook shows how to submit the model trained above.
# 
# 
# ## Acknowledge
# 
# I was provided Z8G5 Fury workstation with dual RTX6000 ada GPU.
# 
# For training a single model the above notebook shows, it takes less than 40min!

# %%
import re
import os
import json
import torch
import argparse
from pathlib import Path
from itertools import chain

import numpy as np
import pandas as pd
from transformers import AutoTokenizer, AutoModelForTokenClassification, AutoModel, AutoConfig, DataCollatorForTokenClassification
from datasets import Dataset
from spacy.lang.en import English

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# %%
!ls ../input/

# %%
backbone = 'microsoft/deberta-v3-large'
model_dir = '/kaggle/input/pii-dbtv3-large-9688'
th = 0.5

id2label = {0: 'B-EMAIL', 1: 'B-ID_NUM', 2: 'B-NAME_STUDENT', 3: 'B-PHONE_NUM', 4: 'B-STREET_ADDRESS', 5: 'B-URL_PERSONAL', 6: 'B-USERNAME', 7: 'I-ID_NUM', 8: 'I-NAME_STUDENT', 9: 'I-PHONE_NUM', 10: 'I-STREET_ADDRESS', 11: 'I-URL_PERSONAL', 12: 'O'}

# %%
def tokenize(example, tokenizer):
    text = []
    token_map = []
    
    idx = 0
    
    for t, ws in zip(example["tokens"], example["trailing_whitespace"]):
        
        text.append(t)
        token_map.extend([idx]*len(t))
        if ws:
            text.append(" ")
            token_map.append(-1)
            
        idx += 1
        
        
    tokenized = tokenizer("".join(text), return_offsets_mapping=True, truncation=False)
    
        
    return {
        **tokenized,
        "token_map": token_map,
    }

# %%
with open("/kaggle/input/pii-detection-removal-from-educational-data/test.json", "r") as f:
    data = json.load(f)

ds = Dataset.from_dict({
    "full_text": [x["full_text"] for x in data],
    "document": [x["document"] for x in data],
    "tokens": [x["tokens"] for x in data],
    "trailing_whitespace": [x["trailing_whitespace"] for x in data],
})

tokenizer = AutoTokenizer.from_pretrained(os.path.join(model_dir, 'fold0'))
collator = DataCollatorForTokenClassification(tokenizer)
ds = ds.map(tokenize, fn_kwargs={"tokenizer": tokenizer}, num_proc=2)

# %%
ds[0].keys()

# %%
def get_model(conf):
    return AutoModelForTokenClassification.from_config(
        conf
    )

fold = 0
conf = AutoConfig.from_pretrained(os.path.join(model_dir, f'fold{fold}'))
model = get_model(conf)
# model_file = os.path.join(model_dir, f'fold{fold}', f'genv1_noext1_ds3_ht1024_dbtv3_large_wu0.1_w10_ado4_wd0_bs2_lr1e5_5f_5epo_fold{fold}_best.pth')
model_file = "/kaggle/input/single_model_only_comp_data/pytorch/training_length_4096/1/deberta_v3_large_w10_ado4_wd0_bs2_lr1e5_5epo_fold0_best.pth"
sd = torch.load(model_file, map_location='cpu')
model.load_state_dict(sd, strict=True)
model.eval()
model.cuda()


# %%
!nvidia-smi

# %%
preds_final = []
with torch.no_grad():
    for sample in ds:
        input_ids = torch.tensor([sample['input_ids']]).to(device)
        attention_mask = torch.tensor([sample['attention_mask']]).to(device)

        preds = model(input_ids, attention_mask=attention_mask, labels=None)['logits'].softmax(2)
        preds[:, :, -1][preds[:, :, -1] < th] = 0
        preds = torch.argmax(preds, axis=2).cpu().numpy()
        preds_final.append(preds[0])


# %%
!nvidia-smi

# %%
pairs = []
processed = []
for p, token_map, offsets, tokens, doc in zip(preds_final, ds["token_map"],
                                              ds["offset_mapping"], ds["tokens"],
                                              ds["document"]):
    for token_pred, (start_idx, end_idx) in zip(p, offsets):
        label_pred = id2label[token_pred]

        if start_idx + end_idx == 0:
            continue

        if token_map[start_idx] == -1:
            start_idx += 1

        # ignore "\n\n"
        while start_idx < len(token_map) and tokens[token_map[start_idx]].isspace():
            start_idx += 1

        if start_idx >= len(token_map):
            break

        token_id = token_map[start_idx]

        # ignore "O" predictions and whitespace preds
        if label_pred not in ("O") and token_id != -1:
#         if label_pred not in ("O", "B-EMAIL", "B-PHONE_NUM", "I-PHONE_NUM") and token_id != -1:
            pair = (doc, token_id)

            if pair not in pairs:
                processed.append({"document": doc, "token": token_id, "label": label_pred, "token_str": tokens[token_id]})
                pairs.append(pair)

# %%
df = pd.DataFrame(processed)
df["row_id"] = list(range(len(df)))
display(df.head(100))

# %%
df[["row_id", "document", "token", "label"]].to_csv("submission.csv", index=False)

# %%