Update README.md

README.md

@@ -92,7 +92,8 @@ def predict_personality(model, encoded_input):
     model.eval()  # Set the model to evaluation mode
     with torch.no_grad():
         outputs = model(**encoded_input)
-    return outputs.logits.squeeze()
+    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
+    return predictions[0].tolist()
 
 def print_predictions(predictions, trait_names):
     for trait, score in zip(trait_names, predictions):
@@ -117,10 +118,11 @@ def main():
     predictions = predict_personality(model, encoded_input)
 
     trait_names = ["Agreeableness", "Openness", "Conscientiousness", "Extraversion", "Neuroticism"]
-    print_predictions(predictions.numpy(), trait_names)
+    print_predictions(predictions, trait_names)
 
 if __name__ == "__main__":
     main()
+
 ```
 ```bash
 python script_name.py --input "Your text here"