removing output highlighting

app.py

@@ -1,30 +1,9 @@
 import gradio as gr
-import re
 from transformers import pipeline, AutoTokenizer
 # Load the Hugging Face model
 model_path = "patrixtano/mt5-base-anaphora_czech_6e"
 model_pipeline = pipeline("text2text-generation", model=model_path)
 tokenizer = AutoTokenizer.from_pretrained("patrixtano/mt5-base-anaphora_czech_6e")
-
-def split_sentence_with_tags(sentence):
-    # Regular expression to match <ana></ana>, <ant></ant>, or regular words
-    pattern = r'<ana>.*?</ana>|<ant>.*?</ant>|\S+'
-
-    # Find all matches in the sentence
-    words = re.findall(pattern, sentence)
-
-    # Create categories list
-    categories = []
-    for word in words:
-        if word.startswith('<ana>') and word.endswith('</ana>'):
-            categories.append("ANA")
-        elif word.startswith('<ant>') and word.endswith('</ant>'):
-            categories.append("ANT")
-        else:
-            categories.append("-")
-
-    return words, categories
-
 def predict(text_input):
     """
     Generate a prediction for the given input text using the Hugging Face model.
@@ -37,7 +16,7 @@ def predict(text_input):
     try:
         result = model_pipeline(text_input, **generation_parameters)
         # Extract and return the generated text
-        return split_sentence_with_tags(result[0]["generated_text"])
+        return result[0]["generated_text"]
     except Exception as e:
         return f"Error: {str(e)}"
 
@@ -49,7 +28,7 @@ examples = ["""Miluji ženu s vařečkou, <ana>která</ana> umí vařit.""",
 interface = gr.Interface(
     fn=predict,
     inputs=gr.Textbox(lines=5, label="Input Text"),
-    outputs=gr.HighlightedText(label="Model Output", show_legend=True, color_map={"ANA": "#f7a7a3", "ANT": "#87CEFF"}),
+    outputs=gr.Textbox(label="Model Output"),
     title="Anaphora resolution demo",
     description="""Enter text into the \"Input Text\" box, include <ana> </ana> tags around the anaphora 
 which is to be resolved. The model generates a copy of the text with <ant> </ant> tags marking the