Update app.py

app.py

@@ -153,13 +153,13 @@ def evolve_emotions():
     emotions['ideal_state']['percentage'] = ideal_state
 
 def predict_emotion(context):
-    emotion_prediction_pipeline = pipeline('text-classification', model=emotion_prediction_model, tokenizer=emotion_prediction_tokenizer, top_k=None)
+    emotion_prediction_pipeline = pipeline('text-classification', model=emotion_prediction_model, tokenizer=emotion_prediction_tokenizer, return_all_scores=True)
     predictions = emotion_prediction_pipeline(context)
-    emotion_scores = {prediction['label']: prediction['score'] for prediction in predictions}
+    emotion_scores = {prediction['label']: prediction['score'] for prediction in predictions[0]}
     emotion_pred = max(emotion_scores, key=emotion_scores.get)
     return emotion_pred
 
-def generate_text(prompt, max_length=100, emotion=None):
+def generate_text(prompt, emotion=None, max_length=100):
     finetuned_lm_tokenizer, finetuned_lm_model = get_finetuned_lm_model()
     input_ids = finetuned_lm_tokenizer.encode(prompt, return_tensors='pt')
     attention_mask = torch.ones(input_ids.shape, dtype=torch.long)