from transformers import GPT2ForSequenceClassification
from classifier import GPT2Classifier
import torch
from torch import nn
import pytorch_lightning as pl
import torchmetrics
import numpy as np


class Embedder(pl.LightningModule):
    def __init__(
        self,
        checkpoint_path: str = "~/search-query-classification/pretrain_model/pretrain-checkpoints/model-epoch=01-val"
        "/loss=0.00.ckpt",
        lr: float = 1e-4,
        weights: list = None,
        num_labels=7,
        use_pretrained=True,
    ):
        super().__init__()
        transformer = GPT2Classifier.load_from_checkpoint(checkpoint_path)
        self.embedder = transformer.transformer.transformer

        if weights:
            weights = torch.tensor(weights)
            self.loss = nn.CrossEntropyLoss(weight=weights)
        else:
            self.loss = nn.CrossEntropyLoss()
        self.learning_rate = lr
        self.train_acc = torchmetrics.Accuracy()
        self.valid_acc = torchmetrics.Accuracy()

    def forward(self, input_ids, attention_mask):
        return self.embedder(input_ids=input_ids, attention_mask=attention_mask)[0][
            :, -1, :
        ]

    def predict_step(self, batch, batch_idx):
        input_ids, attention_mask = batch["input_ids"], batch["attention_mask"]
        label = batch["label"][0]
        preds = self(input_ids, attention_mask)
        with open(f"preds/{label}_{batch_idx}.npy", "wb") as f:
            np.save(f, preds.cpu())
        del preds
