Update app.py

app.py

@@ -40,21 +40,11 @@ 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 = []
@@ -84,39 +74,14 @@ 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]:
-    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
@@ -128,7 +93,6 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
         "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
@@ -139,20 +103,11 @@ def generate_charts(ner_output_bin: dict, ner_output_ext: dict) -> Tuple[go.Figu
         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(
-        xaxis_title='Entity Type',
-        yaxis_title='Count',
-        template='plotly_dark',
-        plot_bgcolor='rgba(0,0,0,0)',
-        paper_bgcolor='rgba(0,0,0,0)'
-    )
     
     # Generate word cloud
     wordcloud_image = generate_wordcloud(ner_output_ext['entities'], ext_color_map)
     
-    return fig1, fig2, wordcloud_image
+    return fig1, wordcloud_image
 
 
 def generate_wordcloud(entities: List[Dict], color_map: Dict[str, str]) -> np.ndarray:
@@ -192,15 +147,11 @@ 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)
+    pie_chart, wordcloud_image = 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, wordcloud_image)
+    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.'],
@@ -210,11 +161,6 @@ 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
@@ -227,11 +173,7 @@ iface = gr.Interface(
                                "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")
     ],
     title="Scoring Demo",