2024-07-18 14:04:51 -07:00
|
|
|
import os
|
|
|
|
|
import sentence_transformers
|
2024-07-30 16:23:23 -07:00
|
|
|
from gliner import GLiNER
|
2024-09-16 19:20:07 -07:00
|
|
|
from transformers import pipeline
|
2024-07-18 14:04:51 -07:00
|
|
|
|
2024-08-06 23:40:06 -07:00
|
|
|
def load_transformers(models = os.getenv("MODELS", "BAAI/bge-large-en-v1.5")):
|
2024-07-18 14:04:51 -07:00
|
|
|
transformers = {}
|
|
|
|
|
|
|
|
|
|
for model in models.split(','):
|
|
|
|
|
transformers[model] = sentence_transformers.SentenceTransformer(model)
|
|
|
|
|
|
|
|
|
|
return transformers
|
2024-07-30 16:23:23 -07:00
|
|
|
|
|
|
|
|
def load_ner_models(models = os.getenv("NER_MODELS", "urchade/gliner_large-v2.1")):
|
|
|
|
|
ner_models = {}
|
|
|
|
|
|
|
|
|
|
for model in models.split(','):
|
|
|
|
|
ner_models[model] = GLiNER.from_pretrained(model)
|
|
|
|
|
|
|
|
|
|
return ner_models
|
2024-09-16 19:20:07 -07:00
|
|
|
|
|
|
|
|
def load_zero_shot_models(models = os.getenv("ZERO_SHOT_MODELS", "tasksource/deberta-base-long-nli")):
|
|
|
|
|
zero_shot_models = {}
|
|
|
|
|
|
|
|
|
|
for model in models.split(','):
|
|
|
|
|
zero_shot_models[model] = pipeline("zero-shot-classification",model=model)
|
|
|
|
|
|
|
|
|
|
return zero_shot_models
|