from transformers import pipeline
from transformers import DistilBertForTokenClassification, DistilBertTokenizer
DRIVE_BASE_PATH = "model/"
model_path = f"{DRIVE_BASE_PATH}"

model = DistilBertForTokenClassification.from_pretrained(model_path)
tokenizer = DistilBertTokenizer.from_pretrained(model_path)

ner_pipeline = pipeline("ner", model=model,tokenizer=tokenizer, aggregation_strategy='simple')

def predict_ner(input_text):
    # Make predictions using the NER pipeline
    results = ner_pipeline(input_text)

    # Extract relevant information for Gradio output
    entities_all = [(result["word"], result.get("entity", None)) for result in results]

    return entities_all