from models import LMTraceSingle, LMTraceSingle
from transformers import ElectraTokenizer, AutoTokenizer, AutoModel, AutoConfig
import torch

model_name = 'google/electra-small-discriminator'
config = AutoConfig.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = LMTraceSingle.from_pretrained(model_name, config = config)
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
model.num_labels = 2
labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
outputs = model(**inputs, relation_label = labels)
model.save_pretrained("test_save")
print(outputs)
