app.py

from __future__ import annotations
from typing import Iterable, List, Dict, Tuple

import gradio as gr
from gradio.themes.base import Base
from gradio.themes.soft import Soft
from gradio.themes.monochrome import Monochrome
from gradio.themes.default import Default
from gradio.themes.utils import colors, fonts, sizes

import spaces
import torch
import os
import io
import re
import colorsys

import numpy as np

from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, pipeline

import matplotlib.pyplot as plt
import plotly.graph_objects as go
from wordcloud import WordCloud
from PIL import Image

ext_color_map = {
    "INTemothou": "#CD5C5C",
    "INTpercept": "#F08080", 
    "INTtime": "#FA8072", 
    "INTplace": "#E9967A", 
    "INTevent": "#8B0000", 
    "EXTsemantic": "#6495ED", 
    "EXTrepetition": "#7B68EE",
    "EXTother":  "#00008B",
}


def hex_to_rgb(hex_color: str) -> tuple[int, int, int]:
    hex_color = hex_color.lstrip('#')
    return tuple(int(hex_color[i:i+2], 16) for i in (0, 2, 4))

def rgb_to_hex(rgb_color: tuple[int, int, int]) -> str:
    return "#{:02x}{:02x}{:02x}".format(*rgb_color)

def adjust_brightness(rgb_color: tuple[int, int, int], factor: float) -> tuple[int, int, int]:
    hsv_color = colorsys.rgb_to_hsv(*[v / 255.0 for v in rgb_color])
    new_v = max(0, min(hsv_color[2] * factor, 1))
    new_rgb = colorsys.hsv_to_rgb(hsv_color[0], hsv_color[1], new_v)
    return tuple(int(v * 255) for v in new_rgb)

monochrome = Monochrome()

auth_token = os.environ['HF_TOKEN']

tokenizer_ext = AutoTokenizer.from_pretrained("AlGe/deberta-v3-large_AIS-token", token=auth_token)
model_ext = AutoModelForTokenClassification.from_pretrained("AlGe/deberta-v3-large_AIS-token", token=auth_token)
tokenizer_ext.model_max_length = 512
pipe_ext = pipeline("ner", model=model_ext, tokenizer=tokenizer_ext)

def process_ner(text: str, pipeline) -> dict:
    output = pipeline(text)
    entities = []
    current_entity = None
    for token in output:
        entity_type = token['entity'][2:]
        entity_prefix = token['entity'][:1]
        if current_entity is None or entity_type != current_entity['entity'] or (entity_prefix == 'B' and entity_type == current_entity['entity']):
            if current_entity is not None:
                entities.append(current_entity)
            current_entity = {
                "entity": entity_type,
                "start": token['start'],
                "end": token['end'],
                "scores": [token['score']],
                "tokens": [token['word']]
            }
        else:
            current_entity['end'] = token['end']
            current_entity['scores'].append(token['score'])
            current_entity['tokens'].append(token['word'])
    if current_entity is not None:
        entities.append(current_entity)
    for entity in entities:
        entity['average_score'] = sum(entity['scores']) / len(entity['scores'])
    return {"text": text, "entities": entities}

def generate_charts(ner_output_ext: dict) -> Tuple[go.Figure, np.ndarray]:
    entities_ext = [entity['entity'] for entity in ner_output_ext['entities']]
    entity_counts_ext = {entity: entities_ext.count(entity) for entity in set(entities_ext)}
    ext_labels = list(entity_counts_ext.keys())
    ext_sizes = list(entity_counts_ext.values())
    ext_colors = [ext_color_map.get(label, "#FFFFFF") for label in ext_labels]
    fig1 = go.Figure(data=[go.Pie(labels=ext_labels, values=ext_sizes, textinfo='label+percent', hole=.3, marker=dict(colors=ext_colors))])
    fig1.update_layout(
        template='plotly_dark',
        plot_bgcolor='rgba(0,0,0,0)',
        paper_bgcolor='rgba(0,0,0,0)'
    )
    wordcloud_image = generate_wordcloud(ner_output_ext['entities'], ext_color_map, "dh3.png")
    return fig1, wordcloud_image

def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str], file_path: str) -> np.ndarray:
    # Construct the absolute path
    base_path = os.path.dirname(os.path.abspath(__file__))
    image_path = os.path.join(base_path, file_path)
    if not os.path.exists(image_path):
        raise FileNotFoundError(f"Mask image file not found: {image_path}")

    mask_image = np.array(Image.open(image_path))
    mask_height, mask_width = mask_image.shape[:2]

    word_details = []

    for entity in entities:
        for token in entity['tokens']:
            # Process each token
            token_text = token.replace("▁", " ").strip()
            if token_text:  # Ensure token is not empty
                word_details.append({
                    'text': token_text,
                    'score': entity.get('average_score', 0.5),
                    'entity': entity['entity']
                })

    # Calculate word frequency weighted by score
    word_freq = {}
    for detail in word_details:
        if detail['text'] in word_freq:
            word_freq[detail['text']]['score'] += detail['score']
            word_freq[detail['text']]['count'] += 1
        else:
            word_freq[detail['text']] = {'score': detail['score'], 'count': 1, 'entity': detail['entity']}

    # Average the scores and prepare final frequency dictionary
    final_word_freq = {word: details['score'] / details['count'] for word, details in word_freq.items()}

    # Prepare entity type mapping for color function
    word_to_entity = {word: details['entity'] for word, details in word_freq.items()}

    def color_func(word, font_size, position, orientation, random_state=None, **kwargs):
        entity_type = word_to_entity.get(word, None)
        return color_map.get(entity_type, "#FFFFFF")

    wordcloud = WordCloud(width=mask_width, height=mask_height, background_color='#121212', mask=mask_image, color_func=color_func).generate_from_frequencies(final_word_freq)

    plt.figure(figsize=(mask_width/100, mask_height/100))
    plt.imshow(wordcloud, interpolation='bilinear')
    plt.axis('off')
    plt.tight_layout(pad=0)

    plt_image = plt.gcf()
    plt_image.canvas.draw()
    image_array = np.frombuffer(plt_image.canvas.tostring_rgb(), dtype=np.uint8)
    image_array = image_array.reshape(plt_image.canvas.get_width_height()[::-1] + (3,))
    plt.close()

    return image_array

@spaces.GPU
def all(text: str):
    ner_output_ext = process_ner(text, pipe_ext)
    
    pie_chart, wordcloud_image = generate_charts(ner_output_ext)
    
    return (ner_output_ext, pie_chart, wordcloud_image)

examples = [
    ['Bevor ich meinen Hund kaufte bin ich immer alleine durch den Park gelaufen. Gestern war ich aber mit dem Hund losgelaufen. Das Wetter war sehr schön, nicht wie sonst im Winter. Ich weiß nicht genau. Mir fällt sonst nichts dazu ein. Wir trafen auf mehrere Spaziergänger. Ein Mann mit seinem Kind. Das Kind hat ein Eis gegessen.'],
]

iface = gr.Interface(
    fn=all,
    inputs=gr.Textbox(lines=5, label="Input Text", placeholder="Write about how your breakfast went or anything else that happened or might happen to you ..."),
    outputs=[
        gr.HighlightedText(label="Extended Sequence Classification",
                           color_map = ext_color_map
                          ),
        gr.Plot(label="Extended SeqClass Entity Distribution Pie Chart"),
        gr.Image(label="Entity Word Cloud")
    ],
    title="Scoring Demo",
    description="Autobiographical Memory Analysis: This demo combines two text - and two sequence classification models to showcase our automated Autobiographical Interview scoring method. Submit a narrative to see the results.",
    examples=examples,
    theme=monochrome
)

iface.launch()