from transformers import BertModel, BertTokenizer, Trainer, TrainingArguments
import torch
import torch.nn as nn
from datasets import load_dataset
from torch.utils.data import Dataset
from torch.nn.functional import cosine_similarity
import json
import numpy as np
import sys
sys.path.append("./")
from tqdm import tqdm
from embedder.train_classification_retriever import ClassificationNetwork

embedder_dir = 'embedder/classification_embedder/checkpoint-200'

num_labels = 9
model = ClassificationNetwork(num_labels=num_labels).to("cuda")
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# model = BertModel.from_pretrained('embedder/classification_embedder/checkpoint-100/')

# # Load the weights
model.load_state_dict(torch.load(embedder_dir + "/pytorch_model.bin"))
model.eval()  # for inference


dataset_path = "data/finetune/classification_preprocess_test.json"

ds = load_dataset("json", data_files=dataset_path)
dataset = ds['train']

success = 0
for idx, item in tqdm(enumerate(dataset), total=len(dataset)):
    # if idx == 1000:
    #     break
    # print("item", item)
    tokenized_input = tokenizer(item["query_prompt"], padding='max_length', truncation=True, max_length=512, return_tensors="pt")
    input_ids = tokenized_input["input_ids"].to("cuda")
    attention_mask = tokenized_input["attention_mask"].to("cuda")

    ground_truth = item["cluster_label"]
    with torch.no_grad():
        # print("input_ids", input_ids)
        # print("attention_mask", attention_mask)
        query_embedding = model(input_ids, attention_mask)#.pooler_output

        # print("query_embedding", query_embedding)
        prediction_label = torch.argmax(query_embedding)
        # print("prediction_label: ", prediction_label)
        # print("ground_truth: ", ground_truth)

        success = success + 1 if prediction_label == ground_truth else success

print("Success rate: ", success / len(dataset))
        # input()
