|
import torch |
|
from transformers import AutoModelForSequenceClassification, AutoTokenizer |
|
|
|
class LanguageModel: |
|
def __init__(self): |
|
self.model = AutoModelForSequenceClassification.from_pretrained("bert-base-uncased") |
|
self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") |
|
|
|
def predict(self, text, image): |
|
|
|
inputs = self.tokenizer.encode_plus( |
|
text, |
|
add_special_tokens=True, |
|
max_length=512, |
|
return_attention_mask=True, |
|
return_tensors='pt' |
|
) |
|
outputs = self.model(inputs['input_ids'], attention_mask=inputs['attention_mask']) |
|
return torch.argmax(outputs.logits) |