Improve prompt target intent matching (#51)

This commit is contained in:
Adil Hafeez 2024-09-16 19:20:07 -07:00 committed by GitHub
parent 8565462ec4
commit 9e50957f22
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
14 changed files with 461 additions and 415 deletions

View file

@ -1,6 +1,7 @@
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 = {}
@ -17,3 +18,11 @@ def load_ner_models(models = os.getenv("NER_MODELS", "urchade/gliner_large-v2.1"
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