From 8fbaf077f79018016e7c9d3cbdf0a523bbe714b2 Mon Sep 17 00:00:00 2001 From: Tillman Elser Date: Tue, 12 Sep 2023 22:00:24 -0700 Subject: [PATCH] update endpoint to use new model files --- src/seer/seer.py | 2 +- src/seer/severity/severity_inference.py | 13 +++++-------- 2 files changed, 6 insertions(+), 9 deletions(-) diff --git a/src/seer/seer.py b/src/seer/seer.py index 444f36d2e..6464ea267 100644 --- a/src/seer/seer.py +++ b/src/seer/seer.py @@ -68,7 +68,7 @@ def def_severity_endpoint(): try: data = request.get_json() severity = embeddings_model.severity_score(data.get("message", "")) - results = {"severity": str(severity[0][1])} + results = {"severity": str(severity)} return results except Exception as e: app.logger.exception("Error processing request") diff --git a/src/seer/severity/severity_inference.py b/src/seer/severity/severity_inference.py index 47d098747..2569877a3 100644 --- a/src/seer/severity/severity_inference.py +++ b/src/seer/severity/severity_inference.py @@ -1,22 +1,19 @@ import sentry_sdk import torch -from transformers import BertForSequenceClassification, BertTokenizerFast +from transformers import RobertaForSequenceClassification, RobertaTokenizerFast from joblib import load class SeverityInference: def __init__(self, embeddings_path, tokenizer_path, classifier_path): """Initialize the inference class with pre-trained models and tokenizer.""" - #TODO: needs to read from GCS - self.embeddings_model = BertForSequenceClassification.from_pretrained( + self.embeddings_model = RobertaForSequenceClassification.from_pretrained( embeddings_path ) - #TODO: needs to read from GCS - self.tokenizer = BertTokenizerFast.from_pretrained(tokenizer_path) + self.tokenizer = RobertaTokenizerFast.from_pretrained(tokenizer_path) self.device = ( torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") ) - #TODO: needs to read from GCS self.classifier = load(classifier_path) def get_embeddings(self, text, max_len=128): @@ -49,5 +46,5 @@ def severity_score(self, text): with sentry_sdk.start_span( op="model.severity", description="predict_proba" ): - pred = self.classifier.predict_proba(embeddings.reshape(1, -1)) - return pred + pred = self.classifier.predict(embeddings.reshape(1, -1))[0] + return min(1.0, max(0.0, pred))