Sentiment / main.py
Msfast's picture
Upload 3 files
c4732f8 verified
raw
history blame contribute delete
No virus
1.01 kB
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.
@app.post("/sentiment")
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}"