This repository has been archived by the owner on Sep 30, 2022. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 0
/
model.py
52 lines (42 loc) · 1.47 KB
/
model.py
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
"""Module to handle model loading and using it to predict urls"""
import joblib
from metrics_calculators import calculate_metrics, transform
from typing import Optional
class Model():
"""
Examples
--------
>>> model = Model()
>>> print(model.predict("https://a11egro.pl"))
'bad'
"""
def __init__(self, model_file: Optional[str] = "Dumps/model.sav",
scaler_file: Optional[str] = "Dumps/scaler.sav",
encoder_file: Optional[str] = "Dumps/encoder.sav"):
"""
Creates model from given files.
Parameters
----------
model_file: str
Name of file to witch model was dumped. Default: "Dumps/model.sav"
scaler_file : str
Name of file to witch StandardScaler was dumped. Default: "Dumps/scaler.sav"
encoder_file : str
Name of file to witch LabelEncoder was dumped. Default: "Dumps/encoder.sav"
"""
self.model = joblib.load(model_file)
self.scaler = joblib.load(scaler_file)
self.encoder = joblib.load(encoder_file)
def predict(self, url: str) -> str:
"""
Predicts whether url is bad or good.
Parameters
----------
url : str
Url to predict.
Returns
-------
str
Either 'good' or 'bad'.
"""
return self.encoder.inverse_transform(self.model.predict(self.scaler.transform(calculate_metrics(transform(url)))))[0]