import torch.nn.functional as F | |
from torch import Tensor | |
from transformers import AutoTokenizer, AutoModel, AutoProcessor | |
from torch import cuda | |
class EndpointHandler(): | |
def __init__(self, path=""): | |
#self.processor = AutoProcessor.from_pretrained(path) | |
#self.model = AutoModel.from_pretrained(path, trust_remote_code=True) | |
#self.device = "cuda" if cuda.is_available() else "cpu" | |
#self.model.to(self.device) | |
def __call__(self, data: Dict[str, Any]) -> List[List[int]]: | |
#image = data.pop("inputs",data) | |
#processed = self.processor(images=image, return_tensors="pt").to(self.device) | |
#prediction = self.model(processed["pixel_values"]) | |
return "OK"#prediction.item() |