updated the model labels and api code

app.py

@@ -6,14 +6,13 @@ model = get_model()
 
 @app.post("/predict")
 async def predict(issue: str):
-    predictions = model.predict(issue)
+    labels = ["No", "Yes"]
+    predictions,label = model.predict(issue)
     print(f"Predictions: {predictions}")
     id = predictions
     print(f"Predictions: {predictions}")
     return {
         "priority 1": str(predictions[0]),
         "priority 2": str(predictions[1]),
-        "priority 3": str(predictions[2]),
-        "priority 4": str(predictions[3]),
-        "priority 5": str(predictions[4])
+        "Label":label
     }

classifier/Albert_latest.py

@@ -1,5 +1,6 @@
 from transformers import AlbertTokenizer, AlbertForSequenceClassification
 import torch
+import numpy as np
 
 class Model:
     def __init__(self, model_weights):
@@ -27,16 +28,15 @@ class Model:
             outputs = self.model(**inputs)
         
         logits = outputs.logits
-        print(f"logits: {logits}")
         predictions = torch.nn.functional.softmax(logits, dim=-1)
-        return predictions[0].tolist()
+        Labels = ["No", "Yes"]
+        return predictions[0].tolist()[:2],Labels[np.argmax(predictions)]
 
 model_instance = None
-model_weights = "assets/albert_sentiment_checkpoint_58.pt"
+model_weights = "../assets/albert_sentiment_checkpoint_58.pt"
+
 def get_model():
     global model_instance
     if model_instance is None:
         model_instance = Model(model_weights)
     return model_instance
-
-