import torch
from ABSA.model import BERT_BiLSTM_CRF  # Same model class you defined
from transformers import AutoTokenizer, AutoModel
import json

# Load tokenizer and base model
model_path = "saved_model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
base_model = AutoModel.from_pretrained(model_path)

# Load label mappings
with open(f"{model_path}/label2id.json") as f:
    label2id = json.load(f)
with open(f"{model_path}/id2label.json") as f:
    id2label = {int(k): v for k, v in json.load(f).items()}

# Init and load model
num_labels = len(label2id)
model = BERT_BiLSTM_CRF(base_model, num_labels)
model.load_state_dict(torch.load(f"{model_path}/full_model.pth", map_location="cpu"))
model.eval()