Update app.py

app.py

@@ -108,7 +108,7 @@ def get_word_probabilities(text):
         word_colors.append(color)
     
     word_probs = [float(p) for p in word_probs]
-    return word_probs, 
+    return word_probs, word_colors
 
 # def get_word_classifications(text):
 #     text = " ".join(text.split(" ")[:2048])
@@ -170,7 +170,7 @@ def setup_hf_dataset():
 
 # Main inference + logging function
 def infer_and_log(text_input):
-    word_tags = get_word_probabilities(text_input)
+    word_probs, word_colors = get_word_probabilities(text_input)
     timestamp = datetime.datetime.now().isoformat()
     submission_id = str(uuid.uuid4())
 
@@ -199,8 +199,7 @@ def infer_and_log(text_input):
         except Exception as e:
             print(f"Error uploading log: {e}")
 
-    return json.dumps(word_probs, indent=2)
-
+    return word_colors
 
 def clear_fields():
     return "", ""