# https://blog.knowledgator.com/meet-the-new-zero-shot-ner-architecture-30ffc2cb1ee0
# https://labelstud.io/blog/fine-tuning-generalist-models-for-named-entity-recognition/
import psutil
from gliner import GLiNER
import streamlit as st

LABELS = ["eventTitle", "eventLocation", "date", "time", "street", "city"]

class GlinerHandler:
    def __init__(self, model_name="urchade/gliner_multi-v2.1"):
        self.model = GLiNER.from_pretrained(model_name)
        st.info("Using NER Model Gliner")

    def extract_entities(self, text, labels=None, threshold=0.3):
        if labels is None:
            labels = LABELS
        entities = self.model.predict_entities(text, labels, threshold=threshold)
        return entities

# gliner = GlinerHandler()
# entities = gliner.extract_entities("Test 20.03.10, In Nürnberg")
# print(entities)