app.py

import gradio as gr
from transformers import pipeline

ner_pipeline = pipeline("ner", model="LendeaViva/bertimbau_base_pos_2")

def classify_tokens(text):
    results = ner_pipeline(text)

    pun_words = []
    pun_tokens_positions = []
    pun_save = ''
    found_pun = False

    for index, entity in enumerate(results):
        word = entity['word']
        entity_type = entity['entity']

        if entity_type == "Pun" and not word.startswith('##'):
            found_pun = True
            pun_save = word
            pun_tokens_positions.append(index)
        elif word.startswith('##') and results[index - 1]['entity'] == 'Pun':
            pun_save += word.lstrip('##')
            pun_tokens_positions.append(index)
            if index == len(results) - 1:
                pun_words.append(pun_save)
        elif found_pun:
            pun_words.append(pun_save)
            found_pun = False
            pun_save = ''

    pun_set = set(pun_words)

    highlighted_text = text
    for pun_word in pun_set:
        highlighted_text = highlighted_text.replace(pun_word, f'<span style="background-color: #FF79C6; font-weight: bold;">{pun_word}</span>')

    return highlighted_text


demo = gr.Interface(
    fn=classify_tokens,
    inputs=gr.Textbox(label="Input text", placeholder="Enter your pun here..."),
    outputs=gr.HTML(label="Highlighted Text"),
    title="Pun Location in Portuguese",
    description="This app highlights pun words in the input text.",
    theme='freddyaboulton/dracula_revamped',
    allow_flagging="never",
)

demo.launch()