Views
No views yet
1import torch
2from transformers import AutoModelForSequenceClassification, AutoTokenizer
3
4MODEL_NAME = "swarogthehater/IMAGE_INTENT"
5model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, load_in_8bit=True)
6model.eval()
7tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
8
9text = "show me yourself"
10IMAGE_INTENT = "[IMG]"
11
12input_text = text+"[SEP]"+IMAGE_INTENT+"[SEP]"
13
14device = torch.device("cpu")
15batch = tokenizer.encode_plus(input_text, return_tensors="pt")
16input_ids = batch['input_ids'].to(device)
17attention_mask = batch['attention_mask'].to(device)
18token_type_ids = batch['token_type_ids'].to(device)
19outputs = model(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)
20
21#labels in common nli terms (0: entailment, 1: neutral, 2: contradiction)
22print(outputs.logits.argmax().item())
23#labels for img intent
24label = 1 if outputs.logits.argmax().item() == 0 else 0
25print(label)
26#scores
27print(outputs.logits.float().softmax(dim=-1).detach().numpy())