interfaces/classification.py

from gliner import GLiNER
import gradio as gr

model = GLiNER.from_pretrained("knowledgator/gliner-multitask-v1.0").to("cpu")

PROMPT_TEMPLATE = """Classify the given text having the following classes: {}"""
classification_examples = [
    [
        """
"I recently purchased the Sony WH-1000XM4 Wireless Noise-Canceling Headphones from Amazon and I must say, I'm thoroughly impressed. The package arrived in New York within 2 days, thanks to Amazon Prime's expedited shipping.

The headphones themselves are remarkable. The noise-canceling feature works like a charm in the bustling city environment, and the 30-hour battery life means I don't have to charge them every day. Connecting them to my Samsung Galaxy S21 was a breeze, and the sound quality is second to none.

I also appreciated the customer service from Amazon when I had a question about the warranty. They responded within an hour and provided all the information I needed.

However, the headphones did not come with a hard case, which was listed in the product description. I contacted Amazon, and they offered a 10% discount on my next purchase as an apology.

Overall, I'd give these headphones a 4.5/5 rating and highly recommend them to anyone looking for top-notch quality in both product and service.""",
        "positive review, negative review, neutral review",
        0.5
    ],
    [
        "I really enjoyed the pizza we had for dinner last night.",
        "Food, Weather, Sports",
        0.5
    ],
    [
        "Das Kind spielt im Park und genießt die frische Luft.",
        "Nature, Technology, Politics",
        0.5
    ],
        [
        """
"Last night, we visited the new Italian restaurant downtown. The Margherita pizza was absolutely delightful, with a perfectly crisp crust and fresh basil. 

However, the service was slow; it took over 20 minutes to take our order. The pasta arrived lukewarm, which was disappointing given the hype around this place. 

On the bright side, the ambiance was cozy, and the wine selection was impressive. Overall, it was a mixed experience, but I might give it another try on a quieter evening."
        """,
        "Food Quality, Technology, Politics",
        0.5
    ],
    [
        """
"Das Kind verbrachte den Nachmittag im Park und entdeckte einen kleinen Teich mit Enten. Es war wunderschön zu sehen, wie es die Natur erkundete. 

Doch plötzlich störten die Geräusche einer Baustelle die ruhige Atmosphäre. Trotzdem spielte das Kind weiter, und ich genoss die frische Luft.

Solche Momente zeigen, wie wichtig es ist, Kinder in der Natur aufwachsen zu lassen."
        """,
        "Outdoor Activities, Gaming, Artificial Intelligence",
        0.5
    ],
    [
        """
"I recently attended a healthcare technology conference. The keynote speaker demonstrated how AI is revolutionizing diagnostics, making it possible to detect rare diseases with incredible accuracy.

However, concerns about data privacy and ethical implications were also heavily discussed. Despite these challenges, the energy in the room was palpable as experts envisioned a future where AI saves millions of lives.

It was an inspiring event that showcased the potential of combining technology and healthcare innovation."
        """,
        "Artificial Intelligence, Music, Sports",
        0.5
    ]
]

def prepare_prompts(text, labels):
    labels_str = ', '.join(labels)
    return PROMPT_TEMPLATE.format(labels_str) + "\n" + text

def process(text, labels, threshold):
    if not text.strip() or not labels.strip():
        return {"text": text, "entities": []}
    
    labels = [label.strip() for label in labels.split(",")]
    prompt = prepare_prompts(text, labels)
    
    predictions = model.run([prompt], ["match"], threshold=threshold)
    entities = []

    if predictions and predictions[0]:
        for pred in predictions[0]:
            entities.append({
                "entity": "match",
                "word": pred["text"],
                "start": pred["start"],
                "end": pred["end"],
                "score": pred["score"]
            })
    
    return {"text": prompt, "entities": entities}

with gr.Blocks(title="Text Classification with Highlighted Labels") as classification_interface:
    gr.Markdown("# Text Classification with Highlighted Labels")
    
    input_text = gr.Textbox(label="Input Text", placeholder="Enter text for classification")
    input_labels = gr.Textbox(label="Labels (Comma-Separated)", placeholder="Enter labels separated by commas (e.g., Positive, Negative, Neutral)")
    threshold = gr.Slider(0, 1, value=0.5, step=0.01, label="Threshold")
    
    output = gr.HighlightedText(label="Classification Results")
    
    submit_btn = gr.Button("Classify")
    
    examples = gr.Examples(
        examples=classification_examples,
        inputs=[input_text, input_labels, threshold],
        outputs=output,
        fn=process,
        cache_examples=True
    )
    theme=gr.themes.Base()


    input_text.submit(fn=process, inputs=[input_text, input_labels, threshold], outputs=output)
    threshold.release(fn=process, inputs=[input_text, input_labels, threshold], outputs=output)
    submit_btn.click(fn=process, inputs=[input_text, input_labels, threshold], outputs=output)

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