app.py

import gradio as gr
import spacy
import spacy_transformers  # Ensure transformer components register
from huggingface_hub import snapshot_download
import os

# Download the model from Hugging Face Hub and load it.
model_repo = "hksung/ASC_tagger_v2"
model_path = snapshot_download(model_repo)  # Assumes the repo is public; add token if needed
nlp = spacy.load(os.path.join(model_path, 'model-best'))

# If sentence segmentation is missing, add a sentencizer.
if 'parser' not in nlp.pipe_names and 'senter' not in nlp.pipe_names:
    nlp.add_pipe('sentencizer')

# Define the inline color mapping for each ASC tag.
color_scheme = {
    "ATTR": "#f6e9e9",      # Light Red
    "INTRAN_S": "#e7957f",   # Light Blue
    "INTRAN_MOT": "#f6f6aa", # Light Green
    "INTRAN_RES": "#f0aed6", # Light Yellow
    "CAUS_MOT": "#cdefb5",   # Light Pink
    "TRAN_S": "#a0d4f7",     # Light Blue
    "TRAN_RES": "#c7aefa",   # Light Purple
    "DITRAN": "#b3f0f7",     # Light Teal
    "PASSIVE": "#c3c0c0"     # Light Gray
}

def get_highlighted_text(doc):
    """
    Wraps detected ASCs in each sentence with a span tag containing an inline style.
    In addition, appends the entity's tag name in an extra span so the label is always visible.
    """
    highlighted_sentences = []
    for sent in doc.sents:
        s = sent.text
        # Find all entities fully contained within this sentence.
        ents_in_sent = [ent for ent in doc.ents if ent.start_char >= sent.start_char and ent.end_char <= sent.end_char]
        if ents_in_sent:
            # Process entities in reverse order so the character offsets remain valid.
            ents_in_sent = sorted(ents_in_sent, key=lambda x: x.start_char, reverse=True)
            for ent in ents_in_sent:
                ent_start = ent.start_char - sent.start_char
                ent_end = ent.end_char - sent.start_char
                # Get the color based on the tag; default to light gray if not found.
                color = color_scheme.get(ent.label_, "#cccccc")
                # Wrap the entity text with an inline style and add the tag label in a separate span.
                s = (
                    s[:ent_start]
                    + f'<span style="background-color: {color}; font-weight: bold;" title="{ent.label_}">'
                    + s[ent_start:ent_end]
                    + f'<span style="display:inline-block; vertical-align: super; font-size: 0.8em; margin-left: 2px; color: black;">{ent.label_}</span>'
                    + '</span>'
                    + s[ent_end:]
                )
        highlighted_sentences.append(s)
    # Combine sentences with HTML breaks.
    return "<br><br>".join(highlighted_sentences)

def process_text(input_text):
    """
    Process input text to detect and tag ASCs.
    Returns an HTML string with inline styles applied to each entity.
    """
    if not input_text.strip():
        return "No text provided. Please enter some text."
    
    doc = nlp(input_text)
    
    if not list(doc.sents):
        return "Please enter at least one sentence."
    
    if not doc.ents:
        return "No ASCs were detected."
    
    return get_highlighted_text(doc)

# Define the sample text to be injected.
sample_text = (
    "Once upon a time in a small, enchanted village nestled between towering mountains, there lived a baker named Oliver. "
    "But this wasn't just any baker—Oliver had a magical secret. Every loaf of bread he baked had the power to grant a wish, "
    "but only if it was made with true joy in his heart. The villagers adored Oliver’s bread, though they had no idea about its magic. "
    "They just knew it tasted unlike anything else—warm, soft, and with a hint of something... extraordinary. Only a few people had ever discovered "
    "its secret, and even fewer had made wishes, because the key was in asking for something truly selfless. "
    "One crisp autumn morning, a traveling musician named Lyla passed through the village. She had a worn-out lute, tattered clothes, "
    "and a bright spark of hope in her eyes despite her hardships. Drawn by the smell of freshly baked bread, she wandered into Oliver’s bakery. "
    "Oliver, busy dusting flour off his hands, smiled warmly. “What can I do for you?” he asked. "
    "Lyla’s stomach growled in response, and she laughed, a sound as melodic as her lute playing. “Just one loaf of bread, please,” she said, "
    "placing a single coin on the counter, all she had left. "
    "Oliver, feeling the joy in the simple act of serving someone in need, baked the loaf with care. As the bread came out of the oven, a golden glow "
    "shimmered around it for just a moment. Lyla didn’t notice, but Oliver did, and he smiled to himself. "
    "As Lyla took her first bite outside the shop, she felt warmth spread through her whole being. She thought of how she missed playing her lute to "
    "cheer others up, but it was broken and beyond repair. Without knowing the magic at work, she whispered under her breath, 'I just wish I could bring "
    "music to everyone again.' "
    "That evening, as Lyla set up by the village fountain to play, something miraculous happened. Her old, broken lute began to transform, glowing just "
    "like the bread had. In her hands now was the most beautiful lute she had ever seen, shimmering with gold and silver strings. When she strummed it, "
    "the music that flowed out wasn’t just sound—it was pure joy. "
    "The village gathered around her, drawn by the enchanting melodies that seemed to fill their hearts with warmth and happiness. People laughed, danced, "
    "and even sang along, forgetting their worries. "
    "Oliver watched from his bakery window, smiling to himself. He had no idea what Lyla had wished for, but he could see its magic at work. He knew his "
    "bread had once again brought a little bit of magic to the world. "
    "And so, every day after that, Lyla would play her magical lute in the village square, bringing joy to everyone who heard her music. People came from "
    "distant lands just to hear her play, and with every note, the village grew a little brighter, a little happier. "
    "As for Oliver, he continued to bake with joy in his heart, knowing that sometimes the smallest acts of kindness—like baking a loaf of bread—could "
    "make the biggest difference in someone’s life. "
    "And so, the village became known not just for its delicious bread, but for its music, magic, and the happiness that lingered in every corner. "
    "(Generated by ChatGPT 4o 🤖)"
)

def fill_sample_text():
    return sample_text

# Build the Gradio interface.
with gr.Blocks() as demo:
    gr.Markdown("# ASC tagger demo")
    gr.Markdown("Enter some text to have ASCs tagged with a custom inline color scheme. Use the button below to fill in sample text.")
    
    input_textbox = gr.Textbox(lines=10, label="Input Text", placeholder="Enter text here...")
    output_html = gr.HTML(label="Tagged Text")
    tag_btn = gr.Button("Tag ASCs")
    fill_btn = gr.Button("Fill Sample Text")

    tag_btn.click(fn=process_text, inputs=input_textbox, outputs=output_html)
    # When "Fill Sample Text" is clicked, populate the textbox.
    fill_btn.click(fn=fill_sample_text, inputs=[], outputs=input_textbox)
    # When "Tag ASCs" is clicked, process the input text.

if __name__ == "__main__":
    demo.launch()