Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 from transformers import pipeline
+import re
 
 # Load chatbot menggunakan InferenceClient
 client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3")
@@ -8,16 +9,20 @@ client = InferenceClient(model="mistralai/Mistral-7B-Instruct-v0.3")
 # Load model Named Entity Recognition (NER)
 ner_pipeline = pipeline("ner", model="d4data/biomedical-ner-all")
 
+# Entitas yang dianggap penting
+important_entities = {"Disease_disorder", "Sign_symptom", "Diagnostic_procedure", "Therapeutic_procedure", "Medication", "Dosage"}
+
 # Fungsi untuk ekstraksi entitas medis dari teks
 def extract_entities(text):
     entities = ner_pipeline(text)
-
+    
     merged_entities = []
     current_word = ""
     current_entity = None
-
+    
     for ent in entities:
         word = ent["word"]
+        entity_type = ent["entity"].split("-")[-1]  # Ambil tipe entitas tanpa B- atau I-
         
         # Gabungkan token subword dengan kata sebelumnya jika diawali "##"
         if word.startswith("#"):
@@ -27,16 +32,17 @@ def extract_entities(text):
                 merged_entities.append({"word": current_word, "entity": current_entity})  # Simpan kata sebelumnya
             
             current_word = word  # Mulai kata baru
-            current_entity = ent["entity"]  # Simpan tipe entity
-        
+            current_entity = entity_type  # Simpan tipe entity
+    
     if current_word and current_entity:  # Tambahkan kata terakhir yang sudah digabung
         merged_entities.append({"word": current_word, "entity": current_entity})
-
-    return merged_entities
-
+    
+    # Filter hanya entitas yang relevan
+    filtered_entities = [ent for ent in merged_entities if ent["entity"] in important_entities]
+    
+    return filtered_entities
 
 # Fungsi untuk highlight teks dan menampilkan daftar entitas yang dikenali
-import re
 def highlight_text(text, entities):
     entities = sorted(entities, key=lambda x: text.lower().find(x["word"].lower()), reverse=True)  # Urutkan dari belakang
 
@@ -59,8 +65,6 @@ def highlight_text(text, entities):
 
     return text + "<br><br>" + entity_list
 
-
-
 # Fungsi chatbot dengan NER
 def chat_with_ner(message, history):
     entities = extract_entities(message)