48 lines
1.1 KiB
Python
48 lines
1.1 KiB
Python
import torch
|
|
|
|
from app.ml.model_manager import model_manager
|
|
from app.config.logging import logger
|
|
from app.ml.prediction_result import PredictionResult
|
|
|
|
class TextClassifier:
|
|
def predict(self, text: str):
|
|
inputs = model_manager.text_tokenizer(
|
|
text,
|
|
return_tensors="pt",
|
|
truncation=True,
|
|
max_length=512,
|
|
padding=True
|
|
)
|
|
|
|
inputs = {
|
|
key: value.to(model_manager.device)
|
|
for key, value in inputs.items()
|
|
}
|
|
|
|
print(inputs)
|
|
|
|
with torch.no_grad():
|
|
outputs = model_manager.text_model(**inputs)
|
|
|
|
print(outputs)
|
|
|
|
probabilities = torch.softmax(outputs.logits, dim=1)
|
|
print("Labels:", model_manager.text_model.config.id2label)
|
|
print("Logits:", outputs.logits)
|
|
print("Probabilities:", probabilities)
|
|
|
|
print(probabilities)
|
|
raw_scores = probabilities[0].tolist()
|
|
score = raw_scores[1]
|
|
|
|
label = (
|
|
model_manager.text_model.config.id2label.get(1, "LABEL_1")
|
|
)
|
|
|
|
return PredictionResult(
|
|
label=label,
|
|
score=score,
|
|
approved=False,
|
|
raw_scores=raw_scores
|
|
)
|