app.py

import gradio as gr
from transformers import pipeline

# Load your model from Hugging Face
checkpoint = "Amit234/distilbert-finetuned-med-NER"  # Use the path to your Hugging Face model
token_classifier = pipeline("token-classification", model=checkpoint, aggregation_strategy="simple")

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=final_function, inputs="text", outputs="json", title="NER for medical data")

# Launch the interface
if __name__ == "__main__":
    interf.launch(inline=False)