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 import plotly.graph_objects as go from typing import Tuple import plotly.io as pio from wordcloud import WordCloud import io 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_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 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'], "score": token['score'] } else: current_entity['end'] = token['end'] current_entity['score'] = max(current_entity['score'], token['score']) if current_entity is not None: entities.append(current_entity) 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(): outputs1 = model1(**inputs1) outputs2 = model2(**inputs1) prediction1 = outputs1[0].item() prediction2 = outputs2[0].item() score = prediction1 / (prediction2 + prediction1) return f"{round(prediction1, 1)}", f"{round(prediction2, 1)}", f"{round(score, 2)}" import plotly.graph_objects as go from typing import Tuple def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figure, go.Figure]: entities_bin = [entity['entity'] for entity in ner_output_bin['entities']] entities_ext = [entity['entity'] for entity in ner_output_ext['entities']] # Counting entities for binary classification entity_counts_bin = {entity: entities_bin.count(entity) for entity in set(entities_bin)} bin_labels = list(entity_counts_bin.keys()) bin_sizes = list(entity_counts_bin.values()) # Counting entities for extended classification 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()) # Define color mapping bin_color_map = { "External": "#6ad5bc", "Internal": "#ee8bac" } ext_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 } bin_colors = [bin_color_map[label] for label in bin_labels] ext_colors = [ext_color_map[label] for label in ext_labels] # Create pie chart for extended classification 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( #title_text='Extended Sequence Classification Subclasses', template='plotly_dark', plot_bgcolor='rgba(0,0,0,0)', paper_bgcolor='rgba(0,0,0,0)' ) # Create bar chart for binary classification fig2 = go.Figure(data=[go.Bar(x=bin_labels, y=bin_sizes, marker=dict(color=bin_colors))]) fig2.update_layout( #title='Binary Sequence Classification Classes', xaxis_title='Entity Type', yaxis_title='Count', template='plotly_dark', plot_bgcolor='rgba(0,0,0,0)', paper_bgcolor='rgba(0,0,0,0)' ) # Create word cloud for extended classification entities wordcloud_fig = generate_wordcloud(ner_output_ext['entities'], 'tab10') return fig1, fig2, wordcloud_fig def generate_wordcloud(entities: list, color_map: dict) -> plt.Figure: word_freq = {entity['entity']: entity['score'] for entity in entities} wc = WordCloud(width=800, height=400, background_color='black', colormap=color_map).generate_from_frequencies(word_freq) fig, ax = plt.subplots(figsize=(10, 5)) ax.imshow(wc, interpolation='bilinear') ax.axis('off') return fig @spaces.GPU def all(text: str) -> Tuple[gr.HighlightedText, gr.HighlightedText, int, int, float, go.Figure, go.Figure, plt.Figure]: ner_output_bin = process_ner(text, pipe_bin) ner_output_ext = process_ner(text, pipe_ext) binary_entities = [{"entity": ent["entity"], "score": ent["score"]} for ent in ner_output_bin["entities"]] extended_entities = [{"entity": ent["entity"], "score": ent["score"]} for ent in ner_output_ext["entities"]] bin_text = [{"text": text[ent["start"]:ent["end"]], "entity": ent["entity"]} for ent in ner_output_bin["entities"]] ext_text = [{"text": text[ent["start"]:ent["end"]], "entity": ent["entity"]} for ent in ner_output_ext["entities"]] fig1, fig2, wordcloud_fig = generate_charts(ner_output_bin, ner_output_ext) # Convert the word cloud figure to an image for Gradio buf = io.BytesIO() wordcloud_fig.savefig(buf, format='png') buf.seek(0) internal_count = sum(1 for ent in ner_output_bin["entities"] if ent["entity"] == "Internal") external_count = sum(1 for ent in ner_output_bin["entities"] if ent["entity"] == "External") internal_ratio = internal_count / (internal_count + external_count) if (internal_count + external_count) > 0 else 0 return bin_text, ext_text, internal_count, external_count, internal_ratio, fig1, fig2, buf 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="Extended SeqClass Entity Distribution Pie Chart"), gr.Plot(label="Binary SeqClass Entity Count Bar Chart"), gr.Image(label="Word Cloud of Extended Entities") ], 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()