Upload 2 files

dtos.py

@@ -0,0 +1,10 @@
+from typing import List
+from pydantic import BaseModel
+
+
+class SentimentAnalysisRequestDto(BaseModel):
+    reviews: List[str]
+
+
+class SentimentAnalysisResponseDto(BaseModel):
+    scores: List[int]

handler.py

@@ -1,6 +1,7 @@
 from typing import  Dict, List, Any
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
+from dtos import SentimentAnalysisRequestDto, SentimentAnalysisResponseDto
 
 
 class EndpointHandler():
@@ -11,20 +12,15 @@ class EndpointHandler():
         self.sentiment_model = AutoModelForSequenceClassification.from_pretrained("Christian2903/amazon-review-sentiment-analysis").to(self.device)
 
 
-    def __call__(self, reviews: List[str]) -> List[int]:
-        """
-        Args:
-            data (:obj:):
-                includes the input data and the parameters for the inference.
-        Return:
-            A :obj:`list`:. The object returned should be a list of one list like [[{"label": 0.9939950108528137}]] containing :
-                - "label": A string representing what the label/class is. There can be multiple labels.
-                - "score": A score between 0 and 1 describing how confident the model is for this label/class.
-        """
+    def __call__(self, reviews: SentimentAnalysisRequestDto) -> SentimentAnalysisResponseDto:
+        reviews = [review for review in reviews.reviews]
 
         inputs = self.tokenizer(reviews, return_tensors="pt", truncation=True, padding="max_length", max_length=256).to(self.device)
         outputs = self.sentiment_model(**inputs)
         logits = outputs.logits.detach()
         predicted_scores = [max(min(int(score + 0.5),5),1) for score in logits]
-        # postprocess the prediction
-        return predicted_scores
+
+        response = SentimentAnalysisResponseDto(
+            scores=predicted_scores
+        )
+        return response