Update app.py

app.py

@@ -2,15 +2,37 @@ import gradio as gr
 from transformers import pipeline
 
 # Load your model from Hugging Face
-checkpoint = "Amit234/distilbert-finetuned-med-ner/checkpoint-3924"  
+checkpoint = "Amit234/distilbert-finetuned-med-NER"  # Use the path to your Hugging Face model
 token_classifier = pipeline("token-classification", model=checkpoint, aggregation_strategy="simple")
 
-# Define prediction function
-def classify_text(text):
-    return token_classifier(text)
+def merge_tokens(entities):
+    merged_entities = []
+    current_entity = None
+
+    for entity in entities:
+        if current_entity is None:
+            current_entity = entity
+        elif "##" in entity['word']:
+            current_entity['word'] += entity['word'][2:]  # Remove "##" and concatenate the rest
+            current_entity['end'] = entity['end']
+        else:
+            merged_entities.append(current_entity)
+            current_entity = entity
+
+    if current_entity is not None:
+        merged_entities.append(current_entity)
+
+    return merged_entities
+
+
+
+def final_function(text):
+  entities = token_classifier(text)
+  merged_entities = merge_tokens(entities)
+  return merged_entities
 
 # Create Gradio interface
-interf = gr.Interface(fn=classify_text, inputs="text", outputs="json", title="NER for medical data")
+interf = gr.Interface(fn=final_function, inputs="text", outputs="json", title="NER for medical data")
 
 # Launch the interface
 if __name__ == "__main__":