from fastapi import FastAPI | |
from pydantic import BaseModel | |
from transformers import pipeline | |
# You can check any other model in the Hugging Face Hub. In my case I chose this one to classify text by positive and negative sentiment. | |
pipe = pipeline(model="distilbert/distilbert-base-uncased-finetuned-sst-2-english") | |
# We define the app | |
app = FastAPI() | |
# We define that we expect our input to be a string | |
class RequestModel(BaseModel): | |
input: str | |
# Now we define that we accept post requests | |
# -> In APIs, requests are made to ask the API to perform a certain task — in this case to analyze a piece of text. | |
def get_response(request: RequestModel): | |
# We get the input prompt | |
prompt = request.input | |
# We use the hf model to classify the prompt | |
response = pipe(prompt) | |
# We get both the label and the score from the input | |
label = response[0]["label"] | |
score = response[0]["score"] | |
return f"The '{prompt}' input is {label} with a score of {score}" | |