Update app.py

app.py

@@ -25,19 +25,23 @@ if st.button("Predict"):
             outputs = model(**inputs)
         
         # Extract the predictions
-        predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
-        predictions = predictions[0].tolist()
+        predictions = outputs.logits.squeeze()
+
+        # Convert to numpy array if necessary
+        predicted_scores = predictions.numpy()
+        #predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        #predictions = predictions[0].tolist()
         
         # Convert predictions to a NumPy array for the calculations
-       # predictions_np = np.array(predictions)
+        #predictions_np = np.array(predictions)
         
         # Scale the predictions
-       # scaled_scores = 2.25 * predictions_np - 1.25
-       # rounded_scores = [round(score * 2) / 2 for score in scaled_scores]  # Round to nearest 0.5
+        scaled_scores = 2.25 * predictions_scores - 1.25
+        rounded_scores = [round(score * 2) / 2 for score in scaled_scores]  # Round to nearest 0.5
 
         # Display the predictions
         labels = ["cohesion", "syntax", "vocabulary", "phraseology", "grammar", "conventions"]
-        for label, score in zip(labels, predictions):
+        for label, score in zip(labels, rounded_scores):
             st.write(f"{label}: {score:.4f}")
     else:
         st.write("Please enter some text to get scores.")