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
from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification, pipeline
import os

import colorsys
import matplotlib.pyplot as plt

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))
    return tuple(int(v * 255) for v in new_rgb)

monochrome = Monochrome()

auth_token = os.environ['HF_TOKEN']


tokenizer_bin = AutoTokenizer.from_pretrained("AlGe/deberta-v3-large_token", token=auth_token)
model_bin = AutoModelForTokenClassification.from_pretrained("AlGe/deberta-v3-large_token", token=auth_token)
tokenizer_bin.model_max_length = 512
pipe_bin = pipeline("ner", model=model_bin, tokenizer=tokenizer_bin)


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)


model1 = AutoModelForSequenceClassification.from_pretrained("AlGe/deberta-v3-large_Int_segment", num_labels=1, token=auth_token)
tokenizer1 = AutoTokenizer.from_pretrained("AlGe/deberta-v3-large_Int_segment", token=auth_token)

model2 = AutoModelForSequenceClassification.from_pretrained("AlGe/deberta-v3-large_seq_ext", num_labels=1, token=auth_token)

def process_ner(text: str, pipeline) -> dict:

    output = pipeline(text)
    entities = []
    current_entity = None

    return {"text": text, "entities": entities}

def process_classification(text: str, model1, model2, tokenizer1) -> Tuple[str, str, str]:
    inputs1 = tokenizer1(text, max_length=512, return_tensors='pt', truncation=True, padding=True)
    
    with torch.no_grad():
    prediction2 = outputs2[0].item()
    score = prediction1 / (prediction2 + prediction1)

    return f"{round(prediction1, 1)}", f"{round(prediction2, 1)}", f"{round(score, 2)}"
    
def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[plt.Figure, plt.Figure]:
    entities_bin = [entity['entity'] for entity in ner_output_bin['entities']]
    entities_ext = [entity['entity'] for entity in ner_output_ext['entities']]
    
    entity_counts_bin = {entity: entities_bin.count(entity) for entity in set(entities_bin)}
    entity_counts_ext = {entity: entities_ext.count(entity) for entity in set(entities_ext)}
    

    pie_labels_bin = list(entity_counts_bin.keys())
    pie_sizes_bin = list(entity_counts_bin.values())
    pie_labels_ext = list(entity_counts_ext.keys())
    pie_sizes_ext = list(entity_counts_ext.values())

    fig1, ax1 = plt.subplots()
    ax1.pie(pie_sizes_ext, labels=pie_labels_ext, autopct='%1.1f%%', startangle=90)
    ax1.axis('equal')
    
    fig2, ax2 = plt.subplots()
    ax2.bar(pie_labels_bin, pie_sizes_bin)
    ax2.set_ylabel('Count')
    ax2.set_xlabel('Entity Type')
    ax2.set_title('Entity Counts')
    
    return fig1, fig2

@spaces.GPU
def all(text: str):
    ner_output_bin = process_ner(text, pipe_bin)
    ner_output_ext = process_ner(text, pipe_ext)
    classification_output = process_classification(text, model1, model2, tokenizer1)
    
    pie_chart, bar_chart = generate_charts(ner_output_bin, ner_output_ext)
    
    return (ner_output_bin, ner_output_ext, 
            classification_output[0], classification_output[1], classification_output[2],
            pie_chart, bar_chart)



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="Binary Sequence Classification",
                           color_map={
                               "External": "#6ad5bcff",
                               "Internal": "#ee8bacff"}
                          ),
        gr.HighlightedText(label="Extended Sequence Classification",
                           color_map={
                               "INTemothou": "#FF7F50",  # Coral
                               "INTpercept": "#FF4500",  # OrangeRed
                               "INTtime": "#FF6347",     # Tomato
                               "INTplace": "#FFD700",    # Gold
                               "INTevent": "#FFA500",    # Orange
                        
                               "EXTsemantic": "#4682B4", # SteelBlue
                               "EXTrepetition": "#5F9EA0", # CadetBlue
                               "EXTother": "#00CED1",    # DarkTurquoise
                           }
                          ),
        gr.Label(label="Internal Detail Count"),
        gr.Label(label="External Detail Count"),
        gr.Label(label="Approximated Internal Detail Ratio"),
        gr.Plot(label="Entity Distribution Pie Chart"),
        gr.Plot(label="Entity Count Bar Chart")
    ],
    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()