|
|
import torch |
|
|
import torch.nn as nn |
|
|
import math |
|
|
from transformers import BertModel, BertTokenizer, BertPreTrainedModel |
|
|
|
|
|
class BertForRegression(BertPreTrainedModel): |
|
|
def __init__(self, config): |
|
|
super().__init__(config) |
|
|
self.bert = BertModel(config) |
|
|
self.regressor = nn.Linear(config.hidden_size, 1) |
|
|
|
|
|
self.init_weights() |
|
|
|
|
|
def forward(self, input_ids, attention_mask=None, token_type_ids=None, labels=None): |
|
|
outputs = self.bert(input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids) |
|
|
logits = self.regressor(outputs.pooler_output) |
|
|
|
|
|
loss = None |
|
|
if labels is not None: |
|
|
loss_fct = nn.MSELoss() |
|
|
loss = loss_fct(logits.squeeze(), labels.float()) |
|
|
|
|
|
return (loss, logits) if loss is not None else logits |
|
|
|
|
|
class EndpointHandler(): |
|
|
def __init__(self, path: str): |
|
|
self.model = BertForRegression.from_pretrained(path) |
|
|
self.tokenizer = BertTokenizer.from_pretrained(path) |
|
|
|
|
|
def __call__(self, data): |
|
|
self.model.eval() |
|
|
|
|
|
|
|
|
if isinstance(data, dict): |
|
|
text = data.get("inputs", "") |
|
|
else: |
|
|
text = data |
|
|
if not isinstance(text, str): |
|
|
raise ValueError("Input text must be a string under the 'inputs' key.") |
|
|
|
|
|
inputs = self.tokenizer(text, return_tensors="pt", padding=True, truncation=True, max_length=512) |
|
|
|
|
|
with torch.no_grad(): |
|
|
logits = self.model(**inputs) |
|
|
prediction = logits[0].item() |
|
|
prediction = math.trunc(prediction * 100) / 100 |
|
|
return str(prediction) |
