crarojasca/TaxonomyAugmentedCARDS

# Models
MAX_LEN = 256
BINARY_MODEL_DIR = "crarojasca/BinaryAugmentedCARDS"
TAXONOMY_MODEL_DIR = "crarojasca/TaxonomyAugmentedCARDS"

# Loading tokenizer
tokenizer = AutoTokenizer.from_pretrained(
    BINARY_MODEL_DIR,
    max_length = MAX_LEN, padding = "max_length", 
    return_token_type_ids = True
)

# Loading Models
## 1. Binary Model
print("Loading binary model: {}".format(BINARY_MODEL_DIR))
config = AutoConfig.from_pretrained(BINARY_MODEL_DIR)
binary_model = AutoModelForSequenceClassification.from_pretrained(BINARY_MODEL_DIR, config=config)
binary_model.to(device)

## 2. Taxonomy Model
print("Loading taxonomy model: {}".format(TAXONOMY_MODEL_DIR))
config = AutoConfig.from_pretrained(TAXONOMY_MODEL_DIR)
taxonomy_model = AutoModelForSequenceClassification.from_pretrained(TAXONOMY_MODEL_DIR, config=config)
taxonomy_model.to(device)

# Load Dataset
id2label = {
    0: '1_1', 1: '1_2', 2: '1_3', 3: '1_4', 4: '1_6', 5: '1_7', 6: '2_1', 
    7: '2_3', 8: '3_1', 9: '3_2', 10: '3_3', 11: '4_1', 12: '4_2', 13: '4_4', 
    14: '4_5', 15: '5_1', 16: '5_2', 17: '5_3'
}


text = "Climate change is just a natural phenomenon"

tokenized_text = tokenizer(text, return_tensors = "pt")


# Running Binary Model
outputs = binary_model(**tokenized_text)
binary_score = outputs.logits.softmax(dim = 1)
binary_prediction = torch.argmax(outputs.logits, axis=1)
binary_predictions = binary_prediction.to('cpu').item()

# Running Taxonomy Model
outputs = taxonomy_model(**tokenized_text)
taxonomy_score = outputs.logits.softmax(dim = 1)
taxonomy_prediction = torch.argmax(outputs.logits, axis=1)
taxonomy_prediction = taxonomy_prediction.to('cpu').item()


prediction = "0_0" if binary_prediction==0 else id2label[taxonomy_prediction]
prediction