File size: 763 Bytes
6766ca8
 
0cfb1e0
5ba6210
ec19c57
 
6766ca8
0cfb1e0
6766ca8
 
 
ec19c57
 
 
 
6766ca8
ec19c57
 
 
 
 
 
 
6766ca8
 
 
 
ec19c57
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
from typing import Dict, Any

import en_core_web_trf
from environs import Env
from huggingface_hub import hf_hub_download
from joblib import load

SPACY_MODEL = en_core_web_trf.load()



class EndpointHandler:
    def __init__(self, path: str):
        env = Env()
        env.read_env()

        model_path = env.str("MODEL_PATH")
        downloaded_model_path = hf_hub_download(
            repo_id="PDAP/url-relevance-models",
            subfolder=model_path,
            filename="model.joblib"
        )
        self.model = load(downloaded_model_path)

    def __call__(self, inputs: Dict[str, Any]) -> Dict[str, str]:
        # Expecting input like: {"inputs": "<html>...</html>"}
        html = inputs["inputs"]
        return {"label": str(self.model)}