plano/model_server/app/load_models.py
Adil Hafeez 060a0d665e
improve service names (#54)
- embedding-server => model_server
- public-types => public_types
- chatbot-ui => chatbot_ui
- function-calling => function_calling
2024-09-17 08:47:35 -07:00

28 lines
844 B
Python

import os
import sentence_transformers
from gliner import GLiNER
from transformers import pipeline
def load_transformers(models = os.getenv("MODELS", "BAAI/bge-large-en-v1.5")):
transformers = {}
for model in models.split(','):
transformers[model] = sentence_transformers.SentenceTransformer(model)
return transformers
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
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