GLiClassReranker / interfaces /initialize_models.py
BioMike's picture
fixed (#4)
9127367 verified
raw
history blame contribute delete
662 Bytes
import torch
import os
from sentence_transformers import CrossEncoder
from gliclass import GLiClassModel, ZeroShotClassificationPipeline
from transformers import AutoTokenizer
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model = GLiClassModel.from_pretrained(os.getenv("GLICLASS_MODEL_PATH")).eval().to(device)
tokenizer = AutoTokenizer.from_pretrained(os.getenv("GLICLASS_MODEL_PATH"))
multi_label_pipeline = ZeroShotClassificationPipeline(model, tokenizer, classification_type='multi-label',
device=device)
st = CrossEncoder("cross-encoder/ms-marco-MiniLM-L6-v2")