from transformers import BertTokenizer, BertModel
import torch
import numpy as np

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

class BERTClass(torch.nn.Module):
    def __init__(self):
        super(BERTClass, self).__init__()
        self.bert_model = BertModel.from_pretrained('bert-base-uncased', return_dict=True)
        self.dropout = torch.nn.Dropout(0.3)
        self.linear = torch.nn.Linear(768, 8)

    def forward(self, input_ids, attn_mask, token_type_ids):
        output = self.bert_model(
            input_ids,
            attention_mask=attn_mask,
            token_type_ids=token_type_ids
        )
        output_dropout = self.dropout(output.pooler_output)
        output = self.linear(output_dropout)
        return output

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

# Hyperparameters
MAX_LEN = 256
THRESHOLD = 0.5

# Loading pretrained model (best model)
model = BERTClass()
model.load_state_dict(torch.load(r"model\MLTC_model_state.bin", map_location=device))
model = model.to(device)

# raw text
raw_text = """
১০০% আসল প্রোডাক্ট। সিলেটের মধ্যে ৮ দিনের মধ্যে ডেলিভারি হয়েছে। বিক্রেতা খুবই সহানুভূতিশীল এবং ভালো ছিলেন। এই প্রোডাক্টটি এই বিক্রেতার কাছ থেকে কেনার জন্য অত্যন্ত সুপারিশ করছি।
"""

encoded_text = tokenizer.encode_plus(
    raw_text,
    max_length=MAX_LEN,
    add_special_tokens=True,
    return_token_type_ids=True,
    pad_to_max_length=True,
    return_attention_mask=True,
    return_tensors='pt',
)

input_ids = encoded_text['input_ids'].to(device)
attention_mask = encoded_text['attention_mask'].to(device)
token_type_ids = encoded_text['token_type_ids'].to(device)

output = model(input_ids, attention_mask, token_type_ids)
output = torch.sigmoid(output).detach().cpu()
output = output.flatten().round().numpy()

target_list = ['price', 'packaging', 'product', 'rider', 'delivery', 'shelf', 'service', 'seller']

print(f"Title: {raw_text}")
for idx, p in enumerate(output):
    if p == 1:
        print(f"Label: {target_list[idx]}")