Update app.py

app.py

@@ -45,16 +45,6 @@ model_bin = AutoModelForTokenClassification.from_pretrained("AlGe/deberta-v3-lar
 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 = []
@@ -84,60 +74,22 @@ def process_ner(text: str, pipeline) -> dict:
 
     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)}"
-
-
-def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figure, go.Figure, np.ndarray]:
+def generate_charts(ner_output_bin: dict) -> Tuple[go.Figure, np.ndarray]:
     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())
+
 
     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.get(label, "#FFFFFF") for label in bin_labels]
-    ext_colors = [ext_color_map.get(label, "#FFFFFF") 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(
-        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))])
@@ -150,9 +102,9 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
     )
     
     # Generate word cloud
-    wordcloud_image = generate_wordcloud(ner_output_ext['entities'], ext_color_map)
+    wordcloud_image = generate_wordcloud(ner_output_bin['entities'], bin_color_map)
     
-    return fig1, fig2, wordcloud_image
+    return fig2, wordcloud_image
 
 
 def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.ndarray:
@@ -193,14 +145,10 @@ def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.nd
 @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, wordcloud_image = generate_charts(ner_output_bin, ner_output_ext)
+
+    bar_chart, wordcloud_image = generate_charts(ner_output_bin)
     
-    return (ner_output_bin, ner_output_ext, 
-            classification_output[0], classification_output[1], classification_output[2],
-            pie_chart, bar_chart, wordcloud_image)
+    return (ner_output_bin, bar_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.'],
@@ -215,22 +163,6 @@ iface = gr.Interface(
                                "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="Entity Word Cloud")
     ],