feat: fastapi python to consume model

Author: AbraaoAlves

requirements.txt

@@ -0,0 +1,37 @@
+appdirs==1.4.4
+black==20.8b1
+certifi==2020.12.5
+chardet==4.0.0
+click==7.1.2
+fastapi==0.63.0
+filelock==3.0.12
+flake8==3.8.4
+gdown==3.12.2
+h11==0.12.0
+idna==2.10
+isort==5.7.0
+joblib==1.0.1
+mccabe==0.6.1
+mypy-extensions==0.4.3
+numpy==1.20.1
+packaging==20.9
+pathspec==0.8.1
+pycodestyle==2.6.0
+pydantic==1.8.1
+pyflakes==2.2.0
+pyparsing==2.4.7
+PySocks==1.7.1
+regex==2020.11.13
+requests==2.25.1
+sacremoses==0.0.43
+six==1.15.0
+starlette==0.13.6
+tokenizers==0.10.1
+toml==0.10.2
+torch==1.8.0
+tqdm==4.59.0
+transformers==4.3.3
+typed-ast==1.4.2
+typing-extensions==3.7.4.3
+urllib3==1.26.3
+uvicorn==0.13.4

src/api.py

@@ -0,0 +1,41 @@
+from typing import Dict
+
+from fastapi import Depends, FastAPI
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+
+from .classifier.model import Model, get_model
+
+app = FastAPI()
+
+#Allowing CORS
+origins = ["*"]
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+class SentimentRequest(BaseModel):
+    text:str
+
+class SentimentResponse(BaseModel):
+    probabilities: Dict[str, float]
+    sentiment: str
+    confidence: float
+
+
+@app.post("/predict", response_model = SentimentResponse)
+def predict(request: SentimentRequest, model: Model = Depends(get_model)):
+    '''
+    The injection of the parameter model is done by the function get_model
+    '''
+    sentiment, confidence, probabilities = model.predict(request.text)
+    return SentimentResponse(
+        sentiment = sentiment,
+        confidence = confidence,
+        probabilities = probabilities
+    )

src/classifier/model.py

@@ -0,0 +1,59 @@
+import json
+import torch
+import torch.nn.functional as F
+from transformers import BertTokenizer
+
+from .sentiment_classifier import SentimentClassifier
+
+with open("config.json") as json_file:
+    config = json.load(json_file)
+
+
+class Model:
+    def __init__(self):
+        self.device = torch.device("cuda:0" if torch.cuda.is_available()
+                                    else "cpu")
+        self.tokenizer = BertTokenizer.from_pretrained(config["BERT_MODEL"])
+
+        classifier = SentimentClassifier(len(config["CLASS_NAMES"]))
+        classifier.load_state_dict(
+            torch.load(config['PRE_TRAINED_MODEL'], map_location=self.device)
+        )
+        classifier = classifier.eval()
+        self.classifier = classifier.to(self.device)
+
+    def predict(self, text):
+        encoded_text = \
+            self.tokenizer.encode_plus( text,
+                                        max_lenght = config["MAX_SEQUENCE_LEN"],
+                                        add_special_tokens = True,
+                                        return_token_type_ids = False,
+                                        pad_to_max_length = True,
+                                        return_attetion_mask = True,
+                                        return_tensors = "pt")
+        input_ids = encoded_text["input_ids"].to(self.device)
+        attention_mask = encoded_text["attention_mask"].to(self.device)
+
+        with torch.no_grad():
+            #Getting the prediction
+            probabilities = F.softmax(self.classifier(input_ids, attention_mask),
+                            dim = 1)
+
+        #Taking the most confident result
+        confidence, predicted_class = torch.max(probabilities, dim = 1)
+        predicted_class = predicted_class.cpu().item()
+        probabilities = probabilities.flatten().cpu().numpy().tolist()
+
+        return(
+            config["CLASS_NAMES"][predicted_class],
+            confidence,
+            dict(zip(config["CLASS_NAMES"], probabilities))
+        )
+
+'''
+[init model is heavy] use the singleton pattern to avoid more one model initialization 
+'''
+
+model = Model()
+def get_model():
+    return model

src/classifier/sentiment_classifier.py

@@ -0,0 +1,22 @@
+import json
+
+from torch import nn
+from transformers import BertModel
+
+with open("config.json") as json_file:
+    config = json.load(json_file)
+
+
+class SentimentClassifier(nn.Module):
+    def __init__(self, n_classes):
+        #TODO: O que é esse super?
+        super(SentimentClassifier, self).__init__()
+        self.bert = BertModel.from_pretrained(config["BERT_MODEL"],
+                                              return_dict=False)
+        self.drop = nn.Dropout(p=0.3)
+        self.out = nn.Linear(self.bert.config.hidden_size, n_classes)
+
+    def forward(self, input_ids, attention_mask):
+        _, pooled_output = self.bert(input_ids=input_ids, attention_mask=attention_mask)
+        output = self.drop(pooled_output)
+        return self.out(output)