FastAPI Server

fastapi.py

import torch
from fastapi import FastAPI
from transformers import AutoTokenizer, BatchEncoding, TensorType, AutoModelForSequenceClassification


application = FastAPI()
tokenizer = AutoTokenizer.from_pretrained("philschmid/MiniLM-L6-H384-uncased-sst2")
model = AutoModelForSequenceClassification.from_pretrained("philschmid/MiniLM-L6-H384-uncased-sst2").to('cuda:0')


@application.get("/predict")
def predict(query: str):
    inputs: BatchEncoding = tokenizer(
        text=query,
        max_length=128,
        truncation=True,
        return_tensors='pt',
    ).to('cuda:0')
    with torch.no_grad():
        logits = model(**inputs).logits
    return logits.to('cpu').tolist()
  return go(f, seed, [])
}