File size: 1,042 Bytes
1320c7e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Hugging Face Inference Endpoints entry point: deploy this repo as a CPU/GPU chat API.

Request:  {"inputs": "How fares the king?"}
          {"inputs": {"message": "...", "history": [["user", "bot"], ...]},
           "parameters": {"max_new_tokens": 200, "temperature": 0.8, "top_k": 40, "seed": 1}}
Response: [{"generated_text": "<reply>"}]
"""
import sys
from pathlib import Path

HERE = Path(__file__).resolve().parent
sys.path.insert(0, str(HERE))
import model as M  # noqa: E402


class EndpointHandler:
    def __init__(self, path: str = ""):
        self.predictor = M.load(path or HERE, "cuda" if M.cuda_available() else "cpu")

    def __call__(self, data: dict):
        inputs = data.pop("inputs", data)
        parameters = data.pop("parameters", None) or {}
        if isinstance(inputs, dict):
            message, history = inputs.get("message", ""), inputs.get("history")
        else:
            message, history = inputs, None
        return [{"generated_text": self.predictor.predict(message, history, **parameters)}]