Upload folder using huggingface_hub

Dockerfile

@@ -6,7 +6,8 @@ COPY . /app
 
 WORKDIR /app
 
-RUN pip install -r requirements.txt && \
+RUN pip install --upgrade pip && \
+    pip install -r requirements.txt && \
     pip install -r model/requirements.txt
 
 EXPOSE 7860

main/helper.py

@@ -2,6 +2,7 @@
 
 import yaml
 import joblib
+import pandas as pd
 from pathlib import Path
 
 # load yaml files to get model meta data.
@@ -32,4 +33,11 @@ def get_model_registry():
 def get_model_version():
     """ Fetches the model version and returns it. """
     model_version = model_metadata['model_version']
-    return model_version
+    return model_version
+
+
+def format_model_input(tweet: str) -> pd.DataFrame:
+    df = pd.DataFrame({
+        "comments" : tweet
+    })
+    return df

main/model_inference.py

@@ -1,6 +1,6 @@
 from fastapi import FastAPI
 from fastapi.responses import JSONResponse
-from main.validate_schema import UserInput, APIResponse
+from main.validate_schema import InputData, APIResponse
 from datetime import datetime
 from main.helper import *
 import uuid, time
@@ -10,20 +10,29 @@ model = load_model()
 # Initializing fastapi
 inference_api = FastAPI()
 
-@inference_api.post('/api', response_model=APIResponse)
-def api(payload: UserInput):
+@inference_api.get("/")
+def root():
+    return JSONResponse(content={
+        "status": 200,
+        "message": "Inference API is running."
+        })
+
+
+@inference_api.post('/get_prediction', response_model=APIResponse)
+def api(payload: InputData):
     timestamp = datetime.now().astimezone().isoformat()
     request_id = str(uuid.uuid4())
-    
     start_time = time.perf_counter()
+
     tweet = payload.comment
-    model_response = model.predict(tweet)
+    model_input = format_model_input(tweet)
+    model_response = model.predict(model_input)
 
     label = int(model_response["class_label"][0])
     probability_scores = model_response["class_probability_scores"]
     proba_class0 = float(probability_scores[0][0])
     proba_class1 = float(probability_scores[0][1])
-
+   
     end_time = time.perf_counter()
 
     if proba_class1 > 0.70:
@@ -35,9 +44,8 @@ def api(payload: UserInput):
     else:
         toxic_level = "none"
 
-
     response = {
-        "response": {
+        "prediction": {
             "class_label": label,
             "confidence": round(abs(proba_class0 - proba_class1), 4),
             "toxic_level": toxic_level,

main/validate_schema.py

@@ -2,11 +2,11 @@ from pydantic import BaseModel, Field
 from typing import Annotated, Dict
 
 
-class UserInput(BaseModel):
+class InputData(BaseModel):
     comment: Annotated[str, Field(..., description="User tweet or comment to be classified")]
 
 
-class ResponseData(BaseModel):
+class Prediction(BaseModel):
     class_label: int
     confidence: float
     toxic_level: str
@@ -28,5 +28,5 @@ class MetaData(BaseModel):
 
 
 class APIResponse(BaseModel):
-    response: ResponseData
+    response: Prediction
     metadata: MetaData