INNER CODE UNIT · Python

TextClassificationAPI

CerebriumAI/examples · 10-batching/2-litserve-batching-gpu/main.py:5

class TextClassificationAPI(ls.LitAPI):
    def setup(self, device):
        print("setup")
        self.model = pipeline(
            "sentiment-analysis", model="stevhliu/my_awesome_model", device=device
        )

    def decode_request(self, request):
        print("decode_request")
        return request["text"]

    def predict(self, x):
        print("predict")
        result = self.model(x)
        return result

    def encode_response(self, output):
        print("encode_response")

View source record →