Views
No views yet
1from __future__ import annotations
2from transformers import RobertaConfig, RobertaModel, RobertaTokenizer, AutoModel, AutoTokenizer
3import torch
4
5# Add a custom regression head to RoBERTa
6class SITCC(torch.nn.Module):
7 def __init__(self, model, config):
8 super(SITCC, self).__init__()
9 self.roberta = model
10 self.regressor = torch.nn.Linear(config.hidden_size, 1) # Outputs a single value
11
12 def forward(self, input_ids, attention_mask):
13 outputs = self.roberta(input_ids=input_ids, attention_mask=attention_mask)
14 sequence_output = outputs[1] # The last hidden-state is the first element of the output tuple
15 logits = self.regressor(sequence_output)
16 return logits
17
18def init_model() -> SITCC:
19 # Load the model from huggingface
20 model_name = "KameronB/sitcc-roberta"
21 tokenizer = AutoTokenizer.from_pretrained(model_name, from_tf=False)
22 config = RobertaConfig.from_pretrained(model_name,)
23
24 # create the model based on the RoBERTa base model
25 model = SITCC(RobertaModel(config), config)
26
27 # fetch the statedict to apply the fine-tuned weights
28 state_dict = torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_name}/resolve/main/pytorch_model.bin")
29 # if running on cpu
30 # state_dict = torch.hub.load_state_dict_from_url(f"https://huggingface.co/{model_name}/resolve/main/pytorch_model.bin", map_location=torch.device('cpu'))
31
32 model.load_state_dict(state_dict)
33 return model, tokenizer
34
35model, tokenizer = init_model()
36
37def predict(sentences):
38 model.eval()
39 inputs = tokenizer(sentences, padding=True, truncation=True, max_length=512, return_tensors="pt")
40 input_ids = inputs['input_ids']
41 attention_mask = inputs['attention_mask']
42
43 with torch.no_grad():
44 outputs = model(input_ids, attention_mask)
45
46 return outputs