|
from typing import Dict, List, Any |
|
from transformers import ClapModel, ClapProcessor |
|
import gc |
|
|
|
class EndpointHandler: |
|
def __init__(self, path=""): |
|
model_name = "laion/larger_clap_general" |
|
self.model = ClapModel.from_pretrained(model_name) |
|
self.processor = ClapProcessor.from_pretrained(model_name) |
|
|
|
def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]: |
|
""" |
|
data args: |
|
inputs (:obj: `str`) |
|
date (:obj: `str`) |
|
Return: |
|
A :obj:`list` | `dict`: will be serialized and returned |
|
""" |
|
print(type(data)) |
|
query = data['inputs'] |
|
text_inputs = self.processor(text=query, return_tensors="pt") |
|
text_embed = self.model.get_text_features(**text_inputs)[0] |
|
gc.collect() |
|
return text_embed |