Spaces:
Sleeping
Sleeping
File size: 2,456 Bytes
0782370 20df6e4 0782370 20df6e4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 |
import numpy as np
import torch
from torch.utils.data import DataLoader, SequentialSampler, TensorDataset
from transformers import BertTokenizer
class UpvotePredictor:
def __init__(self, model_path: str):
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.upvote_ml_model = torch.load(
model_path, map_location=torch.device("cpu"), weights_only=False
)
self.tokenizer = BertTokenizer.from_pretrained(
"bert-base-uncased", do_lower_case=True
)
self.upvote_ml_model.to(self.device)
self.upvote_ml_model.eval()
def get_upvote_prediction(
self, question: str, answer: str, question_context: str = None
) -> int:
llm_response_input_ids = []
llm_response_attention_masks = []
encoded_dict = self.tokenizer.encode_plus(
answer,
add_special_tokens=True,
max_length=256,
padding="max_length",
truncation=True,
return_attention_mask=True,
return_tensors="pt",
)
llm_response_input_ids.append(encoded_dict["input_ids"])
llm_response_attention_masks.append(encoded_dict["attention_mask"])
llm_response_input_ids = torch.cat(llm_response_input_ids, dim=0)
llm_response_attention_masks = torch.cat(llm_response_attention_masks, dim=0)
test_dataset = TensorDataset(
llm_response_input_ids, llm_response_attention_masks
)
test_dataloader = DataLoader(
test_dataset, # The validation samples.
sampler=SequentialSampler(test_dataset), # Pull out batches sequentially.
batch_size=1, # Evaluate with this batch size.
)
predictions = []
for batch in test_dataloader:
b_input_ids = batch[0].to(self.device)
b_input_mask = batch[1].to(self.device)
with torch.no_grad():
output = self.upvote_ml_model(
b_input_ids, token_type_ids=None, attention_mask=b_input_mask
)
logits = output.logits
logits = logits.detach().cpu().numpy()
pred_flat = np.argmax(logits, axis=1).flatten()
predictions.extend(list(pred_flat))
if predictions[0] == 0:
return "Not credible suggestion"
else:
return "Credible suggestion"
|