Update app.py

app.py

@@ -12,49 +12,76 @@ 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))
+
+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)
 
@@ -64,21 +91,18 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[plt.Fig
     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)}
+    all_entities = entities_bin + entities_ext
+    entity_counts = {entity: all_entities.count(entity) for entity in set(all_entities)}
     
-
-    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())
+    pie_labels = list(entity_counts.keys())
+    pie_sizes = list(entity_counts.values())
 
     fig1, ax1 = plt.subplots()
-    ax1.pie(pie_sizes_ext, labels=pie_labels_ext, autopct='%1.1f%%', startangle=90)
+    ax1.pie(pie_sizes, labels=pie_labels, autopct='%1.1f%%', startangle=90)
     ax1.axis('equal')
     
     fig2, ax2 = plt.subplots()
-    ax2.bar(pie_labels_bin, pie_sizes_bin)
+    ax2.bar(entity_counts.keys(), entity_counts.values())
     ax2.set_ylabel('Count')
     ax2.set_xlabel('Entity Type')
     ax2.set_title('Entity Counts')
@@ -97,13 +121,10 @@ def all(text: str):
             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 ..."),
@@ -138,4 +159,4 @@ iface = gr.Interface(
     theme=monochrome
 )
 
-iface.launch()
+iface.launch()