formating and mointoring change (#136)

This commit is contained in:
Co Tran 2024-10-07 15:21:05 -07:00 committed by GitHub
parent 976b2eaae0
commit 93abe553e3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 83 additions and 47 deletions

View file

@ -5,10 +5,11 @@ import torch
import pkg_resources
import yaml
def load_yaml_config(file_name):
# Load the YAML file from the package
yaml_path = pkg_resources.resource_filename('app', file_name)
with open(yaml_path, 'r') as yaml_file:
yaml_path = pkg_resources.resource_filename("app", file_name)
with open(yaml_path, "r") as yaml_file:
return yaml.safe_load(yaml_file)
@ -29,6 +30,7 @@ def split_text_into_chunks(text, max_words=300):
def softmax(x):
return np.exp(x) / np.exp(x).sum(axis=0)
class PredictionHandler:
def __init__(self, model, tokenizer, device, task="toxic", hardware_config="cpu"):
self.model = model