Update app.py

app.py

@@ -110,53 +110,6 @@ def get_word_probabilities(text):
     word_probs = [float(p) for p in word_probs]
     return word_probs, word_colors
 
-# def get_word_classifications(text):
-#     text = " ".join(text.split(" ")[:2048])
-#     inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
-#     inputs = {k: v.to(device) for k, v in inputs.items()}
-#     tokens = tokenizer.convert_ids_to_tokens(inputs["input_ids"][0])
-    
-#     with torch.no_grad():
-#         tags, emissions = model(input_ids=inputs["input_ids"], attention_mask=inputs["attention_mask"])
-
-#     word_tags = []
-#     color_output = []
-#     current_word = ""
-#     current_prob = 0.0
-
-#     for token, prob in zip(tokens, tags[0]):
-#         if token in ["<s>", "</s>"]:
-#             continue
-#         if token.startswith("▁"):
-#             if current_word:
-#                 word_tags.append(round(current_prob, 3))
-#                 color = (
-#                     "green" if current_prob < 0.25 else
-#                     "yellow" if current_prob < 0.5 else
-#                     "orange" if current_prob < 0.75 else
-#                     "red"
-#                 )
-#                 color_output.append(f'<span style="color:{color}">{current_word}</span>')
-#             current_word = token[1:] if token != "▁" else ""
-#             current_prob = prob
-#         else:
-#             current_word += token
-#             current_prob = max(current_prob, prob)
-
-#     if current_word:
-#         word_tags.append(round(current_prob, 3))
-#         color = (
-#             "green" if current_prob < 0.25 else
-#             "yellow" if current_prob < 0.5 else
-#             "orange" if current_prob < 0.75 else
-#             "red"
-#         )
-#         color_output.append(f'<span style="color:{color}">{current_word}</span>')
-
-#     output = " ".join(color_output)
-#     return output, word_tags
-
-
 # HF logging setup
 def setup_hf_dataset():
     global DATASET_CREATED
@@ -199,7 +152,13 @@ def infer_and_log(text_input):
         except Exception as e:
             print(f"Error uploading log: {e}")
 
-    return "".join(word_colors)
+    tokens = text_input.split()
+    formatted_output = " ".join(
+        f'<span style="color:{color}">{token}</span>' for token, color in zip(tokens, word_colors)
+    )
+    
+    return formatted_output
+
 
 def clear_fields():
     return "", ""
@@ -213,7 +172,7 @@ with gr.Blocks() as app:
 
     with gr.Row():
         input_box = gr.Textbox(label="Input Text", lines=10)
-        output_box = gr.Textbox(label="Output Text", lines=10, interactive=False)
+        output_box = gr.HTML(label="Output Text")
 
     with gr.Row():
         submit_btn = gr.Button("Submit")