File size: 821 Bytes
6766ca8
 
c0210b5
5ba6210
ec19c57
 
6766ca8
c0210b5
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 spacy
from environs import Env
from huggingface_hub import hf_hub_download
from joblib import load

SPACY_MODEL = spacy.load('en_core_web_trf', disable=['parser'])  # Largest, slowest, most accurate model



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)}