precision    recall  f1-score   support

         AGE       1.00      1.00      1.00       583
        CITY       1.00      1.00      1.00      5244
     DISEASE       0.99      1.00      1.00      6569
        DRUG       1.00      1.00      1.00      8220
      GENDER       1.00      1.00      1.00       664
  SPECIALITY       1.00      0.98      0.99      4207
      SUBWAY       1.00      1.00      1.00      1084
     SYMPTOM       1.00      1.00      1.00      8979

   micro avg       1.00      1.00      1.00     35550
   macro avg       1.00      1.00      1.00     35550
weighted avg       1.00      1.00      1.00     35550

from transformers import pipeline

pipe = pipeline(task="ner", model='Mykes/rubert_ner_SDDCS', tokenizer='Mykes/rubert_ner_SDDCS', aggregation_strategy="max")
# I made the misspelled words on purpose
query = "У ребенка треога и норушения сна, потеря сознания, раньше ставили паническое расстройство. Подскажи психиатра в районе метро Октбрьской."
pipe(query.lower())

[{'entity_group': 'AGE',
  'score': 0.99993,
  'word': 'ребенка',
  'start': 2,
  'end': 9},
 {'entity_group': 'SYMPTOM',
  'score': 0.9885457,
  'word': 'треога',
  'start': 10,
  'end': 16},
 {'entity_group': 'SYMPTOM',
  'score': 0.9934536,
  'word': 'норушения сна',
  'start': 19,
  'end': 32},
 {'entity_group': 'SYMPTOM',
  'score': 0.9999765,
  'word': 'потеря сознания',
  'start': 34,
  'end': 49},
 {'entity_group': 'DISEASE',
  'score': 0.999972,
  'word': 'паническое расстройство',
  'start': 66,
  'end': 89},
 {'entity_group': 'SPECIALITY',
  'score': 0.85958296,
  'word': 'психиатра',
  'start': 100,
  'end': 109},
 {'entity_group': 'SUBWAY',
  'score': 0.9955049,
  'word': 'октбрьской',
  'start': 125,
  'end': 135}]

import spacy
from spacy import displacy

def convert_to_displacy_format(text, ner_results):
    entities = []
    for result in ner_results:
        # Convert the Hugging Face output into the format displacy expects
        entities.append({
            "start": result['start'],
            "end": result['end'],
            "label": result['entity_group']
        })
    return {
        "text": text,
        "ents": entities,
        "title": None
    }
query = "У ребенка треога и норушения сна, потеря сознания, раньше ставили паническое расстройство, принимал атаракс. Подскажи хорошего психиатра в районе метро Октбрьской."
ner_results = pipe(query.lower())
displacy_data = convert_to_displacy_format(query, ner_results)
colors = {
    "SPECIALITY": "linear-gradient(90deg, #aa9cfc, #fc9ce7)",
    "CITY": "linear-gradient(90deg, #feca57, #ff9f43)",
    "DRUG": "linear-gradient(90deg, #55efc4, #81ecec)",
    "DISEASE": "linear-gradient(90deg, #fab1a0, #ff7675)",
    "SUBWAY": "linear-gradient(90deg, #00add0, #0039a6)",
    "AGE": "linear-gradient(90deg, #f39c12, #e67e22)",
    "SYMPTOM": "linear-gradient(90deg, #e74c3c, #c0392b)"
}
options = {"ents": ["SPECIALITY", "CITY", "DRUG", "DISEASE", "SYMPTOM", "AGE", "SUBWAY"], "colors": colors}
html = displacy.render(displacy_data, style="ent", manual=True, options=options, jupyter=False)
with open("ner_visualization_with_colors.html", "w", encoding="utf-8") as f:
    f.write(html)
from IPython.display import display, HTML
display(HTML(html))